feat: support systemd socket activation for daemon (#2039)

This avoids issues with clients attempting to connect to the daemon
while it's starting, systemd creates the socket early and will queue
connections up until the daemon is ready to accept them.

1 files changed, 48 insertions, 7 deletions

diff --git a/crates/atuin-daemon/src/server.rs b/crates/atuin-daemon/src/server.rs
index 77824f60..1cfcef51 100644
--- a/crates/atuin-daemon/src/server.rs
+++ b/crates/atuin-daemon/src/server.rs
@@ -133,7 +133,7 @@ impl HistorySvc for HistoryService {
 }
 
 #[cfg(unix)]
-async fn shutdown_signal(socket: PathBuf) {
+async fn shutdown_signal(socket: Option<PathBuf>) {
     let mut term = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
         .expect("failed to register sigterm handler");
     let mut int = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())
@@ -145,7 +145,9 @@ async fn shutdown_signal(socket: PathBuf) {
     }
 
     eprintln!("Removing socket...");
-    std::fs::remove_file(socket).expect("failed to remove socket");
+    if let Some(socket) = socket {
+        std::fs::remove_file(socket).expect("failed to remove socket");
+    }
     eprintln!("Shutting down...");
 }
 
@@ -163,15 +165,54 @@ async fn start_server(settings: Settings, history: HistoryService) -> Result<()>
     use tokio::net::UnixListener;
     use tokio_stream::wrappers::UnixListenerStream;
 
-    let socket = settings.daemon.socket_path.clone();
+    let socket_path = settings.daemon.socket_path;
 
-    let uds = UnixListener::bind(socket.clone())?;
-    let uds_stream = UnixListenerStream::new(uds);
+    let (uds, cleanup) = if cfg!(target_os = "linux") && settings.daemon.systemd_socket {
+        #[cfg(target_os = "linux")]
+        {
+            use eyre::OptionExt;
+            tracing::info!("getting systemd socket");
+            let listener = listenfd::ListenFd::from_env()
+                .take_unix_listener(0)?
+                .ok_or_eyre("missing systemd socket")?;
+            listener.set_nonblocking(true)?;
+            let actual_path = listener
+                .local_addr()
+                .context("getting systemd socket's path")
+                .and_then(|addr| {
+                    addr.as_pathname()
+                        .ok_or_eyre("systemd socket missing path")
+                        .map(|path| path.to_owned())
+                });
+            match actual_path {
+                Ok(actual_path) => {
+                    tracing::info!("listening on systemd socket: {actual_path:?}");
+                    if actual_path != std::path::Path::new(&socket_path) {
+                        tracing::warn!(
+                            "systemd socket is not at configured client path: {socket_path:?}"
+                        );
+                    }
+                }
+                Err(err) => {
+                    tracing::warn!("could not detect systemd socket path, ensure that it's at the configured path: {socket_path:?}, error: {err:?}");
+                }
+            }
+            (UnixListener::from_std(listener)?, false)
+        }
+        #[cfg(not(target_os = "linux"))]
+        unreachable!()
+    } else {
+        tracing::info!("listening on unix socket {socket_path:?}");
+        (UnixListener::bind(socket_path.clone())?, true)
+    };
 
-    tracing::info!("listening on unix socket {:?}", socket);
+    let uds_stream = UnixListenerStream::new(uds);
     Server::builder()
         .add_service(HistoryServer::new(history))
-        .serve_with_incoming_shutdown(uds_stream, shutdown_signal(socket.into()))
+        .serve_with_incoming_shutdown(
+            uds_stream,
+            shutdown_signal(cleanup.then_some(socket_path.into())),
+        )
         .await?;
     Ok(())
 }