odem_rs_sync/channel/

shared.rs

//! This module provides dynamically allocated, reference-counted channels for
//! exchanging values between concurrent processes in a discrete-event
//! simulation.
//!
//! Unlike stack-bound channels, these channels use [`Rc`] for shared ownership,
//! allowing senders and receivers to be freely cloned and distributed.
//!
//! ## Components
//!
//! - **[`channel`]**: Creates a new channel, returning a [`Sender`] and a
//!   [`Receiver`] pair. Messages are delivered in the order they were sent,
//!   and the channel closes automatically when both ends are dropped.
//! - **[`Sender`]**: The sending half of the channel, which can be cloned.
//!   Messages sent through it become available to receivers.
//! - **[`Receiver`]**: The receiving half of the channel, which can be cloned.
//!   Receives messages sent by any sender.
//!
//! ## Features
//!
//! - **Dynamically Allocated**: Uses [`Rc`] and [`Weak`] for reference-counted
//!   channel management.
//! - **Non-blocking**: Sending and receiving operations return futures that
//!   can be awaited asynchronously within a simulation.
//! - **Reference Counting**: The channel closes automatically when all senders
//!   and receivers are dropped.

use core::{
	future::Future,
	pin::Pin,
	task::{Context, Poll},
};

use alloc::rc::{Rc, Weak};

use super::{
	Inner, RingBuf, SendError, TryRecvError, TrySendError,
	pinned::{SendFutureState, SendFutureStateProject},
};

#[doc(inline)]
pub use super::pinned::RecvFuture;

/* **************************************************** Free Channel Function */

/// Creates a new channel, returning the sender/receiver halves.
///
/// All data sent on the [`Sender`] half will become available on the
/// [`Receiver`] in the same order as it was sent. If the underlying [`RingBuf`]
/// is dynamically sized, no [`send`] will block. The [`recv`]-method will block
/// until a message is available while there is at least one `Sender` alive.
///
/// Both `Sender` and `Receiver` may be cloned, allowing to receive and send
/// messages over the same channel from different parts of the simulation.
///
/// The underlying channel is closed automatically when the last `Sender`
/// or `Receiver` drops.
///
/// ## Examples
///
/// An ordinary channel transferring items between two jobs.
///
/// ```
/// # use core::pin::pin;
/// # use odem_rs_core::simulator::{Sim, simulation};
/// # use odem_rs_sync::{channel::{channel, Rendezvous}, fork::ForkExt};
/// #[derive(Debug)]
/// struct Message(usize);
///
/// async fn sim_main(sim: &Sim) {
///     let (sx, rx) = channel(Rendezvous::new());
///
///     sim.fork(async move {
///         sx.send(Message(42)).await.unwrap();
///     })
///     .and(async move {
///         let msg = rx.recv().await.unwrap();
///         println!("Received {:?}", msg);
///     })
///     .await;
/// }
/// # fn main() { simulation(sim_main).ok(); }
/// ```
///
/// Sending references to items is allowed, as long as they outlive the channel.
///
/// ```
/// # use core::pin::pin;
/// # use odem_rs_core::simulator::{Sim, simulation};
/// # use odem_rs_sync::{channel::{channel, Rendezvous}, fork::ForkExt};
/// async fn sim_main(sim: &Sim) {
///     let item = 3;
///     let (sx, rx) = channel(Rendezvous::new());
///
///     sim.fork(async {
///         sx.send(&item).await.unwrap();
///     })
///     .and(async {
///         let msg = rx.recv().await.unwrap();
///         println!("Received {:?}", msg);
///     })
///     .await;
/// }
/// # fn main() { simulation(sim_main).ok(); }
/// ```
///
/// [`send`]: Sender::send
/// [`recv`]: Receiver::recv
pub fn channel<R: RingBuf>(buffer: R) -> (Sender<R::Item>, Receiver<R::Item>) {
	use alloc::rc::Rc;
	let channel = Rc::new(Inner::new(buffer));

	// SAFETY: Any lifetime bound on `R` can only result from a bound on
	// `R::Item`, since `Rc<Channel<R>>` safely erases the bound on the
	// underlying ring buffer itself. The bound on `R::Item` is still
	// retained within the `Sender`/`Receiver` pairs generic argument,
	// making it subject to the borrow checker.
	// Therefore, erasing the lifetime here is safe.
	let channel = unsafe {
		core::mem::transmute::<
			Rc<Inner<R::Item, dyn RingBuf<Item = R::Item> + '_>>,
			Rc<Inner<R::Item, dyn RingBuf<Item = R::Item> + 'static>>,
		>(channel)
	};

	(Sender(Rc::downgrade(&channel)), Receiver(channel))
}

