macltok

public
Star0
Watch0
120 commits 3 branches 0 tags 0 releases3.0 MiB
mod feed;
#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
mod index;
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
pub(crate) mod store;
mod text;

use std::{
    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},
    sync::Arc,
    time::Duration,
};

use futures_util::{
    StreamExt,
    future::Abortable,
    stream::{AbortHandle, FuturesUnordered, SelectAll, abortable},
};
use matrix_sdk::{
    Client, Room, RoomState, SendOutsideWasm,
    deserialized_responses::TimelineEvent as SdkTimelineEvent,
    event_cache::{RoomEventCacheSubscriber, RoomEventCacheUpdate},
    room::{Messages, MessagesOptions},
};
use matrix_sdk_base::RoomInfoNotableUpdateReasons;
use matrix_sdk_ui::eyeball_im::VectorDiff;
use ruma::{
    OwnedRoomId, OwnedUserId, UInt,
    api::client::filter::RoomEventFilter,
    events::{
        ignored_user_list::IgnoredUserListEventContent, relation::RelationType,
        room::encrypted::OriginalSyncRoomEncryptedEvent,
    },
    serde::Raw,
};
use serde_json::Value;
use sha2::{Digest, Sha256};
use tokio::sync::{Mutex, broadcast, mpsc, oneshot};

use self::{
    feed::{Effect, Observed, Redactions},
    index::{Changes, Doc, Hit, KeywordIndex},
};
use crate::{
    CommandOutput, EngineConfig, Error, ErrorKind, Event, RoomSearchCoverage, SearchBackfill,
    SearchCloseReason, SearchCommand, SearchCoverage, SearchHit, SearchInMemory, SearchMatch,
    SearchPage, SearchScope, SearchSession,
    engine::{MediaService, SearchStore, context, executor::EngineBoxFuture},
    normalise::{MediaRegistration, MediaScope},
    transport::EventSender,
};

const MAX_LIMIT: u16 = 50;
/// a consumer that never closes its searches still holds at most this many
const MAX_SEARCHES: usize = 8;
/// hits read back at once while a page is put together
const PREVIEWS_AT_ONCE: usize = 8;
const DECRYPTS_AT_ONCE: usize = 8;
const COMMIT_AFTER_MS: u64 = 5_000;
const COMMIT_BUFFERED: usize = 2_000;
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
const COMMIT_RETRY_MS: u64 = 60_000;
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
const COMMIT_ATTEMPTS: u32 = 5;
const IDLE_SUBSCRIPTION_MS: u64 = 5 * 60_000;
const BACKFILL_RETRY_MS: u64 = 60_000;
const BACKFILL_RETRY_MAX_MS: u64 = 15 * 60_000;
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
const SCHEMA: u64 = 2;
/// waiting undecryptable events are saved in this many blobs, so one arriving rewrites a sixteenth
const UTD_SHARDS: u8 = 16;
const BACKFILL_TYPES: [&str; 6] = [
    "m.room.message",
    "m.room.encrypted",
    "m.room.redaction",
    "m.sticker",
    "m.poll.start",
    "org.matrix.msc3381.poll.start",
];

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

/// a local index of message text. keyword search over tantivy is the only one today, an
/// embedding index can stand beside it later without the public api noticing
#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
pub(crate) trait MessageIndex: SendOutsideWasm {
    fn get(&self, key: &str) -> Option<Doc>;
    fn by_content(&self, content: &str) -> Option<Doc>;
    /// replaces whatever `doc.key` held
    fn put(&mut self, doc: Doc);
    fn remove(&mut self, key: &str) -> bool;
    /// every provisional doc waiting on `target`, whoever sent its edit
    fn remove_provisional(&mut self, target: &str) -> bool;
    fn buffered(&self) -> usize;
    /// newest first, strictly older than `before`. also hands back the analysed query terms
    fn search(
        &mut self,
        query: &str,
        rooms: Option<&[String]>,
        before: Option<(u64, u64)>,
        limit: usize,
    ) -> Result<(Vec<Hit>, Vec<String>), Error>;
    fn maintain(&mut self) -> Result<(), Error>;
    fn changes(&mut self) -> Result<Changes, Error>;
    fn committed(&mut self, changes: &Changes);
    /// stops tracking what a commit would need, once nothing will be committed again
    fn detach(&mut self);
}

#[derive(Clone, Debug, Default)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct RoomProgress {
    oldest: Option<u64>,
    newest: Option<u64>,
    /// where the next backfill page starts, `None` before the first one
    tail: Option<String>,
    started: bool,
    start_reached: bool,
}

#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
struct Pending {
    room: String,
    session_id: String,
    ts: u64,
    raw: String,
}

#[derive(Clone, Copy)]
struct Budget {
    pages_per_room: u16,
    page_size: u16,
    interval_ms: u64,
    max_pending: usize,
}

impl Budget {
    fn from_config(config: &EngineConfig) -> Self {
        Self {
            pages_per_room: config.search_backfill_pages_per_room,
            page_size: config.search_backfill_page_size.max(1),
            interval_ms: config.search_backfill_interval_ms,
            max_pending: config.max_search_undecrypted,
        }
    }
}

struct RawPage {
    hits: Vec<Hit>,
    terms: Vec<String>,
    more: bool,
    coverage: SearchCoverage,
}

enum Request {
    Search {
        query: String,
        rooms: Option<Vec<String>>,
        before: Option<(u64, u64)>,
        limit: usize,
        coverage: Vec<OwnedRoomId>,
        reply: oneshot::Sender<Result<RawPage, Error>>,
    },
    Coverage(Vec<OwnedRoomId>, oneshot::Sender<SearchCoverage>),
    /// a hit that turned out redacted when its preview was read
    Forget(String),
    Stop(oneshot::Sender<()>),
}

#[derive(Clone)]
pub(crate) struct Handle {
    requests: mpsc::Sender<Request>,
    client: Client,
    media: Arc<MediaService>,
    sessions: Arc<Sessions>,
}

