use std::fmt;
use tokio::sync::{mpsc, oneshot};
pub struct WithAck<T> {
pub data: T,
ack: Option<oneshot::Sender<()>>,
}
impl<T> WithAck<T> {
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() {
let _ = ack.send(());
}
}
}
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 }
}
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(())
}
}
#[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 {}
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
}
}
pub fn channel<T>(capacity: usize) -> (Sender<T>, Receiver<T>) {
let (tx, rx) = mpsc::channel(capacity);
(Sender::new(tx), Receiver::new(rx))
}