Skip to main content

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::{
25    EventId, MilliSecondsSinceUnixEpoch, OwnedEventId, OwnedRoomId, OwnedUserId, RoomId, UInt,
26};
27use tantivy::{
28    Index, IndexReader, ReloadPolicy, TantivyDocument,
29    collector::{Count, MultiCollector, TopDocs},
30    directory::error::OpenDirectoryError,
31    query::QueryParser,
32    schema::Value,
33};
34use tracing::{debug, error, warn};
35
36use crate::{
37    OpStamp, TANTIVY_INDEX_MEMORY_BUDGET,
38    error::IndexError,
39    schema::{MatrixSearchIndexSchema, RoomMessageSchema},
40    writer::SearchIndexWriter,
41};
42
43/// The result type for a search operation.
44#[derive(Debug)]
45pub struct SearchResult {
46    /// The total count of search results that were found.
47    ///
48    /// Might be bigger than the length of [`SearchResult::events`] as the
49    /// search is paginated.
50    pub total_count: usize,
51
52    /// A list of search results.
53    ///
54    /// Each entry contains the search score and the ID of the event that
55    /// matched the search query. The list may be incomplete because the search
56    /// results are paginated; use [`SearchResult::total_count`] to determine
57    /// the total number of matching results.
58    pub events: Vec<(f32, OwnedEventId)>,
59}
60
61/// The subset of an event's data required to index it and later retrieve it.
62///
63/// Produced by the matrix-sdk layer, which knows how to extract searchable text
64/// from each event type. This crate stays agnostic to Matrix event content.
65#[derive(Clone)]
66pub struct IndexableEvent {
67    /// The event's own id (primary key).
68    pub(crate) event_id: OwnedEventId,
69    /// The id used as the deletion key: the original event id for edits,
70    /// otherwise the event's own id.
71    pub(crate) original_event_id: OwnedEventId,
72    /// The sender of the event.
73    pub(crate) sender: OwnedUserId,
74    /// The origin server timestamp of the event.
75    ///
76    /// Please use the `matrix_sdk_common::TimelineEvent::timestamp` as much as
77    /// possible as it protects against malformed `origin_server_ts`. At worst,
78    /// use the `matrix_sdk_common::serde_helpers::extract_timestamp` function.
79    pub(crate) timestamp: Option<MilliSecondsSinceUnixEpoch>,
80    /// The text to index for this event.
81    pub(crate) body: String,
82}
83
84impl fmt::Debug for IndexableEvent {
85    /// Don't log bodies
86    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
87        f.debug_struct("IndexableEvent")
88            .field("event_id", &self.event_id)
89            .field("original_event_id", &self.original_event_id)
90            .field("sender", &self.sender)
91            .field("timestamp", &self.timestamp)
92            .finish_non_exhaustive()
93    }
94}
95
96/// Maximum value for the timestamp to not overflow when converted to
97/// nanoseconds by Tantivy. See [`IndexableEvent::new`] to learn more.
98const MAX_MILLISECONDS: u64 = (i64::MAX / 1_000_000).cast_unsigned();
99
100impl IndexableEvent {
101    /// Create a new [`IndexableEvent`].
102    pub fn new(
103        event_id: OwnedEventId,
104        original_event_id: OwnedEventId,
105        sender: OwnedUserId,
106        mut timestamp: Option<MilliSecondsSinceUnixEpoch>,
107        body: String,
108    ) -> Self {
109        // Tantivy will transform the number of milliseconds to nanoseconds by
110        // multiplying by 1_000_000 ([1]). If the number of milliseconds is too
111        // big, the multiplication will overflow.
112        //
113        // To avoid this panic, we cap the number of milliseconds to a maximum
114        // value.
115        //
116        // [1]: https://github.com/quickwit-oss/tantivy/blob/31ca1a8ba290b425f871d2e2384592045ec01b8d/common/src/datetime.rs#L62-L67
117        if let Some(timestamp) = &mut timestamp {
118            *timestamp = MilliSecondsSinceUnixEpoch(
119                timestamp.get().min(UInt::new_saturating(MAX_MILLISECONDS)),
120            );
121        }
122
123        Self { event_id, original_event_id, sender, timestamp, body }
124    }
125}
126
127/// A struct to represent the operations on a [`RoomIndex`]
128#[derive(Debug, Clone)]
129pub enum RoomIndexOperation {
130    /// Add this event to the index.
131    Add(IndexableEvent),
132    /// Remove all documents in the index where
133    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id.
134    Remove(OwnedEventId),
135    /// Replace all documents in the index where
136    /// `MatrixSearchIndexSchema::deletion_key()` matches this event id with the
137    /// new event.
138    Edit(OwnedEventId, IndexableEvent),
139    /// Do nothing.
140    Noop,
141}
142
143/// A struct that holds all data pertaining to a particular room's message
144/// index.
145pub struct RoomIndex {
146    index: Index,
147    schema: RoomMessageSchema,
148    query_parser: QueryParser,
149    room_id: OwnedRoomId,
150    /// Events added but not yet committed, mapping each document's primary key
151    /// (event id) to its deletion key (original event id). The deletion key is
152    /// needed so that a [`RoomIndex::remove`] in the same uncommitted batch,
153    /// which deletes by that key, can reconcile these entries too.
154    uncommitted_adds: HashMap<OwnedEventId, OwnedEventId>,
155    uncommitted_removes: HashSet<OwnedEventId>,
156    /// Cached [`IndexReader`].
157    ///
158    /// It is costly to create one, so let's keep it in memory when needed; see
159    /// [`RoomIndex::reader`] to learn more.
160    reader: OnceCell<IndexReader>,
161}
162
163impl fmt::Debug for RoomIndex {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        f.debug_struct("RoomIndex")
166            .field("schema", &self.schema)
167            .field("room_id", &self.room_id)
168            .finish()
169    }
170}
171
172impl RoomIndex {
173    pub(crate) fn new_with(index: Index, schema: RoomMessageSchema, room_id: &RoomId) -> RoomIndex {
174        let query_parser = QueryParser::for_index(&index, schema.default_search_fields());
175        Self {
176            index,
177            schema,
178            query_parser,
179            room_id: room_id.to_owned(),
180            uncommitted_adds: HashMap::new(),
181            uncommitted_removes: HashSet::new(),
182            reader: OnceCell::new(),
183        }
184    }
185
186    /// Get a [`SearchIndexWriter`] for this index.
187    fn writer(&self) -> Result<SearchIndexWriter, IndexError> {
188        let writer = self.index.writer(TANTIVY_INDEX_MEMORY_BUDGET)?;
189        Ok(SearchIndexWriter::new(writer, self.schema.clone()))
190    }
191
192    /// Get or create the cached [`IndexReader`] for this index.
193    ///
194    /// The reload policy is [`ReloadPolicy::Manual`]: we are the only writer of
195    /// this index, and every commit goes through
196    /// [`RoomIndex::commit_and_reload`], which reloads explicitly. Tantivy's
197    /// default policy would instead spawn one meta file watcher thread per
198    /// index, i.e. one per room, and panics if that thread cannot be spawned.
199    fn reader(&self) -> Result<&IndexReader, IndexError> {
200        self.reader.get_or_try_init(|| {
201            Ok(self.index.reader_builder().reload_policy(ReloadPolicy::Manual).try_into()?)
202        })
203    }
204
205    /// Commit added events to [`RoomIndex`]. The changes are not reflected in
206    /// the search results until the serchers are reloaded.
207    ///
208    /// Use [`RoomIndex::commit_and_reload`] for this purpose.
209    fn commit(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
210        let last_commit_opstamp = writer.commit()?; // TODO: This is blocking. Handle it.
211        self.uncommitted_adds.clear();
212        self.uncommitted_removes.clear();
213        Ok(last_commit_opstamp)
214    }
215
216    /// Commit added events to [`RoomIndex`] and update searchers so that they
217    /// reflect the state of the last `.commit()`.
218    ///
219    /// Every commit should be rapidly reflected on your `IndexReader` and you
220    /// should not need to call `reload()` at all.
221    ///
222    /// This automatic reload can take 10s of milliseconds to kick in however,
223    /// and in unit tests it can be nice to deterministically force the reload
224    /// of searchers.
225    fn commit_and_reload(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
226        debug!(
227            "RoomIndex: committing and reloading: uncommitted: {:?}, {:?}",
228            self.uncommitted_adds, self.uncommitted_removes
229        );
230        let last_commit_opstamp = self.commit(writer)?;
231        self.reader()?.reload()?;
232        Ok(last_commit_opstamp)
233    }
234
235    /// Search the [`RoomIndex`] for some query. Returns a list of results with
236    /// a maximum given length. If `pagination_offset` is set then the results
237    /// will start there, i.e.
238    ///
239    /// if `max_number_of_results = 3` and `pagination_offset = 10` (and there
240    /// are a surplus of results) then this will return results `11, 12, 13`
241    pub fn search(
242        &self,
243        query: &str,
244        max_number_of_results: usize,
245        pagination_offset: Option<usize>,
246    ) -> Result<SearchResult, IndexError> {
247        let query = self.query_parser.parse_query(query)?;
248        let searcher = self.reader()?.searcher();
249
250        let offset = pagination_offset.unwrap_or(0);
251
252        let top_docs =
253            TopDocs::with_limit(max_number_of_results).and_offset(offset).order_by_score();
254
255        let mut multi = MultiCollector::new();
256        let count_handle = multi.add_collector(Count);
257        let top_docs_handle = multi.add_collector(top_docs);
258
259        let mut results = searcher.search(&query, &multi)?;
260        let mut ret: Vec<(f32, OwnedEventId)> = Vec::new();
261        let pk = self.schema.primary_key();
262
263        let result_count = count_handle.extract(&mut results);
264
265        for (score, doc_address) in top_docs_handle.extract(&mut results) {
266            let retrieved_doc: TantivyDocument = searcher.doc(doc_address)?;
267            match retrieved_doc.get_first(pk).and_then(|maybe_value| maybe_value.as_str()) {
268                Some(value) => match OwnedEventId::try_from(value) {
269                    Ok(event_id) => ret.push((score, event_id)),
270                    Err(err) => error!("error while parsing event_id from search result: {err:?}"),
271                },
272                _ => error!("unexpected value type while searching documents"),
273            }
274        }
275
276        Ok(SearchResult { total_count: result_count, events: ret })
277    }
278
279    fn events_to_be_removed(&self, event_id: &EventId) -> Result<Vec<OwnedEventId>, IndexError> {
280        Ok(self
281            .search(
282                format!(
283                    "{}:\"{event_id}\"",
284                    self.schema.get_field_name(self.schema.deletion_key())
285                )
286                .as_str(),
287                10000,
288                None,
289            )?
290            .events
291            .into_iter()
292            .map(|(_, id)| id)
293            .collect())
294    }
295
296    fn add(
297        &mut self,
298        writer: &mut SearchIndexWriter,
299        event: IndexableEvent,
300    ) -> Result<(), IndexError> {
301        if !self.contains(&event.event_id) {
302            writer.add(self.schema.make_doc(event.clone())?)?;
303        }
304        self.uncommitted_removes.remove(&event.event_id);
305        self.uncommitted_adds.insert(event.event_id, event.original_event_id);
306        Ok(())
307    }
308
309    fn remove(
310        &mut self,
311        writer: &mut SearchIndexWriter,
312        event_id: OwnedEventId,
313    ) -> Result<(), IndexError> {
314        let events = self.events_to_be_removed(&event_id)?;
315
316        writer.remove(&event_id);
317
318        // Committed documents matching the deletion key.
319        for event in events.into_iter() {
320            self.uncommitted_adds.remove(&event);
321            self.uncommitted_removes.insert(event);
322        }
323
324        // Uncommitted documents added in this same batch also get deleted by
325        // the term above, so reconcile them too. Otherwise `contains` would
326        // still report them as present and a subsequent re-add (e.g. from an
327        // edit) would be wrongly skipped, leaving the document deleted.
328        let uncommitted: Vec<_> = self
329            .uncommitted_adds
330            .iter()
331            .filter(|(_, deletion_key)| **deletion_key == event_id)
332            .map(|(primary_key, _)| primary_key.clone())
333            .collect();
334        for event in uncommitted {
335            self.uncommitted_adds.remove(&event);
336            self.uncommitted_removes.insert(event);
337        }
338
339        Ok(())
340    }
341
342    fn execute_impl(
343        &mut self,
344        writer: &mut SearchIndexWriter,
345        operation: &RoomIndexOperation,
346    ) -> Result<(), IndexError> {
347        debug!("INDEX: executing {operation:?}");
348        match operation.clone() {
349            RoomIndexOperation::Add(event) => {
350                self.add(writer, event)?;
351            }
352            RoomIndexOperation::Remove(event_id) => {
353                self.remove(writer, event_id)?;
354            }
355            RoomIndexOperation::Edit(remove_event_id, event) => {
356                self.remove(writer, remove_event_id)?;
357                self.add(writer, event)?;
358            }
359            RoomIndexOperation::Noop => {}
360        }
361        Ok(())
362    }
363
364    /// Execute [`RoomIndexOperation`] with retry
365    fn execute_with_retry(
366        &mut self,
367        writer: &mut SearchIndexWriter,
368        operation: &RoomIndexOperation,
369        retries: usize,
370    ) -> Result<(), IndexError> {
371        let mut num_tries = 0;
372
373        while let Err(err) = self.execute_impl(writer, operation) {
374            if num_tries == retries {
375                return Err(err);
376            }
377            match err {
378                // Retry
379                IndexError::TantivyError(_)
380                | IndexError::IndexSchemaError(_)
381                | IndexError::IndexWriteError(_)
382                | IndexError::IO(_) => {
383                    num_tries += 1;
384                }
385                IndexError::OpenDirectoryError(ref e) => match e {
386                    // Retry
387                    OpenDirectoryError::IoError { io_error: _, directory_path: _ } => {
388                        num_tries += 1;
389                    }
390                    // Bubble
391                    OpenDirectoryError::DoesNotExist(_)
392                    | OpenDirectoryError::FailedToCreateTempDir(_)
393                    | OpenDirectoryError::NotADirectory(_) => return Err(err),
394                },
395                // Bubble
396                IndexError::QueryParserError(_) => return Err(err),
397                // Ignore
398                IndexError::CannotIndexRedactedMessage
399                | IndexError::EmptyMessage
400                | IndexError::MessageTypeNotSupported => break,
401            }
402            debug!("Failed to execute operation in room index (try {num_tries}): {err}");
403        }
404        Ok(())
405    }
406
407    /// Execute [`RoomIndexOperation`]
408    ///
409    /// If an error occurs, retry 5 times if possible.
410    ///
411    /// This which will add/remove/edit an event in the index based on the
412    /// operation.
413    ///
414    /// Prefer [`RoomIndex::bulk_execute`] for multiple operations.
415    pub fn execute(&mut self, operation: RoomIndexOperation) -> Result<(), IndexError> {
416        let mut writer = self.writer()?;
417        self.execute_with_retry(&mut writer, &operation, 5)?;
418        self.commit_and_reload(&mut writer)?;
419        Ok(())
420    }
421
422    /// Bulk execute [`RoomIndexOperation`]s
423    ///
424    /// If an error occurs in the batch it retries 5 times if possible.
425    ///
426    /// This which will add/remove/edit an events in the index based on the
427    /// operations.
428    ///
429    /// Waits for background merge threads to prevent racing with subsequent
430    /// calls.
431    pub fn bulk_execute(&mut self, operations: Vec<RoomIndexOperation>) -> Result<(), IndexError> {
432        let mut writer = self.writer()?;
433        let mut operations = operations.into_iter();
434        let mut next_operation = operations.next();
435
436        while let Some(ref operation) = next_operation {
437            self.execute_with_retry(&mut writer, operation, 5)?;
438            next_operation = operations.next();
439        }
440
441        self.commit_and_reload(&mut writer)?;
442        writer.wait_merging_threads()?;
443
444        Ok(())
445    }
446
447    fn contains(&self, event_id: &EventId) -> bool {
448        let search_result = self.search(
449            format!("{}:\"{event_id}\"", self.schema.get_field_name(self.schema.primary_key()))
450                .as_str(),
451            1,
452            None,
453        );
454        match search_result {
455            Ok(results) => {
456                !self.uncommitted_removes.contains(event_id)
457                    && (!results.events.is_empty() || self.uncommitted_adds.contains_key(event_id))
458            }
459            Err(err) => {
460                warn!("Failed to check if event has been indexed, assuming it has: {err}");
461                true
462            }
463        }
464    }
465}
466
467#[cfg(test)]
468mod tests {
469    use std::{collections::HashSet, error::Error};
470
471    use matrix_sdk_test::event_factory::EventFactory;
472    use ruma::{
473        EventId, event_id,
474        events::{
475            AnySyncMessageLikeEvent,
476            room::message::{
477                MessageType, OriginalSyncRoomMessageEvent, Relation,
478                RoomMessageEventContentWithoutRelation,
479            },
480        },
481        room_id, user_id,
482    };
483
484    use super::{
485        IndexableEvent, MAX_MILLISECONDS, MilliSecondsSinceUnixEpoch, RoomIndex,
486        RoomIndexOperation, UInt, builder::RoomIndexBuilder,
487    };
488    use crate::error::IndexError;
489
490    /// Build an [`IndexableEvent`] from a text room message (tests only handle
491    /// text).
492    fn to_indexable(event: &OriginalSyncRoomMessageEvent) -> IndexableEvent {
493        let MessageType::Text(content) = &event.content.msgtype else {
494            panic!("test helper only supports text messages")
495        };
496        let original_event_id = match &event.content.relates_to {
497            Some(Relation::Replacement(replacement)) => replacement.event_id.clone(),
498            _ => event.event_id.clone(),
499        };
500
501        IndexableEvent::new(
502            event.event_id.clone(),
503            original_event_id,
504            event.sender.clone(),
505            Some(event.origin_server_ts),
506            content.body.clone(),
507        )
508    }
509
510    /// Helper function to add a regular message to the index
511    ///
512    /// # Panic
513    ///
514    /// Panics when event is not an [`OriginalSyncRoomMessageEvent`] with no
515    /// relations.
516    fn index_message(
517        index: &mut RoomIndex,
518        event: AnySyncMessageLikeEvent,
519    ) -> Result<(), IndexError> {
520        if let AnySyncMessageLikeEvent::RoomMessage(ev) = event
521            && let Some(ev) = ev.as_original()
522            && ev.content.relates_to.is_none()
523        {
524            return index.execute(RoomIndexOperation::Add(to_indexable(ev)));
525        }
526        panic!("Event was not a relationless OriginalSyncRoomMessageEvent.")
527    }
528
529    /// Helper function to remove events to the index
530    fn index_remove(index: &mut RoomIndex, event_id: &EventId) -> Result<(), IndexError> {
531        index.execute(RoomIndexOperation::Remove(event_id.to_owned()))
532    }
533
534    /// Helper function to edit events in index
535    ///
536    /// Edit event with `event_id` into new [`OriginalSyncRoomMessageEvent`]
537    fn index_edit(
538        index: &mut RoomIndex,
539        event_id: &EventId,
540        new: OriginalSyncRoomMessageEvent,
541    ) -> Result<(), IndexError> {
542        index.execute(RoomIndexOperation::Edit(event_id.to_owned(), to_indexable(&new)))
543    }
544
545    #[test]
546    fn test_add_event() {
547        let room_id = room_id!("!room_id:localhost");
548        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
549
550        let event = EventFactory::new()
551            .text_msg("event message")
552            .event_id(event_id!("$event_id:localhost"))
553            .room(room_id)
554            .sender(user_id!("@user_id:localhost"))
555            .into_any_sync_message_like_event();
556
557        index_message(&mut index, event).expect("failed to add event");
558    }
559
560    #[test]
561    fn test_add_event_with_no_timestamp() {
562        let mut index = RoomIndexBuilder::new_in_memory(room_id!("!r")).build();
563
564        index
565            .execute(RoomIndexOperation::Add(IndexableEvent::new(
566                event_id!("$ev").to_owned(),
567                event_id!("$ev").to_owned(),
568                user_id!("@mnt_io:matrix.org").to_owned(),
569                None,
570                "body".to_owned(),
571            )))
572            .expect("failed to add event");
573    }
574
575    #[test]
576    fn test_add_event_with_raw_malformed_timestamp_must_not_panic() {
577        let mut index = RoomIndexBuilder::new_in_memory(room_id!("!r")).build();
578
579        index
580            .execute(RoomIndexOperation::Add(IndexableEvent::new(
581                event_id!("$ev").to_owned(),
582                event_id!("$ev").to_owned(),
583                user_id!("@mnt_io:matrix.org").to_owned(),
584                Some(MilliSecondsSinceUnixEpoch(UInt::new(151393755000000).unwrap())),
585                "body".to_owned(),
586            )))
587            .expect("failed to add event");
588    }
589
590    #[test]
591    fn test_add_event_with_malformed_timestamp_must_not_panic() {
592        let room_id = room_id!("!r");
593        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
594
595        let body = "body".to_owned();
596        let event = EventFactory::new()
597            .server_ts(151393755000000)
598            .text_msg(body.clone())
599            .event_id(event_id!("$ev"))
600            .room(room_id)
601            .sender(user_id!("@mnt_io:matrix.org"))
602            .into_event();
603
604        index
605            .execute(RoomIndexOperation::Add(IndexableEvent::new(
606                event.event_id().unwrap().to_owned(),
607                event.event_id().unwrap().to_owned(),
608                event.sender().unwrap(),
609                event.timestamp(),
610                body,
611            )))
612            .expect("failed to add event");
613    }
614
615    #[test]
616    fn test_indexable_event_timestamp_is_capped() {
617        let event_id = event_id!("$ev").to_owned();
618        let sender = user_id!("@mnt_io:matrix.org").to_owned();
619        let body = "body".to_owned();
620
621        // Not capped.
622        let event = IndexableEvent::new(
623            event_id.clone(),
624            event_id.clone(),
625            sender.clone(),
626            Some(MilliSecondsSinceUnixEpoch(UInt::new(MAX_MILLISECONDS).unwrap())),
627            body.clone(),
628        );
629        assert_eq!(event.timestamp.unwrap().get(), UInt::new(MAX_MILLISECONDS).unwrap());
630
631        // Capped!
632        let event = IndexableEvent::new(
633            event_id.clone(),
634            event_id,
635            sender,
636            Some(MilliSecondsSinceUnixEpoch(UInt::new(MAX_MILLISECONDS + 42).unwrap())),
637            body,
638        );
639        assert_eq!(event.timestamp.unwrap().get(), UInt::new(MAX_MILLISECONDS).unwrap());
640    }
641
642    #[test]
643    fn test_search_populated_index() -> Result<(), Box<dyn Error>> {
644        let room_id = room_id!("!room_id:localhost");
645        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
646
647        let event_id_1 = event_id!("$event_id_1:localhost");
648        let event_id_2 = event_id!("$event_id_2:localhost");
649        let event_id_3 = event_id!("$event_id_3:localhost");
650        let user_id = user_id!("@user_id:localhost");
651        let f = EventFactory::new().room(room_id).sender(user_id);
652
653        index_message(
654            &mut index,
655            f.text_msg("This is a sentence")
656                .event_id(event_id_1)
657                .into_any_sync_message_like_event(),
658        )?;
659
660        index_message(
661            &mut index,
662            f.text_msg("All new words").event_id(event_id_2).into_any_sync_message_like_event(),
663        )?;
664
665        index_message(
666            &mut index,
667            f.text_msg("A similar sentence")
668                .event_id(event_id_3)
669                .into_any_sync_message_like_event(),
670        )?;
671
672        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
673        let result: HashSet<_> = result.events.iter().map(|(_, id)| id).collect();
674
675        let true_value = [event_id_1.to_owned(), event_id_3.to_owned()];
676        let true_value: HashSet<_> = true_value.iter().collect();
677
678        assert_eq!(result, true_value, "search result not correct: {result:?}");
679
680        Ok(())
681    }
682
683    #[test]
684    fn test_search_empty_index() -> Result<(), Box<dyn Error>> {
685        let room_id = room_id!("!room_id:localhost");
686        let index = RoomIndexBuilder::new_in_memory(room_id).build();
687
688        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
689
690        assert!(result.events.is_empty(), "search result not empty: {result:?}");
691
692        Ok(())
693    }
694
695    #[test]
696    fn test_index_contains_false() {
697        let room_id = room_id!("!room_id:localhost");
698        let index = RoomIndexBuilder::new_in_memory(room_id).build();
699
700        let event_id = event_id!("$event_id:localhost");
701
702        assert!(!index.contains(event_id), "Index should not contain event");
703    }
704
705    #[test]
706    fn test_index_contains_true() -> Result<(), Box<dyn Error>> {
707        let room_id = room_id!("!room_id:localhost");
708        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
709
710        let event_id = event_id!("$event_id:localhost");
711        let event = EventFactory::new()
712            .text_msg("This is a sentence")
713            .event_id(event_id)
714            .room(room_id)
715            .sender(user_id!("@user_id:localhost"))
716            .into_any_sync_message_like_event();
717
718        index_message(&mut index, event)?;
719
720        assert!(index.contains(event_id), "Index should contain event");
721
722        Ok(())
723    }
724
725    #[test]
726    fn test_index_add_idempotency() -> Result<(), Box<dyn Error>> {
727        let room_id = room_id!("!room_id:localhost");
728        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
729
730        let event_id = event_id!("$event_id:localhost");
731        let event = EventFactory::new()
732            .text_msg("This is a sentence")
733            .event_id(event_id)
734            .room(room_id)
735            .sender(user_id!("@user_id:localhost"))
736            .into_any_sync_message_like_event();
737
738        index_message(&mut index, event.clone())?;
739
740        assert!(index.contains(event_id), "Index should contain event");
741
742        // indexing again should do nothing
743        index_message(&mut index, event)?;
744
745        assert!(index.contains(event_id), "Index should still contain event");
746
747        let result = index.search("sentence", 10, None).expect("search failed with: {result:?}");
748
749        assert_eq!(result.events.len(), 1, "Index should have ignored second indexing");
750
751        Ok(())
752    }
753
754    #[test]
755    fn test_remove_event() -> Result<(), Box<dyn Error>> {
756        let room_id = room_id!("!room_id:localhost");
757        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
758
759        let event_id = event_id!("$event_id:localhost");
760        let user_id = user_id!("@user_id:localhost");
761        let f = EventFactory::new().room(room_id).sender(user_id);
762
763        let event =
764            f.text_msg("This is a sentence").event_id(event_id).into_any_sync_message_like_event();
765
766        index_message(&mut index, event)?;
767
768        assert!(index.contains(event_id), "Index should contain event");
769
770        index_remove(&mut index, event_id)?;
771
772        assert!(!index.contains(event_id), "Index should not contain event");
773
774        Ok(())
775    }
776
777    #[test]
778    fn test_edit_removes_old_and_adds_new_event() -> Result<(), Box<dyn Error>> {
779        let room_id = room_id!("!room_id:localhost");
780        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
781
782        let old_event_id = event_id!("$old_event_id:localhost");
783        let user_id = user_id!("@user_id:localhost");
784        let f = EventFactory::new().room(room_id).sender(user_id);
785
786        let old_event = f
787            .text_msg("This is a sentence")
788            .event_id(old_event_id)
789            .into_any_sync_message_like_event();
790
791        index_message(&mut index, old_event)?;
792
793        assert!(index.contains(old_event_id), "Index should contain event");
794
795        let new_event_id = event_id!("$new_event_id:localhost");
796        let edit = f
797            .text_msg("This is a brand new sentence!")
798            .edit(
799                old_event_id,
800                RoomMessageEventContentWithoutRelation::text_plain("This is a brand new sentence!"),
801            )
802            .event_id(new_event_id)
803            .into_original_sync_room_message_event();
804
805        index_edit(&mut index, old_event_id, edit)?;
806
807        assert!(!index.contains(old_event_id), "Index should not contain old event");
808        assert!(index.contains(new_event_id), "Index should contain edited event");
809
810        Ok(())
811    }
812
813    #[test]
814    fn test_bulk_add_then_edit_same_event_keeps_it_indexed() -> Result<(), Box<dyn Error>> {
815        let room_id = room_id!("!room_id:localhost");
816        let mut index = RoomIndexBuilder::new_in_memory(room_id).build();
817
818        let original_id = event_id!("$original:localhost");
819        let edit_id = event_id!("$edit:localhost");
820        let user_id = user_id!("@user_id:localhost");
821        let f = EventFactory::new().room(room_id).sender(user_id);
822
823        let edit = f
824            .text_msg("* brand new sentence")
825            .edit(
826                original_id,
827                RoomMessageEventContentWithoutRelation::text_plain("brand new sentence"),
828            )
829            .event_id(edit_id)
830            .into_original_sync_room_message_event();
831        let edit = to_indexable(&edit);
832
833        // An original and its edit arriving in the same batch produce an `Add`
834        // and an `Edit` of the same document. The `Edit`'s removal must not
835        // drop the document added earlier in the same uncommitted batch.
836        index.bulk_execute(vec![
837            RoomIndexOperation::Add(edit.clone()),
838            RoomIndexOperation::Edit(original_id.to_owned(), edit),
839        ])?;
840
841        assert!(index.contains(edit_id), "Edited document should be indexed");
842
843        let result = index.search("sentence", 10, None)?;
844        assert_eq!(result.total_count, 1, "Search should return the total count of results it got");
845        assert_eq!(
846            result.events.len(),
847            1,
848            "Search should find the edited document, got {result:?}"
849        );
850        assert_eq!(result.events[0].1, edit_id, "unexpected event id: {result:?}");
851
852        Ok(())
853    }
854}