enum Work {
    Backfill {
        room: OwnedRoomId,
        head: bool,
        from: Option<String>,
        result: Result<Messages, Error>,
    },
    /// the room was left while its page was in flight
    BackfillCancelled,
    Decrypted {
        key: String,
        room: OwnedRoomId,
        value: Option<Value>,
    },
    /// the original first, then whatever edits of it the cache or server still has
    Refetched {
        room: OwnedRoomId,
        values: Vec<Value>,
    },
}

enum RoomUpdate {
    Events(Vec<SdkTimelineEvent>),
    Lagged,
    Closed,
}

struct Failure {
    message: String,
    retry_at_ms: u64,
    delay_ms: u64,
}

enum Persistence {
    Store(SearchStore),
    Memory(SearchInMemory),
}

pub(crate) struct Service {
    client: Client,
    index: Box<dyn MessageIndex>,
    persistence: Persistence,
    #[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
    commit_failures: u32,
    /// storage keys a load found nothing pointing at, dropped with the next commit
    #[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
    orphans: Vec<String>,
    rooms: BTreeMap<String, RoomProgress>,
    /// each room's newest indexed message when this engine started, what a head pass looks for
    baselines: HashMap<String, u64>,
    pending: BTreeMap<String, Pending>,
    dirty_shards: BTreeSet<u8>,
    decrypt_queue: VecDeque<String>,
    decrypting: usize,
    redactions: Redactions,
    dirty_since: Option<u64>,
    commit_after: u64,
    budget: Budget,
    requests: mpsc::Receiver<Request>,
    subscriptions: HashMap<OwnedRoomId, (AbortHandle, u64)>,
    updates: SelectAll<BoxStream<(OwnedRoomId, RoomUpdate)>>,
    work: FuturesUnordered<EngineBoxFuture<'static, Work>>,
    pages: HashMap<OwnedRoomId, u16>,
    heads: VecDeque<OwnedRoomId>,
    headed: HashSet<OwnedRoomId>,
    head_tokens: HashMap<OwnedRoomId, String>,
    backfilling: Option<(OwnedRoomId, AbortHandle)>,
    next_backfill: u64,
    failures: HashMap<OwnedRoomId, Failure>,
    turn: usize,
}

fn now() -> u64 {
    ruma::MilliSecondsSinceUnixEpoch::now().get().into()
}

fn value_of(event: &SdkTimelineEvent) -> Option<Value> {
    serde_json::from_str(event.raw().json().get()).ok()
}

fn encode_cursor((ts, tie): (u64, u64)) -> String {
    format!("{ts:x}.{tie:x}")
}

fn decode_cursor(cursor: &str) -> Option<(u64, u64)> {
    let (ts, tie) = cursor.split_once('.')?;
    Some((
        u64::from_str_radix(ts, 16).ok()?,
        u64::from_str_radix(tie, 16).ok()?,
    ))
}

fn shard(key: &str) -> u8 {
    Sha256::digest(key.as_bytes())[0] % UTD_SHARDS
}

/// a 400 to a page with a saved token means the server no longer takes the token
fn token_rejected(error: &Error) -> bool {
    error
        .server
        .as_ref()
        .is_some_and(|server| server.status == 400 || server.errcode == "M_INVALID_PARAM")
}

impl Handle {
    pub(crate) async fn execute(&self, command: SearchCommand) -> Result<CommandOutput, Error> {
        match command {
            SearchCommand::Search {
                query,
                scope,
                limit,
            } => {
                let media_scope = self.media.context_scope();
                let prefix = media_prefix(&media_scope);
                let (page, media) = self.search(&query, &scope, None, limit, &prefix).await?;
                let search_id = self
                    .sessions
                    .open(&self.media, query, scope, media_scope, media)
                    .await;
                Ok(CommandOutput::SearchSession(SearchSession {
                    search_id,
                    page,
                }))
            }
            SearchCommand::More {
                search_id,
                cursor,
                limit,
            } => {
                let (query, scope, prefix) = self.sessions.get(&search_id).await?;
                let (page, media) = self
                    .search(&query, &scope, Some(cursor.as_str()), limit, &prefix)
                    .await?;
                self.sessions
                    .register(&self.media, &search_id, media)
                    .await?;
                Ok(CommandOutput::SearchPage(page))
            }
            SearchCommand::Close { search_id } => {
                self.sessions.close(&self.media, &search_id).await?;
                Ok(CommandOutput::Done)
            }
            SearchCommand::Coverage(scope) => {
                let (reply, response) = oneshot::channel();
                self.requests
                    .send(Request::Coverage(self.scope_rooms(&scope), reply))
                    .await
                    .map_err(|_| Error::interrupted())?;
                Ok(CommandOutput::SearchCoverage(
                    response.await.map_err(|_| Error::interrupted())?,
                ))
            }
        }
    }

    /// the media service is cleared beside this, so nothing is released one by one
    pub(crate) async fn stop(&self) {
        let (reply, response) = oneshot::channel();
        if self.requests.send(Request::Stop(reply)).await.is_ok() {
            let _ = response.await;
        }
        self.sessions.open.lock().await.clear();
    }

    fn scope_rooms(&self, scope: &SearchScope) -> Vec<OwnedRoomId> {
        match scope {
            SearchScope::Room(room_id) => vec![room_id.clone()],
            SearchScope::Rooms(room_ids) => room_ids.clone(),
            SearchScope::All => self
                .client
                .joined_rooms()
                .into_iter()
                .map(|room| room.room_id().to_owned())
                .collect(),
        }
    }

