Use handle locks in WalStore. (fixes #974)

src/commons/eventsourcing/locks.rs

@@ -0,0 +1,91 @@
+//! Support locking on Handles so that updates can be
+//! performed sequentially. Useful for both event sourced
+//! types (Aggregates) as well as write-ahead logging
+//! types.
+
+//------------ HandleLocks ---------------------------------------------------
+
+use std::{
+    collections::HashMap,
+    sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
+};
+
+use rpki::ca::idexchange::MyHandle;
+
+#[derive(Debug, Default)]
+struct HandleLockMap(HashMap<MyHandle, RwLock<()>>);
+
+impl HandleLockMap {
+    fn create_handle_lock(&mut self, handle: MyHandle) {
+        self.0.insert(handle, RwLock::new(()));
+    }
+
+    fn has_handle(&self, handle: &MyHandle) -> bool {
+        self.0.contains_key(handle)
+    }
+
+    fn drop_handle_lock(&mut self, handle: &MyHandle) {
+        self.0.remove(handle);
+    }
+}
+
+pub struct HandleLock<'a> {
+    // Needs a read reference to the map that holds the RwLock
+    // for the handle.
+    map: RwLockReadGuard<'a, HandleLockMap>,
+    handle: MyHandle,
+}
+
+impl HandleLock<'_> {
+    // panics if there is no entry for the handle.
+    pub fn read(&self) -> RwLockReadGuard<'_, ()> {
+        self.map.0.get(&self.handle).unwrap().read().unwrap()
+    }
+
+    // panics if there is no entry for the handle.
+    pub fn write(&self) -> RwLockWriteGuard<'_, ()> {
+        self.map.0.get(&self.handle).unwrap().write().unwrap()
+    }
+}
+
+/// This structure is used to ensure that we have unique access to an instance for a [`Handle`]
+/// managed in an [`AggregateStore`] or [`WalStore`]. Currently uses a `std::sync::RwLock`, but
+/// this should be improved to use an async lock instead (e.g. `tokio::sync::RwLock`).
+/// This has not been done yet, because that change is quite pervasive.
+#[derive(Debug, Default)]
+pub struct HandleLocks {
+    locks: RwLock<HandleLockMap>,
+}
+
+impl HandleLocks {
+    pub fn for_handle(&self, handle: MyHandle) -> HandleLock<'_> {
+        {
+            // Return the lock *if* there is an entry for the handle
+            let map = self.locks.read().unwrap();
+            if map.has_handle(&handle) {
+                return HandleLock { map, handle };
+            }
+        }
+
+        {
+            // There was no entry.. try to create an entry for the
+            // handle.
+            let mut map = self.locks.write().unwrap();
+
+            // But.. first check again, because we could have had a
+            // race condition if two threads call this function.
+            if !map.has_handle(&handle) {
+                map.create_handle_lock(handle.clone());
+            }
+        }
+
+        // Entry probably exists now, but recurse in case the entry
+        // was dropped immediately after creation.
+        self.for_handle(handle)
+    }
+
+    pub fn drop_handle(&self, handle: &MyHandle) {
+        let mut map = self.locks.write().unwrap();
+        map.drop_handle_lock(handle);
+    }
+}

src/commons/eventsourcing/mod.rs

@@ -18,6 +18,8 @@ pub use self::store::*;
 mod listener;
 pub use self::listener::{EventCounter, PostSaveEventListener, PreSaveEventListener};
 
+pub mod locks;
+
 mod kv;
 pub use self::kv::*;
 

src/commons/eventsourcing/store.rs

@@ -3,26 +3,25 @@ use std::{
     fmt,
     path::Path,
     str::FromStr,
-    sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
+    sync::{Arc, RwLock},
 };
 
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use rpki::{ca::idexchange::MyHandle, repository::x509::Time};
 
-use crate::commons::eventsourcing::{
-    cmd::{Command, StoredCommandBuilder},
-    Aggregate, Event, KeyStoreKey, KeyValueError, KeyValueStore, PostSaveEventListener, StoredCommand,
-    WithStorableDetails,
-};
 use crate::commons::{
     api::{CommandHistory, CommandHistoryCriteria, CommandHistoryRecord, Label},
     error::KrillIoError,
+    eventsourcing::{
+        cmd::{Command, StoredCommandBuilder},
+        locks::HandleLocks,
+        Aggregate, Event, KeyStoreKey, KeyValueError, KeyValueStore, PostSaveEventListener, PreSaveEventListener,
+        StoredCommand, WithStorableDetails,
+    },
     util::KrillVersion,
 };
 
