atagen

macltok

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