    async fn search(
        &self,
        query: &str,
        scope: &SearchScope,
        cursor: Option<&str>,
        limit: u16,
        media_prefix: &str,
    ) -> Result<(SearchPage, Vec<MediaRegistration>), Error> {
        if limit == 0 || limit > MAX_LIMIT {
            return Err(Error::new(
                ErrorKind::Invalid,
                "search limit must be between 1 and 50",
            ));
        }
        let before = match cursor {
            Some(cursor) => Some(
                decode_cursor(cursor)
                    .ok_or_else(|| Error::new(ErrorKind::Invalid, "search cursor is invalid"))?,
            ),
            None => None,
        };
        let rooms = self.scope_rooms(scope);
        let filter = Some(rooms.iter().map(ToString::to_string).collect());
        let (reply, response) = oneshot::channel();
        self.requests
            .send(Request::Search {
                query: query.to_owned(),
                rooms: filter,
                before,
                limit: usize::from(limit),
                coverage: rooms,
                reply,
            })
            .await
            .map_err(|_| Error::interrupted())?;
        let raw = response.await.map_err(|_| Error::interrupted())??;
        let next_cursor = raw
            .more
            .then(|| raw.hits.last().map(|hit| encode_cursor((hit.ts, hit.tie))))
            .flatten();
        let ignored = self.ignored_users().await;
        let current_user_id = self.client.user_id().map(ToOwned::to_owned);
        let terms = &raw.terms;
        let previews: Vec<_> = futures_util::stream::iter(
            raw.hits
                .into_iter()
                .filter(|hit| !ignored.iter().any(|user| user.as_str() == hit.sender)),
        )
        .map(|hit| self.preview(hit, current_user_id.clone(), terms, media_prefix))
        .buffered(PREVIEWS_AT_ONCE)
        .collect()
        .await;
        let mut hits = Vec::new();
        let mut media = Vec::new();
        for preview in previews {
            match preview {
                Preview::Shown(hit, registrations) => {
                    hits.push(hit);
                    media.extend(registrations);
                }
                Preview::Redacted(key) => {
                    let _ = self.requests.send(Request::Forget(key)).await;
                }
                Preview::Unavailable => {}
            }
        }
        let page = SearchPage {
            hits,
            next_cursor,
            coverage: raw.coverage,
        };
        Ok((page, media))
    }

    async fn ignored_users(&self) -> Vec<OwnedUserId> {
        self.client
            .account()
            .account_data::<IgnoredUserListEventContent>()
            .await
            .ok()
            .flatten()
            .and_then(|raw| raw.deserialize().ok())
            .map(|content| content.ignored_users.into_keys().collect())
            .unwrap_or_default()
    }

    async fn preview(
        &self,
        hit: Hit,
        current_user_id: Option<OwnedUserId>,
        terms: &[String],
        media_prefix: &str,
    ) -> Preview {
        let (Ok(room_id), Ok(event_id)) = (
            OwnedRoomId::try_from(hit.room.as_str()),
            ruma::OwnedEventId::try_from(hit.key.as_str()),
        ) else {
            return Preview::Unavailable;
        };
        let Some(room) = self.client.get_room(&room_id) else {
            return Preview::Unavailable;
        };
        let Some((mut original, highlight)) = load(&room, &hit.key).await else {
            return Preview::Unavailable;
        };
        if original.pointer("/unsigned/redacted_because").is_some() {
            return Preview::Redacted(hit.key);
        }
        // the index already settled which edit is newest, and a bundled one can be older
        if let Some(relations) = original
            .pointer_mut("/unsigned/m.relations")
            .and_then(Value::as_object_mut)
        {
            relations.remove("m.replace");
        }
        let mut values = vec![original];
        if hit.content != hit.key
            && let Some((edit, _)) = load(&room, &hit.content).await
        {
            values.push(edit);
        }
        let Some((event, registrations)) = context::preview(
            &room,
            current_user_id.as_deref(),
            &values,
            highlight,
            media_prefix,
        )
        .await
        else {
            return Preview::Unavailable;
        };
        let text = match &event.content {
            crate::TimelineContent::Message(message) => Some(message.body.as_str()),
            crate::TimelineContent::Sticker { body, .. } => Some(body.as_str()),
            crate::TimelineContent::Poll { question } => Some(question.as_str()),
            _ => None,
        };
        let matches = text
            .map(|text| text::match_ranges(text, terms))
            .unwrap_or_default()
            .into_iter()
            .filter_map(|(start, end)| {
                Some(SearchMatch {
                    start: u32::try_from(start).ok()?,
                    end: u32::try_from(end).ok()?,
                })
            })
            .collect();
        Preview::Shown(
            SearchHit {
                room_id,
                event_id,
                timestamp_ms: event.timestamp_ms,
                event,
                matches,
            },
            registrations,
        )
    }
}

enum Preview {
    Shown(SearchHit, Vec<MediaRegistration>),
    Redacted(String),
    Unavailable,
}

/// open searches, oldest first, each owning the media of the hits it handed out
struct Sessions {
    open: Mutex<VecDeque<Session>>,
    events: EventSender,
}

/// media ids carry the search, so two owners of one event never share an entry
fn media_prefix(scope: &MediaScope) -> String {
    format!("{}#", scope.nonce())
}

struct Session {
    id: String,
    query: String,
    scope: SearchScope,
    media: MediaScope,
}

impl Sessions {
    async fn open(
        &self,
        media: &MediaService,
        query: String,
        scope: SearchScope,
        media_scope: MediaScope,
        registrations: Vec<MediaRegistration>,
    ) -> String {
        let id = format!("search-{}", media_scope.nonce());
        let mut evicted = Vec::new();
        {
            let mut open = self.open.lock().await;
            media.register(&media_scope, registrations).await;
            open.push_back(Session {
                id: id.clone(),
                query,
                scope,
                media: media_scope,
            });
            while open.len() > MAX_SEARCHES {
                if let Some(oldest) = open.pop_front() {
                    media.release(&oldest.media).await;
                    evicted.push(oldest.id);
                }
            }
        }
        for search_id in evicted {
            let event = Event::SearchClosed {
                search_id,
                reason: SearchCloseReason::Evicted,
            };
            if self.events.send(event).await.is_err() {
                break;
            }
        }
        id
    }

