From 1d476c6d4d582cd9ba0d0500a98f97b756fae090 Mon Sep 17 00:00:00 2001
From: Ben McClelland <ben.mcclelland@versity.com>
Date: Sun, 25 Jun 2023 10:01:54 -0700
Subject: [PATCH] add signal handler for clean shutdown

---
 Makefile                 |  3 +++
 cmd/versitygw/main.go    | 26 ++++++++++++++++++++++---
 cmd/versitygw/posix.go   |  2 +-
 cmd/versitygw/scoutfs.go |  2 +-
 cmd/versitygw/singal.go  | 42 ++++++++++++++++++++++++++++++++++++++++
 5 files changed, 70 insertions(+), 5 deletions(-)
 create mode 100644 cmd/versitygw/singal.go

diff --git a/Makefile b/Makefile
index 3100315..81509c9 100644
--- a/Makefile
+++ b/Makefile
@@ -34,6 +34,9 @@ build: $(BIN)
 $(BIN):
 	$(GOBUILD) $(LDFLAGS) -o $(BIN) cmd/$(BIN)/*.go
 
+testbin:
+	$(GOBUILD) $(LDFLAGS) -o $(BIN) -cover -race cmd/$(BIN)/*.go
+
 .PHONY: test
 test: 
 	$(GOTEST) ./...
diff --git a/cmd/versitygw/main.go b/cmd/versitygw/main.go
index b307ffa..0f9bb06 100644
--- a/cmd/versitygw/main.go
+++ b/cmd/versitygw/main.go
@@ -15,6 +15,7 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"log"
@@ -47,6 +48,8 @@ var (
 )
 
 func main() {
+	setupSignalHandler()
+
 	app := initApp()
 
 	app.Commands = []*cli.Command{
@@ -56,7 +59,14 @@ func main() {
 		testCommand(),
 	}
 
-	if err := app.Run(os.Args); err != nil {
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		<-sigDone
+		fmt.Fprintf(os.Stderr, "terminating signal caught, shutting down\n")
+		cancel()
+	}()
+
+	if err := app.RunContext(ctx, os.Args); err != nil {
 		log.Fatal(err)
 	}
 }
@@ -134,7 +144,7 @@ func initFlags() []cli.Flag {
 	}
 }
 
-func runGateway(be backend.Backend, s auth.Storer) error {
+func runGateway(ctx *cli.Context, be backend.Backend, s auth.Storer) error {
 	app := fiber.New(fiber.Config{
 		AppName:      "versitygw",
 		ServerHeader: "VERSITYGW",
@@ -180,5 +190,15 @@ func runGateway(be backend.Backend, s auth.Storer) error {
 		return fmt.Errorf("init gateway: %v", err)
 	}
 
-	return srv.Serve()
+	c := make(chan error, 1)
+	go func() { c <- srv.Serve() }()
+
+	select {
+	case <-ctx.Done():
+		be.Shutdown()
+		return ctx.Err()
+	case err := <-c:
+		be.Shutdown()
+		return err
+	}
 }
diff --git a/cmd/versitygw/posix.go b/cmd/versitygw/posix.go
index ebd8204..8928b8f 100644
--- a/cmd/versitygw/posix.go
+++ b/cmd/versitygw/posix.go
@@ -49,5 +49,5 @@ func runPosix(ctx *cli.Context) error {
 		return fmt.Errorf("init posix: %v", err)
 	}
 
-	return runGateway(be, be)
+	return runGateway(ctx, be, be)
 }
diff --git a/cmd/versitygw/scoutfs.go b/cmd/versitygw/scoutfs.go
index 047b137..7327d4c 100644
--- a/cmd/versitygw/scoutfs.go
+++ b/cmd/versitygw/scoutfs.go
@@ -69,5 +69,5 @@ func runScoutfs(ctx *cli.Context) error {
 		return fmt.Errorf("init scoutfs: %v", err)
 	}
 
-	return runGateway(be, be)
+	return runGateway(ctx, be, be)
 }
diff --git a/cmd/versitygw/singal.go b/cmd/versitygw/singal.go
new file mode 100644
index 0000000..147e55e
--- /dev/null
+++ b/cmd/versitygw/singal.go
@@ -0,0 +1,42 @@
+// Copyright 2023 Versity Software
+// This file is licensed under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+var (
+	sigDone = make(chan bool, 1)
+)
+
+func setupSignalHandler() {
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
+
+	go func() {
+		for sig := range sigs {
+			fmt.Fprintf(os.Stderr, "caught signal %v\n", sig)
+			switch sig {
+			case syscall.SIGINT, syscall.SIGTERM:
+				sigDone <- true
+			case syscall.SIGHUP:
+			}
+		}
+	}()
+}