macltok

public
Star0
Watch0
136 commits 3 branches 0 tags 0 releases3.0 MiB
use std::{
    collections::{HashMap, HashSet, VecDeque},
    sync::Arc,
    time::Duration,
};

use futures_util::StreamExt;
use matrix_sdk::{
    Client, Room,
    event_handler::EventHandlerHandle,
    ruma::events::{AnyMessageLikeEventContent, room::message::MessageType},
    send_queue::{LocalEchoContent, RoomSendQueueUpdate, SendQueueUpdate},
    serde_helpers::extract_thread_root,
};
use matrix_sdk_base::RoomInfoNotableUpdateReasons;
use matrix_sdk_ui::sync_service::{State as SyncState, SyncService};
use ruma::{
    OwnedEventId, OwnedRoomId, OwnedTransactionId,
    events::{
        AnyGlobalAccountDataEvent, AnyRoomAccountDataEvent, AnySyncStateEvent,
        AnySyncTimelineEvent, receipt::SyncReceiptEvent,
    },
    serde::Raw,
    time::Instant,
};
use tokio::sync::{
    Notify,
    broadcast::error::{RecvError, TryRecvError},
    mpsc, watch,
};

use super::executor::TaskSet;
use crate::{
    ConnectionState, Event, UploadOutcome,
    encryption::encryption_state,
    error::classify_sdk_error,
    normalise::{ENABLED_ROOMS, ROOM_PACKS, USER_PACK, recount_thread_unreads},
    space_order,
    transport::EventSender,
};

const SYNC_RESTART_MIN: Duration = Duration::from_secs(1);
const SYNC_RESTART_MAX: Duration = Duration::from_secs(60);
const SEND_QUEUE_RETRY: Duration = Duration::from_secs(5);
const SYNC_STABLE: Duration = Duration::from_secs(30);

pub(super) async fn room_subscriptions(
    sync: Arc<SyncService>,
    mut rooms: watch::Receiver<Vec<ruma::OwnedRoomId>>,
) {
    while rooms.changed().await.is_ok() {
        let room_ids = rooms.borrow_and_update().clone();
        let room_ids = room_ids.iter().map(AsRef::as_ref).collect::<Vec<_>>();
        sync.room_list_service().subscribe_to_rooms(&room_ids).await;
    }
}

/// `states` must be taken before sync starts, a late subscriber never hears the transition it missed
pub(super) async fn connection(
    mut states: impl futures_util::Stream<Item = SyncState> + Unpin,
    sync: Arc<SyncService>,
    client: Client,
    events: EventSender,
    connection: watch::Sender<ConnectionState>,
) {
    let mut backoff = SYNC_RESTART_MIN;
    let mut online_since: Option<Instant> = None;
    while let Some(state) = states.next().await {
        // only a run that lasted counts as recovered, or a server that answers /versions
        // but refuses sync would have us bounce between offline and online forever
        let stable = online_since
            .take()
            .map(|since| since.elapsed() >= SYNC_STABLE);
        if stable == Some(true) {
            backoff = SYNC_RESTART_MIN;
        }
        let (state, restart) = match state {
            // only ever us stopping it for a restart, nothing a client needs to hear
            SyncState::Idle => continue,
            SyncState::Running => {
                online_since = Some(Instant::now());
                // the sdk parks a room queue on any send error and never unparks it
                client.send_queue().set_enabled(true).await;
                (ConnectionState::Online, false)
            }
            SyncState::Offline => (ConnectionState::Offline, stable == Some(false)),
            SyncState::Terminated => (ConnectionState::Offline, true),
            SyncState::Error(error) => {
                tracing::warn!(%error, "sync service reported an error");
                (ConnectionState::Offline, true)
            }
        };
        if *connection.borrow() != state {
            connection.send_replace(state);
            if events.send(Event::ConnectionChanged(state)).await.is_err() {
                return;
            }
        }
        if restart {
            sync.stop().await;
            matrix_sdk::sleep::sleep(backoff).await;
            backoff = (backoff * 2).min(SYNC_RESTART_MAX);
            sync.start().await;
        }
    }
}

