macltok

public
Star0
Watch0
60 commits 3 branches 0 tags 0 releases3.0 MiB
use std::{
    collections::HashMap,
    fmt,
    pin::Pin,
    sync::{Arc, Mutex, MutexGuard, PoisonError},
    task::{Context, Poll},
    time::Duration,
};

use futures_util::{Stream, StreamExt};
use imbl::Vector;
use matrix_sdk::{
    SendOutsideWasm,
    executor::{AbortHandle, spawn},
};
use ruma::{OwnedDeviceId, OwnedRoomId, OwnedTransactionId, OwnedUserId};
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;

use crate::{
    ClientHandle, Command, ConnectionState, EncryptionState, EngineFailure, ErrorKind, Event,
    HistoryRepairReport, HistoryState, ImagePacks, RoomSummary, TimelineCommand, TimelineDiff,
    TimelineId, TimelineItem, TypingState, UploadOutcome, VerificationState, VerificationTarget,
    engine::executor::EngineBoxFuture, transport::PublicationRegistration,
};

#[cfg(not(target_family = "wasm"))]
type EngineBoxStream<T> = futures_util::stream::BoxStream<'static, T>;
#[cfg(target_family = "wasm")]
type EngineBoxStream<T> = futures_util::stream::LocalBoxStream<'static, T>;

const DEFAULT_CAPACITY: usize = 256;
const RESYNC_BACKOFF: Duration = Duration::from_millis(100);
const RESYNC_BACKOFF_CAP: Duration = Duration::from_secs(10);

/// everything the engine has said so far, folded together
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[non_exhaustive]
pub struct State {
    pub user_id: Option<OwnedUserId>,
    pub session: SessionStatus,
    pub connection: ConnectionState,
    pub encryption: Option<EncryptionState>,
    pub verification: Option<VerificationState>,
    pub verification_target: Option<VerificationTarget>,
    pub mentions: Option<crate::MentionsPage>,
    pub rooms: Vector<RoomSummary>,
    pub image_packs: ImagePacks,
    pub timelines: imbl::HashMap<TimelineId, TimelineState>,
    pub repairs: imbl::HashMap<OwnedRoomId, HistoryRepairReport>,
    pub typing: imbl::HashMap<OwnedRoomId, Vec<OwnedUserId>>,
    pub uploads: imbl::HashMap<OwnedTransactionId, UploadState>,
    /// what stopped the engine, if something went wrong
    pub engine_failure: Option<EngineFailure>,
    /// the last save of the session record that failed, until one succeeds
    pub session_save_failure: Option<String>,
}

#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[non_exhaustive]
pub struct TimelineState {
    pub items: Vector<TimelineItem>,
    pub history: HistoryState,
    /// a batch didn't fit, so `items` is the last good copy until the next reset
    pub diverged: bool,
}

#[derive(Clone, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[non_exhaustive]
pub struct UploadState {
    pub room_id: OwnedRoomId,
    pub uploaded_bytes: u64,
    pub total_bytes: u64,
    pub complete: bool,
    pub failure: Option<UploadFailure>,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[non_exhaustive]
pub struct UploadFailure {
    pub kind: ErrorKind,
    pub will_retry: bool,
}

/// the `Termination` it'll end in, as far as the events tell
#[derive(Clone, Debug, Eq, PartialEq)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum SessionStatus {
    Active,
    SoftLoggedOut {
        device_id: OwnedDeviceId,
    },
    AccessRevoked,
    LoggedOut,
    /// stopped with the session still good
    Stopped,
}

#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct Update {
    /// one more than the update before it, and a `State` carries the one it's current as of
    pub sequence: u64,
    pub change: StateChange,
}

#[derive(Clone, Debug)]
pub enum StateChange {
    /// start over from this. a subscription opens with one and gets another after lagging
    State(Arc<State>),
    Authenticated(OwnedUserId),
    Session(SessionStatus),
    /// the record on disk may now hold a spent refresh token
    SessionSaveFailed(String),
    /// saved again after `SessionSaveFailed`
    SessionSaved,
    Connection(ConnectionState),
    Encryption(EncryptionState),
    Verification(VerificationState),
    VerificationTarget(VerificationTarget),
    Mentions(crate::MentionsUpdate),
    /// room modes arrive with the `Rooms` that follows
    NotificationSettings,
    Rooms(Vector<RoomSummary>),
    ImagePacks(ImagePacks),
    Timeline {
        timeline_id: TimelineId,
        change: TimelineChange,
    },
    History {
        timeline_id: TimelineId,
        state: HistoryState,
    },
    HistoryRepair(HistoryRepairReport),
    Typing(TypingState),
    Upload {
        transaction_id: OwnedTransactionId,
        upload: UploadState,
    },
    UploadFinished {
        room_id: OwnedRoomId,
        transaction_id: OwnedTransactionId,
        outcome: UploadOutcome,
    },
    Failed(EngineFailure),
}

