Fix pool opening race and closing stutter

sync-server/src/app_state.rs

@@ -1,5 +1,6 @@
 pub mod cursors;
 pub mod database;
+pub mod weak_slots;
 pub mod websocket;
 
 use anyhow::Result;

sync-server/src/app_state/database.rs

@@ -23,6 +23,7 @@ mod queries;
 mod tests;
 
 use crate::{
+    app_state::weak_slots::get_or_create,
     config::database_config::DatabaseConfig,
     utils::normalize_vault_id::{normalize_string, validate_vault_id},
 };
@@ -147,28 +148,44 @@ impl Database {
     async fn get_connection_pool(&self, vault: &VaultId) -> Result<Pool<Sqlite>> {
         let vault = normalize_string(vault);
         validate_vault_id(&vault)?;
-        let mut pools = self.connection_pools.lock().await;
 
-        if !pools.contains_key(&vault) {
+        if let Some(pool) = self.existing_pool(&vault).await {
+            return Ok(pool);
+        }
+
+        let opening = get_or_create(self.opening_pools.lock(), vault.clone(), || {
+            Arc::new(Mutex::new(()))
+        })
+        .await;
+
+        let _opening = opening.lock().await;
+
+        if let Some(pool) = self.existing_pool(&vault).await {
+            return Ok(pool);
+        }
+
         let file_name = Self::database_path(&self.config, &vault);
         let pool = Self::open_database(&self.config, &file_name).await?;
-            pools.insert(
-                vault.clone(),
+
+        self.connection_pools.lock().await.insert(
+            vault,
             PoolWithTimestamp {
-                    pool,
+                pool: pool.clone(),
                 last_accessed: Instant::now(),
             },
         );
+        Ok(pool)
     }
 
-        let pool_with_timestamp = pools
-            .get_mut(&vault)
-            .expect("Pool was just inserted or already exists");
-
-        // Update last accessed time
-        pool_with_timestamp.last_accessed = Instant::now();
-
-        Ok(pool_with_timestamp.pool.clone())
+    async fn existing_pool(&self, vault: &VaultId) -> Option<Pool<Sqlite>> {
+        self.connection_pools
+            .lock()
+            .await
+            .get_mut(vault)
+            .map(|entry| {
+                entry.last_accessed = Instant::now();
+                entry.pool.clone()
+            })
     }
 
     /// Attempting to write from this transaction might result in a
@@ -203,16 +220,25 @@ impl Database {
             .iter()
             .filter(|(_, pool_with_timestamp)| {
                 now.duration_since(pool_with_timestamp.last_accessed) > idle_timeout
+                    && pool_with_timestamp.pool.num_idle()
+                        == pool_with_timestamp.pool.size() as usize
             })
             .map(|(vault_id, _)| vault_id.clone())
             .collect();
 
-        // Close and remove idle pools
-        for vault_id in &vaults_to_remove {
-            if let Some(pool_with_timestamp) = pools.remove(vault_id) {
-                info!("Closing idle database connection pool for vault `{vault_id}`");
-                pool_with_timestamp.pool.close().await;
-            }
+        let closing: Vec<_> = vaults_to_remove
+            .into_iter()
+            .filter_map(|vault| pools.remove(&vault).map(|entry| (vault, entry.pool)))
+            .collect();
+
+        // Don't hold the pools mutex while closing the removed vaults
+        drop(pools);
+
+        // A checkout racing cleanup can delay close. It must never retain the
+        // mutex used to access every other vault in the process.
+        for (vault, pool) in closing {
+            info!("Closing idle database connection pool for vault `{vault}`");
+            pool.close().await;
         }
     }
 

sync-server/src/app_state/weak_slots.rs

@@ -0,0 +1,26 @@
+use std::{
+    collections::HashMap,
+    future::Future,
+    hash::Hash,
+    ops::DerefMut,
+    sync::{Arc, Weak},
+};
+
+/// Get-or-create an `Arc<T>` for `key` in a weak map behind a mutex
+pub async fn get_or_create<K, T, G, L>(lock: L, key: K, create: impl FnOnce() -> Arc<T>) -> Arc<T>
+where
+    K: Clone + Eq + Hash,
+    G: DerefMut<Target = HashMap<K, Weak<T>>>,
+    L: Future<Output = G>,
+{
+    let mut map = lock.await;
+    map.retain(|_, value| value.strong_count() > 0);
+
+    if let Some(value) = map.get(&key).and_then(Weak::upgrade) {
+        value
+    } else {
+        let value = create();
+        map.insert(key, Arc::downgrade(&value));
+        value
+    }
+}

sync-server/src/app_state/websocket/broadcasts.rs

@@ -1,5 +1,11 @@
 use super::models::CursorPositionFromServer;
-use crate::{app_state::database::models::VaultId, config::server_config::ServerConfig};
+use crate::{
+    app_state::{
+        database::models::VaultId,
+        weak_slots::get_or_create,
+    },
+    config::server_config::ServerConfig,
+};
 use std::{collections::HashMap, sync::Arc};
 use tokio::sync::{Mutex, broadcast};
 
@@ -25,19 +31,13 @@ impl Broadcasts {
         }
     }
 
-    pub fn try_admit(&self, vault: &VaultId) -> Option<tokio::sync::OwnedSemaphorePermit> {
-        let mut limits = self.admission.lock().unwrap();
-
-        limits.retain(|_, limit| limit.strong_count() > 0);
-
-        let limit = limits
-            .get(vault)
-            .and_then(std::sync::Weak::upgrade)
-            .unwrap_or_else(|| {
-                let limit = Arc::new(tokio::sync::Semaphore::new(self.capacity));
-                limits.insert(vault.clone(), Arc::downgrade(&limit));
-                limit
-            });
+    pub async fn try_admit(&self, vault: &VaultId) -> Option<tokio::sync::OwnedSemaphorePermit> {
+        let limit = get_or_create(
+            async { self.admission.lock().unwrap() },
+            vault.clone(),
+            || Arc::new(tokio::sync::Semaphore::new(self.capacity)),
+        )
+        .await;
 
         limit.try_acquire_owned().ok()
     }