-use super::PreSaveEventListener;
-
 pub type StoreResult<T> = Result<T, AggregateStoreError>;
 
 //------------ Storable ------------------------------------------------------
@@ -125,86 +124,6 @@ impl fmt::Display for CommandKeyError {
     }
 }
 
-//------------ AggregateLocks ------------------------------------------------
-
-#[derive(Debug, Default)]
-struct AggregateLockMap(HashMap<MyHandle, RwLock<()>>);
-
-impl AggregateLockMap {
-    fn create_handle_lock(&mut self, handle: MyHandle) {
-        self.0.insert(handle, RwLock::new(()));
-    }
-
-    fn has_handle(&self, handle: &MyHandle) -> bool {
-        self.0.contains_key(handle)
-    }
-
-    fn drop_handle_lock(&mut self, handle: &MyHandle) {
-        self.0.remove(handle);
-    }
-}
-
-struct AggregateLock<'a> {
-    // Needs a read reference to the map that holds the RwLock
-    // for the handle for this aggregate.
-    map: RwLockReadGuard<'a, AggregateLockMap>,
-    handle: MyHandle,
-}
-
-impl AggregateLock<'_> {
-    // panics if there is no entry for the handle.
-    fn read(&self) -> RwLockReadGuard<'_, ()> {
-        self.map.0.get(&self.handle).unwrap().read().unwrap()
-    }
-
-    // panics if there is no entry for the handle.
-    fn write(&self) -> RwLockWriteGuard<'_, ()> {
-        self.map.0.get(&self.handle).unwrap().write().unwrap()
-    }
-}
-
-/// This structure is used to ensure that we have unique access to aggregates
-/// managed in an [`AggregateStore`]. Currently uses a `std::sync::RwLock`, but
-/// this should be improved to use an async lock instead (e.g. `tokio::sync::RwLock`).
-/// This has not been done yet, because that change is quite pervasive.
-#[derive(Debug, Default)]
-struct AggregateLocks {
-    locks: RwLock<AggregateLockMap>,
-}
-
-impl AggregateLocks {
-    fn for_aggregate(&self, handle: MyHandle) -> AggregateLock<'_> {
-        {
-            // Return the lock *if* there is an entry for the handle
-            let map = self.locks.read().unwrap();
-            if map.has_handle(&handle) {
-                return AggregateLock { map, handle };
-            }
-        }
-
-        {
-            // There was no entry.. try to create an entry for the
-            // handle.
-            let mut map = self.locks.write().unwrap();
-
-            // But.. first check again, because we could have had a
-            // race condition if two threads call this function.
-            if !map.has_handle(&handle) {
-                map.create_handle_lock(handle.clone());
-            }
-        }
-
-        // Entry exists now, so return the lock
-        let map = self.locks.read().unwrap();
-        AggregateLock { map, handle }
-    }
-
-    fn drop_aggregate(&self, handle: &MyHandle) {
-        let mut map = self.locks.write().unwrap();
-        map.drop_handle_lock(handle);
-    }
-}
-
 //------------ AggregateStore ------------------------------------------------
 
 /// This type is responsible for managing aggregates.
@@ -213,7 +132,7 @@ pub struct AggregateStore<A: Aggregate> {
     cache: RwLock<HashMap<MyHandle, Arc<A>>>,
     pre_save_listeners: Vec<Arc<dyn PreSaveEventListener<A>>>,
     post_save_listeners: Vec<Arc<dyn PostSaveEventListener<A>>>,
-    locks: AggregateLocks,
+    locks: HandleLocks,
 }
 
 /// # Starting up
@@ -232,7 +151,7 @@ where
         let cache = RwLock::new(HashMap::new());
         let pre_save_listeners = vec![];
         let post_save_listeners = vec![];
