Skip to content

Commit

Permalink
Introduce reimbursement
Browse files Browse the repository at this point in the history
  • Loading branch information
techraed committed Aug 19, 2024
1 parent b710fb8 commit 636ca1c
Show file tree
Hide file tree
Showing 3 changed files with 158 additions and 25 deletions.
139 changes: 129 additions & 10 deletions core-processor/src/ext.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1199,15 +1199,37 @@ 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)?;
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)
}
Expand Down Expand Up @@ -1409,11 +1431,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 {
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
);
}
}
11 changes: 5 additions & 6 deletions core/src/gas.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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: ReimburseUnreserved) -> bool {
match self.left.checked_add(amount) {
None => false,
Some(new_left) => {
self.left = new_left;
true
}
}
}*/
}

/// Reduce gas by `amount`.
///
Expand Down
33 changes: 24 additions & 9 deletions core/src/reservation.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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<ReimburseUnreserved>), ReservationError> {
// Docs error case #1.
let state = self
.states
Expand All @@ -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(ReimburseUnreserved(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.
Expand Down Expand Up @@ -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<ReservationId, GasReservationState>;

Expand Down Expand Up @@ -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());
Expand Down

0 comments on commit 636ca1c

Please sign in to comment.