#[derive(Clone, Debug)]
pub enum TimelineChange {
    Reset {
        items: Vector<TimelineItem>,
    },
    /// `items` is what applying `diffs` left
    Changed {
        diffs: Arc<[TimelineDiff]>,
        items: Vector<TimelineItem>,
    },
    /// a batch didn't fit. a resync is on its way and nothing changes until its reset
    Diverged,
    Closed,
}

/// drains the event stream into a `State` that any number of `MirrorHandle`s watch
pub struct Mirror {
    handle: ClientHandle,
    events: EngineBoxStream<Event>,
    shared: Arc<Shared>,
    _publication: PublicationRegistration,
    resyncs: HashMap<TimelineId, AbortHandle>,
    rooms_fetch: Option<AbortHandle>,
    packs_fetch: Option<AbortHandle>,
}

impl fmt::Debug for Mirror {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_struct("Mirror").finish_non_exhaustive()
    }
}

impl Mirror {
    /// takes any stream of the engine's events, so a host can still look at them on the way past
    pub fn new(
        handle: ClientHandle,
        events: impl Stream<Item = Event> + SendOutsideWasm + 'static,
    ) -> (Self, MirrorHandle) {
        Self::with_capacity(handle, events, DEFAULT_CAPACITY)
    }

    /// `capacity` updates can queue for a subscriber before it lags and starts over. a zero
    /// capacity counts as one
    pub fn with_capacity(
        handle: ClientHandle,
        events: impl Stream<Item = Event> + SendOutsideWasm + 'static,
        capacity: usize,
    ) -> (Self, MirrorHandle) {
        let state = State {
            user_id: None,
            session: SessionStatus::Active,
            connection: *handle.connection_state().borrow(),
            encryption: handle.encryption_state(),
            verification: None,
            verification_target: None,
            mentions: None,
            rooms: Vector::new(),
            image_packs: ImagePacks::default(),
            timelines: imbl::HashMap::new(),
            repairs: imbl::HashMap::new(),
            typing: imbl::HashMap::new(),
            uploads: imbl::HashMap::new(),
            engine_failure: None,
            session_save_failure: None,
        };
        let shared = Arc::new(Shared {
            published: Mutex::new(Published {
                state: Arc::new(state),
                sequence: 0,
                updates: Some(broadcast::channel(capacity.max(1)).0),
                rooms_fetch: Fetch::Idle,
                packs_fetch: Fetch::Idle,
            }),
        });
        let publication = handle.register_publication_mirror();
        #[cfg(not(target_family = "wasm"))]
        let events = events.boxed();
        #[cfg(target_family = "wasm")]
        let events = events.boxed_local();
        let mirror = Self {
            handle,
            events,
            shared: shared.clone(),
            _publication: publication,
            resyncs: HashMap::new(),
            rooms_fetch: None,
            packs_fetch: None,
        };
        (mirror, MirrorHandle { shared })
    }

