chore(test-suite): improved nonce manager

eudelins-zama · eudelins-zama · commit 9614d413fbf1 · 2025-11-12T15:56:06.000+01:00
diff --git a/test-suite/gateway-stress/src/blockchain/mod.rs b/test-suite/gateway-stress/src/blockchain/mod.rs
@@ -1,4 +1,5 @@
 pub mod manager;
+pub mod nonce_manager;
 pub mod provider;
 pub mod wallet;
 
diff --git a/test-suite/gateway-stress/src/blockchain/nonce_manager.rs b/test-suite/gateway-stress/src/blockchain/nonce_manager.rs
@@ -0,0 +1,120 @@
+use alloy::{
+    network::Network,
+    primitives::Address,
+    providers::{Provider, fillers::NonceManager},
+    transports::TransportResult,
+};
+use async_trait::async_trait;
+use futures::lock::{Mutex, MutexGuard};
+use std::collections::{BTreeSet, HashMap, hash_map::Entry};
+use std::sync::Arc;
+use tracing::debug;
+
+/// A robust, in-memory nonce manager for a scalable transaction engine.
+#[derive(Clone, Debug, Default)]
+pub struct ZamaNonceManager {
+    /// Nonce state for each account, shared across all tasks/threads using the nonce manager.
+    accounts: Arc<Mutex<HashMap<Address, AccountState>>>,
+}
+
+/// Represents the complete nonce state for a single account.
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
+pub struct AccountState {
+    /// The "high-water mark" nonce. Used only when no gaps are available.
+    pub next_nonce: u64,
+    /// Nonces that have been dispatched but not yet confirmed or rejected.
+    pub locked_nonces: BTreeSet<u64>,
+    /// Nonces that were previously locked but have been released, creating gaps.
+    pub available_nonces: BTreeSet<u64>,
+}
+
+impl ZamaNonceManager {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// The primary logic for acquiring and locking the next valid nonce.
+    ///
+    /// The logic prioritizes filling gaps from `available_nonces` before
+    /// incrementing the main `next_nonce` counter.
+    pub async fn get_increase_and_lock_nonce<P, N>(
+        &self,
+        provider: &P,
+        address: Address,
+    ) -> TransportResult<u64>
+    where
+        P: Provider<N>,
+        N: Network,
+    {
+        let mut accounts_guard = self.accounts.lock().await;
+        let account =
+            Self::get_or_init_account_state(&mut accounts_guard, provider, address).await?;
+        let nonce_to_use =
+            if let Some(available_nonce) = account.available_nonces.iter().next().copied() {
+                account.available_nonces.remove(&available_nonce);
+                debug!(%address, nonce = available_nonce, "Reusing available nonce");
+                available_nonce
+            } else {
+                let next = account.next_nonce;
+                account.next_nonce += 1;
+                debug!(%address, nonce = next, "Using next sequential nonce");
+                next
+            };
+
+        account.locked_nonces.insert(nonce_to_use);
+        Ok(nonce_to_use)
+    }
+
+    /// Releases a locked nonce, making it available for reuse.
+    pub async fn release_nonce(&self, address: Address, nonce: u64) {
+        let mut accounts = self.accounts.lock().await;
+        if let Some(account) = accounts.get_mut(&address)
+            && account.locked_nonces.remove(&nonce)
+        {
+            account.available_nonces.insert(nonce);
+        }
+    }
+
+    /// Confirms a nonce has been used on-chain, removing it permanently.
+    pub async fn confirm_nonce(&self, address: Address, nonce: u64) {
+        let mut accounts = self.accounts.lock().await;
+        if let Some(account) = accounts.get_mut(&address) {
+            account.locked_nonces.remove(&nonce);
+        }
+    }
+
+    /// Helper to retrieve or initialize the `AccountState` for an address.
+    async fn get_or_init_account_state<'a, P, N>(
+        accounts_guard: &'a mut MutexGuard<'_, HashMap<Address, AccountState>>,
+        provider: &P,
+        address: Address,
+    ) -> TransportResult<&'a mut AccountState>
+    where
+        P: Provider<N>,
+        N: Network,
+    {
+        let account = match accounts_guard.entry(address) {
+            Entry::Occupied(entry) => entry.into_mut(),
+            Entry::Vacant(entry) => {
+                let initial_nonce = provider.get_transaction_count(address).await?;
+                entry.insert(AccountState {
+                    next_nonce: initial_nonce,
+                    ..Default::default()
+                })
+            }
+        };
+        Ok(account)
+    }
+}
+
+// Implements the `NonceManager` trait for seamless integration with Alloy's provider stack.
+#[async_trait]
+impl NonceManager for ZamaNonceManager {
+    async fn get_next_nonce<P, N>(&self, provider: &P, address: Address) -> TransportResult<u64>
+    where
+        P: Provider<N>,
+        N: Network,
+    {
+        self.get_increase_and_lock_nonce(provider, address).await
+    }
+}
diff --git a/test-suite/gateway-stress/src/blockchain/provider.rs b/test-suite/gateway-stress/src/blockchain/provider.rs
@@ -1,3 +1,4 @@
+use crate::blockchain::nonce_manager::ZamaNonceManager;
 use alloy::{
     consensus::Account,
     eips::{BlockId, BlockNumberOrTag, Encodable2718},
@@ -7,29 +8,25 @@ use alloy::{
         U256,
     },
     providers::{
-        EthCall, EthCallMany, EthGetBlock, FilterPollerBuilder, GetSubscription,
-        PendingTransaction, PendingTransactionBuilder, PendingTransactionConfig,
-        PendingTransactionError, Provider, ProviderCall, RootProvider, RpcWithBlock, SendableTx,
-        fillers::{
-            BlobGasFiller, CachedNonceManager, FillProvider, GasFiller, JoinFill, NonceManager,
-            TxFiller,
-        },
+        EthCall, EthCallMany, EthGetBlock, FilterPollerBuilder, PendingTransaction,
+        PendingTransactionBuilder, PendingTransactionConfig, PendingTransactionError, Provider,
+        ProviderCall, RootProvider, RpcWithBlock, SendableTx,
+        fillers::{BlobGasFiller, FillProvider, GasFiller, JoinFill, NonceManager, TxFiller},
     },
     rpc::{
         client::NoParams,
+        json_rpc::ErrorPayload,
         types::{
             AccessListResult, Bundle, EIP1186AccountProofResponse, EthCallResponse, FeeHistory,
             Filter, FilterChanges, Index, Log, SyncStatus, TransactionReceipt, TransactionRequest,
             erc4337::TransactionConditional,
-            pubsub::{Params, SubscriptionKind},
             simulate::{SimulatePayload, SimulatedBlock},
         },
     },
-    transports::{TransportError, TransportResult},
+    transports::{RpcError, TransportError, TransportErrorKind, TransportResult},
 };
