From 8ca5983093e5f43e06e2ef7016f229d83ab7cf38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?K=C3=A9vin=20Commaille?= Date: Tue, 21 Jan 2025 18:03:52 +0100 Subject: [PATCH] feat(sqlite): Implement EventCacheStoreMedia for SqliteEventCacheStore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Kévin Commaille --- .../event_cache_store/004_ignore_policy.sql | 3 + .../src/event_cache_store.rs | 367 +++++++++++++++--- crates/matrix-sdk-sqlite/src/utils.rs | 29 +- 3 files changed, 344 insertions(+), 55 deletions(-) create mode 100644 crates/matrix-sdk-sqlite/migrations/event_cache_store/004_ignore_policy.sql diff --git a/crates/matrix-sdk-sqlite/migrations/event_cache_store/004_ignore_policy.sql b/crates/matrix-sdk-sqlite/migrations/event_cache_store/004_ignore_policy.sql new file mode 100644 index 000000000..b556f6a90 --- /dev/null +++ b/crates/matrix-sdk-sqlite/migrations/event_cache_store/004_ignore_policy.sql @@ -0,0 +1,3 @@ +-- Add an ignore_policy column, defaulting to FALSE for all media content. +ALTER TABLE "media" + ADD COLUMN "ignore_policy" BOOLEAN NOT NULL DEFAULT FALSE; diff --git a/crates/matrix-sdk-sqlite/src/event_cache_store.rs b/crates/matrix-sdk-sqlite/src/event_cache_store.rs index 2afef54ef..ab2b2991f 100644 --- a/crates/matrix-sdk-sqlite/src/event_cache_store.rs +++ b/crates/matrix-sdk-sqlite/src/event_cache_store.rs @@ -19,23 +19,40 @@ use std::{borrow::Cow, fmt, path::Path, sync::Arc}; use async_trait::async_trait; use deadpool_sqlite::{Object as SqliteAsyncConn, Pool as SqlitePool, Runtime}; use matrix_sdk_base::{ - event_cache::{store::EventCacheStore, Event, Gap}, + event_cache::{ + store::{ + media::{ + EventCacheStoreMedia, IgnoreMediaRetentionPolicy, MediaRetentionPolicy, + MediaService, + }, + EventCacheStore, + }, + Event, Gap, + }, linked_chunk::{ChunkContent, ChunkIdentifier, RawChunk, Update}, media::{MediaRequestParameters, UniqueKey}, }; use matrix_sdk_store_encryption::StoreCipher; -use ruma::{MilliSecondsSinceUnixEpoch, RoomId}; -use rusqlite::{OptionalExtension, Transaction, TransactionBehavior}; +use ruma::{time::SystemTime, MilliSecondsSinceUnixEpoch, MxcUri, RoomId}; +use rusqlite::{params_from_iter, OptionalExtension, Transaction, TransactionBehavior}; use tokio::fs; +#[cfg(not(test))] +use tracing::warn; use tracing::{debug, trace}; use crate::{ error::{Error, Result}, - utils::{Key, SqliteAsyncConnExt, SqliteKeyValueStoreAsyncConnExt, SqliteKeyValueStoreConnExt}, + utils::{ + repeat_vars, time_to_timestamp, Key, SqliteAsyncConnExt, SqliteKeyValueStoreAsyncConnExt, + SqliteKeyValueStoreConnExt, SqliteTransactionExt, + }, OpenStoreError, }; mod keys { + // Entries in Key-value store + pub const MEDIA_RETENTION_POLICY: &str = "media_retention_policy"; + // Tables pub const LINKED_CHUNKS: &str = "linked_chunks"; pub const MEDIA: &str = "media"; @@ -46,7 +63,7 @@ mod keys { /// This is used to figure whether the SQLite database requires a migration. /// Every new SQL migration should imply a bump of this number, and changes in /// the [`run_migrations`] function. -const DATABASE_VERSION: u8 = 3; +const DATABASE_VERSION: u8 = 4; /// The string used to identify a chunk of type events, in the `type` field in /// the database. @@ -60,6 +77,7 @@ const CHUNK_TYPE_GAP_TYPE_STRING: &str = "G"; pub struct SqliteEventCacheStore { store_cipher: Option>, pool: SqlitePool, + media_service: Arc, } #[cfg(not(tarpaulin_include))] @@ -96,7 +114,11 @@ impl SqliteEventCacheStore { None => None, }; - Ok(Self { store_cipher, pool }) + let media_service = MediaService::new(); + let media_retention_policy = media_retention_policy(&conn).await?; + media_service.restore(media_retention_policy); + + Ok(Self { store_cipher, pool, media_service: Arc::new(media_service) }) } fn encode_value(&self, value: Vec) -> Result> { @@ -303,6 +325,16 @@ async fn run_migrations(conn: &SqliteAsyncConn, version: u8) -> Result<()> { .await?; } + if version < 4 { + conn.with_transaction(|txn| { + txn.execute_batch(include_str!( + "../migrations/event_cache_store/004_ignore_policy.sql" + ))?; + txn.set_db_version(4) + }) + .await?; + } + Ok(()) } @@ -592,18 +624,9 @@ impl EventCacheStore for SqliteEventCacheStore { request: &MediaRequestParameters, content: Vec, ) -> Result<()> { - let uri = self.encode_key(keys::MEDIA, request.source.unique_key()); - let format = self.encode_key(keys::MEDIA, request.format.unique_key()); - let data = self.encode_value(content)?; - - let conn = self.acquire().await?; - conn.execute( - "INSERT OR REPLACE INTO media (uri, format, data, last_access) VALUES (?, ?, ?, CAST(strftime('%s') as INT))", - (uri, format, data), - ) - .await?; - - Ok(()) + self.media_service + .add_media_content(self, request, content, IgnoreMediaRetentionPolicy::No) + .await } async fn replace_media_key( @@ -619,8 +642,7 @@ impl EventCacheStore for SqliteEventCacheStore { let conn = self.acquire().await?; conn.execute( - r#"UPDATE media SET uri = ?, format = ?, last_access = CAST(strftime('%s') as INT) - WHERE uri = ? AND format = ?"#, + r#"UPDATE media SET uri = ?, format = ? WHERE uri = ? AND format = ?"#, (new_uri, new_format, prev_uri, prev_format), ) .await?; @@ -629,31 +651,7 @@ impl EventCacheStore for SqliteEventCacheStore { } async fn get_media_content(&self, request: &MediaRequestParameters) -> Result>> { - let uri = self.encode_key(keys::MEDIA, request.source.unique_key()); - let format = self.encode_key(keys::MEDIA, request.format.unique_key()); - - let conn = self.acquire().await?; - let data = conn - .with_transaction::<_, rusqlite::Error, _>(move |txn| { - // Update the last access. - // We need to do this first so the transaction is in write mode right away. - // See: https://sqlite.org/lang_transaction.html#read_transactions_versus_write_transactions - txn.execute( - "UPDATE media SET last_access = CAST(strftime('%s') as INT) \ - WHERE uri = ? AND format = ?", - (&uri, &format), - )?; - - txn.query_row::, _, _>( - "SELECT data FROM media WHERE uri = ? AND format = ?", - (&uri, &format), - |row| row.get(0), - ) - .optional() - }) - .await?; - - data.map(|v| self.decode_value(&v).map(Into::into)).transpose() + self.media_service.get_media_content(self, request).await } async fn remove_media_content(&self, request: &MediaRequestParameters) -> Result<()> { @@ -668,9 +666,102 @@ impl EventCacheStore for SqliteEventCacheStore { async fn get_media_content_for_uri( &self, - uri: &ruma::MxcUri, + uri: &MxcUri, ) -> Result>, Self::Error> { + self.media_service.get_media_content_for_uri(self, uri).await + } + + async fn remove_media_content_for_uri(&self, uri: &MxcUri) -> Result<()> { let uri = self.encode_key(keys::MEDIA, uri); + + let conn = self.acquire().await?; + conn.execute("DELETE FROM media WHERE uri = ?", (uri,)).await?; + + Ok(()) + } +} + +#[cfg_attr(target_arch = "wasm32", async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait)] +impl EventCacheStoreMedia for SqliteEventCacheStore { + type Error = Error; + + async fn media_retention_policy_inner( + &self, + ) -> Result, Self::Error> { + let conn = self.acquire().await?; + media_retention_policy(&conn).await + } + + async fn set_media_retention_policy_inner( + &self, + policy: MediaRetentionPolicy, + ) -> Result<(), Self::Error> { + let conn = self.acquire().await?; + + let serialized_policy = rmp_serde::to_vec_named(&policy)?; + conn.set_kv(keys::MEDIA_RETENTION_POLICY, serialized_policy).await?; + + Ok(()) + } + + async fn add_media_content_inner( + &self, + request: &MediaRequestParameters, + data: Vec, + last_access: SystemTime, + policy: MediaRetentionPolicy, + ignore_policy: IgnoreMediaRetentionPolicy, + ) -> Result<(), Self::Error> { + let ignore_policy = ignore_policy.is_yes(); + let data = self.encode_value(data)?; + + if !ignore_policy && policy.exceeds_max_file_size(data.len()) { + return Ok(()); + } + + let uri = self.encode_key(keys::MEDIA, request.source.unique_key()); + let format = self.encode_key(keys::MEDIA, request.format.unique_key()); + let timestamp = time_to_timestamp(last_access); + + let conn = self.acquire().await?; + conn.execute( + "INSERT OR REPLACE INTO media (uri, format, data, last_access, ignore_policy) VALUES (?, ?, ?, ?, ?)", + (uri, format, data, timestamp, ignore_policy), + ) + .await?; + + Ok(()) + } + + async fn set_ignore_media_retention_policy_inner( + &self, + request: &MediaRequestParameters, + ignore_policy: IgnoreMediaRetentionPolicy, + ) -> Result<(), Self::Error> { + let uri = self.encode_key(keys::MEDIA, request.source.unique_key()); + let format = self.encode_key(keys::MEDIA, request.format.unique_key()); + let ignore_policy = ignore_policy.is_yes(); + + let conn = self.acquire().await?; + conn.execute( + r#"UPDATE media SET ignore_policy = ? WHERE uri = ? AND format = ?"#, + (ignore_policy, uri, format), + ) + .await?; + + Ok(()) + } + + async fn get_media_content_inner( + &self, + request: &MediaRequestParameters, + current_time: SystemTime, + ) -> Result>, Self::Error> { + let uri = self.encode_key(keys::MEDIA, request.source.unique_key()); + let format = self.encode_key(keys::MEDIA, request.format.unique_key()); + let timestamp = time_to_timestamp(current_time); + let conn = self.acquire().await?; let data = conn .with_transaction::<_, rusqlite::Error, _>(move |txn| { @@ -678,11 +769,38 @@ impl EventCacheStore for SqliteEventCacheStore { // We need to do this first so the transaction is in write mode right away. // See: https://sqlite.org/lang_transaction.html#read_transactions_versus_write_transactions txn.execute( - "UPDATE media SET last_access = CAST(strftime('%s') as INT) \ - WHERE uri = ?", - (&uri,), + "UPDATE media SET last_access = ? WHERE uri = ? AND format = ?", + (timestamp, &uri, &format), )?; + txn.query_row::, _, _>( + "SELECT data FROM media WHERE uri = ? AND format = ?", + (&uri, &format), + |row| row.get(0), + ) + .optional() + }) + .await?; + + data.map(|v| self.decode_value(&v).map(Into::into)).transpose() + } + + async fn get_media_content_for_uri_inner( + &self, + uri: &MxcUri, + current_time: SystemTime, + ) -> Result>, Self::Error> { + let uri = self.encode_key(keys::MEDIA, uri); + let timestamp = time_to_timestamp(current_time); + + let conn = self.acquire().await?; + let data = conn + .with_transaction::<_, rusqlite::Error, _>(move |txn| { + // Update the last access. + // We need to do this first so the transaction is in write mode right away. + // See: https://sqlite.org/lang_transaction.html#read_transactions_versus_write_transactions + txn.execute("UPDATE media SET last_access = ? WHERE uri = ?", (timestamp, &uri))?; + txn.query_row::, _, _>( "SELECT data FROM media WHERE uri = ?", (&uri,), @@ -695,11 +813,140 @@ impl EventCacheStore for SqliteEventCacheStore { data.map(|v| self.decode_value(&v).map(Into::into)).transpose() } - async fn remove_media_content_for_uri(&self, uri: &ruma::MxcUri) -> Result<()> { - let uri = self.encode_key(keys::MEDIA, uri); + async fn clean_up_media_cache_inner( + &self, + policy: MediaRetentionPolicy, + current_time: SystemTime, + ) -> Result<(), Self::Error> { + if !policy.has_limitations() { + // We can safely skip all the checks. + return Ok(()); + } let conn = self.acquire().await?; - conn.execute("DELETE FROM media WHERE uri = ?", (uri,)).await?; + let removed = conn + .with_transaction::<_, Error, _>(move |txn| { + let mut removed = false; + + // First, check media content that exceed the max filesize. + if let Some(max_file_size) = policy.computed_max_file_size() { + let count = txn.execute( + "DELETE FROM media WHERE ignore_policy IS FALSE AND length(data) > ?", + (max_file_size,), + )?; + + if count > 0 { + removed = true; + } + } + + // Then, clean up expired media content. + if let Some(last_access_expiry) = policy.last_access_expiry { + let current_timestamp = time_to_timestamp(current_time); + let expiry_secs = last_access_expiry.as_secs(); + let count = txn.execute( + "DELETE FROM media WHERE ignore_policy IS FALSE AND (? - last_access) >= ?", + (current_timestamp, expiry_secs), + )?; + + if count > 0 { + removed = true; + } + } + + // Finally, if the cache size is too big, remove old items until it fits. + if let Some(max_cache_size) = policy.max_cache_size { + // i64 is the integer type used by SQLite, use it here to avoid usize overflow + // during the conversion of the result. + let cache_size_int = txn + .query_row( + "SELECT sum(length(data)) FROM media WHERE ignore_policy IS FALSE", + (), + |row| { + // `sum()` returns `NULL` if there are no rows. + row.get::<_, Option>(0) + }, + )? + .unwrap_or_default(); + let cache_size_usize = usize::try_from(cache_size_int); + + // If the cache size is overflowing or bigger than max cache size, clean up. + if cache_size_usize.is_err() + || cache_size_usize.is_ok_and(|cache_size| cache_size > max_cache_size) + { + // Get the sizes of the media contents ordered by last access. + let mut cached_stmt = txn.prepare_cached( + "SELECT rowid, length(data) FROM media \ + WHERE ignore_policy IS FALSE ORDER BY last_access DESC", + )?; + let content_sizes = cached_stmt + .query(())? + .mapped(|row| Ok((row.get::<_, i64>(0)?, row.get::<_, usize>(1)?))); + + let mut accumulated_items_size = 0usize; + let mut limit_reached = false; + let mut rows_to_remove = Vec::new(); + + for result in content_sizes { + let (row_id, size) = match result { + Ok(content_size) => content_size, + Err(error) => { + return Err(error.into()); + } + }; + + if limit_reached { + rows_to_remove.push(row_id); + continue; + } + + match accumulated_items_size.checked_add(size) { + Some(acc) if acc > max_cache_size => { + // We can stop accumulating. + limit_reached = true; + rows_to_remove.push(row_id); + } + Some(acc) => accumulated_items_size = acc, + None => { + // The accumulated size is overflowing but the setting cannot be + // bigger than usize::MAX, we can stop accumulating. + limit_reached = true; + rows_to_remove.push(row_id); + } + }; + } + + if !rows_to_remove.is_empty() { + removed = true; + } + + txn.chunk_large_query_over(rows_to_remove, None, |txn, row_ids| { + let sql_params = repeat_vars(row_ids.len()); + let query = format!("DELETE FROM media WHERE rowid IN ({sql_params})"); + txn.prepare(&query)?.execute(params_from_iter(row_ids))?; + Ok(Vec::<()>::new()) + })?; + } + } + + Ok(removed) + }) + .await?; + + // If we removed media, use the VACUUM command to defragment the + // database and free space on the filesystem. + if removed { + if let Err(error) = conn.execute("VACUUM", ()).await { + // Since this is an optimisation step, do not propagate the error + // but log it. + #[cfg(not(test))] + warn!("Failed to vacuum database: {error}"); + + // We want to know if there is an error with this step during tests. + #[cfg(test)] + return Err(error.into()); + } + } Ok(()) } @@ -786,6 +1033,17 @@ fn insert_chunk( Ok(()) } +/// Get the persisted [`MediaRetentionPolicy`] with the given connection. +async fn media_retention_policy( + conn: &SqliteAsyncConn, +) -> Result, Error> { + let Some(bytes) = conn.get_kv(keys::MEDIA_RETENTION_POLICY).await? else { + return Ok(None); + }; + + Ok(Some(rmp_serde::from_slice(&bytes)?)) +} + #[cfg(test)] mod tests { use std::{ @@ -803,6 +1061,7 @@ mod tests { Gap, }, event_cache_store_integration_tests, event_cache_store_integration_tests_time, + event_cache_store_media_integration_tests, linked_chunk::{ChunkContent, ChunkIdentifier, Position, Update}, media::{MediaFormat, MediaRequestParameters, MediaThumbnailSettings}, }; @@ -828,6 +1087,7 @@ mod tests { event_cache_store_integration_tests!(); event_cache_store_integration_tests_time!(); + event_cache_store_media_integration_tests!(with_media_size_tests); async fn get_event_cache_store_content_sorted_by_last_access( event_cache_store: &SqliteEventCacheStore, @@ -1476,7 +1736,7 @@ mod encrypted_tests { use matrix_sdk_base::{ event_cache::store::EventCacheStoreError, event_cache_store_integration_tests, - event_cache_store_integration_tests_time, + event_cache_store_integration_tests_time, event_cache_store_media_integration_tests, }; use once_cell::sync::Lazy; use tempfile::{tempdir, TempDir}; @@ -1502,4 +1762,5 @@ mod encrypted_tests { event_cache_store_integration_tests!(); event_cache_store_integration_tests_time!(); + event_cache_store_media_integration_tests!(); } diff --git a/crates/matrix-sdk-sqlite/src/utils.rs b/crates/matrix-sdk-sqlite/src/utils.rs index 7daa9358b..28ba52167 100644 --- a/crates/matrix-sdk-sqlite/src/utils.rs +++ b/crates/matrix-sdk-sqlite/src/utils.rs @@ -19,6 +19,7 @@ use async_trait::async_trait; use deadpool_sqlite::Object as SqliteAsyncConn; use itertools::Itertools; use matrix_sdk_store_encryption::StoreCipher; +use ruma::time::SystemTime; use rusqlite::{limits::Limit, OptionalExtension, Params, Row, Statement, Transaction}; use crate::{ @@ -187,7 +188,7 @@ impl SqliteAsyncConnExt for SqliteAsyncConn { } pub(crate) trait SqliteTransactionExt { - fn chunk_large_query_over( + fn chunk_large_query_over( &self, keys_to_chunk: Vec, result_capacity: Option, @@ -199,7 +200,7 @@ pub(crate) trait SqliteTransactionExt { } impl SqliteTransactionExt for Transaction<'_> { - fn chunk_large_query_over( + fn chunk_large_query_over( &self, mut keys_to_chunk: Vec, result_capacity: Option, @@ -380,8 +381,23 @@ pub(crate) fn repeat_vars(count: usize) -> impl fmt::Display { iter::repeat("?").take(count).format(",") } +/// Convert the given `SystemTime` to a timestamp, as the number of seconds +/// since Unix Epoch. +/// +/// Returns an `i64` as it is the numeric type used by SQLite. +pub(crate) fn time_to_timestamp(time: SystemTime) -> i64 { + time.duration_since(SystemTime::UNIX_EPOCH) + .ok() + .and_then(|d| d.as_secs().try_into().ok()) + // It is unlikely to happen unless the time on the system is seriously wrong, but we always + // need a value. + .unwrap_or(0) +} + #[cfg(test)] mod unit_tests { + use std::time::Duration; + use super::*; #[test] @@ -396,4 +412,13 @@ mod unit_tests { fn generating_zero_vars_panics() { repeat_vars(0); } + + #[test] + fn test_time_to_timestamp() { + assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH), 0); + assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH + Duration::from_secs(60)), 60); + + // Fallback value on overflow. + assert_eq!(time_to_timestamp(SystemTime::UNIX_EPOCH - Duration::from_secs(60)), 0); + } }