fix: close daemon runtime cleanly

crates/geth-node/src/daemon.rs

@@ -7,35 +7,77 @@ use crate::{
 use geth_config::{GethConfig, GethPaths, RelayMode};
 use geth_iroh::{EndpointStatus, GethIrohConfig, GethIrohEndpoint, GethRelayMode};
 use geth_store::Store;
-use std::{path::Path, time::Duration};
+use std::time::Duration;
 use tokio::net::UnixListener;
+use tokio::task::JoinHandle;
 
 pub async fn run_daemon(paths: GethPaths) -> Result<(), NodeError> {
     let mut node = init_node(&paths)?;
-    let _iroh_endpoint = start_daemon_iroh_endpoint(&mut node).await?;
-    if let Some(endpoint) = node
-        .iroh_endpoint
-        .lock()
-        .map_err(|_| NodeError::RuntimeLockPoisoned)?
-        .clone()
-    {
-        spawn_iroh_control_accept_loop(node.clone(), endpoint);
-        let config = GethConfig::load(&node.paths.config_file())?;
+    let config = GethConfig::load(&node.paths.config_file())?;
+    let iroh_endpoint = start_daemon_iroh_endpoint(&mut node).await?;
+    let mut background_tasks = Vec::new();
+    if let Some(endpoint) = iroh_endpoint.clone() {
+        background_tasks.push(spawn_iroh_control_accept_loop(node.clone(), endpoint));
         if config.sync.live_sync_enabled {
-            spawn_background_live_sync(
+            background_tasks.push(spawn_background_live_sync(
                 node.clone(),
                 Duration::from_millis(config.sync.live_sync_interval_ms),
-            );
+            ));
         }
     }
-    let _peer_card_lan_discovery = start_peer_card_lan_discovery(&node).await;
-    if Path::new(&paths.socket_path()).exists() {
-        std::fs::remove_file(paths.socket_path())?;
-    }
-    let listener = UnixListener::bind(paths.socket_path())?;
-    tracing::info!(socket = %paths.socket_path().display(), "geth daemon listening");
 
-    serve_local_control(node, listener).await
+    let serve_result = async {
+        let _peer_card_lan_discovery = start_peer_card_lan_discovery(&node).await;
+        remove_control_socket(&paths)?;
+        let listener = UnixListener::bind(paths.socket_path())?;
+        tracing::info!(socket = %paths.socket_path().display(), "geth daemon listening");
+        serve_local_control(node.clone(), listener).await
+    }
+    .await;
+
+    for task in background_tasks {
+        task.abort();
+        let _ = task.await;
+    }
+    let shutdown_result = shutdown_iroh_runtime(&node, iroh_endpoint).await;
+    let socket_cleanup_result = remove_control_socket(&paths);
+    serve_result?;
+    shutdown_result?;
+    socket_cleanup_result
+}
+
+fn remove_control_socket(paths: &GethPaths) -> Result<(), NodeError> {
+    match std::fs::remove_file(paths.socket_path()) {
+        Ok(()) => Ok(()),
+        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
+        Err(error) => Err(error.into()),
+    }
+}
+
+async fn shutdown_iroh_runtime(
+    node: &LocalNode,
+    endpoint: Option<GethIrohEndpoint>,
+) -> Result<(), NodeError> {
+    node.iroh_docs
+        .lock()
+        .map_err(|_| NodeError::RuntimeLockPoisoned)?
+        .take();
+    node.iroh_gossip
+        .lock()
+        .map_err(|_| NodeError::RuntimeLockPoisoned)?
+        .take();
+    node.iroh_blob_store
+        .lock()
+        .map_err(|_| NodeError::RuntimeLockPoisoned)?
+        .take();
+    node.iroh_endpoint
+        .lock()
+        .map_err(|_| NodeError::RuntimeLockPoisoned)?
+        .take();
+    if let Some(endpoint) = endpoint {
+        endpoint.shutdown().await;
+    }
+    Ok(())
 }
 
 async fn serve_local_control(node: LocalNode, listener: UnixListener) -> Result<(), NodeError> {
@@ -59,7 +101,10 @@ async fn serve_local_control(node: LocalNode, listener: UnixListener) -> Result<
     }
 }
 
-pub(crate) fn spawn_iroh_control_accept_loop(node: LocalNode, endpoint: GethIrohEndpoint) {
+pub(crate) fn spawn_iroh_control_accept_loop(
+    node: LocalNode,
+    endpoint: GethIrohEndpoint,
+) -> JoinHandle<()> {
     let raw_endpoint = endpoint.endpoint();
     tokio::spawn(async move {
         tracing::debug!("iroh accept loop started");
@@ -73,10 +118,13 @@ pub(crate) fn spawn_iroh_control_accept_loop(node: LocalNode, endpoint: GethIroh
             });
         }
         tracing::debug!("iroh accept loop ended");
-    });
+    })
 }
 
