From 0b4f59b82b158afc3533fa16b4660eb48f99781d Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kradalby@kradalby.no>
Date: Tue, 31 May 2022 10:57:20 +0200
Subject: [PATCH] Improve signal handling

This commit starts to wire up better signal handling, it starts with
handling shutdown a bit better, using the graceful shutdown for all the
listeners we use.

It also adds the initial switch case for handling config and acl reload,
which is to be implemented.
---
 app.go | 55 ++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 42 insertions(+), 13 deletions(-)

diff --git a/app.go b/app.go
index a78a6b0e..cc4728d3 100644
--- a/app.go
+++ b/app.go
@@ -568,19 +568,6 @@ func (h *Headscale) Serve() error {
 		return fmt.Errorf("failed change permission of gRPC socket: %w", err)
 	}
 
-	// Handle common process-killing signals so we can gracefully shut down:
-	sigc := make(chan os.Signal, 1)
-	signal.Notify(sigc, os.Interrupt, syscall.SIGTERM)
-	go func(c chan os.Signal) {
-		// Wait for a SIGINT or SIGKILL:
-		sig := <-c
-		log.Printf("Caught signal %s: shutting down.", sig)
-		// Stop listening (and unlink the socket if unix type):
-		socketListener.Close()
-		// And we're done:
-		os.Exit(0)
-	}(sigc)
-
 	grpcGatewayMux := runtime.NewServeMux()
 
 	// Make the grpc-gateway connect to grpc over socket
@@ -725,6 +712,48 @@ func (h *Headscale) Serve() error {
 	log.Info().
 		Msgf("listening and serving metrics on: %s", h.cfg.MetricsAddr)
 
+	// Handle common process-killing signals so we can gracefully shut down:
+	sigc := make(chan os.Signal, 1)
+	signal.Notify(sigc,
+		syscall.SIGHUP,
+		syscall.SIGINT,
+		syscall.SIGTERM,
+		syscall.SIGQUIT,
+		syscall.SIGHUP)
+	go func(c chan os.Signal) {
+		// Wait for a SIGINT or SIGKILL:
+		sig := <-c
+		switch sig {
+		case syscall.SIGHUP:
+			log.Info().
+				Str("signal", sig.String()).
+				Msg("Received SIGHUP, reloading ACL and Config")
+
+		// TODO(kradalby): Reload config on SIGHUP
+
+		default:
+			log.Info().
+				Str("signal", sig.String()).
+				Msg("Received signal to stop, shutting down gracefully")
+
+			// Gracefully shut down servers
+			promHTTPServer.Shutdown(ctx)
+			httpServer.Shutdown(ctx)
+			grpcSocket.GracefulStop()
+
+			// Close network listeners
+			promHTTPListener.Close()
+			httpListener.Close()
+			grpcGatewayConn.Close()
+
+			// Stop listening (and unlink the socket if unix type):
+			socketListener.Close()
+
+			// And we're done:
+			os.Exit(0)
+		}
+	}(sigc)
+
 	return errorGroup.Wait()
 }