From 0e0bef15fb23209a2e004c884c0eacc3b267d3c8 Mon Sep 17 00:00:00 2001
From: Zhaofeng Li <hello@zhaofeng.li>
Date: Fri, 19 May 2023 02:48:52 -0600
Subject: [PATCH] Support forking into background

This makes the JS action simpler.
---
 Cargo.lock                    | 10 +++++++
 nix-actions-cache/Cargo.toml  |  1 +
 nix-actions-cache/src/main.rs | 53 +++++++++++++++++++++++++----------
 3 files changed, 49 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 7b7acd2..9f0b8bf 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -251,6 +251,15 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "daemonize"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ab8bfdaacb3c887a54d41bdf48d3af8873b3f5566469f8ba21b92057509f116e"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "derivative"
 version = "2.2.0"
@@ -713,6 +722,7 @@ version = "0.1.0"
 dependencies = [
  "axum",
  "clap",
+ "daemonize",
  "gha-cache",
  "rand",
  "serde_json",
diff --git a/nix-actions-cache/Cargo.toml b/nix-actions-cache/Cargo.toml
index ca9a3dd..ce71b96 100644
--- a/nix-actions-cache/Cargo.toml
+++ b/nix-actions-cache/Cargo.toml
@@ -18,6 +18,7 @@ thiserror = "1.0.40"
 tokio-stream = "0.1.14"
 tokio-util = { version = "0.7.8", features = ["io"] }
 rand = "0.8.5"
+daemonize = "0.5.0"
 
 [dependencies.tokio]
 version = "1.28.0"
diff --git a/nix-actions-cache/src/main.rs b/nix-actions-cache/src/main.rs
index c7ea6f1..8a918fc 100644
--- a/nix-actions-cache/src/main.rs
+++ b/nix-actions-cache/src/main.rs
@@ -16,13 +16,16 @@
 mod api;
 mod error;
 
+use std::fs::{self, File};
 use std::net::SocketAddr;
+use std::os::fd::OwnedFd;
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use axum::{extract::Extension, routing::get, Router};
 use clap::Parser;
-use tokio::fs;
+use daemonize::Daemonize;
+use tokio::runtime::Runtime;
 use tracing_subscriber::EnvFilter;
 
 use gha_cache::{Api, Credentials};
@@ -58,6 +61,12 @@ struct Args {
     /// instead.
     #[arg(long)]
     upstream: Option<String>,
+
+    /// Daemonize the server.
+    ///
+    /// This is for use in the GitHub Action only.
+    #[arg(long, hide = true)]
+    daemon_dir: Option<PathBuf>,
 }
 
 /// The global server state.
@@ -67,17 +76,14 @@ struct StateInner {
     upstream: Option<String>,
 }
 
-#[tokio::main]
-async fn main() {
-    let args = Args::parse();
-
+fn main() {
     init_logging();
 
+    let args = Args::parse();
+
     let credentials = if let Some(credentials_file) = &args.credentials_file {
         tracing::info!("Loading credentials from {:?}", credentials_file);
-        let bytes = fs::read(credentials_file)
-            .await
-            .expect("Failed to read credentials file");
+        let bytes = fs::read(credentials_file).expect("Failed to read credentials file");
 
         serde_json::from_slice(&bytes).expect("Failed to deserialize credentials file")
     } else {
@@ -88,13 +94,13 @@ async fn main() {
 
     let mut api = Api::new(credentials).expect("Failed to initialize GitHub Actions Cache API");
 
-    if let Some(cache_version) = args.cache_version {
+    if let Some(cache_version) = &args.cache_version {
         api.mutate_version(cache_version.as_bytes());
     }
 
     let state = Arc::new(StateInner {
         api,
-        upstream: args.upstream,
+        upstream: args.upstream.clone(),
     });
 
     let app = Router::new().route("/", get(root)).merge(api::get_router());
@@ -106,11 +112,28 @@ async fn main() {
 
     let app = app.layer(Extension(state));
 
-    tracing::info!("listening on {}", args.listen);
-    axum::Server::bind(&args.listen)
-        .serve(app.into_make_service())
-        .await
-        .unwrap();
+    if args.daemon_dir.is_some() {
+        let dir = args.daemon_dir.as_ref().unwrap();
+        let logfile: OwnedFd = File::create(dir.join("daemon.log")).unwrap().into();
+        let daemon = Daemonize::new()
+            .pid_file(dir.join("daemon.pid"))
+            .stdout(File::from(logfile.try_clone().unwrap()))
+            .stderr(File::from(logfile));
+
+        daemon.start().expect("Failed to fork into the background");
+    }
+
+    let rt = Runtime::new().unwrap();
+    rt.block_on(async move {
+        tracing::info!("listening on {}", args.listen);
+        axum::Server::bind(&args.listen)
+            .serve(app.into_make_service())
+            .with_graceful_shutdown(async move {
+                shutdown_receiver.await.ok();
+            })
+            .await
+            .unwrap()
+    });
 }
 
 fn init_logging() {