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//! under msim, parking the OS thread would hang the single-threaded simulator, so it
11//! instead polls in a loop, yielding the simulated thread quantum between checks. This
12//! would busy-wait on a real system, but the simulator only wakes the thread at a
13//! controlled rate, and only blocking-pool threads may wait this way.
14
15/// Create a new oneshot channel.
16pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
17    let (sender, receiver) = oneshot::channel();
18    (Sender(sender), Receiver(receiver))
19}
20
21pub struct Sender<T>(oneshot::Sender<T>);
22pub struct Receiver<T>(oneshot::Receiver<T>);
23
24/// Error returned by `blocking_recv` when the sender was dropped without sending.
25#[derive(Debug, PartialEq, Eq)]
26pub struct RecvError;
27
28/// Error returned by `try_recv` when no value is available.
29#[derive(Debug, PartialEq, Eq)]
30pub enum TryRecvError {
31    /// The sender has not sent a value yet.
32    Empty,
33    /// The sender was dropped without sending.
34    Closed,
35}
36
37impl<T> Sender<T> {
38    /// Send a value, consuming the sender. Returns the value back if the receiver was
39    /// already dropped.
40    pub fn send(self, value: T) -> Result<(), T> {
41        self.0.send(value).map_err(|e| e.into_inner())
42    }
43
44    /// Whether the receiver has been dropped (i.e. a send would fail).
45    pub fn is_closed(&self) -> bool {
46        self.0.is_closed()
47    }
48}
49
50impl<T> Receiver<T> {
51    /// Wait until a value is sent (or the sender is dropped), blocking the calling
52    /// thread.
53    ///
54    /// Must not be called from an async context. Under msim it may only be called from
55    /// a blocking-pool thread (e.g. inside `spawn_blocking`); it yields the thread's
56    /// quantum between readiness checks.
57    pub fn blocking_recv(self) -> Result<T, RecvError> {
58        #[cfg(msim)]
59        loop {
60            match self.0.try_recv() {
61                Ok(value) => return Ok(value),
62                Err(oneshot::TryRecvError::Disconnected) => return Err(RecvError),
63                Err(oneshot::TryRecvError::Empty) => {
64                    crate::assert_reachable!("blocking_recv parks its thread to wait");
65                    msim::task::yield_blocking()
66                }
67            }
68        }
69
70        #[cfg(not(msim))]
71        self.0.recv().map_err(|_| RecvError)
72    }
73
74    /// Return the value if one has been sent, without blocking.
75    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
76        match self.0.try_recv() {
77            Ok(value) => Ok(value),
78            Err(oneshot::TryRecvError::Empty) => Err(TryRecvError::Empty),
79            Err(oneshot::TryRecvError::Disconnected) => Err(TryRecvError::Closed),
80        }
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87    #[cfg(not(msim))]
88    use std::time::Duration;
89
90    #[test]
91    fn send_then_recv() {
92        let (tx, rx) = channel();
93        tx.send(42).unwrap();
94        assert_eq!(rx.blocking_recv(), Ok(42));
95    }
96
97    // These tests block a real OS thread and rely on it being unparked by another
98    // thread. Under msim, `blocking_recv` yields via `msim::task::yield_blocking`, which
99    // must run on a simulator blocking-pool thread rather than a raw `std::thread`, so
100    // they only run outside the simulator. The msim blocking path is exercised by the
101    // execution-layer simtests that use these primitives.
102    #[cfg(not(msim))]
103    #[test]
104    fn recv_blocks_until_send() {
105        let (tx, rx) = channel();
106        let handle = std::thread::spawn(move || rx.blocking_recv());
107        std::thread::sleep(Duration::from_millis(50));
108        tx.send(7).unwrap();
109        assert_eq!(handle.join().unwrap(), Ok(7));
110    }
111
112    #[cfg(not(msim))]
113    #[test]
114    fn sender_drop_unblocks_recv() {
115        let (tx, rx) = channel::<u64>();
116        let handle = std::thread::spawn(move || rx.blocking_recv());
117        std::thread::sleep(Duration::from_millis(50));
118        drop(tx);
119        assert_eq!(handle.join().unwrap(), Err(RecvError));
120    }
121
122    #[test]
123    fn receiver_drop_closes_sender() {
124        let (tx, rx) = channel::<u64>();
125        assert!(!tx.is_closed());
126        drop(rx);
127        assert!(tx.is_closed());
128        assert_eq!(tx.send(1), Err(1));
129    }
130
131    #[test]
132    fn try_recv() {
133        let (tx, mut rx) = channel();
134        assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
135        tx.send(5).unwrap();
136        assert_eq!(rx.try_recv(), Ok(5));
137    }
138}
mysten_common/sync/oneshot.rs

mysten_common/sync/
oneshot.rs
