From 66ca5131bcb94f968feb3dca36392da5aaf62b3e Mon Sep 17 00:00:00 2001
From: Theo Butler <theodusbutler@gmail.com>
Date: Fri, 26 Jan 2024 09:47:48 -0500
Subject: [PATCH] fix: graceful shutdown (#564)

---
 Cargo.lock                | 10 ++++++++++
 Cargo.toml                |  1 +
 graph-gateway/src/main.rs | 30 ++++++++++++++++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index dded9287..f83f794d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4445,6 +4445,15 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "signature"
 version = "2.2.0"
@@ -4924,6 +4933,7 @@ dependencies = [
  "num_cpus",
  "parking_lot",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
  "tokio-macros",
  "windows-sys 0.48.0",
diff --git a/Cargo.toml b/Cargo.toml
index 4a94df60..76304e37 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -45,6 +45,7 @@ tokio = { version = "1.35", features = [
     "macros",
     "parking_lot",
     "rt-multi-thread",
+    "signal",
     "sync",
     "time",
 ] }
diff --git a/graph-gateway/src/main.rs b/graph-gateway/src/main.rs
index a6709d04..a6216e66 100644
--- a/graph-gateway/src/main.rs
+++ b/graph-gateway/src/main.rs
@@ -26,6 +26,7 @@ use thegraph::{
     client as subgraph_client,
     types::{attestation, DeploymentId},
 };
+use tokio::signal::unix::SignalKind;
 use tokio::spawn;
 use toolshed::{
     buffer_queue::{self, QueueWriter},
@@ -408,8 +409,37 @@ async fn main() {
     // disable Nagel's algorithm
     .tcp_nodelay(true)
     .serve(router.into_make_service_with_connect_info::<SocketAddr>())
+    .with_graceful_shutdown(await_shutdown_signals())
     .await
     .expect("Failed to start API server");
+    tracing::warn!("shutdown");
+}
+
+async fn await_shutdown_signals() {
+    #[cfg(unix)]
+    let sigint = async {
+        tokio::signal::unix::signal(SignalKind::interrupt())
+            .expect("install SIGINT handler")
+            .recv()
+            .await;
+    };
+    #[cfg(not(unix))]
+    let sigint = std::future::pending::<()>();
+
+    #[cfg(unix)]
+    let sigterm = async {
+        tokio::signal::unix::signal(SignalKind::terminate())
+            .expect("install SIGTERM handler")
+            .recv()
+            .await;
+    };
+    #[cfg(not(unix))]
+    let sigterm = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = sigint => (),
+        _ = sigterm => (),
+    }
 }
 
 fn update_from_eventual<V, F>(eventual: Eventual<V>, writer: QueueWriter<Update>, f: F)