improvement: locks

src/database/admin.rs

@@ -10,7 +10,7 @@ use ruma::{
     events::{room::message, EventType},
     UserId,
 };
-use tokio::sync::{RwLock, RwLockReadGuard};
+use tokio::sync::{MutexGuard, RwLock, RwLockReadGuard};
 
 pub enum AdminCommand {
     RegisterAppservice(serde_yaml::Value),
@@ -48,38 +48,51 @@ impl Admin {
                 )
                 .unwrap();
 
-            if conduit_room.is_none() {
-                warn!("Conduit instance does not have an #admins room. Logging to that room will not work. Restart Conduit after creating a user to fix this.");
-            }
+            let conduit_room = match conduit_room {
+                None => {
+                    warn!("Conduit instance does not have an #admins room. Logging to that room will not work. Restart Conduit after creating a user to fix this.");
+                    return;
+                }
+                Some(r) => r,
+            };
 
             drop(guard);
 
-            let send_message =
-                |message: message::MessageEventContent, guard: RwLockReadGuard<'_, Database>| {
-                    if let Some(conduit_room) = &conduit_room {
-                        guard
-                            .rooms
-                            .build_and_append_pdu(
-                                PduBuilder {
-                                    event_type: EventType::RoomMessage,
-                                    content: serde_json::to_value(message)
-                                        .expect("event is valid, we just created it"),
-                                    unsigned: None,
-                                    state_key: None,
-                                    redacts: None,
-                                },
-                                &conduit_user,
-                                &conduit_room,
-                                &guard,
-                            )
-                            .unwrap();
-                    }
-                };
+            let send_message = |message: message::MessageEventContent,
+                                guard: RwLockReadGuard<'_, Database>,
+                                mutex_lock: &MutexGuard<'_, ()>| {
+                guard
+                    .rooms
+                    .build_and_append_pdu(
+                        PduBuilder {
+                            event_type: EventType::RoomMessage,
+                            content: serde_json::to_value(message)
+                                .expect("event is valid, we just created it"),
+                            unsigned: None,
+                            state_key: None,
+                            redacts: None,
+                        },
+                        &conduit_user,
+                        &conduit_room,
+                        &guard,
+                        mutex_lock,
+                    )
+                    .unwrap();
+            };
 
             loop {
                 tokio::select! {
                     Some(event) = receiver.next() => {
                         let guard = db.read().await;
+                        let mutex = Arc::clone(
+                            guard.globals
+                                .roomid_mutex
+                                .write()
+                                .unwrap()
+                                .entry(conduit_room.clone())
+                                .or_default(),
+                        );
+                        let mutex_lock = mutex.lock().await;
 
                         match event {
                             AdminCommand::RegisterAppservice(yaml) => {
@@ -93,15 +106,17 @@ impl Admin {
                                         count,
                                         appservices.into_iter().filter_map(|r| r.ok()).collect::<Vec<_>>().join(", ")
                                     );
-                                    send_message(message::MessageEventContent::text_plain(output), guard);
+                                    send_message(message::MessageEventContent::text_plain(output), guard, &mutex_lock);
                                 } else {
-                                    send_message(message::MessageEventContent::text_plain("Failed to get appservices."), guard);
+                                    send_message(message::MessageEventContent::text_plain("Failed to get appservices."), guard, &mutex_lock);
                                 }
                             }
                             AdminCommand::SendMessage(message) => {
-                                send_message(message, guard)
+                                send_message(message, guard, &mutex_lock);
                             }
                         }
+
+                        drop(mutex_lock);
                     }
                 }
             }

src/database/globals.rs

@@ -46,6 +46,7 @@ pub struct Globals {
     pub servername_ratelimiter: Arc<RwLock<BTreeMap<Box<ServerName>, Arc<Semaphore>>>>,
     pub sync_receivers: RwLock<BTreeMap<(UserId, Box<DeviceId>), SyncHandle>>,
     pub roomid_mutex: RwLock<BTreeMap<RoomId, Arc<Mutex<()>>>>,
+    pub roomid_mutex_federation: RwLock<BTreeMap<RoomId, Arc<Mutex<()>>>>, // this lock will be held longer
     pub rotate: RotationHandler,
 }
 
@@ -199,6 +200,7 @@ impl Globals {
             bad_signature_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())),
             servername_ratelimiter: Arc::new(RwLock::new(BTreeMap::new())),
             roomid_mutex: RwLock::new(BTreeMap::new()),
+            roomid_mutex_federation: RwLock::new(BTreeMap::new()),
             sync_receivers: RwLock::new(BTreeMap::new()),
             rotate: RotationHandler::new(),
         };

src/database/rooms.rs

@@ -2,6 +2,7 @@ mod edus;
 
 pub use edus::RoomEdus;
 use member::MembershipState;
+use tokio::sync::MutexGuard;
 
 use crate::{pdu::PduBuilder, utils, Database, Error, PduEvent, Result};
 use log::{debug, error, warn};
@@ -1207,6 +1208,7 @@ impl Rooms {
         sender: &UserId,
         room_id: &RoomId,
         db: &Database,
+        _mutex_lock: &MutexGuard<'_, ()>, // Take mutex guard to make sure users get the room mutex
     ) -> Result<EventId> {
         let PduBuilder {
             event_type,
@@ -1216,7 +1218,6 @@ impl Rooms {
             redacts,
         } = pdu_builder;
 
-        // TODO: Make sure this isn't called twice in parallel
         let prev_events = self
             .get_pdu_leaves(&room_id)?
             .into_iter()
@@ -1790,6 +1791,16 @@ impl Rooms {
                 db,
             )?;
         } else {
+            let mutex = Arc::clone(
+                db.globals
+                    .roomid_mutex
+                    .write()
+                    .unwrap()
+                    .entry(room_id.clone())
+                    .or_default(),
+            );
+            let mutex_lock = mutex.lock().await;
+
             let mut event = serde_json::from_value::<Raw<member::MemberEventContent>>(
                 self.room_state_get(room_id, &EventType::RoomMember, &user_id.to_string())?
                     .ok_or(Error::BadRequest(
@@ -1817,6 +1828,7 @@ impl Rooms {
                 user_id,
                 room_id,
                 db,
+                &mutex_lock,
             )?;
         }