From cfe5cb78c5a2d78ccc69c837c3d0ea715aa0b612 Mon Sep 17 00:00:00 2001 From: Luc Perkins Date: Fri, 17 May 2024 13:22:58 -0300 Subject: [PATCH 1/5] Add file-based notification mechanism --- magic-nix-cache/Cargo.toml | 35 +++++++++++++++++++++++------------ magic-nix-cache/src/main.rs | 25 +++++++++++++++++++++++++ 2 files changed, 48 insertions(+), 12 deletions(-) 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..c9c8c0a 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -33,6 +33,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 +113,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 { @@ -127,6 +133,12 @@ impl Args { ))); } + if self.startup_notification_url.is_none() && self.startup_notification_file.is_none() { + return Err(error::Error::Config(String::from( + "you must specify either --startup-notification-url or --startup-notification-file", + ))); + } + Ok(()) } } @@ -364,7 +376,10 @@ async fn main_cli() -> Result<()> { tracing::info!("Listening on {}", args.listen); + // 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,6 +405,16 @@ async fn main_cli() -> Result<()> { } } + // Notify of startup via file + if let Some(startup_notification_file_path) = args.startup_notification_file { + 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(b"1").await?; + + tracing::debug!("Created startup notification file at {startup_notification_file_path:?}"); + } + let ret = axum::Server::bind(&args.listen) .serve(app.into_make_service()) .with_graceful_shutdown(async move { From c0b7181ddc1b876c92cf497d4668b72d2711d5d2 Mon Sep 17 00:00:00 2001 From: Luc Perkins Date: Fri, 17 May 2024 13:32:43 -0300 Subject: [PATCH 2/5] Allow for both startup options to be None --- magic-nix-cache/src/main.rs | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs index c9c8c0a..39386ae 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -133,12 +133,6 @@ impl Args { ))); } - if self.startup_notification_url.is_none() && self.startup_notification_file.is_none() { - return Err(error::Error::Config(String::from( - "you must specify either --startup-notification-url or --startup-notification-file", - ))); - } - Ok(()) } } @@ -405,12 +399,14 @@ async fn main_cli() -> Result<()> { } } - // Notify of startup via file + // 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(b"1").await?; + notification_file.write_all(file_contents).await?; tracing::debug!("Created startup notification file at {startup_notification_file_path:?}"); } From d67f3303979e12f816517b9d5aaf43216b3a7092 Mon Sep 17 00:00:00 2001 From: Luc Perkins Date: Fri, 17 May 2024 14:07:12 -0300 Subject: [PATCH 3/5] Spawn daemon in separate process --- magic-nix-cache/src/main.rs | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs index 39386ae..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; @@ -370,6 +371,21 @@ 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}"); @@ -411,20 +427,11 @@ async fn main_cli() -> Result<()> { tracing::debug!("Created startup notification file at {startup_notification_file_path:?}"); } - 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 diagnostics endpoint if let Some(diagnostic_endpoint) = diagnostic_endpoint { state.metrics.send(diagnostic_endpoint).await; } - ret?; - Ok(()) } From 736bd0c01924ca361e46a5c0903b12bae47dc2e5 Mon Sep 17 00:00:00 2001 From: Luc Perkins Date: Fri, 17 May 2024 14:22:01 -0300 Subject: [PATCH 4/5] Delete notification file if server fails to start up --- magic-nix-cache/src/main.rs | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/magic-nix-cache/src/main.rs b/magic-nix-cache/src/main.rs index fb97e3c..a63d7bf 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -382,6 +382,16 @@ async fn main_cli() -> Result<()> { tokio::spawn(async move { if let Err(e) = server.await { tracing::error!("failed to start up daemon: {e}"); + + // Delete the notification file if it was created + if let Some(startup_notification_file_path) = args.startup_notification_file_path { + if File::metadata(startup_notification_file_path).await.is_ok() { + if let Err(e) = tokio::fs::remove_file(startup_notification_file_path).await { + tracing::error!("failed to remove stray startup notification file at {startup_notification_file_path:?}; you may need to delete it manually"); + } + } + } + exit(1); } }); From 5da333f97bf38891ff72ccc256c8d4efc6447ecf Mon Sep 17 00:00:00 2001 From: Luc Perkins Date: Fri, 17 May 2024 14:32:40 -0300 Subject: [PATCH 5/5] Add missing pkg-config dependency to dev shell --- flake.nix | 1 + magic-nix-cache/src/main.rs | 10 ---------- 2 files changed, 1 insertion(+), 10 deletions(-) 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/src/main.rs b/magic-nix-cache/src/main.rs index a63d7bf..fb97e3c 100644 --- a/magic-nix-cache/src/main.rs +++ b/magic-nix-cache/src/main.rs @@ -382,16 +382,6 @@ async fn main_cli() -> Result<()> { tokio::spawn(async move { if let Err(e) = server.await { tracing::error!("failed to start up daemon: {e}"); - - // Delete the notification file if it was created - if let Some(startup_notification_file_path) = args.startup_notification_file_path { - if File::metadata(startup_notification_file_path).await.is_ok() { - if let Err(e) = tokio::fs::remove_file(startup_notification_file_path).await { - tracing::error!("failed to remove stray startup notification file at {startup_notification_file_path:?}; you may need to delete it manually"); - } - } - } - exit(1); } });