mysten_common/sync/
oneshot.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4//! A oneshot channel whose receiver blocks the calling OS thread.
5//!
6//! This is the synchronous analogue of `tokio::sync::oneshot`, for code that must wait
7//! for a value produced elsewhere without being async (e.g. execution threads blocking
8//! until an object version is committed). It is a thin wrapper over the [`oneshot`]
9//! crate that adds a [`Receiver::blocking_recv`] tailored to Sui's execution model:
10//!
11//! * The wait always begins with a non-blocking `try_recv`, so a value that is already
12//!   available is returned without any side effects.
13//! * If it must block, it first [releases the thread's execution permit] so other
14//!   execution can proceed (see that module for the deadlock rationale).
15//! * Under msim, parking the OS thread would hang the single-threaded simulator, so it
16//!   instead polls in a loop, yielding the simulated thread quantum between checks. This
17//!   would busy-wait on a real system, but the simulator only wakes the thread at a
18//!   controlled rate, and only blocking-pool threads may wait this way.
19//!
20//! [releases the thread's execution permit]: crate::sync::execution_permit::release_execution_permit
21
22use crate::sync::execution_permit::release_execution_permit;
23
24/// Create a new oneshot channel.
25pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
26    let (sender, receiver) = oneshot::channel();
27    (Sender(sender), Receiver(receiver))
28}
29
30pub struct Sender<T>(oneshot::Sender<T>);
31pub struct Receiver<T>(oneshot::Receiver<T>);
32
33/// Error returned by `blocking_recv` when the sender was dropped without sending.
34#[derive(Debug, PartialEq, Eq)]
35pub struct RecvError;
36
37/// Error returned by `try_recv` when no value is available.
38#[derive(Debug, PartialEq, Eq)]
39pub enum TryRecvError {
40    /// The sender has not sent a value yet.
41    Empty,
42    /// The sender was dropped without sending.
43    Closed,
44}
45
46impl<T> Sender<T> {
47    /// Send a value, consuming the sender. Returns the value back if the receiver was
48    /// already dropped.
49    pub fn send(self, value: T) -> Result<(), T> {
50        self.0.send(value).map_err(|e| e.into_inner())
51    }
52
53    /// Whether the receiver has been dropped (i.e. a send would fail).
54    pub fn is_closed(&self) -> bool {
55        self.0.is_closed()
56    }
57}
58
59impl<T> Receiver<T> {
60    /// Wait until a value is sent (or the sender is dropped), blocking the calling
61    /// thread.
62    ///
63    /// Must not be called from an async context. Under msim it may only be called from
64    /// a blocking-pool thread (e.g. inside `spawn_blocking`); it yields the thread's
65    /// quantum between readiness checks.
66    pub fn blocking_recv(self) -> Result<T, RecvError> {
67        // Fast path: never release the execution permit if a value is already available.
68        match self.0.try_recv() {
69            Ok(value) => return Ok(value),
70            Err(oneshot::TryRecvError::Disconnected) => return Err(RecvError),
71            Err(oneshot::TryRecvError::Empty) => {}
72        }
73
74        // We are about to block; give up our execution permit (if any) so that other
75        // execution can make progress. Blocking while holding it risks deadlock under
76        // limited execution concurrency.
77        release_execution_permit();
78
79        #[cfg(msim)]
80        loop {
81            match self.0.try_recv() {
82                Ok(value) => return Ok(value),
83                Err(oneshot::TryRecvError::Disconnected) => return Err(RecvError),
84                Err(oneshot::TryRecvError::Empty) => msim::task::yield_blocking(),
85            }
86        }
87
88        #[cfg(not(msim))]
89        self.0.recv().map_err(|_| RecvError)
90    }
91
92    /// Return the value if one has been sent, without blocking.
93    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
94        match self.0.try_recv() {
95            Ok(value) => Ok(value),
96            Err(oneshot::TryRecvError::Empty) => Err(TryRecvError::Empty),
97            Err(oneshot::TryRecvError::Disconnected) => Err(TryRecvError::Closed),
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use crate::sync::execution_permit::set_execution_permit;
106    use std::sync::Arc;
107    use std::sync::atomic::{AtomicBool, Ordering};
108    #[cfg(not(msim))]
109    use std::time::Duration;
110
111    struct DropFlag(Arc<AtomicBool>);
112    impl Drop for DropFlag {
113        fn drop(&mut self) {
114            self.0.store(true, Ordering::SeqCst);
115        }
116    }
117
118    #[test]
119    fn send_then_recv() {
120        let (tx, rx) = channel();
121        tx.send(42).unwrap();
122        assert_eq!(rx.blocking_recv(), Ok(42));
123    }
124
125    // These tests block a real OS thread and rely on it being unparked by another
126    // thread. Under msim, `blocking_recv` yields via `msim::task::yield_blocking`, which
127    // must run on a simulator blocking-pool thread rather than a raw `std::thread`, so
128    // they only run outside the simulator. The msim blocking path is exercised by the
129    // execution-layer simtests that use these primitives.
130    #[cfg(not(msim))]
131    #[test]
132    fn recv_blocks_until_send() {
133        let (tx, rx) = channel();
134        let handle = std::thread::spawn(move || rx.blocking_recv());
135        std::thread::sleep(Duration::from_millis(50));
136        tx.send(7).unwrap();
137        assert_eq!(handle.join().unwrap(), Ok(7));
138    }
139
140    #[cfg(not(msim))]
141    #[test]
142    fn sender_drop_unblocks_recv() {
143        let (tx, rx) = channel::<u64>();
144        let handle = std::thread::spawn(move || rx.blocking_recv());
145        std::thread::sleep(Duration::from_millis(50));
146        drop(tx);
147        assert_eq!(handle.join().unwrap(), Err(RecvError));
148    }
149
150    #[test]
151    fn receiver_drop_closes_sender() {
152        let (tx, rx) = channel::<u64>();
153        assert!(!tx.is_closed());
154        drop(rx);
155        assert!(tx.is_closed());
156        assert_eq!(tx.send(1), Err(1));
157    }
158
159    #[test]
160    fn try_recv() {
161        let (tx, mut rx) = channel();
162        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
163        tx.send(5).unwrap();
164        assert_eq!(rx.try_recv(), Ok(5));
165    }
166
167    #[test]
168    fn keeps_permit_when_value_ready() {
169        let (tx, rx) = channel::<u8>();
170        tx.send(9).unwrap();
171        let released = Arc::new(AtomicBool::new(false));
172        let _guard = set_execution_permit(Box::new(DropFlag(released.clone())));
173        assert_eq!(rx.blocking_recv(), Ok(9));
174        assert!(
175            !released.load(Ordering::SeqCst),
176            "permit must be kept when the value is already available"
177        );
178    }
179
180    #[cfg(not(msim))]
181    #[test]
182    fn releases_permit_when_blocking() {
183        let (tx, rx) = channel::<u8>();
184        let released = Arc::new(AtomicBool::new(false));
185        let released_recv = released.clone();
186        let handle = std::thread::spawn(move || {
187            let _guard = set_execution_permit(Box::new(DropFlag(released_recv)));
188            rx.blocking_recv()
189        });
190        // The receiver is now blocked; it should already have released its permit.
191        std::thread::sleep(Duration::from_millis(50));
192        assert!(
193            released.load(Ordering::SeqCst),
194            "permit must be released while blocked"
195        );
196        tx.send(3).unwrap();
197        assert_eq!(handle.join().unwrap(), Ok(3));
198    }
199}
mysten_common/sync/oneshot.rs

mysten_common/sync/
oneshot.rs
