matrix_sdk_base/media/store/
memory_store.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
15use std::{
16    collections::HashMap,
17    num::NonZeroUsize,
18    sync::{Arc, RwLock as StdRwLock},
19};
20
21use async_trait::async_trait;
22use matrix_sdk_common::{
23    cross_process_lock::{
24        CrossProcessLockGeneration,
25        memory_store_helper::{Lease, try_take_leased_lock},
26    },
27    ring_buffer::RingBuffer,
28};
29use ruma::{MxcUri, OwnedMxcUri, time::SystemTime};
30
31use super::Result;
32use crate::media::{
33    MediaRequestParameters, UniqueKey as _,
34    store::{
35        IgnoreMediaRetentionPolicy, MediaRetentionPolicy, MediaService, MediaStore,
36        MediaStoreError, MediaStoreInner,
37    },
38};
39
40/// In-memory, non-persistent implementation of the `MediaStore`.
41///
42/// Default if no other is configured at startup.
43#[derive(Debug, Clone)]
44pub struct MemoryMediaStore {
45    inner: Arc<StdRwLock<MemoryMediaStoreInner>>,
46    media_service: MediaService,
47}
48
49#[derive(Debug)]
50struct MemoryMediaStoreInner {
51    media: RingBuffer<MediaContent>,
52    leases: HashMap<String, Lease>,
53    media_retention_policy: Option<MediaRetentionPolicy>,
54    last_media_cleanup_time: SystemTime,
55}
56
57/// A media content in the `MemoryStore`.
58#[derive(Debug)]
59struct MediaContent {
60    /// The URI of the content.
61    uri: OwnedMxcUri,
62
63    /// The unique key of the content.
64    key: String,
65
66    /// The bytes of the content.
67    data: Vec<u8>,
68
69    /// Whether we should ignore the [`MediaRetentionPolicy`] for this content.
70    ignore_policy: bool,
71
72    /// The time of the last access of the content.
73    last_access: SystemTime,
74}
75
76const NUMBER_OF_MEDIAS: NonZeroUsize = NonZeroUsize::new(20).unwrap();
77
78impl Default for MemoryMediaStore {
79    fn default() -> Self {
80        // Given that the store is empty, we won't need to clean it up right away.
81        let last_media_cleanup_time = SystemTime::now();
82        let media_service = MediaService::new();
83        media_service.restore(None, Some(last_media_cleanup_time));
84
85        Self {
86            inner: Arc::new(StdRwLock::new(MemoryMediaStoreInner {
87                media: RingBuffer::new(NUMBER_OF_MEDIAS),
88                leases: Default::default(),
89                media_retention_policy: None,
90                last_media_cleanup_time,
91            })),
92            media_service,
93        }
94    }
95}
96
97impl MemoryMediaStore {
98    /// Create a new empty MemoryMediaStore
99    pub fn new() -> Self {
100        Self::default()
101    }
102}
103
104#[cfg_attr(target_family = "wasm", async_trait(?Send))]
105#[cfg_attr(not(target_family = "wasm"), async_trait)]
106impl MediaStore for MemoryMediaStore {
107    type Error = MediaStoreError;
108
109    async fn try_take_leased_lock(
110        &self,
111        lease_duration_ms: u32,
112        key: &str,
113        holder: &str,
114    ) -> Result<Option<CrossProcessLockGeneration>, Self::Error> {
115        let mut inner = self.inner.write().unwrap();
116
117        Ok(try_take_leased_lock(&mut inner.leases, lease_duration_ms, key, holder))
118    }
119
120    async fn add_media_content(
121        &self,
122        request: &MediaRequestParameters,
123        data: Vec<u8>,
124        ignore_policy: IgnoreMediaRetentionPolicy,
125    ) -> Result<(), Self::Error> {
126        self.media_service.add_media_content(self, request, data, ignore_policy).await
127    }
128
129    async fn replace_media_key(
130        &self,
131        from: &MediaRequestParameters,
132        to: &MediaRequestParameters,
133    ) -> Result<(), Self::Error> {
134        let expected_key = from.unique_key();
135
136        let mut inner = self.inner.write().unwrap();
137
138        if let Some(media_content) =
139            inner.media.iter_mut().find(|media_content| media_content.key == expected_key)
140        {
141            media_content.uri = to.uri().to_owned();
142            media_content.key = to.unique_key();
143        }
144
145        Ok(())
146    }
147
148    async fn get_media_content(
149        &self,
150        request: &MediaRequestParameters,
151    ) -> Result<Option<Vec<u8>>, Self::Error> {
152        self.media_service.get_media_content(self, request).await
153    }
154
155    async fn remove_media_content(
156        &self,
157        request: &MediaRequestParameters,
158    ) -> Result<(), Self::Error> {
159        let expected_key = request.unique_key();
160
161        let mut inner = self.inner.write().unwrap();
162
163        let Some(index) =
164            inner.media.iter().position(|media_content| media_content.key == expected_key)
165        else {
166            return Ok(());
167        };
168
169        inner.media.remove(index);
170
171        Ok(())
172    }
173
174    async fn remove_media_content_for_uri(&self, uri: &MxcUri) -> Result<(), Self::Error> {
175        let mut inner = self.inner.write().unwrap();
176
177        let positions = inner
178            .media
179            .iter()
180            .enumerate()
181            .filter_map(|(position, media_content)| (media_content.uri == uri).then_some(position))
182            .collect::<Vec<_>>();
183
184        // Iterate in reverse-order so that positions stay valid after first removals.
185        for position in positions.into_iter().rev() {
186            inner.media.remove(position);
187        }
188
189        Ok(())
190    }
191
192    async fn set_media_retention_policy(
193        &self,
194        policy: MediaRetentionPolicy,
195    ) -> Result<(), Self::Error> {
196        self.media_service.set_media_retention_policy(self, policy).await
197    }
198
199    fn media_retention_policy(&self) -> MediaRetentionPolicy {
200        self.media_service.media_retention_policy()
201    }
202
203    async fn set_ignore_media_retention_policy(
204        &self,
205        request: &MediaRequestParameters,
206        ignore_policy: IgnoreMediaRetentionPolicy,
207    ) -> Result<(), Self::Error> {
208        self.media_service.set_ignore_media_retention_policy(self, request, ignore_policy).await
209    }
210
211    async fn clean(&self) -> Result<(), Self::Error> {
212        self.media_service.clean(self).await
213    }
214
215    async fn optimize(&self) -> Result<(), Self::Error> {
216        Ok(())
217    }
218
219    async fn get_size(&self) -> Result<Option<usize>, Self::Error> {
220        Ok(None)
221    }
222
223    async fn close(&self) -> Result<(), Self::Error> {
224        Ok(())
225    }
226
227    async fn reopen(&self) -> Result<(), Self::Error> {
228        Ok(())
229    }
230}
231
232#[cfg_attr(target_family = "wasm", async_trait(?Send))]
233#[cfg_attr(not(target_family = "wasm"), async_trait)]
234impl MediaStoreInner for MemoryMediaStore {
235    type Error = MediaStoreError;
236
237    async fn media_retention_policy_inner(
238        &self,
239    ) -> Result<Option<MediaRetentionPolicy>, Self::Error> {
240        Ok(self.inner.read().unwrap().media_retention_policy)
241    }
242
243    async fn set_media_retention_policy_inner(
244        &self,
245        policy: MediaRetentionPolicy,
246    ) -> Result<(), Self::Error> {
247        self.inner.write().unwrap().media_retention_policy = Some(policy);
248        Ok(())
249    }
250
251    async fn add_media_content_inner(
252        &self,
253        request: &MediaRequestParameters,
254        data: Vec<u8>,
255        last_access: SystemTime,
256        policy: MediaRetentionPolicy,
257        ignore_policy: IgnoreMediaRetentionPolicy,
258    ) -> Result<(), Self::Error> {
259        // Avoid duplication. Let's try to remove it first.
260        self.remove_media_content(request).await?;
261
262        let ignore_policy = ignore_policy.is_yes();
263
264        if !ignore_policy && policy.exceeds_max_file_size(data.len() as u64) {
265            // Do not store it.
266            return Ok(());
267        }
268
269        // Now, let's add it.
270        let mut inner = self.inner.write().unwrap();
271        inner.media.push(MediaContent {
272            uri: request.uri().to_owned(),
273            key: request.unique_key(),
274            data,
275            ignore_policy,
276            last_access,
277        });
278
279        Ok(())
280    }
281
282    async fn set_ignore_media_retention_policy_inner(
283        &self,
284        request: &MediaRequestParameters,
285        ignore_policy: IgnoreMediaRetentionPolicy,
286    ) -> Result<(), Self::Error> {
287        let mut inner = self.inner.write().unwrap();
288        let expected_key = request.unique_key();
289
290        if let Some(media_content) = inner.media.iter_mut().find(|media| media.key == expected_key)
291        {
292            media_content.ignore_policy = ignore_policy.is_yes();
293        }
294
295        Ok(())
296    }
297
298    async fn get_media_content_inner(
299        &self,
300        request: &MediaRequestParameters,
301        current_time: SystemTime,
302    ) -> Result<Option<Vec<u8>>, Self::Error> {
303        let mut inner = self.inner.write().unwrap();
304        let expected_key = request.unique_key();
305
306        // First get the content out of the buffer, we are going to put it back at the
307        // end.
308        let Some(index) = inner.media.iter().position(|media| media.key == expected_key) else {
309            return Ok(None);
310        };
311        let Some(mut content) = inner.media.remove(index) else {
312            return Ok(None);
313        };
314
315        // Clone the data.
316        let data = content.data.clone();
317
318        // Update the last access time.
319        content.last_access = current_time;
320
321        // Put it back in the buffer.
322        inner.media.push(content);
323
324        Ok(Some(data))
325    }
326
327    async fn clean_inner(
328        &self,
329        policy: MediaRetentionPolicy,
330        current_time: SystemTime,
331    ) -> Result<(), Self::Error> {
332        if !policy.has_limitations() {
333            // We can safely skip all the checks.
334            return Ok(());
335        }
336
337        let mut inner = self.inner.write().unwrap();
338
339        // First, check media content that exceed the max filesize.
340        if policy.computed_max_file_size().is_some() {
341            inner.media.retain(|content| {
342                content.ignore_policy || !policy.exceeds_max_file_size(content.data.len() as u64)
343            });
344        }
345
346        // Then, clean up expired media content.
347        if policy.last_access_expiry.is_some() {
348            inner.media.retain(|content| {
349                content.ignore_policy
350                    || !policy.has_content_expired(current_time, content.last_access)
351            });
352        }
353
354        // Finally, if the cache size is too big, remove old items until it fits.
355        if let Some(max_cache_size) = policy.max_cache_size {
356            // Reverse the iterator because in case the cache size is overflowing, we want
357            // to count the number of old items to remove. Items are sorted by last access
358            // and old items are at the start.
359            let (_, items_to_remove) = inner.media.iter().enumerate().rev().fold(
360                (0u64, Vec::with_capacity(NUMBER_OF_MEDIAS.into())),
361                |(mut cache_size, mut items_to_remove), (index, content)| {
362                    if content.ignore_policy {
363                        // Do not count it.
364                        return (cache_size, items_to_remove);
365                    }
366
367                    let remove_item = if items_to_remove.is_empty() {
368                        // We have not reached the max cache size yet.
369                        if let Some(sum) = cache_size.checked_add(content.data.len() as u64) {
370                            cache_size = sum;
371                            // Start removing items if we have exceeded the max cache size.
372                            cache_size > max_cache_size
373                        } else {
374                            // The cache size is overflowing, remove the remaining items, since the
375                            // max cache size cannot be bigger than
376                            // usize::MAX.
377                            true
378                        }
379                    } else {
380                        // We have reached the max cache size already, just remove it.
381                        true
382                    };
383
384                    if remove_item {
385                        items_to_remove.push(index);
386                    }
387
388                    (cache_size, items_to_remove)
389                },
390            );
391
392            // The indexes are already in reverse order so we can just iterate in that order
393            // to remove them starting by the end.
394            for index in items_to_remove {
395                inner.media.remove(index);
396            }
397        }
398
399        inner.last_media_cleanup_time = current_time;
400
401        Ok(())
402    }
403
404    async fn last_media_cleanup_time_inner(&self) -> Result<Option<SystemTime>, Self::Error> {
405        Ok(Some(self.inner.read().unwrap().last_media_cleanup_time))
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::{MemoryMediaStore, Result};
412    use crate::{
413        media_store_inner_integration_tests, media_store_integration_tests,
414        media_store_integration_tests_time,
415    };
416
417    async fn get_media_store() -> Result<MemoryMediaStore> {
418        Ok(MemoryMediaStore::new())
419    }
420
421    media_store_inner_integration_tests!();
422    media_store_integration_tests!();
423    media_store_integration_tests_time!();
424}
matrix_sdk_base/media/store/memory_store.rs

matrix_sdk_base/media/store/
memory_store.rs
