Use auth_cache to avoid db, save state for every event when joining

src/client_server/membership.rs

@@ -643,8 +643,6 @@ async fn join_room_by_id_helper(
         )
         .expect("iterative auth check failed on resolved events");
 
-        let mut state = HashMap::new();
-
         // filter the events that failed the auth check keeping the remaining events
         // sorted correctly
         for ev_id in sorted_event_ids
@@ -660,24 +658,20 @@ async fn join_room_by_id_helper(
             let mut pdu_id = room_id.as_bytes().to_vec();
             pdu_id.push(0xff);
             pdu_id.extend_from_slice(&count.to_be_bytes());
+
+            let hash = db.rooms.append_to_state(&pdu_id, &pdu, &db.globals)?;
+
             db.rooms.append_pdu(
                 &pdu,
                 utils::to_canonical_object(&**pdu).expect("Pdu is valid canonical object"),
                 count,
                 pdu_id.clone().into(),
-                // TODO: can we simplify the DAG or should we copy it exactly??
                 &pdu.prev_events,
                 &db,
             )?;
 
-            if state_events.contains(ev_id) {
-                if let Some(key) = &pdu.state_key {
-                    state.insert((pdu.kind(), key.to_string()), pdu_id);
-                }
-            }
+            db.rooms.set_room_state(room_id, &hash)?;
         }
