ptr.pet
at protocol indexer with flexible filtering, xrpc queries, and a cursor-backed event stream, built on fjall
at-protocol
atproto
indexer
rust
fjall
[api] cancel stream handler thread properly
+25
-3
src/api/stream.rs
+25
-3
src/api/stream.rs
···
31
31
32
32
async fn handle_socket(mut socket: WebSocket, state: Arc<AppState>, query: StreamQuery) {
33
33
let (tx, mut rx) = mpsc::channel(500);
34
+
let (cancel_tx, cancel_rx) = tokio::sync::oneshot::channel::<()>();
34
35
35
-
std::thread::Builder::new()
36
+
let runtime = tokio::runtime::Handle::current();
37
+
38
+
let thread = std::thread::Builder::new()
36
39
.name(format!(
37
40
"stream-handler-{}",
38
41
std::time::SystemTime::UNIX_EPOCH
···
41
44
.as_secs()
42
45
))
43
46
.spawn(move || {
47
+
let mut cancel_rx = cancel_rx;
44
48
let db = &state.db;
45
49
let mut event_rx = db.event_tx.subscribe();
46
50
let ks = db.events.clone();
···
51
55
max_id.saturating_sub(1)
52
56
}
53
57
};
58
+
let _runtime_guard = runtime.enter();
54
59
55
60
loop {
56
61
// 1. catch up from DB
···
146
151
}
147
152
148
153
// 2. wait for live events
149
-
match event_rx.blocking_recv() {
154
+
let next_event = runtime.block_on(async {
155
+
tokio::select! {
156
+
res = event_rx.recv() => Some(res),
157
+
_ = &mut cancel_rx => None,
158
+
}
159
+
});
160
+
161
+
let Some(next_event) = next_event else {
162
+
break;
163
+
};
164
+
165
+
match next_event {
150
166
Ok(BroadcastEvent::Persisted(_)) => {
151
167
// just wake up and run catch-up loop again
152
168
}
···
176
192
.expect("failed to spawn stream handler thread");
177
193
178
194
while let Some(msg) = rx.recv().await {
179
-
if socket.send(msg).await.is_err() {
195
+
if let Err(e) = socket.send(msg).await {
196
+
error!("failed to send ws message: {e}");
180
197
break;
181
198
}
199
+
}
200
+
201
+
let _ = cancel_tx.send(());
202
+
if let Err(e) = thread.join() {
203
+
error!("stream handler thread panicked: {e:?}");
182
204
}
183
205
}