    /// the query, scope and media prefix the search's pages use
    async fn get(&self, id: &str) -> Result<(String, SearchScope, String), Error> {
        let open = self.open.lock().await;
        let session = open
            .iter()
            .find(|session| session.id == id)
            .ok_or_else(not_open)?;
        Ok((
            session.query.clone(),
            session.scope.clone(),
            media_prefix(&session.media),
        ))
    }

    /// under the lock, so a page that lands after its close registers nothing
    async fn register(
        &self,
        media: &MediaService,
        id: &str,
        registrations: Vec<MediaRegistration>,
    ) -> Result<(), Error> {
        let open = self.open.lock().await;
        let session = open
            .iter()
            .find(|session| session.id == id)
            .ok_or_else(not_open)?;
        media.register(&session.media, registrations).await;
        Ok(())
    }

    async fn close(&self, media: &MediaService, id: &str) -> Result<(), Error> {
        let mut open = self.open.lock().await;
        let index = open
            .iter()
            .position(|session| session.id == id)
            .ok_or_else(not_open)?;
        if let Some(session) = open.remove(index) {
            media.release(&session.media).await;
        }
        Ok(())
    }
}

fn not_open() -> Error {
    Error::new(
        ErrorKind::SearchNotOpen,
        "search is closed or was never opened",
    )
}

async fn load(room: &Room, event_id: &str) -> Option<(Value, bool)> {
    let event_id = ruma::EventId::parse(event_id).ok()?;
    let event = room
        .load_or_fetch_event(&event_id, None)
        .await
        .inspect_err(|error| tracing::debug!(%event_id, %error, "loading a search hit failed"))
        .ok()?;
    let highlight = event
        .push_actions()
        .is_some_and(|actions| actions.iter().any(ruma::push::Action::is_highlight));
    Some((value_of(&event)?, highlight))
}

/// the original and every edit of it still around, so a redacted latest edit falls back to
/// the one before it rather than to the original
async fn load_with_edits(room: &Room, event_id: &str) -> Vec<Value> {
    let Ok(event_id) = ruma::EventId::parse(event_id) else {
        return Vec::new();
    };
    match room
        .load_or_fetch_event_with_relations(&event_id, Some(vec![RelationType::Replacement]), None)
        .await
    {
        Ok((original, edits)) => std::iter::once(original)
            .chain(edits)
            .filter_map(|event| value_of(&event))
            .collect(),
        Err(error) => {
            tracing::debug!(%event_id, %error, "reading an original again failed");
            Vec::new()
        }
    }
}

impl Service {
    pub(crate) fn new(
        client: Client,
        media: Arc<MediaService>,
        store: Option<Result<SearchStore, Error>>,
        events: EventSender,
        config: &EngineConfig,
    ) -> (Handle, Self) {
        let (sender, requests) = mpsc::channel(config.command_capacity.max(1));
        let persistence = match store {
            Some(Ok(store)) => Persistence::Store(store),
            Some(Err(error)) => {
                tracing::error!(%error, "search store unavailable, the index lives in memory");
                Persistence::Memory(SearchInMemory::OpenFailed)
            }
            None => Persistence::Memory(SearchInMemory::NoStore),
        };
        let service = Self {
            client: client.clone(),
            index: Box::new(KeywordIndex::new(matches!(
                persistence,
                Persistence::Store(_)
            ))),
            persistence,
            commit_failures: 0,
            orphans: Vec::new(),
            rooms: BTreeMap::new(),
            baselines: HashMap::new(),
            pending: BTreeMap::new(),
            dirty_shards: BTreeSet::new(),
            decrypt_queue: VecDeque::new(),
            decrypting: 0,
            redactions: Redactions::default(),
            dirty_since: None,
            commit_after: 0,
            budget: Budget::from_config(config),
            requests,
            subscriptions: HashMap::new(),
            updates: SelectAll::new(),
            work: FuturesUnordered::new(),
            pages: HashMap::new(),
            heads: VecDeque::new(),
            headed: HashSet::new(),
            head_tokens: HashMap::new(),
            backfilling: None,
            next_backfill: 0,
            failures: HashMap::new(),
            turn: 0,
        };
        (
            Handle {
                requests: sender,
                client,
                media,
                sessions: Arc::new(Sessions {
                    open: Mutex::default(),
                    events,
                }),
            },
            service,
        )
    }

    fn in_memory(&self) -> Option<SearchInMemory> {
        match self.persistence {
            Persistence::Store(_) => None,
            Persistence::Memory(reason) => Some(reason),
        }
    }

    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
    fn fall_back(&mut self, reason: SearchInMemory) {
        self.persistence = Persistence::Memory(reason);
        self.index.detach();
        self.orphans.clear();
    }

    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
    async fn load(&mut self) {
        let Persistence::Store(store) = &self.persistence else {
            return;
        };
        let mut blobs = match store.load().await {
            Ok(blobs) => blobs,
            Err(error) => {
                tracing::error!(%error, "reading the search store failed, the index lives in memory");
                self.fall_back(SearchInMemory::ReadFailed);
                return;
            }
        };
        if blobs.is_empty() {
            return;
        }
        let restored = (|| {
            let meta: Value = serde_json::from_slice(&store.take(&mut blobs, "meta")?).ok()?;
            if meta["schema"].as_u64() != Some(SCHEMA) {
                return None;
            }
            let rooms: BTreeMap<String, RoomProgress> =
                serde_json::from_value(meta["rooms"].clone()).ok()?;
            let mut pending = BTreeMap::new();
            for shard in 0..UTD_SHARDS {
                if let Some(bytes) = store.take(&mut blobs, &format!("utd/{shard}")) {
                    let part: BTreeMap<String, Pending> = serde_json::from_slice(&bytes).ok()?;
                    pending.extend(part);
                }
            }
            let index = KeywordIndex::load(&meta["index"], |name| store.take(&mut blobs, name))
                .inspect_err(|error| tracing::warn!(%error, "search segments are unreadable"))
                .ok()?;
            Some((index, rooms, pending))
        })();
        match restored {
            Some((index, rooms, pending)) => {
                self.index = Box::new(index);
                self.baselines = rooms
                    .iter()
                    .filter_map(|(room, progress)| Some((room.clone(), progress.newest?)))
                    .collect();
                self.rooms = rooms;
                self.pending = pending;
                self.orphans = blobs.leftover();
                if !self.orphans.is_empty() {
                    self.touch();
                }
            }
            None => {
                tracing::warn!("resetting an unreadable search index");
                if let Err(error) = store.reset().await {
                    tracing::error!(%error, "search store unavailable, the index lives in memory");
                    self.fall_back(SearchInMemory::ReadFailed);
                }
            }
        }
    }

