atagen

macltok

public
Star0
Watch0
dc609022bb53f98a7c12d4f5ae0df872fc521de8
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/watchers.rsRust531 lines18.3 KB
1
use std::{collections::HashSet, sync::Arc, time::Duration};
2
3
use futures_util::StreamExt;
4
use matrix_sdk::{
5
    Client, Room,
6
    ruma::events::{AnyMessageLikeEventContent, room::message::MessageType},
7
    send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendQueueUpdate},
8
};
9
use matrix_sdk_ui::sync_service::{State as SyncState, SyncService};
10
use ruma::{OwnedRoomId, OwnedTransactionId, time::Instant};
11
use tokio::sync::{broadcast::error::RecvError, mpsc, watch};
12
13
use super::executor::TaskSet;
14
use crate::{
15
    ConnectionState, Event, UploadOutcome, encryption::encryption_state, error::classify_sdk_error,
16
    transport::EventSender,
17
};
18
19
const SYNC_RESTART_MIN: Duration = Duration::from_secs(1);
20
const SYNC_RESTART_MAX: Duration = Duration::from_secs(60);
21
const SEND_QUEUE_RETRY: Duration = Duration::from_secs(5);
22
const SYNC_STABLE: Duration = Duration::from_secs(30);
23
24
pub(super) async fn room_subscriptions(
25
    sync: Arc<SyncService>,
26
    mut rooms: watch::Receiver<Vec<ruma::OwnedRoomId>>,
27
) {
28
    while rooms.changed().await.is_ok() {
29
        let room_ids = rooms.borrow_and_update().clone();
30
        let room_ids = room_ids.iter().map(AsRef::as_ref).collect::<Vec<_>>();
31
        sync.room_list_service().subscribe_to_rooms(&room_ids).await;
32
    }
33
}
34
35
/// `states` must be taken before sync starts, a late subscriber never hears the transition it missed
36
pub(super) async fn connection(
37
    mut states: impl futures_util::Stream<Item = SyncState> + Unpin,
38
    sync: Arc<SyncService>,
39
    client: Client,
40
    events: EventSender,
41
    connection: watch::Sender<ConnectionState>,
42
) {
43
    let mut backoff = SYNC_RESTART_MIN;
44
    let mut online_since: Option<Instant> = None;
45
    while let Some(state) = states.next().await {
46
        // only a run that lasted counts as recovered, or a server that answers /versions
47
        // but refuses sync would have us bounce between offline and online forever
48
        let stable = online_since
49
            .take()
50
            .map(|since| since.elapsed() >= SYNC_STABLE);
51
        if stable == Some(true) {
52
            backoff = SYNC_RESTART_MIN;
53
        }
54
        let (state, restart) = match state {
55
            // only ever us stopping it for a restart, nothing a client needs to hear
56
            SyncState::Idle => continue,
57
            SyncState::Running => {
58
                online_since = Some(Instant::now());
59
                // the sdk parks a room queue on any send error and never unparks it
60
                client.send_queue().set_enabled(true).await;
61
                (ConnectionState::Online, false)
62
            }
63
            SyncState::Offline => (ConnectionState::Offline, stable == Some(false)),
64
            SyncState::Terminated => (ConnectionState::Offline, true),
65
            SyncState::Error(error) => {
66
                tracing::warn!(%error, "sync service reported an error");
67
                (ConnectionState::Offline, true)
68
            }
69
        };
70
        if *connection.borrow() != state {
71
            connection.send_replace(state);
72
            if events.send(Event::ConnectionChanged(state)).await.is_err() {
73
                return;
74
            }
75
        }
76
        if restart {
77
            sync.stop().await;
78
            matrix_sdk::sleep::sleep(backoff).await;
79
            backoff = (backoff * 2).min(SYNC_RESTART_MAX);
80
            sync.start().await;
81
        }
82
    }
83
}
84
85
/// a recoverable failure while we're online is a hiccup, so give the room another go
86
pub(super) async fn send_queue_errors(
87
    client: Client,
88
    mut errors: tokio::sync::broadcast::Receiver<matrix_sdk::send_queue::SendQueueRoomError>,
89
    online: watch::Receiver<ConnectionState>,
90
) {
91
    let mut retries = TaskSet::new();
92
    loop {
93
        let error = tokio::select! {
94
            error = errors.recv() => error,
95
            Some(_) = retries.join_next(), if !retries.is_empty() => continue,
96
        };
97
        match error {
98
            Ok(error) if error.is_recoverable => {
99
                let client = client.clone();
100
                let online = online.clone();
101
                retries.spawn(async move {
102
                    matrix_sdk::sleep::sleep(SEND_QUEUE_RETRY).await;
103
                    if *online.borrow() == ConnectionState::Online
104
                        && let Some(room) = client.get_room(&error.room_id)
105
                    {
106
                        room.send_queue().set_enabled(true);
107
                    }
108
                });
109
            }
110
            Ok(_) | Err(RecvError::Lagged(_)) => {}
111
            Err(RecvError::Closed) => return,
112
        }
113
    }
114
}
115
116
/// upload progress only goes out on the room's own channel, never the client-wide one.
117
/// the room's sends and cancels come along behind it, so an upload's end can't
118
/// overtake its own progress
119
pub(super) async fn room_uploads(
120
    room: Room,
121
    tracker: mpsc::Sender<RoomUpload>,
122
    ready: watch::Sender<bool>,
123
) {
124
    let room_id = room.room_id().to_owned();
125
    let mut updates = match room.send_queue().subscribe().await {
126
        Ok((_, updates)) => updates,
127
        Err(error) => {
128
            tracing::warn!(%room_id, %error, "reading the media send queue failed");
129
            return;
130
        }
131
    };
132
    ready.send_replace(true);
133
    loop {
134
        let report = match updates.recv().await {
135
            Ok(RoomSendQueueUpdate::MediaUpload {
136
                related_to,
137
                file,
138
                progress,
139
                ..
140
            }) => RoomUpload::Progress {
141
                upload: (room_id.clone(), related_to),
142
                uploaded_bytes: progress.current as u64,
143
                total_bytes: progress.total as u64,
144
                complete: file.is_some() && progress.current == progress.total,
145
            },
146
            Ok(RoomSendQueueUpdate::SentEvent {
147
                transaction_id,
148
                event_id,
149
            }) => RoomUpload::Ended(
150
                (room_id.clone(), transaction_id),
151
                UploadOutcome::Sent { event_id },
152
            ),
153
            Ok(RoomSendQueueUpdate::CancelledLocalEvent { transaction_id }) => {
154
                RoomUpload::Ended((room_id.clone(), transaction_id), UploadOutcome::Cancelled)
155
            }
156
            Ok(_) => continue,
157
            Err(RecvError::Lagged(_)) => RoomUpload::Lagged,
158
            Err(RecvError::Closed) => return,
159
        };
160
        if tracker.send(report).await.is_err() {
161
            return;
162
        }
163
    }
164
}
165
166
pub(super) enum RoomUpload {
167
    Progress {
168
        upload: Upload,
169
        uploaded_bytes: u64,
170
        total_bytes: u64,
171
        complete: bool,
172
    },
173
    Ended(Upload, UploadOutcome),
174
    /// an end may have been lost with the ticks
175
    Lagged,
176
}
177
178
/// the one place upload events leave from, so an upload is never reported after it
179
/// ended. progress comes in from the room forwarders, the rest from the client-wide
180
/// channel, which exists before any upload does
181
pub(super) async fn uploads(
182
    client: Client,
183
    mut updates: tokio::sync::broadcast::Receiver<SendQueueUpdate>,
184
    mut rooms_reports: mpsc::Receiver<RoomUpload>,
185
    events: EventSender,
186
) {
187
    let (_, mut rooms) = client.rooms_stream();
188
    let mut tracker = UploadTracker {
189
        live: attachment_echoes(&client).await,
190
        ended: HashSet::new(),
191
    };
192
    loop {
193
        let reports = tokio::select! {
194
            update = updates.recv() => match update {
195
                Ok(update) => tracker.update(update).into_iter().collect(),
196
                Err(RecvError::Lagged(skipped)) => {
197
                    tracing::warn!(skipped, "upload tracker lagged");
198
                    tracker.resnapshot(attachment_echoes(&client).await)
199
                }
200
                Err(RecvError::Closed) => return,
201
            },
202
            Some(report) = rooms_reports.recv() => match report {
203
                RoomUpload::Lagged => tracker.resnapshot(attachment_echoes(&client).await),
204
                report => tracker.room(report).into_iter().collect(),
205
            },
206
            update = rooms.next() => {
207
                if update.is_none() {
208
                    return;
209
                }
210
                tracker.finish_where(
211
                    |(room_id, _)| client.get_room(room_id).is_none(),
212
                    &UploadOutcome::RoomForgotten,
213
                )
214
            }
215
        };
216
        for event in reports {
217
            if events.send(event).await.is_err() {
218
                return;
219
            }
220
        }
221
    }
222
}
223
224
type Upload = (OwnedRoomId, OwnedTransactionId);
225
226
/// everything reported gets exactly one `MediaUploadFinished`. an upload ended from
227
/// anywhere but its room's forwarder goes in `ended` until the forwarder's own end for
228
/// it arrives, since until then the forwarder can still be holding progress for it
229
struct UploadTracker {
230
    live: HashSet<Upload>,
231
    ended: HashSet<Upload>,
232
}
233
234
impl UploadTracker {
235
    fn update(&mut self, update: SendQueueUpdate) -> Option<Event> {
236
        let room_id = update.room_id;
237
        let (transaction_id, outcome) = match update.update {
238
            RoomSendQueueUpdate::NewLocalEvent(echo) => {
239
                let upload = (room_id, echo.transaction_id);
240
                if is_attachment_echo(&echo.content) && !self.ended.contains(&upload) {
241
                    self.live.insert(upload);
242
                }
243
                return None;
244
            }
245
            RoomSendQueueUpdate::SendError {
246
                transaction_id,
247
                error,
248
                is_recoverable,
249
            } => {
250
                let upload = (room_id, transaction_id);
251
                if !self.live.contains(&upload) {
252
                    return None;
253
                }
254
                let (room_id, transaction_id) = upload;
255
                return Some(Event::MediaUploadFailed {
256
                    room_id,
257
                    transaction_id,
258
                    kind: classify_sdk_error(&error),
259
                    will_retry: is_recoverable,
260
                });
261
            }
262
            RoomSendQueueUpdate::SentEvent {
263
                transaction_id,
264
                event_id,
265
            } => (transaction_id, UploadOutcome::Sent { event_id }),
266
            RoomSendQueueUpdate::CancelledLocalEvent { transaction_id } => {
267
                (transaction_id, UploadOutcome::Cancelled)
268
            }
269
            _ => return None,
270
        };
271
        self.finish((room_id, transaction_id), outcome, true)
272
    }
273
274
    fn room(&mut self, report: RoomUpload) -> Option<Event> {
275
        match report {
276
            RoomUpload::Progress {
277
                upload,
278
                uploaded_bytes,
279
                total_bytes,
280
                complete,
281
            } => {
282
                if self.ended.contains(&upload) {
283
                    return None;
284
                }
285
                self.live.insert(upload.clone());
286
                let (room_id, transaction_id) = upload;
287
                Some(Event::MediaUploadProgress {
288
                    room_id,
289
                    transaction_id,
290
                    uploaded_bytes,
291
                    total_bytes,
292
                    complete,
293
                })
294
            }
295
            RoomUpload::Ended(upload, outcome) => {
296
                self.ended.remove(&upload);
297
                self.finish(upload, outcome, false)
298
            }
299
            RoomUpload::Lagged => None,
300
        }
301
    }
302
303
    fn resnapshot(&mut self, current: HashSet<Upload>) -> Vec<Event> {
304
        let reports = self.finish_where(|upload| !current.contains(upload), &UploadOutcome::Missed);
305
        let ended = &self.ended;
306
        self.live
307
            .extend(current.into_iter().filter(|upload| !ended.contains(upload)));
308
        reports
309
    }
310
311
    fn finish_where(
312
        &mut self,
313
        ends: impl Fn(&Upload) -> bool,
314
        outcome: &UploadOutcome,
315
    ) -> Vec<Event> {
316
        let ending = self
317
            .live
318
            .iter()
319
            .filter(|upload| ends(upload))
320
            .cloned()
321
            .collect::<Vec<_>>();
322
        ending
323
            .into_iter()
324
            .filter_map(|upload| self.finish(upload, outcome.clone(), true))
325
            .collect()
326
    }
327
328
    fn finish(&mut self, upload: Upload, outcome: UploadOutcome, remember: bool) -> Option<Event> {
329
        if !self.live.remove(&upload) {
330
            return None;
331
        }
332
        if remember {
333
            self.ended.insert(upload.clone());
334
        }
335
        let (room_id, transaction_id) = upload;
336
        Some(Event::MediaUploadFinished {
337
            room_id,
338
            transaction_id,
339
            outcome,
340
        })
341
    }
342
}
343
344
async fn attachment_echoes(client: &Client) -> HashSet<Upload> {
345
    let Ok(echoes) = client.send_queue().local_echoes().await else {
346
        return HashSet::new();
347
    };
348
    echoes
349
        .into_iter()
350
        .flat_map(|(room_id, echoes)| {
351
            echoes
352
                .into_iter()
353
                .filter(|echo| is_attachment_echo(&echo.content))
354
                .map(move |echo| (room_id.clone(), echo.transaction_id))
355
        })
356
        .collect()
357
}
358
359
fn is_attachment_echo(content: &LocalEchoContent) -> bool {
360
    let LocalEchoContent::Event {
361
        serialized_event, ..
362
    } = content
363
    else {
364
        return false;
365
    };
366
    matches!(
367
        serialized_event.deserialize(),
368
        Ok(AnyMessageLikeEventContent::RoomMessage(content))
369
            if matches!(
370
                content.msgtype,
371
                MessageType::Audio(_) | MessageType::File(_) | MessageType::Image(_) | MessageType::Video(_)
372
            )
373
    )
374
}
375
376
pub(super) async fn encryption(client: Client, events: EventSender) {
377
    let Some(device_id) = client.device_id().map(ToOwned::to_owned) else {
378
        tracing::warn!("encryption watcher stopped because the client has no device id");
379
        return;
380
    };
381
    let recovery = client.encryption().recovery();
382
    let backups = client.encryption().backups();
383
    let mut recovery_stream = recovery.state_stream();
384
    let mut backup_stream = backups.state_stream();
385
    let mut verification_stream = client.encryption().verification_state();
386
    let mut recovery_state = recovery.state();
387
    let mut backup_state = backups.state();
388
    let mut verification_state = verification_stream.get();
389
    loop {
390
        let state = encryption_state(
391
            device_id.clone(),
392
            recovery_state,
393
            verification_state,
394
            backup_state,
395
        );
396
        if events.send(Event::EncryptionChanged(state)).await.is_err() {
397
            return;
398
        }
399
        tokio::select! {
400
            Some(update) = recovery_stream.next() => recovery_state = update,
401
            Some(update) = verification_stream.next() => verification_state = update,
402
            Some(update) = backup_stream.next() => {
403
                backup_state = update.unwrap_or_else(|_| backups.state());
404
            }
405
            else => return,
406
        }
407
    }
408
}
409
410
pub(super) async fn rooms(client: Client, events: EventSender) {
411
    let (_, mut rooms) = client.rooms_stream();
412
    let mut info = client.room_info_notable_update_receiver();
413
    loop {
414
        if events.send(Event::RoomsChanged).await.is_err() {
415
            return;
416
        }
417
        tokio::select! {
418
            update = rooms.next() => if update.is_none() { return },
419
            update = info.recv() => if matches!(update, Err(RecvError::Closed)) { return },
420
        }
421
        // a sync touching fifty rooms is one invalidation, not fifty
422
        while let Ok(_) | Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) =
423
            info.try_recv()
424
        {}
425
    }