/* ***************************************************** Sender/Receiver Pair */

/// The sending half of the channel. This half can be cloned and is
/// bound to the hosting channel through the reference counted [`Rc`].
pub struct Sender<T>(pub(super) Weak<Inner<T, dyn RingBuf<Item = T>>>);

impl<T> Sender<T> {
	/// Attempt to send the value without waiting.
	///
	/// This can fail for two reasons:
	/// 1. The channel is closed, i.e. [`close`] has been called.
	/// 2. The message buffer is full.
	///
	/// The `Err`-result contains the initially provided item in either case.
	///
	/// [`close`]: Self::close
	pub fn try_send(&self, item: T) -> Result<(), TrySendError<T>> {
		match self.0.upgrade() {
			Some(inner) => inner.try_send(item),
			_ => Err(TrySendError::Disconnected(item)),
		}
	}

	/// Returns a future that gets fulfilled when the value has been written to
	/// the channel. If the channel gets closed while the sending is in
	/// progress, sending the value will fail, and the future will deliver the
	/// value back.
	pub fn send(&self, item: T) -> SendFuture<'_, T> {
		SendFuture::new(self, item)
	}

	/// Returns the number of elements currently inside the channel of this
	/// sender.
	pub fn count(&self) -> isize {
		self.0.upgrade().map_or(0, |inner| inner.count())
	}

	/// Closes the channel from the sending side. Already sent values may
	/// still be received.
	pub fn close(&self) {
		if let Some(inner) = self.0.upgrade() {
			inner.close();
		}
	}
}

impl<T> Clone for Sender<T> {
	fn clone(&self) -> Self {
		Sender(self.0.clone())
	}
}

impl<T> Drop for Sender<T> {
	fn drop(&mut self) {
		if Weak::weak_count(&self.0) == 1 {
			if let Some(inner) = self.0.upgrade() {
				inner.close();
			}
		}
	}
}

/// The receiving half of the channel. This half can be cloned and is
/// bound to the hosting channel through the reference counted [`Rc`].
pub struct Receiver<T>(pub(super) Rc<Inner<T, dyn RingBuf<Item = T>>>);

impl<T> Receiver<T> {
	/// Attempts to receive a value of the channel without waiting.
	///
	/// This can fail for two reasons:
	/// 1. The channel is closed, i.e. [`close`] has been called.
	/// 2. The message buffer is empty.
	///
	/// [`close`]: Self::close
	pub fn try_recv(&self) -> Result<T, TryRecvError> {
		self.0.try_recv()
	}

	/// Returns a future that gets fulfilled when a value is written to the
	/// channel. If the channels get closed, the future will resolve to an
	/// `Err` variant.
	pub fn recv(&self) -> RecvFuture<'_, dyn RingBuf<Item = T>> {
		RecvFuture::new(&self.0)
	}

	/// Returns the number of elements currently inside the channel of this
	/// receiver.
	pub fn count(&self) -> isize {
		self.0.count()
	}

	/// Closes the channel from the receiving side. Already sent values may
	/// still be received.
	pub fn close(&self) {
		self.0.close();
	}
}

impl<T> Clone for Receiver<T> {
	fn clone(&self) -> Self {
		Receiver(self.0.clone())
	}
}

impl<T> Drop for Receiver<T> {
	fn drop(&mut self) {
		// notify all the senders of the channel's imminent closing
		if Rc::strong_count(&self.0) == 1 {
			self.close();
		}
	}
}

/* *************************************************** Specialized SendFuture */

/// Future for the [send](Sender::send) operation of the shared channel.
///
/// We cannot re-use the SendFuture of the channel because our sender only
/// contains a weak reference to the channel.
#[pin_project::pin_project(PinnedDrop)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct SendFuture<'r, T> {
	/// The sender used to send the message.
	chan: &'r Sender<T>,
	/// Either the message or the link into the pending chain.
	#[pin]
	state: SendFutureState<T>,
}