-        let locks = AggregateLocks::default();
+        let locks = HandleLocks::default();
 
         let store = AggregateStore {
             kv,
@@ -441,7 +360,7 @@ where
     /// an AggregateStoreError::UnknownAggregate in case the aggregate
     /// does not exist.
     pub fn get_latest(&self, handle: &MyHandle) -> StoreResult<Arc<A>> {
-        let agg_lock = self.locks.for_aggregate(handle.clone());
+        let agg_lock = self.locks.for_handle(handle.clone());
         let _read_lock = agg_lock.read();
         self.get_latest_no_lock(handle)
     }
@@ -450,7 +369,7 @@ where
     pub fn add(&self, init: A::InitEvent) -> StoreResult<Arc<A>> {
         let handle = init.handle().clone();
 
-        let agg_lock = self.locks.for_aggregate(handle.clone());
+        let agg_lock = self.locks.for_handle(handle.clone());
         let _write_lock = agg_lock.write();
 
         self.store_event(&init)?;
@@ -485,7 +404,7 @@ where
         debug!("Processing command {}", cmd);
         let handle = cmd.handle().clone();
 
-        let agg_lock = self.locks.for_aggregate(handle.clone());
+        let agg_lock = self.locks.for_handle(handle.clone());
         let _write_lock = agg_lock.write();
 
         let mut info = self.get_info(&handle)?;
@@ -1056,7 +975,7 @@ where
     pub fn drop_aggregate(&self, id: &MyHandle) -> Result<(), AggregateStoreError> {
         {
             // First get write access - ensure that no one is using this
-            let agg_lock = self.locks.for_aggregate(id.clone());
+            let agg_lock = self.locks.for_handle(id.clone());
             let _write_lock = agg_lock.write();
 
             self.cache_remove(id);
@@ -1065,7 +984,7 @@ where
 
         // Then drop the lock for this aggregate immediately. The write lock is
         // out of scope now, to ensure we do not get into a deadlock.
-        self.locks.drop_aggregate(id);
+        self.locks.drop_handle(id);
         Ok(())
     }
 

src/commons/eventsourcing/wal.rs

@@ -8,7 +8,7 @@ use std::{
 
 use rpki::ca::idexchange::MyHandle;
 
-use super::{KeyStoreKey, KeyValueError, KeyValueStore, Storable};
+use crate::commons::eventsourcing::{locks::HandleLocks, KeyStoreKey, KeyValueError, KeyValueStore, Storable};
 
 //------------ WalSupport ----------------------------------------------------
 
@@ -120,6 +120,7 @@ impl<T: WalSupport> WalSet<T> {
 pub struct WalStore<T: WalSupport> {
     kv: KeyValueStore,
     cache: RwLock<HashMap<MyHandle, Arc<T>>>,
+    locks: HandleLocks,
 }
 
 impl<T: WalSupport> WalStore<T> {
@@ -131,8 +132,9 @@ impl<T: WalSupport> WalStore<T> {
 
         let kv = KeyValueStore::disk(krill_data_dir, name_space)?;
         let cache = RwLock::new(HashMap::new());
+        let locks = HandleLocks::default();
 
-        Ok(WalStore { kv, cache })
+        Ok(WalStore { kv, cache, locks })
     }
 
     /// Warms up the store: caches all instances.
@@ -149,6 +151,9 @@ impl<T: WalSupport> WalStore<T> {
 
     /// Add a new entity for the given handle. Fails if the handle is in use.
     pub fn add(&self, handle: &MyHandle, instance: T) -> WalStoreResult<()> {
+        let handle_lock = self.locks.for_handle(handle.clone());
+        let _write = handle_lock.write();
+
         let instance = Arc::new(instance);
         let key = Self::key_for_snapshot(handle);
         self.kv.store_new(&key, &instance)?; // Fails if this key exists
@@ -168,6 +173,17 @@ impl<T: WalSupport> WalStore<T> {
     /// from the keystore. Then it will check whether there are any further
     /// changes.
     pub fn get_latest(&self, handle: &MyHandle) -> WalStoreResult<Arc<T>> {
+        let handle_lock = self.locks.for_handle(handle.clone());
+        let _read = handle_lock.read();
+
+        self.get_latest_no_lock(handle)
+    }
+
+    /// Get the latest revision without using a lock.
+    ///
+    /// Intended to be used by public functions which manage the locked read/write access
+    /// to this instance for this handle.
+    fn get_latest_no_lock(&self, handle: &MyHandle) -> WalStoreResult<Arc<T>> {
         let mut instance = match self.cache.read().unwrap().get(handle).cloned() {
             None => Arc::new(self.get_snapshot(handle)?),
             Some(instance) => instance,
@@ -209,8 +225,23 @@ impl<T: WalSupport> WalStore<T> {
         if !self.has(handle)? {
             Err(WalStoreError::Unknown(handle.clone()))
         } else {
-            self.cache.write().unwrap().remove(handle);
-            self.kv.drop_scope(handle.as_str())?;
+            {
+                // First get a lock and remove the object
+                let handle_lock = self.locks.for_handle(handle.clone());
+                let _write = handle_lock.write();
+                self.cache.write().unwrap().remove(handle);
+                self.kv.drop_scope(handle.as_str())?;
+            }
+
+            // Then drop the lock for it as well. We could not do this
+            // while holding the write lock.
+            //
+            // Note that the corresponding entity was removed from the key
+            // value store while we had a write lock for its handle.
+            // So, even if another concurrent thread would now try to update
+            // this same entity, that update would fail because the entity
+            // no longer exists.
+            self.locks.drop_handle(handle);
             Ok(())
         }
     }
@@ -246,7 +277,11 @@ impl<T: WalSupport> WalStore<T> {
     ///
     pub fn send_command(&self, command: T::Command) -> Result<Arc<T>, T::Error> {
         let handle = command.handle().clone();
-        let mut latest = self.get_latest(&handle)?;
+
+        let handle_lock = self.locks.for_handle(handle.clone());
+        let _write = handle_lock.write();
+
+        let mut latest = self.get_latest_no_lock(&handle)?;
 
         let summary = command.to_string();
         let revision = latest.revision();
@@ -285,6 +320,17 @@ impl<T: WalSupport> WalStore<T> {
     /// This is a separate function because serializing a large instance can
     /// be expensive.
     pub fn update_snapshot(&self, handle: &MyHandle, archive: bool) -> WalStoreResult<()> {
+        // Note that we do not need to keep a lock for the instance when we update the snapshot.
+        // This function just updates the latest snapshot in the key value store, and it removes
+        // or archives all write-ahead log ("wal-") changes predating the new snapshot.
+        //
+        // It is fine if another thread gets the entity for this handle and updates it while we
+        // do do this. As it turns out, writing snapshots can be expensive for large objects, so
+        // we do not want block updates while we do this.
+        //
+        // This function is intended to be called in the back-ground at regular (slow) intervals
+        // so any updates that were just missed will simply be folded in to the new snapshot when
+        // this function is called again.
         let latest = self.get_latest(handle)?;
         let key = Self::key_for_snapshot(handle);
         self.kv.store(&key, &latest)?;

src/test.rs

@@ -686,6 +686,22 @@ pub async fn wait_for_nr_cas_under_testbed(nr: usize) -> bool {
     false
 }
 
+pub async fn wait_for_nr_cas_under_publication_server(publishers_expected: usize) {
+    let mut publishers_found = list_publishers().await.publishers().len();
+    for _ in 0..300 {
+        if publishers_found == publishers_expected {
+            return;
+        }
+        sleep_seconds(1).await;
+        publishers_found = list_publishers().await.publishers().len();
+    }
+
+    panic!(
+        "Expected {} publishers, but found {}",
+        publishers_expected, publishers_found
+    );
+}
+
 pub async fn list_publishers() -> PublisherList {
     match krill_embedded_pubd_admin(PubServerCommand::PublisherList).await {
         ApiResponse::PublisherList(pub_list) => pub_list,

tests/benchmark.rs

@@ -18,6 +18,8 @@ async fn benchmark() {
     start_krill(config).await;
 
     assert!(wait_for_nr_cas_under_testbed(cas).await);
+    // We expect all CAs, plus the testbed and the ta as publishers
+    wait_for_nr_cas_under_publication_server(cas + 2).await;
 
     let _ = fs::remove_dir_all(dir);
 }