From 8e2a2c4dbc4ce539ff244375867d2fc5bb4801eb Mon Sep 17 00:00:00 2001
From: TwinProduction <twin@twinnation.org>
Date: Fri, 5 Feb 2021 20:45:28 -0500
Subject: [PATCH] Implement graceful shutdown - Shutdown the HTTP server before
 exiting - Persist data to store before exiting, if applicable

---
 controller/controller.go            | 11 ++++++++++-
 controller/controller_test.go       | 10 ++++++++++
 main.go                             | 22 +++++++++++++++++++++-
 storage/storage.go                  | 14 +++++++++++++-
 storage/store/memory/memory.go      | 17 +++--------------
 storage/store/memory/memory_test.go | 19 +++++++++++++++++++
 storage/store/store.go              |  3 +++
 7 files changed, 79 insertions(+), 17 deletions(-)

diff --git a/controller/controller.go b/controller/controller.go
index 8c65f124..6a1f05b5 100644
--- a/controller/controller.go
+++ b/controller/controller.go
@@ -3,6 +3,7 @@ package controller
 import (
 	"bytes"
 	"compress/gzip"
+	"context"
 	"encoding/json"
 	"fmt"
 	"log"
@@ -59,7 +60,15 @@ func Handle() {
 	if os.Getenv("ROUTER_TEST") == "true" {
 		return
 	}
-	log.Fatal(server.ListenAndServe())
+	log.Println("[controller][Handle]", server.ListenAndServe())
+}
+
+// Shutdown stops the server
+func Shutdown() {
+	if server != nil {
+		_ = server.Shutdown(context.TODO())
+		server = nil
+	}
 }
 
 // CreateRouter creates the router for the http server
diff --git a/controller/controller_test.go b/controller/controller_test.go
index e07975b4..bd828bdd 100644
--- a/controller/controller_test.go
+++ b/controller/controller_test.go
@@ -156,6 +156,7 @@ func TestHandle(t *testing.T) {
 	config.Set(cfg)
 	_ = os.Setenv("ROUTER_TEST", "true")
 	_ = os.Setenv("ENVIRONMENT", "dev")
+	defer os.Clearenv()
 	Handle()
 	request, _ := http.NewRequest("GET", "/health", nil)
 	responseRecorder := httptest.NewRecorder()
@@ -167,3 +168,12 @@ func TestHandle(t *testing.T) {
 		t.Fatal("server should've been set (but because we set ROUTER_TEST, it shouldn't have been started)")
 	}
 }
+
+func TestShutdown(t *testing.T) {
+	// Pretend that we called controller.Handle(), which initializes the server variable
+	server = &http.Server{}
+	Shutdown()
+	if server != nil {
+		t.Error("server should've been shut down")
+	}
+}
diff --git a/main.go b/main.go
index 4fca3e91..b3d9d1e4 100644
--- a/main.go
+++ b/main.go
@@ -1,17 +1,37 @@
 package main
 
 import (
+	"log"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/TwinProduction/gatus/config"
 	"github.com/TwinProduction/gatus/controller"
+	"github.com/TwinProduction/gatus/storage"
 	"github.com/TwinProduction/gatus/watchdog"
 )
 
 func main() {
 	cfg := loadConfiguration()
 	go watchdog.Monitor(cfg)
-	controller.Handle()
+	go controller.Handle()
+	// Wait for termination signal
+	sig := make(chan os.Signal, 1)
+	done := make(chan bool, 1)
+	signal.Notify(sig, os.Interrupt, os.Kill, syscall.SIGTERM)
+	go func() {
+		<-sig
+		log.Println("Received interruption signal, attempting to gracefully shut down")
+		controller.Shutdown()
+		err := storage.Get().Save()
+		if err != nil {
+			log.Println("Failed to save storage provider:", err.Error())
+		}
+		done <- true
+	}()
+	<-done
+	log.Println("Shutting down")
 }
 
 func loadConfiguration() *config.Config {
diff --git a/storage/storage.go b/storage/storage.go
index abd51e16..8811ebd0 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -42,7 +42,19 @@ func Initialize(cfg *Config) error {
 		if err != nil {
 			return err
 		}
-		go provider.(*memory.Store).AutoSave(7 * time.Minute)
+		go autoSave(7 * time.Minute)
 	}
 	return nil
 }
+
+// autoSave automatically calls the Save function of the provider at every interval
+func autoSave(interval time.Duration) {
+	for {
+		time.Sleep(interval)
+		log.Printf("[storage][autoSave] Saving")
+		err := provider.Save()
+		if err != nil {
+			log.Println("[storage][autoSave] Save failed:", err.Error())
+		}
+	}
+}
diff --git a/storage/store/memory/memory.go b/storage/store/memory/memory.go
index 75c991c3..bb4f65a1 100644
--- a/storage/store/memory/memory.go
+++ b/storage/store/memory/memory.go
@@ -3,8 +3,6 @@ package memory
 import (
 	"encoding/gob"
 	"encoding/json"
-	"log"
-	"time"
 
 	"github.com/TwinProduction/gatus/core"
 	"github.com/TwinProduction/gatus/util"
@@ -94,17 +92,8 @@ func (s *Store) Clear() {
 
 // Save persists the cache to the store file
 func (s *Store) Save() error {
-	return s.cache.SaveToFile(s.file)
-}
-
-// AutoSave automatically calls the Save function at every interval
-func (s *Store) AutoSave(interval time.Duration) {
-	for {
-		time.Sleep(interval)
-		log.Printf("[memory][AutoSave] Persisting data to file")
-		err := s.Save()
-		if err != nil {
-			log.Printf("[memory][AutoSave] failed to save to file=%s: %s", s.file, err.Error())
-		}
+	if len(s.file) > 0 {
+		return s.cache.SaveToFile(s.file)
 	}
+	return nil
 }
diff --git a/storage/store/memory/memory_test.go b/storage/store/memory/memory_test.go
index 3c92d40c..20d1a851 100644
--- a/storage/store/memory/memory_test.go
+++ b/storage/store/memory/memory_test.go
@@ -250,3 +250,22 @@ func TestStore_DeleteAllServiceStatusesNotInKeys(t *testing.T) {
 		t.Error("firstService should still exist")
 	}
 }
+
+func TestStore_Save(t *testing.T) {
+	files := []string{
+		"",
+		t.TempDir() + "/test.db",
+	}
+	for _, file := range files {
+		t.Run(file, func(t *testing.T) {
+			store, err := NewStore(file)
+			if err != nil {
+				t.Fatal("expected no error, got", err.Error())
+			}
+			err = store.Save()
+			if err != nil {
+				t.Fatal("expected no error, got", err.Error())
+			}
+		})
+	}
+}
diff --git a/storage/store/store.go b/storage/store/store.go
index f969e227..076ddfc2 100644
--- a/storage/store/store.go
+++ b/storage/store/store.go
@@ -26,6 +26,9 @@ type Store interface {
 
 	// Clear deletes everything from the store
 	Clear()
+
+	// Save persists the data if and where it needs to be persisted
+	Save() error
 }
 
 var (