fix: race in Signer::fetch_tx_nonce (#101)

* fix: race in Signer::fetch_tx_nonce

* fix: get rid of AtomicU64 behind Mutex

* chore: make clippy happy

* chore: track_caller

* chore: make clippy happy

Author: mitinarseny

api/src/errors.rs

@@ -160,7 +160,7 @@ pub enum ArgumentValidationError {
     #[error("Account creation error: {0}")]
     AccountCreationError(#[from] AccountCreationError),
     #[error("Multiple errors: {0:?}")]
-    MultipleErrors(Vec<ArgumentValidationError>),
+    MultipleErrors(Vec<Self>),
 }
 
 impl ArgumentValidationError {

api/src/signer/mod.rs

@@ -110,7 +110,7 @@ use std::{
     collections::HashMap,
     path::{Path, PathBuf},
     sync::{
-        atomic::{AtomicU64, AtomicUsize, Ordering},
+        atomic::{AtomicUsize, Ordering},
         Arc,
     },
 };
@@ -126,7 +126,7 @@ use near_api_types::{
 use borsh::{BorshDeserialize, BorshSerialize};
 use serde::{Deserialize, Serialize};
 use slipped10::BIP32Path;
-use tracing::{debug, info, instrument, trace, warn};
+use tracing::{debug, instrument, warn};
 
 use crate::{
     config::NetworkConfig,
@@ -394,7 +394,7 @@ pub trait SignerTrait {
 /// It provides an access key pooling and a nonce caching mechanism to improve transaction throughput.
 pub struct Signer {
     pool: tokio::sync::RwLock<HashMap<PublicKey, Box<dyn SignerTrait + Send + Sync + 'static>>>,
-    nonce_cache: tokio::sync::RwLock<HashMap<(AccountId, PublicKey), AtomicU64>>,
+    nonce_cache: futures::lock::Mutex<HashMap<(AccountId, PublicKey), u64>>,
     current_public_key: AtomicUsize,
 }
 
@@ -410,7 +410,7 @@ impl Signer {
                 public_key,
                 Box::new(signer) as Box<dyn SignerTrait + Send + Sync + 'static>,
             )])),
-            nonce_cache: tokio::sync::RwLock::new(HashMap::new()),
+            nonce_cache: futures::lock::Mutex::new(HashMap::new()),
             current_public_key: AtomicUsize::new(0),
         }))
     }
@@ -545,36 +545,20 @@ impl Signer {
         network: &NetworkConfig,
     ) -> Result<(Nonce, CryptoHash, BlockHeight), SignerError> {
         debug!(target: SIGNER_TARGET, "Fetching transaction nonce");
+
         let nonce_data = crate::account::Account(account_id.clone())
             .access_key(public_key)
             .fetch_from(network)
             .await
             .map_err(|e| SignerError::FetchNonceError(Box::new(e)))?;
-        let nonce_cache = self.nonce_cache.read().await;
-
-        if let Some(nonce) = nonce_cache.get(&(account_id.clone(), public_key)) {
-            let nonce = nonce.fetch_add(1, Ordering::SeqCst);
-            drop(nonce_cache);
-            trace!(target: SIGNER_TARGET, "Nonce fetched from cache");
-            return Ok((nonce + 1, nonce_data.block_hash, nonce_data.block_height));
-        } else {
-            drop(nonce_cache);
-        }
-
-        // It's initialization, so it's better to take write lock, so other will wait
 
-        // case where multiple writers end up at the same lock acquisition point and tries
-        // to overwrite the cached value that a previous writer already wrote.
-        let nonce = self
-            .nonce_cache
-            .write()
-            .await
-            .entry((account_id.clone(), public_key))
-            .or_insert_with(|| AtomicU64::new(nonce_data.data.nonce.0 + 1))
-            .fetch_max(nonce_data.data.nonce.0 + 1, Ordering::SeqCst)
-            .max(nonce_data.data.nonce.0 + 1);
+        let nonce = {
+            let mut nonce_cache = self.nonce_cache.lock().await;
+            let nonce = nonce_cache.entry((account_id, public_key)).or_default();
+            *nonce = (*nonce).max(nonce_data.data.nonce.0) + 1;
+            *nonce
+        };
 
-        info!(target: SIGNER_TARGET, "Nonce fetched and cached");
         Ok((nonce, nonce_data.block_hash, nonce_data.block_height))
     }
 

types/src/transaction/result.rs

@@ -39,6 +39,7 @@ pub struct Execution<T> {
 }
 
 impl<T> Execution<T> {
+    #[track_caller]
     pub fn assert_success(self) -> T {
         #[allow(clippy::unwrap_used)]
         self.into_result().unwrap()
@@ -251,11 +252,13 @@ impl ExecutionFinalResult {
     ///
     /// Because this function may panic, its use is generally discouraged. Instead, prefer
     /// to call into [`into_result`](ExecutionFinalResult::into_result) then pattern matching and handle the Err case explicitly.
+    #[track_caller]
     pub fn assert_success(self) -> ExecutionSuccess {
         #[allow(clippy::unwrap_used)]
         self.into_result().unwrap()
     }
 
+    #[track_caller]
     pub fn assert_failure(self) -> ExecutionResult<TxExecutionError> {
         #[allow(clippy::unwrap_used)]
         self.into_result().unwrap_err()