atagen

macltok

public
Star0
Watch0
67bf96a80270e409b3d380fe8ea02c7cd9258338
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/engine/watchers.rsRust489 lines17.1 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
pub(super) async fn room_uploads(
150
    room: Room,
151
    progress: mpsc::Sender<UploadProgress>,
152
    ready: watch::Sender<bool>,
153
) {
154
    let room_id = room.room_id().to_owned();
155
    let mut updates = match room.send_queue().subscribe().await {
156
        Ok((_, updates)) => updates,
157
        Err(error) => {
158
            tracing::warn!(%room_id, %error, "reading the media send queue failed");
159
            return;
160
        }
161
    };
162
    ready.send_replace(true);
163
    loop {
164
        let update = match updates.recv().await {
165
            Ok(update) => update,
166
            // a missed tick is fine, the next one says where it's got to
167
            Err(RecvError::Lagged(_)) => continue,
168
            Err(RecvError::Closed) => return,
169
        };
170
        let RoomSendQueueUpdate::MediaUpload {
171
            related_to,
172
            file,
173
            progress: done,
174
            ..
175
        } = update
176
        else {
177
            continue;
178
        };
179
        let tick = UploadProgress {
180
            upload: (room_id.clone(), related_to),
181
            uploaded_bytes: done.current as u64,
182
            total_bytes: done.total as u64,
183
            complete: file.is_some() && done.current == done.total,
184
        };
185
        if progress.send(tick).await.is_err() {
186
            return;
187
        }
188
    }
189
}
190
191
pub(super) struct UploadProgress {
192
    upload: Upload,
193
    uploaded_bytes: u64,
194
    total_bytes: u64,
195
    complete: bool,
196
}
197
198
/// the one place upload events leave from, so an upload is never reported after it
199
/// ended. progress comes in from the room forwarders, the rest from the client-wide
200
/// channel, which exists before any upload does
201
pub(super) async fn uploads(
202
    client: Client,
203
    mut updates: tokio::sync::broadcast::Receiver<SendQueueUpdate>,
204
    mut progress: mpsc::Receiver<UploadProgress>,
205
    events: EventSender,
206
) {
207
    let (_, mut rooms) = client.rooms_stream();
208
    let mut tracker = UploadTracker {
209
        live: attachment_echoes(&client).await,
210
        ended: HashSet::new(),
211
    };
212
    loop {
213
        let reports = tokio::select! {
214
            update = updates.recv() => match update {
215
                Ok(update) => tracker.update(update).into_iter().collect(),
216
                Err(RecvError::Lagged(skipped)) => {
217
                    tracing::warn!(skipped, "upload watcher lagged");
218
                    let current = attachment_echoes(&client).await;
219
                    tracker.resnapshot(current)
220
                }
221
                Err(RecvError::Closed) => return,
222
            },
223
            Some(tick) = progress.recv() => tracker.progress(tick).into_iter().collect(),
224
            update = rooms.next() => {
225
                if update.is_none() {
226
                    return;
227
                }
228
                tracker.finish_where(
229
                    |(room_id, _)| client.get_room(room_id).is_none(),
230
                    &UploadOutcome::RoomForgotten,
231
                )
232
            }
233
        };
234
        for event in reports {
235
            if events.send(event).await.is_err() {
236
                return;
237
            }
238
        }
239
    }
240
}
241
242
type Upload = (OwnedRoomId, OwnedTransactionId);
243
244
/// everything reported gets exactly one `MediaUploadFinished`. `ended` is kept for good,
245
/// since a room forwarder can still be holding progress for an upload the client-wide
246
/// channel has already seen sent
247
struct UploadTracker {
248
    live: HashSet<Upload>,
249
    ended: HashSet<Upload>,
250
}
251
252
impl UploadTracker {
253
    fn update(&mut self, update: SendQueueUpdate) -> Option<Event> {
254
        let room_id = update.room_id;
255
        let (transaction_id, outcome) = match update.update {
256
            RoomSendQueueUpdate::NewLocalEvent(echo) => {
257
                let upload = (room_id, echo.transaction_id);
258
                if is_attachment_echo(&echo.content) && !self.ended.contains(&upload) {
259
                    self.live.insert(upload);
260
                }
261
                return None;
262
            }
263
            RoomSendQueueUpdate::SendError {
264
                transaction_id,
265
                error,
266
                is_recoverable,
267
            } => {
268
                let upload = (room_id, transaction_id);
269
                if !self.live.contains(&upload) {
270
                    return None;
271
                }
272
                let (room_id, transaction_id) = upload;
273
                return Some(Event::MediaUploadFailed {
274
                    room_id,
275
                    transaction_id,
276
                    kind: classify_sdk_error(&error),
277
                    will_retry: is_recoverable,
278
                });
279
            }
280
            RoomSendQueueUpdate::SentEvent {
281
                transaction_id,
282
                event_id,
283
            } => (transaction_id, UploadOutcome::Sent { event_id }),
284
            RoomSendQueueUpdate::CancelledLocalEvent { transaction_id } => {
285
                (transaction_id, UploadOutcome::Cancelled)
286
            }
287
            _ => return None,
288
        };
289
        self.finish((room_id, transaction_id), outcome)
290
    }
291
292
    fn progress(&mut self, tick: UploadProgress) -> Option<Event> {
293
        if self.ended.contains(&tick.upload) {
294
            return None;
295
        }
296
        self.live.insert(tick.upload.clone());
297
        let (room_id, transaction_id) = tick.upload;
298
        Some(Event::MediaUploadProgress {
299
            room_id,
300
            transaction_id,
301
            uploaded_bytes: tick.uploaded_bytes,
302
            total_bytes: tick.total_bytes,
303
            complete: tick.complete,
304
        })
305
    }
306
307
    fn resnapshot(&mut self, current: HashSet<Upload>) -> Vec<Event> {
308
        let reports = self.finish_where(|upload| !current.contains(upload), &UploadOutcome::Missed);
309
        let ended = &self.ended;
310
        self.live
311
            .extend(current.into_iter().filter(|upload| !ended.contains(upload)));
312
        reports
313
    }
314
315
    fn finish_where(
316
        &mut self,
317
        ends: impl Fn(&Upload) -> bool,
318
        outcome: &UploadOutcome,
319
    ) -> Vec<Event> {
320
        let ending = self
321
            .live
322
            .iter()
323
            .filter(|upload| ends(upload))
324
            .cloned()
325
            .collect::<Vec<_>>();
326
        ending
327
            .into_iter()
328
            .filter_map(|upload| self.finish(upload, outcome.clone()))
329
            .collect()
330
    }
331
332
    fn finish(&mut self, upload: Upload, outcome: UploadOutcome) -> Option<Event> {
333
        if !self.live.remove(&upload) {
334
            return None;
335
        }
336
        self.ended.insert(upload.clone());
337
        let (room_id, transaction_id) = upload;
338
        Some(Event::MediaUploadFinished {
339
            room_id,
340
            transaction_id,
341
            outcome,
342
        })
343
    }
344
}
345
346
async fn attachment_echoes(client: &Client) -> HashSet<Upload> {
347
    let Ok(echoes) = client.send_queue().local_echoes().await else {
348
        return HashSet::new();
349
    };
350
    echoes
351
        .into_iter()
352
        .flat_map(|(room_id, echoes)| {
353
            echoes
354
                .into_iter()
355
                .filter(|echo| is_attachment_echo(&echo.content))
356
                .map(move |echo| (room_id.clone(), echo.transaction_id))
357
        })
358
        .collect()
359
}
360
361
fn is_attachment_echo(content: &LocalEchoContent) -> bool {
362
    let LocalEchoContent::Event {
363
        serialized_event, ..
364
    } = content
365
    else {
366
        return false;
367
    };
368
    matches!(
369
        serialized_event.deserialize(),
370
        Ok(AnyMessageLikeEventContent::RoomMessage(content))
371
            if matches!(
372
                content.msgtype,
373
                MessageType::Audio(_) | MessageType::File(_) | MessageType::Image(_) | MessageType::Video(_)
374
            )
375
    )
376
}
377
378
pub(super) async fn encryption(client: Client, events: EventSender) {
379
    let Some(device_id) = client.device_id().map(ToOwned::to_owned) else {
380
        tracing::warn!("encryption watcher stopped because the client has no device id");
381
        return;
382
    };
383
    let recovery = client.encryption().recovery();
384
    let backups = client.encryption().backups();
385
    let mut recovery_stream = recovery.state_stream();
386
    let mut backup_stream = backups.state_stream();
387
    let mut verification_stream = client.encryption().verification_state();
388
    let mut recovery_state = recovery.state();
389
    let mut backup_state = backups.state();
390
    let mut verification_state = verification_stream.get();
391
    loop {
392
        let state = encryption_state(
393
            device_id.clone(),
394
            recovery_state,
395
            verification_state,
396
            backup_state,
397
        );
398
        if events.send(Event::EncryptionChanged(state)).await.is_err() {
399
            return;
400
        }
401
        tokio::select! {
402
            Some(update) = recovery_stream.next() => recovery_state = update,
403
            Some(update) = verification_stream.next() => verification_state = update,
404
            Some(update) = backup_stream.next() => {
405
                backup_state = update.unwrap_or_else(|_| backups.state());
406
            }
407
            else => return,
408
        }
409
    }
410
}
411
412
pub(super) async fn rooms(client: Client, events: EventSender) {
413
    let (_, mut rooms) = client.rooms_stream();
414
    let mut info = client.room_info_notable_update_receiver();
415
    loop {
416
        if events.send(Event::RoomsChanged).await.is_err() {
417
            return;
418
        }
419
        tokio::select! {
420
            update = rooms.next() => if update.is_none() { return },
421
            update = info.recv() => if matches!(update, Err(RecvError::Closed)) { return },
422
        }
423
        // a sync touching fifty rooms is one invalidation, not fifty
424
        while let Ok(_) | Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) =
425
            info.try_recv()
426
        {}
427
    }
