From 636ca1c940b7077227435c21d6fcd7c41c071d90 Mon Sep 17 00:00:00 2001 From: Sabaun Taraki Date: Mon, 19 Aug 2024 14:04:50 +0300 Subject: [PATCH] Introduce reimbursement --- core-processor/src/ext.rs | 139 +++++++++++++++++++++++++++++++++++--- core/src/gas.rs | 11 ++- core/src/reservation.rs | 33 ++++++--- 3 files changed, 158 insertions(+), 25 deletions(-) diff --git a/core-processor/src/ext.rs b/core-processor/src/ext.rs index 3c0aec6621a..34572aa36d1 100644 --- a/core-processor/src/ext.rs +++ b/core-processor/src/ext.rs @@ -1199,15 +1199,37 @@ impl Externalities for Ext { } fn unreserve_gas(&mut self, id: ReservationId) -> Result { - let amount = self.context.gas_reserver.unreserve(id)?; + let (amount, reimburse) = self.context.gas_reserver.unreserve(id)?; - // This statement is like an op that increases "left" counter, but do not affect "burned" counter, - // because we don't actually refund, we just rise "left" counter during unreserve - // and it won't affect gas allowance counter because we don't make any actual calculations - // TODO: uncomment when unreserving in current message features is discussed - /*if !self.context.gas_counter.increase(amount) { - return Err(some_charge_error.into()); - }*/ + if let Some(reimburse) = reimburse { + let current_gas_amount = self.gas_amount(); + + // Basically amount of the reseravtion and the cost for the hold duration. + let reimburse_amount = self + .context + .costs + .rent + .reservation + .cost_for( + self.context + .reserve_for + .saturating_add(reimburse.duration()) + .into(), + ) + .saturating_add(amount); + self.context + .gas_counter + .increase(reimburse_amount, reimburse) + .then_some(()) + .unwrap_or_else(|| { + unreachable!( + "Ext::unreserve_gas: failed to reimburse unreserved gas to left counter. \ + Current gas amount - {}, reimburse amount - {}", + current_gas_amount.left(), + amount, + ) + }); + } Ok(amount) } @@ -1409,11 +1431,12 @@ impl Externalities for Ext { #[cfg(test)] mod tests { use super::*; + use crate::configs::RentCosts; use alloc::vec; use gear_core::{ - costs::SyscallCosts, + costs::{CostOf, SyscallCosts}, message::{ContextSettings, IncomingDispatch, Payload, MAX_PAYLOAD_SIZE}, - reservation::GasReservationState, + reservation::{GasReservationMap, GasReservationSlot, GasReservationState}, }; struct MessageContextBuilder { @@ -1501,6 +1524,12 @@ mod tests { self } + + fn with_reservations_map(mut self, map: GasReservationMap) -> Self { + self.0.gas_reserver = GasReserver::new(&Default::default(), map, 256); + + self + } } // Invariant: Refund never occurs in `free` call. @@ -2182,4 +2211,94 @@ mod tests { // Sending value below ED is also fine assert!(msg.is_ok()); } + + #[test] + fn test_unreserve_no_reimbursement() { + let costs = ExtCosts { + rent: RentCosts { + reservation: CostOf::new(10), + ..Default::default() + }, + ..Default::default() + }; + + // Create "pre-reservation". + let (id, gas_reservation_map) = { + let mut m = BTreeMap::new(); + let id = ReservationId::generate(MessageId::new([5; 32]), 10); + + m.insert( + id, + GasReservationSlot { + amount: 1_000_000, + start: 0, + finish: 10, + }, + ); + + (id, m) + }; + let mut ext = Ext::new( + ProcessorContextBuilder::new() + .with_gas(GasCounter::new(u64::MAX)) + .with_message_context(MessageContextBuilder::new().build()) + .with_existential_deposit(500) + .with_reservations_map(gas_reservation_map) + .with_costs(costs.clone()) + .build(), + ); + + // Check all the reseravations are in "existing" state. + assert!(ext + .context + .gas_reserver + .states() + .iter() + .all(|(_, state)| matches!(state, GasReservationState::Exists { .. }))); + + // Unreserving existing and checking no gas reimbursed. + let gas_before = ext.gas_amount(); + assert!(ext.unreserve_gas(id).is_ok()); + let gas_after = ext.gas_amount(); + + assert_eq!(gas_after.left(), gas_before.left()); + } + + #[test] + fn test_unreserve_with_reimbursement() { + let costs = ExtCosts { + rent: RentCosts { + reservation: CostOf::new(10), + ..Default::default() + }, + ..Default::default() + }; + let mut ext = Ext::new( + ProcessorContextBuilder::new() + .with_gas(GasCounter::new(u64::MAX)) + .with_message_context(MessageContextBuilder::new().build()) + .with_existential_deposit(500) + .with_costs(costs.clone()) + .build(), + ); + + // Reserving and unreserving immediately + let reservation_amount = 1_000_000; + let duration = 10; + let duration_cost = costs + .rent + .reservation + .cost_for(ext.context.reserve_for.saturating_add(duration).into()); + let id = ext + .reserve_gas(reservation_amount, duration) + .expect("internal error: failed reservation"); + let gas_before = ext.gas_amount(); + assert!(ext.unreserve_gas(id).is_ok()); + let gas_after = ext.gas_amount(); + + assert_eq!( + gas_after.left(), + gas_before.left() + reservation_amount + duration_cost + ); + } } diff --git a/core/src/gas.rs b/core/src/gas.rs index 22c33aadbc4..58b2f24e134 100644 --- a/core/src/gas.rs +++ b/core/src/gas.rs @@ -18,7 +18,7 @@ //! Gas module. -use crate::costs::CostToken; +use crate::{costs::CostToken, reservation::ReimburseUnreserved}; use enum_iterator::Sequence; use scale_info::{ scale::{Decode, Encode}, @@ -99,12 +99,11 @@ impl GasCounter { } } - /// Increase gas by `amount`. + /// Increase left gas by `amount`. /// /// Called when gas unreservation is occurred. - // We don't decrease `burn` counter because `GasTree` manipulation is handled by separated function - // TODO: uncomment when unreserving in current message features is discussed - /*pub fn increase(&mut self, amount: u64) -> bool { + /// We don't decrease `burn` counter because `GasTree` manipulation is handled by separated function + pub fn increase(&mut self, amount: u64, _token: ReimburseUnreserved) -> bool { match self.left.checked_add(amount) { None => false, Some(new_left) => { @@ -112,7 +111,7 @@ impl GasCounter { true } } - }*/ + } /// Reduce gas by `amount`. /// diff --git a/core/src/reservation.rs b/core/src/reservation.rs index 5466245bfee..3f080dd2ff8 100644 --- a/core/src/reservation.rs +++ b/core/src/reservation.rs @@ -178,7 +178,6 @@ impl GasReserver { let id = ReservationId::generate(self.message_id, self.nonce.fetch_inc()); - // TODO #2773 let maybe_reservation = self.states.insert( id, GasReservationState::Created { @@ -209,7 +208,10 @@ impl GasReserver { /// 1. Reservation doesn't exist. /// 2. Reservation was "unreserved", so in [`GasReservationState::Removed`] state. /// 3. Reservation was marked used. - pub fn unreserve(&mut self, id: ReservationId) -> Result { + pub fn unreserve( + &mut self, + id: ReservationId, + ) -> Result<(u64, Option), ReservationError> { // Docs error case #1. let state = self .states @@ -229,13 +231,15 @@ impl GasReserver { let state = self.states.remove(&id).unwrap(); - let amount = match state { + Ok(match state { GasReservationState::Exists { amount, finish, .. } => { self.states .insert(id, GasReservationState::Removed { expiration: finish }); - amount + (amount, None) } - GasReservationState::Created { amount, .. } => amount, + GasReservationState::Created { + amount, duration, .. + } => (amount, Some(ReimburseUnreserved(duration))), GasReservationState::Removed { .. } => { let err_msg = "GasReserver::unreserve: `Removed` variant is unreachable, checked above"; @@ -243,9 +247,7 @@ impl GasReserver { log::error!("{err_msg}"); unreachable!("{err_msg}") } - }; - - Ok(amount) + }) } /// Marks reservation as used. @@ -331,6 +333,19 @@ impl GasReserver { } } +/// Safety token returned when unreserved gas can be returned back to the gas counter. +/// +/// Wraps duration for the newly created reservation. +#[derive(Debug, PartialEq, Eq)] +pub struct ReimburseUnreserved(u32); + +impl ReimburseUnreserved { + /// Returns duration for the newly created unreserved reservation. + pub fn duration(&self) -> u32 { + self.0 + } +} + /// Gas reservations states. pub type GasReservationStates = BTreeMap; @@ -527,7 +542,7 @@ mod tests { let mut reserver = GasReserver::new(&Default::default(), map, 256); - let amount = reserver.unreserve(id).expect("Shouldn't fail"); + let (amount, _) = reserver.unreserve(id).expect("Shouldn't fail"); assert_eq!(amount, slot.amount); assert!(reserver.unreserve(id).is_err());