atagen

macltok

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