From 48dd39d86d9d0e780322d16fce7c931ea03e8235 Mon Sep 17 00:00:00 2001
From: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
Date: Wed, 26 Jun 2024 01:35:33 +0200
Subject: [PATCH] chore(provider): simplify nonce filler

---
 crates/provider/src/fillers/nonce.rs | 66 +++++++++++++++++++---------
 1 file changed, 46 insertions(+), 20 deletions(-)

diff --git a/crates/provider/src/fillers/nonce.rs b/crates/provider/src/fillers/nonce.rs
index 237cb7d95836..5f281fc31d3d 100644
--- a/crates/provider/src/fillers/nonce.rs
+++ b/crates/provider/src/fillers/nonce.rs
@@ -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.
 ///
@@ -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 {
@@ -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();