Fix race conditions between indexer and diagnostics

Author: lionel-

crates/ark/src/lsp/main_loop.rs

@@ -11,11 +11,15 @@ use std::path::PathBuf;
 use std::pin::Pin;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
+use std::sync::LazyLock;
 use std::sync::RwLock;
 
 use anyhow::anyhow;
+use futures::stream::FuturesUnordered;
 use futures::StreamExt;
+use tokio::sync::mpsc;
 use tokio::sync::mpsc::unbounded_channel as tokio_unbounded_channel;
+use tokio::task;
 use tokio::task::JoinHandle;
 use tower_lsp::lsp_types;
 use tower_lsp::lsp_types::Diagnostic;
@@ -30,9 +34,10 @@ use crate::lsp::backend::LspNotification;
 use crate::lsp::backend::LspRequest;
 use crate::lsp::backend::LspResponse;
 use crate::lsp::capabilities::Capabilities;
-use crate::lsp::diagnostics;
+use crate::lsp::diagnostics::generate_diagnostics;
 use crate::lsp::documents::Document;
 use crate::lsp::handlers;
+use crate::lsp::indexer;
 use crate::lsp::inputs::library::Library;
 use crate::lsp::state::WorldState;
 use crate::lsp::state_handlers;
@@ -641,27 +646,6 @@ where
     send_auxiliary(AuxiliaryEvent::SpawnedTask(handle));
 }
 
-pub(crate) fn spawn_diagnostics_refresh(uri: Url, document: Document, state: WorldState) {
-    lsp::spawn_blocking(move || {
-        let _s = tracing::info_span!("diagnostics_refresh", uri = %uri).entered();
-
-        let version = document.version;
-        let diagnostics = diagnostics::generate_diagnostics(document, state);
-
-        Ok(Some(AuxiliaryEvent::PublishDiagnostics(
-            uri,
-            diagnostics,
-            version,
-        )))
-    })
-}
-
-pub(crate) fn spawn_diagnostics_refresh_all(state: WorldState) {
-    for (url, document) in state.documents.iter() {
-        spawn_diagnostics_refresh(url.clone(), document.clone(), state.clone())
-    }
-}
-
 pub(crate) fn publish_diagnostics(uri: Url, diagnostics: Vec<Diagnostic>, version: Option<i32>) {
     send_auxiliary(AuxiliaryEvent::PublishDiagnostics(
         uri,
@@ -692,3 +676,191 @@ impl std::fmt::Debug for TraceKernelNotification<'_> {
         }
     }
 }
