1use 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
26pub struct Channel<T, const N: usize> {
33 freeq: UnsafeCell<Deque<u8, N>>,
35 readyq: UnsafeCell<Deque<u8, N>>,
37 receiver_waker: WakerRegistration,
39 slots: [UnsafeCell<MaybeUninit<T>>; N],
41 wait_queue: WaitQueue,
43 receiver_dropped: UnsafeCell<bool>,
45 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 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 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 const fn size_check() {
79 assert!(N < 256, "This queue support a maximum of 255 entries");
82 }
83
84 #[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 #[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 pub fn split(&mut self) -> (Sender<'_, T, N>, Receiver<'_, T, N>) {
116 debug_assert!(self.readyq.as_mut().is_empty(),);
118
119 let freeq = self.freeq.as_mut();
120
121 freeq.clear();
122
123 for idx in 0..N as u8 {
125 debug_assert!(!freeq.is_full());
126
127 unsafe {
130 freeq.push_back_unchecked(idx);
131 }
132 }
133
134 debug_assert!(freeq.is_full());
135
136 *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 unsafe fn return_free_slot(&self, slot: u8) {
156 critical_section::with(|cs| {
157 fence(Ordering::SeqCst);
158
159 if let Some((wait_head, mut freeq_slot)) = self.wait_queue.pop() {
161 unsafe { freeq_slot.replace(Some(slot), cs) };
164 wait_head.wake();
165 } else {
166 unsafe {
168 self.freeq(cs, |freeq| {
169 debug_assert!(!freeq.is_full());
170 freeq.push_back_unchecked(slot);
172 });
173 }
174 }
175 })
176 }
177
178 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#[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 #[allow(static_mut_refs)]
208 unsafe {
209 CHANNEL.split()
210 }
211 }};
212}
213
214#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
218pub struct NoReceiver<T>(pub T);
219
220#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
222pub enum TrySendError<T> {
223 NoReceiver(T),
225 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
264pub 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#[derive(Clone)]
272struct LinkPtr(*mut Option<Link<WaitQueueData>>);
273
274impl LinkPtr {
275 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#[derive(Clone)]
288struct SlotPtr(*mut Option<u8>);
289
290impl SlotPtr {
291 unsafe fn replace(
296 &mut self,
297 new_value: Option<u8>,
298 _cs: critical_section::CriticalSection,
299 ) -> Option<u8> {
300 self.replace_exclusive(new_value)
303 }
304
305 unsafe fn replace_exclusive(&mut self, new_value: Option<u8>) -> Option<u8> {
311 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 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 critical_section::with(|cs| {
346 unsafe {
348 self.0.readyq(cs, |readyq| {
349 debug_assert!(!readyq.is_full());
350 readyq.push_back_unchecked(idx);
352 });
353 }
354 });
355
356 fence(Ordering::SeqCst);
357
358 self.0.receiver_waker.wake();
360 }
361
362 pub fn try_send(&mut self, val: T) -> Result<(), TrySendError<T>> {
364 if !self.0.wait_queue.is_empty() {
366 return Err(TrySendError::Full(val));
367 }
368
369 if self.is_closed() {
371 return Err(TrySendError::NoReceiver(val));
372 }
373
374 let free_slot = critical_section::with(|cs| unsafe {
375 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 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 let mut link_ptr = LinkPtr(core::ptr::addr_of_mut!(link_ptr));
399 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 if let Some(link) = unsafe { link_ptr2.get() } {
409 link.remove_from_list(&self.0.wait_queue);
410 }
411
412 if let Some(freed_slot) = unsafe { free_slot_ptr2.replace_exclusive(None) } {
416 unsafe { self.0.return_free_slot(freed_slot) };
420 }
421 });
422
423 let idx = poll_fn(|cx| {
424 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 let freeq_empty = unsafe { self.0.freeq(cs, |q| q.is_empty()) };
433
434 let link = unsafe { link_ptr.get() };
437
438 if let Some(queue_link) = link {
440 if queue_link.is_popped() {
441 let slot = unsafe { free_slot_ptr.replace(None, cs) };
443
444 link.take();
448
449 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 else if !wq_empty || freeq_empty {
464 let link_ref =
466 link.insert(Link::new((cx.waker().clone(), free_slot_ptr.clone())));
467
468 unsafe { self.0.wait_queue.push(Pin::new_unchecked(link_ref)) };
475
476 Poll::Pending
477 }
478 else {
480 unsafe {
482 self.0.freeq(cs, |freeq| {
483 debug_assert!(!freeq.is_empty());
484 let slot = freeq.pop_back_unchecked();
486 Poll::Ready(Ok(slot))
487 })
488 }
489 }
490 })
491 })
492 .await;
493
494 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 pub fn is_closed(&self) -> bool {
508 critical_section::with(|cs| unsafe {
509 self.0.receiver_dropped(cs, |v| *v)
511 })
512 }
513
514 pub fn is_full(&self) -> bool {
516 critical_section::with(|cs| unsafe {
517 self.0.freeq(cs, |v| v.is_empty())
519 })
520 }
521
522 pub fn is_empty(&self) -> bool {
524 critical_section::with(|cs| unsafe {
525 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 let num_senders = critical_section::with(|cs| {
535 unsafe {
536 self.0.num_senders(cs, |s| {
538 *s -= 1;
539 *s
540 })
541 }
542 });
543
544 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 critical_section::with(|cs| unsafe {
555 self.0.num_senders(cs, |v| *v += 1);
557 });
558
559 Self(self.0)
560 }
561}
562
563pub 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#[cfg_attr(feature = "defmt-03", derive(defmt::Format))]
585#[derive(Debug, PartialEq, Eq, Clone, Copy)]
586pub enum ReceiveError {
587 NoSender,
589 Empty,
591}
592
593impl<T, const N: usize> Receiver<'_, T, N> {
594 pub fn try_recv(&mut self) -> Result<T, ReceiveError> {
596 let ready_slot = critical_section::with(|cs| unsafe {
598 self.0.readyq(cs, |q| q.pop_front())
600 });
601
602 if let Some(rs) = ready_slot {
603 let r = unsafe { self.0.read_slot(rs) };
607
608 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 pub async fn recv(&mut self) -> Result<T, ReceiveError> {
624 poll_fn(|cx| {
626 self.0.receiver_waker.register(cx.waker());
629
630 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 pub fn is_closed(&self) -> bool {
648 critical_section::with(|cs| unsafe {
649 self.0.num_senders(cs, |v| *v == 0)
651 })
652 }
653
654 pub fn is_full(&self) -> bool {
656 critical_section::with(|cs| unsafe {
657 self.0.readyq(cs, |v| v.is_full())
659 })
660 }
661
662 pub fn is_empty(&self) -> bool {
664 critical_section::with(|cs| unsafe {
665 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 critical_section::with(|cs| unsafe {
675 self.0.receiver_dropped(cs, |v| *v = true);
677 });
678
679 let ready_slot = || {
680 critical_section::with(|cs| unsafe {
681 self.0.readyq(cs, |q| q.pop_back())
683 })
684 };
685
686 while let Some(slot) = ready_slot() {
687 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 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 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 boxed.split()
929 }};
930 }
931
932 #[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}