no_alloc_channel/
lib.rs

1#![cfg_attr(not(test), no_std)]
2//! A wrapper of [`Queue`] which allows for asynchronous receiving.
3
4use core::{
5    future::Future,
6    marker::PhantomPinned,
7    pin::Pin,
8    ptr::null_mut,
9    sync::atomic::{AtomicPtr, Ordering},
10    task::{Context, Poll, Waker},
11};
12
13use heapless::spsc::{Consumer, Producer, Queue};
14
15/// A thread-safe channel supporting `async` receiving, with a buffer of up to
16/// `N - 1` elements of type `T`.
17pub struct Channel<T: Sync + 'static, const N: usize> {
18    queue: Queue<T, N>,
19    waker: AtomicPtr<Waker>,
20}
21
22impl<T: Send + Sync + 'static, const N: usize> Channel<T, N> {
23    /// Construct a new, empty.
24    pub const fn new() -> Self {
25        Self {
26            queue: Queue::new(),
27            waker: AtomicPtr::new(null_mut()),
28        }
29    }
30
31    /// Split the channel into a [`SendHalf`] and a [`ReceiveHalf`].
32    pub fn split(&mut self) -> (SendHalf<'_, T, N>, ReceiveHalf<'_, T, N>) {
33        let (producer, consumer) = self.queue.split();
34        let waker = &self.waker;
35        (
36            SendHalf { producer, waker },
37            ReceiveHalf { consumer, waker },
38        )
39    }
40}
41
42/// The sending half of a [`Channel`].
43pub struct SendHalf<'a, T: 'static, const N: usize> {
44    producer: Producer<'a, T, N>,
45    waker: &'a AtomicPtr<Waker>,
46}
47
48impl<'a, T: 'static, const N: usize> SendHalf<'a, T, N> {
49    /// Try to send a value on the channel, returning [`Err`] if the channel is
50    /// full.
51    pub fn try_send(&mut self, packet: T) -> Result<(), SendError> {
52        self.producer.enqueue(packet).map_err(|_| SendError::Full)?;
53        if let Some(waker) = unsafe { self.waker.load(Ordering::Acquire).as_ref() } {
54            waker.wake_by_ref();
55        }
56        Ok(())
57    }
58}
59
60/// The receiving half of a [`Channel`].
61pub struct ReceiveHalf<'a, T: Send + 'static, const N: usize> {
62    consumer: Consumer<'a, T, N>,
63    waker: &'a AtomicPtr<Waker>,
64}
65
66impl<'a, T: Send + 'static, const N: usize> ReceiveHalf<'a, T, N> {
67    /// Obtain a future which resolves with a value received from the channel.
68    ///
69    /// The receiving operation is atomic, and will only actually remove the
70    /// element from the channel in the process of resolving to a
71    /// [`Poll::Ready`] value.
72    pub fn receive(&mut self) -> impl Future<Output = T> + Send + '_ {
73        struct ReceiveFuture<'a, 'b, T: Send + 'static, const N: usize> {
74            receive: &'a mut ReceiveHalf<'b, T, N>,
75            waker: Option<Waker>,
76            _phantom: PhantomPinned,
77        }
78
79        impl<'a, 'b, T: Send + 'static, const N: usize> Drop for ReceiveFuture<'a, 'b, T, N> {
80            fn drop(&mut self) {
81                // The waker pointer is known to currently be either null or point to the waker
82                // inside self.waker (if that is Some(...)). Therefore, no matter what, we want
83                // it to be null after this.
84                self.receive.waker.store(null_mut(), Ordering::Release);
85            }
86        }
87
88        impl<'a, 'b, T: Send + 'static, const N: usize> Future for ReceiveFuture<'a, 'b, T, N> {
89            type Output = T;
90
91            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
92                let self_mut = unsafe { self.get_unchecked_mut() };
93                if let Some(value) = self_mut.receive.try_receive() {
94                    Poll::Ready(value)
95                } else {
96                    if self_mut.waker.is_none() {
97                        let waker = self_mut.waker.get_or_insert(cx.waker().clone());
98                        self_mut.receive.waker.store(waker, Ordering::Release);
99                    }
100                    Poll::Pending
101                }
102            }
103        }
104
105        ReceiveFuture {
106            receive: self,
107            waker: None,
108            _phantom: PhantomPinned,
109        }
110    }
111
112    /// Try to receive a value, returning [`None`] if the channel is empty.
113    pub fn try_receive(&mut self) -> Option<T> {
114        self.consumer.dequeue()
115    }
116}
117
118/// An error that can arise from sending on a channel.
119#[derive(Debug)]
120pub enum SendError {
121    /// The channel is full.
122    Full,
123}
124
125#[cfg(test)]
126#[tokio::test]
127async fn simple() {
128    let mut ch: Channel<u32, 4> = Channel::new();
129    let (mut tx, mut rx) = ch.split();
130
131    assert!(matches!(tx.try_send(0), Ok(())));
132    assert!(matches!(tx.try_send(1), Ok(())));
133    assert!(matches!(tx.try_send(2), Ok(())));
134    assert!(matches!(tx.try_send(3), Err(SendError::Full)));
135
136    assert_eq!(rx.receive().await, 0);
137    assert_eq!(rx.receive().await, 1);
138    assert_eq!(rx.receive().await, 2);
139    assert!(matches!(rx.try_receive(), None));
140}
no_alloc_channel/lib.rs

no_alloc_channel/
lib.rs