    #[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
    async fn load(&mut self) {}

    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
    async fn commit(&mut self) {
        if !matches!(self.persistence, Persistence::Store(_)) {
            if let Err(error) = self.index.maintain() {
                tracing::warn!(%error, "maintaining the search index failed");
            }
            self.dirty_since = None;
            self.dirty_shards.clear();
            return;
        }
        let mut changes = match self.index.changes() {
            Ok(changes) => changes,
            Err(error) => {
                tracing::warn!(%error, "preparing a search commit failed");
                self.commit_failed();
                return;
            }
        };
        let mut puts = std::mem::take(&mut changes.puts);
        let mut deletes = changes.deletes.clone();
        let meta = serde_json::json!({
            "schema": SCHEMA,
            "index": changes.manifest,
            "rooms": self.rooms,
        });
        puts.push(("meta".to_owned(), meta.to_string().into_bytes()));
        for shard in &self.dirty_shards {
            let part: BTreeMap<&String, &Pending> = self
                .pending
                .iter()
                .filter(|(key, _)| self::shard(key) == *shard)
                .collect();
            let name = format!("utd/{shard}");
            if part.is_empty() {
                deletes.push(name);
            } else {
                puts.push((name, serde_json::to_vec(&part).unwrap_or_default()));
            }
        }
        let Persistence::Store(store) = &self.persistence else {
            return;
        };
        match store.commit(puts, deletes, self.orphans.clone()).await {
            Ok(()) => {
                self.index.committed(&changes);
                self.dirty_shards.clear();
                self.orphans.clear();
                self.dirty_since = None;
                self.commit_failures = 0;
            }
            Err(error) => {
                tracing::warn!(%error, "saving the search index failed");
                self.commit_failed();
            }
        }
    }

    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
    fn commit_failed(&mut self) {
        self.commit_failures += 1;
        if self.commit_failures >= COMMIT_ATTEMPTS {
            tracing::error!("search commits keep failing, the index lives in memory from here");
            self.fall_back(SearchInMemory::CommitFailed);
            self.dirty_since = None;
            self.dirty_shards.clear();
        } else {
            self.commit_after = now() + COMMIT_RETRY_MS;
        }
    }

    #[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
    async fn commit(&mut self) {
        if let Err(error) = self.index.maintain() {
            tracing::warn!(%error, "maintaining the search index failed");
        }
        self.dirty_since = None;
        self.dirty_shards.clear();
    }

    fn touch(&mut self) {
        self.dirty_since.get_or_insert_with(now);
    }

    fn unpend(&mut self, key: &str) {
        if self.pending.remove(key).is_some() {
            self.dirty_shards.insert(shard(key));
            self.touch();
        }
    }

    fn observe(&mut self, room: &OwnedRoomId, value: &Value) {
        let observed = feed::observe(value);
        let key = match &observed {
            Observed::Message { key, .. }
            | Observed::Redacted { key }
            | Observed::Encrypted { key, .. } => Some(key.clone()),
            Observed::Edit(edit) => Some(edit.id.clone()),
            Observed::Redaction { .. } | Observed::Other => None,
        };
        if let Observed::Encrypted {
            key,
            session_id,
            ts,
        } = &observed
        {
            if self.index.get(key).is_none() && !self.pending.contains_key(key) {
                self.pending.insert(
                    key.clone(),
                    Pending {
                        room: room.to_string(),
                        session_id: session_id.clone(),
                        ts: *ts,
                        raw: value.to_string(),
                    },
                );
                self.dirty_shards.insert(shard(key));
                while self.pending.len() > self.budget.max_pending {
                    let Some(oldest) = self
                        .pending
                        .iter()
                        .min_by_key(|(_, pending)| pending.ts)
                        .map(|(key, _)| key.clone())
                    else {
                        break;
                    };
                    self.unpend(&oldest);
                }
                self.touch();
            }
            return;
        }
        if let Some(key) = key {
            self.unpend(&key);
        }
        match feed::apply(
            self.index.as_mut(),
            &mut self.redactions,
            room.as_str(),
            observed,
        ) {
            Effect::Indexed(ts) => {
                let progress = self.rooms.entry(room.to_string()).or_default();
                progress.oldest = Some(progress.oldest.map_or(ts, |oldest| oldest.min(ts)));
                progress.newest = Some(progress.newest.map_or(ts, |newest| newest.max(ts)));
                self.touch();
            }
            Effect::Removed | Effect::Held => self.touch(),
            Effect::Refetch(key) => {
                self.touch();
                let room_id = room.clone();
                let client = self.client.clone();
                self.work.push(Box::pin(async move {
                    let values = match client.get_room(&room_id) {
                        Some(room) => load_with_edits(&room, &key).await,
                        None => Vec::new(),
                    };
                    Work::Refetched {
                        room: room_id,
                        values,
                    }
                }));
            }
            Effect::Nothing => {}
        }
    }

    fn observe_all(&mut self, room: &OwnedRoomId, events: &[SdkTimelineEvent]) {
        for event in events {
            if let Some(value) = value_of(event) {
                self.observe(room, &value);
            }
        }
    }

