Skip to main content

mahler_core/sync/
channel.rs

1use std::fmt;
2use tokio::sync::{mpsc, oneshot};
3
4/// A message that requires reception acknowledgement
5pub struct WithAck<T> {
6    pub data: T,
7    ack: Option<oneshot::Sender<()>>,
8}
9
10impl<T> WithAck<T> {
11    /// Manually acknowledge the message
12    pub fn ack(mut self) {
13        if let Some(ack) = self.ack.take() {
14            let _ = ack.send(());
15        }
16    }
17}
18
19impl<T> Drop for WithAck<T> {
20    fn drop(&mut self) {
21        if let Some(ack) = self.ack.take() {
22            // Ack if the message is dropped to avoid
23            // blocking the sender
24            let _ = ack.send(());
25        }
26    }
27}
28
29/// An acknowledged sender
30pub struct Sender<T> {
31    inner: mpsc::Sender<WithAck<T>>,
32}
33
34impl<T> Clone for Sender<T> {
35    fn clone(&self) -> Self {
36        Self {
37            inner: self.inner.clone(),
38        }
39    }
40}
41
42impl<T> Sender<T> {
43    fn new(inner: mpsc::Sender<WithAck<T>>) -> Self {
44        Sender { inner }
45    }
46
47    /// Sends a message and waits for acknowledgment
48    pub async fn send(&self, data: T) -> Result<(), SendError> {
49        let (ack_tx, ack_rx) = oneshot::channel();
50        self.inner
51            .send(WithAck {
52                data,
53                ack: Some(ack_tx),
54            })
55            .await
56            .map_err(|_| SendError)?;
57        ack_rx.await.map_err(|_| SendError)?;
58        Ok(())
59    }
60}
61
62/// Possible errors when sending
63#[derive(Debug)]
64pub struct SendError;
65
66impl fmt::Display for SendError {
67    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
68        write!(fmt, "send error")
69    }
70}
71
72impl std::error::Error for SendError {}
73
74/// A receiver of acknowledged messages
75pub struct Receiver<T> {
76    inner: mpsc::Receiver<WithAck<T>>,
77}
78
79impl<T> Receiver<T> {
80    fn new(inner: mpsc::Receiver<WithAck<T>>) -> Self {
81        Receiver { inner }
82    }
83
84    pub async fn recv(&mut self) -> Option<WithAck<T>> {
85        self.inner.recv().await
86    }
87}
88
89/// Create a new acknowledged channel with the specified capacity
90///
91/// A sender in an acknowledged channel will wait for the sent message
92/// to be acknowledged before releasing the sender
93pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
94    let (tx, rx) = mpsc::channel(capacity);
95    (Sender::new(tx), Receiver::new(rx))
96}