simplify salsa db management with Clone + Arc<Mutex<Session>> (#194)

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<RwLock<Option<Session>>>,
+    session: Arc<Mutex<Session>>,
     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<F, R>(&self, f: F) -> R
     where
         F: FnOnce(&Session) -> R,
-        R: Default,
     {
-        let session = self.session.read().await;
+        let session = self.session.lock().await;
-        if let Some(s) = &*session {
+        f(&session)
-            f(s)
-        } else {
-            tracing::error!("Attempted to access session before initialization");
-            R::default()
-        }
     }
 
     pub async fn with_session_mut<F, R>(&self, f: F) -> R
     where
         F: FnOnce(&mut Session) -> R,
-        R: Default,
     {
-        let mut session = self.session.write().await;
+        let mut session = self.session.lock().await;
-        if let Some(s) = &mut *session {
+        f(&mut session)
-            f(s)
-        } else {
-            tracing::error!("Attempted to access session before initialization");
-            R::default()
-        }
     }
 
     pub async fn with_session_task<F, Fut>(&self, f: F)
     where
-        F: FnOnce(Arc<RwLock<Option<Session>>>) -> Fut + Send + 'static,
+        F: FnOnce(Arc<Mutex<Session>>) -> Fut + Send + 'static,
         Fut: Future<Output = anyhow::Result<()>> + 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;
+            let mut session_lock = self.session.lock().await;
-            *session_lock = Some(session);
+            *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;
+                let session_lock = session_arc.lock().await;
-                match &*session_lock {
+                session_lock.project().map(|p| {
-                    Some(session) => session.project().map(|p| {
+                    (
-                        (
+                        p.path().display().to_string(),
-                            p.path().display().to_string(),
+                        session_lock
-                            session
+                            .settings()
-                                .settings()
+                            .venv_path()
-                                .venv_path()
+                            .map(std::string::ToString::to_string),
-                                .map(std::string::ToString::to_string),
+                    )
-                        )
+                })
-                    }),
-                    None => None,
-                }
             };
 
             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;
+                    let mut session_lock = session_arc.lock().await;
-                    match &mut *session_lock {
+                    if let Some(project) = session_lock.project_mut().as_mut() {
-                        Some(session) => {
+                        project.initialize(venv_path.as_deref())
-                            if let Some(project) = session.project_mut().as_mut() {
+                    } else {
-                                project.initialize(venv_path.as_deref())
+                        // Project was removed between read and write locks
-                            } else {
+                        Ok(())
-                                // Project was removed between read and write locks
-                                Ok(())
-                            }
-                        }
-                        None => Ok(()),
                     }
                 };
 
@@ -189,10 +169,8 @@ impl LanguageServer for DjangoLanguageServer {
                         );
 
                         // Clear project on error
-                        let mut session_lock = session_arc.write().await;
+                        let mut session_lock = session_arc.lock().await;
-                        if let Some(session) = &mut *session_lock {
+                        *session_lock.project_mut() = None;
-                            *session.project_mut() = None;
-                        }
                     }
                 }
             } else {

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
+/// encapsulated Workspace, which provides Salsa database management
-/// management with proper mutation safety through `StorageHandleGuard`.
+/// using the built-in Clone pattern for thread safety.
 pub struct Session {
     /// The Django project configuration
     project: Option<DjangoProject>,

crates/djls-workspace/src/db.rs

@@ -111,20 +111,6 @@ impl Database {
         }
     }
 
-    pub fn from_storage(
-        storage: salsa::Storage<Self>,
-        file_system: Arc<dyn FileSystem>,
-        files: Arc<DashMap<PathBuf, SourceFile>>,
-    ) -> 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<String> {
         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> {
-        &self.storage
-    }
 }
 
 #[salsa::db]

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;

crates/djls-workspace/src/storage.rs

@@ -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<Database>),
-    /// 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<Database>) -> 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<Database> {
-        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<Database>) {
-        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<Database> {
-        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<Database> },
-    /// 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<Database> {
-        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<Database>) {
-        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> {
-        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
-    }
-}

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,
+/// including in-memory buffers, file system abstraction, and the Salsa database.
-/// and the Salsa database handle.
 pub struct Workspace {
     /// Thread-safe shared buffer storage for open documents
     buffers: Buffers,
-    /// File system abstraction with buffer interception
+    /// Salsa database for incremental computation
-    file_system: Arc<WorkspaceFileSystem>,
+    db: Database,
-    /// Shared file tracking across all Database instances
-    files: Arc<DashMap<PathBuf, SourceFile>>,
-    /// Thread-safe Salsa database handle for incremental computation with safe mutation management
-    db_handle: SafeStorageHandle,
 }
 
 impl Workspace {
@@ -47,49 +39,25 @@ impl Workspace {
             buffers.clone(),
             Arc::new(OsFileSystem),
         ));
-        let handle = Database::new(file_system.clone(), files.clone())
+        let db = Database::new(file_system, files);
-            .storage()
-            .clone()
-            .into_zalsa_handle();
 
-        Self {
+        Self { buffers, db }
-            buffers,
-            file_system,
-            files,
-            db_handle: SafeStorageHandle::new(handle),
-        }
     }
 
     /// 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<F, R>(&self, f: F) -> R
     where
         F: FnOnce(&Database) -> R,
     {
-        let handle = self.db_handle.clone_for_read();
+        f(&self.db)
-        let storage = handle.into_storage();
-        let db = Database::from_storage(storage, self.file_system.clone(), self.files.clone());
-        f(&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<F, R>(&mut self, f: F) -> R
     where
         F: FnOnce(&mut Database) -> R,
     {
-        let mut guard = self.db_handle.take_guarded();
+        f(&mut self.db)
-        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
     }
 
     /// 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() {
+    fn test_multiple_reads() {
-        // Multiple with_db calls can run simultaneously
+        // Multiple with_db calls work correctly with Clone pattern
-        let workspace = Arc::new(Workspace::new());
+        let workspace = Workspace::new();
 
-        let w1 = workspace.clone();
+        // Multiple reads work fine
-        let w2 = workspace.clone();
+        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
+        // Both should return false since no files were created
-        let handle1 =
+        assert!(!result1);
-            std::thread::spawn(move || w1.with_db(|db| db.has_file(&PathBuf::from("file1.py"))));
+        assert!(!result2);
-
-        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());
     }
 
     #[test]