Add support for asynchronous scanning

Author: lionel-

crates/ark/src/lsp/main_loop.rs

@@ -6,6 +6,7 @@
 //
 
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::future;
 use std::path::Path;
 use std::path::PathBuf;
@@ -15,12 +16,17 @@ use std::sync::atomic::Ordering;
 use std::sync::LazyLock;
 use std::sync::RwLock;
 
+use aether_url::UrlId;
 use anyhow::anyhow;
 use futures::stream::FuturesUnordered;
 use futures::StreamExt;
 use oak_db::OakDatabase;
 use oak_scan::DbExt;
+use oak_scan::ScanCompleted;
+use oak_scan::ScanRequest;
+use oak_scan::ScanScheduler;
 use oak_semantic::library::Library;
+use stdext::result::ResultExt;
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::unbounded_channel as tokio_unbounded_channel;
 use tokio::task;
@@ -88,6 +94,7 @@ type TaskList<T> = futures::stream::FuturesUnordered<Pin<Box<dyn AnyhowJoinHandl
 pub(crate) enum Event {
     Lsp(LspMessage),
     Kernel(KernelNotification),
+    OakScanCompleted(ScanCompleted),
 }
 
 #[derive(Debug)]
@@ -151,8 +158,9 @@ pub(crate) struct GlobalState {
     events_rx: TokioUnboundedReceiver<Event>,
 }
 