    /// resolves once the event stream ends, which is just after `Stopped`
    pub fn run(mut self) -> EngineBoxFuture<'static, ()> {
        Box::pin(async move {
            self.refetch_rooms();
            self.refetch_image_packs();
            while let Some(event) = self.events.next().await {
                self.apply(event);
            }
        })
    }

    fn apply(&mut self, event: Event) {
        match event {
            Event::ConnectionChanged(connection) => self.publish(|state| {
                state.connection = connection;
                Some(StateChange::Connection(connection))
            }),
            Event::Authenticated { user_id } => self.publish(|state| {
                state.user_id = Some(user_id.clone());
                Some(StateChange::Authenticated(user_id))
            }),
            Event::SessionSaveFailed { message } => self.publish(|state| {
                state.session_save_failure = Some(message.clone());
                Some(StateChange::SessionSaveFailed(message))
            }),
            Event::SessionTokensChanged => self.publish(|state| {
                state.session_save_failure.take()?;
                Some(StateChange::SessionSaved)
            }),
            Event::SessionAccessRejected { soft_logout } => self.publish(|state| {
                state.mentions = None;
                state.session = match soft_logout {
                    Some(device_id) => SessionStatus::SoftLoggedOut { device_id },
                    None => SessionStatus::AccessRevoked,
                };
                Some(StateChange::Session(state.session.clone()))
            }),
            Event::LoggedOut => self.publish(|state| {
                state.mentions = None;
                state.session = SessionStatus::LoggedOut;
                Some(StateChange::Session(SessionStatus::LoggedOut))
            }),
            Event::Stopped => {
                self.cancel_resyncs();
                self.publish(|state| {
                    if state.session != SessionStatus::Active {
                        return None;
                    }
                    state.session = SessionStatus::Stopped;
                    Some(StateChange::Session(SessionStatus::Stopped))
                });
            }
            Event::Failed(failure) => self.publish(|state| {
                state.engine_failure = Some(failure.clone());
                Some(StateChange::Failed(failure))
            }),
            Event::NotificationSettingsChanged => {
                self.publish(|_| Some(StateChange::NotificationSettings));
            }
            Event::RoomsChanged => self.refetch_rooms(),
            Event::MentionsChanged(update) => self.publish(|state| {
                if state.session != SessionStatus::Active
                    || state
                        .user_id
                        .as_ref()
                        .is_some_and(|account| account != &update.snapshot.boundary.account)
                {
                    return None;
                }
                if state
                    .mentions
                    .as_ref()
                    .map_or(false, |current| update.snapshot.revision < current.revision)
                {
                    return None;
                }
                state.mentions = Some(update.snapshot.clone());
                Some(StateChange::Mentions(update))
            }),
            Event::ImagePacksChanged => self.refetch_image_packs(),
            Event::EncryptionChanged(encryption) => self.publish(|state| {
                state.encryption = Some(encryption.clone());
                Some(StateChange::Encryption(encryption))
            }),
            Event::VerificationChanged(verification) => self.publish(|state| {
                state.verification = Some(verification.clone());
                Some(StateChange::Verification(verification))
            }),
            Event::VerificationTargetChanged(target) => self.publish(|state| {
                state.verification_target = Some(target.clone());
                Some(StateChange::VerificationTarget(target))
            }),
            Event::TimelineReset { timeline_id, items } => {
                self.cancel_resync(&timeline_id);
                self.publish(|state| {
                    let history = state
                        .timelines
                        .get(&timeline_id)
                        .map_or(HistoryState::Idle, |timeline| timeline.history.clone());
                    let timeline = TimelineState {
                        items: items.clone(),
                        history,
                        diverged: false,
                    };
                    state.timelines.insert(timeline_id.clone(), timeline);
                    Some(StateChange::Timeline {
                        timeline_id,
                        change: TimelineChange::Reset { items },
                    })
                });
            }
            Event::TimelineChanged { timeline_id, diffs } => self.apply_diffs(timeline_id, diffs),
            Event::HistoryChanged {
                timeline_id,
                state: history,
            } => self.publish(|state| {
                state.timelines.get_mut(&timeline_id)?.history = history.clone();
                Some(StateChange::History {
                    timeline_id,
                    state: history,
                })
            }),
            Event::HistoryRepairChanged(report) => self.publish(|state| {
                state.repairs.insert(report.room_id.clone(), report.clone());
                Some(StateChange::HistoryRepair(report))
            }),
            Event::HistoryRepairPublished {
                report,
                acknowledgement,
            } => {
                self.publish(|state| {
                    state.repairs.insert(report.room_id.clone(), report.clone());
                    Some(StateChange::HistoryRepair(report))
                });
                if let Some(acknowledgement) = acknowledgement {
                    acknowledgement.complete(true);
                }
            }
            Event::TimelineClosed { timeline_id } => {
                self.cancel_resync(&timeline_id);
                self.publish(|state| {
                    state.timelines.remove(&timeline_id)?;
                    Some(StateChange::Timeline {
                        timeline_id,
                        change: TimelineChange::Closed,
                    })
                });
            }
            Event::TypingChanged(typing) => self.publish(|state| {
                if typing.users.is_empty() {
                    state.typing.remove(&typing.room_id);
                } else {
                    state
                        .typing
                        .insert(typing.room_id.clone(), typing.users.clone());
                }
                Some(StateChange::Typing(typing))
            }),
            Event::MediaUploadProgress {
                room_id,
                transaction_id,
                uploaded_bytes,
                total_bytes,
                complete,
            } => self.publish(|state| {
                let upload = UploadState {
                    room_id,
                    uploaded_bytes,
                    total_bytes,
                    complete,
                    failure: None,
                };
                state.uploads.insert(transaction_id.clone(), upload.clone());
                Some(StateChange::Upload {
                    transaction_id,
                    upload,
                })
            }),
            Event::MediaUploadFailed {
                room_id,
                transaction_id,
                kind,
                will_retry,
            } => self.publish(|state| {
                let upload = state
                    .uploads
                    .entry(transaction_id.clone())
                    .or_insert_with(|| UploadState {
                        room_id,
                        uploaded_bytes: 0,
                        total_bytes: 0,
                        complete: false,
                        failure: None,
                    });
                upload.failure = Some(UploadFailure { kind, will_retry });
                Some(StateChange::Upload {
                    transaction_id,
                    upload: upload.clone(),
                })
            }),
            Event::MediaUploadFinished {
                room_id,
                transaction_id,
                outcome,
            } => self.publish(|state| {
                state.uploads.remove(&transaction_id);
                Some(StateChange::UploadFinished {
                    room_id,
                    transaction_id,
                    outcome,
                })
            }),
        }
    }

    fn apply_diffs(&mut self, timeline_id: TimelineId, diffs: Vec<TimelineDiff>) {
        let mut diverged = false;
        self.publish(|state| {
            let Some(timeline) = state.timelines.get_mut(&timeline_id) else {
                diverged = true;
                state.timelines.insert(
                    timeline_id.clone(),
                    TimelineState {
                        items: Vector::new(),
                        history: HistoryState::Idle,
                        diverged: true,
                    },
                );
                return Some(StateChange::Timeline {
                    timeline_id: timeline_id.clone(),
                    change: TimelineChange::Diverged,
                });
            };
            if timeline.diverged {
                return None;
            }
            let diffs: Arc<[TimelineDiff]> = diffs.into();
            let change = match TimelineDiff::try_apply_all(diffs.to_vec(), &mut timeline.items) {
                Ok(()) => TimelineChange::Changed {
                    diffs,
                    items: timeline.items.clone(),
                },
                Err(error) => {
                    tracing::debug!(?timeline_id, length = error.length, "timeline diverged");
                    timeline.diverged = true;
                    diverged = true;
                    TimelineChange::Diverged
                }
            };
            Some(StateChange::Timeline {
                timeline_id: timeline_id.clone(),
                change,
            })
        });
        if diverged {
            let task = spawn(resync(self.handle.clone(), timeline_id.clone()));
            self.resyncs.insert(timeline_id, task.abort_handle());
        }
    }

    fn refetch_rooms(&mut self) {
        let fetch = fetch_rooms(self.handle.clone(), self.shared.clone());
        if let Some(task) = self
            .shared
            .start(|published| &mut published.rooms_fetch, fetch)
        {
            self.rooms_fetch = Some(task);
        }
    }

    fn refetch_image_packs(&mut self) {
        let fetch = fetch_image_packs(self.handle.clone(), self.shared.clone());
        if let Some(task) = self
            .shared
            .start(|published| &mut published.packs_fetch, fetch)
        {
            self.packs_fetch = Some(task);
        }
    }

    fn cancel_resync(&mut self, timeline_id: &TimelineId) {
        if let Some(task) = self.resyncs.remove(timeline_id) {
            task.abort();
        }
    }

    fn cancel_resyncs(&mut self) {
        for (_, task) in self.resyncs.drain() {
            task.abort();
        }
    }

    fn publish(&self, change: impl FnOnce(&mut State) -> Option<StateChange>) {
        self.shared.publish(change);
    }
}

