Merge branch 'custom-store'

This commit is contained in:
Damir Jelić 2021-09-09 12:24:27 +02:00
commit ed198a22b6
3 changed files with 73 additions and 1 deletions

View file

@ -66,6 +66,7 @@ pub struct MemoryStore {
room_event_receipts:
Arc<DashMap<RoomId, DashMap<String, DashMap<EventId, DashMap<UserId, Receipt>>>>>,
media: Arc<Mutex<LruCache<String, Vec<u8>>>>,
custom: Arc<DashMap<Vec<u8>, Vec<u8>>>,
}
impl MemoryStore {
@ -90,6 +91,7 @@ impl MemoryStore {
room_user_receipts: Default::default(),
room_event_receipts: Default::default(),
media: Arc::new(Mutex::new(LruCache::new(100))),
custom: DashMap::new().into(),
}
}
@ -407,6 +409,14 @@ impl MemoryStore {
.unwrap_or_else(Vec::new))
}
async fn get_custom_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
Ok(self.custom.get(key).map(|e| e.value().clone()))
}
async fn set_custom_value(&self, key: &[u8], value: Vec<u8>) -> Result<Option<Vec<u8>>> {
Ok(self.custom.insert(key.to_vec(), value))
}
async fn add_media_content(&self, request: &MediaRequest, data: Vec<u8>) -> Result<()> {
self.media.lock().await.put(request.unique_key(), data);
@ -563,6 +573,14 @@ impl StateStore for MemoryStore {
self.get_event_room_receipt_events(room_id, receipt_type, event_id).await
}
async fn get_custom_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
self.get_custom_value(key).await
}
async fn set_custom_value(&self, key: &[u8], value: Vec<u8>) -> Result<Option<Vec<u8>>> {
self.set_custom_value(key, value).await
}
async fn add_media_content(&self, request: &MediaRequest, data: Vec<u8>) -> Result<()> {
self.add_media_content(request, data).await
}

View file

@ -263,6 +263,22 @@ pub trait StateStore: AsyncTraitDeps {
event_id: &EventId,
) -> Result<Vec<(UserId, Receipt)>>;
/// Get arbitrary data from the custom store
///
/// # Arguments
///
/// * `key` - The key to fetch data for
async fn get_custom_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>>;
/// Put arbitrary data into the custom store
///
/// # Arguments
///
/// * `key` - The key to insert data into
///
/// * `value` - The value to insert
async fn set_custom_value(&self, key: &[u8], value: Vec<u8>) -> Result<Option<Vec<u8>>>;
/// Add a media file's content in the media store.
///
/// # Arguments

View file

@ -189,6 +189,7 @@ pub struct SledStore {
room_user_receipts: Tree,
room_event_receipts: Tree,
media: Tree,
custom: Tree,
}
impl std::fmt::Debug for SledStore {
@ -226,6 +227,8 @@ impl SledStore {
let media = db.open_tree("media")?;
let custom = db.open_tree("custom")?;
Ok(Self {
path,
inner: db,
@ -247,6 +250,7 @@ impl SledStore {
room_user_receipts,
room_event_receipts,
media,
custom,
})
}
@ -762,6 +766,17 @@ impl SledStore {
.map(|m| m.to_vec()))
}
async fn get_custom_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
Ok(self.custom.get(key)?.map(|v| v.to_vec()))
}
async fn set_custom_value(&self, key: &[u8], value: Vec<u8>) -> Result<Option<Vec<u8>>> {
let ret = self.custom.insert(key, value)?.map(|v| v.to_vec());
self.inner.flush_async().await?;
Ok(ret)
}
async fn remove_media_content(&self, request: &MediaRequest) -> Result<()> {
self.media.remove(
(request.media_type.unique_key().as_str(), request.format.unique_key().as_str())
@ -899,6 +914,14 @@ impl StateStore for SledStore {
self.get_event_room_receipt_events(room_id, receipt_type, event_id).await
}
async fn get_custom_value(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
self.get_custom_value(key).await
}
async fn set_custom_value(&self, key: &[u8], value: Vec<u8>) -> Result<Option<Vec<u8>>> {
self.set_custom_value(key, value).await
}
async fn add_media_content(&self, request: &MediaRequest, data: Vec<u8>) -> Result<()> {
self.add_media_content(request, data).await
}
@ -939,7 +962,7 @@ mod test {
};
use serde_json::json;
use super::{SledStore, StateChanges};
use super::{Result, SledStore, StateChanges};
use crate::{
deserialized_responses::MemberEvent,
media::{MediaFormat, MediaRequest, MediaThumbnailSize, MediaType},
@ -1155,4 +1178,19 @@ mod test {
assert!(store.get_media_content(&request_file).await.unwrap().is_none());
assert!(store.get_media_content(&request_thumbnail).await.unwrap().is_none());
}
#[async_test]
async fn test_custom_storage() -> Result<()> {
let key = "my_key";
let value = &[0, 1, 2, 3];
let store = SledStore::open()?;
store.set_custom_value(key.as_bytes(), value.to_vec()).await?;
let read = store.get_custom_value(key.as_bytes()).await?;
assert_eq!(Some(value.as_ref()), read.as_deref());
Ok(())
}
}