From 81b4a08a59febe4c8c12d8b6f5e7a9d56fbe251e Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Tue, 22 Sep 2026 20:04:19 +0100 Subject: [PATCH] Fix pool opening race and closing stutter --- sync-server/src/app_state.rs | 1 + sync-server/src/app_state/database.rs | 72 +++++++++++++------ sync-server/src/app_state/weak_slots.rs | 26 +++++++ .../src/app_state/websocket/broadcasts.rs | 28 ++++---- 4 files changed, 90 insertions(+), 37 deletions(-) create mode 100644 sync-server/src/app_state/weak_slots.rs diff --git a/sync-server/src/app_state.rs b/sync-server/src/app_state.rs index 910564f9..5293ee1c 100644 --- a/sync-server/src/app_state.rs +++ b/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; diff --git a/sync-server/src/app_state/database.rs b/sync-server/src/app_state/database.rs index 29118240..e5c492cb 100644 --- a/sync-server/src/app_state/database.rs +++ b/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> { let vault = normalize_string(vault); validate_vault_id(&vault)?; - let mut pools = self.connection_pools.lock().await; - if !pools.contains_key(&vault) { - let file_name = Self::database_path(&self.config, &vault); - let pool = Self::open_database(&self.config, &file_name).await?; - pools.insert( - vault.clone(), - PoolWithTimestamp { - pool, - last_accessed: Instant::now(), - }, - ); + if let Some(pool) = self.existing_pool(&vault).await { + return Ok(pool); } - let pool_with_timestamp = pools - .get_mut(&vault) - .expect("Pool was just inserted or already exists"); + let opening = get_or_create(self.opening_pools.lock(), vault.clone(), || { + Arc::new(Mutex::new(())) + }) + .await; - // Update last accessed time - pool_with_timestamp.last_accessed = Instant::now(); + let _opening = opening.lock().await; - Ok(pool_with_timestamp.pool.clone()) + 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?; + + self.connection_pools.lock().await.insert( + vault, + PoolWithTimestamp { + pool: pool.clone(), + last_accessed: Instant::now(), + }, + ); + Ok(pool) + } + + async fn existing_pool(&self, vault: &VaultId) -> Option> { + 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; } } diff --git a/sync-server/src/app_state/weak_slots.rs b/sync-server/src/app_state/weak_slots.rs new file mode 100644 index 00000000..bda3db1b --- /dev/null +++ b/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` for `key` in a weak map behind a mutex +pub async fn get_or_create(lock: L, key: K, create: impl FnOnce() -> Arc) -> Arc +where + K: Clone + Eq + Hash, + G: DerefMut>>, + L: Future, +{ + 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 + } +} diff --git a/sync-server/src/app_state/websocket/broadcasts.rs b/sync-server/src/app_state/websocket/broadcasts.rs index 4980ca10..09b27405 100644 --- a/sync-server/src/app_state/websocket/broadcasts.rs +++ b/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 { - 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 { + 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() }