impl Drop for Mirror {
    fn drop(&mut self) {
        self.cancel_resyncs();
        for task in [self.rooms_fetch.take(), self.packs_fetch.take()]
            .into_iter()
            .flatten()
        {
            task.abort();
        }
        self.shared.lock().updates = None;
    }
}

/// asks until the engine answers, since a busy room or a full queue is only ever for now.
/// a reset or a close for the timeline aborts it, and the engine sends a reset before it answers
async fn resync(handle: ClientHandle, timeline_id: TimelineId) {
    let mut backoff = RESYNC_BACKOFF;
    loop {
        let command = Command::Timeline(TimelineCommand::Resync {
            timeline_id: timeline_id.clone(),
        });
        match handle.execute(command).await {
            Ok(_) => return,
            Err(error) if error.kind == ErrorKind::Interrupted => return,
            Err(error) => {
                tracing::debug!(%error, ?timeline_id, "resync refused, trying again");
                matrix_sdk::sleep::sleep(backoff).await;
                backoff = (backoff * 2).min(RESYNC_BACKOFF_CAP);
            }
        }
    }
}

async fn fetch_rooms(handle: ClientHandle, shared: Arc<Shared>) {
    loop {
        let rooms: Vector<RoomSummary> = handle.rooms().await.into();
        shared.publish(|state| {
            state.rooms = rooms.clone();
            Some(StateChange::Rooms(rooms))
        });
        if !shared.fetched(|published| &mut published.rooms_fetch) {
            return;
        }
    }
}