/// a recoverable failure while we're online is a hiccup, so give the room another go
pub(super) async fn send_queue_errors(
    client: Client,
    mut errors: tokio::sync::broadcast::Receiver<matrix_sdk::send_queue::SendQueueRoomError>,
    online: watch::Receiver<ConnectionState>,
) {
    let mut retries = TaskSet::new();
    loop {
        let error = tokio::select! {
            error = errors.recv() => error,
            Some(_) = retries.join_next(), if !retries.is_empty() => continue,
        };
        match error {
            Ok(error) if error.is_recoverable => {
                let client = client.clone();
                let online = online.clone();
                retries.spawn(async move {
                    matrix_sdk::sleep::sleep(SEND_QUEUE_RETRY).await;
                    if *online.borrow() == ConnectionState::Online
                        && let Some(room) = client.get_room(&error.room_id)
                    {
                        room.send_queue().set_enabled(true);
                    }
                });
            }
            Ok(_) | Err(RecvError::Lagged(_)) => {}
            Err(RecvError::Closed) => return,
        }
    }
}

/// the one place upload events leave from, so an upload is never reported after it
/// ended. the client-wide channel exists before any upload does. it holds 32 updates and
/// progress ticks share it, so this keeps draining while the event stream is full
pub(super) async fn uploads(
    client: Client,
    mut updates: tokio::sync::broadcast::Receiver<SendQueueUpdate>,
    events: EventSender,
) {
    let (_, mut rooms) = client.rooms_stream();
    let mut tracker = UploadTracker::new(attachment_echoes(&client).await);
    let mut pending = VecDeque::new();
    loop {
        let reports: Vec<Event> = tokio::select! {
            biased;
            update = updates.recv() => match update {
                Ok(update) => tracker.update(update).into_iter().collect(),
                Err(RecvError::Lagged(skipped)) => {
                    tracing::warn!(skipped, "upload tracker lagged");
                    tracker.resnapshot(attachment_echoes(&client).await)
                }
                Err(RecvError::Closed) => return,
            },
            update = rooms.next() => {
                if update.is_none() {
                    return;
                }
                tracker.finish_where(
                    |(room_id, _)| client.get_room(room_id).is_none(),
                    &UploadOutcome::RoomForgotten,
                )
            }
            permit = events.reserve(), if !pending.is_empty() => {
                let Ok(permit) = permit else { return };
                permit.send(pending.pop_front().expect("pending was checked"));
                continue;
            }
        };
        for event in reports {
            queue_upload_event(&mut pending, event);
        }
    }
}

/// a progress tick replaces a waiting one for the same upload, unless something else about
/// that upload is queued after it
fn queue_upload_event(pending: &mut VecDeque<Event>, event: Event) {
    if matches!(event, Event::MediaUploadProgress { .. })
        && let Some(waiting) = pending
            .iter_mut()
            .rev()
            .find(|waiting| upload_of(waiting) == upload_of(&event))
        && matches!(waiting, Event::MediaUploadProgress { .. })
    {
        *waiting = event;
        return;
    }
    pending.push_back(event);
}

fn upload_of(event: &Event) -> Option<(&OwnedRoomId, &OwnedTransactionId)> {
    match event {
        Event::MediaUploadProgress {
            room_id,
            transaction_id,
            ..
        }
        | Event::MediaUploadFailed {
            room_id,
            transaction_id,
            ..
        }
        | Event::MediaUploadFinished {
            room_id,
            transaction_id,
            ..
        } => Some((room_id, transaction_id)),
        _ => None,
    }
}

type Upload = (OwnedRoomId, OwnedTransactionId);

/// the sdk sends progress ticks from a detached task, so a tick can land after its
/// upload's end. it trails by one task hop, so a short memory of ends covers it
const REMEMBERED_ENDS: usize = 256;

/// everything reported gets exactly one `MediaUploadFinished` and nothing after it
struct UploadTracker {
    live: HashSet<Upload>,
    ended: HashSet<Upload>,
    ended_order: VecDeque<Upload>,
}

impl UploadTracker {
    fn new(live: HashSet<Upload>) -> Self {
        Self {
            live,
            ended: HashSet::new(),
            ended_order: VecDeque::new(),
        }
    }

    fn update(&mut self, update: SendQueueUpdate) -> Option<Event> {
        let room_id = update.room_id;
        let (transaction_id, outcome) = match update.update {
            RoomSendQueueUpdate::NewLocalEvent(echo) => {
                let upload = (room_id, echo.transaction_id);
                if is_attachment_echo(&echo.content) && !self.ended.contains(&upload) {
                    self.live.insert(upload);
                }
                return None;
            }
            RoomSendQueueUpdate::MediaUpload {
                related_to,
                file,
                progress,
                ..
            } => {
                let upload = (room_id, related_to);
                if self.ended.contains(&upload) {
                    return None;
                }
                self.live.insert(upload.clone());
                let (room_id, transaction_id) = upload;
                return Some(Event::MediaUploadProgress {
                    room_id,
                    transaction_id,
                    uploaded_bytes: progress.current as u64,
                    total_bytes: progress.total as u64,
                    complete: file.is_some() && progress.current == progress.total,
                });
            }
            RoomSendQueueUpdate::SendError {
                transaction_id,
                error,
                is_recoverable,
            } => {
                let upload = (room_id, transaction_id);
                if !self.live.contains(&upload) {
                    return None;
                }
                let (room_id, transaction_id) = upload;
                return Some(Event::MediaUploadFailed {
                    room_id,
                    transaction_id,
                    kind: classify_sdk_error(&error),
                    will_retry: is_recoverable,
                });
            }
            RoomSendQueueUpdate::SentEvent {
                transaction_id,
                event_id,
            } => (transaction_id, UploadOutcome::Sent { event_id }),
            RoomSendQueueUpdate::CancelledLocalEvent { transaction_id } => {
                (transaction_id, UploadOutcome::Cancelled)
            }
            _ => return None,
        };
        self.finish((room_id, transaction_id), outcome)
    }

