fix: Add SIGHUP and SIGTERM handling to coder server (#3543)

mafredri · web-flow · commit 6a0f8ae9ccda · 2022-08-18T16:25:32.000+03:00
* 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
diff --git 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
@@ -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
@@ -0,0 +1,9 @@
+//go:build windows
+
+package cli
+
+import (
+	"os"
+)
+
+var interruptSignals = []os.Signal{os.Interrupt}