-pub(crate) fn spawn_background_live_sync(node: LocalNode, interval_duration: Duration) {
+pub(crate) fn spawn_background_live_sync(
+    node: LocalNode,
+    interval_duration: Duration,
+) -> JoinHandle<()> {
     tokio::spawn(async move {
         if let Err(error) = run_live_sync_once(&node).await {
             tracing::debug!(%error, "initial live sync tick failed");
@@ -88,7 +136,7 @@ pub(crate) fn spawn_background_live_sync(node: LocalNode, interval_duration: Dur
                 tracing::debug!(%error, "live sync tick failed");
             }
         }
-    });
+    })
 }
 
 pub(crate) async fn start_daemon_iroh_endpoint(
@@ -186,3 +234,61 @@ fn config_relay_mode_to_iroh(
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn control_socket_cleanup_is_idempotent() {
+        let dir = tempfile::tempdir().expect("tempdir");
+        let paths = GethPaths::from_home(dir.path());
+        paths.ensure_base_dirs().expect("base dirs");
+        std::fs::write(paths.socket_path(), b"stale").expect("stale socket fixture");
+
+        remove_control_socket(&paths).expect("remove existing socket");
+        remove_control_socket(&paths).expect("repeat socket cleanup");
+
+        assert!(!paths.socket_path().exists());
+    }
+
+    #[tokio::test]
+    async fn daemon_shutdown_closes_and_releases_iroh_runtime() {
+        if std::env::var_os("GETH_TEST_SKIP_IROH").is_some() {
+            eprintln!("skipping Iroh shutdown test because GETH_TEST_SKIP_IROH is set");
+            return;
+        }
+
+        let dir = tempfile::tempdir().expect("tempdir");
+        let paths = GethPaths::from_home(dir.path());
+        let mut node = init_node(&paths).expect("init node");
+        std::fs::write(
+            paths.config_file(),
+            "[iroh]\nrelay_mode = \"disabled\"\nlocal_discovery = false\n",
+        )
+        .expect("offline config");
+        let Some(endpoint) = start_daemon_iroh_endpoint(&mut node)
+            .await
+            .expect("start daemon endpoint")
+        else {
+            eprintln!("skipping Iroh shutdown assertion because UDP bind is unavailable");
+            return;
+        };
+        let raw_endpoint = endpoint.endpoint();
+
+        shutdown_iroh_runtime(&node, Some(endpoint))
+            .await
+            .expect("shutdown daemon endpoint");
+
+        assert!(node.iroh_endpoint.lock().expect("endpoint lock").is_none());
+        assert!(node.iroh_blob_store.lock().expect("blob lock").is_none());
+        assert!(node.iroh_docs.lock().expect("docs lock").is_none());
+        assert!(node.iroh_gossip.lock().expect("gossip lock").is_none());
+        assert!(
+            tokio::time::timeout(Duration::from_secs(1), raw_endpoint.accept())
+                .await
+                .expect("closed endpoint accept resolves")
+                .is_none()
+        );
+    }
+}