atagen

macltok

public
Star0
Watch0
cff1b0a4eb30290a6fe32fea825c816b696fe4f4
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/search/mod.rsRust1234 lines42.3 KB
1
mod feed;
2
#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
3
mod index;
4
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
5
pub(crate) mod store;
6
7
use std::{
8
    collections::{BTreeMap, HashMap, HashSet, VecDeque},
9
    sync::Arc,
10
    time::Duration,
11
};
12
13
use futures_util::{
14
    StreamExt,
15
    future::join_all,
16
    stream::{AbortHandle, FuturesUnordered, SelectAll, abortable},
17
};
18
use matrix_sdk::{
19
    Client, Room, RoomState, SendOutsideWasm,
20
    deserialized_responses::TimelineEvent as SdkTimelineEvent,
21
    event_cache::{RoomEventCacheSubscriber, RoomEventCacheUpdate},
22
    room::MessagesOptions,
23
};
24
use matrix_sdk_base::RoomInfoNotableUpdateReasons;
25
use matrix_sdk_ui::eyeball_im::VectorDiff;
26
use ruma::{
27
    OwnedRoomId, OwnedUserId, UInt,
28
    events::{
29
        ignored_user_list::IgnoredUserListEventContent,
30
        room::encrypted::OriginalSyncRoomEncryptedEvent,
31
    },
32
    serde::Raw,
33
};
34
use serde_json::Value;
35
use tokio::sync::{Mutex, broadcast, mpsc, oneshot};
36
37
use self::{
38
    feed::{Effect, Observed, Redactions},
39
    index::{Changes, Doc, Hit, KeywordIndex},
40
};
41
use crate::{
42
    CommandOutput, EngineConfig, Error, ErrorKind, RoomSearchCoverage, SearchBackfill,
43
    SearchCommand, SearchCoverage, SearchHit, SearchMatch, SearchPage, SearchScope, SearchSession,
44
    engine::{MediaService, SearchStore, context, executor::EngineBoxFuture},
45
    normalise::{MediaRegistration, MediaScope},
46
};
47
48
const MAX_LIMIT: u16 = 50;
49
/// a consumer that never closes its searches still holds at most this many
50
const MAX_SEARCHES: usize = 8;
51
const COMMIT_AFTER_MS: u64 = 5_000;
52
const COMMIT_BUFFERED: usize = 2_000;
53
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
54
const COMMIT_RETRY_MS: u64 = 60_000;
55
const IDLE_SUBSCRIPTION_MS: u64 = 5 * 60_000;
56
const BACKFILL_RETRY_MS: u64 = 60_000;
57
const BACKFILL_RETRY_MAX_MS: u64 = 15 * 60_000;
58
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
59
const SCHEMA: u64 = 1;
60
61
#[cfg(not(target_family = "wasm"))]
62
type BoxStream<T> = futures_util::stream::BoxStream<'static, T>;
63
#[cfg(target_family = "wasm")]
64
type BoxStream<T> = futures_util::stream::LocalBoxStream<'static, T>;
65
66
/// a local index of message text. keyword search over tantivy is the only one today, an
67
/// embedding index can stand beside it later without the public api noticing
68
#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]
69
pub(crate) trait MessageIndex: SendOutsideWasm {
70
    fn get(&self, key: &str) -> Option<Doc>;
71
    fn by_content(&self, content: &str) -> Option<Doc>;
72
    /// replaces whatever `doc.key` held
73
    fn put(&mut self, doc: Doc);
74
    fn remove(&mut self, key: &str) -> bool;
75
    fn buffered(&self) -> usize;
76
    /// newest first, strictly older than `before`. also hands back the analysed query terms
77
    fn search(
78
        &mut self,
79
        query: &str,
80
        rooms: Option<&[String]>,
81
        before: Option<(u64, u64)>,
82
        limit: usize,
83
    ) -> Result<(Vec<Hit>, Vec<String>), Error>;
84
    fn maintain(&mut self) -> Result<(), Error>;
85
    fn changes(&mut self) -> Result<Changes, Error>;
86
    fn committed(&mut self, changes: &Changes);
87
    fn manifest(&self) -> Value;
88
}
89
90
#[derive(Clone, Debug, Default)]
91
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
92
struct RoomProgress {
93
    oldest: Option<u64>,
94
    newest: Option<u64>,
95
    /// where the next backfill page starts, `None` before the first one
96
    tail: Option<String>,
97
    started: bool,
98
    start_reached: bool,
99
}
100
101
#[derive(Clone, Debug)]
102
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
103
struct Pending {
104
    room: String,
105
    session_id: String,
106
    ts: u64,
107
    raw: String,
108
}
109
110
#[derive(Clone, Copy)]
111
struct Budget {
112
    pages_per_room: u16,
113
    page_size: u16,
114
    interval_ms: u64,
115
    max_pending: usize,
116
}
117
118
impl Budget {
119
    fn from_config(config: &EngineConfig) -> Self {
120
        Self {
121
            pages_per_room: config.search_backfill_pages_per_room,
122
            page_size: config.search_backfill_page_size.max(1),
123
            interval_ms: config.search_backfill_interval_ms,
124
            max_pending: config.max_search_undecrypted,
125
        }
126
    }
127
}
128
129
struct RawPage {
130
    hits: Vec<Hit>,
131
    terms: Vec<String>,
132
    more: bool,
133
    coverage: SearchCoverage,
134
}
135
136
enum Request {
137
    Search {
138
        query: String,
139
        rooms: Option<Vec<String>>,
140
        before: Option<(u64, u64)>,
141
        limit: usize,
142
        coverage: Vec<OwnedRoomId>,
143
        reply: oneshot::Sender<Result<RawPage, Error>>,
144
    },
145
    Coverage(Vec<OwnedRoomId>, oneshot::Sender<SearchCoverage>),
146
    /// a hit that turned out redacted when its preview was read
147
    Forget(String),
148
    Stop(oneshot::Sender<()>),
149
}
150
151
#[derive(Clone)]
152
pub(crate) struct Handle {
153
    requests: mpsc::Sender<Request>,
154
    client: Client,
155
    media: Arc<MediaService>,
156
    sessions: Arc<Sessions>,
157
}
158
159
enum Work {
160
    Backfill {
161
        room: OwnedRoomId,
162
        head: bool,
163
        result: Result<matrix_sdk::room::Messages, Error>,
164
    },
165
    Decrypted {
166
        key: String,
167
        room: OwnedRoomId,
168
        value: Option<Value>,
169
    },
170
    Refetched {
171
        room: OwnedRoomId,
172
        value: Option<Value>,
173
    },
174
}
175
176
struct Failure {
177
    message: String,
178
    retry_at_ms: u64,
179
}
180
181
pub(crate) struct Service {
182
    client: Client,
183
    index: Box<dyn MessageIndex>,
184
    store: Option<Result<SearchStore, Error>>,
185
    rooms: BTreeMap<String, RoomProgress>,
186
    pending: BTreeMap<String, Pending>,
187
    pending_dirty: bool,
188
    redactions: Redactions,
189
    dirty_since: Option<u64>,
190
    commit_after: u64,
191
    budget: Budget,
192
    requests: mpsc::Receiver<Request>,
193
    subscriptions: HashMap<OwnedRoomId, (AbortHandle, u64)>,
194
    updates: SelectAll<BoxStream<(OwnedRoomId, Option<Vec<SdkTimelineEvent>>)>>,
195
    work: FuturesUnordered<EngineBoxFuture<'static, Work>>,
196
    pages: HashMap<OwnedRoomId, u16>,
197
    heads: VecDeque<OwnedRoomId>,
198
    headed: HashSet<OwnedRoomId>,
199
    head_tokens: HashMap<OwnedRoomId, String>,
200
    backfilling: Option<OwnedRoomId>,
201
    next_backfill: u64,
202
    retry_delay: u64,
203
    failures: HashMap<OwnedRoomId, Failure>,
204
    turn: usize,
205
}
206
207
fn now() -> u64 {
208
    ruma::MilliSecondsSinceUnixEpoch::now().get().into()
209
}
210
211
fn value_of(event: &SdkTimelineEvent) -> Option<Value> {
212
    serde_json::from_str(event.raw().json().get()).ok()
213
}
214
215
fn encode_cursor((ts, tie): (u64, u64)) -> String {
216
    format!("{ts:x}.{tie:x}")
217
}
218
219
fn decode_cursor(cursor: &str) -> Option<(u64, u64)> {
220
    let (ts, tie) = cursor.split_once('.')?;
221
    Some((
222
        u64::from_str_radix(ts, 16).ok()?,
223
        u64::from_str_radix(tie, 16).ok()?,
224
    ))
225
}
226
227
impl Handle {
228
    pub(crate) async fn execute(&self, command: SearchCommand) -> Result<CommandOutput, Error> {
229
        match command {
230
            SearchCommand::Search {
231
                query,
232
                scope,
233
                limit,
234
            } => {
235
                let (page, media) = self.search(&query, &scope, None, limit).await?;
236
                let search_id = self.sessions.open(&self.media, query, scope, media).await;
237
                Ok(CommandOutput::SearchSession(SearchSession { search_id, page }))
238
            }
239
            SearchCommand::More {
240
                search_id,
241
                cursor,
242
                limit,
243
            } => {
244
                let (query, scope) = self.sessions.get(&search_id).await?;
245
                let (page, media) = self.search(&query, &scope, Some(cursor.as_str()), limit).await?;
246
                self.sessions.register(&self.media, &search_id, media).await?;
247
                Ok(CommandOutput::SearchPage(page))
248
            }
249
            SearchCommand::Close { search_id } => {
250
                self.sessions.close(&self.media, &search_id).await?;
251
                Ok(CommandOutput::Done)
252
            }
253
            SearchCommand::Coverage(scope) => {
254
                let (reply, response) = oneshot::channel();
255
                self.requests
256
                    .send(Request::Coverage(self.scope_rooms(&scope), reply))
257
                    .await
258
                    .map_err(|_| Error::interrupted())?;
259
                Ok(CommandOutput::SearchCoverage(
260
                    response.await.map_err(|_| Error::interrupted())?,
261
                ))
262
            }
263
        }
264
    }
265
266
    /// the media service is cleared beside this, so nothing is released one by one
267
    pub(crate) async fn stop(&self) {
268
        let (reply, response) = oneshot::channel();
269
        if self.requests.send(Request::Stop(reply)).await.is_ok() {
270
            let _ = response.await;
271
        }
272
        self.sessions.open.lock().await.clear();
273
    }
274
275
    fn scope_rooms(&self, scope: &SearchScope) -> Vec<OwnedRoomId> {
276
        match scope {
277
            SearchScope::Room(room_id) => vec![room_id.clone()],
278
            SearchScope::Rooms(room_ids) => room_ids.clone(),
279
            SearchScope::All => self
280
                .client
281
                .joined_rooms()
282
                .into_iter()
283
                .map(|room| room.room_id().to_owned())
284
                .collect(),
285
        }
286
    }
287
288
    async fn search(
289
        &self,
290
        query: &str,
291
        scope: &SearchScope,
292
        cursor: Option<&str>,
293
        limit: u16,
294
    ) -> Result<(SearchPage, Vec<MediaRegistration>), Error> {
295
        if limit == 0 || limit > MAX_LIMIT {
296
            return Err(Error::new(
297
                ErrorKind::Invalid,
298
                "search limit must be between 1 and 50",
299
            ));
300
        }
301
        if query.trim().is_empty() {
302
            return Err(Error::new(ErrorKind::Invalid, "search query is empty"));
303
        }
304
        let before = match cursor {
305
            Some(cursor) => Some(
306
                decode_cursor(cursor)
307
                    .ok_or_else(|| Error::new(ErrorKind::Invalid, "search cursor is invalid"))?,
308
            ),
309
            None => None,
310
        };
311
        let rooms = self.scope_rooms(scope);
312
        let filter = Some(rooms.iter().map(ToString::to_string).collect());
313
        let (reply, response) = oneshot::channel();
314
        self.requests
315
            .send(Request::Search {
316
                query: query.to_owned(),
317
                rooms: filter,
318
                before,
319
                limit: usize::from(limit),
320
                coverage: rooms,
321
                reply,
322
            })
323
            .await
324
            .map_err(|_| Error::interrupted())?;
325
        let raw = response.await.map_err(|_| Error::interrupted())??;
326
        let next_cursor = raw
327
            .more
328
            .then(|| raw.hits.last().map(|hit| encode_cursor((hit.ts, hit.tie))))
329
            .flatten();
330
        let ignored = self.ignored_users().await;
331
        let current_user_id = self.client.user_id().map(ToOwned::to_owned);
332
        let previews = join_all(
333
            raw.hits
334
                .into_iter()
335
                .filter(|hit| !ignored.iter().any(|user| user.as_str() == hit.sender))
336
                .map(|hit| self.preview(hit, current_user_id.clone(), &raw.terms)),
337
        )
338
        .await;
339
        let mut hits = Vec::new();
340
        let mut media = Vec::new();
341
        for preview in previews {
342
            match preview {
343
                Preview::Shown(hit, registrations) => {
344
                    hits.push(hit);
345
                    media.extend(registrations);
346
                }
347
                Preview::Redacted(key) => {
348
                    let _ = self.requests.send(Request::Forget(key)).await;
349
                }
350
                Preview::Unavailable => {}
351
            }
352
        }
353
        let page = SearchPage {
354
            hits,
355
            next_cursor,
356
            coverage: raw.coverage,
357
        };
358
        Ok((page, media))
359
    }
360
361
    async fn ignored_users(&self) -> Vec<OwnedUserId> {
362
        self.client
363
            .account()
364
            .account_data::<IgnoredUserListEventContent>()
365
            .await
366
            .ok()
367
            .flatten()
368
            .and_then(|raw| raw.deserialize().ok())
369
            .map(|content| content.ignored_users.into_keys().collect())
370
            .unwrap_or_default()
371
    }
372
373
    async fn preview(
374
        &self,
375
        hit: Hit,
376
        current_user_id: Option<OwnedUserId>,
377
        terms: &[String],
378
    ) -> Preview {
379
        let (Ok(room_id), Ok(event_id)) = (
380
            OwnedRoomId::try_from(hit.room.as_str()),
381
            ruma::OwnedEventId::try_from(hit.key.as_str()),
382
        ) else {
383
            return Preview::Unavailable;
384
        };
385
        let Some(room) = self.client.get_room(&room_id) else {
386
            return Preview::Unavailable;
387
        };
388
        let Some((original, highlight)) = load(&room, &hit.key).await else {
389
            return Preview::Unavailable;
390
        };
391
        if original.pointer("/unsigned/redacted_because").is_some() {
392
            return Preview::Redacted(hit.key);
393
        }
394
        let mut values = vec![original];
395
        if hit.content != hit.key
396
            && let Some((edit, _)) = load(&room, &hit.content).await
397
        {
398
            values.push(edit);
399
        }
400
        let Some((event, registrations)) =
401
            context::preview(&room, current_user_id.as_deref(), &values, highlight).await
402
        else {
403
            return Preview::Unavailable;
404
        };
405
        let text = match &event.content {
406
            crate::TimelineContent::Message(message) => Some(message.body.as_str()),
407
            crate::TimelineContent::Sticker { body, .. } => Some(body.as_str()),
408
            crate::TimelineContent::Poll { question } => Some(question.as_str()),
409
            _ => None,
410
        };
411
        let matches = text
412
            .map(|text| index::match_ranges(text, terms))
413
            .unwrap_or_default()
414
            .into_iter()
415
            .filter_map(|(start, end)| {
416
                Some(SearchMatch {
417
                    start: u32::try_from(start).ok()?,
418
                    end: u32::try_from(end).ok()?,
419
                })
420
            })
421
            .collect();
422
        Preview::Shown(
423
            SearchHit {
424
                room_id,
425
                event_id,
426
                timestamp_ms: event.timestamp_ms,
427
                event,
428
                matches,
429
            },
430
            registrations,
431
        )
432
    }
433
}
434
435
enum Preview {
436
    Shown(SearchHit, Vec<MediaRegistration>),
437
    Redacted(String),
438
    Unavailable,
439
}
440
441
/// open searches, oldest first, each owning the media of the hits it handed out
442
#[derive(Default)]
443
struct Sessions {
444
    open: Mutex<VecDeque<Session>>,
445
}
446
447
struct Session {
448
    id: String,
449
    query: String,
450
    scope: SearchScope,
451
    media: MediaScope,
452
}
453
454
impl Sessions {
455
    async fn open(
456
        &self,
457
        media: &MediaService,
458
        query: String,
459
        scope: SearchScope,
460
        registrations: Vec<MediaRegistration>,
461
    ) -> String {
462
        let media_scope = media.context_scope();
463
        let id = format!("search-{}", media_scope.nonce());
464
        let mut open = self.open.lock().await;
465
        media.register(&media_scope, registrations).await;
466
        open.push_back(Session {
467
            id: id.clone(),
468
            query,
469
            scope,
470
            media: media_scope,
471
        });
472
        while open.len() > MAX_SEARCHES {
473
            if let Some(oldest) = open.pop_front() {
474
                media.release(&oldest.media).await;
475
            }
476
        }
477
        id
478
    }
479
480
    async fn get(&self, id: &str) -> Result<(String, SearchScope), Error> {
481
        let open = self.open.lock().await;
482
        let session = open.iter().find(|session| session.id == id).ok_or_else(not_open)?;
483
        Ok((session.query.clone(), session.scope.clone()))
484
    }
485
486
    /// under the lock, so a page that lands after its close registers nothing
487
    async fn register(
488
        &self,
489
        media: &MediaService,
490
        id: &str,
491
        registrations: Vec<MediaRegistration>,
492
    ) -> Result<(), Error> {
493
        let open = self.open.lock().await;
494
        let session = open.iter().find(|session| session.id == id).ok_or_else(not_open)?;
495
        media.register(&session.media, registrations).await;
496
        Ok(())
497
    }
498
499
    async fn close(&self, media: &MediaService, id: &str) -> Result<(), Error> {
500
        let mut open = self.open.lock().await;
501
        let index = open
502
            .iter()
503
            .position(|session| session.id == id)
504
            .ok_or_else(not_open)?;
505
        if let Some(session) = open.remove(index) {
506
            media.release(&session.media).await;
507
        }
508
        Ok(())
509
    }
510
}
511
512
fn not_open() -> Error {
513
    Error::new(ErrorKind::SearchNotOpen, "search is closed or was never opened")
514
}
515
516
async fn load(room: &Room, event_id: &str) -> Option<(Value, bool)> {
517
    let event_id = ruma::EventId::parse(event_id).ok()?;
518
    let event = room
519
        .load_or_fetch_event(&event_id, None)
520
        .await
521
        .inspect_err(|error| tracing::debug!(%event_id, %error, "loading a search hit failed"))
522
        .ok()?;
523
    let highlight = event
524
        .push_actions()
525
        .is_some_and(|actions| actions.iter().any(ruma::push::Action::is_highlight));
526
    Some((value_of(&event)?, highlight))
527
}
528
529
#[cfg(any(feature = "sqlite", feature = "indexeddb"))]
530
fn decode_meta(
531
    store: &SearchStore,
532
    blobs: &BTreeMap<String, Vec<u8>>,
533
) -> Option<(
534
    Value,
535
    BTreeMap<String, RoomProgress>,
536
    BTreeMap<String, Pending>,
537
)> {
538
    let meta: Value = serde_json::from_slice(&store.open_blob("meta", blobs.get("meta")?)?).ok()?;
539
    if meta["schema"].as_u64() != Some(SCHEMA) {
540
        return None;
541
    }
542
    let rooms = serde_json::from_value(meta["rooms"].clone()).ok()?;
543
    let pending = match blobs.get("utd") {
544
        Some(sealed) => serde_json::from_slice(&store.open_blob("utd", sealed)?).ok()?,
545
        None => BTreeMap::new(),
546
    };
547
    Some((meta["index"].clone(), rooms, pending))
548
}
549
550
impl Service {
551
    pub(crate) fn new(
552
        client: Client,
553
        media: Arc<MediaService>,
554
        store: Option<Result<SearchStore, Error>>,
555
        config: &EngineConfig,
556
    ) -> (Handle, Self) {
557
        let (sender, requests) = mpsc::channel(config.command_capacity.max(1));
558
        let service = Self {
559
            client: client.clone(),
560
            index: Box::new(KeywordIndex::new(store.as_ref().is_some_and(Result::is_ok))),
561
            store,
562
            rooms: BTreeMap::new(),
563
            pending: BTreeMap::new(),
564
            pending_dirty: false,
565
            redactions: Redactions::default(),
566
            dirty_since: None,
567
            commit_after: 0,
568
            budget: Budget::from_config(config),
569
            requests,
570
            subscriptions: HashMap::new(),
571
            updates: SelectAll::new(),
572
            work: FuturesUnordered::new(),
573
            pages: HashMap::new(),
574
            heads: VecDeque::new(),
575
            headed: HashSet::new(),
576
            head_tokens: HashMap::new(),
577
            backfilling: None,
578
            next_backfill: 0,
579
            retry_delay: BACKFILL_RETRY_MS,
580
            failures: HashMap::new(),
581
            turn: 0,
582
        };
583
        (
584
            Handle {
585
                requests: sender,
586
                client,
587
                media,
588
                sessions: Arc::default(),
589
            },
590
            service,
591
        )
592
    }
593
594
    fn persistent(&self) -> bool {
595
        matches!(self.store, Some(Ok(_)))
596
    }
597
598
    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
599
    async fn load(&mut self) {
600
        let Some(Ok(store)) = &self.store else {
601
            if let Some(Err(error)) = &self.store {
602
                tracing::error!(%error, "search store unavailable, the index lives in memory");
603
            }
604
            return;
605
        };
606
        let blobs = match store.load().await {
607
            Ok(blobs) => blobs,
608
            Err(error) => {
609
                tracing::error!(%error, "reading the search store failed, the index lives in memory");
610
                self.store = Some(Err(error));
611
                self.index = Box::new(KeywordIndex::new(false));
612
                return;
613
            }
614
        };
615
        if blobs.is_empty() {
616
            return;
617
        }
618
        let restored = decode_meta(store, &blobs).and_then(|(manifest, rooms, pending)| {
619
            let index = KeywordIndex::load(&manifest, |key| {
620
                blobs
621
                    .get(key)
622
                    .and_then(|sealed| store.open_blob(key, sealed))
623
            })
624
            .inspect_err(|error| tracing::warn!(%error, "search segments are unreadable"))
625
            .ok()?;
626
            Some((index, rooms, pending))
627
        });
628
        drop(blobs);
629
        match restored {
630
            Some((index, rooms, pending)) => {
631
                self.index = Box::new(index);
632
                self.rooms = rooms;
633
                self.pending = pending;
634
            }
635
            None => {
636
                tracing::warn!("resetting an unreadable search index");
637
                if let Err(error) = store.reset().await {
638
                    tracing::error!(%error, "search store unavailable, the index lives in memory");
639
                    self.store = Some(Err(error));
640
                    self.index = Box::new(KeywordIndex::new(false));
641
                }
642
            }
643
        }
644
    }
645
646
    #[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
647
    async fn load(&mut self) {
648
        if let Some(Err(error)) = &self.store {
649
            tracing::error!(%error, "search store unavailable, the index lives in memory");
650
        }
651
    }
652
653
    #[cfg(any(feature = "sqlite", feature = "indexeddb"))]
654
    async fn commit(&mut self) {
655
        if !self.persistent() {
656
            if let Err(error) = self.index.maintain() {
657
                tracing::warn!(%error, "maintaining the search index failed");
658
            }
659
            self.dirty_since = None;
660
            return;
661
        }
662
        let changes = match self.index.changes() {
663
            Ok(changes) => changes,
664
            Err(error) => {
665
                tracing::warn!(%error, "preparing a search commit failed");
666
                self.commit_after = now() + COMMIT_RETRY_MS;
667
                return;
668
            }
669
        };
670
        let mut puts = changes.puts.clone();
671
        let meta = serde_json::json!({
672
            "schema": SCHEMA,
673
            "index": self.index.manifest(),
674
            "rooms": self.rooms,
675
        });
676
        puts.push(("meta".to_owned(), meta.to_string().into_bytes()));
677
        if self.pending_dirty {
678
            puts.push((
679
                "utd".to_owned(),
680
                serde_json::to_vec(&self.pending).unwrap_or_default(),
681
            ));
682
        }
683
        let Some(Ok(store)) = &self.store else {
684
            return;
685
        };
686
        match store.commit(puts, changes.deletes.clone()).await {
687
            Ok(()) => {
688
                self.index.committed(&changes);
689
                self.pending_dirty = false;
690
                self.dirty_since = None;
691
            }
692
            Err(error) => {
693
                tracing::warn!(%error, "saving the search index failed, retrying later");
694
                self.commit_after = now() + COMMIT_RETRY_MS;
695
            }
696
        }
697
    }
698
699
    #[cfg(not(any(feature = "sqlite", feature = "indexeddb")))]
700
    async fn commit(&mut self) {
701
        if let Err(error) = self.index.maintain() {
702
            tracing::warn!(%error, "maintaining the search index failed");
703
        }
704
        self.dirty_since = None;
705
    }
706
707
    fn touch(&mut self) {
708
        self.dirty_since.get_or_insert_with(now);
709
    }
710
711
    fn observe(&mut self, room: &OwnedRoomId, value: &Value) {
712
        let observed = feed::observe(value);
713
        let key = match &observed {
714
            Observed::Message { key, .. }
715
            | Observed::Redacted { key }
716
            | Observed::Encrypted { key, .. } => Some(key.clone()),
717
            Observed::Edit(edit) => Some(edit.id.clone()),
718
            Observed::Redaction { .. } | Observed::Other => None,
719
        };
720
        if let Observed::Encrypted {
721
            key,
722
            session_id,
723
            ts,
724
        } = &observed
725
        {
726
            if self.index.get(key).is_none() && !self.pending.contains_key(key) {
727
                self.pending.insert(
728
                    key.clone(),
729
                    Pending {
730
                        room: room.to_string(),
731
                        session_id: session_id.clone(),
732
                        ts: *ts,
733
                        raw: value.to_string(),
734
                    },
735
                );
736
                while self.pending.len() > self.budget.max_pending {
737
                    let Some(oldest) = self
738
                        .pending
739
                        .iter()
740
                        .min_by_key(|(_, pending)| pending.ts)
741
                        .map(|(key, _)| key.clone())
742
                    else {
743
                        break;
744
                    };
745
                    self.pending.remove(&oldest);
746
                }
747
                self.pending_dirty = true;
748
                self.touch();
749
            }
750
            return;
751
        }
752
        if let Some(key) = key
753
            && self.pending.remove(&key).is_some()
754
        {
755
            self.pending_dirty = true;
756
            self.touch();
757
        }
758
        match feed::apply(
759
            self.index.as_mut(),
760
            &mut self.redactions,
761
            room.as_str(),
762
            observed,
763
        ) {
764
            Effect::Indexed(ts) => {
765
                let progress = self.rooms.entry(room.to_string()).or_default();
766
                progress.oldest = Some(progress.oldest.map_or(ts, |oldest| oldest.min(ts)));
767
                progress.newest = Some(progress.newest.map_or(ts, |newest| newest.max(ts)));
768
                self.touch();
769
            }
770
            Effect::Removed => self.touch(),
771
            Effect::Refetch(key) => {
772
                self.touch();
773
                let room_id = room.clone();
774
                let client = self.client.clone();
775
                self.work.push(Box::pin(async move {
776
                    let value = match client.get_room(&room_id) {
777
                        Some(room) => load(&room, &key).await.map(|(value, _)| value),
778
                        None => None,
779
                    };
780
                    Work::Refetched {
781
                        room: room_id,
782
                        value,
783
                    }
784
                }));
785
            }
786
            Effect::Nothing => {}
787
        }
788
    }
789
790
    fn observe_all(&mut self, room: &OwnedRoomId, events: &[SdkTimelineEvent]) {
791
        for event in events {
792
            if let Some(value) = value_of(event) {
793
                self.observe(room, &value);
794
            }
795
        }
796
    }
797
798
    async fn subscribe(&mut self, room_id: OwnedRoomId) {
799
        if let Some((_, seen)) = self.subscriptions.get_mut(&room_id) {
800
            *seen = now();
801
            return;
802
        }
803
        let Some(room) = self.client.get_room(&room_id) else {
804
            return;
805
        };
806
        if room.state() != RoomState::Joined {
807
            return;
808
        }
809
        let Ok((cache, drop_handles)) = room.event_cache().await else {
810
            return;
811
        };
812
        let Ok((initial, subscriber)) = cache.subscribe().await else {
813
            return;
814
        };
815
        self.observe_all(&room_id, &initial);
816
        let stream = updates(room_id.clone(), subscriber, drop_handles);
817
        let (stream, abort) = abortable(stream);
818
        self.updates.push(Box::pin(stream));
819
        self.subscriptions.insert(room_id.clone(), (abort, now()));
820
        if self.headed.insert(room_id.clone()) {
821
            self.heads.push_back(room_id);
822
        }
823
    }
824
825
    fn drop_idle(&mut self) {
826
        let cutoff = now().saturating_sub(IDLE_SUBSCRIPTION_MS);
827
        self.subscriptions.retain(|_, (abort, seen)| {
828
            let keep = *seen >= cutoff;
829
            if !keep {
830
                abort.abort();
831
            }
832
            keep
833
        });
834
    }
835
836
    fn coverage(&self, rooms: &[OwnedRoomId]) -> SearchCoverage {
837
        let rooms = rooms
838
            .iter()
839
            .map(|room_id| {
840
                let progress = self
841
                    .rooms
842
                    .get(room_id.as_str())
843
                    .cloned()
844
                    .unwrap_or_default();
845
                let backfill = if progress.start_reached {
846
                    SearchBackfill::Complete
847
                } else if self.budget.pages_per_room == 0 {
848
                    SearchBackfill::Disabled
849
                } else if self.backfilling.as_ref() == Some(room_id) {
850
                    SearchBackfill::Running
851
                } else if let Some(Failure {
852
                    message,
853
                    retry_at_ms,
854
                }) = self.failures.get(room_id)
855
                {
856
                    SearchBackfill::Failed {
857
                        message: message.clone(),
858
                        retry_at_ms: *retry_at_ms,
859
                    }
860
                } else if self.pages.get(room_id).copied().unwrap_or(0)
861
                    >= self.budget.pages_per_room
862
                {
863
                    SearchBackfill::Paused
864
                } else {
865
                    SearchBackfill::Waiting
866
                };
867
                RoomSearchCoverage {
868
                    room_id: room_id.clone(),
869
                    oldest_indexed_ms: progress.oldest,
870
                    newest_indexed_ms: progress.newest,
871
                    start_reached: progress.start_reached,
872
                    backfill,
873
                }
874
            })
875
            .collect();
876
        SearchCoverage {
877
            rooms,
878
            persistent: self.persistent(),
879
            undecrypted: self.pending.len(),
880
        }
881
    }
882
883
    /// head passes first, so what arrived while we were away is found before anything older
884
    fn schedule_backfill(&mut self) {
885
        if self.backfilling.is_some()
886
            || self.budget.pages_per_room == 0
887
            || now() < self.next_backfill
888
        {
889
            return;
890
        }
891
        let budget = self.budget.pages_per_room;
892
        let spent = |pages: &HashMap<OwnedRoomId, u16>, room: &OwnedRoomId| {
893
            pages.get(room).copied().unwrap_or(0) >= budget
894
        };
895
        let blocked = |failures: &HashMap<OwnedRoomId, Failure>, room: &OwnedRoomId| {
896
            failures
897
                .get(room)
898
                .is_some_and(|failure| failure.retry_at_ms > now())
899
        };
900
        let mut chosen = None;
901
        while let Some(room_id) = self.heads.pop_front() {
902
            let started = self
903
                .rooms
904
                .get(room_id.as_str())
905
                .is_some_and(|progress| progress.started);
906
            if started && !spent(&self.pages, &room_id) && !blocked(&self.failures, &room_id) {
907
                chosen = Some((room_id, true));
908
                break;
909
            }
910
        }
911
        if chosen.is_none() {
912
            let mut joined: Vec<_> = self
913
                .client
914
                .joined_rooms()
915
                .into_iter()
916
                .map(|room| room.room_id().to_owned())
917
                .collect();
918
            joined.sort();
919
            let count = joined.len();
920
            for offset in 0..count {
921
                let room_id = &joined[(self.turn + offset) % count];
922
                let done = self
923
                    .rooms
924
                    .get(room_id.as_str())
925
                    .is_some_and(|progress| progress.start_reached);
926
                if !done && !spent(&self.pages, room_id) && !blocked(&self.failures, room_id) {
927
                    self.turn = (self.turn + offset + 1) % count;
928
                    chosen = Some((room_id.clone(), false));
929
                    break;
930
                }
931
            }
932
        }
933
        let Some((room_id, head)) = chosen else {
934
            return;
935
        };
936
        let Some(room) = self.client.get_room(&room_id) else {
937
            return;
938
        };
939
        let from = if head {
940
            self.head_tokens.get(&room_id).cloned()
941
        } else {
942
            self.rooms
943
                .get(room_id.as_str())
944
                .and_then(|progress| progress.tail.clone())
945
        };
946
        *self.pages.entry(room_id.clone()).or_default() += 1;
947
        self.backfilling = Some(room_id.clone());
948
        let mut options = MessagesOptions::backward();
949
        options.from = from;
950
        options.limit = UInt::from(self.budget.page_size);
951
        self.work.push(Box::pin(async move {
952
            let result = room
953
                .messages(options)
954
                .await
955
                .map_err(|error| Error::from_sdk("backfilling search history failed", error));
956
            Work::Backfill {
957
                room: room_id,
958
                head,
959
                result,
960
            }
961
        }));
962
    }
963
964
    fn backfilled(
965
        &mut self,
966
        room_id: OwnedRoomId,
967
        head: bool,
968
        result: Result<matrix_sdk::room::Messages, Error>,
969
    ) {
970
        self.backfilling = None;
971
        let messages = match result {
972
            Ok(messages) => messages,
973
            Err(error) => {
974
                let retry_at_ms = now() + self.retry_delay;
975
                self.next_backfill = retry_at_ms;
976
                self.retry_delay = (self.retry_delay * 2).min(BACKFILL_RETRY_MAX_MS);
977
                self.failures.insert(
978
                    room_id,
979
                    Failure {
980
                        message: error.message.clone(),
981
                        retry_at_ms,
982
                    },
983
                );
984
                return;
985
            }
986
        };
987
        self.retry_delay = BACKFILL_RETRY_MS;
988
        self.failures.remove(&room_id);
989
        self.next_backfill = now() + self.budget.interval_ms;
990
        let mut overlapped = false;
991
        for event in &messages.chunk {
992
            let Some(value) = value_of(event) else {
993
                continue;
994
            };
995
            if head
996
                && let Some(key) = value.get("event_id").and_then(Value::as_str)
997
                && self.index.get(key).is_some_and(|doc| !doc.provisional)
998
            {
999
                overlapped = true;
1000
            }
1001
            self.observe(&room_id, &value);
1002
        }
1003
        let ended = messages.end.is_none() || messages.chunk.is_empty();
1004
        if head {
1005
            match messages.end {
1006
                Some(end) if !overlapped && !ended => {
1007
                    self.head_tokens.insert(room_id.clone(), end);
1008
                    self.heads.push_front(room_id);
1009
                }
1010
                _ => {
1011
                    self.head_tokens.remove(&room_id);
1012
                }
1013
            }
1014
            return;
1015
        }
1016
        let progress = self.rooms.entry(room_id.to_string()).or_default();
1017
        progress.started = true;
1018
        progress.tail = messages.end;
1019
        if ended {
1020
            progress.start_reached = true;
1021
            progress.tail = None;
1022
        }
1023
        self.touch();
1024
    }
1025
1026
    fn retry_pending(&mut self, room_id: Option<&OwnedRoomId>, session_id: Option<&str>) {
1027
        let due: Vec<_> = self
1028
            .pending
1029
            .iter()
1030
            .filter(|(_, pending)| {
1031
                room_id.is_none_or(|room_id| pending.room == room_id.as_str())
1032
                    && session_id.is_none_or(|session_id| pending.session_id == session_id)
1033
            })
1034
            .map(|(key, pending)| (key.clone(), pending.clone()))
1035
            .collect();
1036
        for (key, pending) in due {
1037
            let Ok(room_id) = OwnedRoomId::try_from(pending.room.as_str()) else {
1038
                continue;
1039
            };
1040
            let Some(room) = self.client.get_room(&room_id) else {
1041
                continue;
1042
            };
1043
            self.work.push(Box::pin(async move {
1044
                let value = async {
1045
                    let raw: Raw<OriginalSyncRoomEncryptedEvent> =
1046
                        Raw::from_json_string(pending.raw).ok()?;
1047
                    let event = room.decrypt_event(&raw, None).await.ok()?;
1048
                    value_of(&event).filter(|value| value["type"] != "m.room.encrypted")
1049
                }
1050
                .await;
1051
                Work::Decrypted {
1052
                    key,
1053
                    room: room_id,
1054
                    value,
1055
                }
1056
            }));
1057
        }
1058
    }
1059
1060
    fn complete(&mut self, work: Work) {
1061
        match work {
1062
            Work::Backfill { room, head, result } => self.backfilled(room, head, result),
1063
            Work::Decrypted { key, room, value } => {
1064
                if let Some(value) = value {
1065
                    self.pending.remove(&key);
1066
                    self.pending_dirty = true;
1067
                    self.observe(&room, &value);
1068
                    self.touch();
1069
                }
1070
            }
1071
            Work::Refetched { room, value } => {
1072
                if let Some(value) = value {
1073
                    self.observe(&room, &value);
1074
                }
1075
            }
1076
        }
1077
    }
1078
1079
    fn search(
1080
        &mut self,
1081
        query: &str,
1082
        rooms: Option<&[String]>,
1083
        before: Option<(u64, u64)>,
1084
        limit: usize,
1085
        coverage: &[OwnedRoomId],
1086
    ) -> Result<RawPage, Error> {
1087
        let (mut hits, terms) = self.index.search(query, rooms, before, limit + 1)?;
1088
        let more = hits.len() > limit;
1089
        hits.truncate(limit);
1090
        Ok(RawPage {
1091
            hits,
1092
            terms,
1093
            more,
1094
            coverage: self.coverage(coverage),
1095
        })
1096
    }
1097
1098
    pub(crate) async fn run(mut self) {
1099
        self.load().await;
1100
        let client = self.client.clone();
1101
        let mut generic = client.event_cache().subscribe_to_room_generic_updates();
1102
        let mut keys = client
1103
            .encryption()
1104
            .room_keys_received_stream()
1105
            .await
1106
            .map(Box::pin);
1107
        let mut room_info = client.room_info_notable_update_receiver();
1108
        self.retry_pending(None, None);
1109
        let mut tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1110
        loop {
1111
            self.schedule_backfill();
1112
            tokio::select! {
1113
                biased;
1114
                request = self.requests.recv() => match request {
1115
                    Some(Request::Stop(reply)) => {
1116
                        if self.dirty_since.is_some() || self.index.buffered() > 0 {
1117
                            self.commit().await;
1118
                        }
1119
                        for (abort, _) in self.subscriptions.values() {
1120
                            abort.abort();
1121
                        }
1122
                        let _ = reply.send(());
1123
                        return;
1124
                    }
1125
                    Some(Request::Search { query, rooms, before, limit, coverage, reply }) => {
1126
                        let _ = reply.send(self.search(&query, rooms.as_deref(), before, limit, &coverage));
1127
                    }
1128
                    Some(Request::Coverage(rooms, reply)) => {
1129
                        let _ = reply.send(self.coverage(&rooms));
1130
                    }
1131
                    Some(Request::Forget(key)) => {
1132
                        let observed = Observed::Redacted { key };
1133
                        if feed::apply(self.index.as_mut(), &mut self.redactions, "", observed) == Effect::Removed {
1134
                            self.touch();
1135
                        }
1136
                    }
1137
                    None => return,
1138
                },
1139
                Some(work) = self.work.next(), if !self.work.is_empty() => self.complete(work),
1140
                Some((room_id, events)) = self.updates.next(), if !self.updates.is_empty() => {
1141
                    if let Some((_, seen)) = self.subscriptions.get_mut(&room_id) {
1142
                        *seen = now();
1143
                    }
1144
                    match events {
1145
                        Some(events) => self.observe_all(&room_id, &events),
1146
                        None => {
1147
                            if let Some((abort, _)) = self.subscriptions.remove(&room_id) {
1148
                                abort.abort();
1149
                            }
1150
                            self.subscribe(room_id).await;
1151
                        }
1152
                    }
1153
                }
1154
                update = generic.recv() => match update {
1155
                    Ok(update) => self.subscribe(update.room_id).await,
1156
                    Err(broadcast::error::RecvError::Lagged(_)) => {}
1157
                    Err(broadcast::error::RecvError::Closed) => return,
1158
                },
1159
                Some(update) = async { keys.as_mut()?.next().await }, if keys.is_some() => match update {
1160
                    Ok(infos) => {
1161
                        for info in infos {
1162
                            self.retry_pending(Some(&info.room_id), Some(&info.session_id));
1163
                        }
1164
                    }
1165
                    Err(_) => self.retry_pending(None, None),
1166
                },
1167
                update = room_info.recv() => match update {
1168
                    Ok(update) if update.reasons.contains(RoomInfoNotableUpdateReasons::MEMBERSHIP) => {
1169
                        let left = self.client.get_room(&update.room_id).is_none_or(|room| room.state() != RoomState::Joined);
1170
                        if left && let Some((abort, _)) = self.subscriptions.remove(&update.room_id) {
1171
                            abort.abort();
1172
                        }
1173
                    }
1174
                    Ok(_) | Err(broadcast::error::RecvError::Lagged(_)) => {}
1175
                    Err(broadcast::error::RecvError::Closed) => return,
1176
                },
1177
                () = &mut tick => {
1178
                    tick = Box::pin(matrix_sdk::sleep::sleep(Duration::from_secs(1)));
1179
                    self.drop_idle();
1180
                    let due = self.dirty_since.is_some_and(|since| now() >= since + COMMIT_AFTER_MS)
1181
                        || self.index.buffered() >= COMMIT_BUFFERED;
1182
                    if due && now() >= self.commit_after {
1183
                        self.commit().await;
1184
                    }
1185
                }
1186
            }
1187
        }
1188
    }
1189
}
1190
1191
fn updates(
1192
    room_id: OwnedRoomId,
1193
    subscriber: RoomEventCacheSubscriber,
1194
    drop_handles: Arc<matrix_sdk::event_cache::EventCacheDropHandles>,
1195
) -> BoxStream<(OwnedRoomId, Option<Vec<SdkTimelineEvent>>)> {
1196
    Box::pin(futures_util::stream::unfold(
1197
        (subscriber, drop_handles),
1198
        move |(mut subscriber, drop_handles)| {
1199
            let room_id = room_id.clone();
1200
            async move {
1201
                loop {
1202
                    match subscriber.recv().await {
1203
                        Ok(RoomEventCacheUpdate::UpdateTimelineEvents(update)) => {
1204
                            let events = update.diffs.into_iter().flat_map(added).collect();
1205
                            return Some(((room_id, Some(events)), (subscriber, drop_handles)));
1206
                        }
1207
                        Ok(_) => {}
1208
                        Err(broadcast::error::RecvError::Lagged(_)) => {
1209
                            return Some(((room_id, None), (subscriber, drop_handles)));
1210
                        }
1211
                        Err(broadcast::error::RecvError::Closed) => return None,
1212
                    }
1213
                }
1214
            }
1215
        },
1216
    ))
1217
}
1218
1219
fn added(diff: VectorDiff<SdkTimelineEvent>) -> Vec<SdkTimelineEvent> {
1220
    match diff {
1221
        VectorDiff::Append { values } | VectorDiff::Reset { values } => {
1222
            values.into_iter().collect()
1223
        }
1224
        VectorDiff::PushFront { value }
1225
        | VectorDiff::PushBack { value }
1226
        | VectorDiff::Insert { value, .. }
1227
        | VectorDiff::Set { value, .. } => vec![value],
1228
        VectorDiff::Clear
1229
        | VectorDiff::PopFront
1230
        | VectorDiff::PopBack
1231
        | VectorDiff::Remove { .. }
1232
        | VectorDiff::Truncate { .. } => Vec::new(),
1233
    }
1234
}