From ee0c4c62d8b48d87092c28118616d966b3caa6d5 Mon Sep 17 00:00:00 2001
From: Ben Johnson <benbjohnson@yahoo.com>
Date: Sun, 21 Mar 2021 09:35:52 -0600
Subject: [PATCH] Catch sigterm & add shutdown logging

This commit changes the signal handler for `replicate` to catch
`syscall.SIGTERM` for non-Windows installations. It also adds some
logging to indicat when a shutdown has been initiated and when it
has finished.
---
 cmd/litestream/main.go            | 10 +++++++---
 cmd/litestream/main_notwindows.go |  9 +++++++++
 cmd/litestream/main_windows.go    |  7 +++++++
 3 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/cmd/litestream/main.go b/cmd/litestream/main.go
index ac1e6e9..cad28b4 100644
--- a/cmd/litestream/main.go
+++ b/cmd/litestream/main.go
@@ -79,8 +79,7 @@ func (m *Main) Run(ctx context.Context, args []string) (err error) {
 
 		// Setup signal handler.
 		ctx, cancel := context.WithCancel(ctx)
-		ch := make(chan os.Signal, 1)
-		signal.Notify(ch, os.Interrupt)
+		ch := signalChan()
 		go func() { <-ch; cancel() }()
 
 		if err := c.Run(ctx); err != nil {
@@ -90,9 +89,14 @@ func (m *Main) Run(ctx context.Context, args []string) (err error) {
 		// Wait for signal to stop program.
 		<-ctx.Done()
 		signal.Reset()
+		fmt.Println("signal received, litestream shutting down")
 
 		// Gracefully close.
-		return c.Close()
+		if err := c.Close(); err != nil {
+			return err
+		}
+		fmt.Println("litestream shut down")
+		return nil
 
 	case "restore":
 		return (&RestoreCommand{}).Run(ctx, args)
diff --git a/cmd/litestream/main_notwindows.go b/cmd/litestream/main_notwindows.go
index 8304106..c7713c0 100644
--- a/cmd/litestream/main_notwindows.go
+++ b/cmd/litestream/main_notwindows.go
@@ -4,6 +4,9 @@ package main
 
 import (
 	"context"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 const defaultConfigPath = "/etc/litestream.yml"
@@ -15,3 +18,9 @@ func isWindowsService() (bool, error) {
 func runWindowsService(ctx context.Context) error {
 	panic("cannot run windows service as unix process")
 }
+
+func signalChan() <-chan os.Signal {
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+	return ch
+}
diff --git a/cmd/litestream/main_windows.go b/cmd/litestream/main_windows.go
index 31f1a0f..7f7c764 100644
--- a/cmd/litestream/main_windows.go
+++ b/cmd/litestream/main_windows.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"os/signal"
 
 	"golang.org/x/sys/windows"
 	"golang.org/x/sys/windows/svc"
@@ -103,3 +104,9 @@ func (w *eventlogWriter) Write(p []byte) (n int, err error) {
 	elog := (*eventlog.Log)(w)
 	return 0, elog.Info(1, string(p))
 }
+
+func signalChan() <-chan os.Signal {
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, os.Interrupt)
+	return ch
+}