refactor(core-processor, core): Reimburse reduced for the reservation…

… gas, if it was unreserved immediately (#4165)

core-processor/src/ext.rs

@@ -751,6 +751,15 @@ impl<LP: LazyPagesInterface> Ext<LP> {
         }
         Ok(())
     }
+
+    fn cost_for_reservation(&self, amount: u64, duration: u32) -> u64 {
+        self.context
+            .costs
+            .rent
+            .reservation
+            .cost_for(self.context.reserve_for.saturating_add(duration).into())
+            .saturating_add(amount)
+    }
 }
 
 impl<LP: LazyPagesInterface> CountersOwner for Ext<LP> {
@@ -1176,14 +1185,7 @@ impl<LP: LazyPagesInterface> Externalities for Ext<LP> {
                 return Err(ReservationError::ReservationBelowMailboxThreshold.into());
             }
 
-            let reserve = mutator
-                .context
-                .costs
-                .rent
-                .reservation
-                .cost_for(mutator.context.reserve_for.saturating_add(duration).into());
-
-            let reduce_amount = amount.saturating_add(reserve);
+            let reduce_amount = mutator.cost_for_reservation(amount, duration);
 
             mutator
                 .reduce_gas(reduce_amount)
@@ -1199,15 +1201,29 @@ impl<LP: LazyPagesInterface> Externalities for Ext<LP> {
     }
 
     fn unreserve_gas(&mut self, id: ReservationId) -> Result<u64, Self::FallibleError> {
-        let amount = 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());
-        }*/
+        let (amount, reimburse) = self.context.gas_reserver.unreserve(id)?;
+
+        if let Some(reimbursement) = reimburse {
+            let current_gas_amount = self.gas_amount();
+
+            // Basically amount of the reseravtion and the cost for the hold duration.
+            let reimbursement_amount = self.cost_for_reservation(amount, reimbursement.duration());
+            self.context
+                .gas_counter
+                .increase(reimbursement_amount, reimbursement)
+                .then_some(())
+                .unwrap_or_else(|| {
+                    let err_msg = format!(
+                        "Ext::unreserve_gas: failed to reimburse unreserved gas to left counter. \
+                        Current gas amount - {}, reimburse amount - {}",
+                        current_gas_amount.left(),
+                        amount,
+                    );
+
+                    log::error!("{err_msg}");
+                    unreachable!("{err_msg}");
+                });
+        }
 
         Ok(amount)
     }
@@ -1409,11 +1425,12 @@ impl<LP: LazyPagesInterface> Externalities for Ext<LP> {
 #[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 +1518,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 +2205,106 @@ 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(),
+        );
+
+        // Define params
+        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 reservation_total_cost = reservation_amount + duration_cost;
+
+        let gas_before_reservation = ext.gas_amount();
+        assert_eq!(gas_before_reservation.left(), u64::MAX);
+
+        let id = ext
+            .reserve_gas(reservation_amount, duration)
+            .expect("internal error: failed reservation");
+
+        let gas_after_reservation = ext.gas_amount();
+        assert_eq!(
+            gas_before_reservation.left(),
+            gas_after_reservation.left() + reservation_total_cost
+        );
+
+        assert!(ext.unreserve_gas(id).is_ok());
+
+        let gas_after_unreserve = ext.gas_amount();
+        assert_eq!(
+            gas_after_unreserve.left(),
+            gas_after_reservation.left() + reservation_total_cost
+        );
+        assert_eq!(gas_after_unreserve.left(), gas_before_reservation.left());
+    }
 }

core/src/gas.rs

@@ -18,7 +18,7 @@
 
 //! Gas module.
 
-use crate::costs::CostToken;
+use crate::{costs::CostToken, reservation::UnreservedReimbursement};
 use enum_iterator::Sequence;
 use scale_info::{
     scale::{Decode, Encode},
@@ -99,20 +99,19 @@ 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: UnreservedReimbursement) -> bool {
         match self.left.checked_add(amount) {
             None => false,
             Some(new_left) => {
                 self.left = new_left;
                 true
             }
         }
-    }*/
+    }
 
     /// Reduce gas by `amount`.
     ///

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<u64, ReservationError> {
+    pub fn unreserve(
+        &mut self,
+        id: ReservationId,
+    ) -> Result<(u64, Option<UnreservedReimbursement>), ReservationError> {
         // Docs error case #1.
         let state = self
             .states
@@ -229,23 +231,23 @@ 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(UnreservedReimbursement(duration))),
             GasReservationState::Removed { .. } => {
                 let err_msg =
                     "GasReserver::unreserve: `Removed` variant is unreachable, checked above";
 
                 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 UnreservedReimbursement(u32);
+
+impl UnreservedReimbursement {
+    /// Returns duration for the newly created unreserved reservation.
+    pub fn duration(&self) -> u32 {
+        self.0
+    }
+}
+
 /// Gas reservations states.
 pub type GasReservationStates = BTreeMap<ReservationId, GasReservationState>;
 
@@ -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());