forge
ExploreLog in

Tahinli/forge

public
0 watching
★ 0
6.5 KiB
1//! Live unread counts for the bell. [`NotifyHub`] keeps one `watch` channel per signed-in user who
2//! has an open `GET /notifications/stream`. A write that changes a user's inbox publishes the new
3//! count, and each open stream of that user sends it. Nothing polls: a stream waits on its channel.
4//! The hub is in-process only, so a second server process would not see the change.
5//!
6//! A channel with no receivers left is removed, so the map holds only users who are connected.
7
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
10
11use tokio::sync::watch;
12
13use crate::state::AppState;
14
15/// The users with an open stream, each with the channel that carries its unread count.
16pub struct NotifyHub {
17 users: Mutex<HashMap<String, watch::Sender<u64>>>,
18 /// Set once on server shutdown. Every open stream ends when it sees this.
19 closing: watch::Sender<bool>,
20}
21
22impl NotifyHub {
23 pub fn new() -> Self {
24 Self {
25 users: Mutex::new(HashMap::new()),
26 closing: watch::channel(false).0,
27 }
28 }
29
30 fn users(&self) -> MutexGuard<'_, HashMap<String, watch::Sender<u64>>> {
31 self.users.lock().unwrap_or_else(PoisonError::into_inner)
32 }
33
34 /// Opens a subscription to the counts published for `user_id`. The first count is not sent
35 /// here: the caller reads it from the store after subscribing, so a change in between is
36 /// seen by the subscription and sent again rather than lost.
37 pub fn subscribe(self: &Arc<Self>, user_id: &str) -> Subscription {
38 let rx = self
39 .users()
40 .entry(user_id.to_owned())
41 .or_insert_with(|| watch::channel(0).0)
42 .subscribe();
43 Subscription {
44 rx: Some(rx),
45 closing: self.closing.subscribe(),
46 hub: Arc::clone(self),
47 user_id: user_id.to_owned(),
48 }
49 }
50
51 /// True when `user_id` has an open stream. A write skips the count query for anyone else.
52 pub fn is_watched(&self, user_id: &str) -> bool {
53 self.users()
54 .get(user_id)
55 .is_some_and(|tx| tx.receiver_count() > 0)
56 }
57
58 /// Sends `count` to every open stream of `user_id`. Drops the channel when none is left.
59 pub fn publish(&self, user_id: &str, count: u64) {
60 let mut users = self.users();
61 let Some(tx) = users.get(user_id) else {
62 return;
63 };
64 if tx.receiver_count() == 0 {
65 users.remove(user_id);
66 return;
67 }
68 tx.send_replace(count);
69 }
70
71 /// Ends every open stream. Called once when the server shuts down.
72 pub fn close(&self) {
73 self.closing.send_replace(true);
74 }
75
76 /// Removes the user's channel once its last subscription is gone.
77 fn release(&self, user_id: &str) {
78 let mut users = self.users();
79 if users
80 .get(user_id)
81 .is_some_and(|tx| tx.receiver_count() == 0)
82 {
83 users.remove(user_id);
84 }
85 }
86}
87
88/// One open stream's view of its user's count. Dropping it releases the user's channel when no
89/// other stream of that user is open.
90pub struct Subscription {
91 /// Taken in `drop` before the hub is told, so the hub counts this receiver as gone.
92 rx: Option<watch::Receiver<u64>>,
93 closing: watch::Receiver<bool>,
94 hub: Arc<NotifyHub>,
95 user_id: String,
96}
97
98impl Subscription {
99 /// Waits for the next published count. `None` once the hub is closing (server shutdown).
100 pub async fn next_count(&mut self) -> Option<u64> {
101 loop {
102 if *self.closing.borrow() {
103 return None;
104 }
105 let rx = self.rx.as_mut()?;
106 tokio::select! {
107 changed = rx.changed() => {
108 changed.ok()?;
109 return Some(*rx.borrow_and_update());
110 }
111 changed = self.closing.changed() => {
112 changed.ok()?;
113 }
114 }
115 }
116 }
117}
118
119impl Drop for Subscription {
120 fn drop(&mut self) {
121 self.rx = None;
122 self.hub.release(&self.user_id);
123 }
124}
125
126/// Publishes `user_id`'s current unread count to that user's open streams. Call it after a write
127/// that changes the inbox. Skips the count query when nobody is listening. Never fails: problems
128/// are logged.
129pub async fn publish_unread(state: &AppState, user_id: &str) {
130 if !state.notify_hub.is_watched(user_id) {
131 return;
132 }
133 match state.store.unread_count(user_id).await {
134 Ok(count) => state.notify_hub.publish(user_id, count),
135 Err(e) => tracing::error!(error = %e, user_id, "notify: unread count failed"),
136 }
137}
138
139#[cfg(test)]
140mod tests {
141 use std::time::Duration;
142
143 use super::*;
144
145 #[tokio::test]
146 async fn publish_reaches_the_subscriber_of_that_user_only() {
147 let hub = Arc::new(NotifyHub::new());
148 let mut alice = hub.subscribe("alice");
149 let mut bob = hub.subscribe("bob");
150
151 hub.publish("alice", 3);
152 assert_eq!(alice.next_count().await, Some(3));
153
154 // Bob heard nothing: his wait times out.
155 let quiet = tokio::time::timeout(Duration::from_millis(50), bob.next_count()).await;
156 assert!(quiet.is_err());
157 }
158
159 #[tokio::test]
160 async fn publish_without_subscribers_is_a_no_op_and_nothing_is_kept() {
161 let hub = Arc::new(NotifyHub::new());
162 assert!(!hub.is_watched("alice"));
163 hub.publish("alice", 1);
164 assert!(!hub.is_watched("alice"));
165 assert!(hub.users().is_empty());
166 }
167
168 #[tokio::test]
169 async fn the_channel_is_released_when_the_last_subscription_drops() {
170 let hub = Arc::new(NotifyHub::new());
171 let first = hub.subscribe("alice");
172 let second = hub.subscribe("alice");
173 assert!(hub.is_watched("alice"));
174
175 drop(first);
176 assert!(hub.is_watched("alice"));
177
178 drop(second);
179 assert!(!hub.is_watched("alice"));
180 assert!(hub.users().is_empty());
181 }
182
183 #[tokio::test]
184 async fn every_open_stream_of_a_user_gets_the_count() {
185 let hub = Arc::new(NotifyHub::new());
186 let mut tab_one = hub.subscribe("alice");
187 let mut tab_two = hub.subscribe("alice");
188
189 hub.publish("alice", 5);
190 assert_eq!(tab_one.next_count().await, Some(5));
191 assert_eq!(tab_two.next_count().await, Some(5));
192 }
193
194 #[tokio::test]
195 async fn close_ends_every_subscription() {
196 let hub = Arc::new(NotifyHub::new());
197 let mut sub = hub.subscribe("alice");
198
199 hub.close();
200 assert_eq!(sub.next_count().await, None);
201 }
202}