atagen

macltok

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