Skip to content

Commit

Permalink
chore(provider): simplify nonce filler
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Jun 26, 2024
1 parent 1936907 commit 4807a4a
Showing 1 changed file with 74 additions and 20 deletions.
94 changes: 74 additions & 20 deletions crates/provider/src/fillers/nonce.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ use alloy_network::{Network, TransactionBuilder};
use alloy_primitives::Address;
use alloy_transport::{Transport, TransportResult};
use dashmap::DashMap;
use futures::lock::Mutex;
use std::sync::Arc;
use tokio::sync::Mutex;

/// A [`TxFiller`] that fills nonces on transactions.
///
Expand Down Expand Up @@ -41,7 +41,7 @@ use tokio::sync::Mutex;
/// ```
#[derive(Clone, Debug, Default)]
pub struct NonceFiller {
nonces: DashMap<Address, Arc<Mutex<Option<u64>>>>,
nonces: DashMap<Address, Arc<Mutex<u64>>>,
}

impl<N: Network> TxFiller<N> for NonceFiller {
Expand Down Expand Up @@ -86,29 +86,31 @@ impl<N: Network> TxFiller<N> for NonceFiller {

impl NonceFiller {
/// Get the next nonce for the given account.
async fn get_next_nonce<P, T, N>(&self, provider: &P, from: Address) -> TransportResult<u64>
async fn get_next_nonce<P, T, N>(&self, provider: &P, address: Address) -> TransportResult<u64>
where
P: Provider<T, N>,
N: Network,
T: Transport + Clone,
{
// locks dashmap internally for a short duration to clone the `Arc`
let mutex = Arc::clone(self.nonces.entry(from).or_default().value());

// locks the value (does not lock dashmap)
let mut nonce = mutex.lock().await;
match *nonce {
Some(ref mut nonce) => {
*nonce += 1;
Ok(*nonce)
}
None => {
// initialize the nonce if we haven't seen this account before
let initial_nonce = provider.get_transaction_count(from).await?;
*nonce = Some(initial_nonce);
Ok(initial_nonce)
}
}
// Use `u64::MAX` as a sentinel value to indicate that the nonce has not been fetched yet.
const NONE: u64 = u64::MAX;

// Locks dashmap internally for a short duration to clone the `Arc`.
// We also don't want to hold the dashmap lock through the await point below.
let nonce = {
let rm = self.nonces.entry(address).or_insert_with(|| Arc::new(Mutex::new(NONE)));
Arc::clone(rm.value())
};

let mut nonce = nonce.lock().await;
let new_nonce = if *nonce == NONE {
// Initialize the nonce if we haven't seen this account before.
provider.get_transaction_count(address).await?
} else {
*nonce + 1
};
*nonce = new_nonce;
Ok(new_nonce)
}
}

Expand All @@ -119,6 +121,58 @@ mod tests {
use alloy_primitives::{address, U256};
use alloy_rpc_types_eth::TransactionRequest;

async fn check_nonces<P, T, N>(filler: &NonceFiller, provider: &P, address: Address, start: u64)
where
P: Provider<T, N>,
N: Network,
T: Transport + Clone,
{
for i in start..start + 5 {
let nonce = filler.get_next_nonce(&provider, address).await.unwrap();
assert_eq!(nonce, i);
}
}

#[tokio::test]
async fn smoke_test() {
let filler = NonceFiller::default();
let provider = ProviderBuilder::new().on_anvil();
let address = Address::ZERO;
check_nonces(&filler, &provider, address, 0).await;

#[cfg(feature = "anvil-api")]
{
use crate::ext::AnvilApi;
filler.nonces.clear();
provider.anvil_set_nonce(address, U256::from(69)).await.unwrap();
check_nonces(&filler, &provider, address, 69).await;
}
}

#[tokio::test]
async fn concurrency() {
let filler = Arc::new(NonceFiller::default());
let provider = Arc::new(ProviderBuilder::new().on_anvil());
let address = Address::ZERO;
let tasks = (0..5)
.map(|_| {
let filler = Arc::clone(&filler);
let provider = Arc::clone(&provider);
tokio::spawn(async move { filler.get_next_nonce(&provider, address).await })
})
.collect::<Vec<_>>();

let mut ns = Vec::new();
for task in tasks {
ns.push(task.await.unwrap().unwrap());
}
ns.sort_unstable();
assert_eq!(ns, (0..5).collect::<Vec<_>>());

assert_eq!(filler.nonces.len(), 1);
assert_eq!(*filler.nonces.get(&address).unwrap().value().lock().await, 4);
}

#[tokio::test]
async fn no_nonce_if_sender_unset() {
let provider = ProviderBuilder::new().with_nonce_management().on_anvil();
Expand Down

0 comments on commit 4807a4a

Please sign in to comment.