+
+#[derive(Debug)]
+pub(crate) enum IndexerQueueTask {
+    Indexer(IndexerTask),
+    Diagnostics(RefreshDiagnosticsTask),
+}
+
+#[derive(Debug)]
+pub enum IndexerTask {
+    Start { folders: Vec<String> },
+    Update { document: Document, uri: Url },
+}
+
+#[derive(Debug)]
+pub(crate) struct RefreshDiagnosticsTask {
+    uri: Url,
+    state: WorldState,
+}
+
+#[derive(Debug)]
+struct RefreshDiagnosticsResult {
+    uri: Url,
+    diagnostics: Vec<Diagnostic>,
+    version: Option<i32>,
+}
+
+static INDEXER_QUEUE: LazyLock<tokio::sync::mpsc::UnboundedSender<IndexerQueueTask>> =
+    LazyLock::new(|| {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        tokio::spawn(process_indexer_queue(rx));
+        tx
+    });
+
+/// Process indexer and diagnostics tasks
+///
+/// Diagnostics need an up-to-date index to be accurate, so we synchronise
+/// indexing and diagnostics tasks using a simple queue.
+///
+/// - We make sure to refresh diagnostics after every indexer updates.
+/// - Indexer tasks are batched together, same for diagnostics tasks.
+/// - Cancellation is simply dealt with by deduplicating tasks for the same URI,
+///   retaining only the most recent one.
+///
+/// Ideally we'd process indexer tasks continually without making them dependent
+/// on diagnostics tasks. The current setup blocks the queue loop while
+/// diagnostics are running, but it has the benefit that rounds of diagnostic
+/// refreshes don't race against each other. The frontend will receive all
+/// results in order, ensuring that diagnostics for an outdated version are
+/// eventually replaced by the most up-to-date diagnostics.
+async fn process_indexer_queue(mut rx: mpsc::UnboundedReceiver<IndexerQueueTask>) {
+    while let Some(task) = rx.recv().await {
+        // Drain all available tasks
+        let mut tasks = vec![task];
+        while let Ok(next_task) = rx.try_recv() {
+            tasks.push(next_task);
+        }
+
+        // Separate by type
+        let mut diagnostics_batch = Vec::new();
+        let mut indexer_batch = Vec::new();
+
+        for task in tasks {
+            match task {
+                IndexerQueueTask::Indexer(indexer_task) => indexer_batch.push(indexer_task),
+                IndexerQueueTask::Diagnostics(diagnostic_task) => {
+                    diagnostics_batch.push(diagnostic_task)
+                },
+            }
+        }
+
+        // Process indexer tasks first so diagnostics tasks work with an up-to-date index
+        process_indexer_batch(indexer_batch).await;
+        process_diagnostics_batch(diagnostics_batch).await;
+    }
+}
+
+async fn process_indexer_batch(batch: Vec<IndexerTask>) {
+    // Deduplicate tasks by key. We use a `HashMap` so only the last insertion
+    // is retained. `Update` tasks use URI as key, `Start` tasks use None (we
+    // only expect one though). This is effectively a way of cancelling `Update`
+    // tasks for outdated documents.
+    let batch: std::collections::HashMap<_, _> = batch
+        .into_iter()
+        .map(|task| match &task {
+            IndexerTask::Update { uri, .. } => (Some(uri.clone()), task),
+            IndexerTask::Start { .. } => (None, task),
+        })
+        .collect();
+
+    let mut handles = Vec::new();
+
+    for (_, task) in batch {
+        handles.push(tokio::task::spawn_blocking(move || match task {
+            IndexerTask::Start { folders } => {
+                indexer::start(folders);
+            },
+            IndexerTask::Update { document, uri } => {
+                let result = if let Ok(path) = uri.to_file_path() {
+                    indexer::update(&document, &path)
+                } else {
+                    Err(anyhow!("Failed to convert URI to file path: {uri}"))
+                };
+                if let Err(err) = result {
+                    log::error!("Indexer update failed: {err}");
+                }
+            },
+        }));
+    }
+
+    for handle in handles {
+        let _ = handle.await;
+    }
+}
+
+async fn process_diagnostics_batch(batch: Vec<RefreshDiagnosticsTask>) {
+    // Deduplicate tasks by keeping only the last one for each URI. We use a
+    // `HashMap` so only the last insertion is retained. This is effectively a
+    // way of cancelling diagnostics tasks for outdated documents.
+    let batch: std::collections::HashMap<_, _> = batch
+        .into_iter()
+        .map(|task| (task.uri, task.state))
+        .collect();
+
+    let mut futures = FuturesUnordered::new();
+
+    for (uri, state) in batch {
+        futures.push(task::spawn_blocking(move || {
+            let _span = tracing::info_span!("diagnostics_refresh", uri = %uri).entered();
+
+            if let Some(document) = state.documents.get(&uri) {
+                let diagnostics = generate_diagnostics(document.clone(), state.clone());
+                Some(RefreshDiagnosticsResult {
+                    uri,
+                    diagnostics,
+                    version: document.version,
+                })
+            } else {
+                None
+            }
+        }));
+    }
+
+    // Publish results as they complete
+    while let Some(Ok(Some(result))) = futures.next().await {
+        publish_diagnostics(result.uri, result.diagnostics, result.version);
+    }
+}
+
+pub(crate) fn index_start(folders: Vec<String>, state: WorldState) {
+    INDEXER_QUEUE
+        .send(IndexerQueueTask::Indexer(IndexerTask::Start { folders }))
+        .unwrap_or_else(|err| lsp::log_error!("Failed to queue initial indexing: {err}"));
+
+    diagnostics_refresh_all(state);
+}
+
+pub(crate) fn index_update(uri: Url, document: Document, state: WorldState) {
+    INDEXER_QUEUE
+        .send(IndexerQueueTask::Indexer(IndexerTask::Update {
+            document,
+            uri: uri.clone(),
+        }))
+        .unwrap_or_else(|err| lsp::log_error!("Failed to queue index update: {err}"));
+
+    diagnostics_refresh(uri, state.clone());
+}
+
+pub(crate) fn diagnostics_refresh(uri: Url, state: WorldState) {
+    INDEXER_QUEUE
+        .send(IndexerQueueTask::Diagnostics(RefreshDiagnosticsTask {
+            uri,
+            state,
+        }))
+        .unwrap_or_else(|err| lsp::log_error!("Failed to queue diagnostics refresh: {err}"));
+}
+
+pub(crate) fn diagnostics_refresh_all(state: WorldState) {
+    // Expand RefreshAll into individual RefreshDiagnostics tasks
+    // This allows the deduplication logic to work uniformly
+    for (uri, _document) in state.documents.iter() {
+        INDEXER_QUEUE
+            .send(IndexerQueueTask::Diagnostics(RefreshDiagnosticsTask {
+                uri: uri.clone(),
+                state: state.clone(),
+            }))
+            .unwrap_or_else(|err| lsp::log_error!("Failed to queue diagnostics refresh: {err}"));
+    }
+}