426
}
427
428
pub(super) async fn notification_settings(client: Client, events: EventSender) {
429
    let settings = client.notification_settings().await;
430
    let mut changes = settings.subscribe_to_changes();
431
    while let Ok(()) | Err(RecvError::Lagged(_)) = changes.recv().await {
432
        // room summaries carry the effective mode, and push rule changes never touch room info
433
        for event in [Event::NotificationSettingsChanged, Event::RoomsChanged] {
434
            if events.send(event).await.is_err() {
435
                return;
436
            }
437
        }
438
    }
439
}
440
441
#[cfg(test)]
442
mod tests {
443
    use matrix_sdk::send_queue::{RoomSendQueueUpdate, SendQueueUpdate};
444
    use ruma::{OwnedTransactionId, event_id, owned_room_id};
445
446
    use super::{HashSet, RoomUpload, Upload, UploadTracker};
447
    use crate::{Event, UploadOutcome};
448
449
    fn upload() -> Upload {
450
        (
451
            owned_room_id!("!room:localhost"),
452
            OwnedTransactionId::from("upload"),
453
        )
454
    }
455
456
    fn tick() -> RoomUpload {
457
        RoomUpload::Progress {
458
            upload: upload(),
459
            uploaded_bytes: 1,
460
            total_bytes: 2,
461
            complete: false,
462
        }
463
    }
464
465
    fn sent() -> UploadOutcome {
466
        UploadOutcome::Sent {
467
            event_id: event_id!("$sent:localhost").to_owned(),
468
        }
469
    }
470
471
    fn tracker() -> UploadTracker {
472
        UploadTracker {
473
            live: HashSet::new(),
474
            ended: HashSet::new(),
475
        }
476
    }
477
478
    fn finished(event: Option<Event>) -> Option<UploadOutcome> {
479
        match event? {
480
            Event::MediaUploadFinished { outcome, .. } => Some(outcome),
481
            _ => None,
482
        }
483
    }
484
485
    #[test]
486
    fn a_stale_tick_after_a_missed_end_changes_nothing() {
487
        let mut tracker = tracker();
488
        assert!(tracker.room(tick()).is_some());
489
        let missed = tracker.resnapshot(HashSet::new());
490
        assert!(matches!(
491
            missed.as_slice(),
492
            [Event::MediaUploadFinished {
493
                outcome: UploadOutcome::Missed,
494
                ..
495
            }]
496
        ));
497
498
        assert!(tracker.room(tick()).is_none());
499
        let (room_id, transaction_id) = upload();
500
        let global = SendQueueUpdate {
501
            room_id,
502
            update: RoomSendQueueUpdate::SentEvent {
503
                transaction_id,
504
                event_id: event_id!("$sent:localhost").to_owned(),
505
            },
506
        };
507
        assert!(tracker.update(global).is_none());
508
        assert!(tracker.room(RoomUpload::Ended(upload(), sent())).is_none());
509
        assert!(tracker.live.is_empty() && tracker.ended.is_empty());
510
    }
511
512
    #[test]
513
    fn an_end_the_tracker_never_knew_about_still_follows_its_progress() {
514
        let mut tracker = tracker();
515
        let (room_id, transaction_id) = upload();
516
        let global = SendQueueUpdate {
517
            room_id,
518
            update: RoomSendQueueUpdate::SentEvent {
519
                transaction_id,
520
                event_id: event_id!("$sent:localhost").to_owned(),
521
            },
522
        };
523
        assert!(tracker.update(global).is_none());
524
        assert!(tracker.room(tick()).is_some());
525
        assert_eq!(
526
            finished(tracker.room(RoomUpload::Ended(upload(), sent()))),
527
            Some(sent())
528
        );
529
        assert!(tracker.live.is_empty() && tracker.ended.is_empty());
530
    }
531
}