molo_agent/message_channel/watch.rs
1//! A latest-value change notification channel: an implementation based on tokio watch.
2
3use super::{ChannelError, IncomingMessage, MessageChannel};
4use tokio::sync::{Mutex, watch};
5
6/// Observes changes of the latest state value: holds and updates one value, and all observers
7/// ([`WatchChannel::subscribe`]) are notified when it changes.
8///
9/// The semantics are "latest value", not a message queue:
10///
11/// - observers are only guaranteed to eventually see the latest value, not every intermediate change — values
12/// overwritten between two `notify` calls are never re-sent;
13/// - unbounded: a new value simply replaces the old one; there is no queue capacity, and no backpressure;
14/// - observers are independent of each other; how fast one consumes does not affect the others.
15///
16/// Driven by the [`MessageChannel`] interface:
17///
18/// - [`MessageChannel::notify`] — update the value to the latest message; all observers get notified;
19/// - [`MessageChannel::ask`] — not supported (watch has no "reply" concept), returns
20/// [`ChannelError::NotSupported`];
21/// - receive side: [`WatchReceiver::recv`] waits for a change and takes the latest value.
22///
23/// Implements [`Clone`] — every clone is a sender of the same channel, and any clone's `notify`
24/// notifies all observers.
25///
26/// # Example
27///
28/// ```rust
29/// # extern crate molo_agent as molo;
30/// use molo::{MessageChannel, WatchChannel};
31///
32/// # #[tokio::main]
33/// # async fn main() -> Result<(), molo::ChannelError> {
34/// let status = WatchChannel::new();
35/// let sub_a = status.subscribe();
36/// let sub_b = status.subscribe();
37///
38/// status.notify("ready").await?;
39/// assert_eq!(sub_a.recv().await?.text(), "ready");
40/// assert_eq!(sub_b.recv().await?.text(), "ready");
41/// # Ok(())
42/// # }
43/// ```
44///
45/// # Notes
46///
47/// - the initial value is an empty string: before the first `notify`, `recv` keeps waiting;
48/// - when there are no observers (never `subscribe`d, or all observers have been dropped), `notify`
49/// returns [`ChannelError::NoReceiver`](crate::ChannelError::NoReceiver) — unlike
50/// [`BroadcastChannel`](crate::BroadcastChannel), which succeeds silently even with nobody listening,
51/// so `subscribe` first, then `notify`; the channel itself is not closed, and `notify` works again
52/// after re-subscribing.
53#[derive(Debug, Clone)]
54pub struct WatchChannel {
55 tx: watch::Sender<String>,
56}
57
58/// The receive end of a watch channel: observes value changes.
59///
60/// Does not implement [`Clone`] — a single receive end can only be driven by one waiter at a time
61/// (serialized by an internal lock). For multiple observers, call [`subscribe`](WatchChannel::subscribe)
62/// on the [`WatchChannel`] separately for each.
63#[derive(Debug)]
64pub struct WatchReceiver {
65 rx: Mutex<watch::Receiver<String>>,
66}
67
68impl WatchChannel {
69 /// Creates a watch channel whose initial value is an empty string.
70 ///
71 /// Before the first `notify`, a subscriber's `recv` keeps waiting.
72 pub fn new() -> Self {
73 let (tx, _rx) = watch::channel(String::new());
74 Self { tx }
75 }
76
77 /// Subscribes to value changes; each observer gets an independent receive end, unaffected by the others.
78 pub fn subscribe(&self) -> WatchReceiver {
79 WatchReceiver {
80 rx: Mutex::new(self.tx.subscribe()),
81 }
82 }
83}
84
85impl Default for WatchChannel {
86 fn default() -> Self {
87 Self::new()
88 }
89}
90
91impl WatchReceiver {
92 /// Waits for a change and takes the latest value.
93 ///
94 /// # Errors
95 ///
96 /// Returns [`ChannelError::Closed`] when the sender ([`WatchChannel`]) has been dropped.
97 ///
98 /// # Cancellation semantics
99 ///
100 /// Being cancelled while waiting has no side effects: watch is not a queue and has no "take one item" concept,
101 /// the next `recv` starts from the latest value and keeps waiting for the next change.
102 pub async fn recv(&self) -> Result<IncomingMessage, ChannelError> {
103 let mut rx = self.rx.lock().await;
104 rx.changed().await.map_err(|_| ChannelError::Closed)?;
105 let message = rx.borrow().clone();
106 Ok(IncomingMessage {
107 text: message,
108 reply_tx: None,
109 })
110 }
111}
112
113#[async_trait::async_trait]
114impl MessageChannel for WatchChannel {
115 async fn ask(&self, _message: &str) -> Result<String, ChannelError> {
116 Err(ChannelError::NotSupported(
117 "watch channel does not support ask".into(),
118 ))
119 }
120
121 async fn notify(&self, message: &str) -> Result<(), ChannelError> {
122 // watch's send only fails when there are no receivers: the channel is not closed and re-subscribing
123 // restores it, so use NoReceiver rather than the semantically misleading Closed.
124 self.tx
125 .send(message.to_string())
126 .map_err(|_| ChannelError::NoReceiver)
127 }
128}
129
130#[cfg(test)]
131mod tests {
132 use super::*;
133
134 #[tokio::test]
135 async fn notify_then_recv_latest() {
136 let channel = WatchChannel::new();
137 let receiver = channel.subscribe();
138 // Alternate notify / recv: each recv waits for one change and takes the latest value.
139 channel.notify("v1").await.unwrap();
140 assert_eq!(receiver.recv().await.unwrap().text(), "v1");
141 channel.notify("v2").await.unwrap();
142 assert_eq!(receiver.recv().await.unwrap().text(), "v2");
143 }
144
145 #[tokio::test]
146 async fn notify_before_recv_wakes_all_observers() {
147 let channel = WatchChannel::new();
148 let sub_a = channel.subscribe();
149 let sub_b = channel.subscribe();
150 channel.notify("start").await.unwrap();
151 let (a, b) = tokio::join!(sub_a.recv(), sub_b.recv());
152 assert_eq!(a.unwrap().text(), "start");
153 assert_eq!(b.unwrap().text(), "start");
154 }
155
156 #[tokio::test]
157 async fn ask_not_supported() {
158 let channel = WatchChannel::new();
159 assert!(matches!(
160 channel.ask("q").await,
161 Err(ChannelError::NotSupported(_))
162 ));
163 }
164}