atagen

macltok

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