From 88c60605b457d0163b0b8d427e51cd07b0dd1f4c Mon Sep 17 00:00:00 2001
From: Devin Ragotzy <devin.ragotzy@gmail.com>
Date: Mon, 18 Jan 2021 19:41:38 -0500
Subject: [PATCH] Add ability to update room leaves with multiple eventIds

Tokio seems a bit broken with Rocket...
---
 src/client_server/membership.rs |  2 ++
 src/database/rooms.rs           | 37 ++++++++++++++++-----------------
 src/server_server.rs            |  7 ++-----
 3 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/src/client_server/membership.rs b/src/client_server/membership.rs
index 70bb480c..11591854 100644
--- a/src/client_server/membership.rs
+++ b/src/client_server/membership.rs
@@ -665,6 +665,8 @@ async fn join_room_by_id_helper(
                 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,
             )?;
 
diff --git a/src/database/rooms.rs b/src/database/rooms.rs
index 665e3287..a3f3aab8 100644
--- a/src/database/rooms.rs
+++ b/src/database/rooms.rs
@@ -397,8 +397,11 @@ impl Rooms {
         Ok(events)
     }
 
-    /// Force an update to the leaves of a room.
-    pub fn force_pdu_leaves(&self, room_id: &RoomId, event_ids: &[EventId]) -> Result<()> {
+    /// Replace the leaves of a room.
+    ///
+    /// The provided `event_ids` become the new leaves, this enables an event having multiple
+    /// `prev_events`.
+    pub fn replace_pdu_leaves(&self, room_id: &RoomId, event_ids: &[EventId]) -> Result<()> {
         let mut prefix = room_id.as_bytes().to_vec();
         prefix.push(0xff);
 
@@ -415,21 +418,6 @@ impl Rooms {
         Ok(())
     }
 
-    /// Replace the leaves of a room with a new event.
-    pub fn replace_pdu_leaves(&self, room_id: &RoomId, event_id: &EventId) -> Result<()> {
-        let mut prefix = room_id.as_bytes().to_vec();
-        prefix.push(0xff);
-
-        for key in self.roomid_pduleaves.scan_prefix(&prefix).keys() {
-            self.roomid_pduleaves.remove(key?)?;
-        }
-
-        prefix.extend_from_slice(event_id.as_bytes());
-        self.roomid_pduleaves.insert(&prefix, event_id.as_bytes())?;
-
-        Ok(())
-    }
-
     /// Returns the pdu from the outlier tree.
     pub fn get_pdu_outlier(&self, event_id: &EventId) -> Result<Option<PduEvent>> {
         self.eventid_outlierpdu
@@ -465,6 +453,7 @@ impl Rooms {
         mut pdu_json: CanonicalJsonObject,
         count: u64,
         pdu_id: IVec,
+        leaves: &[EventId],
         db: &Database,
     ) -> Result<()> {
         // Make unsigned fields correct. This is not properly documented in the spec, but state
@@ -497,7 +486,7 @@ impl Rooms {
         // We no longer keep this pdu as an outlier
         self.eventid_outlierpdu.remove(pdu.event_id().as_bytes())?;
 
-        self.replace_pdu_leaves(&pdu.room_id, &pdu.event_id)?;
+        self.replace_pdu_leaves(&pdu.room_id, leaves)?;
 
         // Mark as read first so the sending client doesn't get a notification even if appending
         // fails
@@ -943,7 +932,17 @@ impl Rooms {
         // pdu without it's state. This is okay because append_pdu can't fail.
         let statehashid = self.append_to_state(&pdu_id, &pdu, &db.globals)?;
 
-        self.append_pdu(&pdu, pdu_json, count, pdu_id.clone().into(), db)?;
+        // remove the
+        self.append_pdu(
+            &pdu,
+            pdu_json,
+            count,
+            pdu_id.clone().into(),
+            // Since this PDU references all pdu_leaves we can update the leaves
+            // of the room
+            &[pdu.event_id.clone()],
+            db,
+        )?;
 
         // We set the room state after inserting the pdu, so that we never have a moment in time
         // where events in the current room state do not exist
diff --git a/src/server_server.rs b/src/server_server.rs
index 16a1a8e9..f782ad5a 100644
--- a/src/server_server.rs
+++ b/src/server_server.rs
@@ -24,7 +24,7 @@ use ruma::{
 };
 use state_res::{Event, EventMap, StateMap};
 use std::{
-    collections::{BTreeMap, BTreeSet, HashSet},
+    collections::{BTreeMap, BTreeSet},
     convert::TryFrom,
     fmt::Debug,
     future::Future,
@@ -1245,13 +1245,10 @@ fn append_state(db: &Database, pdu: &PduEvent, new_room_leaves: &[EventId]) -> R
         utils::to_canonical_object(pdu).expect("Pdu is valid canonical object"),
         count,
         pdu_id.clone().into(),
+        &new_room_leaves,
         &db,
     )?;
 
-    // If we update the room leaves after calling append_pdu it will stick since append_pdu
-    // calls replace_pdu_leaves with only the given event.
-    db.rooms.force_pdu_leaves(pdu.room_id(), new_room_leaves)?;
-
     // We set the room state after inserting the pdu, so that we never have a moment in time
     // where events in the current room state do not exist
     db.rooms.set_room_state(&pdu.room_id, &statehashid)?;