From af529b171981adbcf5eab53d1f0663b06f68c676 Mon Sep 17 00:00:00 2001
From: Ellie Huxtable <ellie@posthog.com>
Date: Thu, 14 Sep 2023 15:30:06 +0100
Subject: [PATCH] Add graceful shutdown handler (#27)

I rolled out the new image, saw the old took ages to kill, and realised
we hadn't setup a SIGTERM handler. Add it!
---
 capture-server/src/main.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/capture-server/src/main.rs b/capture-server/src/main.rs
index 88ffc94..e2232d5 100644
--- a/capture-server/src/main.rs
+++ b/capture-server/src/main.rs
@@ -2,6 +2,22 @@ use std::env;
 use std::net::SocketAddr;
 
 use capture::{router, sink, time};
+use tokio::signal;
+
+async fn shutdown() {
+    let mut term = signal::unix::signal(signal::unix::SignalKind::terminate())
+        .expect("failed to register SIGTERM handler");
+
+    let mut interrupt = signal::unix::signal(signal::unix::SignalKind::interrupt())
+        .expect("failed to register SIGINT handler");
+
+    tokio::select! {
+        _ = term.recv() => {},
+        _ = interrupt.recv() => {},
+    };
+
+    tracing::info!("Shutting down gracefully...");
+}
 
 #[tokio::main]
 async fn main() {
@@ -29,6 +45,7 @@ async fn main() {
 
     axum::Server::bind(&address.parse().unwrap())
         .serve(app.into_make_service_with_connect_info::<SocketAddr>())
+        .with_graceful_shutdown(shutdown())
         .await
         .unwrap();
 }