matrix_sdk_search/index/
mod.rs

1// Copyright 2024 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15/// A module for building a [`RoomIndex`]
16pub mod builder;
17
18use std::{
19    collections::{HashMap, HashSet},
20    fmt,
21};
22
23use once_cell::sync::OnceCell;
24use ruma::{EventId, MilliSecondsSinceUnixEpoch, OwnedEventId, OwnedRoomId, OwnedUserId, RoomId};
25use tantivy::{
26    Index, IndexReader, TantivyDocument, collector::TopDocs, directory::error::OpenDirectoryError,
27    query::QueryParser, schema::Value,
28};
29use tracing::{debug, error, warn};
30
31use crate::{
32    OpStamp, TANTIVY_INDEX_MEMORY_BUDGET,
33    error::IndexError,
34    schema::{MatrixSearchIndexSchema, RoomMessageSchema},
35    writer::SearchIndexWriter,
36};
37
38/// The subset of an event's data required to index it and later retrieve it.
39///
40/// Produced by the matrix-sdk layer, which knows how to extract searchable text
41/// from each event type. This crate stays agnostic to Matrix event content.
42#[derive(Debug, Clone)]
43pub struct IndexableEvent {
44    /// The event's own id (primary key).
45    pub event_id: OwnedEventId,
46    /// The id used as the deletion key: the original event id for edits,
47    /// otherwise the event's own id.
48    pub original_event_id: OwnedEventId,
49    /// The sender of the event.
50    pub sender: OwnedUserId,
51    /// The origin server timestamp of the event.
52    pub timestamp: MilliSecondsSinceUnixEpoch,
53    /// The text to index for this event.
54    pub body: String,
55}
56
57/// A struct to represent the operations on a [`RoomIndex`]
58#[derive(Debug, Clone)]
59pub enum RoomIndexOperation {
60    /// Add this event to the index.
61    Add(IndexableEvent),
62    /// Remove all documents in the index where
63    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id.
64    Remove(OwnedEventId),
65    /// Replace all documents in the index where
66    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id with
67    /// the new event.
68    Edit(OwnedEventId, IndexableEvent),
69    /// Do nothing.
70    Noop,
71}
72
73/// A struct that holds all data pertaining to a particular room's
74/// message index.
75pub struct RoomIndex {
76    index: Index,
77    schema: RoomMessageSchema,
78    query_parser: QueryParser,
79    room_id: OwnedRoomId,
80    /// Events added but not yet committed, mapping each document's primary key
81    /// (event id) to its deletion key (original event id). The deletion key is
82    /// needed so that a [`RoomIndex::remove`] in the same uncommitted batch,
83    /// which deletes by that key, can reconcile these entries too.
84    uncommitted_adds: HashMap<OwnedEventId, OwnedEventId>,
85    uncommitted_removes: HashSet<OwnedEventId>,
86    /// Cached [`IndexReader`].
87    ///
88    /// It is costly to create one, so let's keep it in memory when needed; see
89    /// [`RoomIndex::reader`] to learn more.
90    reader: OnceCell<IndexReader>,
91}
92
93impl fmt::Debug for RoomIndex {
94    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95        f.debug_struct("RoomIndex")
96            .field("schema", &self.schema)
97            .field("room_id", &self.room_id)
98            .finish()
99    }
100}
101
102impl RoomIndex {
103    pub(crate) fn new_with(index: Index, schema: RoomMessageSchema, room_id: &RoomId) -> RoomIndex {
104        let query_parser = QueryParser::for_index(&index, schema.default_search_fields());
105        Self {
106            index,
107            schema,
108            query_parser,
109            room_id: room_id.to_owned(),
110            uncommitted_adds: HashMap::new(),
111            uncommitted_removes: HashSet::new(),
112            reader: OnceCell::new(),
113        }
114    }
115
116    /// Get a [`SearchIndexWriter`] for this index.
117    fn writer(&self) -> Result<SearchIndexWriter, IndexError> {
118        let writer = self.index.writer(TANTIVY_INDEX_MEMORY_BUDGET)?;
119        Ok(SearchIndexWriter::new(writer, self.schema.clone()))
120    }
121
122    /// Get or create the cached [`IndexReader`] for this index.
123    fn reader(&self) -> Result<&IndexReader, IndexError> {
124        self.reader.get_or_try_init(|| Ok(self.index.reader_builder().try_into()?))
125    }
126
127    /// Commit added events to [`RoomIndex`]. The changes are not reflected in
128    /// the search results until the serchers are reloaded.
129    ///
130    /// Use [`RoomIndex::commit_and_reload`] for this purpose.
131    fn commit(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
132        let last_commit_opstamp = writer.commit()?; // TODO: This is blocking. Handle it.
133        self.uncommitted_adds.clear();
134        self.uncommitted_removes.clear();
135        Ok(last_commit_opstamp)
136    }
137
138    /// Commit added events to [`RoomIndex`] and
139    /// update searchers so that they reflect the state of the last
140    /// `.commit()`.
141    ///
142    /// Every commit should be rapidly reflected on your `IndexReader` and you
143    /// should not need to call `reload()` at all.
144    ///
145    /// This automatic reload can take 10s of milliseconds to kick in however,
146    /// and in unit tests it can be nice to deterministically force the
147    /// reload of searchers.
148    fn commit_and_reload(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
149        debug!(
150            "RoomIndex: committing and reloading: uncommitted: {:?}, {:?}",
151            self.uncommitted_adds, self.uncommitted_removes
152        );
153        let last_commit_opstamp = self.commit(writer)?;
154        self.reader()?.reload()?;
155        Ok(last_commit_opstamp)
156    }
157
158    /// Search the [`RoomIndex`] for some query. Returns a list of
159    /// results with a maximum given length. If `pagination_offset` is
160    /// set then the results will start there, i.e.
161    ///
162    /// if `max_number_of_results = 3` and `pagination_offset = 10`
163    /// (and there are a surplus of results)
164    /// then this will return results `11, 12, 13`
165    pub fn search(
166        &self,
167        query: &str,
168        max_number_of_results: usize,
169        pagination_offset: Option<usize>,
170    ) -> Result<Vec<(f32, OwnedEventId)>, IndexError> {
171        let query = self.query_parser.parse_query(query)?;
172        let searcher = self.reader()?.searcher();
173
174        let offset = pagination_offset.unwrap_or(0);
175
176        let results = searcher.search(
177            &query,
178            &TopDocs::with_limit(max_number_of_results).and_offset(offset).order_by_score(),
179        )?;
180        let mut ret: Vec<(f32, OwnedEventId)> = Vec::new();
181        let pk = self.schema.primary_key();
182
183        for (score, doc_address) in results {
184            let retrieved_doc: TantivyDocument = searcher.doc(doc_address)?;
185            match retrieved_doc.get_first(pk).and_then(|maybe_value| maybe_value.as_str()) {
186                Some(value) => match OwnedEventId::try_from(value) {
187                    Ok(event_id) => ret.push((score, event_id)),
188                    Err(err) => error!("error while parsing event_id from search result: {err:?}"),
189                },
190                _ => error!("unexpected value type while searching documents"),
191            }
192        }
193
194        Ok(ret)
195    }
196
197    fn events_to_be_removed(&self, event_id: &EventId) -> Result<Vec<OwnedEventId>, IndexError> {
198        Ok(self
199            .search(
200                format!(
201                    "{}:\"{event_id}\"",
202                    self.schema.get_field_name(self.schema.deletion_key())
203                )
204                .as_str(),
205                10000,
206                None,
207            )?
208            .into_iter()
209            .map(|(_, id)| id)
210            .collect())
211    }
212
213    fn add(
214        &mut self,
215        writer: &mut SearchIndexWriter,
216        event: IndexableEvent,
217    ) -> Result<(), IndexError> {
218        if !self.contains(&event.event_id) {
219            writer.add(self.schema.make_doc(event.clone())?)?;
220        }
221        self.uncommitted_removes.remove(&event.event_id);
222        self.uncommitted_adds.insert(event.event_id, event.original_event_id);
223        Ok(())
224    }
225
226    fn remove(
227        &mut self,
228        writer: &mut SearchIndexWriter,
229        event_id: OwnedEventId,
230    ) -> Result<(), IndexError> {
231        let events = self.events_to_be_removed(&event_id)?;
232
233        writer.remove(&event_id);
234
235        // Committed documents matching the deletion key.
236        for event in events.into_iter() {
237            self.uncommitted_adds.remove(&event);
238            self.uncommitted_removes.insert(event);
239        }
240
241        // Uncommitted documents added in this same batch also get deleted by the
242        // term above, so reconcile them too. Otherwise `contains` would still
243        // report them as present and a subsequent re-add (e.g. from an edit)
244        // would be wrongly skipped, leaving the document deleted.
245        let uncommitted: Vec<_> = self
246            .uncommitted_adds
247            .iter()
248            .filter(|(_, deletion_key)| **deletion_key == event_id)
249            .map(|(primary_key, _)| primary_key.clone())
250            .collect();
251        for event in uncommitted {
252            self.uncommitted_adds.remove(&event);
253            self.uncommitted_removes.insert(event);
254        }
255
256        Ok(())
257    }
258
259    fn execute_impl(
260        &mut self,
261        writer: &mut SearchIndexWriter,
262        operation: &RoomIndexOperation,
263    ) -> Result<(), IndexError> {
264        debug!("INDEX: executing {operation:?}");
265        match operation.clone() {
266            RoomIndexOperation::Add(event) => {
267                self.add(writer, event)?;
268            }
269            RoomIndexOperation::Remove(event_id) => {
270                self.remove(writer, event_id)?;
271            }
272            RoomIndexOperation::Edit(remove_event_id, event) => {
273                self.remove(writer, remove_event_id)?;
274                self.add(writer, event)?;
275            }
276            RoomIndexOperation::Noop => {}
277        }
278        Ok(())
279    }
280
281    /// Execute [`RoomIndexOperation`] with retry
282    fn execute_with_retry(
283        &mut self,
284        writer: &mut SearchIndexWriter,
285        operation: &RoomIndexOperation,
286        retries: usize,
287    ) -> Result<(), IndexError> {
288        let mut num_tries = 0;
289
290        while let Err(err) = self.execute_impl(writer, operation) {
291            if num_tries == retries {
292                return Err(err);
293            }
294            match err {
295                // Retry
296                IndexError::TantivyError(_)
297                | IndexError::IndexSchemaError(_)
298                | IndexError::IndexWriteError(_)
299                | IndexError::IO(_) => {
300                    num_tries += 1;
301                }
302                IndexError::OpenDirectoryError(ref e) => match e {
303                    // Retry
304                    OpenDirectoryError::IoError { io_error: _, directory_path: _ } => {
305                        num_tries += 1;
306                    }
307                    // Bubble
308                    OpenDirectoryError::DoesNotExist(_)
309                    | OpenDirectoryError::FailedToCreateTempDir(_)
310                    | OpenDirectoryError::NotADirectory(_) => return Err(err),
311                },
312                // Bubble
313                IndexError::QueryParserError(_) => return Err(err),
314                // Ignore
315                IndexError::CannotIndexRedactedMessage
316                | IndexError::EmptyMessage
317                | IndexError::MessageTypeNotSupported => break,
318            }
319            debug!("Failed to execute operation in room index (try {num_tries}): {err}");
320        }
321        Ok(())
322    }
323
324    /// Execute [`RoomIndexOperation`]
325    ///
326    /// If an error occurs, retry 5 times if possible.
327    ///
328    /// This which will add/remove/edit an event in the index based on the
329    /// operation.
330    ///
331    /// Prefer [`RoomIndex::bulk_execute`] for multiple operations.
332    pub fn execute(&mut self, operation: RoomIndexOperation) -> Result<(), IndexError> {
333        let mut writer = self.writer()?;
334        self.execute_with_retry(&mut writer, &operation, 5)?;
335        self.commit_and_reload(&mut writer)?;
336        Ok(())
337    }
338
339    /// Bulk execute [`RoomIndexOperation`]s
340    ///
341    /// If an error occurs in the batch it retries 5 times if possible.
342    ///
343    /// This which will add/remove/edit an events in the index based on the
344    /// operations.
345    pub fn bulk_execute(&mut self, operations: Vec<RoomIndexOperation>) -> Result<(), IndexError> {
346        let mut writer = self.writer()?;
347        let mut operations = operations.into_iter();
348        let mut next_operation = operations.next();
349
350        while let Some(ref operation) = next_operation {
351            self.execute_with_retry(&mut writer, operation, 5)?;
352            next_operation = operations.next();
353        }
354
355        self.commit_and_reload(&mut writer)?;
356
357        Ok(())
358    }
359
360    fn contains(&self, event_id: &EventId) -> bool {
361        let search_result = self.search(
362            format!("{}:\"{event_id}\"", self.schema.get_field_name(self.schema.primary_key()))
363                .as_str(),
364            1,
365            None,
366        );
367        match search_result {
368            Ok(results) => {
369                !self.uncommitted_removes.contains(event_id)
370                    && (!results.is_empty() || self.uncommitted_adds.contains_key(event_id))
371            }
372            Err(err) => {
373                warn!("Failed to check if event has been indexed, assuming it has: {err}");
374                true
375            }
376        }
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use std::{collections::HashSet, error::Error};
383
384    use matrix_sdk_test::event_factory::EventFactory;
385    use ruma::{
386        EventId, event_id,
387        events::{
388            AnySyncMessageLikeEvent,
389            room::message::{
390                MessageType, OriginalSyncRoomMessageEvent, Relation,
391                RoomMessageEventContentWithoutRelation,
392            },
393        },
394        room_id, user_id,
395    };
396
397    use crate::{
398        error::IndexError,
399        index::{IndexableEvent, RoomIndex, RoomIndexOperation, builder::RoomIndexBuilder},
400    };
401
402    /// Build an [`IndexableEvent`] from a text room message (tests only handle
403    /// text).
404    fn to_indexable(event: &OriginalSyncRoomMessageEvent) -> IndexableEvent {
405        let MessageType::Text(content) = &event.content.msgtype else {
406            panic!("test helper only supports text messages")
407        };
408        let original_event_id = match &event.content.relates_to {
409            Some(Relation::Replacement(replacement)) => replacement.event_id.clone(),
410            _ => event.event_id.clone(),
411        };
412        IndexableEvent {
413            event_id: event.event_id.clone(),
414            original_event_id,
415            sender: event.sender.clone(),
416            timestamp: event.origin_server_ts,
417            body: content.body.clone(),
418        }
419    }
420
421    /// Helper function to add a regular message to the index
422    ///
423    /// # Panic
424    /// Panics when event is not a [`OriginalSyncRoomMessageEvent`] with no
425    /// relations.
426    fn index_message(
427        index: &mut RoomIndex,
428        event: AnySyncMessageLikeEvent,
429    ) -> Result<(), IndexError> {
430        if let AnySyncMessageLikeEvent::RoomMessage(ev) = event
431            && let Some(ev) = ev.as_original()
432            && ev.content.relates_to.is_none()
433        {
434            return index.execute(RoomIndexOperation::Add(to_indexable(ev)));
435        }
436        panic!("Event was not a relationless OriginalSyncRoomMessageEvent.")
437    }
438
439    /// Helper function to remove events to the index
440    fn index_remove(index: &mut RoomIndex, event_id: &EventId) -> Result<(), IndexError> {
441        index.execute(RoomIndexOperation::Remove(event_id.to_owned()))
442    }
443
444    /// Helper function to edit events in index
445    ///
446    /// Edit event with `event_id` into new [`OriginalSyncRoomMessageEvent`]
447    fn index_edit(
448        index: &mut RoomIndex,
449        event_id: &EventId,
450        new: OriginalSyncRoomMessageEvent,
451    ) -> Result<(), IndexError> {
452        index.execute(RoomIndexOperation::Edit(event_id.to_owned(), to_indexable(&new)))
453    }
454
455    #[test]
456    fn test_add_event() {
457        let room_id = room_id!("!room_id:localhost");
458        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
459
460        let event = EventFactory::new()
461            .text_msg("event message")
462            .event_id(event_id!("$event_id:localhost"))
463            .room(room_id)
464            .sender(user_id!("@user_id:localhost"))
465            .into_any_sync_message_like_event();
466
467        index_message(&mut index, event).expect("failed to add event: {res:?}");
468    }
469
470    #[test]
471    fn test_search_populated_index() -> Result<(), Box<dyn Error>> {
472        let room_id = room_id!("!room_id:localhost");
473        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
474
475        let event_id_1 = event_id!("$event_id_1:localhost");
476        let event_id_2 = event_id!("$event_id_2:localhost");
477        let event_id_3 = event_id!("$event_id_3:localhost");
478        let user_id = user_id!("@user_id:localhost");
479        let f = EventFactory::new().room(room_id).sender(user_id);
480
481        index_message(
482            &mut index,
483            f.text_msg("This is a sentence")
484                .event_id(event_id_1)
485                .into_any_sync_message_like_event(),
486        )?;
487
488        index_message(
489            &mut index,
490            f.text_msg("All new words").event_id(event_id_2).into_any_sync_message_like_event(),
491        )?;
492
493        index_message(
494            &mut index,
495            f.text_msg("A similar sentence")
496                .event_id(event_id_3)
497                .into_any_sync_message_like_event(),
498        )?;
499
500        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
501        let result: HashSet<_> = result.iter().map(|(_, id)| id).collect();
502
503        let true_value = [event_id_1.to_owned(), event_id_3.to_owned()];
504        let true_value: HashSet<_> = true_value.iter().collect();
505
506        assert_eq!(result, true_value, "search result not correct: {result:?}");
507
508        Ok(())
509    }
510
511    #[test]
512    fn test_search_empty_index() -> Result<(), Box<dyn Error>> {
513        let room_id = room_id!("!room_id:localhost");
514        let index = RoomIndexBuilder::new_in_memory(room_id).build();
515
516        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
517
518        assert!(result.is_empty(), "search result not empty: {result:?}");
519
520        Ok(())
521    }
522
523    #[test]
524    fn test_index_contains_false() {
525        let room_id = room_id!("!room_id:localhost");
526        let index = RoomIndexBuilder::new_in_memory(room_id).build();
527
528        let event_id = event_id!("$event_id:localhost");
529
530        assert!(!index.contains(event_id), "Index should not contain event");
531    }
532
533    #[test]
534    fn test_index_contains_true() -> Result<(), Box<dyn Error>> {
535        let room_id = room_id!("!room_id:localhost");
536        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
537
538        let event_id = event_id!("$event_id:localhost");
539        let event = EventFactory::new()
540            .text_msg("This is a sentence")
541            .event_id(event_id)
542            .room(room_id)
543            .sender(user_id!("@user_id:localhost"))
544            .into_any_sync_message_like_event();
545
546        index_message(&mut index, event)?;
547
548        assert!(index.contains(event_id), "Index should contain event");
549
550        Ok(())
551    }
552
553    #[test]
554    fn test_index_add_idempotency() -> Result<(), Box<dyn Error>> {
555        let room_id = room_id!("!room_id:localhost");
556        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
557
558        let event_id = event_id!("$event_id:localhost");
559        let event = EventFactory::new()
560            .text_msg("This is a sentence")
561            .event_id(event_id)
562            .room(room_id)
563            .sender(user_id!("@user_id:localhost"))
564            .into_any_sync_message_like_event();
565
566        index_message(&mut index, event.clone())?;
567
568        assert!(index.contains(event_id), "Index should contain event");
569
570        // indexing again should do nothing
571        index_message(&mut index, event)?;
572
573        assert!(index.contains(event_id), "Index should still contain event");
574
575        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
576
577        assert_eq!(result.len(), 1, "Index should have ignored second indexing");
578
579        Ok(())
580    }
581
582    #[test]
583    fn test_remove_event() -> Result<(), Box<dyn Error>> {
584        let room_id = room_id!("!room_id:localhost");
585        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
586
587        let event_id = event_id!("$event_id:localhost");
588        let user_id = user_id!("@user_id:localhost");
589        let f = EventFactory::new().room(room_id).sender(user_id);
590
591        let event =
592            f.text_msg("This is a sentence").event_id(event_id).into_any_sync_message_like_event();
593
594        index_message(&mut index, event)?;
595
596        assert!(index.contains(event_id), "Index should contain event");
597
598        index_remove(&mut index, event_id)?;
599
600        assert!(!index.contains(event_id), "Index should not contain event");
601
602        Ok(())
603    }
604
605    #[test]
606    fn test_edit_removes_old_and_adds_new_event() -> Result<(), Box<dyn Error>> {
607        let room_id = room_id!("!room_id:localhost");
608        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
609
610        let old_event_id = event_id!("$old_event_id:localhost");
611        let user_id = user_id!("@user_id:localhost");
612        let f = EventFactory::new().room(room_id).sender(user_id);
613
614        let old_event = f
615            .text_msg("This is a sentence")
616            .event_id(old_event_id)
617            .into_any_sync_message_like_event();
618
619        index_message(&mut index, old_event)?;
620
621        assert!(index.contains(old_event_id), "Index should contain event");
622
623        let new_event_id = event_id!("$new_event_id:localhost");
624        let edit = f
625            .text_msg("This is a brand new sentence!")
626            .edit(
627                old_event_id,
628                RoomMessageEventContentWithoutRelation::text_plain("This is a brand new sentence!"),
629            )
630            .event_id(new_event_id)
631            .into_original_sync_room_message_event();
632
633        index_edit(&mut index, old_event_id, edit)?;
634
635        assert!(!index.contains(old_event_id), "Index should not contain old event");
636        assert!(index.contains(new_event_id), "Index should contain edited event");
637
638        Ok(())
639    }
640
641    #[test]
642    fn test_bulk_add_then_edit_same_event_keeps_it_indexed() -> Result<(), Box<dyn Error>> {
643        let room_id = room_id!("!room_id:localhost");
644        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
645
646        let original_id = event_id!("$original:localhost");
647        let edit_id = event_id!("$edit:localhost");
648        let user_id = user_id!("@user_id:localhost");
649        let f = EventFactory::new().room(room_id).sender(user_id);
650
651        let edit = f
652            .text_msg("* brand new sentence")
653            .edit(
654                original_id,
655                RoomMessageEventContentWithoutRelation::text_plain("brand new sentence"),
656            )
657            .event_id(edit_id)
658            .into_original_sync_room_message_event();
659        let edit = to_indexable(&edit);
660
661        // An original and its edit arriving in the same batch produce an `Add`
662        // and an `Edit` of the same document. The `Edit`'s removal must not drop
663        // the document added earlier in the same uncommitted batch.
664        index.bulk_execute(vec![
665            RoomIndexOperation::Add(edit.clone()),
666            RoomIndexOperation::Edit(original_id.to_owned(), edit),
667        ])?;
668
669        assert!(index.contains(edit_id), "Edited document should be indexed");
670
671        let result = index.search("sentence", 10, None)?;
672        assert_eq!(result.len(), 1, "Search should find the edited document, got {result:?}");
673        assert_eq!(result[0].1, edit_id, "unexpected event id: {result:?}");
674
675        Ok(())
676    }
677}
matrix_sdk_search/index/mod.rs

matrix_sdk_search/index/
mod.rs
