1pub 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#[derive(Debug)]
45pub struct SearchResult {
46 pub total_count: usize,
51
52 pub events: Vec<(f32, OwnedEventId)>,
59}
60
61#[derive(Clone)]
66pub struct IndexableEvent {
67 pub(crate) event_id: OwnedEventId,
69 pub(crate) original_event_id: OwnedEventId,
72 pub(crate) sender: OwnedUserId,
74 pub(crate) timestamp: Option<MilliSecondsSinceUnixEpoch>,
80 pub(crate) body: String,
82}
83
84impl fmt::Debug for IndexableEvent {
85 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
96const MAX_MILLISECONDS: u64 = (i64::MAX / 1_000_000).cast_unsigned();
99
100impl IndexableEvent {
101 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 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#[derive(Debug, Clone)]
129pub enum RoomIndexOperation {
130 Add(IndexableEvent),
132 Remove(OwnedEventId),
135 Edit(OwnedEventId, IndexableEvent),
139 Noop,
141}
142
143pub struct RoomIndex {
146 index: Index,
147 schema: RoomMessageSchema,
148 query_parser: QueryParser,
149 room_id: OwnedRoomId,
150 uncommitted_adds: HashMap<OwnedEventId, OwnedEventId>,
155 uncommitted_removes: HashSet<OwnedEventId>,
156 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 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 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 fn commit(&mut self, writer: &mut SearchIndexWriter) -> Result<OpStamp, IndexError> {
210 let last_commit_opstamp = writer.commit()?; self.uncommitted_adds.clear();
212 self.uncommitted_removes.clear();
213 Ok(last_commit_opstamp)
214 }
215
216 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 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 for event in events.into_iter() {
320 self.uncommitted_adds.remove(&event);
321 self.uncommitted_removes.insert(event);
322 }
323
324 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 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 IndexError::TantivyError(_)
380 | IndexError::IndexSchemaError(_)
381 | IndexError::IndexWriteError(_)
382 | IndexError::IO(_) => {
383 num_tries += 1;
384 }
385 IndexError::OpenDirectoryError(ref e) => match e {
386 OpenDirectoryError::IoError { io_error: _, directory_path: _ } => {
388 num_tries += 1;
389 }
390 OpenDirectoryError::DoesNotExist(_)
392 | OpenDirectoryError::FailedToCreateTempDir(_)
393 | OpenDirectoryError::NotADirectory(_) => return Err(err),
394 },
395 IndexError::QueryParserError(_) => return Err(err),
397 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 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 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 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 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 fn index_remove(index: &mut RoomIndex, event_id: &EventId) -> Result<(), IndexError> {
531 index.execute(RoomIndexOperation::Remove(event_id.to_owned()))
532 }
533
534 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 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 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 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 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}