1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
|
use std::collections::{hash_map::Entry, HashMap};
use std::sync::{Arc, Mutex, MutexGuard};
use futures::{future, stream::StreamExt as _, Stream};
use sqlx::sqlite::SqlitePool;
use tokio::sync::broadcast::{channel, Sender};
use tokio_stream::wrappers::{errors::BroadcastStreamRecvError, BroadcastStream};
use super::repo::broadcast;
use crate::repo::channel::{self, Provider as _};
// Clones will share the same senders collection.
#[derive(Clone)]
pub struct Broadcaster {
// The use of std::sync::Mutex, and not tokio::sync::Mutex, follows Tokio's
// own advice: <https://tokio.rs/tokio/tutorial/shared-state>. Methods that
// lock it must be sync.
senders: Arc<Mutex<HashMap<channel::Id, Sender<broadcast::Message>>>>,
}
impl Broadcaster {
pub async fn from_database(db: &SqlitePool) -> Result<Self, sqlx::Error> {
let mut tx = db.begin().await?;
let channels = tx.channels().all().await?;
tx.commit().await?;
let channels = channels.iter().map(|c| &c.id);
let broadcaster = Self::new(channels);
Ok(broadcaster)
}
fn new<'i>(channels: impl IntoIterator<Item = &'i channel::Id>) -> Self {
let senders: HashMap<_, _> = channels
.into_iter()
.cloned()
.map(|id| (id, Self::make_sender()))
.collect();
Self {
senders: Arc::new(Mutex::new(senders)),
}
}
// panic: if ``channel`` is already registered.
pub fn register_channel(&self, channel: &channel::Id) {
match self.senders().entry(channel.clone()) {
// This ever happening indicates a serious logic error.
Entry::Occupied(_) => panic!("duplicate channel registration for channel {channel}"),
Entry::Vacant(entry) => {
entry.insert(Self::make_sender());
}
}
}
// panic: if ``channel`` has not been previously registered, and was not
// part of the initial set of channels.
pub fn broadcast(&self, channel: &channel::Id, message: &broadcast::Message) {
let tx = self.sender(channel);
// Per the Tokio docs, the returned error is only used to indicate that
// there are no receivers. In this use case, that's fine; a lack of
// listening consumers (chat clients) when a message is sent isn't an
// error.
//
// The successful return value, which includes the number of active
// receivers, also isn't that interesting to us.
let _ = tx.send(message.clone());
}
// panic: if ``channel`` has not been previously registered, and was not
// part of the initial set of channels.
pub fn listen(&self, channel: &channel::Id) -> impl Stream<Item = broadcast::Message> {
let rx = self.sender(channel).subscribe();
BroadcastStream::from(rx)
.take_while(|r| {
future::ready(match r {
Ok(_) => true,
// Stop the stream here. This will disconnect SSE clients
// (see `routes.rs`), who will then resume from
// `Last-Event-ID`, allowing them to catch up by reading
// the skipped messages from the database.
Err(BroadcastStreamRecvError::Lagged(_)) => false,
})
})
.map(|r| {
// Since the previous transform stops at the first error, this
// should always hold.
//
// See also <https://users.rust-lang.org/t/taking-from-stream-while-ok/48854>.
r.expect("after filtering, only `Ok` messages should remain")
})
}
// panic: if ``channel`` has not been previously registered, and was not
// part of the initial set of channels.
fn sender(&self, channel: &channel::Id) -> Sender<broadcast::Message> {
self.senders()[channel].clone()
}
fn senders(&self) -> MutexGuard<HashMap<channel::Id, Sender<broadcast::Message>>> {
self.senders.lock().unwrap() // propagate panics when mutex is poisoned
}
fn make_sender() -> Sender<broadcast::Message> {
// Queue depth of 16 chosen entirely arbitrarily. Don't read too much
// into it.
let (tx, _) = channel(16);
tx
}
}
|
