Run the main loop on its own thread

Author: lionel-

crates/ark/src/lsp/backend.rs

@@ -47,6 +47,7 @@ use crate::lsp::input_boundaries::InputBoundariesParams;
 use crate::lsp::input_boundaries::InputBoundariesResponse;
 use crate::lsp::main_loop::Event;
 use crate::lsp::main_loop::GlobalState;
+use crate::lsp::main_loop::LoopHandles;
 use crate::lsp::main_loop::TokioUnboundedSender;
 use crate::lsp::statement_range;
 use crate::lsp::statement_range::StatementRangeParams;
@@ -232,9 +233,9 @@ struct Backend {
     /// Channel for communication with the main loop.
     events_tx: TokioUnboundedSender<Event>,
 
-    /// Handle to main loop. Drop it to cancel the loop, all associated tasks,
-    /// and drop all owned state.
-    _main_loop: tokio::task::JoinSet<()>,
+    /// Handle to the LSP loops. Drop it to shut the loops down and drop all
+    /// owned state.
+    _main_loop: LoopHandles,
 }
 
 impl Backend {

crates/ark/src/lsp/handler.rs

@@ -39,7 +39,8 @@ impl Lsp {
     ) -> Self {
         let rt = Builder::new_multi_thread()
             .enable_all()
-            // One for the main loop and one spare
+            // Workers serve tower-lsp, the auxiliary loop, and the diagnostics
+            // queue. The main loop runs on its own thread.
             .worker_threads(2)
             // Used for diagnostics
             .max_blocking_threads(2)

crates/ark/src/lsp/main_loop.rs

@@ -26,10 +26,11 @@ use oak_scan::ScanRequest;
 use oak_scan::ScanScheduler;
 use oak_semantic::library::Library;
 use stdext::result::ResultExt;
+use stdext::spawn;
 use tokio::runtime::Handle;
-use tokio::runtime::RuntimeFlavor;
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::unbounded_channel as tokio_unbounded_channel;
+use tokio::sync::oneshot;
 use tokio::task::JoinHandle;
 use tower_lsp::jsonrpc;
 use tower_lsp::lsp_types;
@@ -159,6 +160,21 @@ pub(crate) struct GlobalState {
     events_rx: TokioUnboundedReceiver<Event>,
 }
 
+/// Owns the running LSP loops. Dropping it shuts them down.
+///
+/// - The auxiliary loop is a runtime task, aborted when `_aux_loop` drops.
+/// - The main loop runs on its own thread. Dropping `_main_shutdown_tx` closes
+///   the channel the loop selects on, so it breaks, drops the owned
+///   `GlobalState`, and the thread exits. We hold the thread's handle but don't
+///   join it on drop (it winds down on its own), so dropping never blocks the
+///   caller.
+#[derive(Debug)]
+pub(crate) struct LoopHandles {
+    _main_loop: std::thread::JoinHandle<()>,
+    _main_shutdown_tx: oneshot::Sender<()>,
+    _aux_loop: tokio::task::JoinSet<()>,
+}
+
 /// 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.
@@ -267,37 +283,66 @@ impl GlobalState {
         self.events_tx.clone()
     }
 
-    /// Start the main and auxiliary loops
+    /// Start the main and auxiliary loops.
     ///
-    /// Returns a `JoinSet` that holds onto all tasks and state owned by the
-    /// event loop. Drop it to cancel everything and shut down the service.
-    pub(crate) fn start(self) -> tokio::task::JoinSet<()> {
-        let mut set = tokio::task::JoinSet::<()>::new();
-
-        // Spawn latency-sensitive auxiliary loop. Must be first to initialise
-        // global transmission channel.
-        let aux = AuxiliaryState::new(self.client.clone());
-        set.spawn(async move { aux.start().await });
-
-        // Spawn main loop
-        set.spawn(async move { self.main_loop().await });
+    /// The returned [`LoopHandles`] owns everything the loops need. Drop it to
+    /// shut the loops down and release the owned state.
+    pub(crate) fn start(self) -> LoopHandles {
+        let mut aux = tokio::task::JoinSet::<()>::new();
+
+        // The auxiliary loop is fully async and never blocks. Must be started
+        // first to initialise the global transmission channel.
+        let aux_state = AuxiliaryState::new(self.client.clone());
+        aux.spawn(async move { aux_state.start().await });
+
+        // Since the main loop owns the Salsa DB and writes to it, we run on its
+        // own thread instead of a Tokio worker. Salsa writes are potentially
+        // blocking until the writer gains exclusive access. If background tasks
+        // holding clones of the DB are stuck on the same thread as the main
+        // loop, the LSP deadlocks. This can be avoided by wrapping writes in
+        // `block_in_place()` but the safer structure is to have it run on an OS
+        // thread that we're in control of.
+        let (shutdown_tx, shutdown_rx) = oneshot::channel();
+        let handle = Handle::current();
+        let main_loop = spawn!("oak-main-loop", move || {
+            handle.block_on(self.main_loop(shutdown_rx));
+        });
 
-        set
+        LoopHandles {
+            _main_shutdown_tx: shutdown_tx,
+            _aux_loop: aux,
+            _main_loop: main_loop,
+        }
     }
 
     /// Run main loop
     ///
     /// This takes ownership of all global state and handles one by one LSP
     /// requests, notifications, and other internal events.
-    async fn main_loop(mut self) {
+    async fn main_loop(mut self, mut shutdown_rx: oneshot::Receiver<()>) {
         loop {
-            let event = self.next_event().await;
-            if let Err(err) = self.handle_event(event).await {
-                lsp::log_error!("Failure while handling event:\n{err:?}")
+            tokio::select! {
+                _ = &mut shutdown_rx => {
+                    lsp::log_info!("Main loop stopping: handle dropped");
+                    break;
+                },
+                event = self.events_rx.recv() => {
+                    let Some(event) = event else {
+                        lsp::log_info!("Main loop stopping: event channel closed");
+                        break;
+                    };
+                    if let Err(err) = self.handle_event(event).await {
+                        lsp::log_error!("Failure while handling event:\n{err:?}")
+                    }
+                }
             }
         }
     }
 
+    /// Pull the next event off the channel. Only the test pump uses this; the
+    /// running loop selects on the channel directly so it can also watch for
+    /// shutdown.
+    #[cfg(test)]
     async fn next_event(&mut self) -> Event {
         self.events_rx.recv().await.unwrap()
     }
@@ -351,10 +396,10 @@ impl GlobalState {
                             state_handlers::did_change_watched_files(params, &mut self.world, &mut self.lsp_state, &self.events_tx)?;
                         },
                         LspNotification::DidOpenTextDocument(params) => {
-                            block_for_write(|| state_handlers::did_open(params, &mut self.world))?;
+                            state_handlers::did_open(params, &mut self.world)?;
                         },
                         LspNotification::DidChangeTextDocument(params) => {
-                            block_for_write(|| state_handlers::did_change(params, &mut self.lsp_state, &mut self.world))?;
+                            state_handlers::did_change(params, &mut self.lsp_state, &mut self.world)?;
                         },
                         LspNotification::DidSaveTextDocument(_params) => {
                             // Currently ignored
@@ -479,13 +524,11 @@ impl GlobalState {
                 // kicked off. The buffer-drain inside `apply_scan_completed` uses
                 // this set as its watcher-event `skip` argument.
                 let editor_owned: HashSet<FilePath> = self.world.open_files.keys().cloned().collect();
-                let followups = block_for_write(|| {
-                    self.lsp_state.oak_scheduler.apply_scan_completed(
-                        &mut self.world.db,
-                        scan,
-                        &editor_owned,
-                    )
-                });
+                let followups = self.lsp_state.oak_scheduler.apply_scan_completed(
+                    &mut self.world.db,
+                    scan,
+                    &editor_owned,
+                );
                 lsp::log_info!("Dispatching {n} followup scan requests", n = followups.len());
                 dispatch_scan_requests(&self.events_tx, followups);
 
@@ -849,24 +892,6 @@ pub(crate) fn log(level: lsp_types::MessageType, message: String) {
     };
 }
 
-/// Run a blocking Salsa write without stranding the runtime.
-///
-/// Salsa writes may block the current thread until it gains exclusive access to
-/// the DB. This requires processing/cancelling background tasks that hold DB
-/// clones, e.g. diagnostics tasks. If the dispatching of these tasks lives in
-/// tokio tasks, there is a risk of deadlock. We call `block_in_place()` to let
-/// Tokio know the thread is about to potentially block, allowing tasks to be
-/// scheduled on other worker threads.
-///
-/// `block_in_place()` panics on a current-thread runtime, which the tests use, so
-/// fall back to calling `f` directly there.
-fn block_for_write<T>(f: impl FnOnce() -> T) -> T {
-    match Handle::try_current().map(|handle| handle.runtime_flavor()) {
-        Ok(RuntimeFlavor::MultiThread) => tokio::task::block_in_place(f),
-        _ => f(),
-    }
-}
-
 /// Spawn a blocking task
 ///
 /// This runs tasks that do semantic analysis on a separate thread pool to avoid