    async fn subscribe(&mut self, room_id: OwnedRoomId) {
        if let Some((_, seen)) = self.subscriptions.get_mut(&room_id) {
            *seen = now();
            return;
        }
        let Some(room) = self.client.get_room(&room_id) else {
            return;
        };
        if room.state() != RoomState::Joined {
            return;
        }
        let Ok((cache, drop_handles)) = room.event_cache().await else {
            return;
        };
        let Ok((initial, subscriber)) = cache.subscribe().await else {
            return;
        };
        self.observe_all(&room_id, &initial);
        let stream = updates(room_id.clone(), subscriber, drop_handles);
        let (stream, abort) = abortable(stream);
        self.updates.push(Box::pin(stream));
        self.subscriptions.insert(room_id.clone(), (abort, now()));
        let started = self
            .rooms
            .get(room_id.as_str())
            .is_some_and(|progress| progress.started);
        if started
            && self.baselines.contains_key(room_id.as_str())
            && self.headed.insert(room_id.clone())
        {
            self.heads.push_back(room_id);
        }
    }

    /// the generic channel holds 128 updates, and a missed one could be a room we never
    /// follow, so a lag follows every joined room
    async fn generic(&mut self, update: Result<OwnedRoomId, broadcast::error::RecvError>) -> bool {
        match update {
            Ok(room_id) => self.subscribe(room_id).await,
            Err(broadcast::error::RecvError::Lagged(_)) => {
                for room in self.client.joined_rooms() {
                    self.subscribe(room.room_id().to_owned()).await;
                }
            }
            Err(broadcast::error::RecvError::Closed) => return false,
        }
        true
    }

    async fn room_update(&mut self, room_id: OwnedRoomId, update: RoomUpdate) {
        match update {
            RoomUpdate::Events(events) => {
                if let Some((_, seen)) = self.subscriptions.get_mut(&room_id) {
                    *seen = now();
                }
                self.observe_all(&room_id, &events);
            }
            RoomUpdate::Lagged => {
                if let Some((abort, _)) = self.subscriptions.remove(&room_id) {
                    abort.abort();
                }
                self.subscribe(room_id).await;
            }
            RoomUpdate::Closed => {
                self.subscriptions.remove(&room_id);
            }
        }
    }

    fn left(&mut self, room_id: &OwnedRoomId) {
        if let Some((abort, _)) = self.subscriptions.remove(room_id) {
            abort.abort();
        }
        self.heads.retain(|head| head != room_id);
        self.head_tokens.remove(room_id);
        if let Some((backfilling, abort)) = &self.backfilling
            && backfilling == room_id
        {
            abort.abort();
        }
    }

    fn drop_idle(&mut self) {
        let cutoff = now().saturating_sub(IDLE_SUBSCRIPTION_MS);
        self.subscriptions.retain(|_, (abort, seen)| {
            let keep = *seen >= cutoff;
            if !keep {
                abort.abort();
            }
            keep
        });
    }

    fn coverage(&self, rooms: &[OwnedRoomId]) -> SearchCoverage {
        let scoped: HashSet<&str> = rooms.iter().map(|room| room.as_str()).collect();
        let undecrypted = self
            .pending
            .values()
            .filter(|pending| scoped.contains(pending.room.as_str()))
            .count();
        let rooms = rooms
            .iter()
            .map(|room_id| {
                let progress = self
                    .rooms
                    .get(room_id.as_str())
                    .cloned()
                    .unwrap_or_default();
                let backfill = if progress.start_reached {
                    SearchBackfill::Complete
                } else if self.budget.pages_per_room == 0 {
                    SearchBackfill::Disabled
                } else if self
                    .backfilling
                    .as_ref()
                    .is_some_and(|(backfilling, _)| backfilling == room_id)
                {
                    SearchBackfill::Running
                } else if let Some(failure) = self.failures.get(room_id) {
                    SearchBackfill::Failed {
                        message: failure.message.clone(),
                        retry_at_ms: failure.retry_at_ms,
                    }
                } else if self.pages.get(room_id).copied().unwrap_or(0)
                    >= self.budget.pages_per_room
                {
                    SearchBackfill::Paused
                } else {
                    SearchBackfill::Waiting
                };
                RoomSearchCoverage {
                    room_id: room_id.clone(),
                    oldest_indexed_ms: progress.oldest,
                    newest_indexed_ms: progress.newest,
                    start_reached: progress.start_reached,
                    backfill,
                }
            })
            .collect();
        let in_memory = self.in_memory();
        SearchCoverage {
            rooms,
            persistent: in_memory.is_none(),
            in_memory,
            undecrypted,
        }
    }

    fn spent(&self, room_id: &OwnedRoomId) -> bool {
        self.pages.get(room_id).copied().unwrap_or(0) >= self.budget.pages_per_room
    }

    fn blocked(&self, room_id: &OwnedRoomId, at: u64) -> bool {
        self.failures
            .get(room_id)
            .is_some_and(|failure| failure.retry_at_ms > at)
    }

