Skip to content

Commit

Permalink
fix(pool): fix race condition in paymaster tracking
Browse files Browse the repository at this point in the history
  • Loading branch information
@dancoombs
dancoombs committed Jun 17, 2024
1 parent 98f8e2c commit b7fd2cb
Show file tree
Hide file tree
Showing 3 changed files with 193 additions and 141 deletions.
136 changes: 46 additions & 90 deletions crates/pool/src/mempool/paymaster.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ use rundler_types::{
use rundler_utils::cache::LruMap;

use super::MempoolResult;
use crate::chain::{BalanceUpdate, MinedOp};
use crate::chain::MinedOp;

/// Keeps track of current and pending paymaster balances
#[derive(Debug)]
Expand Down Expand Up @@ -94,28 +94,6 @@ where
Ok(stake_status)
}

pub(crate) fn update_paymaster_balances_after_update<'a>(
&self,
entity_balance_updates: impl Iterator<Item = &'a BalanceUpdate>,
unmined_entity_balance_updates: impl Iterator<Item = &'a BalanceUpdate>,
) {
for balance_update in entity_balance_updates {
self.state.write().update_paymaster_balance_from_event(
balance_update.address,
balance_update.amount,
balance_update.is_addition,
)
}

for unmined_balance_update in unmined_entity_balance_updates {
self.state.write().update_paymaster_balance_from_event(
unmined_balance_update.address,
unmined_balance_update.amount,
!unmined_balance_update.is_addition,
)
}
}

pub(crate) async fn paymaster_balance(
&self,
paymaster: Address,
Expand Down Expand Up @@ -174,7 +152,20 @@ where
self.state.write().set_tracking(tracking_enabled);
}

pub(crate) async fn reset_confimed_balances(&self) -> MempoolResult<()> {
pub(crate) async fn reset_confirmed_balances_for(
&self,
addresses: Vec<Address>,
) -> MempoolResult<()> {
let balances = self.entry_point.get_balances(addresses.clone()).await?;

self.state
.write()
.set_confimed_balances(&addresses, &balances);

Ok(())
}

pub(crate) async fn reset_confirmed_balances(&self) -> MempoolResult<()> {
let paymaster_addresses = self.paymaster_addresses();

let balances = self
Expand All @@ -194,6 +185,7 @@ where
.write()
.update_paymaster_balance_from_mined_op(mined_op);
}

pub(crate) fn remove_operation(&self, id: &UserOperationId) {
self.state.write().remove_operation(id);
}
Expand Down Expand Up @@ -324,21 +316,6 @@ impl PaymasterTrackerInner {
keys
}

fn update_paymaster_balance_from_event(
&mut self,
paymaster: Address,
amount: U256,
should_add: bool,
) {
if let Some(paymaster_balance) = self.paymaster_balances.get(&paymaster) {
if should_add {
paymaster_balance.confirmed = paymaster_balance.confirmed.saturating_add(amount);
} else {
paymaster_balance.confirmed = paymaster_balance.confirmed.saturating_sub(amount);
}
}
}

fn paymaster_metadata(&self, paymaster: Address) -> Option<PaymasterMetadata> {
if let Some(paymaster_balance) = self.paymaster_balances.peek(&paymaster) {
return Some(PaymasterMetadata {
Expand Down Expand Up @@ -530,7 +507,7 @@ mod tests {
};

use super::*;
use crate::{chain::BalanceUpdate, mempool::paymaster::PaymasterTracker};
use crate::mempool::paymaster::PaymasterTracker;

fn demo_pool_op(uo: UserOperation) -> PoolOperation {
PoolOperation {
Expand Down Expand Up @@ -609,55 +586,6 @@ mod tests {
assert!(res.is_err());
}

#[tokio::test]
async fn test_update_balance() {
let paymaster_tracker = new_paymaster_tracker();

let paymaster = Address::random();
let pending_op_cost = U256::from(100);
let confirmed_balance = U256::from(1000);

paymaster_tracker.add_new_paymaster(paymaster, confirmed_balance, pending_op_cost);

let deposit = BalanceUpdate {
address: paymaster,
entrypoint: Address::random(),
amount: 100.into(),
is_addition: true,
};

// deposit
paymaster_tracker
.update_paymaster_balances_after_update(vec![&deposit].into_iter(), vec![].into_iter());

let balance = paymaster_tracker
.paymaster_balance(paymaster)
.await
.unwrap();

assert_eq!(balance.confirmed_balance, 1100.into());

let withdrawal = BalanceUpdate {
address: paymaster,
entrypoint: Address::random(),
amount: 50.into(),
is_addition: false,
};

// withdrawal
paymaster_tracker.update_paymaster_balances_after_update(
vec![&withdrawal].into_iter(),
vec![].into_iter(),
);

let balance = paymaster_tracker
.paymaster_balance(paymaster)
.await
.unwrap();

assert_eq!(balance.confirmed_balance, 1050.into());
}

#[tokio::test]
async fn new_uo_not_enough_balance_tracking_disabled() {
let paymaster_tracker = new_paymaster_tracker();
Expand Down Expand Up @@ -731,7 +659,35 @@ mod tests {

assert_eq!(balance_0.confirmed_balance, 1000.into());

let _ = paymaster_tracker.reset_confimed_balances().await;
let _ = paymaster_tracker.reset_confirmed_balances().await;

let balance_0 = paymaster_tracker
.paymaster_balance(paymaster_0)
.await
.unwrap();

assert_eq!(balance_0.confirmed_balance, 50.into());
}

#[tokio::test]
async fn test_reset_balances_for() {
let paymaster_tracker = new_paymaster_tracker();

let paymaster_0 = Address::random();
let paymaster_0_confimed = 1000.into();

paymaster_tracker.add_new_paymaster(paymaster_0, paymaster_0_confimed, 0.into());

let balance_0 = paymaster_tracker
.paymaster_balance(paymaster_0)
.await
.unwrap();

assert_eq!(balance_0.confirmed_balance, 1000.into());

let _ = paymaster_tracker
.reset_confirmed_balances_for(vec![paymaster_0])
.await;

let balance_0 = paymaster_tracker
.paymaster_balance(paymaster_0)
Expand Down
Loading

0 comments on commit b7fd2cb

Please sign in to comment.