#![allow(clippy::module_name_repetitions)]
pub use crate::buffer_pool::{BufferLease, BufferPool};
use core::cell::{Cell, RefCell};
use core::future::{Future, poll_fn};
use core::pin::Pin;
use core::sync::atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering};
use core::task::Poll;
use embassy_sync::blocking_mutex::Mutex as BlockingMutex;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::channel::Channel;
use embassy_sync::waitqueue::{AtomicWaker, MultiWakerRegistration};
const SEND_WAKER_CAP: usize = 8;
use crate::transport::{
MpscRecv, MpscSend, OneshotCancelled, OneshotRecv, OneshotSend, UnboundedRecv, UnboundedSend,
};
const O_SENDER_ALIVE: u8 = 0b001;
const O_RECEIVER_ALIVE: u8 = 0b010;
const O_CANCELLED: u8 = 0b100;
pub struct OneshotSlot<T: Send + 'static> {
chan: Channel<CriticalSectionRawMutex, T, 1>,
cancel_waker: AtomicWaker,
state: AtomicU8,
next_free: AtomicUsize,
}
impl<T: Send + 'static> OneshotSlot<T> {
#[must_use]
pub const fn new() -> Self {
Self {
chan: Channel::new(),
cancel_waker: AtomicWaker::new(),
state: AtomicU8::new(0),
next_free: AtomicUsize::new(0),
}
}
}
impl<T: Send + 'static> Default for OneshotSlot<T> {
fn default() -> Self {
Self::new()
}
}
trait OneshotReclaim<T: Send + 'static>: Send + Sync + 'static {
fn release(&self, slot: &'static OneshotSlot<T>);
}
pub struct OneshotPool<T: Send + 'static, const POOL_SIZE: usize> {
slots: [OneshotSlot<T>; POOL_SIZE],
free_head: BlockingMutex<CriticalSectionRawMutex, Cell<usize>>,
seeded: AtomicBool,
}
impl<T: Send + 'static, const POOL_SIZE: usize> OneshotPool<T, POOL_SIZE> {
#[must_use]
pub const fn new() -> Self {
Self {
slots: [const { OneshotSlot::new() }; POOL_SIZE],
free_head: BlockingMutex::new(Cell::new(0)),
seeded: AtomicBool::new(false),
}
}
pub fn claim(&'static self) -> Option<(StaticOneshotSender<T>, StaticOneshotReceiver<T>)> {
self.ensure_seeded();
let slot = self.pop_free()?;
slot.state
.store(O_SENDER_ALIVE | O_RECEIVER_ALIVE, Ordering::Release);
let _ = slot.chan.try_receive();
Some((
StaticOneshotSender {
slot,
pool: self,
sent: false,
},
StaticOneshotReceiver { slot, pool: self },
))
}
fn ensure_seeded(&self) {
if self.seeded.load(Ordering::Acquire) {
return;
}
self.free_head.lock(|h| {
if self.seeded.load(Ordering::Acquire) {
return;
}
for i in 0..POOL_SIZE {
let next = if i + 1 < POOL_SIZE { i + 2 } else { 0 };
self.slots[i].next_free.store(next, Ordering::Release);
}
h.set(1);
self.seeded.store(true, Ordering::Release);
});
}
fn pop_free(&self) -> Option<&OneshotSlot<T>> {
self.free_head.lock(|h| {
let head = h.get();
if head == 0 {
return None;
}
let slot = &self.slots[head - 1];
let next = slot.next_free.load(Ordering::Acquire);
h.set(next);
slot.next_free.store(0, Ordering::Release);
Some(slot)
})
}
}
impl<T: Send + 'static, const POOL_SIZE: usize> Default for OneshotPool<T, POOL_SIZE> {
fn default() -> Self {
Self::new()
}
}
impl<T: Send + 'static, const POOL_SIZE: usize> OneshotReclaim<T> for OneshotPool<T, POOL_SIZE> {
fn release(&self, slot: &'static OneshotSlot<T>) {
let base = self.slots.as_ptr() as usize;
let here = core::ptr::from_ref::<OneshotSlot<T>>(slot) as usize;
let stride = core::mem::size_of::<OneshotSlot<T>>();
debug_assert!(stride > 0, "OneshotSlot must be sized");
debug_assert!(here >= base);
let idx = (here - base) / stride;
debug_assert!(idx < POOL_SIZE, "slot does not belong to this pool");
let _ = slot.chan.try_receive();
slot.cancel_waker.register(core::task::Waker::noop());
slot.state.store(0, Ordering::Release);
self.free_head.lock(|h| {
slot.next_free.store(h.get(), Ordering::Release);
h.set(idx + 1);
});
}
}
pub struct StaticOneshotSender<T: Send + 'static> {
slot: &'static OneshotSlot<T>,
pool: &'static dyn OneshotReclaim<T>,
sent: bool,
}
impl<T: Send + 'static> OneshotSend<T> for StaticOneshotSender<T> {
fn send(mut self, value: T) -> Result<(), T> {
if self.slot.state.load(Ordering::Acquire) & O_RECEIVER_ALIVE == 0 {
return Err(value);
}
match self.slot.chan.try_send(value) {
Ok(()) => {
self.sent = true;
Ok(())
}
Err(embassy_sync::channel::TrySendError::Full(v)) => Err(v),
}
}
}
impl<T: Send + 'static> Drop for StaticOneshotSender<T> {
fn drop(&mut self) {
if !self.sent {
self.slot.state.fetch_or(O_CANCELLED, Ordering::AcqRel);
self.slot.cancel_waker.wake();
}
let prev = self.slot.state.fetch_and(!O_SENDER_ALIVE, Ordering::AcqRel);
let after = prev & !O_SENDER_ALIVE;
if (after & O_RECEIVER_ALIVE) == 0 {
self.pool.release(self.slot);
}
}
}
pub struct StaticOneshotReceiver<T: Send + 'static> {
slot: &'static OneshotSlot<T>,
pool: &'static dyn OneshotReclaim<T>,
}
impl<T: Send + 'static> OneshotRecv<T> for StaticOneshotReceiver<T> {
async fn recv(self) -> Result<T, OneshotCancelled> {
let slot = self.slot;
let result = poll_fn(move |cx| {
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Ok(v));
}
if slot.state.load(Ordering::Acquire) & O_CANCELLED != 0 {
return Poll::Ready(Err(OneshotCancelled));
}
slot.cancel_waker.register(cx.waker());
{
let mut fut = slot.chan.receive();
let pinned = unsafe { Pin::new_unchecked(&mut fut) };
if let Poll::Ready(v) = pinned.poll(cx) {
return Poll::Ready(Ok(v));
}
}
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Ok(v));
}
if slot.state.load(Ordering::Acquire) & O_CANCELLED != 0 {
return Poll::Ready(Err(OneshotCancelled));
}
Poll::Pending
})
.await;
drop(self);
result
}
}
impl<T: Send + 'static> Drop for StaticOneshotReceiver<T> {
fn drop(&mut self) {
let prev = self
.slot
.state
.fetch_and(!O_RECEIVER_ALIVE, Ordering::AcqRel);
let after = prev & !O_RECEIVER_ALIVE;
if (after & O_SENDER_ALIVE) == 0 {
self.pool.release(self.slot);
}
}
}
pub struct MpscSlot<T: Send + 'static, const SLOT_CAP: usize> {
chan: Channel<CriticalSectionRawMutex, T, SLOT_CAP>,
close_waker: AtomicWaker,
send_wakers:
BlockingMutex<CriticalSectionRawMutex, RefCell<MultiWakerRegistration<SEND_WAKER_CAP>>>,
refcount: AtomicUsize,
closed: AtomicBool,
next_free: AtomicUsize,
}
impl<T: Send + 'static, const SLOT_CAP: usize> MpscSlot<T, SLOT_CAP> {
#[must_use]
pub const fn new() -> Self {
Self {
chan: Channel::new(),
close_waker: AtomicWaker::new(),
send_wakers: BlockingMutex::new(RefCell::new(MultiWakerRegistration::new())),
refcount: AtomicUsize::new(0),
closed: AtomicBool::new(false),
next_free: AtomicUsize::new(0),
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> Default for MpscSlot<T, SLOT_CAP> {
fn default() -> Self {
Self::new()
}
}
trait MpscReclaim<T: Send + 'static, const SLOT_CAP: usize>: Send + Sync + 'static {
fn release(&self, slot: &'static MpscSlot<T, SLOT_CAP>);
}
pub struct MpscPool<T: Send + 'static, const POOL_SIZE: usize, const SLOT_CAP: usize> {
slots: [MpscSlot<T, SLOT_CAP>; POOL_SIZE],
free_head: BlockingMutex<CriticalSectionRawMutex, Cell<usize>>,
seeded: AtomicBool,
}
impl<T: Send + 'static, const POOL_SIZE: usize, const SLOT_CAP: usize>
MpscPool<T, POOL_SIZE, SLOT_CAP>
{
#[must_use]
pub const fn new() -> Self {
Self {
slots: [const { MpscSlot::new() }; POOL_SIZE],
free_head: BlockingMutex::new(Cell::new(0)),
seeded: AtomicBool::new(false),
}
}
pub fn claim_bounded(
&'static self,
) -> Option<(
StaticBoundedSender<T, SLOT_CAP>,
StaticBoundedReceiver<T, SLOT_CAP>,
)> {
let slot = self.claim_inner()?;
Some((
StaticBoundedSender { slot, pool: self },
StaticBoundedReceiver { slot, pool: self },
))
}
pub fn claim_unbounded(
&'static self,
) -> Option<(
StaticUnboundedSender<T, SLOT_CAP>,
StaticUnboundedReceiver<T, SLOT_CAP>,
)> {
let slot = self.claim_inner()?;
Some((
StaticUnboundedSender { slot, pool: self },
StaticUnboundedReceiver { slot, pool: self },
))
}
fn claim_inner(&'static self) -> Option<&'static MpscSlot<T, SLOT_CAP>> {
self.ensure_seeded();
let slot = self.pop_free()?;
slot.refcount.store(2, Ordering::Release); slot.closed.store(false, Ordering::Release);
while slot.chan.try_receive().is_ok() {}
Some(slot)
}
fn ensure_seeded(&self) {
if self.seeded.load(Ordering::Acquire) {
return;
}
self.free_head.lock(|h| {
if self.seeded.load(Ordering::Acquire) {
return;
}
for i in 0..POOL_SIZE {
let next = if i + 1 < POOL_SIZE { i + 2 } else { 0 };
self.slots[i].next_free.store(next, Ordering::Release);
}
h.set(1);
self.seeded.store(true, Ordering::Release);
});
}
fn pop_free(&self) -> Option<&MpscSlot<T, SLOT_CAP>> {
self.free_head.lock(|h| {
let head = h.get();
if head == 0 {
return None;
}
let slot = &self.slots[head - 1];
let next = slot.next_free.load(Ordering::Acquire);
h.set(next);
slot.next_free.store(0, Ordering::Release);
Some(slot)
})
}
}
impl<T: Send + 'static, const POOL_SIZE: usize, const SLOT_CAP: usize> Default
for MpscPool<T, POOL_SIZE, SLOT_CAP>
{
fn default() -> Self {
Self::new()
}
}
impl<T: Send + 'static, const POOL_SIZE: usize, const SLOT_CAP: usize> MpscReclaim<T, SLOT_CAP>
for MpscPool<T, POOL_SIZE, SLOT_CAP>
{
fn release(&self, slot: &'static MpscSlot<T, SLOT_CAP>) {
let base = self.slots.as_ptr() as usize;
let here = core::ptr::from_ref::<MpscSlot<T, SLOT_CAP>>(slot) as usize;
let stride = core::mem::size_of::<MpscSlot<T, SLOT_CAP>>();
debug_assert!(stride > 0);
debug_assert!(here >= base);
let idx = (here - base) / stride;
debug_assert!(idx < POOL_SIZE);
while slot.chan.try_receive().is_ok() {}
slot.close_waker.register(core::task::Waker::noop());
slot.send_wakers.lock(|w| w.borrow_mut().wake());
slot.refcount.store(0, Ordering::Release);
slot.closed.store(false, Ordering::Release);
self.free_head.lock(|h| {
slot.next_free.store(h.get(), Ordering::Release);
h.set(idx + 1);
});
}
}
pub struct StaticBoundedSender<T: Send + 'static, const SLOT_CAP: usize> {
slot: &'static MpscSlot<T, SLOT_CAP>,
pool: &'static dyn MpscReclaim<T, SLOT_CAP>,
}
impl<T: Send + 'static, const SLOT_CAP: usize> Clone for StaticBoundedSender<T, SLOT_CAP> {
fn clone(&self) -> Self {
self.slot.refcount.fetch_add(1, Ordering::AcqRel);
Self {
slot: self.slot,
pool: self.pool,
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> Drop for StaticBoundedSender<T, SLOT_CAP> {
fn drop(&mut self) {
let prev = self.slot.refcount.fetch_sub(1, Ordering::AcqRel);
if prev == 2 {
self.slot.closed.store(true, Ordering::Release);
self.slot.close_waker.wake();
} else if prev == 1 {
self.pool.release(self.slot);
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> MpscSend<T> for StaticBoundedSender<T, SLOT_CAP> {
async fn send(&self, value: T) -> Result<(), ()> {
let slot = self.slot;
if slot.closed.load(Ordering::Acquire) {
return Err(());
}
let mut send_fut = core::pin::pin!(slot.chan.send(value));
poll_fn(|cx| {
if slot.closed.load(Ordering::Acquire) {
return Poll::Ready(Err(()));
}
match send_fut.as_mut().poll(cx) {
Poll::Ready(()) => Poll::Ready(Ok(())),
Poll::Pending => {
slot.send_wakers
.lock(|w| w.borrow_mut().register(cx.waker()));
if slot.closed.load(Ordering::Acquire) {
return Poll::Ready(Err(()));
}
Poll::Pending
}
}
})
.await
}
}
pub struct StaticBoundedReceiver<T: Send + 'static, const SLOT_CAP: usize> {
slot: &'static MpscSlot<T, SLOT_CAP>,
pool: &'static dyn MpscReclaim<T, SLOT_CAP>,
}
impl<T: Send + 'static, const SLOT_CAP: usize> Drop for StaticBoundedReceiver<T, SLOT_CAP> {
fn drop(&mut self) {
self.slot.closed.store(true, Ordering::Release);
self.slot.send_wakers.lock(|w| w.borrow_mut().wake());
let prev = self.slot.refcount.fetch_sub(1, Ordering::AcqRel);
if prev == 1 {
self.pool.release(self.slot);
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> MpscRecv<T> for StaticBoundedReceiver<T, SLOT_CAP> {
fn recv(&mut self) -> impl Future<Output = Option<T>> + Send + '_ {
let slot = self.slot;
async move { mpsc_recv_inner(slot).await }
}
fn poll_recv(&mut self, cx: &mut core::task::Context<'_>) -> core::task::Poll<Option<T>> {
mpsc_poll_recv(self.slot, cx)
}
}
pub struct StaticUnboundedSender<T: Send + 'static, const SLOT_CAP: usize> {
slot: &'static MpscSlot<T, SLOT_CAP>,
pool: &'static dyn MpscReclaim<T, SLOT_CAP>,
}
impl<T: Send + 'static, const SLOT_CAP: usize> Clone for StaticUnboundedSender<T, SLOT_CAP> {
fn clone(&self) -> Self {
self.slot.refcount.fetch_add(1, Ordering::AcqRel);
Self {
slot: self.slot,
pool: self.pool,
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> Drop for StaticUnboundedSender<T, SLOT_CAP> {
fn drop(&mut self) {
let prev = self.slot.refcount.fetch_sub(1, Ordering::AcqRel);
if prev == 2 {
self.slot.closed.store(true, Ordering::Release);
self.slot.close_waker.wake();
} else if prev == 1 {
self.pool.release(self.slot);
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> UnboundedSend<T>
for StaticUnboundedSender<T, SLOT_CAP>
{
fn send_now(&self, value: T) -> Result<(), T> {
if self.slot.closed.load(Ordering::Acquire) {
return Err(value);
}
self.slot.chan.try_send(value).map_err(|e| match e {
embassy_sync::channel::TrySendError::Full(v) => v,
})
}
}
pub struct StaticUnboundedReceiver<T: Send + 'static, const SLOT_CAP: usize> {
slot: &'static MpscSlot<T, SLOT_CAP>,
pool: &'static dyn MpscReclaim<T, SLOT_CAP>,
}
impl<T: Send + 'static, const SLOT_CAP: usize> Drop for StaticUnboundedReceiver<T, SLOT_CAP> {
fn drop(&mut self) {
self.slot.closed.store(true, Ordering::Release);
self.slot.send_wakers.lock(|w| w.borrow_mut().wake());
let prev = self.slot.refcount.fetch_sub(1, Ordering::AcqRel);
if prev == 1 {
self.pool.release(self.slot);
}
}
}
impl<T: Send + 'static, const SLOT_CAP: usize> UnboundedRecv<T>
for StaticUnboundedReceiver<T, SLOT_CAP>
{
fn recv(&mut self) -> impl Future<Output = Option<T>> + Send + '_ {
let slot = self.slot;
async move { mpsc_recv_inner(slot).await }
}
}
async fn mpsc_recv_inner<T: Send + 'static, const SLOT_CAP: usize>(
slot: &'static MpscSlot<T, SLOT_CAP>,
) -> Option<T> {
poll_fn(|cx| mpsc_poll_recv(slot, cx)).await
}
fn mpsc_poll_recv<T: Send + 'static, const SLOT_CAP: usize>(
slot: &'static MpscSlot<T, SLOT_CAP>,
cx: &mut core::task::Context<'_>,
) -> core::task::Poll<Option<T>> {
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Some(v));
}
if slot.closed.load(Ordering::Acquire) {
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Some(v));
}
return Poll::Ready(None);
}
slot.close_waker.register(cx.waker());
{
let mut fut = slot.chan.receive();
let pinned = unsafe { Pin::new_unchecked(&mut fut) };
if let Poll::Ready(v) = pinned.poll(cx) {
return Poll::Ready(Some(v));
}
}
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Some(v));
}
if slot.closed.load(Ordering::Acquire) {
if let Ok(v) = slot.chan.try_receive() {
return Poll::Ready(Some(v));
}
return Poll::Ready(None);
}
Poll::Pending
}
impl<T: Send + 'static> core::fmt::Debug for OneshotSlot<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("OneshotSlot")
.field("state", &self.state)
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for OneshotPool<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("OneshotPool").finish_non_exhaustive()
}
}
impl<T: Send + 'static> core::fmt::Debug for StaticOneshotSender<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticOneshotSender")
.field("sent", &self.sent)
.finish_non_exhaustive()
}
}
impl<T: Send + 'static> core::fmt::Debug for StaticOneshotReceiver<T> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticOneshotReceiver")
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for MpscSlot<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("MpscSlot")
.field("refcount", &self.refcount)
.field("closed", &self.closed)
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const P: usize, const N: usize> core::fmt::Debug for MpscPool<T, P, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("MpscPool").finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for StaticBoundedSender<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticBoundedSender")
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for StaticBoundedReceiver<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticBoundedReceiver")
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for StaticUnboundedSender<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticUnboundedSender")
.finish_non_exhaustive()
}
}
impl<T: Send + 'static, const N: usize> core::fmt::Debug for StaticUnboundedReceiver<T, N> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("StaticUnboundedReceiver")
.finish_non_exhaustive()
}
}
pub const UNBOUNDED_DEFAULT_CAP: usize = 128;
#[macro_export]
macro_rules! define_static_channels {
( vis: $vis:vis, name: $name:ident, $($rest:tt)* ) => {
$crate::define_static_channels! { @body $vis, $name, $($rest)* }
};
( name: $name:ident, $($rest:tt)* ) => {
$crate::define_static_channels! { @body pub, $name, $($rest)* }
};
(
@body $vis:vis, $name:ident,
oneshot: [ $( ($ot:ty, $opool:literal) ),* $(,)? ],
bounded: [ $( (($bt:ty, $bcap:literal), $bpool:literal) ),* $(,)? ],
unbounded: [ $( ($ut:ty, $upool:literal) ),* $(,)? ] $(,)?
) => {
#[derive(Clone, Copy, Debug)]
$vis struct $name;
impl $crate::transport::ChannelFactory for $name {
type OneshotSender<T: ::core::marker::Send + 'static> =
$crate::static_channels::StaticOneshotSender<T>;
type OneshotReceiver<T: ::core::marker::Send + 'static> =
$crate::static_channels::StaticOneshotReceiver<T>;
type BoundedSender<T: ::core::marker::Send + 'static, const N: usize> =
$crate::static_channels::StaticBoundedSender<T, N>;
type BoundedReceiver<T: ::core::marker::Send + 'static, const N: usize> =
$crate::static_channels::StaticBoundedReceiver<T, N>;
type UnboundedSender<T: ::core::marker::Send + 'static> =
$crate::static_channels::StaticUnboundedSender<
T,
{ $crate::static_channels::UNBOUNDED_DEFAULT_CAP },
>;
type UnboundedReceiver<T: ::core::marker::Send + 'static> =
$crate::static_channels::StaticUnboundedReceiver<
T,
{ $crate::static_channels::UNBOUNDED_DEFAULT_CAP },
>;
}
$(
impl $crate::transport::OneshotPooled<$name> for $ot {
fn oneshot_pair() -> (
<$name as $crate::transport::ChannelFactory>::OneshotSender<Self>,
<$name as $crate::transport::ChannelFactory>::OneshotReceiver<Self>,
) {
static POOL: $crate::static_channels::OneshotPool<$ot, $opool> =
$crate::static_channels::OneshotPool::new();
POOL.claim().expect(::core::concat!(
"OneshotPool<",
::core::stringify!($ot),
", ",
::core::stringify!($opool),
"> exhausted; increase the pool size declared in define_static_channels!"
))
}
}
)*
$(
impl $crate::transport::BoundedPooled<$name, $bcap> for $bt {
fn bounded_pair() -> (
<$name as $crate::transport::ChannelFactory>::BoundedSender<Self, $bcap>,
<$name as $crate::transport::ChannelFactory>::BoundedReceiver<Self, $bcap>,
) {
static POOL: $crate::static_channels::MpscPool<$bt, $bpool, $bcap> =
$crate::static_channels::MpscPool::new();
POOL.claim_bounded().expect(::core::concat!(
"MpscPool<",
::core::stringify!($bt),
", pool=",
::core::stringify!($bpool),
", slot_cap=",
::core::stringify!($bcap),
"> exhausted; increase the pool size declared in define_static_channels!"
))
}
}
)*
$(
impl $crate::transport::UnboundedPooled<$name> for $ut {
fn unbounded_pair() -> (
<$name as $crate::transport::ChannelFactory>::UnboundedSender<Self>,
<$name as $crate::transport::ChannelFactory>::UnboundedReceiver<Self>,
) {
static POOL: $crate::static_channels::MpscPool<
$ut,
$upool,
{ $crate::static_channels::UNBOUNDED_DEFAULT_CAP },
> = $crate::static_channels::MpscPool::new();
POOL.claim_unbounded().expect(::core::concat!(
"MpscPool<",
::core::stringify!($ut),
", pool=",
::core::stringify!($upool),
", unbounded> exhausted; increase the pool size declared in define_static_channels!"
))
}
}
)*
};
}
#[cfg(test)]
mod tests {
use super::*;
use core::future::Future;
use core::pin::pin;
use core::task::{Context, Poll, Waker};
use std::boxed::Box;
fn poll_once<F: Future>(f: &mut core::pin::Pin<&mut F>) -> Poll<F::Output> {
let waker = Waker::noop();
let mut cx = Context::from_waker(waker);
f.as_mut().poll(&mut cx)
}
static ONESHOT_POOL_4: OneshotPool<u32, 4> = OneshotPool::new();
#[test]
fn oneshot_send_recv_happy_path() {
let (tx, rx) = ONESHOT_POOL_4.claim().expect("pool not empty");
tx.send(42).unwrap();
let mut fut = pin!(rx.recv());
match poll_once(&mut fut) {
Poll::Ready(Ok(v)) => assert_eq!(v, 42),
other => panic!("expected ready ok, got {other:?}"),
}
}
#[test]
fn oneshot_sender_drop_cancels_receiver() {
let (tx, rx) = ONESHOT_POOL_4.claim().expect("pool not empty");
drop(tx);
let mut fut = pin!(rx.recv());
match poll_once(&mut fut) {
Poll::Ready(Err(OneshotCancelled)) => {}
other => panic!("expected cancelled, got {other:?}"),
}
}
#[test]
fn oneshot_claim_release_cycles() {
static POOL: OneshotPool<u32, 4> = OneshotPool::new();
let p1 = POOL.claim().unwrap();
let p2 = POOL.claim().unwrap();
let p3 = POOL.claim().unwrap();
let p4 = POOL.claim().unwrap();
assert!(POOL.claim().is_none(), "5th claim must exhaust");
drop((p1, p2, p3, p4));
let p5 = POOL.claim();
assert!(p5.is_some(), "post-drop claim must succeed");
}
#[test]
fn oneshot_pool_exhaustion_returns_none() {
static POOL_2: OneshotPool<u32, 2> = OneshotPool::new();
let _a = POOL_2.claim().unwrap();
let _b = POOL_2.claim().unwrap();
assert!(POOL_2.claim().is_none(), "third claim must exhaust");
}
#[test]
fn oneshot_concurrent_first_claim_does_not_panic() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as O};
static POOL: OneshotPool<u32, 8> = OneshotPool::new();
let success_count = Arc::new(AtomicUsize::new(0));
let mut handles = std::vec::Vec::new();
for _ in 0..4 {
let s = Arc::clone(&success_count);
handles.push(std::thread::spawn(move || {
if POOL.claim().is_some() {
s.fetch_add(1, O::SeqCst);
}
}));
}
for h in handles {
h.join().unwrap();
}
assert_eq!(
success_count.load(O::SeqCst),
4,
"all 4 concurrent claims should have succeeded against an 8-slot pool",
);
}
#[test]
fn mpsc_bounded_receiver_drop_wakes_all_cloned_senders() {
static POOL: MpscPool<u32, 4, 1> = MpscPool::new();
let (tx, rx) = POOL.claim_bounded().expect("claim");
let mut filler_fut = pin!(tx.send(0));
match poll_once(&mut filler_fut) {
Poll::Ready(Ok(())) => {}
other => panic!("filler send should resolve immediately: {other:?}"),
}
let clones: std::vec::Vec<_> = (0..3).map(|_| tx.clone()).collect();
let mut futs: std::vec::Vec<_> = clones
.iter()
.enumerate()
.map(|(i, c)| Box::pin(c.send(u32::try_from(i).unwrap() + 1)))
.collect();
for f in &mut futs {
match f.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
Poll::Pending => {}
Poll::Ready(other) => panic!("expected Pending, got Ready({other:?})"),
}
}
drop(rx);
for f in &mut futs {
match f.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
Poll::Ready(Err(())) => {}
Poll::Ready(Ok(())) => {
panic!("expected Err after receiver drop on cloned sender, got Ok")
}
Poll::Pending => panic!("expected Err after receiver drop, got Pending"),
}
}
}
#[test]
fn mpsc_concurrent_first_claim_does_not_panic() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering as O};
static POOL: MpscPool<u32, 8, 4> = MpscPool::new();
let success_count = Arc::new(AtomicUsize::new(0));
let mut handles = std::vec::Vec::new();
for _ in 0..4 {
let s = Arc::clone(&success_count);
handles.push(std::thread::spawn(move || {
if POOL.claim_bounded().is_some() {
s.fetch_add(1, O::SeqCst);
}
}));
}
for h in handles {
h.join().unwrap();
}
assert_eq!(
success_count.load(O::SeqCst),
4,
"all 4 concurrent claims should have succeeded against an 8-slot pool",
);
}
static MPSC_POOL: MpscPool<u32, 2, 4> = MpscPool::new();
#[test]
fn mpsc_bounded_send_recv() {
let (tx, mut rx) = MPSC_POOL.claim_bounded().expect("pool not empty");
let mut send_fut = pin!(tx.send(7));
assert!(matches!(poll_once(&mut send_fut), Poll::Ready(Ok(()))));
let mut recv_fut = pin!(rx.recv());
match poll_once(&mut recv_fut) {
Poll::Ready(Some(7)) => {}
other => panic!("expected ready Some(7), got {other:?}"),
}
}
#[test]
fn mpsc_bounded_clone_then_drop_all_closes_receiver() {
static POOL: MpscPool<u32, 1, 2> = MpscPool::new();
let (tx, mut rx) = POOL.claim_bounded().expect("pool not empty");
let tx2 = tx.clone();
drop(tx);
{
let mut recv_fut = pin!(rx.recv());
assert!(matches!(poll_once(&mut recv_fut), Poll::Pending));
}
drop(tx2);
let mut recv_fut = pin!(rx.recv());
match poll_once(&mut recv_fut) {
Poll::Ready(None) => {}
other => panic!("expected ready None, got {other:?}"),
}
}
#[test]
fn unbounded_send_now_returns_full_when_capacity_exhausted() {
static POOL: MpscPool<u32, 1, 2> = MpscPool::new();
let (tx, _rx) = POOL.claim_unbounded().expect("pool not empty");
assert!(tx.send_now(1).is_ok());
assert!(tx.send_now(2).is_ok());
match tx.send_now(3) {
Err(3) => {}
other => panic!("expected Err(3), got {other:?}"),
}
}
crate::define_static_channels! {
name: MacroTestChannels,
oneshot: [
(u32, 4),
(Result<i32, ()>, 2),
],
bounded: [
((u8, 4), 2),
],
unbounded: [
(u16, 1),
],
}
#[test]
fn macro_oneshot_dispatches_through_factory() {
use crate::transport::{ChannelFactory, OneshotSend};
let (tx, rx) = MacroTestChannels::oneshot::<u32>();
tx.send(99).unwrap();
let mut fut = pin!(<_ as crate::transport::OneshotRecv<u32>>::recv(rx));
match poll_once(&mut fut) {
Poll::Ready(Ok(99)) => {}
other => panic!("expected ready Ok(99), got {other:?}"),
}
}
#[test]
fn macro_bounded_dispatches_through_factory() {
use crate::transport::{ChannelFactory, MpscRecv, MpscSend};
let (tx, mut rx) = MacroTestChannels::bounded::<u8, 4>();
{
let mut send_fut = pin!(tx.send(7));
assert!(matches!(poll_once(&mut send_fut), Poll::Ready(Ok(()))));
}
let mut recv_fut = pin!(rx.recv());
match poll_once(&mut recv_fut) {
Poll::Ready(Some(7)) => {}
other => panic!("expected ready Some(7), got {other:?}"),
}
}
#[test]
fn macro_unbounded_dispatches_through_factory() {
use crate::transport::{ChannelFactory, UnboundedSend};
let (tx, _rx) = MacroTestChannels::unbounded::<u16>();
assert!(tx.send_now(1234).is_ok());
}
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering as SAtomic};
struct WakeFlag(AtomicBool);
impl std::task::Wake for WakeFlag {
fn wake(self: Arc<Self>) {
self.0.store(true, SAtomic::Release);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.store(true, SAtomic::Release);
}
}
fn tracking_waker() -> (Arc<WakeFlag>, Waker) {
let flag = Arc::new(WakeFlag(AtomicBool::new(false)));
let waker = Waker::from(flag.clone());
(flag, waker)
}
#[test]
fn oneshot_waker_fires_on_send() {
static POOL: OneshotPool<u32, 2> = OneshotPool::new();
let (tx, rx) = POOL.claim().expect("pool not empty");
let (flag, waker) = tracking_waker();
let mut cx = Context::from_waker(&waker);
let mut fut = pin!(rx.recv());
assert!(matches!(fut.as_mut().poll(&mut cx), Poll::Pending));
tx.send(42u32).unwrap();
assert!(
flag.0.load(SAtomic::Acquire),
"waker must fire when value is sent"
);
let noop = Waker::noop();
let mut cx2 = Context::from_waker(noop);
assert!(matches!(fut.as_mut().poll(&mut cx2), Poll::Ready(Ok(42))));
}
#[test]
fn oneshot_cancel_waker_fires_on_sender_drop() {
static POOL: OneshotPool<u32, 2> = OneshotPool::new();
let (tx, rx) = POOL.claim().expect("pool not empty");
let (flag, waker) = tracking_waker();
let mut cx = Context::from_waker(&waker);
let mut fut = pin!(rx.recv());
assert!(matches!(fut.as_mut().poll(&mut cx), Poll::Pending));
drop(tx);
assert!(
flag.0.load(SAtomic::Acquire),
"waker must fire when sender is dropped (cancel)"
);
let noop = Waker::noop();
let mut cx2 = Context::from_waker(noop);
assert!(matches!(
fut.as_mut().poll(&mut cx2),
Poll::Ready(Err(OneshotCancelled))
));
}
#[test]
fn mpsc_close_waker_fires_on_all_senders_drop() {
static POOL: MpscPool<u32, 1, 4> = MpscPool::new();
let (tx, mut rx) = POOL.claim_bounded().expect("pool not empty");
let tx2 = tx.clone();
let (flag, waker) = tracking_waker();
let mut cx = Context::from_waker(&waker);
let mut fut = pin!(rx.recv());
assert!(matches!(fut.as_mut().poll(&mut cx), Poll::Pending));
drop(tx);
assert!(
!flag.0.load(SAtomic::Acquire),
"waker must not fire until last sender drops"
);
drop(tx2);
assert!(
flag.0.load(SAtomic::Acquire),
"waker must fire when last sender drops"
);
let noop = Waker::noop();
let mut cx2 = Context::from_waker(noop);
assert!(matches!(fut.as_mut().poll(&mut cx2), Poll::Ready(None)));
}
#[test]
fn mpsc_bounded_pool_exhaustion_returns_none() {
static POOL: MpscPool<u32, 1, 4> = MpscPool::new();
let _a = POOL.claim_bounded().expect("pool not empty");
assert!(
POOL.claim_bounded().is_none(),
"second claim must exhaust pool of size 1"
);
}
#[test]
fn oneshot_send_after_receiver_drop_returns_err() {
static POOL: OneshotPool<u32, 2> = OneshotPool::new();
let (tx, rx) = POOL.claim().expect("pool not empty");
drop(rx);
match tx.send(42) {
Err(42) => {}
other => panic!("expected Err(42) after receiver drop, got {other:?}"),
}
}
#[test]
fn unbounded_send_now_after_receiver_drop_returns_err() {
static POOL: MpscPool<u32, 1, 4> = MpscPool::new();
let (tx, rx) = POOL.claim_unbounded().expect("pool not empty");
drop(rx);
match tx.send_now(7) {
Err(7) => {}
other => panic!("expected Err(7) after receiver drop, got {other:?}"),
}
}
#[test]
fn bounded_send_unblocks_with_err_on_receiver_drop() {
static POOL: MpscPool<u32, 1, 1> = MpscPool::new();
let (tx, rx) = POOL.claim_bounded().expect("pool not empty");
{
let mut send_fut = pin!(tx.send(1));
assert!(matches!(poll_once(&mut send_fut), Poll::Ready(Ok(()))));
}
let mut send_fut = pin!(tx.send(2));
let (flag, waker) = tracking_waker();
let mut cx = Context::from_waker(&waker);
assert!(matches!(send_fut.as_mut().poll(&mut cx), Poll::Pending));
drop(rx);
assert!(
flag.0.load(SAtomic::Acquire),
"send_waker must fire when receiver drops while sender is awaiting"
);
let noop = Waker::noop();
let mut cx2 = Context::from_waker(noop);
match send_fut.as_mut().poll(&mut cx2) {
Poll::Ready(Err(())) => {}
other => panic!("expected Err(()) after receiver drop, got {other:?}"),
}
}
#[test]
fn bounded_send_after_receiver_drop_returns_err_fast_path() {
static POOL: MpscPool<u32, 1, 4> = MpscPool::new();
let (tx, rx) = POOL.claim_bounded().expect("pool not empty");
drop(rx);
let mut send_fut = pin!(tx.send(99));
match poll_once(&mut send_fut) {
Poll::Ready(Err(())) => {}
other => panic!("expected Err(()) on closed slot, got {other:?}"),
}
}
}