atagen

macltok

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