Dima/matrix-rust-sdk
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sqlite): Implement EventCacheStoreMedia for SqliteEventCacheStore

Browse Source 
Signed-off-by: Kévin Commaille <zecakeh@tedomum.fr>
This commit is contained in:
Kévin Commaille
2025-01-21 18:03:52 +01:00
parent 144f568a5c
commit 8ca5983093
3 changed files with 344 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/matrix-sdk-sqlite/migrations/event_cache_store/004_ignore_policy.sql
+3
View File
@@ -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;
crates/matrix-sdk-sqlite/src/event_cache_store.rs
+314 -53
View File
@@ -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<Arc<StoreCipher>>,
pool: SqlitePool,
media_service: Arc<MediaService>,
}
#[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<u8>) -> Result<Vec<u8>> {
@@ -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<u8>,
) -> 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<Option<Vec<u8>>> {
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::<Vec<u8>, _, _>(
"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<Option<Vec<u8>>, 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<Option<MediaRetentionPolicy>, 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<u8>,
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<Option<Vec<u8>>, 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::<Vec<u8>, _, _>(
"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<Option<Vec<u8>>, 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::<Vec<u8>, _, _>(
"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<i64>>(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<Option<MediaRetentionPolicy>, 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!();
}
crates/matrix-sdk-sqlite/src/utils.rs
+27 -2
View File
@@ -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<Query, Res>(
fn chunk_large_query_over<Key, Query, Res>(
&self,
keys_to_chunk: Vec<Key>,
result_capacity: Option<usize>,
@@ -199,7 +200,7 @@ pub(crate) trait SqliteTransactionExt {
}
impl SqliteTransactionExt for Transaction<'_> {
fn chunk_large_query_over<Query, Res>(
fn chunk_large_query_over<Key, Query, Res>(
&self,
mut keys_to_chunk: Vec<Key>,
result_capacity: Option<usize>,
@@ -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);
}
}