crates/ark/src/lsp/mod.rs

@@ -63,7 +63,6 @@ pub(crate) use _log;
 pub(crate) use log_error;
 pub(crate) use log_info;
 pub(crate) use log_warn;
+pub(crate) use main_loop::diagnostics_refresh_all;
 pub(crate) use main_loop::publish_diagnostics;
 pub(crate) use main_loop::spawn_blocking;
-pub(crate) use main_loop::spawn_diagnostics_refresh;
-pub(crate) use main_loop::spawn_diagnostics_refresh_all;

crates/ark/src/lsp/state_handlers.rs

@@ -5,8 +5,6 @@
 //
 //
 
-use std::path::Path;
-
 use anyhow::anyhow;
 use tower_lsp::lsp_types;
 use tower_lsp::lsp_types::CompletionOptions;
@@ -44,7 +42,6 @@ use crate::lsp::config::DOCUMENT_SETTINGS;
 use crate::lsp::config::GLOBAL_SETTINGS;
 use crate::lsp::documents::Document;
 use crate::lsp::encoding::get_position_encoding_kind;
-use crate::lsp::indexer;
 use crate::lsp::inputs::package::Package;
 use crate::lsp::inputs::source_root::SourceRoot;
 use crate::lsp::main_loop::DidCloseVirtualDocumentParams;
@@ -121,10 +118,7 @@ pub(crate) fn initialize(
     }
 
     // Start first round of indexing
-    lsp::spawn_blocking(|| {
-        indexer::start(folders);
-        Ok(None)
-    });
+    lsp::main_loop::index_start(folders, state.clone());
 
     Ok(InitializeResult {
         server_info: Some(ServerInfo {
@@ -206,8 +200,7 @@ pub(crate) fn did_open(
     // NOTE: Do we need to call `update_config()` here?
     // update_config(vec![uri]).await;
 
-    update_index(&uri, &document);
-    lsp::spawn_diagnostics_refresh(uri, document, state.clone());
+    lsp::main_loop::index_update(uri.clone(), document.clone(), state.clone());
 
     Ok(())
 }
@@ -219,17 +212,16 @@ pub(crate) fn did_change(
     state: &mut WorldState,
 ) -> anyhow::Result<()> {
     let uri = &params.text_document.uri;
-    let doc = state.get_document_mut(uri)?;
+    let document = state.get_document_mut(uri)?;
 
     let mut parser = lsp_state
         .parsers
         .get_mut(uri)
         .ok_or(anyhow!("No parser for {uri}"))?;
 
-    doc.on_did_change(&mut parser, &params);
+    document.on_did_change(&mut parser, &params);
 
-    update_index(uri, doc);
-    lsp::spawn_diagnostics_refresh(uri.clone(), doc.clone(), state.clone());
+    lsp::main_loop::index_update(uri.clone(), document.clone(), state.clone());
 
     Ok(())
 }
@@ -384,7 +376,7 @@ async fn update_config(
     // Refresh diagnostics if the configuration changed
     if state.config.diagnostics != diagnostics_config {
         tracing::info!("Refreshing diagnostics after configuration changed");
-        lsp::spawn_diagnostics_refresh_all(state.clone());
+        lsp::main_loop::diagnostics_refresh_all(state.clone());
     }
 
     Ok(())
@@ -402,7 +394,7 @@ pub(crate) fn did_change_console_inputs(
     // during package development in conjunction with `devtools::load_all()`.
     // Ideally diagnostics would not rely on these though, and we wouldn't need
     // to refresh from here.
-    lsp::spawn_diagnostics_refresh_all(state.clone());
+    lsp::diagnostics_refresh_all(state.clone());
 
     Ok(())
 }
@@ -425,16 +417,3 @@ pub(crate) fn did_close_virtual_document(
     state.virtual_documents.remove(&params.uri);
     Ok(())
 }
-
-// FIXME: The initial indexer is currently racing against our state notification
-// handlers. The indexer is synchronised through a mutex but we might end up in
-// a weird state. Eventually the index should be moved to WorldState and created
-// on demand with Salsa instrumenting and cancellation.
-fn update_index(uri: &url::Url, doc: &Document) {
-    if let Ok(path) = uri.to_file_path() {
-        let path = Path::new(&path);
-        if let Err(err) = indexer::update(&doc, &path) {
-            lsp::log_error!("{err:?}");
-        }
-    }
-}