atagen

macltok

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