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::{collections::HashSet, fmt};
19
20use ruma::{
21    EventId, OwnedEventId, OwnedRoomId, RoomId, events::room::message::OriginalSyncRoomMessageEvent,
22};
23use tantivy::{
24    Index, IndexReader, TantivyDocument, collector::TopDocs, directory::error::OpenDirectoryError,
25    query::QueryParser, schema::Value,
26};
27use tracing::{debug, error, warn};
28
29use crate::{
30    OpStamp, TANTIVY_INDEX_MEMORY_BUDGET,
31    error::IndexError,
32    schema::{MatrixSearchIndexSchema, RoomMessageSchema},
33    writer::SearchIndexWriter,
34};
35
36/// A struct to represent the operations on a [`RoomIndex`]
37#[derive(Debug, Clone)]
38pub enum RoomIndexOperation {
39    /// Add this event to the index.
40    Add(OriginalSyncRoomMessageEvent),
41    /// Remove all documents in the index where
42    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id.
43    Remove(OwnedEventId),
44    /// Replace all documents in the index where
45    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id with
46    /// the new event.
47    Edit(OwnedEventId, OriginalSyncRoomMessageEvent),
48    /// Do nothing.
49    Noop,
50}
51
52/// A struct that holds all data pertaining to a particular room's
53/// message index.
54pub struct RoomIndex {
55    index: Index,
56    schema: RoomMessageSchema,
57    query_parser: QueryParser,
58    room_id: OwnedRoomId,
59    uncommitted_adds: HashSet<OwnedEventId>,
60    uncommitted_removes: HashSet<OwnedEventId>,
61    reader: Option<IndexReader>,
62}
63
64impl fmt::Debug for RoomIndex {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        f.debug_struct("RoomIndex")
67            .field("schema", &self.schema)
68            .field("room_id", &self.room_id)
69            .finish()
70    }
71}
72
73impl RoomIndex {
74    pub(crate) fn new_with(index: Index, schema: RoomMessageSchema, room_id: &RoomId) -> RoomIndex {
75        let query_parser = QueryParser::for_index(&index, schema.default_search_fields());
76        Self {
77            index,
78            schema,
79            query_parser,
80            room_id: room_id.to_owned(),
81            uncommitted_adds: HashSet::new(),
82            uncommitted_removes: HashSet::new(),
83            reader: None,
84        }
85    }
86
87    /// Get a [`SearchIndexWriter`] for this index.
88    fn get_writer(&self) -> Result<SearchIndexWriter, IndexError> {
89        let writer = self.index.writer(TANTIVY_INDEX_MEMORY_BUDGET)?;
90        Ok(SearchIndexWriter::new(writer, self.schema.clone()))
91    }
92
93    /// Get or create the cached [`IndexReader`] for this index.
94    fn get_reader(&mut self) -> Result<&IndexReader, IndexError> {
95        Ok(match self.reader {
96            Some(ref reader) => reader,
97            None => {
98                let reader = self.index.reader_builder().try_into()?;
99                self.reader.insert(reader)
100            }
101        })
102    }
103
104    /// Commit added events to [`RoomIndex`]. The changes are not reflected in
105    /// the search results until the serchers are reloaded.
106    ///
107    /// Use [`RoomIndex::commit_and_reload`] for this purpose.
108    fn commit(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
109        let last_commit_opstamp = writer.commit()?; // TODO: This is blocking. Handle it.
110        self.uncommitted_adds.clear();
111        self.uncommitted_removes.clear();
112        Ok(last_commit_opstamp)
113    }
114
115    /// Commit added events to [`RoomIndex`] and
116    /// update searchers so that they reflect the state of the last
117    /// `.commit()`.
118    ///
119    /// Every commit should be rapidly reflected on your `IndexReader` and you
120    /// should not need to call `reload()` at all.
121    ///
122    /// This automatic reload can take 10s of milliseconds to kick in however,
123    /// and in unit tests it can be nice to deterministically force the
124    /// reload of searchers.
125    fn commit_and_reload(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
126        debug!(
127            "RoomIndex: committing and reloading: uncommitted: {:?}, {:?}",
128            self.uncommitted_adds, self.uncommitted_removes
129        );
130        let last_commit_opstamp = self.commit(writer)?;
131        self.get_reader()?.reload()?;
132        Ok(last_commit_opstamp)
133    }
134
135    /// Search the [`RoomIndex`] for some query. Returns a list of
136    /// results with a maximum given length. If `pagination_offset` is
137    /// set then the results will start there, i.e.
138    ///
139    /// if `max_number_of_results = 3` and `pagination_offset = 10`
140    /// (and there are a surplus of results)
141    /// then this will return results `11, 12, 13`
142    pub fn search(
143        &mut self,
144        query: &str,
145        max_number_of_results: usize,
146        pagination_offset: Option<usize>,
147    ) -> Result<Vec<(f32, OwnedEventId)>, IndexError> {
148        let query = self.query_parser.parse_query(query)?;
149        let searcher = self.get_reader()?.searcher();
150
151        let offset = pagination_offset.unwrap_or(0);
152
153        let results = searcher.search(
154            &query,
155            &TopDocs::with_limit(max_number_of_results).and_offset(offset).order_by_score(),
156        )?;
157        let mut ret: Vec<(f32, OwnedEventId)> = Vec::new();
158        let pk = self.schema.primary_key();
159
160        for (score, doc_address) in results {
161            let retrieved_doc: TantivyDocument = searcher.doc(doc_address)?;
162            match retrieved_doc.get_first(pk).and_then(|maybe_value| maybe_value.as_str()) {
163                Some(value) => match OwnedEventId::try_from(value) {
164                    Ok(event_id) => ret.push((score, event_id)),
165                    Err(err) => error!("error while parsing event_id from search result: {err:?}"),
166                },
167                _ => error!("unexpected value type while searching documents"),
168            }
169        }
170
171        Ok(ret)
172    }
173
174    fn get_events_to_be_removed(
175        &mut self,
176        event_id: &EventId,
177    ) -> Result<Vec<OwnedEventId>, IndexError> {
178        Ok(self
179            .search(
180                format!(
181                    "{}:\"{event_id}\"",
182                    self.schema.get_field_name(self.schema.deletion_key())
183                )
184                .as_str(),
185                10000,
186                None,
187            )?
188            .into_iter()
189            .map(|(_, id)| id)
190            .collect())
191    }
192
193    fn add(
194        &mut self,
195        writer: &mut SearchIndexWriter,
196        event: OriginalSyncRoomMessageEvent,
197    ) -> Result<(), IndexError> {
198        if !self.contains(&event.event_id) {
199            writer.add(self.schema.make_doc(event.clone())?)?;
200        }
201        self.uncommitted_removes.remove(&event.event_id);
202        self.uncommitted_adds.insert(event.event_id);
203        Ok(())
204    }
205
206    fn remove(
207        &mut self,
208        writer: &mut SearchIndexWriter,
209        event_id: OwnedEventId,
210    ) -> Result<(), IndexError> {
211        let events = self.get_events_to_be_removed(&event_id)?;
212
213        writer.remove(&event_id);
214
215        for event in events.into_iter() {
216            self.uncommitted_adds.remove(&event);
217            self.uncommitted_removes.insert(event);
218        }
219
220        Ok(())
221    }
222
223    fn execute_impl(
224        &mut self,
225        writer: &mut SearchIndexWriter,
226        operation: &RoomIndexOperation,
227    ) -> Result<(), IndexError> {
228        debug!("INDEX: executing {operation:?}");
229        match operation.clone() {
230            RoomIndexOperation::Add(event) => {
231                self.add(writer, event)?;
232            }
233            RoomIndexOperation::Remove(event_id) => {
234                self.remove(writer, event_id)?;
235            }
236            RoomIndexOperation::Edit(remove_event_id, event) => {
237                self.remove(writer, remove_event_id)?;
238                self.add(writer, event)?;
239            }
240            RoomIndexOperation::Noop => {}
241        }
242        Ok(())
243    }
244
245    /// Execute [`RoomIndexOperation`] with retry
246    fn execute_with_retry(
247        &mut self,
248        writer: &mut SearchIndexWriter,
249        operation: &RoomIndexOperation,
250        retries: usize,
251    ) -> Result<(), IndexError> {
252        let mut num_tries = 0;
253
254        while let Err(err) = self.execute_impl(writer, operation) {
255            if num_tries == retries {
256                return Err(err);
257            }
258            match err {
259                // Retry
260                IndexError::TantivyError(_)
261                | IndexError::IndexSchemaError(_)
262                | IndexError::IndexWriteError(_)
263                | IndexError::IO(_) => {
264                    num_tries += 1;
265                }
266                IndexError::OpenDirectoryError(ref e) => match e {
267                    // Retry
268                    OpenDirectoryError::IoError { io_error: _, directory_path: _ } => {
269                        num_tries += 1;
270                    }
271                    // Bubble
272                    OpenDirectoryError::DoesNotExist(_)
273                    | OpenDirectoryError::FailedToCreateTempDir(_)
274                    | OpenDirectoryError::NotADirectory(_) => return Err(err),
275                },
276                // Bubble
277                IndexError::QueryParserError(_) => return Err(err),
278                // Ignore
279                IndexError::CannotIndexRedactedMessage
280                | IndexError::EmptyMessage
281                | IndexError::MessageTypeNotSupported => break,
282            }
283            debug!("Failed to execute operation in room index (try {num_tries}): {err}");
284        }
285        Ok(())
286    }
287
288    /// Execute [`RoomIndexOperation`]
289    ///
290    /// If an error occurs, retry 5 times if possible.
291    ///
292    /// This which will add/remove/edit an event in the index based on the
293    /// operation.
294    ///
295    /// Prefer [`RoomIndex::bulk_execute`] for multiple operations.
296    pub fn execute(&mut self, operation: RoomIndexOperation) -> Result<(), IndexError> {
297        let mut writer = self.get_writer()?;
298        self.execute_with_retry(&mut writer, &operation, 5)?;
299        self.commit_and_reload(&mut writer)?;
300        Ok(())
301    }
302
303    /// Bulk execute [`RoomIndexOperation`]s
304    ///
305    /// If an error occurs in the batch it retries 5 times if possible.
306    ///
307    /// This which will add/remove/edit an events in the index based on the
308    /// operations.
309    pub fn bulk_execute(&mut self, operations: Vec<RoomIndexOperation>) -> Result<(), IndexError> {
310        let mut writer = self.get_writer()?;
311        let mut operations = operations.into_iter();
312        let mut next_operation = operations.next();
313
314        while let Some(ref operation) = next_operation {
315            self.execute_with_retry(&mut writer, operation, 5)?;
316            next_operation = operations.next();
317        }
318
319        self.commit_and_reload(&mut writer)?;
320
321        Ok(())
322    }
323
324    fn contains(&mut self, event_id: &EventId) -> bool {
325        let search_result = self.search(
326            format!("{}:\"{event_id}\"", self.schema.get_field_name(self.schema.primary_key()))
327                .as_str(),
328            1,
329            None,
330        );
331        match search_result {
332            Ok(results) => {
333                !self.uncommitted_removes.contains(event_id)
334                    && (!results.is_empty() || self.uncommitted_adds.contains(event_id))
335            }
336            Err(err) => {
337                warn!("Failed to check if event has been indexed, assuming it has: {err}");
338                true
339            }
340        }
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::{collections::HashSet, error::Error};
347
348    use matrix_sdk_test::event_factory::EventFactory;
349    use ruma::{
350        EventId, event_id,
351        events::{
352            AnySyncMessageLikeEvent,
353            room::message::{OriginalSyncRoomMessageEvent, RoomMessageEventContentWithoutRelation},
354        },
355        room_id, user_id,
356    };
357
358    use crate::{
359        error::IndexError,
360        index::{RoomIndex, RoomIndexOperation, builder::RoomIndexBuilder},
361    };
362
363    /// Helper function to add a regular message to the index
364    ///
365    /// # Panic
366    /// Panics when event is not a [`OriginalSyncRoomMessageEvent`] with no
367    /// relations.
368    fn index_message(
369        index: &mut RoomIndex,
370        event: AnySyncMessageLikeEvent,
371    ) -> Result<(), IndexError> {
372        if let AnySyncMessageLikeEvent::RoomMessage(ev) = event
373            && let Some(ev) = ev.as_original()
374            && ev.content.relates_to.is_none()
375        {
376            return index.execute(RoomIndexOperation::Add(ev.clone()));
377        }
378        panic!("Event was not a relationless OriginalSyncRoomMessageEvent.")
379    }
380
381    /// Helper function to remove events to the index
382    fn index_remove(index: &mut RoomIndex, event_id: &EventId) -> Result<(), IndexError> {
383        index.execute(RoomIndexOperation::Remove(event_id.to_owned()))
384    }
385
386    /// Helper function to edit events in index
387    ///
388    /// Edit event with `event_id` into new [`OriginalSyncRoomMessageEvent`]
389    fn index_edit(
390        index: &mut RoomIndex,
391        event_id: &EventId,
392        new: OriginalSyncRoomMessageEvent,
393    ) -> Result<(), IndexError> {
394        index.execute(RoomIndexOperation::Edit(event_id.to_owned(), new))
395    }
396
397    #[test]
398    fn test_add_event() {
399        let room_id = room_id!("!room_id:localhost");
400        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
401
402        let event = EventFactory::new()
403            .text_msg("event message")
404            .event_id(event_id!("$event_id:localhost"))
405            .room(room_id)
406            .sender(user_id!("@user_id:localhost"))
407            .into_any_sync_message_like_event();
408
409        index_message(&mut index, event).expect("failed to add event: {res:?}");
410    }
411
412    #[test]
413    fn test_search_populated_index() -> Result<(), Box<dyn Error>> {
414        let room_id = room_id!("!room_id:localhost");
415        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
416
417        let event_id_1 = event_id!("$event_id_1:localhost");
418        let event_id_2 = event_id!("$event_id_2:localhost");
419        let event_id_3 = event_id!("$event_id_3:localhost");
420        let user_id = user_id!("@user_id:localhost");
421        let f = EventFactory::new().room(room_id).sender(user_id);
422
423        index_message(
424            &mut index,
425            f.text_msg("This is a sentence")
426                .event_id(event_id_1)
427                .into_any_sync_message_like_event(),
428        )?;
429
430        index_message(
431            &mut index,
432            f.text_msg("All new words").event_id(event_id_2).into_any_sync_message_like_event(),
433        )?;
434
435        index_message(
436            &mut index,
437            f.text_msg("A similar sentence")
438                .event_id(event_id_3)
439                .into_any_sync_message_like_event(),
440        )?;
441
442        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
443        let result: HashSet<_> = result.iter().map(|(_, id)| id).collect();
444
445        let true_value = [event_id_1.to_owned(), event_id_3.to_owned()];
446        let true_value: HashSet<_> = true_value.iter().collect();
447
448        assert_eq!(result, true_value, "search result not correct: {result:?}");
449
450        Ok(())
451    }
452
453    #[test]
454    fn test_search_empty_index() -> Result<(), Box<dyn Error>> {
455        let room_id = room_id!("!room_id:localhost");
456        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
457
458        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
459
460        assert!(result.is_empty(), "search result not empty: {result:?}");
461
462        Ok(())
463    }
464
465    #[test]
466    fn test_index_contains_false() {
467        let room_id = room_id!("!room_id:localhost");
468        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
469
470        let event_id = event_id!("$event_id:localhost");
471
472        assert!(!index.contains(event_id), "Index should not contain event");
473    }
474
475    #[test]
476    fn test_index_contains_true() -> Result<(), Box<dyn Error>> {
477        let room_id = room_id!("!room_id:localhost");
478        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
479
480        let event_id = event_id!("$event_id:localhost");
481        let event = EventFactory::new()
482            .text_msg("This is a sentence")
483            .event_id(event_id)
484            .room(room_id)
485            .sender(user_id!("@user_id:localhost"))
486            .into_any_sync_message_like_event();
487
488        index_message(&mut index, event)?;
489
490        assert!(index.contains(event_id), "Index should contain event");
491
492        Ok(())
493    }
494
495    #[test]
496    fn test_index_add_idempotency() -> Result<(), Box<dyn Error>> {
497        let room_id = room_id!("!room_id:localhost");
498        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
499
500        let event_id = event_id!("$event_id:localhost");
501        let event = EventFactory::new()
502            .text_msg("This is a sentence")
503            .event_id(event_id)
504            .room(room_id)
505            .sender(user_id!("@user_id:localhost"))
506            .into_any_sync_message_like_event();
507
508        index_message(&mut index, event.clone())?;
509
510        assert!(index.contains(event_id), "Index should contain event");
511
512        // indexing again should do nothing
513        index_message(&mut index, event)?;
514
515        assert!(index.contains(event_id), "Index should still contain event");
516
517        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
518
519        assert_eq!(result.len(), 1, "Index should have ignored second indexing");
520
521        Ok(())
522    }
523
524    #[test]
525    fn test_remove_event() -> Result<(), Box<dyn Error>> {
526        let room_id = room_id!("!room_id:localhost");
527        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
528
529        let event_id = event_id!("$event_id:localhost");
530        let user_id = user_id!("@user_id:localhost");
531        let f = EventFactory::new().room(room_id).sender(user_id);
532
533        let event =
534            f.text_msg("This is a sentence").event_id(event_id).into_any_sync_message_like_event();
535
536        index_message(&mut index, event)?;
537
538        assert!(index.contains(event_id), "Index should contain event");
539
540        index_remove(&mut index, event_id)?;
541
542        assert!(!index.contains(event_id), "Index should not contain event");
543
544        Ok(())
545    }
546
547    #[test]
548    fn test_edit_removes_old_and_adds_new_event() -> Result<(), Box<dyn Error>> {
549        let room_id = room_id!("!room_id:localhost");
550        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
551
552        let old_event_id = event_id!("$old_event_id:localhost");
553        let user_id = user_id!("@user_id:localhost");
554        let f = EventFactory::new().room(room_id).sender(user_id);
555
556        let old_event = f
557            .text_msg("This is a sentence")
558            .event_id(old_event_id)
559            .into_any_sync_message_like_event();
560
561        index_message(&mut index, old_event)?;
562
563        assert!(index.contains(old_event_id), "Index should contain event");
564
565        let new_event_id = event_id!("$new_event_id:localhost");
566        let edit = f
567            .text_msg("This is a brand new sentence!")
568            .edit(
569                old_event_id,
570                RoomMessageEventContentWithoutRelation::text_plain("This is a brand new sentence!"),
571            )
572            .event_id(new_event_id)
573            .into_original_sync_room_message_event();
574
575        index_edit(&mut index, old_event_id, edit)?;
576
577        assert!(!index.contains(old_event_id), "Index should not contain old event");
578        assert!(index.contains(new_event_id), "Index should contain edited event");
579
580        Ok(())
581    }
582}
matrix_sdk_search/index/mod.rs

matrix_sdk_search/index/
mod.rs
