From cfe5cb78c5a2d78ccc69c837c3d0ea715aa0b612 Mon Sep 17 00:00:00 2001
From: Luc Perkins <lucperkins@gmail.com>
Date: Fri, 17 May 2024 13:22:58 -0300
Subject: [PATCH] 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<reqwest::Url>,
+
+    /// File to write to when indicating startup.
+    #[arg(long)]
+    startup_notification_file: Option<PathBuf>,
 }
 
 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 {