From 67c5574f377f295720e4756f1a8253dfbc6b8a1c Mon Sep 17 00:00:00 2001 From: Josh Thomas Date: Fri, 5 Sep 2025 18:46:42 -0500 Subject: [PATCH] simplify salsa db management with `Clone` + `Arc>` (#194) --- crates/djls-server/src/server.rs | 78 ++---- crates/djls-server/src/session.rs | 4 +- crates/djls-workspace/src/db.rs | 21 -- crates/djls-workspace/src/lib.rs | 1 - crates/djls-workspace/src/storage.rs | 343 ------------------------- crates/djls-workspace/src/workspace.rs | 76 ++---- 6 files changed, 46 insertions(+), 477 deletions(-) delete mode 100644 crates/djls-workspace/src/storage.rs diff --git a/crates/djls-server/src/server.rs b/crates/djls-server/src/server.rs index e5cc780..d221fcd 100644 --- a/crates/djls-server/src/server.rs +++ b/crates/djls-server/src/server.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use djls_workspace::paths; use djls_workspace::FileKind; -use tokio::sync::RwLock; +use tokio::sync::Mutex; use tower_lsp_server::jsonrpc::Result as LspResult; use tower_lsp_server::lsp_types; use tower_lsp_server::Client; @@ -19,7 +19,7 @@ const SERVER_VERSION: &str = "0.1.0"; pub struct DjangoLanguageServer { #[allow(dead_code)] // will be needed when diagnostics and other features are added client: Client, - session: Arc>>, + session: Arc>, queue: Queue, _log_guard: WorkerGuard, } @@ -29,7 +29,7 @@ impl DjangoLanguageServer { pub fn new(client: Client, log_guard: WorkerGuard) -> Self { Self { client, - session: Arc::new(RwLock::new(None)), + session: Arc::new(Mutex::new(Session::default())), queue: Queue::new(), _log_guard: log_guard, } @@ -38,34 +38,22 @@ impl DjangoLanguageServer { pub async fn with_session(&self, f: F) -> R where F: FnOnce(&Session) -> R, - R: Default, { - let session = self.session.read().await; - if let Some(s) = &*session { - f(s) - } else { - tracing::error!("Attempted to access session before initialization"); - R::default() - } + let session = self.session.lock().await; + f(&session) } pub async fn with_session_mut(&self, f: F) -> R where F: FnOnce(&mut Session) -> R, - R: Default, { - let mut session = self.session.write().await; - if let Some(s) = &mut *session { - f(s) - } else { - tracing::error!("Attempted to access session before initialization"); - R::default() - } + let mut session = self.session.lock().await; + f(&mut session) } pub async fn with_session_task(&self, f: F) where - F: FnOnce(Arc>>) -> Fut + Send + 'static, + F: FnOnce(Arc>) -> Fut + Send + 'static, Fut: Future> + Send + 'static, { let session_arc = Arc::clone(&self.session); @@ -89,8 +77,8 @@ impl LanguageServer for DjangoLanguageServer { let encoding = session.position_encoding(); { - let mut session_lock = self.session.write().await; - *session_lock = Some(session); + let mut session_lock = self.session.lock().await; + *session_lock = session; } Ok(lsp_types::InitializeResult { @@ -137,19 +125,16 @@ impl LanguageServer for DjangoLanguageServer { self.with_session_task(move |session_arc| async move { let project_path_and_venv = { - let session_lock = session_arc.read().await; - match &*session_lock { - Some(session) => session.project().map(|p| { - ( - p.path().display().to_string(), - session - .settings() - .venv_path() - .map(std::string::ToString::to_string), - ) - }), - None => None, - } + let session_lock = session_arc.lock().await; + session_lock.project().map(|p| { + ( + p.path().display().to_string(), + session_lock + .settings() + .venv_path() + .map(std::string::ToString::to_string), + ) + }) }; if let Some((path_display, venv_path)) = project_path_and_venv { @@ -163,17 +148,12 @@ impl LanguageServer for DjangoLanguageServer { } let init_result = { - let mut session_lock = session_arc.write().await; - match &mut *session_lock { - Some(session) => { - if let Some(project) = session.project_mut().as_mut() { - project.initialize(venv_path.as_deref()) - } else { - // Project was removed between read and write locks - Ok(()) - } - } - None => Ok(()), + let mut session_lock = session_arc.lock().await; + if let Some(project) = session_lock.project_mut().as_mut() { + project.initialize(venv_path.as_deref()) + } else { + // Project was removed between read and write locks + Ok(()) } }; @@ -189,10 +169,8 @@ impl LanguageServer for DjangoLanguageServer { ); // Clear project on error - let mut session_lock = session_arc.write().await; - if let Some(session) = &mut *session_lock { - *session.project_mut() = None; - } + let mut session_lock = session_arc.lock().await; + *session_lock.project_mut() = None; } } } else { diff --git a/crates/djls-server/src/session.rs b/crates/djls-server/src/session.rs index c651bd9..ff523b1 100644 --- a/crates/djls-server/src/session.rs +++ b/crates/djls-server/src/session.rs @@ -20,8 +20,8 @@ use url::Url; /// - Workspace operations (delegated to the Workspace facade) /// /// All document lifecycle and database operations are delegated to the -/// encapsulated Workspace, which provides thread-safe Salsa database -/// management with proper mutation safety through `StorageHandleGuard`. +/// encapsulated Workspace, which provides Salsa database management +/// using the built-in Clone pattern for thread safety. pub struct Session { /// The Django project configuration project: Option, diff --git a/crates/djls-workspace/src/db.rs b/crates/djls-workspace/src/db.rs index e98fdfc..7fab4d0 100644 --- a/crates/djls-workspace/src/db.rs +++ b/crates/djls-workspace/src/db.rs @@ -111,20 +111,6 @@ impl Database { } } - pub fn from_storage( - storage: salsa::Storage, - file_system: Arc, - files: Arc>, - ) -> Self { - Self { - storage, - fs: file_system, - files, - #[cfg(test)] - logs: Arc::new(Mutex::new(None)), - } - } - /// Read file content through the file system. pub fn read_file_content(&self, path: &Path) -> std::io::Result { self.fs.read_to_string(path) @@ -200,13 +186,6 @@ impl Database { new_rev ); } - - /// Get a reference to the storage for handle extraction. - /// - /// This is used by `Session` to extract the [`StorageHandle`](salsa::StorageHandle) after mutations. - pub fn storage(&self) -> &salsa::Storage { - &self.storage - } } #[salsa::db] diff --git a/crates/djls-workspace/src/lib.rs b/crates/djls-workspace/src/lib.rs index 532f88c..7eb60b4 100644 --- a/crates/djls-workspace/src/lib.rs +++ b/crates/djls-workspace/src/lib.rs @@ -19,7 +19,6 @@ pub mod encoding; mod fs; mod language; pub mod paths; -mod storage; mod workspace; use std::path::Path; diff --git a/crates/djls-workspace/src/storage.rs b/crates/djls-workspace/src/storage.rs deleted file mode 100644 index 28a25ff..0000000 --- a/crates/djls-workspace/src/storage.rs +++ /dev/null @@ -1,343 +0,0 @@ -use salsa::StorageHandle; - -use crate::db::Database; - -/// Safe wrapper for [`StorageHandle`](salsa::StorageHandle) that prevents misuse through type safety. -/// -/// This enum ensures that database handles can only be in one of two valid states, -/// making invalid states unrepresentable and eliminating the need for placeholder -/// handles during mutations. -/// -/// ## Panic Behavior -/// -/// Methods in this type may panic when the state machine invariants are violated. -/// These panics represent **programming bugs**, not runtime errors that should be -/// handled. They indicate violations of the internal API contract, similar to how -/// `RefCell::borrow_mut()` panics on double borrows. The panics ensure that bugs -/// are caught during development rather than causing silent data corruption. -pub enum SafeStorageHandle { - /// Handle is available for use - Available(StorageHandle), - /// Handle has been taken for mutation - no handle available - TakenForMutation, -} - -impl SafeStorageHandle { - /// Create a new `SafeStorageHandle` in the `Available` state - pub fn new(handle: StorageHandle) -> Self { - Self::Available(handle) - } - - /// Take the handle for mutation, leaving the enum in `TakenForMutation` state. - /// - /// ## Panics - /// - /// Panics if the handle has already been taken for mutation. - pub fn take_for_mutation(&mut self) -> StorageHandle { - match std::mem::replace(self, Self::TakenForMutation) { - Self::Available(handle) => handle, - Self::TakenForMutation => panic!( - "Database handle already taken for mutation. This indicates a programming error - \ - ensure you're not calling multiple mutation operations concurrently or forgetting \ - to restore the handle after a previous mutation." - ), - } - } - - /// Restore the handle after mutation, returning it to `Available` state. - /// - /// ## Panics - /// - /// Panics if the handle is not currently taken for mutation. - pub fn restore_from_mutation(&mut self, handle: StorageHandle) { - match self { - Self::TakenForMutation => { - *self = Self::Available(handle); - } - Self::Available(_) => panic!( - "Cannot restore database handle - handle is not currently taken for mutation. \ - This indicates a programming error in the StorageHandleGuard implementation." - ), - } - } - - /// Get a clone of the handle for read-only operations. - /// - /// ## Panics - /// - /// Panics if the handle is currently taken for mutation. - pub fn clone_for_read(&self) -> StorageHandle { - match self { - Self::Available(handle) => handle.clone(), - Self::TakenForMutation => panic!( - "Cannot access database handle for read - handle is currently taken for mutation. \ - Wait for the current mutation operation to complete." - ), - } - } - - /// Take the handle for mutation with automatic restoration via guard. - /// This ensures the handle is always restored even if the operation panics. - pub fn take_guarded(&mut self) -> StorageHandleGuard { - StorageHandleGuard::new(self) - } -} - -/// State of the [`StorageHandleGuard`] during its lifetime. -/// -/// See [`StorageHandleGuard`] for usage and state machine details. -enum GuardState { - /// Guard holds the handle, ready to be consumed - Active { handle: StorageHandle }, - /// Handle consumed, awaiting restoration - Consumed, - /// Handle restored to [`SafeStorageHandle`] - Restored, -} - -/// RAII guard for safe [`StorageHandle`](salsa::StorageHandle) management during mutations. -/// -/// This guard ensures that database handles are automatically restored even if -/// panics occur during mutation operations. It prevents double-takes and -/// provides clear error messages for misuse. -/// -/// ## State Machine -/// -/// The guard follows these valid state transitions: -/// - `Active` → `Consumed` (via `handle()` method) -/// - `Consumed` → `Restored` (via `restore()` method) -/// -/// ## Invalid Transitions -/// -/// Invalid operations will panic with specific error messages: -/// - `handle()` on `Consumed` state: "[`StorageHandle`](salsa::StorageHandle) already consumed" -/// - `handle()` on `Restored` state: "Cannot consume handle - guard has already been restored" -/// - `restore()` on `Active` state: "Cannot restore handle - it hasn't been consumed yet" -/// - `restore()` on `Restored` state: "Handle has already been restored" -/// -/// ## Drop Behavior -/// -/// The guard will panic on drop unless it's in the `Restored` state: -/// - Drop in `Active` state: "`StorageHandleGuard` dropped without using the handle" -/// - Drop in `Consumed` state: "`StorageHandleGuard` dropped without restoring handle" -/// - Drop in `Restored` state: No panic - proper cleanup completed -/// -/// ## Usage Example -/// -/// ```rust,ignore -/// let mut guard = StorageHandleGuard::new(&mut safe_handle); -/// let handle = guard.handle(); // Active → Consumed -/// // ... perform mutations with handle ... -/// guard.restore(updated_handle); // Consumed → Restored -/// // Guard drops cleanly in Restored state -/// ``` -#[must_use = "StorageHandleGuard must be used - dropping it immediately defeats the purpose"] -pub struct StorageHandleGuard<'a> { - /// Reference to the workspace's `SafeStorageHandle` for restoration - safe_handle: &'a mut SafeStorageHandle, - /// Current state of the guard and handle - state: GuardState, -} - -impl<'a> StorageHandleGuard<'a> { - /// Create a new guard by taking the handle from the `SafeStorageHandle`. - pub fn new(safe_handle: &'a mut SafeStorageHandle) -> Self { - let handle = safe_handle.take_for_mutation(); - Self { - safe_handle, - state: GuardState::Active { handle }, - } - } - - /// Get the [`StorageHandle`](salsa::StorageHandle) for mutation operations. - /// - /// ## Panics - /// - /// Panics if the handle has already been consumed or restored. - pub fn handle(&mut self) -> StorageHandle { - match std::mem::replace(&mut self.state, GuardState::Consumed) { - GuardState::Active { handle } => handle, - GuardState::Consumed => panic!( - "StorageHandle already consumed from guard. Each guard can only provide \ - the handle once - this prevents accidental multiple uses." - ), - GuardState::Restored => panic!( - "Cannot consume handle - guard has already been restored. Once restored, \ - the guard cannot provide the handle again." - ), - } - } - - /// Restore the handle manually before the guard drops. - /// - /// This is useful when you want to restore the handle and continue using - /// the workspace in the same scope. - /// - /// ## Panics - /// - /// Panics if the handle hasn't been consumed yet, or if already restored. - pub fn restore(mut self, handle: StorageHandle) { - match self.state { - GuardState::Consumed => { - self.safe_handle.restore_from_mutation(handle); - self.state = GuardState::Restored; - } - GuardState::Active { .. } => panic!( - "Cannot restore handle - it hasn't been consumed yet. Call guard.handle() \ - first to get the handle, then restore the updated handle after mutations." - ), - GuardState::Restored => { - panic!("Handle has already been restored. Each guard can only restore once.") - } - } - } -} - -impl Drop for StorageHandleGuard<'_> { - fn drop(&mut self) { - // Provide specific error messages based on the exact state - // Avoid double-panic during unwinding - if !std::thread::panicking() { - match &self.state { - GuardState::Active { .. } => { - panic!( - "StorageHandleGuard dropped without using the handle. Either call \ - guard.handle() to consume the handle for mutations, or ensure the \ - guard is properly used in your mutation workflow." - ); - } - GuardState::Consumed => { - panic!( - "StorageHandleGuard dropped without restoring handle. You must call \ - guard.restore(updated_handle) to properly restore the database handle \ - after mutation operations complete." - ); - } - GuardState::Restored => { - // All good - proper cleanup completed - } - } - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use dashmap::DashMap; - - use super::*; - use crate::buffers::Buffers; - use crate::fs::OsFileSystem; - use crate::fs::WorkspaceFileSystem; - - fn create_test_handle() -> StorageHandle { - Database::new( - Arc::new(WorkspaceFileSystem::new( - Buffers::new(), - Arc::new(OsFileSystem), - )), - Arc::new(DashMap::new()), - ) - .storage() - .clone() - .into_zalsa_handle() - } - - #[test] - fn test_handle_lifecycle() { - // Test the happy path: take handle, use it, restore it - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - - let handle = safe_handle.take_for_mutation(); - - // Simulate using the handle to create a database - let storage = handle.into_storage(); - let db = Database::from_storage( - storage, - Arc::new(WorkspaceFileSystem::new( - Buffers::new(), - Arc::new(OsFileSystem), - )), - Arc::new(DashMap::new()), - ); - - // Get new handle after simulated mutation - let new_handle = db.storage().clone().into_zalsa_handle(); - - safe_handle.restore_from_mutation(new_handle); - - // Should be able to take it again - let _handle2 = safe_handle.take_for_mutation(); - } - - #[test] - fn test_guard_auto_restore_on_drop() { - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - - { - let mut guard = safe_handle.take_guarded(); - let handle = guard.handle(); - - // Simulate mutation - let storage = handle.into_storage(); - let db = Database::from_storage( - storage, - Arc::new(WorkspaceFileSystem::new( - Buffers::new(), - Arc::new(OsFileSystem), - )), - Arc::new(DashMap::new()), - ); - let new_handle = db.storage().clone().into_zalsa_handle(); - - guard.restore(new_handle); - } // Guard drops here in Restored state - should be clean - - // Should be able to use handle again after guard drops - let _handle = safe_handle.clone_for_read(); - } - - #[test] - #[should_panic(expected = "Database handle already taken for mutation")] - fn test_panic_on_double_mutation() { - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - - let _handle1 = safe_handle.take_for_mutation(); - // Can't take handle twice, should panic - let _handle2 = safe_handle.take_for_mutation(); - } - - #[test] - #[should_panic(expected = "Cannot access database handle for read")] - fn test_panic_on_read_during_mutation() { - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - - let _handle = safe_handle.take_for_mutation(); - // Can't read while mutating, should panic - let _read = safe_handle.clone_for_read(); - } - - #[test] - #[should_panic(expected = "Cannot restore handle - it hasn't been consumed yet")] - fn test_guard_enforces_consume_before_restore() { - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - let guard = safe_handle.take_guarded(); - - let dummy_handle = create_test_handle(); - // Try to restore without consuming, should panic - guard.restore(dummy_handle); - } - - #[test] - #[should_panic(expected = "StorageHandleGuard dropped without restoring handle")] - fn test_guard_panics_if_dropped_without_restore() { - let mut safe_handle = SafeStorageHandle::new(create_test_handle()); - - { - let mut guard = safe_handle.take_guarded(); - let _handle = guard.handle(); - } // Explicitly drop guard without restore, should panic - } -} diff --git a/crates/djls-workspace/src/workspace.rs b/crates/djls-workspace/src/workspace.rs index dfb9a2d..e17d97e 100644 --- a/crates/djls-workspace/src/workspace.rs +++ b/crates/djls-workspace/src/workspace.rs @@ -5,7 +5,6 @@ //! This provides a clean API boundary between server and workspace layers. use std::path::Path; -use std::path::PathBuf; use std::sync::Arc; use dashmap::DashMap; @@ -14,27 +13,20 @@ use url::Url; use crate::buffers::Buffers; use crate::db::Database; -use crate::db::SourceFile; use crate::document::TextDocument; use crate::fs::OsFileSystem; use crate::fs::WorkspaceFileSystem; use crate::paths::url_to_path; -use crate::storage::SafeStorageHandle; /// Workspace facade that encapsulates all workspace components. /// /// This struct provides a unified interface for managing workspace state, -/// including in-memory buffers, file system abstraction, file tracking, -/// and the Salsa database handle. +/// including in-memory buffers, file system abstraction, and the Salsa database. pub struct Workspace { /// Thread-safe shared buffer storage for open documents buffers: Buffers, - /// File system abstraction with buffer interception - file_system: Arc, - /// Shared file tracking across all Database instances - files: Arc>, - /// Thread-safe Salsa database handle for incremental computation with safe mutation management - db_handle: SafeStorageHandle, + /// Salsa database for incremental computation + db: Database, } impl Workspace { @@ -47,49 +39,25 @@ impl Workspace { buffers.clone(), Arc::new(OsFileSystem), )); - let handle = Database::new(file_system.clone(), files.clone()) - .storage() - .clone() - .into_zalsa_handle(); + let db = Database::new(file_system, files); - Self { - buffers, - file_system, - files, - db_handle: SafeStorageHandle::new(handle), - } + Self { buffers, db } } /// Execute a read-only operation with access to the database. - /// - /// Creates a temporary Database instance from the handle for the closure. - /// This is safe for concurrent read operations. pub fn with_db(&self, f: F) -> R where F: FnOnce(&Database) -> R, { - let handle = self.db_handle.clone_for_read(); - let storage = handle.into_storage(); - let db = Database::from_storage(storage, self.file_system.clone(), self.files.clone()); - f(&db) + f(&self.db) } /// Execute a mutable operation with exclusive access to the database. - /// - /// Uses the `StorageHandleGuard` pattern to ensure the handle is safely restored - /// even if the operation panics. This eliminates the need for placeholder handles. pub fn with_db_mut(&mut self, f: F) -> R where F: FnOnce(&mut Database) -> R, { - let mut guard = self.db_handle.take_guarded(); - let handle = guard.handle(); - let storage = handle.into_storage(); - let mut db = Database::from_storage(storage, self.file_system.clone(), self.files.clone()); - let result = f(&mut db); - let new_handle = db.storage().clone().into_zalsa_handle(); - guard.restore(new_handle); - result + f(&mut self.db) } /// Open a document in the workspace. @@ -181,11 +149,6 @@ impl Workspace { } }); } - - #[must_use] - pub fn db_handle(&self) -> &SafeStorageHandle { - &self.db_handle - } } impl Default for Workspace { @@ -197,7 +160,6 @@ impl Default for Workspace { #[cfg(test)] mod tests { use std::path::PathBuf; - use std::sync::Arc; use tempfile::tempdir; @@ -236,23 +198,17 @@ mod tests { } #[test] - fn test_concurrent_reads() { - // Multiple with_db calls can run simultaneously - let workspace = Arc::new(Workspace::new()); + fn test_multiple_reads() { + // Multiple with_db calls work correctly with Clone pattern + let workspace = Workspace::new(); - let w1 = workspace.clone(); - let w2 = workspace.clone(); + // Multiple reads work fine + let result1 = workspace.with_db(|db| db.has_file(&PathBuf::from("file1.py"))); + let result2 = workspace.with_db(|db| db.has_file(&PathBuf::from("file2.py"))); - // Spawn concurrent reads - let handle1 = - std::thread::spawn(move || w1.with_db(|db| db.has_file(&PathBuf::from("file1.py")))); - - let handle2 = - std::thread::spawn(move || w2.with_db(|db| db.has_file(&PathBuf::from("file2.py")))); - - // Both should complete without issues - assert!(!handle1.join().unwrap()); - assert!(!handle2.join().unwrap()); + // Both should return false since no files were created + assert!(!result1); + assert!(!result2); } #[test]