Support cancellation requests (#18627)

Author: MichaReiser

crates/ruff/src/commands/server.rs

@@ -1,14 +1,7 @@
-use std::num::NonZeroUsize;
-
 use crate::ExitStatus;
 use anyhow::Result;
-use ruff_server::Server;
-
-pub(crate) fn run_server(
-    worker_threads: NonZeroUsize,
-    preview: Option<bool>,
-) -> Result<ExitStatus> {
-    let server = Server::new(worker_threads, preview)?;
 
-    server.run().map(|()| ExitStatus::Success)
+pub(crate) fn run_server(preview: Option<bool>) -> Result<ExitStatus> {
+    ruff_server::run(preview)?;
+    Ok(ExitStatus::Success)
 }

crates/ruff/src/lib.rs

@@ -2,7 +2,6 @@
 
 use std::fs::File;
 use std::io::{self, BufWriter, Write, stdout};
-use std::num::NonZeroUsize;
 use std::path::{Path, PathBuf};
 use std::process::ExitCode;
 use std::sync::mpsc::channel;
@@ -223,13 +222,7 @@ fn analyze_graph(
 }
 
 fn server(args: ServerCommand) -> Result<ExitStatus> {
-    let four = NonZeroUsize::new(4).unwrap();
-
-    // by default, we set the number of worker threads to `num_cpus`, with a maximum of 4.
-    let worker_threads = std::thread::available_parallelism()
-        .unwrap_or(four)
-        .min(four);
-    commands::server::run_server(worker_threads, args.resolve_preview())
+    commands::server::run_server(args.resolve_preview())
 }
 
 pub fn check(args: CheckCommand, global_options: GlobalConfigArgs) -> Result<ExitStatus> {

crates/ruff_server/src/lib.rs

@@ -1,13 +1,15 @@
 //! ## The Ruff Language Server
 
+use std::num::NonZeroUsize;
+
+use anyhow::Context as _;
 pub use edit::{DocumentKey, NotebookDocument, PositionEncoding, TextDocument};
 use lsp_types::CodeActionKind;
-pub use server::Server;
-pub use session::{ClientOptions, DocumentQuery, DocumentSnapshot, GlobalOptions, Session};
+pub use server::{ConnectionSender, MainLoopSender, Server};
+pub use session::{Client, ClientOptions, DocumentQuery, DocumentSnapshot, GlobalOptions, Session};
 pub use workspace::{Workspace, Workspaces};
 
-#[macro_use]
-mod message;
+use crate::server::ConnectionInitializer;
 
 mod edit;
 mod fix;
@@ -37,3 +39,35 @@ pub(crate) type Result<T> = anyhow::Result<T>;
 pub(crate) fn version() -> &'static str {
     ruff_linter::VERSION
 }
+
+pub fn run(preview: Option<bool>) -> Result<()> {
+    let four = NonZeroUsize::new(4).unwrap();
+
+    // by default, we set the number of worker threads to `num_cpus`, with a maximum of 4.
+    let worker_threads = std::thread::available_parallelism()
+        .unwrap_or(four)
+        .min(four);
+
+    let (connection, io_threads) = ConnectionInitializer::stdio();
+
+    let server_result = Server::new(worker_threads, connection, preview)
+        .context("Failed to start server")?
+        .run();
+
+    let io_result = io_threads.join();
+
+    let result = match (server_result, io_result) {
+        (Ok(()), Ok(())) => Ok(()),
+        (Err(server), Err(io)) => Err(server).context(format!("IO thread error: {io}")),
+        (Err(server), _) => Err(server),
+        (_, Err(io)) => Err(io).context("IO thread error"),
+    };
+
+    if let Err(err) = result.as_ref() {
+        tracing::warn!("Server shut down with an error: {err}");
+    } else {
+        tracing::info!("Server shut down");
+    }
+
+    result
+}