    fn resnapshot(&mut self, current: HashSet<Upload>) -> Vec<Event> {
        let reports = self.finish_where(|upload| !current.contains(upload), &UploadOutcome::Missed);
        let ended = &self.ended;
        self.live
            .extend(current.into_iter().filter(|upload| !ended.contains(upload)));
        reports
    }

    fn finish_where(
        &mut self,
        ends: impl Fn(&Upload) -> bool,
        outcome: &UploadOutcome,
    ) -> Vec<Event> {
        let ending = self
            .live
            .iter()
            .filter(|upload| ends(upload))
            .cloned()
            .collect::<Vec<_>>();
        ending
            .into_iter()
            .filter_map(|upload| self.finish(upload, outcome.clone()))
            .collect()
    }

    /// an end for an upload we never saw is remembered too, its ticks may still be coming
    fn finish(&mut self, upload: Upload, outcome: UploadOutcome) -> Option<Event> {
        if self.ended.insert(upload.clone()) {
            self.ended_order.push_back(upload.clone());
            if self.ended_order.len() > REMEMBERED_ENDS
                && let Some(oldest) = self.ended_order.pop_front()
            {
                self.ended.remove(&oldest);
            }
        }
        if !self.live.remove(&upload) {
            return None;
        }
        let (room_id, transaction_id) = upload;
        Some(Event::MediaUploadFinished {
            room_id,
            transaction_id,
            outcome,
        })
    }
}

async fn attachment_echoes(client: &Client) -> HashSet<Upload> {
    let Ok(echoes) = client.send_queue().local_echoes().await else {
        return HashSet::new();
    };
    echoes
        .into_iter()
        .flat_map(|(room_id, echoes)| {
            echoes
                .into_iter()
                .filter(|echo| is_attachment_echo(&echo.content))
                .map(move |echo| (room_id.clone(), echo.transaction_id))
        })
        .collect()
}

fn is_attachment_echo(content: &LocalEchoContent) -> bool {
    let LocalEchoContent::Event {
        serialized_event, ..
    } = content
    else {
        return false;
    };
    matches!(
        serialized_event.deserialize(),
        Ok(AnyMessageLikeEventContent::RoomMessage(content))
            if matches!(
                content.msgtype,
                MessageType::Audio(_) | MessageType::File(_) | MessageType::Image(_) | MessageType::Video(_)
            )
    )
}

pub(super) async fn encryption(
    client: Client,
    events: EventSender,
    mut enabling: watch::Receiver<bool>,
) {
    let Some(device_id) = client.device_id().map(ToOwned::to_owned) else {
        tracing::warn!("encryption watcher stopped because the client has no device id");
        return;
    };
    let recovery = client.encryption().recovery();
    let backups = client.encryption().backups();
    // until this is cached, the timeline labels undecryptable events from a local guess
    matrix_sdk::executor::spawn({
        let backups = backups.clone();
        async move {
            let _ = backups.exists_on_server().await;
        }
    });
    let mut recovery_stream = recovery.state_stream();
    let mut backup_stream = backups.state_stream();
    let mut verification_stream = client.encryption().verification_state();
    let mut recovery_state = recovery.state();
    let mut backup_state = backups.state();
    let mut verification_state = verification_stream.get();
    loop {
        let mut state = encryption_state(
            device_id.clone(),
            recovery_state,
            verification_state,
            backup_state,
        );
        if *enabling.borrow_and_update() {
            state.recovery = crate::RecoveryState::Enabling;
        }
        if events.send(Event::EncryptionChanged(state)).await.is_err() {
            return;
        }
        tokio::select! {
            Some(update) = recovery_stream.next() => recovery_state = update,
            Ok(()) = enabling.changed() => {}
            Some(update) = verification_stream.next() => verification_state = update,
            Some(update) = backup_stream.next() => {
                backup_state = update.unwrap_or_else(|_| backups.state());
            }
            else => return,
        }
    }
}