428
}
429
430
pub(super) async fn notification_settings(client: Client, events: EventSender) {
431
    let settings = client.notification_settings().await;
432
    let mut changes = settings.subscribe_to_changes();
433
    while let Ok(()) | Err(RecvError::Lagged(_)) = changes.recv().await {
434
        // room summaries carry the effective mode, and push rule changes never touch room info
435
        for event in [Event::NotificationSettingsChanged, Event::RoomsChanged] {
436
            if events.send(event).await.is_err() {
437
                return;
438
            }
439
        }
440
    }
441
}
442
443
#[cfg(test)]
444
mod tests {
445
    use matrix_sdk::send_queue::{RoomSendQueueUpdate, SendQueueUpdate};
446
    use ruma::{OwnedTransactionId, event_id, owned_room_id};
447
448
    use super::{HashSet, UploadProgress, UploadTracker};
449
    use crate::{Event, UploadOutcome};
450
451
    #[test]
452
    fn an_upload_only_progress_saw_still_ends_once() {
453
        let upload = (
454
            owned_room_id!("!room:localhost"),
455
            OwnedTransactionId::from("upload"),
456
        );
457
        let tick = || UploadProgress {
458
            upload: upload.clone(),
459
            uploaded_bytes: 1,
460
            total_bytes: 2,
461
            complete: false,
462
        };
463
        let mut tracker = UploadTracker {
464
            live: HashSet::new(),
465
            ended: HashSet::new(),
466
        };
467
        assert!(tracker.progress(tick()).is_some());
468
469
        let missed = tracker.resnapshot(HashSet::new());
470
        assert!(matches!(
471
            missed.as_slice(),
472
            [Event::MediaUploadFinished {
473
                outcome: UploadOutcome::Missed,
474
                ..
475
            }]
476
        ));
477
478
        assert!(tracker.progress(tick()).is_none());
479
        let sent = SendQueueUpdate {
480
            room_id: upload.0.clone(),
481
            update: RoomSendQueueUpdate::SentEvent {
482
                transaction_id: upload.1.clone(),
483
                event_id: event_id!("$sent:localhost").to_owned(),
484
            },
485
        };
486
        assert!(tracker.update(sent).is_none());
487
        assert!(tracker.resnapshot(HashSet::from([upload])).is_empty());
488
    }
489
}