    /// head passes first, so what arrived while we were away is found before anything older
    fn schedule_backfill(&mut self) {
        let at = now();
        if self.backfilling.is_some() || self.budget.pages_per_room == 0 || at < self.next_backfill
        {
            return;
        }
        let mut chosen = None;
        for _ in 0..self.heads.len() {
            let Some(room_id) = self.heads.pop_front() else {
                break;
            };
            if self.spent(&room_id) {
                self.head_tokens.remove(&room_id);
                continue;
            }
            if self.blocked(&room_id, at) {
                self.heads.push_back(room_id);
                continue;
            }
            chosen = Some((room_id, true));
            break;
        }
        if chosen.is_none() {
            let mut joined: Vec<_> = self
                .client
                .joined_rooms()
                .into_iter()
                .map(|room| room.room_id().to_owned())
                .collect();
            joined.sort();
            let count = joined.len();
            for offset in 0..count {
                let room_id = &joined[(self.turn + offset) % count];
                let done = self
                    .rooms
                    .get(room_id.as_str())
                    .is_some_and(|progress| progress.start_reached);
                if !done && !self.spent(room_id) && !self.blocked(room_id, at) {
                    self.turn = (self.turn + offset + 1) % count;
                    chosen = Some((room_id.clone(), false));
                    break;
                }
            }
        }
        let Some((room_id, head)) = chosen else {
            return;
        };
        let Some(room) = self.client.get_room(&room_id) else {
            return;
        };
        let from = if head {
            self.head_tokens.get(&room_id).cloned()
        } else {
            self.rooms
                .get(room_id.as_str())
                .and_then(|progress| progress.tail.clone())
        };
        *self.pages.entry(room_id.clone()).or_default() += 1;
        let mut options = MessagesOptions::backward();
        options.from = from.clone();
        options.limit = UInt::from(self.budget.page_size);
        let mut filter = RoomEventFilter::default();
        filter.types = Some(
            BACKFILL_TYPES
                .iter()
                .map(|kind| (*kind).to_owned())
                .collect(),
        );
        options.filter = filter;
        let (abort, registration) = AbortHandle::new_pair();
        self.backfilling = Some((room_id.clone(), abort));
        let page = Abortable::new(
            async move {
                room.messages(options)
                    .await
                    .map_err(|error| Error::from_sdk("backfilling search history failed", error))
            },
            registration,
        );
        self.work.push(Box::pin(async move {
            match page.await {
                Ok(result) => Work::Backfill {
                    room: room_id,
                    head,
                    from,
                    result,
                },
                Err(_) => Work::BackfillCancelled,
            }
        }));
    }

    fn fail_backfill(&mut self, room_id: OwnedRoomId, message: String, delay_ms: u64) {
        self.failures.insert(
            room_id,
            Failure {
                message,
                retry_at_ms: now() + delay_ms,
                delay_ms,
            },
        );
    }

    fn backfilled(
        &mut self,
        room_id: OwnedRoomId,
        head: bool,
        from: Option<String>,
        result: Result<Messages, Error>,
    ) {
        self.backfilling = None;
        self.next_backfill = now() + self.budget.interval_ms;
        let messages = match result {
            Ok(messages) => messages,
            Err(error) if from.is_some() && token_rejected(&error) => {
                if head {
                    self.head_tokens.remove(&room_id);
                } else {
                    let progress = self.rooms.entry(room_id.to_string()).or_default();
                    progress.tail = None;
                    progress.started = false;
                    self.touch();
                }
                let message = format!("{}, starting again from the newest", error.message);
                self.fail_backfill(room_id, message, self.budget.interval_ms);
                return;
            }
            Err(error) => {
                let delay_ms = self
                    .failures
                    .get(&room_id)
                    .map_or(BACKFILL_RETRY_MS, |failure| {
                        (failure.delay_ms * 2).min(BACKFILL_RETRY_MAX_MS)
                    });
                self.fail_backfill(room_id, error.message.clone(), delay_ms);
                return;
            }
        };
        self.failures.remove(&room_id);
        let baseline = self.baselines.get(room_id.as_str()).copied();
        let mut overlapped = false;
        for event in &messages.chunk {
            let Some(value) = value_of(event) else {
                continue;
            };
            if head && head_overlap(self.index.as_ref(), baseline, &value) {
                overlapped = true;
            }
            self.observe(&room_id, &value);
        }
        let ended = messages.end.is_none() || (messages.chunk.is_empty() && messages.end == from);
        if head {
            match messages.end {
                Some(end) if !overlapped && !ended => {
                    self.head_tokens.insert(room_id.clone(), end);
                    self.heads.push_front(room_id);
                }
                _ => {
                    self.head_tokens.remove(&room_id);
                }
            }
            return;
        }
        let progress = self.rooms.entry(room_id.to_string()).or_default();
        progress.started = true;
        progress.tail = messages.end;
        if ended {
            progress.start_reached = true;
            progress.tail = None;
        }
        self.touch();
    }

    fn retry_pending(&mut self, room_id: Option<&OwnedRoomId>, session_id: Option<&str>) {
        let queued: HashSet<String> = self.decrypt_queue.iter().cloned().collect();
        let due: Vec<_> = self
            .pending
            .iter()
            .filter(|(key, pending)| {
                !queued.contains(*key)
                    && room_id.is_none_or(|room_id| pending.room == room_id.as_str())
                    && session_id.is_none_or(|session_id| pending.session_id == session_id)
            })
            .map(|(key, _)| key.clone())
            .collect();
        self.decrypt_queue.extend(due);
        self.pump_decrypts();
    }

    fn pump_decrypts(&mut self) {
        while self.decrypting < DECRYPTS_AT_ONCE {
            let Some(key) = self.decrypt_queue.pop_front() else {
                return;
            };
            let Some(pending) = self.pending.get(&key).cloned() else {
                continue;
            };
            let Ok(room_id) = OwnedRoomId::try_from(pending.room.as_str()) else {
                continue;
            };
            let Some(room) = self.client.get_room(&room_id) else {
                continue;
            };
            self.decrypting += 1;
            self.work.push(Box::pin(async move {
                let value = async {
                    let raw: Raw<OriginalSyncRoomEncryptedEvent> =
                        Raw::from_json_string(pending.raw).ok()?;
                    let event = room.decrypt_event(&raw, None).await.ok()?;
                    value_of(&event).filter(|value| value["type"] != "m.room.encrypted")
                }
                .await;
                Work::Decrypted {
                    key,
                    room: room_id,
                    value,
                }
            }));
        }
    }

    fn complete(&mut self, work: Work) {
        match work {
            Work::Backfill {
                room,
                head,
                from,
                result,
            } => self.backfilled(room, head, from, result),
            Work::BackfillCancelled => {
                self.backfilling = None;
                self.next_backfill = now() + self.budget.interval_ms;
            }
            Work::Decrypted { key, room, value } => {
                self.decrypting -= 1;
                if let Some(value) = value {
                    self.unpend(&key);
                    self.observe(&room, &value);
                }
                self.pump_decrypts();
            }
            Work::Refetched { room, values } => {
                for value in values {
                    self.observe(&room, &value);
                }
            }
        }
    }

