messing around with arcs

20 files changed, 100 insertions, 47 deletions

diff --git a/src/database/key_value/rooms/alias.rs b/src/database/key_value/rooms/alias.rs
index 0aa8dd4..437902d 100644
--- a/src/database/key_value/rooms/alias.rs
+++ b/src/database/key_value/rooms/alias.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::{RoomId, RoomAliasId, api::client::error::ErrorKind};
 
 use crate::{service, database::KeyValueDatabase, utils, Error, services, Result};
 
-impl service::rooms::alias::Data for KeyValueDatabase {
+impl service::rooms::alias::Data for Arc<KeyValueDatabase> {
     fn set_alias(
         &self,
         alias: &RoomAliasId,
diff --git a/src/database/key_value/rooms/auth_chain.rs b/src/database/key_value/rooms/auth_chain.rs
index 888d472..2dffb04 100644
--- a/src/database/key_value/rooms/auth_chain.rs
+++ b/src/database/key_value/rooms/auth_chain.rs
@@ -1,28 +1,60 @@
-use std::{collections::HashSet, mem::size_of};
+use std::{collections::HashSet, mem::size_of, sync::Arc};
 
 use crate::{service, database::KeyValueDatabase, Result, utils};
 
-impl service::rooms::auth_chain::Data for KeyValueDatabase {
-    fn get_cached_eventid_authchain(&self, shorteventid: u64) -> Result<Option<HashSet<u64>>> {
-        Ok(self.shorteventid_authchain
-            .get(&shorteventid.to_be_bytes())?
-            .map(|chain| {
-                chain
-                    .chunks_exact(size_of::<u64>())
-                    .map(|chunk| {
-                        utils::u64_from_bytes(chunk).expect("byte length is correct")
-                    })
-                    .collect()
-            }))
+impl service::rooms::auth_chain::Data for Arc<KeyValueDatabase> {
+    fn get_cached_eventid_authchain(&self, key: &[u64]) -> Result<Option<Arc<HashSet<u64>>>> {
+        // Check RAM cache
+        if let Some(result) = self.auth_chain_cache.lock().unwrap().get_mut(key) {
+            return Ok(Some(Arc::clone(result)));
+        }
+
+        // We only save auth chains for single events in the db
+        if key.len() == 1 {
+            // Check DB cache
+            let chain = self.shorteventid_authchain
+                .get(&key[0].to_be_bytes())?
+                .map(|chain| {
+                    chain
+                        .chunks_exact(size_of::<u64>())
+                        .map(|chunk| {
+                            utils::u64_from_bytes(chunk).expect("byte length is correct")
+                        })
+                        .collect()
+                });
+
+            if let Some(chain) = chain {
+                let chain = Arc::new(chain);
+
+                // Cache in RAM
+                self.auth_chain_cache
+                    .lock()
+                    .unwrap()
+                    .insert(vec![key[0]], Arc::clone(&chain));
+
+                return Ok(Some(chain));
+            }
+        }
+
+        Ok(None)
+
     }
 
-    fn cache_eventid_authchain(&self, shorteventid: u64, auth_chain: &HashSet<u64>) -> Result<()> {
-        self.shorteventid_authchain.insert(
-            &shorteventid.to_be_bytes(),
-            &auth_chain
-                .iter()
-                .flat_map(|s| s.to_be_bytes().to_vec())
-                .collect::<Vec<u8>>(),
-        )
+    fn cache_auth_chain(&self, key: Vec<u64>, auth_chain: Arc<HashSet<u64>>) -> Result<()> {
+        // Only persist single events in db
+        if key.len() == 1 {
+            self.shorteventid_authchain.insert(
+                &key[0].to_be_bytes(),
+                &auth_chain
+                    .iter()
+                    .flat_map(|s| s.to_be_bytes().to_vec())
+                    .collect::<Vec<u8>>(),
+            )?;
+        }
+
+        // Cache in RAM
+        self.auth_chain_cache.lock().unwrap().insert(key, auth_chain);
+
+        Ok(())
     }
 }
diff --git a/src/database/key_value/rooms/directory.rs b/src/database/key_value/rooms/directory.rs
index 727004e..864e75e 100644
--- a/src/database/key_value/rooms/directory.rs
+++ b/src/database/key_value/rooms/directory.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::RoomId;
 
 use crate::{service, database::KeyValueDatabase, utils, Error, Result};
 
-impl service::rooms::directory::Data for KeyValueDatabase {
+impl service::rooms::directory::Data for Arc<KeyValueDatabase> {
     fn set_public(&self, room_id: &RoomId) -> Result<()> {
         self.publicroomids.insert(room_id.as_bytes(), &[])
     }
diff --git a/src/database/key_value/rooms/edus/mod.rs b/src/database/key_value/rooms/edus/mod.rs
index b5007f8..03e4219 100644
--- a/src/database/key_value/rooms/edus/mod.rs
+++ b/src/database/key_value/rooms/edus/mod.rs
@@ -2,6 +2,8 @@ mod presence;
 mod typing;
 mod read_receipt;
 
+use std::sync::Arc;
+
 use crate::{service, database::KeyValueDatabase};
 
-impl service::rooms::edus::Data for KeyValueDatabase {}
+impl service::rooms::edus::Data for Arc<KeyValueDatabase> {}
diff --git a/src/database/key_value/rooms/edus/presence.rs b/src/database/key_value/rooms/edus/presence.rs
index 1477c28..5aeb147 100644
--- a/src/database/key_value/rooms/edus/presence.rs
+++ b/src/database/key_value/rooms/edus/presence.rs
@@ -1,10 +1,10 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::Arc};
 
 use ruma::{UserId, RoomId, events::presence::PresenceEvent, presence::PresenceState, UInt};
 
 use crate::{service, database::KeyValueDatabase, utils, Error, services, Result};
 
-impl service::rooms::edus::presence::Data for KeyValueDatabase {
+impl service::rooms::edus::presence::Data for Arc<KeyValueDatabase> {
     fn update_presence(
         &self,
         user_id: &UserId,
diff --git a/src/database/key_value/rooms/edus/read_receipt.rs b/src/database/key_value/rooms/edus/read_receipt.rs
index a12e265..7fcb8ac 100644
--- a/src/database/key_value/rooms/edus/read_receipt.rs
+++ b/src/database/key_value/rooms/edus/read_receipt.rs
@@ -1,10 +1,10 @@
-use std::mem;
+use std::{mem, sync::Arc};
 
 use ruma::{UserId, RoomId, events::receipt::ReceiptEvent, serde::Raw, signatures::CanonicalJsonObject};
 
 use crate::{database::KeyValueDatabase, service, utils, Error, services, Result};
 
-impl service::rooms::edus::read_receipt::Data for KeyValueDatabase {
+impl service::rooms::edus::read_receipt::Data for Arc<KeyValueDatabase> {
     fn readreceipt_update(
         &self,
         user_id: &UserId,
diff --git a/src/database/key_value/rooms/edus/typing.rs b/src/database/key_value/rooms/edus/typing.rs
index b7d3596..7f3526d 100644
--- a/src/database/key_value/rooms/edus/typing.rs
+++ b/src/database/key_value/rooms/edus/typing.rs
@@ -1,10 +1,10 @@
-use std::collections::HashSet;
+use std::{collections::HashSet, sync::Arc};
 
 use ruma::{UserId, RoomId};
 
 use crate::{database::KeyValueDatabase, service, utils, Error, services, Result};
 
-impl service::rooms::edus::typing::Data for KeyValueDatabase {
+impl service::rooms::edus::typing::Data for Arc<KeyValueDatabase> {
     fn typing_add(
         &self,
         user_id: &UserId,
diff --git a/src/database/key_value/rooms/lazy_load.rs b/src/database/key_value/rooms/lazy_load.rs
index 133e1d0..b16657a 100644
--- a/src/database/key_value/rooms/lazy_load.rs
+++ b/src/database/key_value/rooms/lazy_load.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::{UserId, DeviceId, RoomId};
 
 use crate::{service, database::KeyValueDatabase, Result};
 
-impl service::rooms::lazy_loading::Data for KeyValueDatabase {
+impl service::rooms::lazy_loading::Data for Arc<KeyValueDatabase> {
     fn lazy_load_was_sent_before(
         &self,
         user_id: &UserId,
diff --git a/src/database/key_value/rooms/metadata.rs b/src/database/key_value/rooms/metadata.rs
index db2bc69..560beb9 100644
--- a/src/database/key_value/rooms/metadata.rs
+++ b/src/database/key_value/rooms/metadata.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::RoomId;
 
 use crate::{service, database::KeyValueDatabase, Result, services};
 
-impl service::rooms::metadata::Data for KeyValueDatabase {
+impl service::rooms::metadata::Data for Arc<KeyValueDatabase> {
     fn exists(&self, room_id: &RoomId) -> Result<bool> {
         let prefix = match services().rooms.short.get_shortroomid(room_id)? {
             Some(b) => b.to_be_bytes().to_vec(),
diff --git a/src/database/key_value/rooms/mod.rs b/src/database/key_value/rooms/mod.rs
index 406943e..97c29e5 100644
--- a/src/database/key_value/rooms/mod.rs
+++ b/src/database/key_value/rooms/mod.rs
@@ -15,6 +15,8 @@ mod state_compressor;
 mod timeline;
 mod user;
 
+use std::sync::Arc;
+
 use crate::{database::KeyValueDatabase, service};
 
-impl service::rooms::Data for KeyValueDatabase {}
+impl service::rooms::Data for Arc<KeyValueDatabase> {}
diff --git a/src/database/key_value/rooms/outlier.rs b/src/database/key_value/rooms/outlier.rs
index aa97544..b1ae816 100644
--- a/src/database/key_value/rooms/outlier.rs
+++ b/src/database/key_value/rooms/outlier.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::{EventId, signatures::CanonicalJsonObject};
 
 use crate::{service, database::KeyValueDatabase, PduEvent, Error, Result};
 
-impl service::rooms::outlier::Data for KeyValueDatabase {
+impl service::rooms::outlier::Data for Arc<KeyValueDatabase> {
     fn get_outlier_pdu_json(&self, event_id: &EventId) -> Result<Option<CanonicalJsonObject>> {
         self.eventid_outlierpdu
             .get(event_id.as_bytes())?
diff --git a/src/database/key_value/rooms/pdu_metadata.rs b/src/database/key_value/rooms/pdu_metadata.rs
index f3ac414..f5e8f76 100644
--- a/src/database/key_value/rooms/pdu_metadata.rs
+++ b/src/database/key_value/rooms/pdu_metadata.rs
@@ -4,7 +4,7 @@ use ruma::{RoomId, EventId};
 
 use crate::{service, database::KeyValueDatabase, Result};
 
-impl service::rooms::pdu_metadata::Data for KeyValueDatabase {
+impl service::rooms::pdu_metadata::Data for Arc<KeyValueDatabase> {
     fn mark_as_referenced(&self, room_id: &RoomId, event_ids: &[Arc<EventId>]) -> Result<()> {
         for prev in event_ids {
             let mut key = room_id.as_bytes().to_vec();
diff --git a/src/database/key_value/rooms/search.rs b/src/database/key_value/rooms/search.rs
index dfbdbc6..7b8d278 100644
--- a/src/database/key_value/rooms/search.rs
+++ b/src/database/key_value/rooms/search.rs
@@ -1,10 +1,10 @@
-use std::mem::size_of;
+use std::{mem::size_of, sync::Arc};
 
 use ruma::RoomId;
 
 use crate::{service, database::KeyValueDatabase, utils, Result, services};
 
-impl service::rooms::search::Data for KeyValueDatabase {
+impl service::rooms::search::Data for Arc<KeyValueDatabase> {
     fn index_pdu<'a>(&self, shortroomid: u64, pdu_id: &[u8], message_body: String) -> Result<()> {
         let mut batch = message_body
             .split_terminator(|c: char| !c.is_alphanumeric())
diff --git a/src/database/key_value/rooms/short.rs b/src/database/key_value/rooms/short.rs
index 9129638..9a302b5 100644
--- a/src/database/key_value/rooms/short.rs
+++ b/src/database/key_value/rooms/short.rs
@@ -1,4 +1,6 @@
+use std::sync::Arc;
+
 use crate::{database::KeyValueDatabase, service};
 
-impl service::rooms::short::Data for KeyValueDatabase {
+impl service::rooms::short::Data for Arc<KeyValueDatabase> {
 }
diff --git a/src/database/key_value/rooms/state.rs b/src/database/key_value/rooms/state.rs
index 405939d..527c240 100644
--- a/src/database/key_value/rooms/state.rs
+++ b/src/database/key_value/rooms/state.rs
@@ -1,11 +1,12 @@
 use ruma::{RoomId, EventId};
+use tokio::sync::MutexGuard;
 use std::sync::Arc;
-use std::{sync::MutexGuard, collections::HashSet};
+use std::collections::HashSet;
 use std::fmt::Debug;
 
 use crate::{service, database::KeyValueDatabase, utils, Error, Result};
 
-impl service::rooms::state::Data for KeyValueDatabase {
+impl service::rooms::state::Data for Arc<KeyValueDatabase> {
     fn get_room_shortstatehash(&self, room_id: &RoomId) -> Result<Option<u64>> {
         self.roomid_shortstatehash
             .get(room_id.as_bytes())?
@@ -48,7 +49,7 @@ impl service::rooms::state::Data for KeyValueDatabase {
     fn set_forward_extremities<'a>(
         &self,
         room_id: &RoomId,
-        event_ids: impl IntoIterator<Item = &'a EventId> + Debug,
+        event_ids: &mut dyn Iterator<Item = &'a EventId>,
         _mutex_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room state mutex
     ) -> Result<()> {
         let mut prefix = room_id.as_bytes().to_vec();
diff --git a/src/database/key_value/rooms/state_accessor.rs b/src/database/key_value/rooms/state_accessor.rs
index 4d5bd4a..9af45db 100644
--- a/src/database/key_value/rooms/state_accessor.rs
+++ b/src/database/key_value/rooms/state_accessor.rs
@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use ruma::{EventId, events::StateEventType, RoomId};
 
 #[async_trait]
-impl service::rooms::state_accessor::Data for KeyValueDatabase {
+impl service::rooms::state_accessor::Data for Arc<KeyValueDatabase> {
     async fn state_full_ids(&self, shortstatehash: u64) -> Result<BTreeMap<u64, Arc<EventId>>> {
         let full_state = services().rooms.state_compressor
             .load_shortstatehash_info(shortstatehash)?
diff --git a/src/database/key_value/rooms/state_cache.rs b/src/database/key_value/rooms/state_cache.rs
index 5f05485..bdb8cf8 100644
--- a/src/database/key_value/rooms/state_cache.rs
+++ b/src/database/key_value/rooms/state_cache.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::{UserId, RoomId, events::{AnyStrippedStateEvent, AnySyncStateEvent}, serde::Raw};
 
 use crate::{service, database::KeyValueDatabase, services, Result};
 
-impl service::rooms::state_cache::Data for KeyValueDatabase {
+impl service::rooms::state_cache::Data for Arc<KeyValueDatabase> {
     fn mark_as_once_joined(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
         let mut userroom_id = user_id.as_bytes().to_vec();
         userroom_id.push(0xff);
diff --git a/src/database/key_value/rooms/state_compressor.rs b/src/database/key_value/rooms/state_compressor.rs
index aee1890..e1c0280 100644
--- a/src/database/key_value/rooms/state_compressor.rs
+++ b/src/database/key_value/rooms/state_compressor.rs
@@ -1,8 +1,8 @@
-use std::{collections::HashSet, mem::size_of};
+use std::{collections::HashSet, mem::size_of, sync::Arc};
 
 use crate::{service::{self, rooms::state_compressor::data::StateDiff}, database::KeyValueDatabase, Error, utils, Result};
 
-impl service::rooms::state_compressor::Data for KeyValueDatabase {
+impl service::rooms::state_compressor::Data for Arc<KeyValueDatabase> {
     fn get_statediff(&self, shortstatehash: u64) -> Result<StateDiff> {
         let value = self
             .shortstatehash_statediff
diff --git a/src/database/key_value/rooms/timeline.rs b/src/database/key_value/rooms/timeline.rs
index a3b6c17..2d334b9 100644
--- a/src/database/key_value/rooms/timeline.rs
+++ b/src/database/key_value/rooms/timeline.rs
@@ -5,7 +5,7 @@ use tracing::error;
 
 use crate::{service, database::KeyValueDatabase, utils, Error, PduEvent, Result, services};
 
-impl service::rooms::timeline::Data for KeyValueDatabase {
+impl service::rooms::timeline::Data for Arc<KeyValueDatabase> {
     fn last_timeline_count(&self, sender_user: &UserId, room_id: &RoomId) -> Result<u64> {
         match self
             .lasttimelinecount_cache
diff --git a/src/database/key_value/rooms/user.rs b/src/database/key_value/rooms/user.rs
index 66681e3..4d20b00 100644
--- a/src/database/key_value/rooms/user.rs
+++ b/src/database/key_value/rooms/user.rs
@@ -1,8 +1,10 @@
+use std::sync::Arc;
+
 use ruma::{UserId, RoomId};
 
 use crate::{service, database::KeyValueDatabase, utils, Error, Result, services};
 
-impl service::rooms::user::Data for KeyValueDatabase {
+impl service::rooms::user::Data for Arc<KeyValueDatabase> {
     fn reset_notification_counts(&self, user_id: &UserId, room_id: &RoomId) -> Result<()> {
         let mut userroom_id = user_id.as_bytes().to_vec();
         userroom_id.push(0xff);