forest/message_pool/
mpool_locker.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use crate::shim::address::Address;
5use ahash::HashMap;
6use parking_lot::Mutex;
7use std::sync::Arc;
8use tokio::sync::OwnedMutexGuard;
9
10/// Per-address async lock for serializing `MpoolPushMessage` RPC calls.
11/// Concurrent pushes for the same sender block on each other, while
12/// different senders proceed in parallel.
13///
14/// This is the *outer* lock in Forest's two-tier locking strategy (analogous to
15/// Lotus's `MpoolLocker` + `MessageSigner.lk`). It covers the entire RPC
16/// critical section -- from gas estimation through the final push -- preventing
17/// a second request from reading stale nonce or balance state while the first
18/// is still in-flight.
19///
20/// See also [`NonceTracker`](super::NonceTracker), the inner nonce-specific
21/// lock.
22#[derive(derive_more::Deref)]
23pub struct MpoolLocker {
24    inner: Mutex<HashMap<Address, Arc<tokio::sync::Mutex<()>>>>,
25}
26
27impl MpoolLocker {
28    /// Create a new `MpoolLocker`.
29    pub fn new() -> Self {
30        Self {
31            inner: Mutex::new(HashMap::default()),
32        }
33    }
34
35    /// Acquire an async lock for the given address. The returned guard must be
36    /// held for the duration of the nonce-assign + sign + push critical section.
37    pub async fn take_lock(&self, addr: Address) -> OwnedMutexGuard<()> {
38        let mutex = {
39            let mut map = self.lock();
40            map.retain(|_, v| Arc::strong_count(v) > 1);
41            map.entry(addr)
42                .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
43                .clone()
44        };
45        mutex.lock_owned().await
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use tokio::sync::{Barrier, oneshot};
53    use tokio::time::{Duration, timeout};
54
55    #[tokio::test]
56    async fn test_take_lock_serializes_same_address() {
57        let locker = Arc::new(MpoolLocker::new());
58        let addr = Address::new_id(1);
59
60        let (first_acquired_tx, first_acquired_rx) = oneshot::channel();
61        let (release_first_tx, release_first_rx) = oneshot::channel();
62        let (second_acquired_tx, second_acquired_rx) = oneshot::channel();
63
64        let locker2 = locker.clone();
65        let t1 = tokio::spawn(async move {
66            let _guard = locker2.take_lock(addr).await;
67            let _ = first_acquired_tx.send(());
68            let _ = release_first_rx.await;
69        });
70
71        // Ensure task 1 is holding the lock before starting task 2.
72        first_acquired_rx.await.unwrap();
73
74        let locker3 = locker.clone();
75        let t2 = tokio::spawn(async move {
76            let _guard = locker3.take_lock(addr).await;
77            let _ = second_acquired_tx.send(());
78        });
79
80        // Task 2 must remain blocked while task 1 holds the lock.
81        assert!(
82            timeout(Duration::from_millis(50), second_acquired_rx)
83                .await
84                .is_err(),
85            "second task should not acquire the same address lock while first holds it"
86        );
87
88        let _ = release_first_tx.send(());
89        t1.await.unwrap();
90        t2.await.unwrap();
91    }
92
93    #[tokio::test]
94    async fn test_take_lock_allows_different_addresses() {
95        let locker = Arc::new(MpoolLocker::new());
96        let addr_a = Address::new_id(1);
97        let addr_b = Address::new_id(2);
98
99        let acquired_barrier = Arc::new(Barrier::new(2));
100
101        let locker2 = locker.clone();
102        let barrier_a = acquired_barrier.clone();
103        let t1 = tokio::spawn(async move {
104            let _guard = locker2.take_lock(addr_a).await;
105            barrier_a.wait().await;
106        });
107
108        let locker3 = locker.clone();
109        let barrier_b = acquired_barrier.clone();
110        let t2 = tokio::spawn(async move {
111            let _guard = locker3.take_lock(addr_b).await;
112            barrier_b.wait().await;
113        });
114
115        timeout(Duration::from_millis(200), async {
116            t1.await.unwrap();
117            t2.await.unwrap();
118        })
119        .await
120        .expect("different address locks should be acquired in parallel");
121    }
122
123    #[tokio::test]
124    async fn test_take_lock_prunes_idle_entries() {
125        let locker = MpoolLocker::new();
126        let addr_a = Address::new_id(1);
127        let addr_b = Address::new_id(2);
128
129        {
130            let _guard = locker.take_lock(addr_a).await;
131            assert_eq!(locker.inner.lock().len(), 1);
132        }
133        // Locking a different address triggers retain, which prunes addr_a as guard is now dropped.
134        let _guard_b = locker.take_lock(addr_b).await;
135        assert_eq!(
136            locker.inner.lock().len(),
137            1,
138            "idle entry for addr_a should have been pruned"
139        );
140    }
141}
forest/message_pool/mpool_locker.rs

forest/message_pool/
mpool_locker.rs
