mahler-core 0.25.5

An automated job orchestration library that builds and executes dynamic workflows
Documentation
use std::fmt;
use tokio::sync::{mpsc, oneshot};

/// A message that requires reception acknowledgement
pub struct WithAck<T> {
    pub data: T,
    ack: Option<oneshot::Sender<()>>,
}

impl<T> WithAck<T> {
    /// Manually acknowledge the message
    pub fn ack(mut self) {
        if let Some(ack) = self.ack.take() {
            let _ = ack.send(());
        }
    }
}

impl<T> Drop for WithAck<T> {
    fn drop(&mut self) {
        if let Some(ack) = self.ack.take() {
            // Ack if the message is dropped to avoid
            // blocking the sender
            let _ = ack.send(());
        }
    }
}

/// An acknowledged sender
pub struct Sender<T> {
    inner: mpsc::Sender<WithAck<T>>,
}

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

impl<T> Sender<T> {
    fn new(inner: mpsc::Sender<WithAck<T>>) -> Self {
        Sender { inner }
    }

    /// Sends a message and waits for acknowledgment
    pub async fn send(&self, data: T) -> Result<(), SendError> {
        let (ack_tx, ack_rx) = oneshot::channel();
        self.inner
            .send(WithAck {
                data,
                ack: Some(ack_tx),
            })
            .await
            .map_err(|_| SendError)?;
        ack_rx.await.map_err(|_| SendError)?;
        Ok(())
    }
}

/// Possible errors when sending
#[derive(Debug)]
pub struct SendError;

impl fmt::Display for SendError {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(fmt, "send error")
    }
}

impl std::error::Error for SendError {}

/// A receiver of acknowledged messages
pub struct Receiver<T> {
    inner: mpsc::Receiver<WithAck<T>>,
}

impl<T> Receiver<T> {
    fn new(inner: mpsc::Receiver<WithAck<T>>) -> Self {
        Receiver { inner }
    }

    pub async fn recv(&mut self) -> Option<WithAck<T>> {
        self.inner.recv().await
    }
}

/// Create a new acknowledged channel with the specified capacity
///
/// A sender in an acknowledged channel will wait for the sent message
/// to be acknowledged before releasing the sender
pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
    let (tx, rx) = mpsc::channel(capacity);
    (Sender::new(tx), Receiver::new(rx))
}