Skip to content

Make it impossible for background threads to mutate Salsa inputs #1314

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lionel- wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Make it impossible for background threads to mutate Salsa inputs #1314

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions crates/ark/src/lsp/completions/completion_context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,18 @@ use crate::lsp::completions::function_context::FunctionContext;
use crate::lsp::completions::sources::composite::pipe::find_pipe_root;
use crate::lsp::completions::sources::composite::pipe::PipeRoot;
use crate::lsp::document_context::DocumentContext;
use crate::lsp::state::WorldState;
use crate::lsp::state::WorldSnapshot;
use crate::treesitter::node_find_containing_call;
pub(crate) struct CompletionContext<'a> {
pub(crate) document_context: &'a DocumentContext<'a>,
pub(crate) state: &'a WorldState,
pub(crate) state: &'a WorldSnapshot,
pipe_root_cell: OnceCell<Option<PipeRoot>>,
containing_call_cell: OnceCell<Option<Node<'a>>>,
function_context_cell: OnceCell<anyhow::Result<FunctionContext>>,
}

impl<'a> CompletionContext<'a> {
pub fn new(document_context: &'a DocumentContext, state: &'a WorldState) -> Self {
pub fn new(document_context: &'a DocumentContext, state: &'a WorldSnapshot) -> Self {
Self {
document_context,
state,
Expand Down
4 changes: 2 additions & 2 deletions crates/ark/src/lsp/completions/provide.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,15 @@ use crate::lsp::completions::completion_context::CompletionContext;
use crate::lsp::completions::sources::composite;
use crate::lsp::completions::sources::unique;
use crate::lsp::document_context::DocumentContext;
use crate::lsp::state::WorldState;
use crate::lsp::state::WorldSnapshot;
use crate::lsp::traits::node::NodeExt;
use crate::treesitter::NodeTypeExt;

// Entry point for completions.
// Must be within an `r_task()`.
pub(crate) fn provide_completions(
document_context: &DocumentContext,
state: &WorldState,
state: &WorldSnapshot,
) -> anyhow::Result<Vec<CompletionItem>> {
log::info!(
"provide_completions() - Completion node text: '{node_text}', Node type: '{node_type:?}'",
Expand Down
4 changes: 2 additions & 2 deletions crates/ark/src/lsp/completions/sources/composite.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,7 +252,7 @@ mod tests {
let (text, point) = point_from_cursor("@");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);

assert!(context.document_context.node.is_program());
Expand All @@ -270,7 +270,7 @@ mod tests {
let (text, point) = point_from_cursor(code);
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);

assert!(context.document_context.node.is_program());
Expand Down
24 changes: 12 additions & 12 deletions crates/ark/src/lsp/completions/sources/composite/call.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ fn completions_from_call(
},
};

completions_from_arguments(&context.state.db, document_context, callee, object)
completions_from_arguments(context.state.db.read(), document_context, callee, object)
}

fn get_first_argument(context: &DocumentContext, node: &Node) -> anyhow::Result<Option<RObject>> {
Expand Down Expand Up @@ -294,7 +294,7 @@ mod tests {
let (text, point) = point_from_cursor("match(tab@)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -307,7 +307,7 @@ mod tests {
let (text, point) = point_from_cursor("match(1, tab@)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -327,7 +327,7 @@ mod tests {
let (text, point) = point_from_cursor("not_a_known_function(@)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap();
assert!(completions.is_none());
Expand All @@ -347,7 +347,7 @@ mod tests {
let (text, point) = point_from_cursor("my_fun(@)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -364,7 +364,7 @@ mod tests {
let (text, point) = point_from_cursor("my_fun@()");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap();
assert!(completions.is_none());
Expand All @@ -373,7 +373,7 @@ mod tests {
let (text, point) = point_from_cursor("my_fun()@");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap();
assert!(completions.is_none());
Expand All @@ -396,7 +396,7 @@ mod tests {
let (text, point) = point_from_cursor("my_fun(@)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();
assert_eq!(completions.len(), 0);
Expand All @@ -413,7 +413,7 @@ mod tests {
let (text, point) = point_from_cursor("match(\n @\n)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -425,7 +425,7 @@ mod tests {
let (text, point) = point_from_cursor("match(\n tab@\n)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -437,7 +437,7 @@ mod tests {
let (text, point) = point_from_cursor("match(\n 1,\n tab@\n)");
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap().unwrap();

Expand All @@ -454,7 +454,7 @@ mod tests {
let (text, point) = point_from_cursor(code_with_cursor);
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);
let completions = completions_from_call(&context).unwrap();

Expand Down
2 changes: 1 addition & 1 deletion crates/ark/src/lsp/completions/sources/composite/workspace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ fn completions_from_workspace(
let token = token.as_str();

// get entries from the index
indexer::map(&state.db, |file, symbol, entry| {
indexer::map(state.db.read(), |file, symbol, entry| {
if !symbol.fuzzy_matches(token) {
return;
}
Expand Down
8 changes: 4 additions & 4 deletions crates/ark/src/lsp/completions/sources/unique/custom.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,7 @@ mod tests {
// Helper functions for testing custom completions
fn assert_has_completion(code_with_cursor: &str, name: &str, expected_insert_text: &str) {
let (text, point) = point_from_cursor(code_with_cursor);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let context = CompletionContext::new(&document_context, &state);
Expand All @@ -235,7 +235,7 @@ mod tests {

fn assert_no_completions(code_with_cursor: &str) {
let (text, point) = point_from_cursor(code_with_cursor);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let context = CompletionContext::new(&document_context, &state);
Expand All @@ -254,7 +254,7 @@ mod tests {
};

let (text, point) = point_from_cursor("library(@)");
let state = WorldState::default();
let state = WorldState::default().snapshot();
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let context = CompletionContext::new(&document_context, &state);
Expand All @@ -268,7 +268,7 @@ mod tests {
assert_eq!(n_compls, n_packages);

let (text, point) = point_from_cursor("library(uti@)");
let state = WorldState::default();
let state = WorldState::default().snapshot();
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let context = CompletionContext::new(&document_context, &state);
Expand Down
2 changes: 1 addition & 1 deletion crates/ark/src/lsp/completions/sources/unique/namespace.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ mod tests {
let (text, point) = point_from_cursor(cursor_text);
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();
let context = CompletionContext::new(&document_context, &state);

completions_from_namespace(&context)
Expand Down
2 changes: 1 addition & 1 deletion crates/ark/src/lsp/completions/sources/unique/string.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,7 @@ mod tests {
assert_match!(res, Some(items) => { assert!(items.is_empty()) });

// Check for same result when consulting (potentially all) unique sources
let state = WorldState::default();
let state = WorldState::default().snapshot();
let completion_context = CompletionContext::new(&context, &state);
let res = unique::get_completions(&completion_context).unwrap();

Expand Down
2 changes: 1 addition & 1 deletion crates/ark/src/lsp/completions/tests/utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pub(crate) fn get_completions_at_cursor(cursor_text: &str) -> anyhow::Result<Vec
let (text, point) = point_from_cursor(cursor_text);
let doc = TestDocument::new(&text);
let document_context = doc.context(point);
let state = WorldState::default();
let state = WorldState::default().snapshot();

match provide_completions(&document_context, &state) {
Ok(completions) => Ok(completions),
Expand Down
31 changes: 31 additions & 0 deletions crates/ark/src/lsp/db.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,37 @@ pub(crate) trait ArkDb: oak_db::Db {}
#[salsa::db]
impl ArkDb for OakDatabase {}

/// Read-only handle to an `OakDatabase` that acts as a read-only snapshot.
///
/// An `Analysis` snapshot holds the DB clone in a private field and lends a
/// shared ref via `read()`. This prevents a read-only background workers
/// from reaching DB setters and deadock with the main loop. This is oak's
/// version of rust-analyzer's `Analysis` facade.
#[derive(Clone, Debug)]
pub(crate) struct Analysis {
db: OakDatabase,
}
Comment on lines +20 to +29

@DavisVaughan DavisVaughan Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a gut feeling this is too complex for us

What if we just had

pub(crate) struct WorldStateSnapshot {
    // Private field. We only give out `&dyn ArkDb` handles on threads to enforce that it is read-only.
    db: OakDatabase
}

impl WorldStateSnapshot {
    /// Read-only access to the `OakDatabase`. Important that we return as `&dyn ArkDb`
    /// to prevent the thread from simply `.clone()`ing to get its own copy that it can call
    /// setters with.
    pub(crate) db(&self) -> &dyn ArkDb {
        &self.db
    }
}

Your little test only cancellation_token() could live on WorldStateSnapshot, I think.

This setup makes a lot of sense to me! And avoids a lot of the layers that seem to make naming this thing hard.


impl Analysis {
pub(crate) fn new(db: OakDatabase) -> Self {
Self { db }
}

/// The database as a shared `&dyn ArkDb`. Only lends `&`, so a reader can
/// query but can't reach a setter through it.
pub(crate) fn read(&self) -> &dyn ArkDb {
&self.db
}

/// The database's salsa cancellation token. Read-side only: it observes and
/// arms cancellation, it doesn't mutate any input, so it's safe to expose
/// on the read-only facade. Only cancellation tests arm it by hand.
#[cfg(test)]
pub(crate) fn cancellation_token(&self) -> salsa::CancellationToken {
salsa::Database::cancellation_token(&self.db)
}
}

/// Extension trait that adds ark-side query methods to `oak_db::File`.
pub(crate) trait FileArkExt {
fn tree_sitter(self, db: &dyn ArkDb) -> &tree_sitter::Tree;
Expand Down
8 changes: 4 additions & 4 deletions crates/ark/src/lsp/diagnostics.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ use crate::lsp::diagnostics_syntax::syntax_diagnostics;
use crate::lsp::indexer;
use crate::lsp::inputs::source_root::SourceRoot;
use crate::lsp::open_file::lsp_range_from_tree_sitter_range;
use crate::lsp::state::WorldState;
use crate::lsp::state::WorldSnapshot;
use crate::lsp::traits::node::NodeExt;
use crate::treesitter::node_has_error_or_missing;
use crate::treesitter::BinaryOperatorType;
Expand Down Expand Up @@ -151,7 +151,7 @@ impl<'a> DiagnosticContext<'a> {

pub(crate) fn generate_diagnostics(
file: File,
state: WorldState,
state: WorldSnapshot,
testthat: bool,
) -> Vec<Diagnostic> {
let mut diagnostics = Vec::new();
Expand All @@ -160,7 +160,7 @@ pub(crate) fn generate_diagnostics(
return diagnostics;
}

let db = &state.db;
let db = state.db.read();

// Check that diagnostics are not disabled in top-level declarations for
// this document
Expand Down Expand Up @@ -1193,7 +1193,7 @@ mod tests {
Some(code.to_string()),
None,
);
super::generate_diagnostics(file, state, false)
super::generate_diagnostics(file, state.snapshot(), false)
}

fn current_state() -> WorldState {
Expand Down
13 changes: 9 additions & 4 deletions crates/ark/src/lsp/handlers.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,10 +233,15 @@ pub(crate) fn handle_completion(
);
lsp::log_info!("Completion context: {:#?}", context);

// TODO(oak/completions): Clone so the closure captures by value. `r_task()`
// sends the closure across threads, and `&WorldState` isn't `Send` because
// `OakDatabase`'s salsa storage keeps thread-local query state.
let state = state.clone();
// Snapshot so the closure captures by value. `r_task()` sends the closure
// across threads, and `&WorldState` isn't `Send` because `OakDatabase`'s
// salsa storage keeps thread-local query state. `snapshot()` hands the
// reader a `WorldSnapshot`, so the background thread can query oak but
// can't call a setter.
// TODO(oak/completions): We don't really need a snapshot here since
// completions are serviced from the main loop, it's only needed for the
// `r_task()`.
let state = state.snapshot();
let completions = r_task(move || provide_completions(&context, &state))?;

if !completions.is_empty() {
Expand Down
13 changes: 7 additions & 6 deletions crates/ark/src/lsp/main_loop.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ use crate::lsp::backend::LspRequest;
use crate::lsp::backend::LspResponse;
use crate::lsp::backend::LspResult;
use crate::lsp::capabilities::Capabilities;
use crate::lsp::db::Analysis;
use crate::lsp::diagnostics::generate_diagnostics;
use crate::lsp::handlers;
use crate::lsp::indexer;
Expand All @@ -58,6 +59,7 @@ use crate::lsp::sources::OakSourceHandler;
use crate::lsp::sources::SourceCompleted;
use crate::lsp::sources::SourceHandler;
use crate::lsp::sources::SourceScheduler;
use crate::lsp::state::WorldSnapshot;
use crate::lsp::state::WorldState;
use crate::lsp::state_handlers;
use crate::lsp::state_handlers::ConsoleInputs;
Expand Down Expand Up @@ -559,7 +561,7 @@ impl GlobalState {
// and the diagnostics passes they trigger force the same
// memos.
if !self.lsp_state.oak_scheduler.has_pending_scans() {
warm_workspace_index(self.world.db.clone());
warm_workspace_index(Analysis::new(self.world.db.clone()));
}
},

Expand Down Expand Up @@ -1005,7 +1007,7 @@ impl std::fmt::Debug for TraceKernelNotification<'_> {
pub(crate) struct RefreshDiagnosticsTask {
/// Snapshot carrying the live oak plus the session context the diagnostics
/// walk reads. See [`WorldState::diagnostics_snapshot`].
state: WorldState,
state: WorldSnapshot,
/// The file to diagnose, built against the live oak at enqueue time.
file: OpenFile,
}
Expand Down Expand Up @@ -1144,11 +1146,11 @@ pub(crate) fn diagnostics_refresh_all(state: &WorldState) {
/// cancelled (`spawn_blocking()` swallows the unwind). A cancelling write can
/// only come from an editor buffer, so a document is open, and the diagnostics
/// passes spawned by that same write force the same memos and finish the job.
fn warm_workspace_index(db: OakDatabase) {
fn warm_workspace_index(analysis: Analysis) {
spawn_blocking(move || {
let now = std::time::Instant::now();
lsp::log_info!("Starting workspace index warmup");
indexer::warm(&db);
indexer::warm(analysis.read());
lsp::log_info!("Finished workspace index warmup ({:.0?})", now.elapsed());
Ok(None)
})
Expand All @@ -1158,7 +1160,6 @@ fn warm_workspace_index(db: OakDatabase) {
mod tests {
use aether_path::FilePath;
use oak_scan::DbScan;
use salsa::Database;
use tower_lsp::jsonrpc;
use url::Url;

Expand Down Expand Up @@ -1222,7 +1223,7 @@ mod tests {
respond(
response_tx,
|| {
let _ = file.tree_sitter(&snapshot.db);
let _ = file.tree_sitter(snapshot.db.read());
Ok(LspResponse::Hover(None))
},
|response| response,
Expand Down
Loading
Loading