/// space orders, pins and power levels all feed summaries without raising a notable room info update
pub(super) async fn rooms(client: Client, events: EventSender) {
    let (_, mut rooms) = client.rooms_stream();
    let mut info = client.room_info_notable_update_receiver();
    let changed = Arc::new(Notify::new());
    let handlers = [
        client.add_event_handler({
            let changed = changed.clone();
            move |event: Raw<AnyRoomAccountDataEvent>| {
                let changed = changed.clone();
                async move {
                    if event_type(&event).as_deref() == Some(space_order::EVENT_TYPE) {
                        changed.notify_one();
                    }
                }
            }
        }),
        client.add_event_handler({
            let changed = changed.clone();
            move |event: Raw<AnySyncStateEvent>| {
                let changed = changed.clone();
                async move {
                    if matches!(
                        event_type(&event).as_deref(),
                        Some("m.room.pinned_events" | "m.room.power_levels")
                    ) {
                        changed.notify_one();
                    }
                }
            }
        }),
    ];
    let (touches, mut touched) = mpsc::unbounded_channel();
    let thread_handlers = thread_handlers(&client, &touches);
    let mut timelines = client.event_cache().subscribe_to_room_generic_updates();
    'invalidate: loop {
        if events.send(Event::RoomsChanged).await.is_err() {
            break;
        }
        // thread events and threaded receipts leave room info alone, so neither reaches `info`
        let mut pending = Touched::default();
        loop {
            let mut invalidate = tokio::select! {
                update = rooms.next() => if update.is_none() { break 'invalidate } else { true },
                update = info.recv() => {
                    if matches!(update, Err(RecvError::Closed)) { break 'invalidate }
                    true
                }
                () = changed.notified() => true,
                Some((room_id, root)) = touched.recv() => {
                    pending.add(room_id, root);
                    false
                }
                update = timelines.recv() => match update {
                    Ok(update) => {
                        pending.add(update.room_id, None);
                        false
                    }
                    Err(RecvError::Lagged(_)) => true,
                    Err(RecvError::Closed) => break 'invalidate,
                },
            };
            // a sync's room info lands before the event cache has stored its events, so
            // give it a moment and take the whole sync as one invalidation
            matrix_sdk::sleep::sleep(SYNC_SETTLE).await;
            loop {
                match info.try_recv() {
                    Ok(_) | Err(TryRecvError::Lagged(_)) => invalidate = true,
                    Err(_) => break,
                }
            }
            while let Ok((room_id, root)) = touched.try_recv() {
                pending.add(room_id, root);
            }
            loop {
                match timelines.try_recv() {
                    Ok(update) => pending.add(update.room_id, None),
                    Err(TryRecvError::Lagged(_)) => invalidate = true,
                    Err(_) => break,
                }
            }
            for (room_id, roots) in pending.0.drain() {
                if let Some(room) = client.get_room(&room_id)
                    && recount_thread_unreads(&room, roots).await.is_some()
                {
                    invalidate = true;
                }
            }
            if invalidate {
                break;
            }
        }
    }
    for handler in handlers.into_iter().chain(thread_handlers) {
        client.remove_event_handler(handler);
    }
}

const SYNC_SETTLE: Duration = Duration::from_millis(100);

/// rooms to recount, with the threads a sync brought news of
#[derive(Default)]
struct Touched(HashMap<OwnedRoomId, HashSet<OwnedEventId>>);

impl Touched {
    fn add(&mut self, room_id: OwnedRoomId, root: Option<OwnedEventId>) {
        let roots = self.0.entry(room_id).or_default();
        roots.extend(root);
    }
}

/// a thread event names its thread before the event cache has it, and only the user's own
/// receipts move their counts
fn thread_handlers(
    client: &Client,
    touches: &mpsc::UnboundedSender<(OwnedRoomId, Option<OwnedEventId>)>,
) -> [EventHandlerHandle; 2] {
    let own_user_id = client.user_id().map(ToOwned::to_owned);
    [
        client.add_event_handler({
            let touches = touches.clone();
            move |event: Raw<AnySyncTimelineEvent>, room: Room| {
                if let Some(root) = extract_thread_root(&event) {
                    let _ = touches.send((room.room_id().to_owned(), Some(root)));
                }
                async {}
            }
        }),
        client.add_event_handler({
            let touches = touches.clone();
            move |event: SyncReceiptEvent, room: Room| {
                let own = own_user_id.as_ref().is_some_and(|own_user_id| {
                    event
                        .content
                        .0
                        .values()
                        .flat_map(|receipts| receipts.values())
                        .any(|users| users.contains_key(own_user_id))
                });
                if own {
                    let _ = touches.send((room.room_id().to_owned(), None));
                }
                async {}
            }
        }),
    ]
}

/// packs live in room state and account data, and leaving a room takes its packs with it
pub(super) async fn image_packs(client: Client, events: EventSender) {
    let changed = Arc::new(Notify::new());
    let handlers = [
        client.add_event_handler({
            let changed = changed.clone();
            move |event: Raw<AnySyncStateEvent>| {
                let changed = changed.clone();
                async move {
                    if event_type(&event).as_deref() == Some(ROOM_PACKS) {
                        changed.notify_one();
                    }
                }
            }
        }),
        client.add_event_handler({
            let changed = changed.clone();
            move |event: Raw<AnyGlobalAccountDataEvent>| {
                let changed = changed.clone();
                async move {
                    if matches!(
                        event_type(&event).as_deref(),
                        Some(USER_PACK | ENABLED_ROOMS)
                    ) {
                        changed.notify_one();
                    }
                }
            }
        }),
    ];
    let mut info = client.room_info_notable_update_receiver();
    loop {
        tokio::select! {
            () = changed.notified() => {}
            update = info.recv() => match update {
                Ok(update) if !update.reasons.contains(RoomInfoNotableUpdateReasons::MEMBERSHIP) => {
                    continue;
                }
                Ok(_) | Err(RecvError::Lagged(_)) => {}
                Err(RecvError::Closed) => break,
            },
        }
        if events.send(Event::ImagePacksChanged).await.is_err() {
            break;
        }
    }
    for handler in handlers {
        client.remove_event_handler(handler);
    }
}

fn event_type<T>(event: &Raw<T>) -> Option<String> {
    event.get_field("type").ok().flatten()
}

pub(super) async fn notification_settings(client: Client, events: EventSender) {
    let settings = client.notification_settings().await;
    let mut changes = settings.subscribe_to_changes();
    while let Ok(()) | Err(RecvError::Lagged(_)) = changes.recv().await {
        // room summaries carry the effective mode, and push rule changes never touch room info
        for event in [Event::NotificationSettingsChanged, Event::RoomsChanged] {
            if events.send(event).await.is_err() {
                return;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use matrix_sdk::send_queue::{AbstractProgress, RoomSendQueueUpdate, SendQueueUpdate};
    use ruma::{OwnedTransactionId, event_id, owned_room_id};

    use super::{HashSet, Upload, UploadTracker};
    use crate::{Event, UploadOutcome};

    fn upload() -> Upload {
        (
            owned_room_id!("!room:localhost"),
            OwnedTransactionId::from("upload"),
        )
    }

    fn update(update: RoomSendQueueUpdate) -> SendQueueUpdate {
        SendQueueUpdate {
            room_id: upload().0,
            update,
        }
    }

    fn tick() -> SendQueueUpdate {
        update(RoomSendQueueUpdate::MediaUpload {
            related_to: upload().1,
            file: None,
            index: 0,
            progress: AbstractProgress {
                current: 1,
                total: 2,
            },
        })
    }

    fn sent() -> SendQueueUpdate {
        update(RoomSendQueueUpdate::SentEvent {
            transaction_id: upload().1,
            event_id: event_id!("$sent:localhost").to_owned(),
        })
    }

    fn finished(event: Option<Event>) -> Option<UploadOutcome> {
        match event? {
            Event::MediaUploadFinished { outcome, .. } => Some(outcome),
            _ => None,
        }
    }

    #[test]
    fn a_tick_trailing_its_end_says_nothing() {
        let mut tracker = UploadTracker::new(HashSet::new());
        assert!(tracker.update(tick()).is_some());
        assert!(matches!(
            finished(tracker.update(sent())),
            Some(UploadOutcome::Sent { .. })
        ));
        assert!(tracker.update(tick()).is_none());
        assert!(tracker.live.is_empty());
    }

    #[test]
    fn a_stale_tick_after_a_missed_end_changes_nothing() {
        let mut tracker = UploadTracker::new(HashSet::new());
        assert!(tracker.update(tick()).is_some());
        let missed = tracker.resnapshot(HashSet::new());
        assert!(matches!(
            missed.as_slice(),
            [Event::MediaUploadFinished {
                outcome: UploadOutcome::Missed,
                ..
            }]
        ));
        assert!(tracker.update(tick()).is_none());
        assert!(tracker.update(sent()).is_none());
        assert!(tracker.live.is_empty());
    }
}