simplify salsa db management with Clone + Arc<Mutex<Session>> (#194)
Some checks are pending
test / Python , Django () (push) Blocked by required conditions
test / tests (push) Blocked by required conditions
zizmor 🌈 / zizmor latest via PyPI (push) Waiting to run
lint / pre-commit (push) Waiting to run
lint / rustfmt (push) Waiting to run
lint / clippy (push) Waiting to run
lint / cargo-check (push) Waiting to run
release / build (push) Waiting to run
release / test (push) Waiting to run
release / release (push) Blocked by required conditions
test / generate-matrix (push) Waiting to run

This commit is contained in:
Josh Thomas 2025-09-05 18:46:42 -05:00 committed by GitHub
parent b90862d72a
commit 67c5574f37
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 46 additions and 477 deletions

View file

@ -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;
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<F, R>(&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<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;
*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 {

View file

@ -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<DjangoProject>,

View file

@ -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]

View file

@ -19,7 +19,6 @@ pub mod encoding;
mod fs;
mod language;
pub mod paths;
mod storage;
mod workspace;
use std::path::Path;

View file

@ -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
}
}

View file

@ -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<WorkspaceFileSystem>,
/// 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,
/// 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<F, R>(&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<F, R>(&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]