use crate::{
fillers::{FillerControlFlow, TxFiller},
provider::SendableTx,
Provider,
};
use alloy_network::{Network, TransactionBuilder};
use alloy_primitives::Address;
use alloy_transport::{Transport, TransportResult};
use async_trait::async_trait;
use dashmap::DashMap;
use futures::lock::Mutex;
use std::sync::Arc;
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
pub trait NonceManager: Clone + Send + Sync + std::fmt::Debug {
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;
}
#[derive(Clone, Debug, Default)]
#[non_exhaustive]
pub struct SimpleNonceManager;
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl NonceManager for SimpleNonceManager {
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,
{
provider.get_transaction_count(address).pending().await
}
}
#[derive(Clone, Debug, Default)]
pub struct CachedNonceManager {
nonces: DashMap<Address, Arc<Mutex<u64>>>,
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl NonceManager for CachedNonceManager {
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,
{
const NONE: u64 = u64::MAX;
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 {
provider.get_transaction_count(address).await?
} else {
*nonce + 1
};
*nonce = new_nonce;
Ok(new_nonce)
}
}
#[derive(Clone, Debug, Default)]
pub struct NonceFiller<M: NonceManager = SimpleNonceManager> {
nonce_manager: M,
}
impl<M: NonceManager> NonceFiller<M> {
pub const fn new(nonce_manager: M) -> Self {
Self { nonce_manager }
}
}
impl<M: NonceManager, N: Network> TxFiller<N> for NonceFiller<M> {
type Fillable = u64;
fn status(&self, tx: &<N as Network>::TransactionRequest) -> FillerControlFlow {
if tx.nonce().is_some() {
return FillerControlFlow::Finished;
}
if tx.from().is_none() {
return FillerControlFlow::missing("NonceManager", vec!["from"]);
}
FillerControlFlow::Ready
}
fn fill_sync(&self, _tx: &mut SendableTx<N>) {}
async fn prepare<P, T>(
&self,
provider: &P,
tx: &N::TransactionRequest,
) -> TransportResult<Self::Fillable>
where
P: Provider<T, N>,
T: Transport + Clone,
{
let from = tx.from().expect("checked by 'ready()'");
self.nonce_manager.get_next_nonce(provider, from).await
}
async fn fill(
&self,
nonce: Self::Fillable,
mut tx: SendableTx<N>,
) -> TransportResult<SendableTx<N>> {
if let Some(builder) = tx.as_mut_builder() {
builder.set_nonce(nonce);
}
Ok(tx)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{ProviderBuilder, WalletProvider};
use alloy_primitives::{address, U256};
use alloy_rpc_types_eth::TransactionRequest;
async fn check_nonces<P, T, N, M>(
filler: &NonceFiller<M>,
provider: &P,
address: Address,
start: u64,
) where
P: Provider<T, N>,
N: Network,
T: Transport + Clone,
M: NonceManager,
{
for i in start..start + 5 {
let nonce = filler.nonce_manager.get_next_nonce(&provider, address).await.unwrap();
assert_eq!(nonce, i);
}
}
#[tokio::test]
async fn smoke_test() {
let filler = NonceFiller::<CachedNonceManager>::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.nonce_manager.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::<CachedNonceManager>::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.nonce_manager.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.nonce_manager.nonces.len(), 1);
assert_eq!(*filler.nonce_manager.nonces.get(&address).unwrap().value().lock().await, 4);
}
#[tokio::test]
async fn no_nonce_if_sender_unset() {
let provider = ProviderBuilder::new().with_cached_nonce_management().on_anvil();
let tx = TransactionRequest {
value: Some(U256::from(100)),
to: Some(address!("d8dA6BF26964aF9D7eEd9e03E53415D37aA96045").into()),
gas_price: Some(20e9 as u128),
gas: Some(21000),
..Default::default()
};
assert!(provider.send_transaction(tx).await.is_err());
}
#[tokio::test]
async fn increments_nonce() {
let provider = ProviderBuilder::new().with_cached_nonce_management().on_anvil_with_wallet();
let from = provider.default_signer_address();
let tx = TransactionRequest {
from: Some(from),
value: Some(U256::from(100)),
to: Some(address!("d8dA6BF26964aF9D7eEd9e03E53415D37aA96045").into()),
gas_price: Some(20e9 as u128),
gas: Some(21000),
..Default::default()
};
let pending = provider.send_transaction(tx.clone()).await.unwrap();
let tx_hash = pending.watch().await.unwrap();
let mined_tx = provider
.get_transaction_by_hash(tx_hash)
.await
.expect("failed to fetch tx")
.expect("tx not included");
assert_eq!(mined_tx.nonce, 0);
let pending = provider.send_transaction(tx).await.unwrap();
let tx_hash = pending.watch().await.unwrap();
let mined_tx = provider
.get_transaction_by_hash(tx_hash)
.await
.expect("fail to fetch tx")
.expect("tx didn't finalize");
assert_eq!(mined_tx.nonce, 1);
}
}