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 25, 2024
1 parent 1936907 commit 48dd39d
Showing 1 changed file with 46 additions and 20 deletions.
66 changes: 46 additions & 20 deletions crates/provider/src/fillers/nonce.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,10 @@ use alloy_network::{Network, TransactionBuilder};
use alloy_primitives::Address;
use alloy_transport::{Transport, TransportResult};
use dashmap::DashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};

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

impl<N: Network> TxFiller<N> for NonceFiller {
Expand Down Expand Up @@ -92,33 +94,57 @@ impl NonceFiller {
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 = Arc::clone(
self.nonces.entry(from).or_insert_with(|| Arc::new(AtomicU64::new(NONE))).value(),
);

let prev_nonce = nonce.fetch_add(1, Ordering::AcqRel);
let current_nonce;
if prev_nonce == NONE {
// Initialize the nonce if we haven't seen this account before.
current_nonce = provider.get_transaction_count(from).await?;
nonce.store(current_nonce, Ordering::Release);
} else {
current_nonce = prev_nonce + 1;
// `nonce` is already incremented in the `fetch_add` call above.
}
Ok(current_nonce)
}
}

#[cfg(test)]
mod tests {
use super::*;
use crate::{ProviderBuilder, WalletProvider};
use crate::{ext::AnvilApi, ProviderBuilder, WalletProvider};
use alloy_primitives::{address, U256};
use alloy_rpc_types_eth::TransactionRequest;

#[tokio::test]
async fn smoke_test() {
let filler = NonceFiller::default();
let provider = ProviderBuilder::new().on_anvil();
let address = Address::ZERO;
for i in 0..5 {
let nonce = filler.get_next_nonce(&provider, address).await.unwrap();
assert_eq!(nonce, i);
}

#[cfg(feature = "anvil-api")]
{
filler.nonces.clear();
provider.anvil_set_nonce(address, U256::from(69)).await.unwrap();
for i in 0..5 {
let nonce = filler.get_next_nonce(&provider, address).await.unwrap();
assert_eq!(nonce, 69 + i);
}
}
}

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

0 comments on commit 48dd39d

Please sign in to comment.