From 9ef4fef8583a4543149055af1e6736bed32caf8f Mon Sep 17 00:00:00 2001
From: Olivier 'reivilibre <olivier@librepush.net>
Date: Mon, 21 Mar 2022 19:11:07 +0000
Subject: [PATCH] Shut down gently on SIGINT or SIGTERM (supposedly)

---
 Cargo.lock                          | 13 +++++++++++-
 quickpeep_raker/Cargo.toml          |  1 +
 quickpeep_raker/src/bin/qp-raker.rs | 33 ++++++++++++++++++++++++++---
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 90996fd..9eff500 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -633,7 +633,7 @@ dependencies = [
  "libc",
  "mio 0.7.14",
  "parking_lot 0.11.2",
- "signal-hook",
+ "signal-hook 0.1.17",
  "winapi",
 ]
 
@@ -3053,6 +3053,7 @@ dependencies = [
  "reqwest",
  "serde",
  "serde_bare",
+ "signal-hook 0.3.13",
  "sitemap",
  "smartstring",
  "tokio",
@@ -3570,6 +3571,16 @@ dependencies = [
  "signal-hook-registry",
 ]
 
+[[package]]
+name = "signal-hook"
+version = "0.3.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "647c97df271007dcea485bb74ffdb57f2e683f1306c854f468a0c244badabf2d"
+dependencies = [
+ "libc",
+ "signal-hook-registry",
+]
+
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.0"
diff --git a/quickpeep_raker/Cargo.toml b/quickpeep_raker/Cargo.toml
index 69a6370..383458e 100644
--- a/quickpeep_raker/Cargo.toml
+++ b/quickpeep_raker/Cargo.toml
@@ -51,6 +51,7 @@ lru = "0.7.3"
 diplomatic-bag = "0.2.0"
 arc-interner = "0.7.0"
 smartstring = "1.0.0"
+signal-hook = "0.3.13"
 
 ### Raking helpers
 # HTTP Requests
diff --git a/quickpeep_raker/src/bin/qp-raker.rs b/quickpeep_raker/src/bin/qp-raker.rs
index 96259fa..f184dba 100644
--- a/quickpeep_raker/src/bin/qp-raker.rs
+++ b/quickpeep_raker/src/bin/qp-raker.rs
@@ -9,12 +9,14 @@ use lru::LruCache;
 use metrics_exporter_prometheus::PrometheusBuilder;
 use reqwest::header::{HeaderMap, HeaderValue, USER_AGENT};
 use reqwest::redirect::Policy;
+use signal_hook::consts::{SIGINT, SIGTERM};
+use signal_hook::iterator::Signals;
 use std::path::PathBuf;
-use std::sync::atomic::AtomicBool;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Mutex, RwLock};
 use std::time::Duration;
 use tokio::fs::File;
-use tokio::sync::{mpsc, oneshot, Semaphore};
+use tokio::sync::{mpsc, oneshot, Notify, Semaphore};
 use tokio::time::MissedTickBehavior;
 
 use quickpeep_raker::config;
@@ -220,7 +222,13 @@ pub async fn main() -> anyhow::Result<()> {
         })
     };
 
-    drop(task_context);
+    let TaskContext {
+        graceful_stop,
+        notify,
+        ..
+    } = task_context;
+    // ^C is SIGINT; systemd sends SIGTERM
+    start_signal_handler(Signals::new([SIGINT, SIGTERM])?, graceful_stop, notify)?;
 
     for task in tasks {
         task.await?;
@@ -230,3 +238,22 @@ pub async fn main() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+fn start_signal_handler(
+    mut signals: Signals,
+    shutdown: Arc<AtomicBool>,
+    shutdown_notify: Arc<Notify>,
+) -> anyhow::Result<()> {
+    std::thread::Builder::new()
+        .name("signals".to_string())
+        .stack_size(512)
+        .spawn(move || {
+            for sig in signals.forever() {
+                eprintln!("Received signal {:?}, shutting down.", sig);
+                shutdown.store(true, Ordering::SeqCst);
+                shutdown_notify.notify_waiters();
+                break;
+            }
+        })?;
+    Ok(())
+}