-use futures::lock::Mutex;
 use serde_json::value::RawValue;
-use std::{borrow::Cow, sync::Arc};
+use std::borrow::Cow;
 
 pub type FillersWithoutNonceManagement = JoinFill<GasFiller, BlobGasFiller>;
 
@@ -46,7 +43,7 @@ where
 {
     inner: FillProvider<F, P, N>,
     signer_address: Address,
-    nonce_manager: Arc<Mutex<CachedNonceManager>>,
+    nonce_manager: ZamaNonceManager,
 }
 
 impl<F, P> NonceManagedProvider<F, P>
@@ -58,19 +55,18 @@ where
         Self {
             inner: provider,
             signer_address,
-            nonce_manager: Default::default(),
+            nonce_manager: ZamaNonceManager::new(),
         }
     }
 
     pub async fn send_transaction_sync(
         &self,
         mut tx: TransactionRequest,
     ) -> TransportResult<TransactionReceipt> {
+        let signer_addr = self.signer_address;
         let nonce = self
             .nonce_manager
-            .lock()
-            .await
-            .get_next_nonce(&self.inner, self.signer_address)
+            .get_next_nonce(&self.inner, signer_addr)
             .await?;
         tx.set_nonce(nonce);
 
@@ -82,15 +78,22 @@ where
             .map_err(|e| TransportError::LocalUsageError(Box::new(e)))?
             .encode_2718(&mut tx_bytes);
 
-        let res = self
+        let send_tx_result = self
             .client()
             .request("eth_sendRawTransactionSync", (Bytes::from(tx_bytes),))
             .await;
-        if res.is_err() {
-            // Reset the nonce manager if the transaction sending failed.
-            *self.nonce_manager.lock().await = Default::default();
+
+        match &send_tx_result {
+            Err(e) if is_nonce_too_low(e) => {
+                self.nonce_manager.confirm_nonce(signer_addr, nonce).await;
+            }
+            Err(_) => {
+                self.nonce_manager.release_nonce(signer_addr, nonce).await;
+            }
+            Ok(_) => self.nonce_manager.confirm_nonce(signer_addr, nonce).await,
         }
-        res
+
+        send_tx_result
     }
 }
 
@@ -109,6 +112,25 @@ where
     }
 }
 
+// See https://ethereum-json-rpc.com/errors
+const ETH_INVALID_INPUT_RPC_ERROR_CODE: i64 = -32000;
+
+/// Returns `true` if the RPC error is "nonce too low" or "already known", `false` otherwise.
+fn is_nonce_too_low(error: &RpcError<TransportErrorKind>) -> bool {
+    match error {
+        RpcError::ErrorResp(ErrorPayload { code, message, .. }) => {
+            if *code == ETH_INVALID_INPUT_RPC_ERROR_CODE {
+                let lowercase_msg = message.to_lowercase();
+                lowercase_msg.starts_with("nonce too low")
+                    || lowercase_msg.starts_with("already known")
+            } else {
+                false
+            }
+        }
+        _ => false,
+    }
+}
+
 #[async_trait::async_trait]
 impl<F, P, N> Provider<N> for NonceManagedProvider<F, P, N>
 where
