diff --git a/flake.nix b/flake.nix index 697c2a7..560d87c 100644 --- a/flake.nix +++ b/flake.nix @@ -51,6 +51,7 @@ packages = with pkgs; [ bashInteractive cranePkgs.rustNightly + pkg-config cargo-bloat cargo-edit diff --git a/magic-nix-cache/Cargo.toml b/magic-nix-cache/Cargo.toml index e853dd3..0253343 100644 --- a/magic-nix-cache/Cargo.toml +++ b/magic-nix-cache/Cargo.toml @@ -7,11 +7,25 @@ license = "Apache-2.0" [dependencies] gha-cache = { path = "../gha-cache" } -axum = { version = "0.6.18", default-features = false, features = ["json", "tokio"] } +axum = { version = "0.6.18", default-features = false, features = [ + "json", + "tokio", +] } axum-macros = "0.3.7" -clap = { version = "4.2.7", default-features = false, features = ["std", "derive", "error-context", "wrap_help"] } +clap = { version = "4.2.7", default-features = false, features = [ + "std", + "derive", + "error-context", + "wrap_help", +] } tracing = "0.1.37" -tracing-subscriber = { version = "0.3.17", default-features = false, features = ["ansi", "env-filter", "fmt", "tracing-log", "smallvec"] } +tracing-subscriber = { version = "0.3.17", default-features = false, features = [ + "ansi", + "env-filter", + "fmt", + "tracing-log", + "smallvec", +] } tower-http = { version = "0.4.0", features = ["trace"] } serde = { version = "1.0.162", features = ["derive"] } serde_json = { version = "1.0.96", default-features = false } @@ -21,7 +35,11 @@ tokio-util = { version = "0.7.8", features = ["io", "compat"] } daemonize = "0.5.0" is_ci = "1.1.1" sha2 = { version = "0.10.6", default-features = false } -reqwest = { version = "0.11.17", default-features = false, features = ["blocking", "rustls-tls-native-roots", "trust-dns"] } +reqwest = { version = "0.11.17", default-features = false, features = [ + "blocking", + "rustls-tls-native-roots", + "trust-dns", +] } netrc-rs = "0.1.2" attic = { git = "https://github.com/DeterminateSystems/attic", branch = "fixups-for-magic-nix-cache" } attic-client = { git = "https://github.com/DeterminateSystems/attic", branch = "fixups-for-magic-nix-cache" } @@ -36,11 +54,4 @@ async-compression = "0.4" [dependencies.tokio] version = "1.28.0" default-features = false -features = [ - "fs", - "macros", - "process", - "rt", - "rt-multi-thread", - "sync", -] +features = ["fs", "macros", "process", "rt", "rt-multi-thread", "sync"] diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs index f7f8d1e..fb97e3c 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -26,6 +26,7 @@ use std::io::Write; use std::net::SocketAddr; use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; +use std::process::exit; use std::sync::Arc; use ::attic::nix_store::NixStore; @@ -33,6 +34,8 @@ use anyhow::{anyhow, Context, Result}; use axum::{extract::Extension, routing::get, Router}; use clap::Parser; use tempfile::NamedTempFile; +use tokio::fs::File; +use tokio::io::AsyncWriteExt; use tokio::process::Command; use tokio::sync::{oneshot, Mutex, RwLock}; use tracing_subscriber::filter::EnvFilter; @@ -111,6 +114,10 @@ struct Args { /// URL to which to post startup notification. #[arg(long)] startup_notification_url: Option, + + /// File to write to when indicating startup. + #[arg(long)] + startup_notification_file: Option, } impl Args { @@ -364,7 +371,25 @@ async fn main_cli() -> Result<()> { tracing::info!("Listening on {}", args.listen); + let server = axum::Server::bind(&args.listen) + .serve(app.into_make_service()) + .with_graceful_shutdown(async move { + shutdown_receiver.await.ok(); + tracing::info!("Shutting down"); + }); + + // Spawn here so that post-startup tasks can proceed + tokio::spawn(async move { + if let Err(e) = server.await { + tracing::error!("failed to start up daemon: {e}"); + exit(1); + } + }); + + // Notify of startup via HTTP if let Some(startup_notification_url) = args.startup_notification_url { + tracing::debug!("Startup notification via HTTP POST to {startup_notification_url}"); + let response = reqwest::Client::new() .post(startup_notification_url) .header(reqwest::header::CONTENT_TYPE, "application/json") @@ -390,20 +415,23 @@ async fn main_cli() -> Result<()> { } } - let ret = axum::Server::bind(&args.listen) - .serve(app.into_make_service()) - .with_graceful_shutdown(async move { - shutdown_receiver.await.ok(); - tracing::info!("Shutting down"); - }) - .await; + // Notify of startup by writing "1" to the specified file + if let Some(startup_notification_file_path) = args.startup_notification_file { + let file_contents: &[u8] = b"1"; + tracing::debug!("Startup notification via file at {startup_notification_file_path:?}"); + + let mut notification_file = File::create(&startup_notification_file_path).await?; + notification_file.write_all(file_contents).await?; + + tracing::debug!("Created startup notification file at {startup_notification_file_path:?}"); + } + + // Notify diagnostics endpoint if let Some(diagnostic_endpoint) = diagnostic_endpoint { state.metrics.send(diagnostic_endpoint).await; } - ret?; - Ok(()) }