Skip to main content

polybox_core/
oneshot.rs

1use futures::{Future, FutureExt};
2use std::{
3    pin::Pin,
4    task::{Context, Poll},
5};
6use tokio::sync::oneshot;
7
8/// Create a new request, consisting of a [`Tx<T>`] and an [`Rx<T>`].
9/// The `Tx` (_transmitter_) can be used to send a single message `T` to the `Rx` (_receiver_).
10///
11/// This is just a wrapper around a [`tokio::sync::oneshot`] channel.
12pub fn new_request<T>() -> (Tx<T>, Rx<T>) {
13    let (tx, rx) = oneshot::channel();
14    (Tx(tx), Rx(rx))
15}
16
17//------------------------------------------------------------------------------------------------
18//  Tx
19//------------------------------------------------------------------------------------------------
20
21/// The transmitter part of a request, created with [`new_request`].
22///
23/// This implements [`MessageDerive<M>`] to be used with the [`derive@Message`] derive macro.
24#[derive(Debug)]
25pub struct Tx<M>(oneshot::Sender<M>);
26
27impl<M> Tx<M> {
28    /// Send a message.
29    pub fn send(self, msg: M) -> Result<(), TxError<M>> {
30        self.0.send(msg).map_err(|msg| TxError(msg))
31    }
32
33    /// Whether the [`Rx`] has closed/dropped the oneshot-channel.
34    pub fn is_closed(&self) -> bool {
35        self.0.is_closed()
36    }
37
38    /// Wait for the [`Rx`] to close/drop the oneshot-channel.
39    pub async fn closed(&mut self) {
40        self.0.closed().await
41    }
42}
43
44//  Rx
45//------------------------------------------------------------------------------------------------
46
47/// The receiver part of a request, created with [`new_request`].
48///
49/// This implements [`MessageDerive<M>`] to be used with the [`derive@Message`] derive macro.
50#[derive(Debug)]
51pub struct Rx<M>(oneshot::Receiver<M>);
52
53impl<M> Rx<M> {
54    /// Attempt to take the message out, if it exists.
55    pub fn try_recv(&mut self) -> Result<M, TryRxError> {
56        self.0.try_recv().map_err(|e| e.into())
57    }
58
59    /// Block the thread while waiting for the message.
60    pub fn recv_blocking(self) -> Result<M, RxError> {
61        self.0.blocking_recv().map_err(|e| e.into())
62    }
63
64    /// Close the oneshot-channel, preventing the [`Tx`] from sending a message.
65    pub fn close(&mut self) {
66        self.0.close()
67    }
68}
69
70impl<M> Unpin for Rx<M> {}
71
72impl<M> Future for Rx<M> {
73    type Output = Result<M, RxError>;
74
75    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
76        self.0.poll_unpin(cx).map_err(|e| e.into())
77    }
78}
79
80//------------------------------------------------------------------------------------------------
81//  Errors
82//------------------------------------------------------------------------------------------------
83
84/// Error returned when receiving a message using an [`Rx`].
85#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash, thiserror::Error)]
86#[error("Failed to receive from Rx because it is closed.")]
87pub struct RxError;
88
89impl From<oneshot::error::RecvError> for RxError {
90    fn from(_: oneshot::error::RecvError) -> Self {
91        Self
92    }
93}
94
95/// Error returned when trying to receive a message using an [`Rx`].
96#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash, thiserror::Error)]
97pub enum TryRxError {
98    #[error("Closed")]
99    Closed,
100    #[error("Empty")]
101    Empty,
102}
103
104impl From<oneshot::error::TryRecvError> for TryRxError {
105    fn from(e: oneshot::error::TryRecvError) -> Self {
106        match e {
107            oneshot::error::TryRecvError::Empty => Self::Empty,
108            oneshot::error::TryRecvError::Closed => Self::Closed,
109        }
110    }
111}
112
113/// Error returned when sending a message using a [`Tx`].
114#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash, thiserror::Error)]
115#[error("Failed to send to Tx because it is closed.")]
116pub struct TxError<M>(pub M);