@@ -124,19 +146,25 @@ where
         &self,
         mut tx: N::TransactionRequest,
     ) -> TransportResult<PendingTransactionBuilder<N>> {
+        let signer_addr = self.signer_address;
         let nonce = self
             .nonce_manager
-            .lock()
-            .await
-            .get_next_nonce(&self.inner, self.signer_address)
+            .get_next_nonce(&self.inner, signer_addr)
             .await?;
         tx.set_nonce(nonce);
-        let res = self.inner.send_transaction(tx).await;
-        if res.is_err() {
-            // Reset the nonce manager if the transaction sending failed.
-            *self.nonce_manager.lock().await = Default::default();
+        let send_tx_result = self.inner.send_transaction(tx).await;
+
+        match &send_tx_result {
+            Err(e) if is_nonce_too_low(e) => {
+                self.nonce_manager.confirm_nonce(signer_addr, nonce).await;
+            }
+            Err(_) => {
+                self.nonce_manager.release_nonce(signer_addr, nonce).await;
+            }
+            Ok(_) => self.nonce_manager.confirm_nonce(signer_addr, nonce).await,
         }
-        res
+
+        send_tx_result
     }
 
     fn get_accounts(&self) -> ProviderCall<NoParams, Vec<Address>> {
@@ -416,28 +444,6 @@ where
         self.inner.sign_transaction(tx).await
     }
 
-    fn subscribe_blocks(&self) -> GetSubscription<(SubscriptionKind,), N::HeaderResponse> {
-        self.inner.subscribe_blocks()
-    }
-
-    fn subscribe_pending_transactions(&self) -> GetSubscription<(SubscriptionKind,), B256> {
-        self.inner.subscribe_pending_transactions()
-    }
-
-    fn subscribe_full_pending_transactions(
-        &self,
-    ) -> GetSubscription<(SubscriptionKind, Params), N::TransactionResponse> {
-        self.inner.subscribe_full_pending_transactions()
-    }
-
-    fn subscribe_logs(&self, filter: &Filter) -> GetSubscription<(SubscriptionKind, Params), Log> {
-        self.inner.subscribe_logs(filter)
-    }
-
-    async fn unsubscribe(&self, id: B256) -> TransportResult<()> {
-        self.inner.unsubscribe(id).await
-    }
-
     fn syncing(&self) -> ProviderCall<NoParams, SyncStatus> {
         self.inner.syncing()
     }
diff --git a/test-suite/gateway-stress/src/decryption/mod.rs b/test-suite/gateway-stress/src/decryption/mod.rs
@@ -62,11 +62,8 @@ where
     for i in 1..=TX_RETRIES {
         overprovision_gas(provider, &mut decryption_call).await;
 
-        trace!("Sending transaction to the Gateway");
-        match provider
-            .send_transaction_sync(decryption_call.clone())
-            .await
-        {
+        debug!("Sending transaction to the Gateway");
+        match send_tx_sync_with_increased_gas_limit(provider, &mut decryption_call).await {
             Ok(receipt) => {
                 let id = extract_id_fn(&receipt)?;
                 id_sender.send(id)?;
@@ -83,6 +80,21 @@ where
     ))
 }
 
+async fn send_tx_sync_with_increased_gas_limit(
+    provider: &AppProvider,
+    call: &mut TransactionRequest,
+) -> anyhow::Result<TransactionReceipt> {
+    // Force a fresh gas estimation on each attempt to account for state drift
+    call.gas = None;
+    overprovision_gas(provider, call).await;
+
+    let receipt = provider.send_transaction_sync(call.clone()).await?;
+    if !receipt.status() {
+        return Err(anyhow!("Tx {} was reverted", receipt.transaction_hash));
+    }
+    Ok(receipt)
+}
+
 async fn overprovision_gas<P: Provider>(provider: &P, call: &mut TransactionRequest) {
     let current_gas = match call.gas {
         Some(gas) => gas,
diff --git a/test-suite/gateway-stress/templates/gw_bench.csv b/test-suite/gateway-stress/templates/gw_bench.csv
@@ -7,11 +7,15 @@ parallel_requests;number_of_measures;decryption_type
 50;50;public
 100;50;public
 200;10;public
+300;5;public
+400;5;public
 1;50;user
 2;50;user
 5;50;user
 10;50;user
 20;50;user
 50;30;user
-100;10;user
-200;5;user
+100;50;user
+200;10;user
+300;5;user
+400;5;user

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`pub mod manager;`
	`2`	`+pub mod nonce_manager;`
`2`	`3`	`pub mod provider;`
`3`	`4`	`pub mod wallet;`
`4`	`5`