nuclei/sync/local/
oneshot.rs

1//! A channel that carries one value from one sender to one receiver, both
2//! on one shard.
3//!
4//! The names, signatures and errors follow [`sync::oneshot`](crate::sync::oneshot),
5//! so code switches by changing the path. The state is one counted block of
6//! the thread's cache (as an `Rc`, so no allocation in steady state) with a
7//! `Cell` of three bits (value sent, receiver closed, sender dropped), the
8//! value in a cell and the two wakers in cells: no atomic and no fence. The
9//! sender sends at most one value and never waits; the receiver is a
10//! [`Future`] that completes with the value, or with a
11//! [`RecvError`] when the sender is dropped without sending, and, like the
12//! atomic one, panics when polled again after it completed. A value that was
13//! sent and never received stays in the channel until the receiver is
14//! dropped, which drops it on the spot (outside any borrow).
15//!
16//! A receive that completes spends one operation of the cooperative budget
17//! (see [`task::coop`](crate::task::coop)).
18
19#![allow(unsafe_code)]
20
21use core::cell::{Cell, UnsafeCell};
22use core::fmt;
23use core::future::Future;
24use core::pin::Pin;
25use core::task::{Context, Poll, Waker};
26
27use crate::mem::BlockRc as Rc;
28pub use crate::sync::oneshot::error;
29use crate::sync::oneshot::error::{RecvError, TryRecvError};
30
31const VALUE_SENT: u8 = 0b001;
32const RX_CLOSED: u8 = 0b010;
33const TX_DROPPED: u8 = 0b100;
34
35struct Inner<T> {
36    state: Cell<u8>,
37    value: UnsafeCell<Option<T>>,
38    rx_waker: Cell<Option<Waker>>,
39    tx_waker: Cell<Option<Waker>>,
40}
41
42impl<T> Inner<T> {
43    fn has(&self, bit: u8) -> bool {
44        self.state.get() & bit != 0
45    }
46
47    fn set(&self, bit: u8) {
48        self.state.set(self.state.get() | bit);
49    }
50
51    fn put(&self, value: T) {
52        // SAFETY: `!Sync` (a `BlockRc`), and this step runs no user code and
53        // calls no other step; the cell is empty until the one send.
54        unsafe { *self.value.get() = Some(value) };
55    }
56
57    fn take(&self) -> Option<T> {
58        // SAFETY: as in `put`.
59        unsafe { (*self.value.get()).take() }
60    }
61}
62
63/// Replaces the waker in `slot` with `waker` unless it already wakes the
64/// same task; the old one is dropped after the cell is set.
65#[inline]
66fn register(slot: &Cell<Option<Waker>>, waker: &Waker) {
67    match slot.take() {
68        Some(old) if old.will_wake(waker) => slot.set(Some(old)),
69        old => {
70            slot.set(Some(waker.clone()));
71            drop(old);
72        }
73    }
74}
75
76#[inline]
77fn wake(slot: &Cell<Option<Waker>>) {
78    if let Some(waker) = slot.take() {
79        waker.wake();
80    }
81}
82
83/// Creates a one-shot channel and answers its two halves.
84///
85/// The halves belong to the tasks of one shard: they are `!Send`.
86///
87/// ```
88/// use nuclei::sync::local::oneshot;
89///
90/// nuclei::runtime::block_on(async {
91///     let (tx, rx) = oneshot::channel();
92///     tx.send(3).unwrap();
93///     assert_eq!(rx.await, Ok(3));
94/// });
95/// ```
96pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
97    let inner = Rc::new(Inner {
98        state: Cell::new(0),
99        value: UnsafeCell::new(None),
100        rx_waker: Cell::new(None),
101        tx_waker: Cell::new(None),
102    });
103    (
104        Sender {
105            inner: Some(Rc::clone(&inner)),
106        },
107        Receiver { inner: Some(inner) },
108    )
109}
110
111/// Sends one value to the [`Receiver`].
112///
113/// Dropping it without sending ends the receiver's wait with a
114/// [`RecvError`].
115pub struct Sender<T> {
116    inner: Option<Rc<Inner<T>>>,
117}
118
119impl<T> Sender<T> {
120    /// Sends the value and wakes the receiver. Never waits.
121    ///
122    /// # Errors
123    ///
124    /// Gives the value back when the receiver is closed or gone.
125    pub fn send(mut self, value: T) -> Result<(), T> {
126        let Some(inner) = self.inner.take() else {
127            return Err(value);
128        };
129        if inner.has(RX_CLOSED) {
130            inner.set(TX_DROPPED);
131            return Err(value);
132        }
133        inner.put(value);
134        inner.set(VALUE_SENT | TX_DROPPED);
135        wake(&inner.rx_waker);
136        Ok(())
137    }
138
139    /// Completes when the receiver has been closed or dropped. Cancel safe.
140    pub async fn closed(&mut self) {
141        core::future::poll_fn(|cx| self.poll_closed(cx)).await;
142    }
143
144    /// Whether the receiver has been closed or dropped.
145    pub fn is_closed(&self) -> bool {
146        self.inner.as_ref().is_none_or(|inner| inner.has(RX_CLOSED))
147    }
148
149    /// Polls whether the receiver has been closed or dropped; the waker of
150    /// `cx` is woken when it is.
151    pub fn poll_closed(&mut self, cx: &mut Context<'_>) -> Poll<()> {
152        let Some(inner) = self.inner.as_ref() else {
153            return Poll::Ready(());
154        };
155        if inner.has(RX_CLOSED) {
156            return Poll::Ready(());
157        }
158        register(&inner.tx_waker, cx.waker());
159        Poll::Pending
160    }
161}
162
163impl<T> Drop for Sender<T> {
164    fn drop(&mut self) {
165        let Some(inner) = self.inner.as_ref() else {
166            return;
167        };
168        inner.set(TX_DROPPED);
169        wake(&inner.rx_waker);
170        drop(inner.tx_waker.take());
171    }
172}
173
174impl<T> fmt::Debug for Sender<T> {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        f.debug_struct("Sender")
177            .field("is_closed", &self.is_closed())
178            .finish()
179    }
180}
181
182/// Receives the value of the [`Sender`]; a future.
183///
184/// Cancel safe: dropping the future before it completes loses nothing the
185/// channel still holds (a value sent later is dropped with the receiver).
186pub struct Receiver<T> {
187    inner: Option<Rc<Inner<T>>>,
188}
189
190impl<T> Receiver<T> {
191    /// Closes the receiving half: a later send fails and gets its value
192    /// back. A value sent before the close stays receivable.
193    pub fn close(&mut self) {
194        let Some(inner) = self.inner.as_ref() else {
195            return;
196        };
197        inner.set(RX_CLOSED);
198        wake(&inner.tx_waker);
199    }
200
201    /// Whether the receiver completed or failed already.
202    pub fn is_terminated(&self) -> bool {
203        self.inner.is_none()
204    }
205
206    /// Whether no value is waiting for the receiver.
207    pub fn is_empty(&self) -> bool {
208        self.inner
209            .as_ref()
210            .is_none_or(|inner| !inner.has(VALUE_SENT))
211    }
212
213    /// Receives the value without waiting.
214    ///
215    /// # Errors
216    ///
217    /// [`TryRecvError::Empty`] if no value was sent yet and the sender is
218    /// still there; [`TryRecvError::Closed`] if the sender dropped without
219    /// sending, the receiver was closed empty, or the value was received.
220    pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
221        let Some(inner) = self.inner.as_ref() else {
222            return Err(TryRecvError::Closed);
223        };
224        if let Some(value) = inner.take() {
225            self.inner = None;
226            return Ok(value);
227        }
228        if inner.has(RX_CLOSED | TX_DROPPED) {
229            self.inner = None;
230            return Err(TryRecvError::Closed);
231        }
232        Err(TryRecvError::Empty)
233    }
234}
235
236impl<T> Drop for Receiver<T> {
237    fn drop(&mut self) {
238        let Some(inner) = self.inner.as_ref() else {
239            return;
240        };
241        inner.set(RX_CLOSED);
242        wake(&inner.tx_waker);
243        drop(inner.rx_waker.take());
244        drop(inner.take());
245    }
246}
247
248impl<T> Future for Receiver<T> {
249    type Output = Result<T, RecvError>;
250
251    #[expect(
252        clippy::panic,
253        reason = "polled after it completed, which the `Future` contract rules out"
254    )]
255    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
256        let this = self.get_mut();
257        let Some(inner) = this.inner.as_ref() else {
258            panic!("called after complete");
259        };
260        if !inner.has(VALUE_SENT | RX_CLOSED | TX_DROPPED) {
261            register(&inner.rx_waker, cx.waker());
262            return Poll::Pending;
263        }
264        if crate::coop::exhausted(cx.waker()) {
265            return Poll::Pending;
266        }
267        let done = inner.take().ok_or(RecvError(()));
268        this.inner = None;
269        crate::coop::spent_one();
270        Poll::Ready(done)
271    }
272}
273
274impl<T> fmt::Debug for Receiver<T> {
275    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
276        f.debug_struct("Receiver")
277            .field("is_terminated", &self.is_terminated())
278            .field("is_empty", &self.is_empty())
279            .finish()
280    }
281}
nuclei/sync/local/oneshot.rs

nuclei/sync/local/
oneshot.rs
