Skip to main content

matrix_sdk_base/event_cache/store/
controlled_memory_store.rs

1// Copyright 2026 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 [`MemoryStore`] whose writes can be blocked or made to fail, for tests.
16
17use std::{
18    sync::{
19        Arc,
20        atomic::{AtomicBool, Ordering},
21    },
22    time::Duration,
23};
24
25use async_trait::async_trait;
26use matrix_sdk_common::{
27    cross_process_lock::CrossProcessLockGeneration,
28    linked_chunk::{
29        ChunkIdentifier, ChunkIdentifierGenerator, ChunkMetadata, LinkedChunkId, Position,
30        RawChunk, Update,
31    },
32    sleep::sleep,
33};
34use ruma::{EventId, OwnedEventId, RoomId, events::relation::RelationType};
35use tokio::sync::Notify;
36
37use super::{
38    super::{Event, Gap, thread::ThreadInfo},
39    EventCacheStore, EventCacheStoreError, MemoryStore, Result,
40};
41
42/// A [`MemoryStore`] whose linked chunk writes can be blocked.
43///
44/// This is meant for tests that need to control when, or whether, the event
45/// cache persists its data: for instance to check what the rest of the SDK
46/// does while the store is slow, or when it returns an error.
47///
48/// Only [`EventCacheStore::handle_linked_chunk_updates`] is affected. All the
49/// other methods are forwarded to the inner [`MemoryStore`] as is.
50#[derive(Debug, Clone, Default)]
51pub struct ControlledMemoryStore {
52    inner: MemoryStore,
53    writes_blocked: Arc<AtomicBool>,
54    write_started: Arc<Notify>,
55    write_finished: Arc<Notify>,
56}
57
58impl ControlledMemoryStore {
59    /// Create a new store, which doesn't block writes.
60    pub fn new() -> Self {
61        Self::default()
62    }
63
64    /// Block linked chunk writes until [`Self::unblock_writes`] is called.
65    ///
66    /// A blocked write still counts as started for
67    /// [`Self::wait_for_write_to_start`].
68    pub fn block_writes(&self) {
69        self.writes_blocked.store(true, Ordering::SeqCst);
70    }
71
72    /// Let the blocked linked chunk writes, and the next ones, go through.
73    pub fn unblock_writes(&self) {
74        self.writes_blocked.store(false, Ordering::SeqCst);
75    }
76
77    /// Wait until a linked chunk write starts.
78    ///
79    /// Returns at once if a write has started since the last call.
80    pub async fn wait_for_write_to_start(&self) {
81        self.write_started.notified().await;
82    }
83
84    /// Wait until a linked chunk write finishes.
85    ///
86    /// Returns at once if a write has finished since the last call.
87    pub async fn wait_for_write_to_finish(&self) {
88        self.write_finished.notified().await;
89    }
90}
91
92#[cfg_attr(target_family = "wasm", async_trait(?Send))]
93#[cfg_attr(not(target_family = "wasm"), async_trait)]
94impl EventCacheStore for ControlledMemoryStore {
95    type Error = EventCacheStoreError;
96
97    async fn handle_linked_chunk_updates(
98        &self,
99        linked_chunk_id: LinkedChunkId<'_>,
100        updates: Vec<Update<Event, Gap>>,
101    ) -> Result<(), Self::Error> {
102        self.write_started.notify_one();
103
104        // Polling is good enough for tests, and simpler than a proper wake up.
105        while self.writes_blocked.load(Ordering::SeqCst) {
106            sleep(Duration::from_millis(10)).await;
107        }
108
109        let result = self.inner.handle_linked_chunk_updates(linked_chunk_id, updates).await;
110
111        self.write_finished.notify_one();
112
113        result
114    }
115
116    async fn close(&self) -> Result<(), Self::Error> {
117        self.inner.close().await
118    }
119
120    async fn reopen(&self) -> Result<(), Self::Error> {
121        self.inner.reopen().await
122    }
123
124    async fn try_take_leased_lock(
125        &self,
126        lease_duration_ms: u32,
127        key: &str,
128        holder: &str,
129    ) -> Result<Option<CrossProcessLockGeneration>, Self::Error> {
130        self.inner.try_take_leased_lock(lease_duration_ms, key, holder).await
131    }
132
133    async fn load_all_chunks(
134        &self,
135        linked_chunk_id: LinkedChunkId<'_>,
136    ) -> Result<Vec<RawChunk<Event, Gap>>, Self::Error> {
137        self.inner.load_all_chunks(linked_chunk_id).await
138    }
139
140    async fn load_all_chunks_metadata(
141        &self,
142        linked_chunk_id: LinkedChunkId<'_>,
143    ) -> Result<Vec<ChunkMetadata>, Self::Error> {
144        self.inner.load_all_chunks_metadata(linked_chunk_id).await
145    }
146
147    async fn load_last_chunk(
148        &self,
149        linked_chunk_id: LinkedChunkId<'_>,
150    ) -> Result<(Option<RawChunk<Event, Gap>>, ChunkIdentifierGenerator), Self::Error> {
151        self.inner.load_last_chunk(linked_chunk_id).await
152    }
153
154    async fn load_previous_chunk(
155        &self,
156        linked_chunk_id: LinkedChunkId<'_>,
157        before_chunk_identifier: ChunkIdentifier,
158    ) -> Result<Option<RawChunk<Event, Gap>>, Self::Error> {
159        self.inner.load_previous_chunk(linked_chunk_id, before_chunk_identifier).await
160    }
161
162    async fn load_thread_info(
163        &self,
164        room_id: &RoomId,
165        thread_id: &EventId,
166        insert_default_if_missing: bool,
167    ) -> Result<Option<ThreadInfo>, Self::Error> {
168        self.inner.load_thread_info(room_id, thread_id, insert_default_if_missing).await
169    }
170
171    async fn update_thread_info(
172        &self,
173        room_id: &RoomId,
174        thread_id: &EventId,
175        thread_info: &ThreadInfo,
176    ) -> Result<(), Self::Error> {
177        self.inner.update_thread_info(room_id, thread_id, thread_info).await
178    }
179
180    async fn clear_all_events(&self, room_id: Option<&RoomId>) -> Result<(), Self::Error> {
181        self.inner.clear_all_events(room_id).await
182    }
183
184    async fn filter_duplicated_events(
185        &self,
186        linked_chunk_id: LinkedChunkId<'_>,
187        events: Vec<OwnedEventId>,
188    ) -> Result<Vec<(OwnedEventId, Position)>, Self::Error> {
189        self.inner.filter_duplicated_events(linked_chunk_id, events).await
190    }
191
192    async fn find_event(
193        &self,
194        room_id: &RoomId,
195        event_id: &EventId,
196    ) -> Result<Option<Event>, Self::Error> {
197        self.inner.find_event(room_id, event_id).await
198    }
199
200    async fn find_event_relations(
201        &self,
202        room_id: &RoomId,
203        event_id: &EventId,
204        filters: Option<&[RelationType]>,
205    ) -> Result<Vec<(Event, Option<Position>)>, Self::Error> {
206        self.inner.find_event_relations(room_id, event_id, filters).await
207    }
208
209    async fn get_room_events(
210        &self,
211        room_id: &RoomId,
212        event_type: Option<&str>,
213        session_id: Option<&str>,
214    ) -> Result<Vec<Event>, Self::Error> {
215        self.inner.get_room_events(room_id, event_type, session_id).await
216    }
217
218    async fn save_event(&self, room_id: &RoomId, event: Event) -> Result<(), Self::Error> {
219        self.inner.save_event(room_id, event).await
220    }
221
222    async fn optimize(&self) -> Result<(), Self::Error> {
223        self.inner.optimize().await
224    }
225
226    async fn get_size(&self) -> Result<Option<usize>, Self::Error> {
227        self.inner.get_size().await
228    }
229}