Shut down help proxy when help comm drops (#1285)

Branched from #1284 

While working on that PR I noticed that we never shut down the proxy
server when the Help comm is tore down. If we start a new one, we'd pile
up a new proxy.

This is theoretical because Help is one of the comms that the frontend
reuses on refresh/reload (unlike the UI comm). But it still seems nice
to handle the shutdown path properly. Covered by a new test.

Author: lionel-

crates/ark/src/help/r_help.rs

@@ -42,11 +42,15 @@ pub struct HelpPorts {
 /// The R Help handler (together with the help proxy) provides the server side
 /// of Positron's Help panel.
 ///
-/// A stateless unit struct. The server and proxy ports live on `Console` as
-/// `help_ports` so they can be reached via `Console::get()` by the
-/// `browseURL()` hook.
-#[derive(Debug)]
-pub struct RHelp;
+/// The server and proxy ports live on `Console` as `help_ports`, not here, so
+/// the `browseURL()` hook can reach them through `Console::get()`. What this
+/// handler does own is the proxy drop guard, which shuts the proxy down when
+/// the help comm closes.
+#[derive(Debug, Default)]
+pub struct RHelp {
+    /// Drop guard to stop the help proxy server on teardown.
+    proxy: Option<help_proxy::ProxyHandle>,
+}
 
 impl RHelp {
     /// Public associated function so that callers can cheaply check if a url is
@@ -230,14 +234,15 @@ impl CommHandler for RHelp {
         };
         log::info!("R help server listening on port {r_port}");
 
-        let proxy_port = match help_proxy::start(r_port) {
-            Ok(port) => port,
+        let (proxy_port, proxy) = match help_proxy::start(r_port) {
+            Ok(proxy) => proxy,
             Err(err) => {
                 log::error!("Could not start R help proxy server: {err:?}");
                 return;
             },
         };
 
+        self.proxy = Some(proxy);
         Console::get_mut().set_help_ports(r_port, proxy_port);
     }
 

crates/ark/src/help_proxy.rs

@@ -44,9 +44,21 @@ struct AppState {
     target_port: u16,
 }
 
+/// Drop guard that ties the proxy server's lifetime to the help comm.
+///
+/// Dropping this stops the `actix_web` server and lets its thread and tokio
+/// runtime exit. We hold it in `RHelp` so the proxy lives exactly as long as
+/// the help comm. The field is never read. Dropping the `Sender` is the whole
+/// mechanism, because that's what the proxy thread waits on.
+#[derive(Debug)]
+pub struct ProxyHandle {
+    _shutdown_tx: tokio::sync::oneshot::Sender<()>,
+}
+
 // Starts the help proxy.
-pub fn start(target_port: u16) -> anyhow::Result<u16> {
+pub fn start(target_port: u16) -> anyhow::Result<(u16, ProxyHandle)> {
     let (port_tx, port_rx) = crossbeam::channel::bounded::<u16>(1);
+    let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
 
     spawn!("ark-help-proxy", move || -> anyhow::Result<()> {
         // Bind to port `0` to allow the OS to assign the port, avoiding any race conditions
@@ -85,9 +97,26 @@ pub fn start(target_port: u16) -> anyhow::Result<u16> {
 
         // Execute the task within the runtime.
         rt.block_on(async {
-            match server.run().await {
-                Ok(value) => log::info!("Help proxy server exited with value: {:?}", value),
-                Err(error) => log::error!("Help proxy server exited unexpectedly: {}", error),
+            let server = server.run();
+            let handle = server.handle();
+
+            // Dropping the help comm must tear down the whole proxy and its
+            // runtime, not just stop the server. The trigger is `ProxyHandle`
+            // (this oneshot's sender) dropping with `RHelp`, which resolves
+            // `shutdown_rx`. The explicit `stop()` is required because dropping
+            // the `ServerHandle` can't stop the server (the `Server` future
+            // owns its own handle). `stop()` lets the `server.await` call
+            // return, allowing the task and runtime to exit.
+            tokio::spawn(async move {
+                // We never send on the sender, so this only resolves (to `Err`)
+                // when `ProxyHandle` drops. Either way it means "shut down".
+                shutdown_rx.await.ok();
+                handle.stop(false).await;
+            });
+
+            match server.await {
+                Ok(value) => log::info!("Help proxy server exited with value: {value:?}"),
+                Err(error) => log::error!("Help proxy server exited unexpectedly: {error}"),
             }
         });
 
@@ -96,7 +125,9 @@ pub fn start(target_port: u16) -> anyhow::Result<u16> {
 
     // Wait for the returned port with an extensive timeout
     match port_rx.recv_timeout(Duration::from_secs(20)) {
-        Ok(port) => Ok(port),
+        Ok(port) => Ok((port, ProxyHandle {
+            _shutdown_tx: shutdown_tx,
+        })),
         Err(err) => Err(anyhow::anyhow!(
             "Help proxy server timed out while waiting for a port: {err:?}"
         )),

crates/ark/src/shell.rs

@@ -347,14 +347,15 @@ fn handle_comm_open_help(
 ) -> amalthea::Result<(bool, Option<Receiver<()>>)> {
     // Register the handler on the R thread, like the UI comm. `RHelp::handle_open`
     // starts the R help server and proxy and records their ports on `Console`, all
-    // on the R thread. The RPC handler itself is stateless.
+    // on the R thread. The handler holds the proxy's drop guard, so the proxy is
+    // torn down when the comm is removed.
     let (done_tx, done_rx) = bounded(0);
     kernel_request_tx
         .send(KernelRequest::CommOpen {
             comm_id: comm.comm_id.clone(),
             comm_name: comm.comm_name.clone(),
             outgoing_tx: comm.outgoing_tx.clone(),
-            handler: Box::new(RHelp),
+            handler: Box::new(RHelp::default()),
             done_tx,
         })
         .map_err(|err| amalthea::Error::SendError(err.to_string()))?;

crates/ark/tests/integration/help.rs

@@ -6,6 +6,8 @@
 //
 
 use core::panic;
+use std::net::TcpStream;
+use std::time::Duration;
 
 use amalthea::comm::comm_channel::CommMsg;
 use amalthea::comm::event::CommEvent;
@@ -68,7 +70,7 @@ impl TestRHelp {
             let (comm_event_tx, _) = bounded::<CommEvent>(10);
             let ctx = CommHandlerContext::new(outgoing_tx, comm_event_tx);
 
-            let mut handler = RHelp;
+            let mut handler = RHelp::default();
             handler.handle_msg(msg, &ctx);
         });
 
@@ -199,3 +201,88 @@ fn test_help_show_help_event() {
     frontend.recv_iopub_idle();
     frontend.recv_shell_execute_reply();
 }
+
+/// Reopening the help comm (as happens on a frontend reload) must not leak the
+/// previous proxy server.
+///
+/// The proxy's lifetime is tied to the help comm via a drop guard held in
+/// `RHelp`. When the frontend opens a second `positron.help` comm, the kernel
+/// replaces the handler, the old `RHelp` drops, and its proxy stops. We observe
+/// this directly. Each proxy binds its own localhost port, so we connect to the
+/// old port and check it stops accepting connections. Before this was wired up,
+/// the old proxy stayed bound forever.
+#[test]
+fn test_help_proxy_torn_down_on_reopen() {
+    let frontend = DummyArkFrontend::lock();
+
+    let comm_id_1 = open_help_comm(&frontend);
+    let port_1 = show_help_and_get_proxy_port(&frontend, &comm_id_1);
+    assert!(proxy_is_listening(port_1));
+
+    // Reopen the help comm, which replaces the handler and drops the old one.
+    let comm_id_2 = open_help_comm(&frontend);
+    let port_2 = show_help_and_get_proxy_port(&frontend, &comm_id_2);
+
+    // The new proxy binds before the old `RHelp` is dropped, so the OS hands out
+    // a different port and we can tell the two apart.
+    assert_ne!(port_1, port_2);
+    assert!(proxy_is_listening(port_2));
+
+    // Teardown is asynchronous (a task inside the proxy's runtime calls
+    // `stop()`), so poll until the old port refuses connections.
+    wait_until_proxy_stops(port_1);
+}
+
+fn open_help_comm(frontend: &DummyArkFrontend) -> String {
+    let comm_id = uuid::Uuid::new_v4().to_string();
+    frontend.send_shell(CommOpen {
+        comm_id: comm_id.clone(),
+        target_name: String::from("positron.help"),
+        data: serde_json::json!({}),
+    });
+    frontend.recv_iopub_busy();
+    frontend.recv_iopub_idle();
+    comm_id
+}
+
+/// Request a help topic and return the proxy port from the `show_help` URL.
+fn show_help_and_get_proxy_port(frontend: &DummyArkFrontend, comm_id: &str) -> u16 {
+    frontend.send_execute_request("?plot", ExecuteRequestOptions::default());
+    frontend.recv_iopub_busy();
+    frontend.recv_iopub_execute_input();
+
+    let msg = frontend.recv_iopub_comm_msg();
+    assert_eq!(msg.comm_id, comm_id);
+    assert_eq!(
+        msg.data.get("method").and_then(|v| v.as_str()),
+        Some("show_help")
+    );
+    let content = msg.data["params"]["content"].as_str().unwrap();
+    let port = proxy_port_from_url(content);
+
+    frontend.recv_iopub_idle();
+    frontend.recv_shell_execute_reply();
+
+    port
+}
+
+/// Pull the port out of a proxy URL like `http://127.0.0.1:<port>/...`.
+fn proxy_port_from_url(url: &str) -> u16 {
+    let after_host = url.strip_prefix("http://127.0.0.1:").unwrap();
+    let end = after_host.find('/').unwrap();
+    after_host[..end].parse().unwrap()
+}
+
+fn proxy_is_listening(port: u16) -> bool {
+    TcpStream::connect(("127.0.0.1", port)).is_ok()
+}
+
+fn wait_until_proxy_stops(port: u16) {
+    for _ in 0..100 {
+        if !proxy_is_listening(port) {
+            return;
+        }
+        std::thread::sleep(Duration::from_millis(50));
+    }
+    panic!("Proxy on port {port} is still accepting connections after teardown");
+}