zlink_smol/notified/
state.rs

1use std::{
2    fmt::Debug,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use crate::Reply;
8use async_broadcast::{
9    InactiveReceiver, Receiver as BroadcastReceiver, Sender as BroadcastSender, broadcast,
10};
11use pin_project_lite::pin_project;
12
13/// A notified state (e.g a field) of a service implementation.
14#[derive(Debug, Clone)]
15pub struct State<T, ReplyParams> {
16    value: T,
17    tx: BroadcastSender<ReplyParams>,
18    // Keep an inactive receiver to prevent the channel from closing.
19    inactive_rx: InactiveReceiver<ReplyParams>,
20}
21
22impl<T, ReplyParams> zlink_core::notified::State<T, ReplyParams> for State<T, ReplyParams>
23where
24    T: Into<ReplyParams> + Clone + Debug + Send,
25    ReplyParams: Clone + Send + 'static + Debug,
26{
27    type Stream = Stream<ReplyParams>;
28
29    /// Create a new notified field.
30    fn new(value: T) -> Self {
31        let (mut tx, rx) = broadcast(1);
32        // Notification broadcast shouldn't await active subscribers.
33        tx.set_await_active(false);
34        // Enable overflow mode because:
35        // 1. We don't need to ensure that subscribers receive all values, as long as they always
36        //    receive the latest value so we don't want the broadcast to wait for receivers.
37        // 2. This would be consistent with the behavior of the `zlink_tokio::notified::State`.
38        tx.set_overflow(true);
39        // Deactivate the initial receiver to keep the channel open without consuming buffer space.
40        let inactive_rx = rx.deactivate();
41
42        Self {
43            value,
44            tx,
45            inactive_rx,
46        }
47    }
48
49    /// Set the value of the notified field and notify all listeners.
50    async fn set(&mut self, value: T) {
51        self.value = value.clone();
52        self.tx
53            .broadcast_direct(value.into())
54            .await
55            // Since we enabled overflow and disabled awaiting active receivers, this can't fail.
56            .expect("Failed to broadcast value");
57    }
58
59    /// The value of the notified field.
60    fn get(&self) -> T {
61        self.value.clone()
62    }
63
64    /// A stream of replies for the notified field.
65    fn stream(&self) -> Stream<ReplyParams> {
66        Stream {
67            inner: self.inactive_rx.activate_cloned(),
68            cached: None,
69            once: false,
70        }
71    }
72
73    /// A stream of replies for this state, that only yields one reply: the current state.
74    fn stream_once(&self) -> Stream<ReplyParams> {
75        Stream {
76            inner: self.inactive_rx.activate_cloned(),
77            cached: Some(self.get().into()),
78            once: true,
79        }
80    }
81}
82
83pin_project! {
84    /// The stream to use as the [`zlink_core::Service::ReplyStream`] in service implementation when
85    /// using [`State`].
86    #[derive(Debug)]
87    pub struct Stream<ReplyParams> {
88        #[pin]
89        inner: BroadcastReceiver<ReplyParams>,
90        cached: Option<ReplyParams>,
91        once: bool,
92    }
93}
94
95impl<ReplyParams> futures_util::Stream for Stream<ReplyParams>
96where
97    ReplyParams: Clone + Send + 'static,
98{
99    type Item = Reply<ReplyParams>;
100
101    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
102        let this = self.project();
103        if *this.once {
104            return Poll::Ready(
105                this.cached
106                    .take()
107                    .map(|reply| Reply::new(Some(reply)).set_continues(Some(false))),
108            );
109        }
110        match futures_util::ready!(this.inner.poll_next(cx)) {
111            Some(reply) => {
112                // Cache and yield immediately with continues=true.
113                *this.cached = Some(reply.clone());
114                Poll::Ready(Some(Reply::new(Some(reply)).set_continues(Some(true))))
115            }
116            // Channel closed - yield cached value with continues=false.
117            None => Poll::Ready(
118                this.cached
119                    .take()
120                    .map(|reply| Reply::new(Some(reply)).set_continues(Some(false))),
121            ),
122        }
123    }
124}
zlink_smol/notified/state.rs

zlink_smol/notified/
state.rs
