Skip to main content

forest/message_pool/
nonce_tracker.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use crate::eth::EthChainId;
5use crate::key_management::{Key, sign_message};
6use crate::message_pool::MessagePool;
7use crate::message_pool::msgpool::provider::Provider;
8use crate::shim::message::Message;
9use tokio::sync::Mutex;
10
11/// Serializes nonce assignment globally. The global mutex prevents concurrent
12/// nonce assignment across all senders, ensuring sequential nonce values.
13///
14/// See also [`MpoolLocker`](super::MpoolLocker), the outer per-sender lock.
15pub struct NonceTracker {
16    lock: Mutex<()>,
17}
18
19impl NonceTracker {
20    pub fn new() -> Self {
21        Self {
22            lock: Mutex::new(()),
23        }
24    }
25
26    /// Acquire the global lock, assign a nonce, sign, and push to `mpool`.
27    pub async fn sign_and_push<T: Provider + Send + Sync + 'static>(
28        &self,
29        mpool: &MessagePool<T>,
30        mut message: Message,
31        key: &Key,
32        eth_chain_id: EthChainId,
33    ) -> anyhow::Result<crate::message::SignedMessage> {
34        let _guard = self.lock.lock().await;
35
36        let nonce = mpool.get_sequence(&message.from).await?;
37        message.sequence = nonce;
38
39        let smsg = sign_message(key, &message, eth_chain_id)?;
40        mpool.push(smsg.clone()).await?;
41        Ok(smsg)
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use crate::key_management::{KeyStore, KeyStoreConfig, Wallet};
49    use crate::message_pool::MessagePool;
50    use crate::message_pool::msgpool::test_provider::TestApi;
51    use crate::prelude::*;
52    use crate::shim::crypto::SignatureType;
53    use crate::shim::{address::Address, econ::TokenAmount};
54    use std::sync::Arc;
55    use tokio::task::JoinSet;
56
57    fn make_test_pool_and_wallet() -> (
58        MessagePool<TestApi>,
59        Wallet,
60        Address,
61        flume::Receiver<crate::libp2p::NetworkMessage>,
62    ) {
63        let keystore = KeyStore::new(KeyStoreConfig::Memory).unwrap();
64        let mut wallet = Wallet::new(keystore);
65        let sender = wallet.generate_addr(SignatureType::Secp256k1).unwrap();
66        let tma = TestApi::default();
67        tma.set_state_sequence(&sender, 0);
68        tma.set_state_balance_raw(&sender, TokenAmount::from_whole(1000));
69        let (tx, rx) = flume::bounded(50);
70        let mut services = JoinSet::new();
71        let mpool = MessagePool::new(
72            tma,
73            tx,
74            Default::default(),
75            Default::default(),
76            &mut services,
77        )
78        .unwrap();
79        (mpool, wallet, sender, rx)
80    }
81
82    fn make_message(from: Address) -> Message {
83        Message {
84            from,
85            to: Address::new_id(99),
86            value: TokenAmount::from_atto(1),
87            method_num: 0,
88            sequence: 0,
89            gas_limit: 10_000_000,
90            gas_fee_cap: TokenAmount::from_atto(10_000),
91            gas_premium: TokenAmount::from_atto(100),
92            ..Default::default()
93        }
94    }
95
96    #[tokio::test]
97    async fn test_sign_and_push_assigns_sequential_nonces() {
98        let tracker = NonceTracker::new();
99        let (mpool, mut wallet, sender, _rx) = make_test_pool_and_wallet();
100
101        let key = wallet.find_key(&sender).unwrap();
102        let eth_chain_id: EthChainId = crate::networks::calibnet::ETH_CHAIN_ID;
103
104        let msg1 = make_message(sender);
105        let smsg1 = tracker
106            .sign_and_push(&mpool, msg1, &key, eth_chain_id)
107            .await
108            .unwrap();
109        assert_eq!(smsg1.message().sequence, 0);
110
111        let msg2 = make_message(sender);
112        let smsg2 = tracker
113            .sign_and_push(&mpool, msg2, &key, eth_chain_id)
114            .await
115            .unwrap();
116        assert_eq!(smsg2.message().sequence, 1);
117    }
118
119    #[tokio::test]
120    async fn test_concurrent_push_no_nonce_duplicates() {
121        const N: usize = 10;
122        let tracker = Arc::new(NonceTracker::new());
123        let (mpool, mut wallet, sender, _rx) = make_test_pool_and_wallet();
124        let key = Arc::new(wallet.find_key(&sender).unwrap());
125        let eth_chain_id: EthChainId = crate::networks::calibnet::ETH_CHAIN_ID;
126
127        let mut tasks = JoinSet::new();
128        for _ in 0..N {
129            let (tracker, mpool, key) = (tracker.clone(), mpool.shallow_clone(), key.clone());
130            tasks.spawn(async move {
131                tracker
132                    .sign_and_push(&mpool, make_message(sender), &key, eth_chain_id)
133                    .await
134                    .unwrap()
135                    .message()
136                    .sequence
137            });
138        }
139
140        let mut nonces: Vec<u64> = tasks.join_all().await;
141        nonces.sort();
142
143        let expected: Vec<u64> = (0..N as u64).collect();
144        assert_eq!(nonces, expected, "nonces must be contiguous 0..{N}");
145    }
146}