diff options
| author | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-07-20 12:57:36 +0200 |
|---|---|---|
| committer | Benedikt Peetz <benedikt.peetz@b-peetz.de> | 2026-07-20 12:57:36 +0200 |
| commit | 73be69cd99f5a4784fe69f5d78f423c72e837284 (patch) | |
| tree | cb47246d2c7edd8cd828abb25a033b732d2289ae /crates/daemon/src/main.rs | |
| parent | chore: Move more stuff out of atuin-client (diff) | |
| download | atuin-73be69cd99f5a4784fe69f5d78f423c72e837284.zip | |
chore: Commit
Diffstat (limited to 'crates/daemon/src/main.rs')
| -rw-r--r-- | crates/daemon/src/main.rs | 377 |
1 files changed, 377 insertions, 0 deletions
diff --git a/crates/daemon/src/main.rs b/crates/daemon/src/main.rs new file mode 100644 index 00000000..26a5cafd --- /dev/null +++ b/crates/daemon/src/main.rs @@ -0,0 +1,377 @@ +#[allow(unused_imports)] +use clap::Parser; +use eyre::{Result, WrapErr, bail, eyre}; +use fs4::fs_std::FileExt; +use std::fs::{self, File, OpenOptions}; +use std::io::{ErrorKind, Write}; +#[cfg(unix)] +use std::os::unix::net::UnixStream as StdUnixStream; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; +use tokio::time::sleep; +use turtle_daemon::{ + DaemonEvent, + aclient::{ + database::ClientSqlite, history::History, record::sqlite_store::SqliteStore, + settings::Settings, + }, + client::{ControlClient, DaemonClientErrorKind, HistoryClient, classify_error}, +}; + +#[derive(Parser, Debug)] +#[command(infer_subcommands = true)] +pub(crate) enum Cmd { + /// Start the daemon server + Start { + #[arg(long, hide = true)] + daemonize: bool, + + /// Also write daemon logs to the console (useful for debugging) + #[arg(long)] + show_logs: bool, + }, + + /// Show the daemon's current status + Status, + + /// Stop the daemon gracefully + Stop, +} + +impl Cmd { + pub(crate) async fn run( + self, + settings: Settings, + store: SqliteStore, + history_db: ClientSqlite, + ) -> Result<()> { + match self { + Cmd::Start { .. } => run(settings, store, history_db).await, + Cmd::Status => status_cmd(&settings).await, + Cmd::Stop => stop_cmd(&settings).await, + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let settings = Settings::new().wrap_err("could not load client settings")?; + let db_path = PathBuf::from(settings.db_path.as_str()); + let record_store_path = PathBuf::from(settings.record_store_path.as_str()); + + let db = ClientSqlite::new(db_path, settings.local_timeout).await?; + let sqlite_store = SqliteStore::new(record_store_path, settings.local_timeout).await?; + + Cmd::parse().run(settings, sqlite_store, db).await +} + +const DAEMON_VERSION: &str = env!("CARGO_PKG_VERSION"); +const DAEMON_PROTOCOL_VERSION: u32 = 1; +const STARTUP_POLL: Duration = Duration::from_millis(40); +const LOCK_POLL: Duration = Duration::from_millis(20); +const LEGACY_DAEMON_RESTART_MESSAGE: &str = "legacy daemon detected; restart daemon manually"; + +struct PidfileGuard { + file: File, +} + +impl PidfileGuard { + fn acquire(path: &Path) -> Result<Self> { + let mut file = open_lock_file(path)?; + + if !file.try_lock_exclusive()? { + bail!( + "daemon already running (pidfile lock busy at {})", + path.display() + ); + } + + file.set_len(0) + .wrap_err_with(|| format!("could not truncate daemon pidfile {}", path.display()))?; + writeln!(file, "{}", std::process::id()) + .and_then(|()| writeln!(file, "{DAEMON_VERSION}")) + .wrap_err_with(|| format!("could not write daemon pidfile {}", path.display()))?; + + Ok(Self { file }) + } +} + +impl Drop for PidfileGuard { + fn drop(&mut self) { + drop(self.file.unlock()); + } +} + +enum Probe { + Ready(HistoryClient), + NeedsRestart(String), + Unreachable(eyre::Report), +} + +fn daemon_matches_expected(version: &str, protocol: u32) -> bool { + version == DAEMON_VERSION && protocol == DAEMON_PROTOCOL_VERSION +} + +fn daemon_mismatch_message(version: &str, protocol: u32) -> String { + if protocol == DAEMON_PROTOCOL_VERSION { + format!("daemon is out of date: expected {DAEMON_VERSION}, got {version}") + } else { + format!("daemon protocol mismatch: expected {DAEMON_PROTOCOL_VERSION}, got {protocol}") + } +} + +fn is_legacy_daemon_error(err: &eyre::Report) -> bool { + matches!(classify_error(err), DaemonClientErrorKind::Unimplemented) +} + +fn open_lock_file(path: &Path) -> Result<File> { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .wrap_err_with(|| format!("could not create lock directory {}", parent.display()))?; + } + + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(path) + .wrap_err_with(|| format!("could not open lock file {}", path.display())) +} + +async fn wait_for_lock(path: &Path, timeout: Duration) -> Result<File> { + let file = open_lock_file(path)?; + let start = Instant::now(); + + loop { + match file.try_lock_exclusive() { + Ok(true) => return Ok(file), + Ok(false) => { + if start.elapsed() >= timeout { + bail!("timed out waiting for lock at {}", path.display()); + } + + sleep(LOCK_POLL).await; + } + Err(err) => { + return Err(eyre!("could not lock {}: {err}", path.display())); + } + } + } +} + +async fn wait_for_pidfile_available(path: &Path, timeout: Duration) -> Result<()> { + let file = wait_for_lock(path, timeout).await?; + file.unlock() + .wrap_err_with(|| format!("failed to unlock {}", path.display()))?; + Ok(()) +} + +async fn connect_client(settings: &Settings) -> Result<HistoryClient> { + HistoryClient::new( + #[cfg(unix)] + settings.daemon.socket_path.clone(), + ) + .await +} + +async fn probe(settings: &Settings) -> Probe { + let mut client = match connect_client(settings).await { + Ok(client) => client, + Err(err) => return Probe::Unreachable(err), + }; + + match client.status().await { + Ok(status) => { + if daemon_matches_expected(&status.version, status.protocol) { + Probe::Ready(client) + } else { + Probe::NeedsRestart(daemon_mismatch_message(&status.version, status.protocol)) + } + } + Err(err) => Probe::Unreachable(err), + } +} + +async fn request_shutdown(settings: &Settings) { + if let Ok(mut client) = connect_client(settings).await { + drop(client.shutdown().await); + } +} + +fn startup_timeout(settings: &Settings) -> Duration { + Duration::from_secs_f64(settings.local_timeout.max(0.5) + 2.0) +} + +pub(crate) async fn start_history(settings: &Settings, history: History) -> Result<String> { + match async { + connect_client(settings) + .await? + .start_history(history.clone()) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(resp.id); + } + + Err(eyre!( + "{}. Restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )) + } + Err(err) => Err(err), + } +} + +pub(crate) async fn end_history( + settings: &Settings, + id: String, + duration: u64, + exit: i64, +) -> Result<()> { + match async { + connect_client(settings) + .await? + .end_history(id.clone(), duration, exit) + .await + } + .await + { + Ok(resp) => { + if daemon_matches_expected(&resp.version, resp.protocol) { + return Ok(()); + } + + Err(eyre!( + "{}. Restart the daemon manually", + daemon_mismatch_message(&resp.version, resp.protocol) + )) + } + Err(err) => Err(err), + } +} + +/// Emit a daemon event. +pub(crate) async fn emit_event(settings: &Settings, event: DaemonEvent) { + // Try to connect and send + match ControlClient::from_settings(settings).await { + Ok(mut client) => { + if let Err(e) = client.send_event(event).await { + tracing::debug!(?e, "failed to send event to daemon"); + } + } + Err(e) => { + tracing::debug!(?e, "daemon not available, skipping event emission"); + } + } +} + +pub(crate) async fn tail_client(settings: &Settings) -> Result<HistoryClient> { + match probe(settings).await { + Probe::Ready(client) => Ok(client), + Probe::NeedsRestart(reason) => { + bail!("{reason}. Restart the daemon manually"); + } + Probe::Unreachable(err) if is_legacy_daemon_error(&err) => { + Err(err.wrap_err(LEGACY_DAEMON_RESTART_MESSAGE)) + } + Probe::Unreachable(err) => Err(err), + } +} + +async fn status_cmd(settings: &Settings) -> Result<()> { + match probe(settings).await { + Probe::Ready(mut client) => { + let status = client.status().await?; + println!("Daemon running"); + println!(" PID: {}", status.pid); + println!(" Version: {}", status.version); + println!(" Protocol: {}", status.protocol); + println!(" Healthy: {}", status.healthy); + #[cfg(unix)] + println!(" Socket: {}", settings.daemon.socket_path); + } + Probe::NeedsRestart(reason) => { + println!("Daemon running (needs restart)"); + println!(" Reason: {reason}"); + } + Probe::Unreachable(_) => { + println!("Daemon is not running"); + } + } + + Ok(()) +} + +async fn stop_cmd(settings: &Settings) -> Result<()> { + let Ok(mut client) = connect_client(settings).await else { + println!("Daemon is not running"); + return Ok(()); + }; + + match client.shutdown().await { + Ok(true) => { + println!("Shutdown requested"); + + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let timeout = Duration::from_secs(5); + match wait_for_pidfile_available(&pidfile_path, timeout).await { + Ok(()) => println!("Daemon stopped"), + Err(_) => println!("Daemon may still be shutting down"), + } + + Ok(()) + } + Ok(false) => bail!("Daemon rejected shutdown request"), + Err(err) => Err(err.wrap_err("Failed to send shutdown request")), + } +} + +async fn run(settings: Settings, store: SqliteStore, history_db: ClientSqlite) -> Result<()> { + let pidfile_path = PathBuf::from(&settings.daemon.pidfile_path); + let _pidfile_guard = PidfileGuard::acquire(&pidfile_path)?; + + turtle_daemon::boot(settings, store, history_db).await?; + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::{ + DAEMON_PROTOCOL_VERSION, DAEMON_VERSION, daemon_matches_expected, daemon_mismatch_message, + }; + + #[test] + fn test_version_matches() { + assert!(daemon_matches_expected( + DAEMON_VERSION, + DAEMON_PROTOCOL_VERSION + )); + } + + #[test] + fn test_version_mismatch() { + assert!(!daemon_matches_expected("0.0.0", DAEMON_PROTOCOL_VERSION)); + assert!(!daemon_matches_expected(DAEMON_VERSION, 999)); + assert!(!daemon_matches_expected("0.0.0", 999)); + } + + #[test] + fn test_mismatch_message_version() { + let msg = daemon_mismatch_message("0.0.0", DAEMON_PROTOCOL_VERSION); + assert!(msg.contains("out of date"), "got: {msg}"); + assert!(msg.contains("0.0.0")); + assert!(msg.contains(DAEMON_VERSION)); + } + + #[test] + fn test_mismatch_message_protocol() { + let msg = daemon_mismatch_message(DAEMON_VERSION, 999); + assert!(msg.contains("protocol mismatch"), "got: {msg}"); + } +} |
