feat(rome_lsp): stop the daemon after the last client disconnects (#3643

)

* feat(rome_lsp): shutdown the automatically spawned daemon instance once the last client disconnect

* stop the server immediately if at least one session was initialized

* address PR review

crates/rome_cli/src/commands/daemon.rs

@@ -21,7 +21,7 @@ use crate::{
 
 pub(crate) fn start(mut session: CliSession) -> Result<(), Termination> {
     let rt = Runtime::new()?;
-    let did_spawn = rt.block_on(ensure_daemon())?;
+    let did_spawn = rt.block_on(ensure_daemon(false))?;
 
     if did_spawn {
         session.app.console.log(markup! {
@@ -60,11 +60,13 @@ pub(crate) fn stop(mut session: CliSession) -> Result<(), Termination> {
     Ok(())
 }
 
-pub(crate) fn run_server() -> Result<(), Termination> {
+pub(crate) fn run_server(mut session: CliSession) -> Result<(), Termination> {
     setup_tracing_subscriber();
 
+    let stop_on_disconnect = session.args.contains("--stop-on-disconnect");
+
     let rt = Runtime::new()?;
-    let factory = ServerFactory::default();
+    let factory = ServerFactory::new(stop_on_disconnect);
     let cancellation = factory.cancellation();
     let span = debug_span!("Running Server", pid = std::process::id());
 
@@ -101,7 +103,7 @@ pub(crate) fn lsp_proxy() -> Result<(), Termination> {
 /// Receives a process via `stdin` and then copy the content to the LSP socket.
 /// Copy to the process on `stdout` when the LSP responds to a message
 async fn start_lsp_proxy(rt: &Runtime) -> Result<(), Termination> {
-    ensure_daemon().await?;
+    ensure_daemon(true).await?;
 
     match open_socket().await? {
         Some((mut owned_read_half, mut owned_write_half)) => {

crates/rome_cli/src/lib.rs

@@ -93,7 +93,7 @@ impl<'app> CliSession<'app> {
             Some("lsp-proxy") => commands::daemon::lsp_proxy(),
 
             // Internal commands
-            Some("__run_server") => commands::daemon::run_server(),
+            Some("__run_server") => commands::daemon::run_server(self),
             Some("__print_socket") => commands::daemon::print_socket(),
 
             // Print the help for known commands called without any arguments, and exit with an error

crates/rome_cli/src/service/unix.rs

@@ -51,12 +51,16 @@ async fn try_connect() -> io::Result<UnixStream> {
 }
 
 /// Spawn the daemon server process in the background
-fn spawn_daemon() -> io::Result<Child> {
+fn spawn_daemon(stop_on_disconnect: bool) -> io::Result<Child> {
     let binary = env::current_exe()?;
 
     let mut cmd = Command::new(binary);
     cmd.arg("__run_server");
 
+    if stop_on_disconnect {
+        cmd.arg("--stop-on-disconnect");
+    }
+
     // Create a new session for the process and make it the leader, this will
     // ensures that the child process is fully detached from its parent and will
     // continue running in the background even after the parent process exits
@@ -102,7 +106,7 @@ pub(crate) async fn open_socket() -> io::Result<Option<(OwnedReadHalf, OwnedWrit
 ///
 /// Returns false if the daemon process was already running or true if it had
 /// to be started
-pub(crate) async fn ensure_daemon() -> io::Result<bool> {
+pub(crate) async fn ensure_daemon(stop_on_disconnect: bool) -> io::Result<bool> {
     let mut current_child: Option<Child> = None;
     let mut last_error = None;
 
@@ -140,7 +144,7 @@ pub(crate) async fn ensure_daemon() -> io::Result<bool> {
                 } else {
                     // Spawn the daemon process and wait a few milliseconds for
                     // it to become ready then retry the connection
-                    current_child = Some(spawn_daemon()?);
+                    current_child = Some(spawn_daemon(stop_on_disconnect)?);
                     time::sleep(Duration::from_millis(50)).await;
                 }
             }
@@ -162,7 +166,7 @@ pub(crate) async fn ensure_daemon() -> io::Result<bool> {
 /// Ensure the server daemon is running and ready to receive connections and
 /// print the global socket name in the standard output
 pub(crate) async fn print_socket() -> io::Result<()> {
-    ensure_daemon().await?;
+    ensure_daemon(true).await?;
     println!("{}", get_socket_name().display());
     Ok(())
 }

crates/rome_cli/src/service/windows.rs

@@ -67,12 +67,16 @@ async fn try_connect() -> io::Result<NamedPipeClient> {
 const CREATE_NEW_PROCESS_GROUP: u32 = 0x00000200;
 
 /// Spawn the daemon server process in the background
-fn spawn_daemon() -> io::Result<()> {
+fn spawn_daemon(stop_on_disconnect: bool) -> io::Result<()> {
     let binary = env::current_exe()?;
 
     let mut cmd = Command::new(binary);
     cmd.arg("__run_server");
 
+    if stop_on_disconnect {
+        cmd.arg("--stop-on-disconnect");
+    }
+
     cmd.creation_flags(CREATE_NEW_PROCESS_GROUP);
 
     cmd.spawn()?;
@@ -164,14 +168,14 @@ impl AsyncWrite for ClientWriteHalf {
 ///
 /// Returns false if the daemon process was already running or true if it had
 /// to be started
-pub(crate) async fn ensure_daemon() -> io::Result<bool> {
+pub(crate) async fn ensure_daemon(stop_on_disconnect: bool) -> io::Result<bool> {
     let mut did_spawn = false;
 
     loop {
         match open_socket().await {
             Ok(Some(_)) => break,
             Ok(None) => {
-                spawn_daemon()?;
+                spawn_daemon(stop_on_disconnect)?;
                 did_spawn = true;
                 time::sleep(Duration::from_millis(50)).await;
             }
@@ -185,7 +189,7 @@ pub(crate) async fn ensure_daemon() -> io::Result<bool> {
 /// Ensure the server daemon is running and ready to receive connections and
 /// print the global pipe name in the standard output
 pub(crate) async fn print_socket() -> io::Result<()> {
-    ensure_daemon().await?;
+    ensure_daemon(true).await?;
     println!("{}", get_pipe_name());
     Ok(())
 }

crates/rome_lsp/src/server.rs

@@ -10,7 +10,7 @@ use rome_fs::CONFIG_NAME;
 use rome_service::workspace::{RageEntry, RageParams, RageResult};
 use rome_service::{workspace, Workspace};
 use std::collections::HashMap;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex};
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio::sync::Notify;
@@ -24,11 +24,27 @@ pub struct LSPServer {
     session: SessionHandle,
     /// Map of all sessions connected to the same [ServerFactory] as this [LSPServer].
     sessions: Sessions,
+    /// If this is true the server will broadcast a shutdown signal once the
+    /// last client disconnected
+    stop_on_disconnect: bool,
+    /// This shared flag is set to true once at least one sessions has been
+    /// initialized on this server instance
+    is_initialized: Arc<AtomicBool>,
 }
 
 impl LSPServer {
-    fn new(session: SessionHandle, sessions: Sessions) -> Self {
-        Self { session, sessions }
+    fn new(
+        session: SessionHandle,
+        sessions: Sessions,
+        stop_on_disconnect: bool,
+        is_initialized: Arc<AtomicBool>,
+    ) -> Self {
+        Self {
+            session,
+            sessions,
+            stop_on_disconnect,
+            is_initialized,
+        }
     }
 
     async fn syntax_tree_request(&self, params: SyntaxTreePayload) -> LspResult<String> {
@@ -172,6 +188,7 @@ impl LanguageServer for LSPServer {
     #[tracing::instrument(level = "trace", skip(self))]
     async fn initialize(&self, params: InitializeParams) -> LspResult<InitializeResult> {
         info!("Starting Rome Language Server...");
+        self.is_initialized.store(true, Ordering::Relaxed);
 
         self.session
             .client_capabilities
@@ -329,8 +346,14 @@ impl Drop for LSPServer {
     fn drop(&mut self) {
         if let Ok(mut sessions) = self.sessions.lock() {
             let _removed = sessions.remove(&self.session.key);
-
             debug_assert!(_removed.is_some(), "Session did not exist.");
+
+            if self.stop_on_disconnect
+                && sessions.is_empty()
+                && self.is_initialized.load(Ordering::Relaxed)
+            {
+                self.session.cancellation.notify_one();
+            }
         }
     }
 }
@@ -356,6 +379,13 @@ pub struct ServerFactory {
 
     /// Session key generator. Stores the key of the next session.
     next_session_key: AtomicU64,
+
+    /// If this is true the server will broadcast a shutdown signal once the
+    /// last client disconnected
+    stop_on_disconnect: bool,
+    /// This shared flag is set to true once at least one sessions has been
+    /// initialized on this server instance
+    is_initialized: Arc<AtomicBool>,
 }
 
 /// Helper method for wrapping a [Workspace] method in a `custom_method` for
@@ -393,6 +423,17 @@ macro_rules! workspace_method {
 }
 
 impl ServerFactory {
+    pub fn new(stop_on_disconnect: bool) -> Self {
+        Self {
+            cancellation: Arc::default(),
+            workspace: None,
+            sessions: Sessions::default(),
+            next_session_key: AtomicU64::new(0),
+            stop_on_disconnect,
+            is_initialized: Arc::default(),
+        }
+    }
+
     /// Create a new [ServerConnection] from this factory
     pub fn create(&self) -> ServerConnection {
         let workspace = self
@@ -408,7 +449,13 @@ impl ServerFactory {
 
             let mut sessions = self.sessions.lock().unwrap();
             sessions.insert(session_key, handle.clone());
-            LSPServer::new(handle, self.sessions.clone())
+
+            LSPServer::new(
+                handle,
+                self.sessions.clone(),
+                self.stop_on_disconnect,
+                self.is_initialized.clone(),
+            )
         });
 
         builder = builder.custom_method(SYNTAX_TREE_REQUEST, LSPServer::syntax_tree_request);

crates/rome_lsp/src/session.rs

@@ -63,7 +63,7 @@ pub(crate) struct Session {
     documents: RwLock<HashMap<lsp_types::Url, Document>>,
     url_interner: RwLock<UrlInterner>,
 
-    cancellation: Arc<Notify>,
+    pub(crate) cancellation: Arc<Notify>,
 }
 
 pub(crate) type SessionHandle = Arc<Session>;

editors/vscode/README.md

@@ -84,20 +84,6 @@ Rome doesn’t yet support loading the `rome.json` file from a directory other t
 
 Rome isn't yet able to pick up the `rome.json` configuration in [multi-root workspaces](https://code.visualstudio.com/docs/editor/multi-root-workspaces) if the configuration isn't in the first root folder ([issue 3538](https://github.com/rome/tools/issues/3538)). You can work around this limitation by making the folder with the configuration the first root folder of the workspace (drag it to the top).
 
-### The extension uses an outdated version of Rome
-
-Make sure to restart Rome’s LSP server after updating the extension. To stop Rome’s LSP, send the `stop` command (VS will start a new LSP instance for you):
-
-```bash
-npx rome stop
-
-# or
-pnpx exec rome stop
-
-# or
-yarn run rome stop
-```
-
 ### Disable Rome for workspaces without a `rome.json` configuration
 
 Rome's VS Code extension is active for every workspace regardless if the workspace contains a `rome.json` configuration ([issue 3506](https://github.com/rome/tools/issues/3506)). That may be surprising to you if you use other extensions like ESLint where the extension is disabled if the configuration is missing. This behavior is intentional because it's Rome's philosophy that the configuration should be optional.