diff --git a/crates/workspace-server/src/feature_storage.rs b/crates/workspace-server/src/feature_storage.rs index 13fb08541..c8ad19ba9 100644 --- a/crates/workspace-server/src/feature_storage.rs +++ b/crates/workspace-server/src/feature_storage.rs @@ -8,6 +8,7 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, Weak}; use std::time::Duration; @@ -221,6 +222,9 @@ impl FeatureStorage { .values() .filter_map(Weak::upgrade) .collect::>(); + for connection in &connections { + connection.fence(); + } state.connections.clear(); connections }; @@ -270,6 +274,7 @@ impl FeatureStorage { let connection = Arc::new(FeatureDatabaseInner { workspace_id: workspace_id.to_string(), feature_id: registration.feature_id.clone(), + accepting_operations: AtomicBool::new(true), connection: Mutex::new(Some(connection)), }); state.connections.insert(key, Arc::downgrade(&connection)); @@ -381,6 +386,9 @@ impl FeatureStorage { .filter(|((candidate, _), _)| candidate == workspace_id) .filter_map(|(_, connection)| connection.upgrade()) .collect::>(); + for connection in &connections { + connection.fence(); + } state .connections .retain(|(candidate, _), _| candidate != workspace_id); @@ -513,6 +521,7 @@ pub struct FeatureDatabase { struct FeatureDatabaseInner { workspace_id: String, feature_id: String, + accepting_operations: AtomicBool, connection: Mutex>, } @@ -531,9 +540,11 @@ impl FeatureDatabase { &self, operation: impl FnOnce(&Connection) -> Result, ) -> Result { + self.inner.ensure_accepting()?; let connection = self.inner.connection.lock().map_err(|_| { FeatureStorageError::Operation("Feature database lock was poisoned".to_string()) })?; + self.inner.ensure_accepting()?; let connection = connection .as_ref() .ok_or(FeatureStorageError::DatabaseClosed)?; @@ -546,9 +557,11 @@ impl FeatureDatabase { &self, operation: impl FnOnce(&Transaction<'_>) -> Result, ) -> Result { + self.inner.ensure_accepting()?; let mut connection = self.inner.connection.lock().map_err(|_| { FeatureStorageError::Operation("Feature database lock was poisoned".to_string()) })?; + self.inner.ensure_accepting()?; let connection = connection .as_mut() .ok_or(FeatureStorageError::DatabaseClosed)?; @@ -572,7 +585,20 @@ impl FeatureDatabase { } impl FeatureDatabaseInner { + fn ensure_accepting(&self) -> Result<()> { + if self.accepting_operations.load(Ordering::Acquire) { + Ok(()) + } else { + Err(FeatureStorageError::DatabaseClosed) + } + } + + fn fence(&self) { + self.accepting_operations.store(false, Ordering::Release); + } + fn close(&self) -> Result<()> { + self.fence(); let mut connection = self.connection.lock().map_err(|_| { FeatureStorageError::Operation("Feature database lock was poisoned".to_string()) })?; @@ -1264,8 +1290,14 @@ mod tests { let manager = manager.clone(); thread::spawn(move || manager.shutdown()) }; - thread::sleep(Duration::from_millis(20)); + while database.inner.accepting_operations.load(Ordering::Acquire) { + thread::yield_now(); + } assert!(!shutdown.is_finished()); + assert!(matches!( + values(&database).unwrap_err(), + FeatureStorageError::DatabaseClosed + )); release.wait(); operation.join().unwrap().unwrap(); shutdown.join().unwrap().unwrap();