forest/message_pool/
nonce_tracker.rs1use 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
11pub struct NonceTracker {
16 lock: Mutex<()>,
17}
18
19impl NonceTracker {
20 pub fn new() -> Self {
21 Self {
22 lock: Mutex::new(()),
23 }
24 }
25
26 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}