ruff server no longer hangs after shutdown (#11222)

## Summary

Fixes https://github.com/astral-sh/ruff/issues/11207.

The server would hang after handling a shutdown request on
`IoThreads::join()` because a global sender (`MESSENGER`, used to send
`window/showMessage` notifications) would remain allocated even after
the event loop finished, which kept the writer I/O thread channel open.

To fix this, I've made a few structural changes to `ruff server`. I've
wrapped the send/receive channels and thread join handle behind a new
struct, `Connection`, which facilitates message sending and receiving,
and also runs `IoThreads::join()` after the event loop finishes. To
control the number of sender channels, the `Connection` wraps the sender
channel in an `Arc` and only allows the creation of a wrapper type,
`ClientSender`, which hold a weak reference to this `Arc` instead of
direct channel access. The wrapper type implements the channel methods
directly to prevent access to the inner channel (which would allow the
channel to be cloned). ClientSender's function is analogous to
[`WeakSender` in
`tokio`](https://docs.rs/tokio/latest/tokio/sync/mpsc/struct.WeakSender.html).
Additionally, the receiver channel cannot be accessed directly - the
`Connection` only exposes an iterator over it.

These changes will guarantee that all channels are closed before the I/O
threads are joined.

## Test Plan

Repeatedly open and close an editor utilizing `ruff server` while
observing the task monitor. The net total amount of open `ruff`
instances should be zero once all editor windows have closed.

The following logs should also appear after the server is shut down:

<img width="835" alt="Screenshot 2024-04-30 at 3 56 22 PM"
src="404b74f5-ef08-4bb4-9fa2-72e72b946695">

This can be tested on VS Code by changing the settings and then checking
`Output`.

crates/ruff_server/src/server.rs

@@ -2,7 +2,6 @@
 
 use std::num::NonZeroUsize;
 
-use lsp::Connection;
 use lsp_server as lsp;
 use lsp_types as types;
 use types::ClientCapabilities;
@@ -18,6 +17,8 @@ use types::TextDocumentSyncOptions;
 use types::WorkDoneProgressOptions;
 use types::WorkspaceFoldersServerCapabilities;
 
+use self::connection::Connection;
+use self::connection::ConnectionInitializer;
 use self::schedule::event_loop_thread;
 use self::schedule::Scheduler;
 use self::schedule::Task;
@@ -28,34 +29,39 @@ use crate::PositionEncoding;
 
 mod api;
 mod client;
+mod connection;
 mod schedule;
 
-pub(crate) use client::ClientSender;
+pub(crate) use connection::ClientSender;
 
 pub(crate) type Result<T> = std::result::Result<T, api::Error>;
 
 pub struct Server {
-    conn: lsp::Connection,
+    connection: Connection,
     client_capabilities: ClientCapabilities,
-    threads: lsp::IoThreads,
     worker_threads: NonZeroUsize,
     session: Session,
 }
 
 impl Server {
     pub fn new(worker_threads: NonZeroUsize) -> crate::Result<Self> {
-        let (conn, threads) = lsp::Connection::stdio();
+        let connection = ConnectionInitializer::stdio();
 
-        crate::message::init_messenger(&conn.sender);
-
-        let (id, params) = conn.initialize_start()?;
-
-        let init_params: types::InitializeParams = serde_json::from_value(params)?;
+        let (id, init_params) = connection.initialize_start()?;
 
         let client_capabilities = init_params.capabilities;
         let position_encoding = Self::find_best_position_encoding(&client_capabilities);
         let server_capabilities = Self::server_capabilities(position_encoding);
 
+        let connection = connection.initialize_finish(
+            id,
+            &server_capabilities,
+            crate::SERVER_NAME,
+            crate::version(),
+        )?;
+
+        crate::message::init_messenger(connection.make_sender());
+
         let AllSettings {
             global_settings,
             mut workspace_settings,
@@ -86,19 +92,8 @@ impl Server {
                 anyhow::anyhow!("Failed to get the current working directory while creating a default workspace.")
             })?;
 
-        let initialize_data = serde_json::json!({
-            "capabilities": server_capabilities,
-            "serverInfo": {
-                "name": crate::SERVER_NAME,
-                "version": crate::version()
-            }
-        });
-
-        conn.initialize_finish(id, initialize_data)?;
-
         Ok(Self {
-            conn,
-            threads,
+            connection,
             worker_threads,
             session: Session::new(
                 &client_capabilities,
@@ -111,17 +106,20 @@ impl Server {
     }
 
     pub fn run(self) -> crate::Result<()> {
-        let result = event_loop_thread(move || {
+        event_loop_thread(move || {
             Self::event_loop(
-                &self.conn,
+                &self.connection,
                 &self.client_capabilities,
                 self.session,
                 self.worker_threads,
-            )
+            )?;
+            self.connection.close()?;
+            // Note: when we start routing tracing through the LSP,
+            // this should be replaced with a log directly to `stderr`.
+            tracing::info!("Server has shut down successfully");
+            Ok(())
         })?
-        .join();
-        self.threads.join()?;
-        result
+        .join()
     }
 
     #[allow(clippy::needless_pass_by_value)] // this is because we aren't using `next_request_id` yet.
@@ -132,22 +130,21 @@ impl Server {
         worker_threads: NonZeroUsize,
     ) -> crate::Result<()> {
         let mut scheduler =
-            schedule::Scheduler::new(&mut session, worker_threads, &connection.sender);
+            schedule::Scheduler::new(&mut session, worker_threads, connection.make_sender());
 
         Self::try_register_capabilities(client_capabilities, &mut scheduler);
-        for msg in &connection.receiver {
+        for msg in connection.incoming() {
+            if connection.handle_shutdown(&msg)? {
+                break;
+            }
             let task = match msg {
-                lsp::Message::Request(req) => {
-                    if connection.handle_shutdown(&req)? {
-                        return Ok(());
-                    }
-                    api::request(req)
-                }
+                lsp::Message::Request(req) => api::request(req),
                 lsp::Message::Notification(notification) => api::notification(notification),
                 lsp::Message::Response(response) => scheduler.response(response),
             };
             scheduler.dispatch(task);
         }
+
         Ok(())
     }