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 (