    fn search(
        &mut self,
        query: &str,
        rooms: Option<&[String]>,
        before: Option<(u64, u64)>,
        limit: usize,
        coverage: &[OwnedRoomId],
    ) -> Result<RawPage, Error> {
        let (mut hits, terms) = self.index.search(query, rooms, before, limit + 1)?;
        let more = hits.len() > limit;
        hits.truncate(limit);
        Ok(RawPage {
            hits,
            terms,
            more,
            coverage: self.coverage(coverage),
        })
    }

    pub(crate) async fn run(mut self) {
        let client = self.client.clone();
        let mut generic = client.event_cache().subscribe_to_room_generic_updates();
        let mut room_info = client.room_info_notable_update_receiver();
        self.load().await;
        let mut keys = client
            .encryption()
            .room_keys_received_stream()
            .await
            .map(Box::pin);
        self.retry_pending(None, None);
        let mut tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
        loop {
            self.schedule_backfill();
            tokio::select! {
                biased;
                request = self.requests.recv() => match request {
                    Some(Request::Stop(reply)) => {
                        if self.dirty_since.is_some() || self.index.buffered() > 0 {
                            self.commit().await;
                        }
                        for (abort, _) in self.subscriptions.values() {
                            abort.abort();
                        }
                        let _ = reply.send(());
                        return;
                    }
                    Some(Request::Search { query, rooms, before, limit, coverage, reply }) => {
                        let _ = reply.send(self.search(&query, rooms.as_deref(), before, limit, &coverage));
                    }
                    Some(Request::Coverage(rooms, reply)) => {
                        let _ = reply.send(self.coverage(&rooms));
                    }
                    Some(Request::Forget(key)) => {
                        let observed = Observed::Redacted { key };
                        if feed::apply(self.index.as_mut(), &mut self.redactions, "", observed) == Effect::Removed {
                            self.touch();
                        }
                    }
                    None => return,
                },
                Some(work) = self.work.next(), if !self.work.is_empty() => self.complete(work),
                Some((room_id, update)) = self.updates.next(), if !self.updates.is_empty() => {
                    self.room_update(room_id, update).await;
                }
                update = generic.recv() => {
                    if !self.generic(update.map(|update| update.room_id)).await {
                        return;
                    }
                }
                Some(update) = async { keys.as_mut()?.next().await }, if keys.is_some() => match update {
                    Ok(infos) => {
                        for info in infos {
                            self.retry_pending(Some(&info.room_id), Some(&info.session_id));
                        }
                    }
                    Err(_) => self.retry_pending(None, None),
                },
                update = room_info.recv() => match update {
                    Ok(update) if update.reasons.contains(RoomInfoNotableUpdateReasons::MEMBERSHIP) => {
                        let left = self.client.get_room(&update.room_id).is_none_or(|room| room.state() != RoomState::Joined);
                        if left {
                            self.left(&update.room_id);
                        }
                    }
                    Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {}
                    Err(broadcast::error::RecvError::Closed) => return,
                },
                () = &mut tick => {
                    tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
                    self.drop_idle();
                    let due = self.dirty_since.is_some_and(|since| now() >= since + COMMIT_AFTER_MS)
                        || self.index.buffered() >= COMMIT_BUFFERED;
                    if due && now() >= self.commit_after {
                        self.commit().await;
                    }
                }
            }
        }
    }
}

/// a head pass stops at what this engine found already indexed, and live events indexed
/// since it started don't count, or the first page would always stop it
fn head_overlap(index: &dyn MessageIndex, baseline: Option<u64>, value: &Value) -> bool {
    let (Some(baseline), Some(key), Some(ts)) = (
        baseline,
        value.get("event_id").and_then(Value::as_str),
        value.get("origin_server_ts").and_then(Value::as_u64),
    ) else {
        return false;
    };
    ts <= baseline && index.get(key).is_some_and(|doc| !doc.provisional)
}

fn updates(
    room_id: OwnedRoomId,
    subscriber: RoomEventCacheSubscriber,
    drop_handles: Arc<matrix_sdk::event_cache::EventCacheDropHandles>,
) -> BoxStream<(OwnedRoomId, RoomUpdate)> {
    Box::pin(futures_util::stream::unfold(
        Some((subscriber, drop_handles)),
        move |state| {
            let room_id = room_id.clone();
            async move {
                let (mut subscriber, drop_handles) = state?;
                loop {
                    match subscriber.recv().await {
                        Ok(RoomEventCacheUpdate::UpdateTimelineEvents(update)) => {
                            let events = update.diffs.into_iter().flat_map(added).collect();
                            return Some((
                                (room_id, RoomUpdate::Events(events)),
                                Some((subscriber, drop_handles)),
                            ));
                        }
                        Ok(_) => {}
                        Err(broadcast::error::RecvError::Lagged(_)) => {
                            return Some((
                                (room_id, RoomUpdate::Lagged),
                                Some((subscriber, drop_handles)),
                            ));
                        }
                        Err(broadcast::error::RecvError::Closed) => {
                            return Some(((room_id, RoomUpdate::Closed), None));
                        }
                    }
                }
            }
        },
    ))
}

fn added(diff: VectorDiff<SdkTimelineEvent>) -> Vec<SdkTimelineEvent> {
    match diff {
        VectorDiff::Append { values } | VectorDiff::Reset { values } => {
            values.into_iter().collect()
        }
        VectorDiff::PushFront { value }
        | VectorDiff::PushBack { value }
        | VectorDiff::Insert { value, .. }
        | VectorDiff::Set { value, .. } => vec![value],
        VectorDiff::Clear
        | VectorDiff::PopFront
        | VectorDiff::PopBack
        | VectorDiff::Remove { .. }
        | VectorDiff::Truncate { .. } => Vec::new(),
    }
}

#[cfg(all(test, feature = "sqlite"))]
#[path = "service_tests.rs"]
mod tests;