1use std::{
18 collections::HashMap,
19 fmt,
20 iter::once,
21 ops::{Deref, Not},
22 path::{Path, PathBuf},
23 sync::Arc,
24};
25
26use async_trait::async_trait;
27use deadpool::managed::PoolConfig;
28use matrix_sdk_base::{
29 cross_process_lock::CrossProcessLockGeneration,
30 deserialized_responses::TimelineEvent,
31 event_cache::{
32 Event, Gap,
33 store::{EventCacheStore, extract_event_relation},
34 thread::ThreadInfo,
35 },
36 linked_chunk::{
37 ChunkContent, ChunkIdentifier, ChunkIdentifierGenerator, ChunkMetadata, LinkedChunkId,
38 Position, RawChunk, Update,
39 },
40 timer,
41};
42use matrix_sdk_store_encryption::StoreCipher;
43use ruma::{
44 EventId, MilliSecondsSinceUnixEpoch, OwnedEventId, RoomId, events::relation::RelationType,
45};
46use rusqlite::{
47 OptionalExtension, ToSql, Transaction, TransactionBehavior, params, params_from_iter,
48};
49use tokio::{
50 fs,
51 sync::{Mutex, OwnedMutexGuard},
52};
53use tracing::{debug, error, instrument, trace};
54
55use crate::{
56 OpenStoreError, RuntimeConfig, Secret, SqliteStoreConfig,
57 connection::{self, Connection as SqliteAsyncConn, Pool as SqlitePool, SqliteConnections},
58 error::{Error, Result},
59 utils::{
60 EncryptableStore, Key, SqliteAsyncConnExt, SqliteKeyValueStoreAsyncConnExt,
61 SqliteKeyValueStoreConnExt, SqliteTransactionExt, host_parameters,
62 },
63};
64
65mod keys {
66 pub const LINKED_CHUNKS: &str = "linked_chunks";
68 pub const EVENTS: &str = "events";
69}
70
71const DATABASE_NAME: &str = "matrix-sdk-event-cache.sqlite3";
73
74const CHUNK_TYPE_EVENT_TYPE_STRING: &str = "E";
77const CHUNK_TYPE_GAP_TYPE_STRING: &str = "G";
80
81struct Encryption {
86 cipher: Option<StoreCipher>,
87}
88
89impl Encryption {
90 fn encode_event(&self, event: &TimelineEvent) -> Result<EncodedEvent> {
91 let serialized = serde_json::to_vec(event)?;
92
93 let raw_event = event.raw();
95 let (relates_to, rel_type) = extract_event_relation(raw_event).unzip();
96
97 let content = self.encode_value(serialized)?;
99
100 Ok(EncodedEvent {
101 content,
102 rel_type,
103 relates_to: relates_to
104 .map(|relates_to| self.encode_event_id(keys::EVENTS, &relates_to)),
105 })
106 }
107
108 fn decode_event(&self, raw_encoded_event: &[u8]) -> Result<Event> {
109 Ok(serde_json::from_slice(&self.decode_value(raw_encoded_event)?)?)
110 }
111
112 fn encode_event_id(&self, table_name: &str, event_id: &EventId) -> Key {
115 self.encode_key(table_name, event_id)
116 }
117
118 fn encode_room_id(&self, table_name: &str, room_id: &RoomId) -> Key {
121 self.encode_key(table_name, room_id)
122 }
123
124 fn encode_linked_chunk(&self, table_name: &str, linked_chunk_id: &LinkedChunkId<'_>) -> Key {
126 self.encode_key(table_name, linked_chunk_id.storage_key())
127 }
128
129 fn encode_thread_id(&self, thread_id: &EventId) -> Result<Vec<u8>> {
131 self.encode_value(String::from(thread_id.as_str()))
132 }
133
134 #[cfg(test)]
136 fn decode_thread_id(&self, encoded_thread_id: &[u8]) -> Result<OwnedEventId> {
137 let as_slice = self.decode_value(encoded_thread_id)?;
138 let as_str = str::from_utf8(as_slice.as_ref())?;
139
140 Ok(EventId::parse(as_str)?)
141 }
142
143 fn encode_thread_info(&self, thread_info: &ThreadInfo) -> Result<Vec<u8>> {
145 self.encode_value(serde_json::to_vec(thread_info)?)
146 }
147
148 fn decode_thread_info(&self, encoded_thread_info: &[u8]) -> Result<ThreadInfo> {
150 Ok(serde_json::from_slice(&self.decode_value(encoded_thread_info)?)?)
151 }
152}
153
154impl EncryptableStore for Encryption {
155 fn get_cypher(&self) -> Option<&StoreCipher> {
156 self.cipher.as_ref()
157 }
158}
159
160#[derive(Clone)]
162pub struct SqliteEventCacheStore {
163 encryption: Arc<Encryption>,
165
166 connections: Arc<Mutex<Option<SqliteConnections>>>,
168
169 db_path: PathBuf,
171
172 pool_config: PoolConfig,
174
175 runtime_config: RuntimeConfig,
177}
178
179#[cfg(not(tarpaulin_include))]
180impl fmt::Debug for SqliteEventCacheStore {
181 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182 f.debug_struct("SqliteEventCacheStore").finish_non_exhaustive()
183 }
184}
185
186impl SqliteEventCacheStore {
187 pub async fn open(
190 path: impl AsRef<Path>,
191 passphrase: Option<&str>,
192 ) -> Result<Self, OpenStoreError> {
193 Self::open_with_config(&SqliteStoreConfig::new(path).passphrase(passphrase)).await
194 }
195
196 pub async fn open_with_key(
199 path: impl AsRef<Path>,
200 key: Option<&[u8; 32]>,
201 ) -> Result<Self, OpenStoreError> {
202 Self::open_with_config(&SqliteStoreConfig::new(path).key(key)).await
203 }
204
205 #[instrument(skip(config), fields(path = ?config.path))]
207 pub async fn open_with_config(config: &SqliteStoreConfig) -> Result<Self, OpenStoreError> {
208 debug!(?config);
209
210 let _timer = timer!("open_with_config");
211
212 fs::create_dir_all(&config.path).await.map_err(OpenStoreError::CreateDir)?;
213
214 let db_path = config.path.join(DATABASE_NAME);
215 let pool_config = config.pool_config();
216 let runtime_config = config.runtime_config();
217
218 let pool = config.build_pool_of_connections(DATABASE_NAME)?;
219
220 let this =
221 Self::open_with_pool(pool, db_path, pool_config, runtime_config, config.secret.clone())
222 .await?;
223
224 this.write().await?.apply_runtime_config(runtime_config).await?;
226
227 Ok(this)
228 }
229
230 async fn open_with_pool(
233 pool: SqlitePool,
234 db_path: PathBuf,
235 pool_config: PoolConfig,
236 runtime_config: RuntimeConfig,
237 secret: Option<Secret>,
238 ) -> Result<Self, OpenStoreError> {
239 let conn = pool.get().await?;
240
241 let version = conn.db_version().await?;
242
243 run_migrations(&conn, version).await?;
244
245 conn.wal_checkpoint().await;
246
247 let cipher = match secret {
248 Some(s) => Some(conn.get_or_create_store_cipher(s).await?),
249 None => None,
250 };
251
252 let connections = SqliteConnections {
253 pool,
254 write_connection: Arc::new(Mutex::new(conn)),
256 };
257
258 Ok(Self {
259 encryption: Arc::new(Encryption { cipher }),
260 connections: Arc::new(Mutex::new(Some(connections))),
261 db_path,
262 pool_config,
263 runtime_config,
264 })
265 }
266
267 #[instrument(skip_all)]
269 async fn read(&self) -> Result<SqliteAsyncConn> {
270 let pool = {
271 let guard = self.connections.lock().await;
272 let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
273 conns.pool.clone()
274 };
275
276 let connection = pool.get().await?;
277
278 connection.execute_batch("PRAGMA foreign_keys = ON;").await?;
283
284 Ok(connection)
285 }
286
287 #[instrument(skip_all)]
289 async fn write(&self) -> Result<OwnedMutexGuard<SqliteAsyncConn>> {
290 let write_connection = {
291 let guard = self.connections.lock().await;
292 let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
293 conns.write_connection.clone()
294 };
295
296 let connection = write_connection.lock_owned().await;
297
298 connection.execute_batch("PRAGMA foreign_keys = ON;").await?;
303
304 Ok(connection)
305 }
306
307 fn map_row_to_chunk(
308 row: &rusqlite::Row<'_>,
309 ) -> Result<(u64, Option<u64>, Option<u64>, String), rusqlite::Error> {
310 Ok((
311 row.get::<_, u64>(0)?,
312 row.get::<_, Option<u64>>(1)?,
313 row.get::<_, Option<u64>>(2)?,
314 row.get::<_, String>(3)?,
315 ))
316 }
317
318 pub async fn vacuum(&self) -> Result<()> {
319 let write_connection = {
320 let guard = self.connections.lock().await;
321 let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
322 conns.write_connection.clone()
323 };
324 write_connection.lock().await.vacuum().await
325 }
326
327 async fn get_db_size(&self) -> Result<Option<usize>> {
328 let pool = {
329 let guard = self.connections.lock().await;
330 let conns = guard.as_ref().ok_or(Error::StoreClosed)?;
331 conns.pool.clone()
332 };
333 Ok(Some(pool.get().await?.get_db_size().await?))
334 }
335
336 pub async fn close(&self) -> Result<()> {
337 connection::close_connections(&self.connections, "Event cache store").await;
338 Ok(())
339 }
340
341 pub async fn reopen(&self) -> Result<()> {
342 connection::reopen_connections(
343 &self.connections,
344 self.db_path.clone(),
345 self.pool_config,
346 self.runtime_config,
347 )
348 .await?;
349 Ok(())
350 }
351
352 #[cfg(test)]
354 async fn pool_max_size(&self) -> Option<usize> {
355 let guard = self.connections.lock().await;
356 guard.as_ref().map(|conns| conns.pool.status().max_size)
357 }
358}
359
360struct EncodedEvent {
361 content: Vec<u8>,
362 rel_type: Option<String>,
363 relates_to: Option<Key>,
364}
365
366trait TransactionExtForLinkedChunks {
367 fn rebuild_chunk(
368 &self,
369 encryption: &Encryption,
370 linked_chunk_id: &Key,
371 previous: Option<u64>,
372 index: u64,
373 next: Option<u64>,
374 chunk_type: &str,
375 ) -> Result<RawChunk<Event, Gap>>;
376
377 fn load_gap_content(
378 &self,
379 encryption: &Encryption,
380 linked_chunk_id: &Key,
381 chunk_id: ChunkIdentifier,
382 ) -> Result<Gap>;
383
384 fn load_events_content(
385 &self,
386 encryption: &Encryption,
387 linked_chunk_id: &Key,
388 chunk_id: ChunkIdentifier,
389 ) -> Result<Vec<Event>>;
390}
391
392impl TransactionExtForLinkedChunks for Transaction<'_> {
393 fn rebuild_chunk(
394 &self,
395 encryption: &Encryption,
396 linked_chunk_id: &Key,
397 previous: Option<u64>,
398 id: u64,
399 next: Option<u64>,
400 chunk_type: &str,
401 ) -> Result<RawChunk<Event, Gap>> {
402 let previous = previous.map(ChunkIdentifier::new);
403 let next = next.map(ChunkIdentifier::new);
404 let id = ChunkIdentifier::new(id);
405
406 match chunk_type {
407 CHUNK_TYPE_GAP_TYPE_STRING => {
408 let gap = self.load_gap_content(encryption, linked_chunk_id, id)?;
410 Ok(RawChunk { content: ChunkContent::Gap(gap), previous, identifier: id, next })
411 }
412
413 CHUNK_TYPE_EVENT_TYPE_STRING => {
414 let events = self.load_events_content(encryption, linked_chunk_id, id)?;
416 Ok(RawChunk {
417 content: ChunkContent::Items(events),
418 previous,
419 identifier: id,
420 next,
421 })
422 }
423
424 other => {
425 Err(Error::InvalidData {
427 details: format!("a linked chunk has an unknown type {other}"),
428 })
429 }
430 }
431 }
432
433 fn load_gap_content(
434 &self,
435 encryption: &Encryption,
436 linked_chunk_id: &Key,
437 chunk_id: ChunkIdentifier,
438 ) -> Result<Gap> {
439 let encoded_prev_token: Vec<u8> = self.query_one(
442 "SELECT prev_token FROM gap_chunks WHERE chunk_id = ? AND linked_chunk_id = ?",
443 (chunk_id.index(), &linked_chunk_id),
444 |row| row.get(0),
445 )?;
446 let prev_token_bytes = encryption.decode_value(&encoded_prev_token)?;
447 let prev_token = String::from_utf8(prev_token_bytes.into_owned())?;
448 Ok(Gap { token: prev_token })
449 }
450
451 fn load_events_content(
452 &self,
453 encryption: &Encryption,
454 linked_chunk_id: &Key,
455 chunk_id: ChunkIdentifier,
456 ) -> Result<Vec<Event>> {
457 let mut events = Vec::new();
459
460 for event_data in self
461 .prepare(
462 "SELECT events.content \
463 FROM event_chunks ec, events \
464 WHERE events.event_id = ec.event_id AND ec.chunk_id = ? AND ec.linked_chunk_id = ? \
465 ORDER BY ec.position ASC",
466 )?
467 .query_map((chunk_id.index(), &linked_chunk_id), |row| row.get::<_, Vec<u8>>(0))?
468 {
469 events.push(encryption.decode_event(&event_data?)?);
470 }
471
472 Ok(events)
473 }
474}
475
476async fn run_migrations(conn: &SqliteAsyncConn, version: u8) -> Result<()> {
478 conn.execute_batch("PRAGMA foreign_keys = ON;").await?;
480
481 if version < 1 {
482 debug!("Creating database");
483 conn.execute_batch("PRAGMA journal_mode = wal;").await?;
486 conn.with_transaction(|txn| {
487 txn.execute_batch(include_str!("../migrations/event_cache_store/001_init.sql"))?;
488 txn.set_db_version(1)
489 })
490 .await?;
491 }
492
493 if version < 2 {
494 debug!("Upgrading database to version 2");
495 conn.with_transaction(|txn| {
496 txn.execute_batch(include_str!("../migrations/event_cache_store/002_lease_locks.sql"))?;
497 txn.set_db_version(2)
498 })
499 .await?;
500 }
501
502 if version < 3 {
503 debug!("Upgrading database to version 3");
504 conn.with_transaction(|txn| {
505 txn.execute_batch(include_str!("../migrations/event_cache_store/003_events.sql"))?;
506 txn.set_db_version(3)
507 })
508 .await?;
509 }
510
511 if version < 4 {
512 debug!("Upgrading database to version 4");
513 conn.with_transaction(|txn| {
514 txn.execute_batch(include_str!(
515 "../migrations/event_cache_store/004_ignore_policy.sql"
516 ))?;
517 txn.set_db_version(4)
518 })
519 .await?;
520 }
521
522 if version < 5 {
523 debug!("Upgrading database to version 5");
524 conn.with_transaction(|txn| {
525 txn.execute_batch(include_str!(
526 "../migrations/event_cache_store/005_events_index_on_event_id.sql"
527 ))?;
528 txn.set_db_version(5)
529 })
530 .await?;
531 }
532
533 if version < 6 {
534 debug!("Upgrading database to version 6");
535 conn.with_transaction(|txn| {
536 txn.execute_batch(include_str!("../migrations/event_cache_store/006_events.sql"))?;
537 txn.set_db_version(6)
538 })
539 .await?;
540 }
541
542 if version < 7 {
543 debug!("Upgrading database to version 7");
544 conn.with_transaction(|txn| {
545 txn.execute_batch(include_str!(
546 "../migrations/event_cache_store/007_event_chunks.sql"
547 ))?;
548 txn.set_db_version(7)
549 })
550 .await?;
551 }
552
553 if version < 8 {
554 debug!("Upgrading database to version 8");
555 conn.with_transaction(|txn| {
556 txn.execute_batch(include_str!(
557 "../migrations/event_cache_store/008_linked_chunk_id.sql"
558 ))?;
559 txn.set_db_version(8)
560 })
561 .await?;
562 }
563
564 if version < 9 {
565 debug!("Upgrading database to version 9");
566 conn.with_transaction(|txn| {
567 txn.execute_batch(include_str!(
568 "../migrations/event_cache_store/009_related_event_index.sql"
569 ))?;
570 txn.set_db_version(9)
571 })
572 .await?;
573 }
574
575 if version < 10 {
576 debug!("Upgrading database to version 10");
577 conn.with_transaction(|txn| {
578 txn.execute_batch(include_str!("../migrations/event_cache_store/010_drop_media.sql"))?;
579 txn.set_db_version(10)
580 })
581 .await?;
582
583 if version >= 1 {
584 conn.vacuum().await?;
587 }
588 }
589
590 if version < 11 {
591 debug!("Upgrading database to version 11");
592 conn.with_transaction(|txn| {
593 txn.execute_batch(include_str!(
594 "../migrations/event_cache_store/011_empty_event_cache.sql"
595 ))?;
596 txn.set_db_version(11)
597 })
598 .await?;
599 }
600
601 if version < 12 {
602 debug!("Upgrading database to version 12");
603 conn.with_transaction(|txn| {
604 txn.execute_batch(include_str!(
605 "../migrations/event_cache_store/012_store_event_type.sql"
606 ))?;
607 txn.set_db_version(12)
608 })
609 .await?;
610 }
611
612 if version < 13 {
613 debug!("Upgrading database to version 13");
614 conn.with_transaction(|txn| {
615 txn.execute_batch(include_str!(
616 "../migrations/event_cache_store/013_lease_locks_with_generation.sql"
617 ))?;
618 txn.set_db_version(13)
619 })
620 .await?;
621 }
622
623 if version < 14 {
624 debug!("Upgrading database to version 14");
625 conn.with_transaction(|txn| {
626 txn.execute_batch(include_str!(
627 "../migrations/event_cache_store/014_event_chunks_event_id_index.sql"
628 ))?;
629 txn.set_db_version(14)
630 })
631 .await?;
632 }
633
634 if version < 15 {
635 debug!("Upgrading database to version 15");
636 conn.with_transaction(|txn| {
637 txn.execute_batch(include_str!(
638 "../migrations/event_cache_store/015_event_ids_are_encoded.sql"
639 ))?;
640 txn.set_db_version(15)
641 })
642 .await?;
643 }
644
645 if version < 16 {
646 debug!("Upgrading database to version 16");
647 conn.with_transaction(|txn| {
648 txn.execute_batch(include_str!("../migrations/event_cache_store/016_threads.sql"))?;
649 txn.set_db_version(16)
650 })
651 .await?;
652 }
653
654 if version < 17 {
655 debug!("Upgrading database to version 17");
656 conn.with_transaction(|txn| {
657 txn.execute_batch(include_str!(
658 "../migrations/event_cache_store/017_threads_with_thread_infos.sql"
659 ))?;
660 txn.set_db_version(17)
661 })
662 .await?;
663 }
664
665 if version < 18 {
666 debug!("Upgrading database to version 18");
667 conn.with_transaction(|txn| {
668 txn.execute_batch(include_str!("../migrations/event_cache_store/018_event_chunks_unique_linked_chunk_id_event_id.sql"))?;
669 txn.set_db_version(18)
670 })
671 .await?;
672 }
673
674 Ok(())
675}
676
677#[async_trait]
678impl EventCacheStore for SqliteEventCacheStore {
679 type Error = Error;
680
681 #[instrument(skip(self))]
682 async fn try_take_leased_lock(
683 &self,
684 lease_duration_ms: u32,
685 key: &str,
686 holder: &str,
687 ) -> Result<Option<CrossProcessLockGeneration>> {
688 let key = key.to_owned();
689 let holder = holder.to_owned();
690
691 let now: u64 = MilliSecondsSinceUnixEpoch::now().get().into();
692 let expiration = now + lease_duration_ms as u64;
693
694 let generation = self
696 .write()
697 .await?
698 .with_transaction(move |txn| {
699 txn.query_one(
700 "INSERT INTO lease_locks (key, holder, expiration) \
701 VALUES (?1, ?2, ?3) \
702 ON CONFLICT (key) \
703 DO \
704 UPDATE SET \
705 holder = excluded.holder, \
706 expiration = excluded.expiration, \
707 generation = \
708 CASE holder \
709 WHEN excluded.holder THEN generation \
710 ELSE generation + 1 \
711 END \
712 WHERE \
713 holder = excluded.holder \
714 OR expiration < ?4 \
715 RETURNING generation",
716 (key, holder, expiration, now),
717 |row| row.get(0),
718 )
719 .optional()
720 })
721 .await?;
722
723 Ok(generation)
724 }
725
726 #[instrument(skip(self, updates))]
727 async fn handle_linked_chunk_updates(
728 &self,
729 linked_chunk_id: LinkedChunkId<'_>,
730 updates: Vec<Update<Event, Gap>>,
731 ) -> Result<(), Self::Error> {
732 let _timer = timer!("method");
733
734 let hashed_linked_chunk_id =
735 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
736 let hashed_room_id =
737 self.encryption.encode_room_id(keys::EVENTS, linked_chunk_id.room_id());
738 let encryption = self.encryption.clone();
739
740 with_immediate_transaction(self, move |txn| {
743 for update in updates {
744 match update {
745 Update::NewItemsChunk { previous, new, next } => {
746 let previous = previous.as_ref().map(ChunkIdentifier::index);
747 let new = new.index();
748 let next = next.as_ref().map(ChunkIdentifier::index);
749
750 trace!("new events chunk (prev={previous:?}, i={new}, next={next:?})");
751
752 insert_chunk(
753 txn,
754 &hashed_linked_chunk_id,
755 previous,
756 new,
757 next,
758 CHUNK_TYPE_EVENT_TYPE_STRING,
759 )?;
760 }
761
762 Update::NewGapChunk { previous, new, next, gap } => {
763 let hashed_prev_token = encryption.encode_value(gap.token)?;
764
765 let previous = previous.as_ref().map(ChunkIdentifier::index);
766 let new = new.index();
767 let next = next.as_ref().map(ChunkIdentifier::index);
768
769 trace!("new gap chunk (prev={previous:?}, i={new}, next={next:?})");
770
771 insert_chunk(
773 txn,
774 &hashed_linked_chunk_id,
775 previous,
776 new,
777 next,
778 CHUNK_TYPE_GAP_TYPE_STRING,
779 )?;
780
781 txn.execute(
783 r#"
784 INSERT INTO gap_chunks(chunk_id, linked_chunk_id, prev_token)
785 VALUES (?, ?, ?)
786 "#,
787 (new, &hashed_linked_chunk_id, hashed_prev_token),
788 )?;
789 }
790
791 Update::RemoveChunk(chunk_identifier) => {
792 let chunk_id = chunk_identifier.index();
793
794 trace!("removing chunk @ {chunk_id}");
795
796 let (previous, next): (Option<usize>, Option<usize>) = txn.query_one(
798 "SELECT previous, next FROM linked_chunks WHERE id = ? AND linked_chunk_id = ?",
799 (chunk_id, &hashed_linked_chunk_id),
800 |row| Ok((row.get(0)?, row.get(1)?))
801 )?;
802
803 if let Some(previous) = previous {
805 txn.execute("UPDATE linked_chunks SET next = ? WHERE id = ? AND linked_chunk_id = ?", (next, previous, &hashed_linked_chunk_id))?;
806 }
807
808 if let Some(next) = next {
810 txn.execute("UPDATE linked_chunks SET previous = ? WHERE id = ? AND linked_chunk_id = ?", (previous, next, &hashed_linked_chunk_id))?;
811 }
812
813 txn.execute("DELETE FROM linked_chunks WHERE id = ? AND linked_chunk_id = ?", (chunk_id, &hashed_linked_chunk_id))?;
816 }
817
818 Update::PushItems { at, items } => {
819 if items.is_empty() {
820 continue;
822 }
823
824 let chunk_id = at.chunk_identifier().index();
825
826 trace!("pushing {} items @ {chunk_id}", items.len());
827
828 let mut chunk_statement = txn.prepare(
829 "INSERT INTO event_chunks(chunk_id, linked_chunk_id, event_id, position) VALUES (?, ?, ?, ?)"
830 )?;
831
832 let mut content_statement = txn.prepare(
837 "INSERT OR REPLACE INTO events(room_id, event_id, event_type, session_id, content, relates_to, rel_type) VALUES (?, ?, ?, ?, ?, ?, ?)"
838 )?;
839
840 let invalid_event = |event: TimelineEvent| {
841 let Some(event_id) = event.event_id() else {
842 error!("Trying to push an event with no ID");
843 return None;
844 };
845
846 let Some(event_type) = event.kind.event_type() else {
847 error!(%event_id, "Trying to save an event with no event type");
848 return None;
849 };
850
851 Some((event_id.to_owned(), event_type, event))
852 };
853
854 for (i, (event_id, event_type, event)) in items.into_iter().filter_map(invalid_event).enumerate() {
855 let hashed_event_id = encryption.encode_event_id(
856 keys::EVENTS,
861 &event_id,
862 );
863
864 {
866 let index = at.index() + i;
867
868 chunk_statement.execute((chunk_id, &hashed_linked_chunk_id, &hashed_event_id, index))?;
869 }
870
871 {
873 let hashed_session_id = event.kind.session_id().map(|s| encryption.encode_key(keys::EVENTS, s));
874 let hashed_event_type = encryption.encode_key(keys::EVENTS, event_type);
875 let encoded_event = encryption.encode_event(&event)?;
876
877 content_statement.execute((
878 &hashed_room_id,
879 &hashed_event_id,
880 hashed_event_type,
881 hashed_session_id,
882 encoded_event.content,
883 encoded_event.relates_to,
884 encoded_event.rel_type
885 ))?;
886 }
887 }
888 }
889
890 Update::ReplaceItem { at, item: event } => {
891 let chunk_id = at.chunk_identifier().index();
892 let index = at.index();
893
894 trace!("replacing item @ {chunk_id}:{index}");
895
896 let Some(event_id) = event.event_id().map(|event_id| event_id.to_owned()) else {
898 error!("Trying to replace an event with a new one that has no ID");
899 continue;
900 };
901
902 let Some(event_type) = event.kind.event_type() else {
903 error!(%event_id, "Trying to save an event with no event type");
904 continue;
905 };
906
907 let hashed_event_id = encryption.encode_event_id(keys::EVENTS, &event_id);
908
909 {
916 let hashed_session_id = event.kind.session_id().map(|s| encryption.encode_key(keys::EVENTS, s));
917 let hashed_event_type = encryption.encode_key(keys::EVENTS, event_type);
918 let encoded_event = encryption.encode_event(&event)?;
919
920 txn.execute(
921 "INSERT OR REPLACE INTO events(room_id, event_id, event_type, session_id, content, relates_to, rel_type) VALUES (?, ?, ?, ?, ?, ?, ?)",
922 (
923 &hashed_room_id,
924 &hashed_event_id,
925 hashed_event_type,
926 hashed_session_id,
927 encoded_event.content,
928 encoded_event.relates_to,
929 encoded_event.rel_type
930 ),
931 )?;
932 }
933
934 {
936 txn.execute(
938 r#"UPDATE event_chunks SET event_id = ? WHERE linked_chunk_id = ? AND chunk_id = ? AND position = ?"#,
939 (&hashed_event_id, &hashed_linked_chunk_id, chunk_id, index)
940 )?;
941 }
942 }
943
944 Update::RemoveItem { at } => {
945 let chunk_id = at.chunk_identifier().index();
946 let index = at.index();
947
948 trace!("removing item @ {chunk_id}:{index}");
949
950 txn.execute("DELETE FROM event_chunks WHERE linked_chunk_id = ? AND chunk_id = ? AND position = ?", (&hashed_linked_chunk_id, chunk_id, index))?;
952
953 txn.execute(
1052 r#"
1053 UPDATE event_chunks
1054 SET position = -(position - 1)
1055 WHERE linked_chunk_id = ? AND chunk_id = ? AND position > ?
1056 "#,
1057 (&hashed_linked_chunk_id, chunk_id, index)
1058 )?;
1059 txn.execute(
1060 r#"
1061 UPDATE event_chunks
1062 SET position = -position
1063 WHERE position < 0 AND linked_chunk_id = ? AND chunk_id = ?
1064 "#,
1065 (&hashed_linked_chunk_id, chunk_id)
1066 )?;
1067
1068 }
1071
1072 Update::DetachLastItems { at } => {
1073 let chunk_id = at.chunk_identifier().index();
1074 let index = at.index();
1075
1076 trace!("truncating items >= {chunk_id}:{index}");
1077
1078 txn.execute("DELETE FROM event_chunks WHERE linked_chunk_id = ? AND chunk_id = ? AND position >= ?", (&hashed_linked_chunk_id, chunk_id, index))?;
1080
1081 }
1084
1085 Update::Clear => {
1086 trace!("clearing items");
1087
1088 txn.execute(
1090 "DELETE FROM linked_chunks WHERE linked_chunk_id = ?",
1091 (&hashed_linked_chunk_id,),
1092 )?;
1093
1094 }
1097
1098 Update::StartReattachItems | Update::EndReattachItems => {
1099 }
1101 }
1102 }
1103
1104 Ok(())
1105 })
1106 .await?;
1107
1108 Ok(())
1109 }
1110
1111 #[instrument(skip(self))]
1112 async fn load_all_chunks(
1113 &self,
1114 linked_chunk_id: LinkedChunkId<'_>,
1115 ) -> Result<Vec<RawChunk<Event, Gap>>, Self::Error> {
1116 let _timer = timer!("method");
1117
1118 let hashed_linked_chunk_id =
1119 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1120 let encryption = self.encryption.clone();
1121
1122 let result = self
1123 .read()
1124 .await?
1125 .with_transaction(move |txn| -> Result<_> {
1126 let mut items = Vec::new();
1127
1128 for data in txn
1130 .prepare(
1131 "SELECT id, previous, next, type FROM linked_chunks WHERE linked_chunk_id = ? ORDER BY id",
1132 )?
1133 .query_map((&hashed_linked_chunk_id,), Self::map_row_to_chunk)?
1134 {
1135 let (id, previous, next, chunk_type) = data?;
1136 let new = txn.rebuild_chunk(
1137 &encryption,
1138 &hashed_linked_chunk_id,
1139 previous,
1140 id,
1141 next,
1142 chunk_type.as_str(),
1143 )?;
1144 items.push(new);
1145 }
1146
1147 Ok(items)
1148 })
1149 .await?;
1150
1151 Ok(result)
1152 }
1153
1154 #[instrument(skip(self))]
1155 async fn load_all_chunks_metadata(
1156 &self,
1157 linked_chunk_id: LinkedChunkId<'_>,
1158 ) -> Result<Vec<ChunkMetadata>, Self::Error> {
1159 let _timer = timer!("method");
1160
1161 let hashed_linked_chunk_id =
1162 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1163
1164 self.read()
1165 .await?
1166 .with_transaction(move |txn| -> Result<_> {
1167 let num_events_by_chunk_ids = txn
1194 .prepare(
1195 "SELECT ec.chunk_id, COUNT(ec.event_id) \
1196 FROM event_chunks as ec \
1197 WHERE ec.linked_chunk_id = ? \
1198 GROUP BY ec.chunk_id",
1199 )?
1200 .query_map((&hashed_linked_chunk_id,), |row| {
1201 Ok((row.get::<_, u64>(0)?, row.get::<_, usize>(1)?))
1202 })?
1203 .collect::<Result<HashMap<_, _>, _>>()?;
1204
1205 txn.prepare(
1206 "SELECT \
1207 lc.id, \
1208 lc.previous, \
1209 lc.next, \
1210 lc.type \
1211 FROM linked_chunks as lc \
1212 WHERE lc.linked_chunk_id = ? \
1213 ORDER BY lc.id",
1214 )?
1215 .query_map((&hashed_linked_chunk_id,), |row| {
1216 Ok((
1217 row.get::<_, u64>(0)?,
1218 row.get::<_, Option<u64>>(1)?,
1219 row.get::<_, Option<u64>>(2)?,
1220 row.get::<_, String>(3)?,
1221 ))
1222 })?
1223 .map(|data| -> Result<_> {
1224 let (id, previous, next, chunk_type) = data?;
1225
1226 let num_items = if chunk_type == CHUNK_TYPE_GAP_TYPE_STRING {
1232 0
1233 } else {
1234 num_events_by_chunk_ids.get(&id).copied().unwrap_or(0)
1235 };
1236
1237 Ok(ChunkMetadata {
1238 identifier: ChunkIdentifier::new(id),
1239 previous: previous.map(ChunkIdentifier::new),
1240 next: next.map(ChunkIdentifier::new),
1241 num_items,
1242 })
1243 })
1244 .collect::<Result<Vec<_>, _>>()
1245 })
1246 .await
1247 }
1248
1249 #[instrument(skip(self))]
1250 async fn load_last_chunk(
1251 &self,
1252 linked_chunk_id: LinkedChunkId<'_>,
1253 ) -> Result<(Option<RawChunk<Event, Gap>>, ChunkIdentifierGenerator), Self::Error> {
1254 let _timer = timer!("method");
1255
1256 let hashed_linked_chunk_id =
1257 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1258 let encryption = self.encryption.clone();
1259
1260 self
1261 .read()
1262 .await?
1263 .with_transaction(move |txn| -> Result<_> {
1264 let (observed_max_identifier, number_of_chunks) = txn
1266 .prepare(
1267 "SELECT MAX(id), COUNT(*) FROM linked_chunks WHERE linked_chunk_id = ?"
1268 )?
1269 .query_one(
1270 (&hashed_linked_chunk_id,),
1271 |row| {
1272 Ok((
1273 row.get::<_, Option<u64>>(0)?,
1278 row.get::<_, u64>(1)?,
1279 ))
1280 }
1281 )?;
1282
1283 let chunk_identifier_generator = match observed_max_identifier {
1284 Some(max_observed_identifier) => {
1285 ChunkIdentifierGenerator::new_from_previous_chunk_identifier(
1286 ChunkIdentifier::new(max_observed_identifier)
1287 )
1288 },
1289 None => ChunkIdentifierGenerator::new_from_scratch(),
1290 };
1291
1292 let Some((chunk_identifier, previous_chunk, chunk_type)) = txn
1294 .prepare(
1295 "SELECT id, previous, type FROM linked_chunks WHERE linked_chunk_id = ? AND next IS NULL"
1296 )?
1297 .query_one(
1298 (&hashed_linked_chunk_id,),
1299 |row| {
1300 Ok((
1301 row.get::<_, u64>(0)?,
1302 row.get::<_, Option<u64>>(1)?,
1303 row.get::<_, String>(2)?,
1304 ))
1305 }
1306 )
1307 .optional()?
1308 else {
1309 if number_of_chunks == 0 {
1312 return Ok((None, chunk_identifier_generator));
1313 }
1314 else {
1319 return Err(Error::InvalidData {
1320 details:
1321 "last chunk is not found but chunks exist: the linked chunk contains a cycle"
1322 .to_owned()
1323 }
1324 )
1325 }
1326 };
1327
1328 let last_chunk = txn.rebuild_chunk(
1330 &encryption,
1331 &hashed_linked_chunk_id,
1332 previous_chunk,
1333 chunk_identifier,
1334 None,
1335 &chunk_type
1336 )?;
1337
1338 Ok((Some(last_chunk), chunk_identifier_generator))
1339 })
1340 .await
1341 }
1342
1343 #[instrument(skip(self))]
1344 async fn load_previous_chunk(
1345 &self,
1346 linked_chunk_id: LinkedChunkId<'_>,
1347 before_chunk_identifier: ChunkIdentifier,
1348 ) -> Result<Option<RawChunk<Event, Gap>>, Self::Error> {
1349 let _timer = timer!("method");
1350
1351 let hashed_linked_chunk_id =
1352 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1353 let encryption = self.encryption.clone();
1354
1355 self
1356 .read()
1357 .await?
1358 .with_transaction(move |txn| -> Result<_> {
1359 let Some((chunk_identifier, previous_chunk, next_chunk, chunk_type)) = txn
1361 .prepare(
1362 "SELECT id, previous, next, type FROM linked_chunks WHERE linked_chunk_id = ? AND next = ?"
1363 )?
1364 .query_one(
1365 (&hashed_linked_chunk_id, before_chunk_identifier.index()),
1366 |row| {
1367 Ok((
1368 row.get::<_, u64>(0)?,
1369 row.get::<_, Option<u64>>(1)?,
1370 row.get::<_, Option<u64>>(2)?,
1371 row.get::<_, String>(3)?,
1372 ))
1373 }
1374 )
1375 .optional()?
1376 else {
1377 return Ok(None);
1379 };
1380
1381 let last_chunk = txn.rebuild_chunk(
1383 &encryption,
1384 &hashed_linked_chunk_id,
1385 previous_chunk,
1386 chunk_identifier,
1387 next_chunk,
1388 &chunk_type
1389 )?;
1390
1391 Ok(Some(last_chunk))
1392 })
1393 .await
1394 }
1395
1396 async fn load_thread_info(
1397 &self,
1398 room_id: &RoomId,
1399 thread_id: &EventId,
1400 ) -> Result<ThreadInfo, Self::Error> {
1401 let linked_chunk_id = LinkedChunkId::Thread(room_id, thread_id);
1402 let hashed_linked_chunk_id =
1403 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1404 let encryption = self.encryption.clone();
1405
1406 let maybe_thread_info = self
1411 .read()
1412 .await?
1413 .with_transaction(move |txn| {
1414 let maybe_encoded_thread_info = txn
1415 .query_one(
1416 "SELECT info FROM threads WHERE linked_chunk_id = ?",
1417 (hashed_linked_chunk_id,),
1418 |row| row.get::<_, Vec<u8>>(0),
1419 )
1420 .optional()?;
1421
1422 maybe_encoded_thread_info
1423 .map(|encoded_thread_info| {
1424 encryption.decode_thread_info(encoded_thread_info.as_slice())
1425 })
1426 .transpose()
1427 })
1428 .await?;
1429
1430 if let Some(thread_info) = maybe_thread_info {
1431 return Ok(thread_info);
1432 }
1433
1434 let thread_info = ThreadInfo::new();
1437 let hashed_linked_chunk_id =
1438 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1439 let hashed_room_id = self.encryption.encode_room_id(keys::EVENTS, room_id);
1440 let hashed_thread_id = self.encryption.encode_thread_id(thread_id)?;
1441 let encoded_thread_id = self.encryption.encode_thread_info(&thread_info)?;
1442
1443 self.write()
1444 .await?
1445 .with_transaction(move |txn| {
1446 txn.execute(
1447 "INSERT INTO threads VALUES (?, ?, ?, ?)",
1448 (hashed_linked_chunk_id, hashed_room_id, hashed_thread_id, encoded_thread_id),
1449 )?;
1450
1451 Ok::<(), Self::Error>(())
1452 })
1453 .await?;
1454
1455 Ok(thread_info)
1456 }
1457
1458 async fn update_thread_info(
1459 &self,
1460 room_id: &RoomId,
1461 thread_id: &EventId,
1462 thread_info: &ThreadInfo,
1463 ) -> Result<(), Self::Error> {
1464 let hashed_linked_chunk_id = self
1465 .encryption
1466 .encode_linked_chunk(keys::LINKED_CHUNKS, &LinkedChunkId::Thread(room_id, thread_id));
1467 let encoded_thread_info = self.encryption.encode_thread_info(thread_info)?;
1468
1469 self.write()
1470 .await?
1471 .with_transaction(move |txn| {
1472 txn.execute(
1473 "UPDATE threads SET info = ? WHERE linked_chunk_id = ?",
1474 (encoded_thread_info, hashed_linked_chunk_id),
1475 )?;
1476
1477 Ok(())
1478 })
1479 .await
1480 }
1481
1482 #[instrument(skip(self))]
1483 async fn clear_all_events(&self, room_id: Option<&RoomId>) -> Result<(), Self::Error> {
1484 let _timer = timer!("method");
1485
1486 match room_id {
1487 None => {
1489 self.write()
1490 .await?
1491 .with_transaction(move |txn| {
1492 txn.execute("DELETE FROM linked_chunks", ())?;
1494
1495 txn.execute("DELETE FROM events", ())?;
1497
1498 Ok(())
1499 })
1500 .await
1501 }
1502
1503 Some(room_id) => {
1505 let encryption = self.encryption.clone();
1506 let room_id = room_id.to_owned();
1507
1508 self.write()
1509 .await?
1510 .with_transaction(move |txn| {
1511 {
1513 let mut delete =
1514 txn.prepare("DELETE FROM linked_chunks WHERE linked_chunk_id = ?")?;
1515
1516 for linked_chunk_id in
1517 [LinkedChunkId::Room(&room_id), LinkedChunkId::PinnedEvents(&room_id)]
1518 {
1519 let linked_chunk_id = encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1520
1521 delete.execute((&linked_chunk_id,))?;
1524 }
1525 }
1526
1527 let encoded_room_id = encryption.encode_room_id(keys::EVENTS, &room_id);
1528
1529 {
1531 txn.execute(
1532 "DELETE FROM linked_chunks WHERE linked_chunk_id IN (SELECT linked_chunk_id FROM threads WHERE room_id = ?)",
1533 (&encoded_room_id,),
1534 )?;
1535 }
1536
1537 txn.execute(
1539 "DELETE FROM events WHERE room_id = ?",
1540 (encoded_room_id,),
1541 )?;
1542
1543 Ok(())
1544 })
1545 .await
1546 }
1547 }
1548 }
1549
1550 #[instrument(skip(self, event_ids))]
1551 async fn filter_duplicated_events(
1552 &self,
1553 linked_chunk_id: LinkedChunkId<'_>,
1554 event_ids: Vec<OwnedEventId>,
1555 ) -> Result<Vec<(OwnedEventId, Position)>, Self::Error> {
1556 let _timer = timer!("method");
1557
1558 if event_ids.is_empty() {
1562 return Ok(Vec::new());
1563 }
1564
1565 let hashed_linked_chunk_id =
1567 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &linked_chunk_id);
1568 let event_ids_and_hashed_event_ids = event_ids
1569 .into_iter()
1570 .map(|event_id| {
1571 let hashed_event_id = self.encryption.encode_event_id(
1572 keys::EVENTS,
1577 &event_id,
1578 );
1579
1580 (event_id, hashed_event_id)
1581 })
1582 .collect::<Vec<_>>();
1583
1584 self.read()
1585 .await?
1586 .with_transaction(move |txn| -> Result<_> {
1587 txn.chunk_large_query_over(
1588 event_ids_and_hashed_event_ids,
1589 None,
1590 move |txn, event_ids_and_hashed_event_ids| {
1591 let query = format!(
1592 "SELECT event_id, chunk_id, position \
1593 FROM event_chunks \
1594 WHERE linked_chunk_id = ? AND event_id IN ({}) \
1595 ORDER BY chunk_id ASC, position ASC",
1596 event_ids_and_hashed_event_ids.host_parameters(),
1597 );
1598
1599 let parameters = params_from_iter(
1600 once(
1602 hashed_linked_chunk_id
1603 .to_sql()
1604 .unwrap(),
1606 )
1607 .chain(
1609 event_ids_and_hashed_event_ids.iter().map(
1610 |(_event_id, hashed_event_id)| {
1611 hashed_event_id
1612 .to_sql()
1613 .unwrap()
1616 },
1617 ),
1618 ),
1619 );
1620
1621 let mut duplicated_events = Vec::new();
1622
1623 for duplicated_event in
1624 txn.prepare(&query)?.query_map(parameters, |row| {
1625 Ok((
1626 row.get::<_, Vec<u8>>(0)?,
1627 row.get::<_, u64>(1)?,
1628 row.get::<_, usize>(2)?,
1629 ))
1630 })?
1631 {
1632 let (duplicated_hashed_event_id, chunk_identifier, index) =
1633 duplicated_event?;
1634
1635 let Some(duplicated_event_id) = event_ids_and_hashed_event_ids
1639 .iter()
1640 .find_map(|(event_id, hashed_event_id)| {
1641 (hashed_event_id.deref() == duplicated_hashed_event_id)
1642 .then_some(event_id.clone())
1643 })
1644 else {
1645 error!(
1646 "Unreachable: found a duplicated event that was not requested"
1647 );
1648 continue;
1649 };
1650
1651 duplicated_events.push((
1652 duplicated_event_id,
1653 Position::new(ChunkIdentifier::new(chunk_identifier), index),
1654 ));
1655 }
1656
1657 Ok(duplicated_events)
1658 },
1659 )
1660 })
1661 .await
1662 }
1663
1664 #[instrument(skip(self, event_id))]
1665 async fn find_event(
1666 &self,
1667 room_id: &RoomId,
1668 event_id: &EventId,
1669 ) -> Result<Option<Event>, Self::Error> {
1670 let _timer = timer!("method");
1671
1672 let encryption = self.encryption.clone();
1673
1674 let hashed_room_id = self.encryption.encode_room_id(keys::EVENTS, room_id);
1675 let hashed_event_id = self.encryption.encode_event_id(keys::EVENTS, event_id);
1676
1677 self.read()
1678 .await?
1679 .with_transaction(move |txn| -> Result<_> {
1680 let Some(event) = txn
1681 .prepare("SELECT content FROM events WHERE event_id = ? AND room_id = ?")?
1682 .query_one((hashed_event_id, hashed_room_id), |row| row.get::<_, Vec<u8>>(0))
1683 .optional()?
1684 else {
1685 return Ok(None);
1687 };
1688
1689 Ok(Some(encryption.decode_event(&event)?))
1690 })
1691 .await
1692 }
1693
1694 #[instrument(skip(self, event_id, filters))]
1695 async fn find_event_relations(
1696 &self,
1697 room_id: &RoomId,
1698 event_id: &EventId,
1699 filters: Option<&[RelationType]>,
1700 ) -> Result<Vec<(Event, Option<Position>)>, Self::Error> {
1701 let _timer = timer!("method");
1702
1703 let hashed_room_id = self.encryption.encode_room_id(keys::EVENTS, room_id);
1704 let hashed_linked_chunk_id =
1705 self.encryption.encode_linked_chunk(keys::LINKED_CHUNKS, &LinkedChunkId::Room(room_id));
1706 let hashed_event_id = self.encryption.encode_event_id(keys::EVENTS, event_id);
1707
1708 let filters = filters.map(ToOwned::to_owned);
1709 let encryption = self.encryption.clone();
1710
1711 self.read()
1712 .await?
1713 .with_transaction(move |txn| -> Result<_> {
1714 find_event_relations_transaction(
1715 &encryption,
1716 hashed_room_id,
1717 hashed_linked_chunk_id,
1718 hashed_event_id,
1719 filters,
1720 txn,
1721 )
1722 })
1723 .await
1724 }
1725
1726 #[instrument(skip(self))]
1727 async fn get_room_events(
1728 &self,
1729 room_id: &RoomId,
1730 event_type: Option<&str>,
1731 session_id: Option<&str>,
1732 ) -> Result<Vec<Event>, Self::Error> {
1733 let _timer = timer!("method");
1734
1735 let encryption = self.encryption.clone();
1736
1737 let hashed_room_id = self.encryption.encode_room_id(keys::EVENTS, room_id);
1738 let hashed_event_type = event_type.map(|e| self.encryption.encode_key(keys::EVENTS, e));
1739 let hashed_session_id = session_id.map(|s| self.encryption.encode_key(keys::EVENTS, s));
1740
1741 self.read()
1742 .await?
1743 .with_transaction(move |txn| -> Result<_> {
1744 #[allow(clippy::redundant_clone)]
1748 let (query, keys) = match (hashed_event_type, hashed_session_id) {
1749 (None, None) => {
1750 ("SELECT content FROM events WHERE room_id = ?", params![hashed_room_id])
1751 }
1752 (None, Some(session_id)) => (
1753 "SELECT content FROM events WHERE room_id = ?1 AND session_id = ?2",
1754 params![hashed_room_id, session_id.to_owned()],
1755 ),
1756 (Some(event_type), None) => (
1757 "SELECT content FROM events WHERE room_id = ? AND event_type = ?",
1758 params![hashed_room_id, event_type.to_owned()]
1759 ),
1760 (Some(event_type), Some(session_id)) => (
1761 "SELECT content FROM events WHERE room_id = ?1 AND event_type = ?2 AND session_id = ?3",
1762 params![hashed_room_id, event_type.to_owned(), session_id.to_owned()],
1763 ),
1764 };
1765
1766 let mut statement = txn.prepare(query)?;
1767
1768 statement
1769 .query_map(keys, |row| row.get::<_, Vec<u8>>(0))?
1770 .map(|maybe_encoded_event| {
1771 encryption.decode_event(&maybe_encoded_event?)
1772 })
1773 .collect::<Result<Vec<_>>>()
1774 })
1775 .await
1776 }
1777
1778 #[instrument(skip(self, event))]
1779 async fn save_event(&self, room_id: &RoomId, event: Event) -> Result<(), Self::Error> {
1780 let _timer = timer!("method");
1781
1782 let Some(event_id) = event.event_id() else {
1783 error!("Trying to save an event with no ID");
1784 return Ok(());
1785 };
1786
1787 let Some(event_type) = event.kind.event_type() else {
1788 error!(%event_id, "Trying to save an event with no event type");
1789 return Ok(());
1790 };
1791
1792 let hashed_event_type = self.encryption.encode_key(keys::EVENTS, event_type);
1793 let hashed_session_id =
1794 event.kind.session_id().map(|s| self.encryption.encode_key(keys::EVENTS, s));
1795
1796 let hashed_room_id = self.encryption.encode_room_id(keys::EVENTS, room_id);
1797 let hashed_event_id = self.encryption.encode_event_id(keys::EVENTS, event_id);
1798 let encoded_event = self.encryption.encode_event(&event)?;
1799
1800 self.write()
1801 .await?
1802 .with_transaction(move |txn| -> Result<_> {
1803 txn.execute(
1804 "INSERT OR REPLACE INTO events(room_id, event_id, event_type, session_id, content, relates_to, rel_type) VALUES (?, ?, ?, ?, ?, ?, ?)",
1805 (
1806 &hashed_room_id,
1807 hashed_event_id,
1808 hashed_event_type,
1809 hashed_session_id,
1810 encoded_event.content,
1811 encoded_event.relates_to,
1812 encoded_event.rel_type
1813 )
1814 )?;
1815
1816 Ok(())
1817 })
1818 .await
1819 }
1820
1821 async fn close(&self) -> Result<(), Self::Error> {
1822 SqliteEventCacheStore::close(self).await
1823 }
1824
1825 async fn reopen(&self) -> Result<(), Self::Error> {
1826 SqliteEventCacheStore::reopen(self).await
1827 }
1828
1829 async fn optimize(&self) -> Result<(), Self::Error> {
1830 Ok(self.vacuum().await?)
1831 }
1832
1833 async fn get_size(&self) -> Result<Option<usize>, Self::Error> {
1834 self.get_db_size().await
1835 }
1836}
1837
1838fn find_event_relations_transaction(
1839 encryption: &Encryption,
1840 hashed_room_id: Key,
1841 hashed_linked_chunk_id: Key,
1842 hashed_event_id: Key,
1843 filters: Option<Vec<RelationType>>,
1844 txn: &Transaction<'_>,
1845) -> Result<Vec<(Event, Option<Position>)>> {
1846 let get_rows = |row: &rusqlite::Row<'_>| {
1847 Ok((
1848 row.get::<_, Vec<u8>>(0)?,
1849 row.get::<_, Option<u64>>(1)?,
1850 row.get::<_, Option<usize>>(2)?,
1851 ))
1852 };
1853
1854 let collect_results = |transaction| {
1856 let mut related = Vec::new();
1857
1858 for result in transaction {
1859 let (event, chunk_id, index): (Vec<u8>, Option<u64>, _) = result?;
1860 let event = encryption.decode_event(&event)?;
1861
1862 let pos = chunk_id
1865 .zip(index)
1866 .map(|(chunk_id, index)| Position::new(ChunkIdentifier::new(chunk_id), index));
1867
1868 related.push((event, pos));
1869 }
1870
1871 Ok(related)
1872 };
1873
1874 if let Some(filters) = filters
1875 && filters.is_empty().not()
1876 {
1877 let query = format!(
1878 "SELECT events.content, event_chunks.chunk_id, event_chunks.position \
1879 FROM events \
1880 LEFT JOIN event_chunks ON events.event_id = event_chunks.event_id AND event_chunks.linked_chunk_id = ? \
1881 WHERE events.relates_to = ? AND events.room_id = ? AND events.rel_type IN ({})",
1882 host_parameters(filters.len())
1883 );
1884
1885 let filter_strings: Vec<_> = filters.iter().map(|f| f.to_string()).collect();
1890 let filters_params: Vec<_> = filter_strings
1891 .iter()
1892 .map(|f| f.to_sql().expect("converting a string to SQL should work"))
1893 .collect();
1894
1895 let parameters = params_from_iter(
1896 [
1897 hashed_linked_chunk_id.to_sql().expect(
1898 "We should be able to convert a hashed linked chunk ID to a SQLite value",
1899 ),
1900 hashed_event_id
1901 .to_sql()
1902 .expect("We should be able to convert an event ID to a SQLite value"),
1903 hashed_room_id
1904 .to_sql()
1905 .expect("We should be able to convert a room ID to a SQLite value"),
1906 ]
1907 .into_iter()
1908 .chain(filters_params),
1909 );
1910
1911 let mut transaction = txn.prepare(&query)?;
1912 let transaction = transaction.query_map(parameters, get_rows)?;
1913
1914 collect_results(transaction)
1915 } else {
1916 let query = "SELECT events.content, event_chunks.chunk_id, event_chunks.position \
1917 FROM events \
1918 LEFT JOIN event_chunks ON events.event_id = event_chunks.event_id AND event_chunks.linked_chunk_id = ? \
1919 WHERE events.relates_to = ? AND events.room_id = ?";
1920 let parameters = (hashed_linked_chunk_id, hashed_event_id, hashed_room_id);
1921
1922 let mut transaction = txn.prepare(query)?;
1923 let transaction = transaction.query_map(parameters, get_rows)?;
1924
1925 collect_results(transaction)
1926 }
1927}
1928
1929async fn with_immediate_transaction<
1934 T: Send + 'static,
1935 F: FnOnce(&Transaction<'_>) -> Result<T, Error> + Send + 'static,
1936>(
1937 this: &SqliteEventCacheStore,
1938 f: F,
1939) -> Result<T, Error> {
1940 this.write()
1941 .await?
1942 .interact(move |conn| -> Result<T, Error> {
1943 conn.set_transaction_behavior(TransactionBehavior::Immediate);
1947
1948 let code = || -> Result<T, Error> {
1949 let txn = conn.transaction()?;
1950 let res = f(&txn)?;
1951 txn.commit()?;
1952 Ok(res)
1953 };
1954
1955 let res = code();
1956
1957 conn.set_transaction_behavior(TransactionBehavior::Deferred);
1960
1961 res
1962 })
1963 .await
1964 .unwrap()
1966}
1967
1968fn insert_chunk(
1969 txn: &Transaction<'_>,
1970 linked_chunk_id: &Key,
1971 previous: Option<u64>,
1972 new: u64,
1973 next: Option<u64>,
1974 type_str: &str,
1975) -> rusqlite::Result<()> {
1976 txn.execute(
1978 r#"
1979 INSERT INTO linked_chunks(id, linked_chunk_id, previous, next, type)
1980 VALUES (?, ?, ?, ?, ?)
1981 "#,
1982 (new, linked_chunk_id, previous, next, type_str),
1983 )?;
1984
1985 if let Some(previous) = previous {
1987 let updated = txn.execute(
1988 r#"
1989 UPDATE linked_chunks
1990 SET next = ?
1991 WHERE id = ? AND linked_chunk_id = ?
1992 "#,
1993 (new, previous, linked_chunk_id),
1994 )?;
1995 if updated < 1 {
1996 return Err(rusqlite::Error::QueryReturnedNoRows);
1997 }
1998 if updated > 1 {
1999 return Err(rusqlite::Error::QueryReturnedMoreThanOneRow);
2000 }
2001 }
2002
2003 if let Some(next) = next {
2005 let updated = txn.execute(
2006 r#"
2007 UPDATE linked_chunks
2008 SET previous = ?
2009 WHERE id = ? AND linked_chunk_id = ?
2010 "#,
2011 (new, next, linked_chunk_id),
2012 )?;
2013 if updated < 1 {
2014 return Err(rusqlite::Error::QueryReturnedNoRows);
2015 }
2016 if updated > 1 {
2017 return Err(rusqlite::Error::QueryReturnedMoreThanOneRow);
2018 }
2019 }
2020
2021 Ok(())
2022}
2023
2024#[cfg(test)]
2025mod tests {
2026 use std::{
2027 path::PathBuf,
2028 sync::{
2029 LazyLock,
2030 atomic::{AtomicU32, Ordering::SeqCst},
2031 },
2032 };
2033
2034 use assert_matches::assert_matches;
2035 use matrix_sdk_base::{
2036 event_cache::store::{
2037 EventCacheStore, EventCacheStoreError, IntoEventCacheStore,
2038 integration_tests::EventCacheStoreIntegrationTests,
2039 },
2040 event_cache_store_integration_tests, event_cache_store_integration_tests_time,
2041 linked_chunk::{ChunkIdentifier, LinkedChunkId, Update},
2042 };
2043 use matrix_sdk_test::{DEFAULT_TEST_ROOM_ID, async_test};
2044 use ruma::{OwnedEventId, event_id};
2045 use tempfile::{TempDir, tempdir};
2046
2047 use super::{SqliteEventCacheStore, keys};
2048 use crate::{SqliteStoreConfig, utils::SqliteAsyncConnExt};
2049
2050 static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
2051 static NUM: AtomicU32 = AtomicU32::new(0);
2052
2053 fn new_event_cache_store_workspace() -> PathBuf {
2054 let name = NUM.fetch_add(1, SeqCst).to_string();
2055 TMP_DIR.path().join(name)
2056 }
2057
2058 async fn get_event_cache_store() -> Result<SqliteEventCacheStore, EventCacheStoreError> {
2059 let tmpdir_path = new_event_cache_store_workspace();
2060
2061 tracing::info!("using event cache store @ {}", tmpdir_path.to_str().unwrap());
2062
2063 Ok(SqliteEventCacheStore::open(tmpdir_path.to_str().unwrap(), None).await.unwrap())
2064 }
2065
2066 event_cache_store_integration_tests!();
2067 event_cache_store_integration_tests_time!();
2068
2069 #[async_test]
2070 async fn test_encryption_encode_decode_thread_id_roundtrip() {
2071 let store = get_event_cache_store().await.expect("creating cache store failed");
2072 let thread_id: OwnedEventId = event_id!("$event").to_owned();
2073
2074 let encoded_thread_id = store.encryption.encode_thread_id(&thread_id).unwrap();
2075 let decoded_thread_id: OwnedEventId =
2076 store.encryption.decode_thread_id(&encoded_thread_id).unwrap();
2077
2078 assert_eq!(thread_id, decoded_thread_id);
2079 }
2080
2081 #[async_test]
2082 async fn test_pool_size() {
2083 let tmpdir_path = new_event_cache_store_workspace();
2084 let store_open_config = SqliteStoreConfig::new(tmpdir_path).pool_max_size(42);
2085
2086 let store = SqliteEventCacheStore::open_with_config(&store_open_config).await.unwrap();
2087
2088 assert_eq!(store.pool_max_size().await, Some(42));
2089 }
2090
2091 #[async_test]
2092 async fn test_linked_chunk_remove_chunk() {
2093 let store = get_event_cache_store().await.expect("creating cache store failed");
2094
2095 store.clone().into_event_cache_store().test_linked_chunk_remove_chunk().await;
2097
2098 let gaps = store
2100 .read()
2101 .await
2102 .unwrap()
2103 .with_transaction(|txn| -> rusqlite::Result<_> {
2104 let mut gaps = Vec::new();
2105 for data in txn
2106 .prepare("SELECT chunk_id FROM gap_chunks ORDER BY chunk_id")?
2107 .query_map((), |row| row.get::<_, u64>(0))?
2108 {
2109 gaps.push(data?);
2110 }
2111 Ok(gaps)
2112 })
2113 .await
2114 .unwrap();
2115
2116 assert_eq!(gaps, vec![42, 44]);
2119 }
2120
2121 #[async_test]
2122 async fn test_linked_chunk_remove_item() {
2123 let store = get_event_cache_store().await.expect("creating cache store failed");
2124
2125 store.clone().into_event_cache_store().test_linked_chunk_remove_item().await;
2127
2128 let room_id = *DEFAULT_TEST_ROOM_ID;
2129 let hashed_linked_chunk_id = store
2130 .encryption
2131 .encode_linked_chunk(keys::LINKED_CHUNKS, &LinkedChunkId::Room(room_id));
2132
2133 let num_rows: u64 = store
2135 .read()
2136 .await
2137 .unwrap()
2138 .with_transaction(move |txn| {
2139 txn.query_one(
2140 "SELECT COUNT(*) FROM event_chunks WHERE chunk_id = 42 AND linked_chunk_id = ? AND position IN (2, 3, 4)",
2141 (hashed_linked_chunk_id,),
2142 |row| row.get(0),
2143 )
2144 })
2145 .await
2146 .unwrap();
2147 assert_eq!(num_rows, 3);
2148 }
2149
2150 #[async_test]
2151 async fn test_linked_chunk_clear() {
2152 let store = get_event_cache_store().await.expect("creating cache store failed");
2153
2154 store.clone().into_event_cache_store().test_linked_chunk_clear().await;
2156
2157 store
2159 .read()
2160 .await
2161 .unwrap()
2162 .with_transaction(|txn| -> rusqlite::Result<_> {
2163 let num_gaps = txn
2164 .prepare("SELECT COUNT(chunk_id) FROM gap_chunks ORDER BY chunk_id")?
2165 .query_one((), |row| row.get::<_, u64>(0))?;
2166 assert_eq!(num_gaps, 0);
2167
2168 let num_events = txn
2169 .prepare("SELECT COUNT(event_id) FROM event_chunks ORDER BY chunk_id")?
2170 .query_one((), |row| row.get::<_, u64>(0))?;
2171 assert_eq!(num_events, 0);
2172
2173 Ok(())
2174 })
2175 .await
2176 .unwrap();
2177 }
2178
2179 #[async_test]
2180 async fn test_linked_chunk_update_is_a_transaction() {
2181 let store = get_event_cache_store().await.expect("creating cache store failed");
2182
2183 let room_id = *DEFAULT_TEST_ROOM_ID;
2184 let linked_chunk_id = LinkedChunkId::Room(room_id);
2185
2186 let err = store
2189 .handle_linked_chunk_updates(
2190 linked_chunk_id,
2191 vec![
2192 Update::NewItemsChunk {
2193 previous: None,
2194 new: ChunkIdentifier::new(42),
2195 next: None,
2196 },
2197 Update::NewItemsChunk {
2198 previous: None,
2199 new: ChunkIdentifier::new(42),
2200 next: None,
2201 },
2202 ],
2203 )
2204 .await
2205 .unwrap_err();
2206
2207 assert_matches!(err, crate::error::Error::Sqlite(err) => {
2209 assert_matches!(err.sqlite_error_code(), Some(rusqlite::ErrorCode::ConstraintViolation));
2210 });
2211
2212 let chunks = store.load_all_chunks(linked_chunk_id).await.unwrap();
2216 assert!(chunks.is_empty());
2217 }
2218}
2219
2220#[cfg(test)]
2221mod encrypted_tests {
2222 use std::sync::{
2223 LazyLock,
2224 atomic::{AtomicU32, Ordering::SeqCst},
2225 };
2226
2227 use matrix_sdk_base::{
2228 event_cache::store::{EventCacheStore, EventCacheStoreError},
2229 event_cache_store_integration_tests, event_cache_store_integration_tests_time,
2230 };
2231 use matrix_sdk_test::{async_test, event_factory::EventFactory};
2232 use ruma::{
2233 event_id,
2234 events::{relation::RelationType, room::message::RoomMessageEventContentWithoutRelation},
2235 room_id, user_id,
2236 };
2237 use tempfile::{TempDir, tempdir};
2238
2239 use super::SqliteEventCacheStore;
2240
2241 static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
2242 static NUM: AtomicU32 = AtomicU32::new(0);
2243
2244 async fn get_event_cache_store() -> Result<SqliteEventCacheStore, EventCacheStoreError> {
2245 let name = NUM.fetch_add(1, SeqCst).to_string();
2246 let tmpdir_path = TMP_DIR.path().join(name);
2247
2248 tracing::info!("using event cache store @ {}", tmpdir_path.to_str().unwrap());
2249
2250 Ok(SqliteEventCacheStore::open(
2251 tmpdir_path.to_str().unwrap(),
2252 Some("default_test_password"),
2253 )
2254 .await
2255 .unwrap())
2256 }
2257
2258 event_cache_store_integration_tests!();
2259 event_cache_store_integration_tests_time!();
2260
2261 #[async_test]
2262 async fn test_no_sqlite_injection_in_find_event_relations() {
2263 let room_id = room_id!("!test:localhost");
2264 let another_room_id = room_id!("!r1:matrix.org");
2265 let sender = user_id!("@alice:localhost");
2266
2267 let store = get_event_cache_store()
2268 .await
2269 .expect("We should be able to create a new, empty, event cache store");
2270
2271 let f = EventFactory::new().room(room_id).sender(sender);
2272
2273 let event_id = event_id!("$DO_NOT_FIND_ME:matrix.org");
2275 let event = f.text_msg("DO NOT FIND").event_id(event_id).into_event();
2276
2277 let edit_id = event_id!("$find_me:matrix.org");
2279 let edit = f
2280 .text_msg("Find me")
2281 .event_id(edit_id)
2282 .edit(event_id, RoomMessageEventContentWithoutRelation::text_plain("jebote"))
2283 .into_event();
2284
2285 let f = f.room(another_room_id);
2287
2288 let another_event_id = event_id!("$DO_NOT_FIND_ME_EITHER:matrix.org");
2289 let another_event =
2290 f.text_msg("DO NOT FIND ME EITHER").event_id(another_event_id).into_event();
2291
2292 store.save_event(room_id, event).await.unwrap();
2294 store.save_event(room_id, edit).await.unwrap();
2295 store.save_event(another_room_id, another_event).await.unwrap();
2296
2297 let filter = Some(vec![RelationType::Replacement, "x\") OR 1=1; --".into()]);
2301
2302 let results = store
2304 .find_event_relations(room_id, event_id, filter.as_deref())
2305 .await
2306 .expect("We should be able to attempt to find event relations");
2307
2308 similar_asserts::assert_eq!(
2310 results.len(),
2311 1,
2312 "We should only have loaded events for the first room {results:#?}"
2313 );
2314
2315 let (found_event, _) = &results[0];
2317 assert_eq!(
2318 found_event.event_id(),
2319 Some(edit_id),
2320 "The single event we found should be the edit event"
2321 );
2322 }
2323}
2324
2325#[cfg(test)]
2326mod close_reopen_tests {
2327 use std::sync::{
2328 LazyLock,
2329 atomic::{AtomicU32, Ordering::SeqCst},
2330 };
2331
2332 use matrix_sdk_base::{event_cache::store::EventCacheStore, linked_chunk::LinkedChunkId};
2333 use matrix_sdk_test::{DEFAULT_TEST_ROOM_ID, async_test};
2334 use tempfile::{TempDir, tempdir};
2335
2336 use super::SqliteEventCacheStore;
2337
2338 static TMP_DIR: LazyLock<TempDir> = LazyLock::new(|| tempdir().unwrap());
2339 static NUM: AtomicU32 = AtomicU32::new(0);
2340
2341 async fn new_store() -> SqliteEventCacheStore {
2342 let name = NUM.fetch_add(1, SeqCst).to_string();
2343 let tmpdir_path = TMP_DIR.path().join(name);
2344 SqliteEventCacheStore::open(tmpdir_path, None).await.unwrap()
2345 }
2346
2347 #[async_test]
2348 async fn test_close_completes_without_timeout() {
2349 let store = new_store().await;
2350
2351 let start = std::time::Instant::now();
2353 store.close().await.unwrap();
2354 let elapsed = start.elapsed();
2355
2356 assert!(
2357 elapsed < std::time::Duration::from_secs(2),
2358 "close() took {elapsed:?}, expected < 2s (no timeout)"
2359 );
2360
2361 let guard = store.connections.lock().await;
2363 assert!(guard.is_none(), "connections should be None after close");
2364 }
2365
2366 #[async_test]
2367 async fn test_reopen_restores_connections() {
2368 let store = new_store().await;
2369
2370 store.close().await.unwrap();
2371
2372 {
2373 let guard = store.connections.lock().await;
2374 assert!(guard.is_none());
2375 }
2376
2377 store.reopen().await.unwrap();
2378
2379 {
2380 let guard = store.connections.lock().await;
2381 assert!(guard.is_some(), "connections should be Some after reopen");
2382 }
2383 }
2384
2385 #[async_test]
2386 async fn test_close_is_idempotent() {
2387 let store = new_store().await;
2388
2389 store.close().await.unwrap();
2390 store.close().await.unwrap();
2392
2393 let guard = store.connections.lock().await;
2394 assert!(guard.is_none());
2395 }
2396
2397 #[async_test]
2398 async fn test_reopen_is_idempotent() {
2399 let store = new_store().await;
2400
2401 store.reopen().await.unwrap();
2403
2404 let guard = store.connections.lock().await;
2405 assert!(guard.is_some());
2406 }
2407
2408 #[async_test]
2409 async fn test_read_fails_when_closed() {
2410 let store = new_store().await;
2411 store.close().await.unwrap();
2412
2413 let err = store.load_all_chunks(LinkedChunkId::Room(*DEFAULT_TEST_ROOM_ID)).await;
2414 assert!(err.is_err(), "read should fail when closed");
2415
2416 let err_msg = err.unwrap_err().to_string();
2417 assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
2418 }
2419
2420 #[async_test]
2421 async fn test_write_fails_when_closed() {
2422 let store = new_store().await;
2423 store.close().await.unwrap();
2424
2425 let err = store.try_take_leased_lock(1000, "test_lock", "holder").await;
2426 assert!(err.is_err(), "write should fail when closed");
2427
2428 let err_msg = err.unwrap_err().to_string();
2429 assert!(err_msg.contains("closed"), "error should mention 'closed', got: {err_msg}");
2430 }
2431
2432 #[async_test]
2433 async fn test_data_persists_across_close_reopen() {
2434 let store = new_store().await;
2435
2436 let result = store.try_take_leased_lock(60_000, "test_lock", "holder").await.unwrap();
2438 assert!(result.is_some(), "should have acquired the lock");
2439
2440 store.close().await.unwrap();
2442 store.reopen().await.unwrap();
2443
2444 let result = store.try_take_leased_lock(60_000, "test_lock", "other_holder").await.unwrap();
2446 assert!(result.is_none(), "lock should still be held by the original holder after reopen");
2447 }
2448
2449 #[async_test]
2450 async fn test_multiple_close_reopen_cycles() {
2451 let store = new_store().await;
2452
2453 for _ in 0..5 {
2454 store.close().await.unwrap();
2455 store.reopen().await.unwrap();
2456
2457 let result = store.load_all_chunks(LinkedChunkId::Room(*DEFAULT_TEST_ROOM_ID)).await;
2459 assert!(result.is_ok(), "store should work after close/reopen cycle");
2460 }
2461 }
2462
2463 #[async_test]
2464 async fn test_pool_is_fully_drained_after_close() {
2465 let store = new_store().await;
2466
2467 let _ = store.load_all_chunks(LinkedChunkId::Room(*DEFAULT_TEST_ROOM_ID)).await;
2469 let _ = store.load_all_chunks(LinkedChunkId::Room(*DEFAULT_TEST_ROOM_ID)).await;
2470
2471 store.close().await.unwrap();
2472
2473 let guard = store.connections.lock().await;
2476 assert!(guard.is_none(), "all connections should be released after close");
2477 }
2478
2479 #[async_test]
2480 async fn test_operations_work_immediately_after_reopen() {
2481 let store = new_store().await;
2482
2483 store.close().await.unwrap();
2484 store.reopen().await.unwrap();
2485
2486 let result = store.load_all_chunks(LinkedChunkId::Room(*DEFAULT_TEST_ROOM_ID)).await;
2488 assert!(result.is_ok(), "read should succeed immediately after reopen");
2489
2490 let result = store.try_take_leased_lock(1000, "test_lock", "holder").await;
2492 assert!(result.is_ok(), "write should succeed immediately after reopen");
2493 }
2494
2495 #[async_test]
2496 async fn test_close_waits_for_held_read_connection_to_drain() {
2497 let store = new_store().await;
2498
2499 let held_conn = store.read().await.unwrap();
2501
2502 let store_clone = store.clone();
2505 let close_handle = tokio::spawn(async move {
2506 store_clone.close().await.unwrap();
2507 });
2508
2509 tokio::time::sleep(std::time::Duration::from_millis(200)).await;
2511
2512 assert!(!close_handle.is_finished(), "close should be waiting for the held connection");
2514
2515 drop(held_conn);
2517
2518 let timeout = tokio::time::timeout(std::time::Duration::from_secs(3), close_handle).await;
2520 assert!(timeout.is_ok(), "close should complete after the held connection is released");
2521 timeout.unwrap().unwrap();
2522
2523 let guard = store.connections.lock().await;
2525 assert!(guard.is_none(), "connections should be None after close");
2526 }
2527}