From dec2ca5228321d55a211b7b4476966a379cdeafb Mon Sep 17 00:00:00 2001
From: Jean-Francois Chevrette <jfchevrette@gmail.com>
Date: Thu, 13 Feb 2020 16:02:41 -0500
Subject: [PATCH] refactor main so we catch SIGTERM

---
 main.go | 31 +++++++++++++++++++++++++++++--
 1 file changed, 29 insertions(+), 2 deletions(-)

diff --git a/main.go b/main.go
index 95b8464..49542c8 100644
--- a/main.go
+++ b/main.go
@@ -3,6 +3,8 @@ package main
 import (
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/credentials"
@@ -27,6 +29,10 @@ var (
 )
 
 func main() {
+	os.Exit(run())
+}
+
+func run() int {
 	awsRegion := os.Getenv("AWS_REGION")
 	if awsRegion == "" {
 		log.Fatalln("AWS_REGION has to be defined")
@@ -64,6 +70,27 @@ func main() {
              </body>
              </html>`))
 	})
-	log.Infoln("Starting HTTP server on", *listenAddress)
-	log.Fatal(http.ListenAndServe(*listenAddress, nil))
+
+	srv := http.Server{Addr: *listenAddress}
+	srvc := make(chan struct{})
+	term := make(chan os.Signal, 1)
+	signal.Notify(term, os.Interrupt, syscall.SIGTERM)
+
+	go func() {
+		log.Infoln("Starting HTTP server on", *listenAddress)
+		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
+			log.Errorf("Error starting HTTP server: %v", err)
+			close(srvc)
+		}
+	}()
+
+	for {
+		select {
+		case <-term:
+			log.Infoln("Received SIGTERM, exiting gracefully...")
+			return 0
+		case <-srvc:
+			return 1
+		}
+	}
 }