impl<'r, T> SendFuture<'r, T> {
	/// Creates a [send](Sender::send)-future and primes it with a message.
	///
	/// The send-operation is attempted once this future is resolved.
	const fn new(chan: &'r Sender<T>, msg: T) -> Self {
		SendFuture {
			chan,
			state: SendFutureState::Primed(Some(msg)),
		}
	}
}

impl<T> Future for SendFuture<'_, T> {
	type Output = Result<(), SendError<T>>;

	fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
		let mut inner = self.project();

		match inner.state.as_mut().project() {
			SendFutureStateProject::Primed(msg) => {
				// extract the message
				let msg = msg.take().unwrap();

				// attempt to get an owning reference to the channel
				match inner.chan.0.upgrade() {
					Some(chan) => {
						// attempt to send the message
						match chan.try_send(msg) {
							Err(TrySendError::Full(msg)) => {
								// if the message queue is already full,
								// create a link and subscribe
								unsafe {
									inner.state.link(&chan.chain, cx.waker().clone(), msg);
								}

								// increase the number of pending senders
								chan.add_pending_sender();

								// defer
								Poll::Pending
							}
							Err(TrySendError::Disconnected(msg)) => {
								// the channel has been closed
								// return an error
								Poll::Ready(Err(SendError(msg)))
							}
							Ok(()) => {
								// the message was sent successfully
								// return with success
								Poll::Ready(Ok(()))
							}
						}
					}
					_ => {
						// failed: the channel is closed and dropped
						// return with an error
						Poll::Ready(Err(SendError(msg)))
					}
				}
			}
			SendFutureStateProject::Linked(link) => {
				// this future has been reactivated after subscribing
				// to the pending chain

				// either the message has been received (success)
				// or the channel has been closed (failure)
				Poll::Ready(match link.note.take() {
					Some(msg) => Err(SendError(msg)),
					_ => Ok(()),
				})
			}
		}
	}
}

#[pin_project::pinned_drop]
impl<T> PinnedDrop for SendFuture<'_, T> {
	fn drop(self: Pin<&mut Self>) {
		// make sure to unlink the soon-to-be-dangling link if the future has
		// been dropped after subscribing but before reactivation
		if let Some(chan) = self.chan.0.upgrade() {
			if unsafe { self.as_ref().project_ref().state.unlink(&chan.chain) } {
				chan.sub_pending_sender();
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use crate::{
		channel::{ArrayBuf, Rendezvous},
		fork::ForkExt as _,
	};
	use odem_rs_core::simulator::{Sim, Simulator};
	use proptest::prelude::*;

	async fn test_insertion_order(
		sim: &Sim,
		buf: impl RingBuf<Item = usize> + 'static,
		elems: usize,
	) {
		let (sx, rx) = channel(buf);

		sim.fork(async move {
			for i in 0..elems {
				sx.send(i).await.unwrap();
				sim.advance(1.0).await;
			}
		})
		.and(async move {
			sim.advance(10.0).await;
			for exp in 0..elems {
				let val = rx.recv().await.unwrap();
				assert_eq!(val, exp);
			}
		})
		.await;
	}

	proptest! {
		#[test]
		#[cfg_attr(miri, ignore)]
		fn rendezvous_channels_respect_insertion_order(elems in 1..100usize) {
			Simulator::default()
				.run(async |sim| {
					test_insertion_order(sim, Rendezvous::default(), elems).await;
				})
				.expect("unexpected deadlock");
		}

		#[test]
		#[cfg_attr(miri, ignore)]
		fn array_channels_respect_insertion_order(elems in 1..100usize) {
			Simulator::default()
				.run(async |sim| {
					test_insertion_order(sim, ArrayBuf::<2,_>::new(), elems).await;
				})
				.expect("unexpected deadlock");
		}

		#[test]
		#[cfg_attr(miri, ignore)]
		fn deque_channels_respect_insertion_order(elems in 1..100usize) {
			use alloc::collections::VecDeque;
			Simulator::default()
				.run(async |sim| {
					test_insertion_order(sim, VecDeque::new(), elems).await;
				})
				.expect("unexpected deadlock");
		}
	}
}