Skip to main content

rtic_sync/
channel.rs

1//! An async aware MPSC channel that can be used on no-alloc systems.
2
3use crate::unsafecell::UnsafeCell;
4use core::{
5    future::poll_fn,
6    mem::MaybeUninit,
7    pin::Pin,
8    ptr,
9    sync::atomic::{fence, Ordering},
10    task::{Poll, Waker},
11};
12#[doc(hidden)]
13pub use critical_section;
14use heapless::Deque;
15use rtic_common::{
16    dropper::OnDrop, wait_queue::DoublyLinkedList, wait_queue::Link,
17    waker_registration::CriticalSectionWakerRegistration as WakerRegistration,
18};
19
20#[cfg(feature = "defmt-03")]
21use crate::defmt;
22
23type WaitQueueData = (Waker, SlotPtr);
24type WaitQueue = DoublyLinkedList<WaitQueueData>;
25
26/// An MPSC channel for use in no-alloc systems. `N` sets the size of the queue.
27///
28/// This channel uses critical sections, however there are extremely small and all `memcpy`
29/// operations of `T` are done without critical sections.
30///
31/// Right now, the size of the queue `N` is limited to 255 elements.
32pub struct Channel<T, const N: usize> {
33    // Here are all indexes that are not used in `slots` and ready to be allocated.
34    freeq: UnsafeCell<Deque<u8, N>>,
35    // Here are wakers and indexes to slots that are ready to be dequeued by the receiver.
36    readyq: UnsafeCell<Deque<u8, N>>,
37    // Waker for the receiver.
38    receiver_waker: WakerRegistration,
39    // Storage for N `T`s, so we don't memcpy around a lot of `T`s.
40    slots: [UnsafeCell<MaybeUninit<T>>; N],
41    // If there is no room in the queue a `Sender`s can wait for there to be place in the queue.
42    wait_queue: WaitQueue,
43    // Keep track of the receiver.
44    receiver_dropped: UnsafeCell<bool>,
45    // Keep track of the number of senders.
46    num_senders: UnsafeCell<usize>,
47}
48
49unsafe impl<T, const N: usize> Send for Channel<T, N> {}
50
51unsafe impl<T, const N: usize> Sync for Channel<T, N> {}
52
53macro_rules! cs_access {
54    ($name:ident, $type:ty) => {
55        /// Access the value mutably.
56        ///
57        /// SAFETY: this function must not be called recursively within `f`.
58        unsafe fn $name<F, R>(&self, _cs: critical_section::CriticalSection, f: F) -> R
59        where
60            F: FnOnce(&mut $type) -> R,
61        {
62            let v = self.$name.get_mut();
63            // SAFETY: we have exclusive access due to the critical section.
64            let v = unsafe { v.deref() };
65            f(v)
66        }
67    };
68}
69
70impl<T, const N: usize> Default for Channel<T, N> {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl<T, const N: usize> Channel<T, N> {
77    // Size check.
78    const fn size_check() {
79        // This limit comes from the the internal `Deque` structures used to track free and ready
80        // slots which have a type of [u8]
81        assert!(N < 256, "This queue support a maximum of 255 entries");
82    }
83
84    /// Create a new channel.
85    #[cfg(not(loom))]
86    pub const fn new() -> Self {
87        const { Self::size_check() };
88        Self {
89            freeq: UnsafeCell::new(Deque::new()),
90            readyq: UnsafeCell::new(Deque::new()),
91            receiver_waker: WakerRegistration::new(),
92            slots: [const { UnsafeCell::new(MaybeUninit::uninit()) }; N],
93            wait_queue: WaitQueue::new(),
94            receiver_dropped: UnsafeCell::new(false),
95            num_senders: UnsafeCell::new(0),
96        }
97    }
98
99    /// Create a new channel.
100    #[cfg(loom)]
101    pub fn new() -> Self {
102        const { Self::size_check() };
103        Self {
104            freeq: UnsafeCell::new(Deque::new()),
105            readyq: UnsafeCell::new(Deque::new()),
106            receiver_waker: WakerRegistration::new(),
107            slots: core::array::from_fn(|_| UnsafeCell::new(MaybeUninit::uninit())),
108            wait_queue: WaitQueue::new(),
109            receiver_dropped: UnsafeCell::new(false),
110            num_senders: UnsafeCell::new(0),
111        }
112    }
113
114    /// Split the queue into a `Sender`/`Receiver` pair.
115    pub fn split(&mut self) -> (Sender<'_, T, N>, Receiver<'_, T, N>) {
116        // NOTE(assert): queue is cleared by dropping the corresponding `Receiver`.
117        debug_assert!(self.readyq.as_mut().is_empty(),);
118
119        let freeq = self.freeq.as_mut();
120
121        freeq.clear();
122
123        // Fill free queue
124        for idx in 0..N as u8 {
125            debug_assert!(!freeq.is_full());
126
127            // SAFETY: This safe as the loop goes from 0 to the capacity of the underlying queue,
128            // and the queue is cleared beforehand.
129            unsafe {
130                freeq.push_back_unchecked(idx);
131            }
132        }
133
134        debug_assert!(freeq.is_full());
135
136        // There is now 1 sender
137        *self.num_senders.as_mut() = 1;
138
139        (Sender(self), Receiver(self))
140    }
141
142    cs_access!(freeq, Deque<u8, N>);
143    cs_access!(readyq, Deque<u8, N>);
144    cs_access!(receiver_dropped, bool);
145    cs_access!(num_senders, usize);
146
147    /// Return free slot `slot` to the channel.
148    ///
149    /// This will do one of two things:
150    /// 1. If there are any waiting `send`-ers, wake the longest-waiting one and hand it `slot`.
151    /// 2. else, insert `slot` into `self.freeq`.
152    ///
153    /// SAFETY: `slot` must be a `u8` that is obtained by dequeueing from [`Self::readyq`], and that `slot`
154    /// is returned at most once.
155    unsafe fn return_free_slot(&self, slot: u8) {
156        critical_section::with(|cs| {
157            fence(Ordering::SeqCst);
158
159            // If someone is waiting in the `wait_queue`, wake the first one up & hand it the free slot.
160            if let Some((wait_head, mut freeq_slot)) = self.wait_queue.pop() {
161                // SAFETY: `freeq_slot` is valid for writes: we are in a critical
162                // section & the `SlotPtr` lives for at least the duration of the wait queue link.
163                unsafe { freeq_slot.replace(Some(slot), cs) };
164                wait_head.wake();
165            } else {
166                // SAFETY: `self.freeq` is not called recursively.
167                unsafe {
168                    self.freeq(cs, |freeq| {
169                        debug_assert!(!freeq.is_full());
170                        // SAFETY: `freeq` is not full.
171                        freeq.push_back_unchecked(slot);
172                    });
173                }
174            }
175        })
176    }
177
178    /// SAFETY: the caller must guarantee that `slot` is an `u8` obtained by dequeueing from [`Self::readyq`],
179    /// and is read at most once.
180    unsafe fn read_slot(&self, slot: u8) -> T {
181        let first_element = self.slots.get_unchecked(slot as usize).get_mut();
182        let ptr = first_element.deref().as_ptr();
183        ptr::read(ptr)
184    }
185}
186
187/// Creates a split channel with `'static` lifetime.
188#[macro_export]
189#[cfg(not(loom))]
190macro_rules! make_channel {
191    ($type:ty, $size:expr) => {{
192        static mut CHANNEL: $crate::channel::Channel<$type, $size> =
193            $crate::channel::Channel::new();
194
195        static CHECK: $crate::portable_atomic::AtomicU8 = $crate::portable_atomic::AtomicU8::new(0);
196
197        $crate::channel::critical_section::with(|_| {
198            if CHECK.load(::core::sync::atomic::Ordering::Relaxed) != 0 {
199                ::core::panic!("call to the same `make_channel` instance twice");
200            }
201
202            CHECK.store(1, ::core::sync::atomic::Ordering::Relaxed);
203        });
204
205        // SAFETY: This is safe as we hide the static mut from others to access it.
206        // Only this point is where the mutable access happens.
207        #[allow(static_mut_refs)]
208        unsafe {
209            CHANNEL.split()
210        }
211    }};
212}
213
214// -------- Sender
215
216/// Error state for when the receiver has been dropped.
217#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
218pub struct NoReceiver<T>(pub T);
219
220/// Errors that 'try_send` can have.
221#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
222pub enum TrySendError<T> {
223    /// Error state for when the receiver has been dropped.
224    NoReceiver(T),
225    /// Error state when the queue is full.
226    Full(T),
227}
228
229impl<T> core::fmt::Debug for NoReceiver<T>
230where
231    T: core::fmt::Debug,
232{
233    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
234        write!(f, "NoReceiver({:?})", self.0)
235    }
236}
237
238impl<T> core::fmt::Debug for TrySendError<T>
239where
240    T: core::fmt::Debug,
241{
242    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
243        match self {
244            TrySendError::NoReceiver(v) => write!(f, "NoReceiver({v:?})"),
245            TrySendError::Full(v) => write!(f, "Full({v:?})"),
246        }
247    }
248}
249
250impl<T> PartialEq for TrySendError<T>
251where
252    T: PartialEq,
253{
254    fn eq(&self, other: &Self) -> bool {
255        match (self, other) {
256            (TrySendError::NoReceiver(v1), TrySendError::NoReceiver(v2)) => v1.eq(v2),
257            (TrySendError::NoReceiver(_), TrySendError::Full(_)) => false,
258            (TrySendError::Full(_), TrySendError::NoReceiver(_)) => false,
259            (TrySendError::Full(v1), TrySendError::Full(v2)) => v1.eq(v2),
260        }
261    }
262}
263
264/// A `Sender` can send to the channel and can be cloned.
265pub struct Sender<'a, T, const N: usize>(&'a Channel<T, N>);
266
267unsafe impl<T, const N: usize> Send for Sender<'_, T, N> {}
268
269/// This is needed to make the async closure in `send` accept that we "share"
270/// the link possible between threads.
271#[derive(Clone)]
272struct LinkPtr(*mut Option<Link<WaitQueueData>>);
273
274impl LinkPtr {
275    /// This will dereference the pointer stored within and give out an `&mut`.
276    unsafe fn get(&mut self) -> &mut Option<Link<WaitQueueData>> {
277        &mut *self.0
278    }
279}
280
281unsafe impl Send for LinkPtr {}
282
283unsafe impl Sync for LinkPtr {}
284
285/// This is needed to make the async closure in `send` accept that we "share"
286/// the link possible between threads.
287#[derive(Clone)]
288struct SlotPtr(*mut Option<u8>);
289
290impl SlotPtr {
291    /// Replace the value of this slot with `new_value`, and return
292    /// the old value.
293    ///
294    /// SAFETY: the pointer in this `SlotPtr` must be valid for writes.
295    unsafe fn replace(
296        &mut self,
297        new_value: Option<u8>,
298        _cs: critical_section::CriticalSection,
299    ) -> Option<u8> {
300        // SAFETY: the critical section guarantees exclusive access, and the
301        // caller guarantees that the pointer is valid.
302        self.replace_exclusive(new_value)
303    }
304
305    /// Replace the value of this slot with `new_value`, and return
306    /// the old value.
307    ///
308    /// SAFETY: the pointer in this `SlotPtr` must be valid for writes, and the caller must guarantee exclusive
309    /// access to the underlying value..
310    unsafe fn replace_exclusive(&mut self, new_value: Option<u8>) -> Option<u8> {
311        // SAFETY: the caller has ensured that we have exclusive access & that
312        // the pointer is valid.
313        unsafe { core::ptr::replace(self.0, new_value) }
314    }
315}
316
317unsafe impl Send for SlotPtr {}
318
319unsafe impl Sync for SlotPtr {}
320
321impl<T, const N: usize> core::fmt::Debug for Sender<'_, T, N> {
322    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
323        write!(f, "Sender")
324    }
325}
326
327#[cfg(feature = "defmt-03")]
328impl<T, const N: usize> defmt::Format for Sender<'_, T, N> {
329    fn format(&self, f: defmt::Formatter) {
330        defmt::write!(f, "Sender",)
331    }
332}
333
334impl<T, const N: usize> Sender<'_, T, N> {
335    #[inline(always)]
336    fn send_footer(&mut self, idx: u8, val: T) {
337        // Write the value to the slots, note; this memcpy is not under a critical section.
338        unsafe {
339            let first_element = self.0.slots.get_unchecked(idx as usize).get_mut();
340            let ptr = first_element.deref().as_mut_ptr();
341            ptr::write(ptr, val)
342        }
343
344        // Write the value into the ready queue.
345        critical_section::with(|cs| {
346            // SAFETY: `self.0.readyq` is not called recursively.
347            unsafe {
348                self.0.readyq(cs, |readyq| {
349                    debug_assert!(!readyq.is_full());
350                    // SAFETY: ready is not full.
351                    readyq.push_back_unchecked(idx);
352                });
353            }
354        });
355
356        fence(Ordering::SeqCst);
357
358        // If there is a receiver waker, wake it.
359        self.0.receiver_waker.wake();
360    }
361
362    /// Try to send a value, non-blocking. If the channel is full this will return an error.
363    pub fn try_send(&mut self, val: T) -> Result<(), TrySendError<T>> {
364        // If the wait queue is not empty, we can't try to push into the queue.
365        if !self.0.wait_queue.is_empty() {
366            return Err(TrySendError::Full(val));
367        }
368
369        // No receiver available.
370        if self.is_closed() {
371            return Err(TrySendError::NoReceiver(val));
372        }
373
374        let free_slot = critical_section::with(|cs| unsafe {
375            // SAFETY: `self.0.freeq` is not called recursively.
376            self.0.freeq(cs, |q| q.pop_front())
377        });
378
379        let idx = if let Some(idx) = free_slot {
380            idx
381        } else {
382            return Err(TrySendError::Full(val));
383        };
384
385        self.send_footer(idx, val);
386
387        Ok(())
388    }
389
390    /// Send a value. If there is no place left in the queue this will wait until there is.
391    /// If the receiver does not exist this will return an error.
392    pub async fn send(&mut self, val: T) -> Result<(), NoReceiver<T>> {
393        let mut free_slot_ptr: Option<u8> = None;
394        let mut link_ptr: Option<Link<WaitQueueData>> = None;
395
396        // Make this future `Drop`-safe.
397        // SAFETY(link_ptr): Shadow the original definition of `link_ptr` so we can't abuse it.
398        let mut link_ptr = LinkPtr(core::ptr::addr_of_mut!(link_ptr));
399        // SAFETY(freed_slot): Shadow the original definition of `free_slot_ptr` so we can't abuse it.
400        let mut free_slot_ptr = SlotPtr(core::ptr::addr_of_mut!(free_slot_ptr));
401
402        let mut link_ptr2 = link_ptr.clone();
403        let mut free_slot_ptr2 = free_slot_ptr.clone();
404        let dropper = OnDrop::new(|| {
405            // SAFETY: We only run this closure and dereference the pointer if we have
406            // exited the `poll_fn` below in the `drop(dropper)` call. The other dereference
407            // of this pointer is in the `poll_fn`.
408            if let Some(link) = unsafe { link_ptr2.get() } {
409                link.remove_from_list(&self.0.wait_queue);
410            }
411
412            // Return our potentially-unused free slot.
413            // Since we are certain that our link has been removed from the list (either
414            // pop-ed or removed just above), we have exclusive access to the free slot pointer.
415            if let Some(freed_slot) = unsafe { free_slot_ptr2.replace_exclusive(None) } {
416                // SAFETY: freed slot is passed to us from `return_free_slot`, which either
417                // directly (through `try_recv`), or indirectly (through another `return_free_slot`)
418                // comes from `readyq`.
419                unsafe { self.0.return_free_slot(freed_slot) };
420            }
421        });
422
423        let idx = poll_fn(|cx| {
424            //  Do all this in one critical section, else there can be race conditions
425            critical_section::with(|cs| {
426                if self.is_closed() {
427                    return Poll::Ready(Err(()));
428                }
429
430                let wq_empty = self.0.wait_queue.is_empty();
431                // SAFETY: `self.0.freeq` is not called recursively.
432                let freeq_empty = unsafe { self.0.freeq(cs, |q| q.is_empty()) };
433
434                // SAFETY: This pointer is only dereferenced here and on drop of the future
435                // which happens outside this `poll_fn`'s stack frame.
436                let link = unsafe { link_ptr.get() };
437
438                // We are already in the wait queue.
439                if let Some(queue_link) = link {
440                    if queue_link.is_popped() {
441                        // SAFETY: `free_slot_ptr` is valid for writes until the end of this future.
442                        let slot = unsafe { free_slot_ptr.replace(None, cs) };
443
444                        // Our link was popped, so it is most definitely not in the list.
445                        // We can safely & correctly `take` it to prevent ourselves from
446                        // redundantly attempting to remove it from the list a 2nd time.
447                        link.take();
448
449                        // If our link is popped, then:
450                        // 1. We were popped by `return_free_lot` and provided us with a slot.
451                        // 2. We were popped by `Receiver::drop` and it did not provide us with a slot, and the channel is closed.
452                        if let Some(slot) = slot {
453                            Poll::Ready(Ok(slot))
454                        } else {
455                            Poll::Ready(Err(()))
456                        }
457                    } else {
458                        Poll::Pending
459                    }
460                }
461                // We are not in the wait queue, but others are, or there is currently no free
462                // slot available.
463                else if !wq_empty || freeq_empty {
464                    // Place the link in the wait queue.
465                    let link_ref =
466                        link.insert(Link::new((cx.waker().clone(), free_slot_ptr.clone())));
467
468                    // SAFETY(new_unchecked): The address to the link is stable as it is defined
469                    // outside this stack frame.
470                    // SAFETY(push): `link_ref` lifetime comes from `link_ptr` and `free_slot_ptr` that
471                    // are shadowed and we make sure in `dropper` that the link is removed from the queue
472                    // before dropping `link_ptr` AND `dropper` makes sure that the shadowed
473                    // `ptr`s live until the end of the stack frame.
474                    unsafe { self.0.wait_queue.push(Pin::new_unchecked(link_ref)) };
475
476                    Poll::Pending
477                }
478                // We are not in the wait queue, no one else is waiting, and there is a free slot available.
479                else {
480                    // SAFETY: `self.0.freeq` is not called recursively.
481                    unsafe {
482                        self.0.freeq(cs, |freeq| {
483                            debug_assert!(!freeq.is_empty());
484                            // SAFETY: `freeq` is non-empty
485                            let slot = freeq.pop_back_unchecked();
486                            Poll::Ready(Ok(slot))
487                        })
488                    }
489                }
490            })
491        })
492        .await;
493
494        // Make sure the link is removed from the queue.
495        drop(dropper);
496
497        if let Ok(idx) = idx {
498            self.send_footer(idx, val);
499
500            Ok(())
501        } else {
502            Err(NoReceiver(val))
503        }
504    }
505
506    /// Returns true if there is no `Receiver`s.
507    pub fn is_closed(&self) -> bool {
508        critical_section::with(|cs| unsafe {
509            // SAFETY: `self.0.receiver_dropped` is not called recursively.
510            self.0.receiver_dropped(cs, |v| *v)
511        })
512    }
513
514    /// Is the queue full.
515    pub fn is_full(&self) -> bool {
516        critical_section::with(|cs| unsafe {
517            // SAFETY: `self.0.freeq` is not called recursively.
518            self.0.freeq(cs, |v| v.is_empty())
519        })
520    }
521
522    /// Is the queue empty.
523    pub fn is_empty(&self) -> bool {
524        critical_section::with(|cs| unsafe {
525            // SAFETY: `self.0.freeq` is not called recursively.
526            self.0.freeq(cs, |v| v.is_full())
527        })
528    }
529}
530
531impl<T, const N: usize> Drop for Sender<'_, T, N> {
532    fn drop(&mut self) {
533        // Count down the reference counter
534        let num_senders = critical_section::with(|cs| {
535            unsafe {
536                // SAFETY: `self.0.num_senders` is not called recursively.
537                self.0.num_senders(cs, |s| {
538                    *s -= 1;
539                    *s
540                })
541            }
542        });
543
544        // If there are no senders, wake the receiver to do error handling.
545        if num_senders == 0 {
546            self.0.receiver_waker.wake();
547        }
548    }
549}
550
551impl<T, const N: usize> Clone for Sender<'_, T, N> {
552    fn clone(&self) -> Self {
553        // Count up the reference counter
554        critical_section::with(|cs| unsafe {
555            // SAFETY: `self.0.num_senders` is not called recursively.
556            self.0.num_senders(cs, |v| *v += 1);
557        });
558
559        Self(self.0)
560    }
561}
562
563// -------- Receiver
564
565/// A receiver of the channel. There can only be one receiver at any time.
566pub struct Receiver<'a, T, const N: usize>(&'a Channel<T, N>);
567
568unsafe impl<T, const N: usize> Send for Receiver<'_, T, N> {}
569
570impl<T, const N: usize> core::fmt::Debug for Receiver<'_, T, N> {
571    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
572        write!(f, "Receiver")
573    }
574}
575
576#[cfg(feature = "defmt-03")]
577impl<T, const N: usize> defmt::Format for Receiver<'_, T, N> {
578    fn format(&self, f: defmt::Formatter) {
579        defmt::write!(f, "Receiver",)
580    }
581}
582
583/// Possible receive errors.
584#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
585#[derive(Debug, PartialEq, Eq, Clone, Copy)]
586pub enum ReceiveError {
587    /// Error state for when all senders has been dropped.
588    NoSender,
589    /// Error state for when the queue is empty.
590    Empty,
591}
592
593impl<T, const N: usize> Receiver<'_, T, N> {
594    /// Receives a value if there is one in the channel, non-blocking.
595    pub fn try_recv(&mut self) -> Result<T, ReceiveError> {
596        // Try to get a ready slot.
597        let ready_slot = critical_section::with(|cs| unsafe {
598            // SAFETY: `self.0.readyq` is not called recursively.
599            self.0.readyq(cs, |q| q.pop_front())
600        });
601
602        if let Some(rs) = ready_slot {
603            // Read the value from the slots, note; this memcpy is not under a critical section.
604            // SAFETY: `rs` is directly obtained from `self.0.readyq` and is read exactly
605            // once.
606            let r = unsafe { self.0.read_slot(rs) };
607
608            // Return the index to the free queue after we've read the value.
609            // SAFETY: `rs` comes directly from `readyq` and is only returned
610            // once.
611            unsafe { self.0.return_free_slot(rs) };
612
613            Ok(r)
614        } else if self.is_closed() {
615            Err(ReceiveError::NoSender)
616        } else {
617            Err(ReceiveError::Empty)
618        }
619    }
620
621    /// Receives a value, waiting if the queue is empty.
622    /// If all senders are dropped this will error with `NoSender`.
623    pub async fn recv(&mut self) -> Result<T, ReceiveError> {
624        // There was nothing in the queue, setup the waiting.
625        poll_fn(|cx| {
626            // Register waker.
627            // TODO: Should it happen here or after the if? This might cause a spurious wake.
628            self.0.receiver_waker.register(cx.waker());
629
630            // Try to dequeue.
631            match self.try_recv() {
632                Ok(val) => {
633                    return Poll::Ready(Ok(val));
634                }
635                Err(ReceiveError::NoSender) => {
636                    return Poll::Ready(Err(ReceiveError::NoSender));
637                }
638                _ => {}
639            }
640
641            Poll::Pending
642        })
643        .await
644    }
645
646    /// Returns true if there are no `Sender`s.
647    pub fn is_closed(&self) -> bool {
648        critical_section::with(|cs| unsafe {
649            // SAFETY: `self.0.num_senders` is not called recursively.
650            self.0.num_senders(cs, |v| *v == 0)
651        })
652    }
653
654    /// Is the queue full.
655    pub fn is_full(&self) -> bool {
656        critical_section::with(|cs| unsafe {
657            // SAFETY: `self.0.readyq` is not called recursively.
658            self.0.readyq(cs, |v| v.is_full())
659        })
660    }
661
662    /// Is the queue empty.
663    pub fn is_empty(&self) -> bool {
664        critical_section::with(|cs| unsafe {
665            // SAFETY: `self.0.readyq` is not called recursively.
666            self.0.readyq(cs, |v| v.is_empty())
667        })
668    }
669}
670
671impl<T, const N: usize> Drop for Receiver<'_, T, N> {
672    fn drop(&mut self) {
673        // Mark the receiver as dropped and wake all waiters
674        critical_section::with(|cs| unsafe {
675            // SAFETY: `self.0.receiver_dropped` is not called recursively.
676            self.0.receiver_dropped(cs, |v| *v = true);
677        });
678
679        let ready_slot = || {
680            critical_section::with(|cs| unsafe {
681                // SAFETY: `self.0.readyq` is not called recursively.
682                self.0.readyq(cs, |q| q.pop_back())
683            })
684        };
685
686        while let Some(slot) = ready_slot() {
687            // SAFETY: `slot` comes from `readyq` and is
688            // read exactly once.
689            drop(unsafe { self.0.read_slot(slot) })
690        }
691
692        while let Some((waker, _)) = self.0.wait_queue.pop() {
693            waker.wake();
694        }
695    }
696}
697
698#[cfg(test)]
699#[cfg(not(loom))]
700mod tests {
701    use core::sync::atomic::AtomicBool;
702    use std::sync::Arc;
703
704    use cassette::Cassette;
705
706    use super::*;
707
708    #[test]
709    fn empty() {
710        let (mut s, mut r) = make_channel!(u32, 10);
711
712        assert!(s.is_empty());
713        assert!(r.is_empty());
714
715        s.try_send(1).unwrap();
716
717        assert!(!s.is_empty());
718        assert!(!r.is_empty());
719
720        r.try_recv().unwrap();
721
722        assert!(s.is_empty());
723        assert!(r.is_empty());
724    }
725
726    #[test]
727    fn full() {
728        let (mut s, mut r) = make_channel!(u32, 3);
729
730        for _ in 0..3 {
731            assert!(!s.is_full());
732            assert!(!r.is_full());
733
734            s.try_send(1).unwrap();
735        }
736
737        assert!(s.is_full());
738        assert!(r.is_full());
739
740        for _ in 0..3 {
741            r.try_recv().unwrap();
742
743            assert!(!s.is_full());
744            assert!(!r.is_full());
745        }
746    }
747
748    #[test]
749    fn send_recieve() {
750        let (mut s, mut r) = make_channel!(u32, 10);
751
752        for i in 0..10 {
753            s.try_send(i).unwrap();
754        }
755
756        assert_eq!(s.try_send(11), Err(TrySendError::Full(11)));
757
758        for i in 0..10 {
759            assert_eq!(r.try_recv().unwrap(), i);
760        }
761
762        assert_eq!(r.try_recv(), Err(ReceiveError::Empty));
763    }
764
765    #[test]
766    fn closed_recv() {
767        let (s, mut r) = make_channel!(u32, 10);
768
769        drop(s);
770
771        assert!(r.is_closed());
772
773        assert_eq!(r.try_recv(), Err(ReceiveError::NoSender));
774    }
775
776    #[test]
777    fn closed_sender() {
778        let (mut s, r) = make_channel!(u32, 10);
779
780        drop(r);
781
782        assert!(s.is_closed());
783
784        assert_eq!(s.try_send(11), Err(TrySendError::NoReceiver(11)));
785    }
786
787    fn make() {
788        let _ = make_channel!(u32, 10);
789    }
790
791    #[test]
792    #[should_panic]
793    fn double_make_channel() {
794        make();
795        make();
796    }
797
798    #[test]
799    fn tuple_channel() {
800        let _ = make_channel!((i32, u32), 10);
801    }
802
803    fn freeq<const N: usize, T, F, R>(channel: &Channel<T, N>, f: F) -> R
804    where
805        F: FnOnce(&mut Deque<u8, N>) -> R,
806    {
807        critical_section::with(|cs| unsafe { channel.freeq(cs, f) })
808    }
809
810    #[test]
811    fn dropping_waked_send_returns_freeq_item() {
812        let (mut tx, mut rx) = make_channel!(u8, 1);
813
814        tx.try_send(0).unwrap();
815        assert!(freeq(&rx.0, |q| q.is_empty()));
816
817        // Running this in a separate thread scope to ensure that `pinned_future` is dropped fully.
818        //
819        // Calling drop explicitly gets hairy because dropping things behind a `Pin` is not easy.
820        std::thread::scope(|scope| {
821            scope.spawn(|| {
822                let pinned_future = core::pin::pin!(tx.send(1));
823                let mut future = Cassette::new(pinned_future);
824
825                future.poll_on();
826
827                assert!(freeq(&rx.0, |q| q.is_empty()));
828                assert!(!rx.0.wait_queue.is_empty());
829
830                assert_eq!(rx.try_recv(), Ok(0));
831
832                assert!(freeq(&rx.0, |q| q.is_empty()));
833            });
834        });
835
836        assert!(!freeq(&rx.0, |q| q.is_empty()));
837
838        // Make sure that rx & tx are alive until here for good measure.
839        drop((tx, rx));
840    }
841
842    #[derive(Debug)]
843    struct SetToTrueOnDrop(Arc<AtomicBool>);
844
845    impl Drop for SetToTrueOnDrop {
846        fn drop(&mut self) {
847            self.0.store(true, Ordering::SeqCst);
848        }
849    }
850
851    #[test]
852    fn non_popped_item_is_dropped() {
853        let mut channel: Channel<SetToTrueOnDrop, 1> = Channel::new();
854
855        let (mut tx, rx) = channel.split();
856
857        let value = Arc::new(AtomicBool::new(false));
858        tx.try_send(SetToTrueOnDrop(value.clone())).unwrap();
859
860        drop((tx, rx));
861
862        assert!(value.load(Ordering::SeqCst));
863    }
864
865    #[test]
866    pub fn splitting_empty_channel_works() {
867        let mut channel: Channel<(), 1> = Channel::new();
868
869        let (mut tx, rx) = channel.split();
870
871        tx.try_send(()).unwrap();
872
873        drop((tx, rx));
874
875        channel.split();
876    }
877}
878
879#[cfg(not(loom))]
880#[cfg(test)]
881mod tokio_tests {
882    #[tokio::test]
883    async fn stress_channel() {
884        const NUM_RUNS: usize = 1_000;
885        const QUEUE_SIZE: usize = 10;
886
887        let (s, mut r) = make_channel!(u32, QUEUE_SIZE);
888        let mut v = std::vec::Vec::new();
889
890        for i in 0..NUM_RUNS {
891            let mut s = s.clone();
892
893            v.push(tokio::spawn(async move {
894                s.send(i as _).await.unwrap();
895            }));
896        }
897
898        let mut map = std::collections::BTreeSet::new();
899
900        for _ in 0..NUM_RUNS {
901            map.insert(r.recv().await.unwrap());
902        }
903
904        assert_eq!(map.len(), NUM_RUNS);
905
906        for v in v {
907            v.await.unwrap();
908        }
909    }
910}
911
912#[cfg(test)]
913#[cfg(loom)]
914mod loom_test {
915    use cassette::Cassette;
916    use loom::thread;
917
918    #[macro_export]
919    #[allow(missing_docs)]
920    macro_rules! make_loom_channel {
921        ($type:ty, $size:expr) => {{
922            let channel: crate::channel::Channel<$type, $size> = super::Channel::new();
923            let boxed = Box::new(channel);
924            let boxed = Box::leak(boxed);
925
926            // SAFETY: This is safe as we hide the static mut from others to access it.
927            // Only this point is where the mutable access happens.
928            boxed.split()
929        }};
930    }
931
932    // This test tests the following scenarios:
933    // 1. Receiver is dropped while concurrent senders are waiting to send.
934    // 2. Concurrent senders are competing for the same free slot.
935    #[test]
936    pub fn concurrent_send_while_full_and_drop() {
937        loom::model(|| {
938            let (mut tx, mut rx) = make_loom_channel!([u8; 20], 1);
939            let mut cloned = tx.clone();
940
941            tx.try_send([1; 20]).unwrap();
942
943            let handle1 = thread::spawn(move || {
944                let future = std::pin::pin!(tx.send([1; 20]));
945                let mut future = Cassette::new(future);
946                if future.poll_on().is_none() {
947                    future.poll_on();
948                }
949            });
950
951            rx.try_recv().ok();
952
953            let future = std::pin::pin!(cloned.send([1; 20]));
954            let mut future = Cassette::new(future);
955            if future.poll_on().is_none() {
956                future.poll_on();
957            }
958
959            drop(rx);
960
961            handle1.join().unwrap();
962        });
963    }
964}