-/// Unlike `WorldState`, `ParserState` cannot be cloned and is only accessed by
-/// exclusive handlers.
+/// Non-cloneable, per-session state mutated only by exclusive handlers.
+/// Sits alongside [`WorldState`] (which is cloneable for snapshot
+/// handlers); state that can't be cloned lives here instead.
 pub(crate) struct LspState {
     /// The set of tree-sitter document parsers managed by the `GlobalState`.
     pub(crate) parsers: HashMap<Url, tree_sitter::Parser>,
@@ -162,6 +170,11 @@ pub(crate) struct LspState {
 
     /// Channel for sending notifications to Console (e.g., document changes for DAP)
     pub(crate) console_notification_tx: TokioUnboundedSender<ConsoleNotification>,
+
+    /// Coordinator for asynchronous workspace scans. Mutated only from
+    /// main-loop handlers. Must be out of [`WorldState`] because the scheduler
+    /// is not clonable.
+    pub(crate) oak_scheduler: ScanScheduler,
 }
 
 /// State for the auxiliary loop
@@ -199,6 +212,7 @@ impl GlobalState {
             parsers: HashMap::new(),
             capabilities: Capabilities::default(),
             console_notification_tx,
+            oak_scheduler: ScanScheduler::new(),
         };
 
         // FIXME: We shouldn't call R code in the kernel to figure this out
@@ -299,13 +313,15 @@ impl GlobalState {
                             handlers::handle_initialized(&self.client, &self.lsp_state).await?;
                         },
                         LspNotification::DidChangeWorkspaceFolders(params) => {
-                            state_handlers::did_change_workspace_folders(params, &mut self.world)?;
+                            let pending = state_handlers::did_change_workspace_folders(params, &mut self.world, &mut self.lsp_state)?;
+                            dispatch_scan_requests(&self.events_tx, pending);
                         },
                         LspNotification::DidChangeConfiguration(params) => {
                             state_handlers::did_change_configuration(params, &self.client, &mut self.world).await?;
                         },
                         LspNotification::DidChangeWatchedFiles(params) => {
-                            state_handlers::did_change_watched_files(params, &mut self.world)?;
+                            let pending = state_handlers::did_change_watched_files(params, &mut self.world, &mut self.lsp_state)?;
+                            dispatch_scan_requests(&self.events_tx, pending);
                         },
                         LspNotification::DidOpenTextDocument(params) => {
                             state_handlers::did_open(params, &mut self.lsp_state, &mut self.world)?;
@@ -336,7 +352,8 @@ impl GlobalState {
 
                     match request {
                         LspRequest::Initialize(params) => {
-                            respond(tx, || state_handlers::initialize(params, &mut self.lsp_state, &mut self.world), LspResponse::Initialize)?;
+                            let pending = respond_with(tx, || state_handlers::initialize(params, &mut self.lsp_state, &mut self.world), LspResponse::Initialize)?;
+                            dispatch_scan_requests(&self.events_tx, pending);
                         },
                         LspRequest::WorkspaceSymbol(params) => {
                             respond(tx, || handlers::handle_symbol(params, &self.world), LspResponse::WorkspaceSymbol)?;
@@ -420,6 +437,25 @@ impl GlobalState {
                     }
                 }
             },
+
+            Event::OakScanCompleted(scan) => {
+                // Recompute editor-owned files at apply time, not at spawn
+                // time: a buffer may have opened or closed since the scan
+                // kicked off. The buffer-drain inside `apply_scan_completed` uses
+                // this set as its watcher-event `skip` argument.
+                let editor_owned: HashSet<UrlId> = self.world
+                    .documents
+                    .keys()
+                    .map(|url| UrlId::from_url(url.clone()))
+                    .collect();
+
+                let followups = self.lsp_state.oak_scheduler.apply_scan_completed(
+                    &mut self.world.oak,
+                    scan,
+                    &editor_owned,
+                );
+                dispatch_scan_requests(&self.events_tx, followups);
+            },
         }
 
         // TODO Make this threshold configurable by the client
@@ -452,6 +488,24 @@ impl GlobalState {
     }
 }
 
+/// Spawn each [`ScanRequest`] on a blocking task. Each task runs the
+/// pure-I/O [`ScanRequest::run`] and ships the [`ScanCompleted`] back
+/// to the main loop as [`Event::OakScanCompleted`], where the scheduler
+/// then applies it.
+pub(super) fn dispatch_scan_requests(
+    events_tx: &TokioUnboundedSender<Event>,
+    requests: Vec<ScanRequest>,
+) {
+    for req in requests {
+        let tx = events_tx.clone();
+        spawn_blocking(move || {
+            let scan = req.run();
+            tx.send(Event::OakScanCompleted(scan)).log_err();
+            Ok(None)
+        });
+    }
+}
+
 /// Respond to a request from the LSP
 ///
 /// We receive requests from the LSP client with a response channel. Once we
@@ -475,11 +529,28 @@ fn respond<T>(
     response: impl FnOnce() -> LspResult<T>,
     into_lsp_response: impl FnOnce(T) -> LspResponse,
 ) -> anyhow::Result<()> {
-    let response = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(response)) {
-        Ok(response) => {
-            let response = response.map(into_lsp_response);
-            RequestResponse::Result(response)
-        },
+    respond_with(
+        response_tx,
+        || response().map(|t| (t, ())),
+        into_lsp_response,
+    )
+}
+
+/// Variant of [`respond`] for handlers that produce a side output alongside
+/// their LSP response. The `S` value is returned to the caller on success;
+/// on handler error or panic the side output is `S::default()` (the caller
+/// gets the "do nothing" value, while the error response still flows to
+/// the client through `response_tx`). Used by `Initialize` to ship the
+/// scheduler's pending `ScanRequest`s back to the main loop without an
+/// out-parameter.
+fn respond_with<T, S: Default>(
+    response_tx: TokioUnboundedSender<RequestResponse>,
+    response: impl FnOnce() -> LspResult<(T, S)>,
+    into_lsp_response: impl FnOnce(T) -> LspResponse,
+) -> anyhow::Result<S> {
+    let (response, side) = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(response)) {
+        Ok(Ok((t, s))) => (RequestResponse::Result(Ok(into_lsp_response(t))), s),
+        Ok(Err(e)) => (RequestResponse::Result(Err(e)), S::default()),
         Err(err) => {
             // Set global crash flag to disable the LSP
             LSP_HAS_CRASHED.store(true, Ordering::Release);
@@ -495,19 +566,22 @@ fn respond<T>(
             // This creates an uninformative backtrace that is reported in the
             // LSP logs. Note that the relevant backtrace is the one created by
             // our panic hook and reported via the _kernel_ logs.
-            RequestResponse::Crashed(anyhow!("Panic occurred while handling request: {msg}"))
+            (
+                RequestResponse::Crashed(anyhow!("Panic occurred while handling request: {msg}")),
+                S::default(),
+            )
         },
     };
 
     let out = match response {
-        RequestResponse::Result(Ok(_)) => Ok(()),
+        RequestResponse::Result(Ok(_)) => Ok(side),
         RequestResponse::Result(Err(ref error)) => {
             // The error has already been sent to the client on `response_tx`
             // as a jsonrpc error, so the user sees the popup. Log here at
             // info level (with `{:?}` for the full debug format including a
             // backtrace) so server logs keep diagnostic context.
             lsp::log_info!("Error while handling request:\n{error:?}");
-            Ok(())
+            Ok(side)
         },
         RequestResponse::Crashed(ref error) => {
             Err(anyhow!("Crashed while handling request:\n{error:?}"))

crates/ark/src/lsp/state_handlers.rs

@@ -13,6 +13,7 @@ use anyhow::anyhow;
 use oak_scan::DbExt;
 use oak_scan::FileEvent;
 use oak_scan::FileEventKind;
+use oak_scan::ScanRequest;
 use oak_semantic::package::Package;
 use stdext::result::ResultExt;
 use tower_lsp::lsp_types;
@@ -94,7 +95,7 @@ pub(crate) fn initialize(
     params: InitializeParams,
     lsp_state: &mut LspState,
     state: &mut WorldState,
-) -> LspResult<InitializeResult> {
+) -> LspResult<(InitializeResult, Vec<ScanRequest>)> {
     let workspace_uris = effective_workspace_uris(&params);
     lsp_state.capabilities = Capabilities::new(params.capabilities);
 
@@ -147,12 +148,14 @@ pub(crate) fn initialize(
         .keys()
         .map(|url| UrlId::from_url(url.clone()))
         .collect();
-    state
-        .oak
-        .set_workspace_paths(&workspace_paths, &editor_owned);
+    let pending = lsp_state.oak_scheduler.set_workspace_paths(
+        &mut state.oak,
+        &workspace_paths,
+        &editor_owned,
+    );
     lsp::main_loop::index_start(folders, state.clone());
 
-    Ok(InitializeResult {
+    let result = InitializeResult {
         server_info: Some(ServerInfo {
             name: "Ark R Kernel".to_string(),
             version: Some(crate::BUILD_VERSION.to_string()),
@@ -234,7 +237,9 @@ pub(crate) fn initialize(
             }),
             ..ServerCapabilities::default()
         },
-    })
+    };
+
+    Ok((result, pending))
 }
 
 /// Resolve the effective workspace folders from `InitializeParams`.
@@ -400,7 +405,8 @@ pub(crate) fn did_rename_files(
 pub(crate) fn did_change_watched_files(
     params: DidChangeWatchedFilesParams,
     state: &mut WorldState,
-) -> anyhow::Result<()> {
+    lsp_state: &mut LspState,
+) -> anyhow::Result<Vec<ScanRequest>> {
     // Editor owns the contents of files it has open: Oak should ignore
     // disk-side events for those URLs.
     let editor_owned: HashSet<UrlId> = state
@@ -426,15 +432,18 @@ pub(crate) fn did_change_watched_files(
         })
         .collect();
 
-    state.oak.apply_watcher_events(events, &editor_owned);
-    Ok(())
+    let pending = lsp_state
+        .oak_scheduler
+        .apply_watcher_events(&mut state.oak, events, &editor_owned);
+    Ok(pending)
 }
 
 #[tracing::instrument(level = "info", skip_all)]
 pub(crate) fn did_change_workspace_folders(
     params: DidChangeWorkspaceFoldersParams,
     state: &mut WorldState,
-) -> anyhow::Result<()> {
+    lsp_state: &mut LspState,
+) -> anyhow::Result<Vec<ScanRequest>> {
     let removed: HashSet<Url> = params.event.removed.iter().map(|f| f.uri.clone()).collect();
     state.workspace.folders.retain(|uri| !removed.contains(uri));
 
@@ -460,10 +469,12 @@ pub(crate) fn did_change_workspace_folders(
         .map(|url| UrlId::from_url(url.clone()))
         .collect();
 
-    state
-        .oak
-        .set_workspace_paths(&workspace_paths, &editor_owned);
-    Ok(())
+    let pending = lsp_state.oak_scheduler.set_workspace_paths(
+        &mut state.oak,
+        &workspace_paths,
+        &editor_owned,
+    );
+    Ok(pending)
 }
 
 fn parse_uri_or_none(uri: &str) -> Option<url::Url> {

crates/ark/src/lsp/tests.rs

@@ -1,5 +1,6 @@
 mod find_references;
 mod goto_definition;
+mod main_loop;
 mod rename;
 mod state_handlers;
 mod utils;

crates/ark/src/lsp/tests/main_loop.rs

@@ -0,0 +1,77 @@
+//! End-to-end smoke test of the async workspace scan path.
+//!
+//! Exercises the real tokio plumbing the LSP main loop relies on:
+//! [`dispatch_scan_requests`] spawning [`ScanRequest::run`] on a
+//! blocking task, the mpsc round-trip back as [`Event::OakScanCompleted`],
+//! and the main-loop apply step. The rest of the scheduler is unit
+//! tested without tokio in `oak_scan`; this test pins the wiring.
+
+use std::collections::HashSet;
+use std::fs;
+use std::time::Duration;
+
+use oak_db::DbInputs;
+use oak_db::OakDatabase;
+use oak_scan::ScanScheduler;
+use tokio::runtime::Runtime;
+use tokio::time::timeout;
+
+use crate::lsp::main_loop::dispatch_scan_requests;
+use crate::lsp::main_loop::init_aux_for_test;
+use crate::lsp::main_loop::Event;
+
+#[test]
+fn test_workspace_scan_round_trip_through_tokio() {
+    // `dispatch_scan_requests` routes through `lsp::spawn_blocking`, which
+    // hands the `JoinHandle` to the aux loop. Without an aux tx the spawn
+    // helper drops the panic-logging path; init a no-op aux for the test.
+    init_aux_for_test();
+
+    let tmp = tempfile::tempdir().unwrap();
+    fs::create_dir_all(tmp.path().join("pkg/R")).unwrap();
+    fs::write(
+        tmp.path().join("pkg/DESCRIPTION"),
+        "Package: pkg\nVersion: 0.0.0\n",
+    )
+    .unwrap();
+    fs::write(tmp.path().join("pkg/R/a.R"), "x <- 1\n").unwrap();
+
+    let rt = Runtime::new().unwrap();
+    rt.block_on(async {
+        let mut db = OakDatabase::new();
+        let mut scheduler = ScanScheduler::new();
+        let (events_tx, mut events_rx) = tokio::sync::mpsc::unbounded_channel::<Event>();
+
+        // Kick off the scan as `did_change_workspace_folders` would.
+        let initial =
+            scheduler.set_workspace_paths(&mut db, &[tmp.path().to_path_buf()], &HashSet::new());
+        assert_eq!(initial.len(), 1);
+        dispatch_scan_requests(&events_tx, initial);
+
+        // Pump events until the scheduler stops issuing followups. Each
+        // iteration exercises the real spawn_blocking + mpsc + apply
+        // round-trip the main loop performs in production.
+        loop {
+            let event = timeout(Duration::from_secs(5), events_rx.recv())
+                .await
+                .expect("scan timed out")
+                .expect("event channel closed");
+            let Event::OakScanCompleted(scan) = event else {
+                panic!("unexpected event variant");
+            };
+            let followups = scheduler.apply_scan_completed(&mut db, scan, &HashSet::new());
+            if followups.is_empty() {
+                break;
+            }
+            dispatch_scan_requests(&events_tx, followups);
+        }
+
+        // Workspace state reflects the on-disk package.
+        let roots = db.workspace_roots().roots(&db).clone();
+        assert_eq!(roots.len(), 1);
+        let packages = roots[0].packages(&db).clone();
+        assert_eq!(packages.len(), 1);
+        assert_eq!(packages[0].name(&db), "pkg");
+        assert_eq!(packages[0].files(&db).len(), 1);
+    });
+}