From 6a0f8ae9ccdae90c54222b84224418aabeae273d Mon Sep 17 00:00:00 2001
From: Mathias Fredriksson <mafredri@gmail.com>
Date: Thu, 18 Aug 2022 16:25:32 +0300
Subject: [PATCH] fix: Add `SIGHUP` and `SIGTERM` handling to `coder server`
 (#3543)

* fix: Add `SIGHUP` and `SIGTERM` handling to `coder server`

To prevent additional signals from aborting program execution, signal
handling was moved to the beginning of the main function, this ensures
that signals stays registered for the entire shutdown procedure.

Fixes #1529
---
 cli/server.go         | 26 +++++++++++++++-----------
 cli/signal_unix.go    | 14 ++++++++++++++
 cli/signal_windows.go |  9 +++++++++
 3 files changed, 38 insertions(+), 11 deletions(-)
 create mode 100644 cli/signal_unix.go
 create mode 100644 cli/signal_windows.go

diff --git a/cli/server.go b/cli/server.go
index 86016d9c9b..126ea0bc82 100644
--- a/cli/server.go
+++ b/cli/server.go
@@ -127,6 +127,19 @@ func server() *cobra.Command {
 			ctx, cancel := context.WithCancel(cmd.Context())
 			defer cancel()
 
+			// Register signals early on so that graceful shutdown can't
+			// be interrupted by additional signals. Note that we avoid
+			// shadowing cancel() (from above) here because notifyStop()
+			// restores default behavior for the signals. This protects
+			// the shutdown sequence from abrubtly terminating things
+			// like: database migrations, provisioner work, workspace
+			// cleanup in dev-mode, etc.
+			//
+			// To get out of a graceful shutdown, the user can send
+			// SIGQUIT with ctrl+\ or SIGKILL with `kill -9`.
+			notifyCtx, notifyStop := signal.NotifyContext(ctx, interruptSignals...)
+			defer notifyStop()
+
 			// Clean up idle connections at the end, e.g.
 			// embedded-postgres can leave an idle connection
 			// which is caught by goleaks.
@@ -521,22 +534,13 @@ func server() *cobra.Command {
 			// such as via the systemd service.
 			_ = config.URL().Write(client.URL.String())
 
-			// Because the graceful shutdown includes cleaning up workspaces in dev mode, we're
-			// going to make it harder to accidentally skip the graceful shutdown by hitting ctrl+c
-			// two or more times.  So the stopChan is unlimited in size and we don't call
-			// signal.Stop() until graceful shutdown finished--this means we swallow additional
-			// SIGINT after the first.  To get out of a graceful shutdown, the user can send SIGQUIT
-			// with ctrl+\ or SIGTERM with `kill`.
-			ctx, stop := signal.NotifyContext(ctx, os.Interrupt)
-			defer stop()
-
 			// Currently there is no way to ask the server to shut
 			// itself down, so any exit signal will result in a non-zero
 			// exit of the server.
 			var exitErr error
 			select {
-			case <-ctx.Done():
-				exitErr = ctx.Err()
+			case <-notifyCtx.Done():
+				exitErr = notifyCtx.Err()
 				_, _ = fmt.Fprintln(cmd.OutOrStdout(), cliui.Styles.Bold.Render(
 					"Interrupt caught, gracefully exiting. Use ctrl+\\ to force quit",
 				))
diff --git a/cli/signal_unix.go b/cli/signal_unix.go
new file mode 100644
index 0000000000..7d2cd0e502
--- /dev/null
+++ b/cli/signal_unix.go
@@ -0,0 +1,14 @@
+//go:build !windows
+
+package cli
+
+import (
+	"os"
+	"syscall"
+)
+
+var interruptSignals = []os.Signal{
+	os.Interrupt,
+	syscall.SIGTERM,
+	syscall.SIGHUP,
+}
diff --git a/cli/signal_windows.go b/cli/signal_windows.go
new file mode 100644
index 0000000000..17652adfb6
--- /dev/null
+++ b/cli/signal_windows.go
@@ -0,0 +1,9 @@
+//go:build windows
+
+package cli
+
+import (
+	"os"
+)
+
+var interruptSignals = []os.Signal{os.Interrupt}