/// a failed read keeps the packs from before, the next change asks again
async fn fetch_image_packs(handle: ClientHandle, shared: Arc<Shared>) {
    loop {
        match handle.image_packs().await {
            Ok(packs) => shared.publish(|state| {
                if state.image_packs == packs {
                    return None;
                }
                state.image_packs = packs.clone();
                Some(StateChange::ImagePacks(packs))
            }),
            Err(error) => tracing::warn!(%error, "reading image packs failed"),
        }
        if !shared.fetched(|published| &mut published.packs_fetch) {
            return;
        }
    }
}

/// cheap to clone, one per frontend
#[derive(Clone, Debug)]
pub struct MirrorHandle {
    shared: Arc<Shared>,
}

impl MirrorHandle {
    #[must_use]
    pub fn snapshot(&self) -> Arc<State> {
        self.shared.lock().state.clone()
    }

    /// opens with the current state, then every update after it. ends once the mirror stops
    #[must_use]
    pub fn subscribe(&self) -> Subscription {
        let (first, receiving) = self.shared.clone().subscribe();
        let rest = futures_util::stream::unfold(receiving, |(shared, receiver)| async move {
            let mut receiver = receiver?;
            match receiver.recv().await {
                Ok(update) => Some((update, (shared, Some(receiver)))),
                Err(broadcast::error::RecvError::Lagged(missed)) => {
                    tracing::debug!(missed, "mirror subscriber lagged, starting it over");
                    Some(shared.subscribe())
                }
                Err(broadcast::error::RecvError::Closed) => None,
            }
        });
        let updates = futures_util::stream::once(std::future::ready(first)).chain(rest);
        #[cfg(not(target_family = "wasm"))]
        let updates = updates.boxed();
        #[cfg(target_family = "wasm")]
        let updates = updates.boxed_local();
        Subscription(updates)
    }
}

pub struct Subscription(EngineBoxStream<Update>);

impl fmt::Debug for Subscription {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("Subscription")
            .finish_non_exhaustive()
    }
}

impl Stream for Subscription {
    type Item = Update;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Update>> {
        self.0.poll_next_unpin(context)
    }
}

#[derive(Debug)]
struct Shared {
    published: Mutex<Published>,
}

#[derive(Debug)]
struct Published {
    state: Arc<State>,
    sequence: u64,
    /// gone once the mirror stops, which ends every subscription
    updates: Option<broadcast::Sender<Update>>,
    rooms_fetch: Fetch,
    packs_fetch: Fetch,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Fetch {
    Idle,
    Running,
    Again,
}

impl Shared {
    fn lock(&self) -> MutexGuard<'_, Published> {
        self.published
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
    }

    /// one fetch at a time, with at most one more owed
    fn start(
        &self,
        fetch: impl FnOnce(&mut Published) -> &mut Fetch,
        task: impl Future<Output = ()> + SendOutsideWasm + 'static,
    ) -> Option<AbortHandle> {
        let mut published = self.lock();
        let fetch = fetch(&mut published);
        match *fetch {
            Fetch::Idle => {
                *fetch = Fetch::Running;
                drop(published);
                Some(spawn(task).abort_handle())
            }
            Fetch::Running => {
                *fetch = Fetch::Again;
                None
            }
            Fetch::Again => None,
        }
    }

    /// whether another fetch was asked for while this one ran
    fn fetched(&self, fetch: impl FnOnce(&mut Published) -> &mut Fetch) -> bool {
        let mut published = self.lock();
        let fetch = fetch(&mut published);
        if *fetch == Fetch::Again {
            *fetch = Fetch::Running;
            return true;
        }
        *fetch = Fetch::Idle;
        false
    }

    /// state, sequence and send move together under the one lock, which is what lets
    /// `subscribe` take a snapshot and a receiver with nothing between them
    fn publish(&self, change: impl FnOnce(&mut State) -> Option<StateChange>) {
        let mut published = self.lock();
        let Some(change) = change(Arc::make_mut(&mut published.state)) else {
            return;
        };
        published.sequence += 1;
        let update = Update {
            sequence: published.sequence,
            change,
        };
        if let Some(updates) = &published.updates {
            let _ = updates.send(update);
        }
    }

    fn subscribe(self: Arc<Self>) -> (Update, (Arc<Self>, Option<broadcast::Receiver<Update>>)) {
        let published = self.lock();
        let update = Update {
            sequence: published.sequence,
            change: StateChange::State(published.state.clone()),
        };
        let receiver = published.updates.as_ref().map(broadcast::Sender::subscribe);
        drop(published);
        (update, (self, receiver))
    }
}