Skip to content

Commit

Permalink
feat: support systemd socket activation for daemon (#2039)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Nemo157 authored May 25, 2024
1 parent 40543eb commit 2e88321
Show file tree
Hide file tree
Showing 7 changed files with 79 additions and 13 deletions.
20 changes: 16 additions & 4 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ readme = "README.md"
async-trait = "0.1.58"
base64 = "0.22"
log = "0.4"
time = { version = "=0.3.34", features = [
time = { version = "0.3.36", features = [
"serde-human-readable",
"macros",
"local-offset",
Expand Down
5 changes: 5 additions & 0 deletions crates/atuin-client/config.toml
Original file line number Diff line number Diff line change
Expand Up @@ -223,5 +223,10 @@ records = true
## windows: Not Supported
# socket_path = "~/.local/share/atuin/atuin.sock"

## Use systemd socket activation rather than opening the given path (the path must still be correct for the client)
## linux: false
## mac/windows: Not Supported
# systemd_socket = false

## The port that should be used for TCP on non unix systems
# tcp_port = 8889
5 changes: 5 additions & 0 deletions crates/atuin-client/src/settings.rs
Original file line number Diff line number Diff line change
Expand Up @@ -350,6 +350,9 @@ pub struct Daemon {
/// The path to the unix socket used by the daemon
pub socket_path: String,

/// Use a socket passed via systemd's socket activation protocol, instead of the path
pub systemd_socket: bool,

/// The port that should be used for TCP on non unix systems
pub tcp_port: u64,
}
Expand All @@ -368,6 +371,7 @@ impl Default for Daemon {
enabled: false,
sync_frequency: 300,
socket_path: "".to_string(),
systemd_socket: false,
tcp_port: 8889,
}
}
Expand Down Expand Up @@ -715,6 +719,7 @@ impl Settings {
.set_default("daemon.sync_frequency", 300)?
.set_default("daemon.enabled", false)?
.set_default("daemon.socket_path", socket_path.to_str())?
.set_default("daemon.systemd_socket", false)?
.set_default("daemon.tcp_port", 8889)?
.set_default(
"prefers_reduced_motion",
Expand Down
3 changes: 3 additions & 0 deletions crates/atuin-daemon/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -32,5 +32,8 @@ prost-types = "0.12"
tokio-stream = {version="0.1.14", features=["net"]}
rand.workspace = true

[target.'cfg(target_os = "linux")'.dependencies]
listenfd = "1.0.1"

[build-dependencies]
tonic-build = "0.11"
55 changes: 48 additions & 7 deletions crates/atuin-daemon/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand All @@ -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...");
}

Expand All @@ -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, 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)
};

let uds = UnixListener::bind(socket.clone())?;
let uds_stream = UnixListenerStream::new(uds);

tracing::info!("listening on unix socket {:?}", socket);
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(())
}
Expand Down
2 changes: 1 addition & 1 deletion ui/backend/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ eyre = "0.6"
tauri = { version = "2.0.0-beta", features = ["tray-icon"] }
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
time = "0.3.34"
time = "0.3.36"
uuid = "1.7.0"
syntect = "5.2.0"

Expand Down

0 comments on commit 2e88321

Please sign in to comment.