From 0f8ecdc625cd5fa66e726d2b47a9fe7f419b5541 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timo=20K=C3=B6sters?= <timo@koesters.xyz>
Date: Tue, 20 Jul 2021 12:41:35 +0200
Subject: [PATCH] fix: state res fixes

---
 Cargo.lock            | 36 +++++++++++++++++-----------------
 Cargo.toml            |  2 +-
 src/database/rooms.rs |  6 +++---
 src/server_server.rs  | 45 ++++++++++++++++++++++---------------------
 4 files changed, 45 insertions(+), 44 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8bfba0e8..a6b87c9d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2015,7 +2015,7 @@ dependencies = [
 [[package]]
 name = "ruma"
 version = "0.2.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "assign",
  "js_int",
@@ -2036,7 +2036,7 @@ dependencies = [
 [[package]]
 name = "ruma-api"
 version = "0.17.1"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "bytes",
  "http",
@@ -2052,7 +2052,7 @@ dependencies = [
 [[package]]
 name = "ruma-api-macros"
 version = "0.17.1"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "proc-macro-crate",
  "proc-macro2",
@@ -2063,7 +2063,7 @@ dependencies = [
 [[package]]
 name = "ruma-appservice-api"
 version = "0.3.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "ruma-api",
  "ruma-common",
@@ -2077,7 +2077,7 @@ dependencies = [
 [[package]]
 name = "ruma-client-api"
 version = "0.11.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "assign",
  "bytes",
@@ -2097,7 +2097,7 @@ dependencies = [
 [[package]]
 name = "ruma-common"
 version = "0.5.4"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "indexmap",
  "js_int",
@@ -2112,7 +2112,7 @@ dependencies = [
 [[package]]
 name = "ruma-events"
 version = "0.23.2"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "indoc",
  "js_int",
@@ -2128,7 +2128,7 @@ dependencies = [
 [[package]]
 name = "ruma-events-macros"
 version = "0.23.2"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "proc-macro-crate",
  "proc-macro2",
@@ -2139,7 +2139,7 @@ dependencies = [
 [[package]]
 name = "ruma-federation-api"
 version = "0.2.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "js_int",
  "ruma-api",
@@ -2154,7 +2154,7 @@ dependencies = [
 [[package]]
 name = "ruma-identifiers"
 version = "0.19.4"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "paste",
  "rand 0.8.4",
@@ -2168,7 +2168,7 @@ dependencies = [
 [[package]]
 name = "ruma-identifiers-macros"
 version = "0.19.4"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "quote",
  "ruma-identifiers-validation",
@@ -2178,12 +2178,12 @@ dependencies = [
 [[package]]
 name = "ruma-identifiers-validation"
 version = "0.4.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 
 [[package]]
 name = "ruma-identity-service-api"
 version = "0.2.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "js_int",
  "ruma-api",
@@ -2196,7 +2196,7 @@ dependencies = [
 [[package]]
 name = "ruma-push-gateway-api"
 version = "0.2.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "js_int",
  "ruma-api",
@@ -2211,7 +2211,7 @@ dependencies = [
 [[package]]
 name = "ruma-serde"
 version = "0.4.1"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "bytes",
  "form_urlencoded",
@@ -2225,7 +2225,7 @@ dependencies = [
 [[package]]
 name = "ruma-serde-macros"
 version = "0.4.1"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "proc-macro-crate",
  "proc-macro2",
@@ -2236,7 +2236,7 @@ dependencies = [
 [[package]]
 name = "ruma-signatures"
 version = "0.8.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "base64 0.13.0",
  "ed25519-dalek",
@@ -2253,7 +2253,7 @@ dependencies = [
 [[package]]
 name = "ruma-state-res"
 version = "0.2.0"
-source = "git+https://github.com/timokoesters/ruma?rev=a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386#a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386"
+source = "git+https://github.com/timokoesters/ruma?rev=74cf83c4ca937fa5e2709fb71e9d11848e72e487#74cf83c4ca937fa5e2709fb71e9d11848e72e487"
 dependencies = [
  "itertools 0.10.1",
  "js_int",
diff --git a/Cargo.toml b/Cargo.toml
index 64d67a14..c62ff821 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,7 +19,7 @@ rocket = { version = "0.5.0-rc.1", features = ["tls"] } # Used to handle request
 
 # Used for matrix spec type definitions and helpers
 #ruma = { git = "https://github.com/ruma/ruma", rev = "c29c2b16ec114fa655e2b70bdd53c82e35859005", features = ["compat", "rand", "appservice-api-c", "client-api", "federation-api", "push-gateway-api-c", "state-res", "unstable-pre-spec", "unstable-exhaustive-types"] }
-ruma = { git = "https://github.com/timokoesters/ruma", rev = "a3fd405d6b331c7bc4c6f366bc1b6ec303b3a386", features = ["compat", "rand", "appservice-api-c", "client-api", "federation-api", "push-gateway-api-c", "state-res", "unstable-pre-spec", "unstable-exhaustive-types"] }
+ruma = { git = "https://github.com/timokoesters/ruma", rev = "74cf83c4ca937fa5e2709fb71e9d11848e72e487", features = ["compat", "rand", "appservice-api-c", "client-api", "federation-api", "push-gateway-api-c", "state-res", "unstable-pre-spec", "unstable-exhaustive-types"] }
 #ruma = { path = "../ruma/crates/ruma", features = ["compat", "rand", "appservice-api-c", "client-api", "federation-api", "push-gateway-api-c", "state-res", "unstable-pre-spec", "unstable-exhaustive-types"] }
 
 # Used for long polling and federation sender, should be the same as rocket::tokio
diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index fa121bda..aad691b4 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -653,9 +653,9 @@ impl Rooms {
         Ok(())
     }
 
-    pub fn is_pdu_referenced(&self, pdu: &PduEvent) -> Result<bool> {
-        let mut key = pdu.room_id().as_bytes().to_vec();
-        key.extend_from_slice(pdu.event_id().as_bytes());
+    pub fn is_event_referenced(&self, room_id: &RoomId, event_id: &EventId) -> Result<bool> {
+        let mut key = room_id.as_bytes().to_vec();
+        key.extend_from_slice(event_id.as_bytes());
         Ok(self.prevevent_parent.get(&key)?.is_some())
     }
 
diff --git a/src/server_server.rs b/src/server_server.rs
index 39a18473..e463bba2 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -6,7 +6,6 @@ use crate::{
 use get_profile_information::v1::ProfileField;
 use http::header::{HeaderValue, AUTHORIZATION, HOST};
 use log::{debug, error, info, trace, warn};
-use lru_cache::LruCache;
 use regex::Regex;
 use rocket::response::content::Json;
 use ruma::{
@@ -1174,6 +1173,9 @@ pub fn handle_incoming_pdu<'a>(
             }
         }
 
+        // Only keep those extremities we don't have in our timeline yet
+        extremities.retain(|id| !matches!(db.rooms.get_non_outlier_pdu_json(id), Ok(Some(_))));
+
         let mut extremity_statehashes = Vec::new();
 
         for id in &extremities {
@@ -1276,7 +1278,6 @@ pub fn handle_incoming_pdu<'a>(
                 })
                 .collect::<Vec<_>>();
 
-            let auth_chain_t = Instant::now();
             let mut auth_chain_sets = Vec::new();
             for state in fork_states {
                 auth_chain_sets.push(
@@ -1284,9 +1285,7 @@ pub fn handle_incoming_pdu<'a>(
                         .map_err(|_| "Failed to load auth chain.".to_owned())?,
                 );
             }
-            dbg!(auth_chain_t.elapsed());
 
-            let state_res_t = Instant::now();
             let state = match state_res::StateResolution::resolve(
                 &room_id,
                 room_version_id,
@@ -1305,7 +1304,6 @@ pub fn handle_incoming_pdu<'a>(
                     return Err("State resolution failed, either an event could not be found or deserialization".into());
                 }
             };
-            dbg!(state_res_t.elapsed());
             state
         };
 
@@ -1726,38 +1724,41 @@ async fn append_incoming_pdu(
 }
 
 fn get_auth_chain(starting_events: Vec<EventId>, db: &Database) -> Result<HashSet<EventId>> {
-    let mut auth_chain_cache = db.rooms.auth_chain_cache();
+    let mut full_auth_chain = HashSet::new();
 
-    let mut auth_chain = HashSet::new();
+    let mut cache = db.rooms.auth_chain_cache();
+    for event_id in starting_events {
+        let auth_chain = if let Some(cached) = cache.get_mut(&event_id) {
+            cached.clone()
+        } else {
+            drop(cache);
+            let auth_chain = get_auth_chain_recursive(&event_id, db)?;
 
-    for event in starting_events {
-        auth_chain.extend(get_auth_chain_recursive(&event, &mut auth_chain_cache, db)?);
+            cache = db.rooms.auth_chain_cache();
+
+            cache.insert(event_id, auth_chain.clone());
+
+            auth_chain
+        };
+
+        full_auth_chain.extend(auth_chain);
     }
 
-    Ok(auth_chain)
+    Ok(full_auth_chain)
 }
 
-fn get_auth_chain_recursive(
-    event_id: &EventId,
-    auth_chain_cache: &mut std::sync::MutexGuard<'_, LruCache<EventId, HashSet<EventId>>>,
-    db: &Database,
-) -> Result<HashSet<EventId>> {
-    if let Some(cached) = auth_chain_cache.get_mut(event_id) {
-        return Ok(cached.clone());
-    }
-
+fn get_auth_chain_recursive(event_id: &EventId, db: &Database) -> Result<HashSet<EventId>> {
     let mut auth_chain = HashSet::new();
 
     if let Some(pdu) = db.rooms.get_pdu(&event_id)? {
+        auth_chain.extend(pdu.auth_events.iter().cloned());
         for auth_event in &pdu.auth_events {
-            auth_chain.extend(get_auth_chain_recursive(&auth_event, auth_chain_cache, db)?);
+            auth_chain.extend(get_auth_chain_recursive(&auth_event, db)?);
         }
     } else {
         warn!("Could not find pdu mentioned in auth events.");
     }
 
-    auth_chain_cache.insert(event_id.clone(), auth_chain.clone());
-
     Ok(auth_chain)
 }