-
-        db.rooms.force_state(room_id, state, &db.globals)?;
     } else {
         let event = member::MemberEventContent {
             membership: member::MembershipState::Join,

src/database.rs

@@ -161,6 +161,7 @@ pub async fn load_or_create(config: Config) -> Result<Self> {
                 roomid_statehash: db.open_tree("roomid_statehash")?,
                 roomeventid_outlierpdu: db.open_tree("roomeventid_outlierpdu")?,
                 prevevent_parent: db.open_tree("prevevent_parent")?,
+                roomeventid_outlierpducount: db.open_tree("roomeventid_outlierpducount")?,
             },
             account_data: account_data::AccountData {
                 roomuserdataid_accountdata: db.open_tree("roomuserdataid_accountdata")?,

src/database/rooms.rs

@@ -69,6 +69,10 @@ pub struct Rooms {
     /// RoomId + EventId -> outlier PDU.
     /// Any pdu that has passed the steps 1-8 in the incoming event /federation/send/txn.
     pub(super) roomeventid_outlierpdu: sled::Tree,
+    /// RoomId + EventId -> count of the last known pdu when the outlier was inserted.
+    /// This allows us to skip any state snapshots that would for sure not have the outlier.
+    pub(super) roomeventid_outlierpducount: sled::Tree,
+
     /// RoomId + EventId -> Parent PDU EventId.
     pub(super) prevevent_parent: sled::Tree,
 }
@@ -323,6 +327,15 @@ pub fn get_pdu_count(&self, event_id: &EventId) -> Result<Option<u64>> {
             .map_or(Ok(None), |pdu_id| self.pdu_count(&pdu_id).map(Some))
     }
 
+    pub fn latest_pdu_count(&self, room_id: &RoomId) -> Result<u64> {
+        self.pduid_pdu
+            .scan_prefix(room_id.as_bytes())
+            .last()
+            .map(|b| self.pdu_count(&b?.0))
+            .transpose()
+            .map(|op| op.unwrap_or_default())
+    }
+
     /// Returns the json of a pdu.
     pub fn get_pdu_json(&self, event_id: &EventId) -> Result<Option<serde_json::Value>> {
         self.eventid_pduid
@@ -490,6 +503,8 @@ pub fn append_pdu_outlier(&self, pdu: &PduEvent) -> Result<()> {
             &key,
             &*serde_json::to_string(&pdu).expect("PduEvent is always a valid String"),
         )?;
+        self.roomeventid_outlierpducount
+            .insert(&key, &self.latest_pdu_count(pdu.room_id())?.to_be_bytes())?;
         Ok(())
     }
 
@@ -537,7 +552,45 @@ pub fn append_pdu(
         let mut key = pdu.room_id().as_bytes().to_vec();
         key.push(0xff);
         key.extend_from_slice(pdu.event_id().as_bytes());
-        self.roomeventid_outlierpdu.remove(key)?;
+        if self.roomeventid_outlierpdu.remove(&key)?.is_some() {
+            if let Some(state_key) = pdu.state_key.as_deref() {
+                let mut statekey = pdu.kind().as_ref().as_bytes().to_vec();
+                statekey.extend_from_slice(state_key.as_bytes());
+
+                let short = match self.statekey_short.get(&statekey)? {
+                    Some(short) => utils::u64_from_bytes(&short).map_err(|_| {
+                        Error::bad_database("Invalid short bytes in statekey_short.")
+                    })?,
+                    None => {
+                        error!(
+                            "This event has been inserted into the state snapshot tree previously."
+                        );
+                        let short = db.globals.next_count()?;
+                        self.statekey_short
+                            .insert(&statekey, &short.to_be_bytes())?;
+                        short
+                    }
+                };
+
+                let mut start = pdu.room_id().as_bytes().to_vec();
+                start.extend_from_slice(
+                    &self
+                        .roomeventid_outlierpducount
+                        .get(&key)?
+                        .unwrap_or_default(),
+                );
+                for hash in self.pduid_statehash.range(start..).values() {
+                    let mut hash = hash?.to_vec();
+                    hash.extend_from_slice(&short.to_be_bytes());
+
+                    let _ = self.stateid_pduid.compare_and_swap(
+                        hash,
+                        Some(pdu.event_id().as_bytes()),
+                        Some(pdu_id.as_ref()),
+                    )?;
+                }
+            }
+        }
 
         // We must keep track of all events that have been referenced.
         for leaf in leaves {

src/server_server.rs

@@ -971,6 +971,7 @@ fn validate_event<'a>(
                     }
                 }
                 Err(_e) => {
+                    error!("{}", _e);
                     return Err("Signature verification failed".to_string());
                 }
             };
@@ -988,7 +989,7 @@ fn validate_event<'a>(
 
         fetch_check_auth_events(db, origin, pub_key_map, &pdu.auth_events, auth_cache)
             .await
-            .map_err(|_| "Event failed auth chain check".to_string())?;
+            .map_err(|e| e.to_string())?;
 
         let pdu = Arc::new(pdu.clone());
 
@@ -1064,6 +1065,7 @@ async fn fetch_check_auth_events(
 /// Find the event and auth it. Once the event is validated (steps 1 - 8)
 /// it is appended to the outliers Tree.
 ///
+/// 0. Look in the auth_cache
 /// 1. Look in the main timeline (pduid_pdu tree)
 /// 2. Look at outlier pdu tree
 /// 3. Ask origin server over federation
@@ -1080,28 +1082,35 @@ pub(crate) async fn fetch_events(
 ) -> Result<Vec<Arc<PduEvent>>> {
     let mut pdus = vec![];
     for id in events {
-        // `get_pdu` checks the outliers tree for us
-        let pdu = match db.rooms.get_pdu(&id)? {
-            Some(pdu) => Arc::new(pdu),
-            None => match db
-                .sending
-                .send_federation_request(
-                    &db.globals,
-                    origin,
-                    get_event::v1::Request { event_id: &id },
-                )
-                .await
-            {
-                Ok(res) => {
-                    let (event_id, value) = crate::pdu::gen_event_id_canonical_json(&res.pdu);
-                    let (pdu, _) = validate_event(db, value, event_id, key_map, origin, auth_cache)
-                        .await
-                        .map_err(|_| Error::Conflict("Authentication of event failed"))?;
-
-                    db.rooms.append_pdu_outlier(&pdu)?;
-                    pdu
-                }
-                Err(_) => return Err(Error::BadServerResponse("Failed to fetch event")),
+        let pdu = match auth_cache.get(id) {
+            Some(pdu) => pdu.clone(),
+            // `get_pdu` checks the outliers tree for us
+            None => match db.rooms.get_pdu(&id)? {
+                Some(pdu) => Arc::new(pdu),
+                None => match db
+                    .sending
+                    .send_federation_request(
+                        &db.globals,
+                        origin,
+                        get_event::v1::Request { event_id: &id },
+                    )
+                    .await
+                {
+                    Ok(res) => {
+                        let (event_id, value) = crate::pdu::gen_event_id_canonical_json(&res.pdu);
+                        let (pdu, _) =
+                            validate_event(db, value, event_id, key_map, origin, auth_cache)
+                                .await
+                                .map_err(|e| {
+                                    error!("{:?}", e);
+                                    Error::Conflict("Authentication of event failed")
+                                })?;
+
+                        db.rooms.append_pdu_outlier(&pdu)?;
+                        pdu
+                    }
+                    Err(_) => return Err(Error::BadServerResponse("Failed to fetch event")),
+                },
             },
         };
         auth_cache.entry(id.clone()).or_insert_with(|| pdu.clone());