Merge pull request #7 from iobear/catch-sigterm

Catch sigterm

cmd/dashgoat/alertmanager.go

@@ -44,6 +44,10 @@ type (
 
 func fromAlertmanager(c echo.Context) error {
 
+	if !isDashGoatReady() {
+		return c.NoContent(http.StatusServiceUnavailable)
+	}
+
 	dec := json.NewDecoder(c.Request().Body)
 	defer c.Request().Body.Close()
 

cmd/dashgoat/buddy.go

@@ -211,6 +211,10 @@ func findBuddy(buddyConfig []Buddy) {
 			}
 		}
 
+		if isDashGoatShutdown() {
+			break
+		}
+
 		if !isDashGoatReady() {
 			setDashGoatReady(true)
 		}

cmd/dashgoat/handler.go

@@ -18,6 +18,10 @@ import (
 // heartBeat update
 func heartBeat(c echo.Context) error {
 
+	if !isDashGoatReady() {
+		return c.NoContent(http.StatusServiceUnavailable)
+	}
+
 	ss.mutex.Lock()
 	defer ss.mutex.Unlock()
 
@@ -89,6 +93,10 @@ func heartBeat(c echo.Context) error {
 // updateStatus - service update
 func updateStatus(c echo.Context) error {
 
+	if !isDashGoatReady() {
+		return c.NoContent(http.StatusServiceUnavailable)
+	}
+
 	ss.mutex.Lock()
 	defer ss.mutex.Unlock()
 

cmd/dashgoat/hostfacts.go

@@ -29,9 +29,14 @@ type (
 		Items HostFact
 		mutex sync.RWMutex
 	}
+	DGshutdown struct {
+		shutdown bool
+		mutex    sync.RWMutex
+	}
 )
 
 var host_facts HostFacts
+var dg_shutdown DGshutdown
 
 func generateHostFacts() {
 	host_facts.mutex.Lock()
@@ -130,3 +135,20 @@ func setDashGoatReady(ready bool) {
 	defer host_facts.mutex.Unlock()
 	host_facts.Items.Ready = ready
 }
+
+func isDashGoatShutdown() bool {
+	dg_shutdown.mutex.RLock()
+	defer dg_shutdown.mutex.RUnlock()
+	return dg_shutdown.shutdown
+}
+
+func setDashGoatShutdown(state bool) {
+	dg_shutdown.mutex.Lock()
+	defer dg_shutdown.mutex.Unlock()
+	dg_shutdown.shutdown = state
+
+	if state {
+		setDashGoatReady(false)
+	}
+
+}

cmd/dashgoat/main.go

@@ -14,8 +14,11 @@ import (
 	"log/slog"
 	"net/http"
 	"os"
+	"os/signal"
 	"strings"
 	"sync"
+	"syscall"
+	"time"
 
 	"github.com/labstack/echo-contrib/echoprometheus"
 	"github.com/labstack/echo/v4"
@@ -87,6 +90,7 @@ func main() {
 	flag.StringVar(&config.IPport, "ipport", ":2000", "Specify <ip>:<port>")
 	flag.StringVar(&config.WebPath, "webpath", "/", "Specify added url http://host:port/<path> Default: /")
 	flag.StringVar(&config.UpdateKey, "updatekey", "changeme", "Specify key to API update")
+	flag.StringVar(&config.UrnKey, "urnkey", "", "Specify key for Heartbeat and Alertmanager")
 	flag.StringVar(&config.DashName, "dashname", "", "Dashboard name")
 	flag.StringVar(&configfile, "configfile", "dashgoat.yaml", "Name of configfile")
 	flag.StringVar(&buddy_cli.Url, "buddyurl", "", "Buddy url")
@@ -161,7 +165,19 @@ func main() {
 	go findBuddy(config.BuddyHosts)
 	go initPagerDuty()
 
+	//catch interrupt or sigterm
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
+	go func() {
+		for sig := range sigs {
+			logger.Error("main", "signal", sig)
+			fmt.Println(sig)
+			setDashGoatShutdown(true)
+			time.Sleep(2 * time.Second)
+			os.Exit(0)
+		}
+	}()
+
 	// Start server
 	e.Logger.Fatal(e.Start(config.IPport))
-
 }

cmd/dashgoat/timeoutprobe.go

@@ -34,6 +34,9 @@ func lostProbeTimer() {
 			}
 		}
 		time.Sleep(time.Duration(interval) * time.Second)
+		if isDashGoatShutdown() {
+			return
+		}
 	}
 
 }

cmd/dashgoat/ttl.go

@@ -41,6 +41,9 @@ func ttlHousekeeping() {
 	defer ticker.Stop()
 
 	for range ticker.C {
+		if isDashGoatShutdown() {
+			break
+		}
 		statusList := readStatusList()
 		currentUnixTimestamp := time.Now().Unix()
 		for key, serviceState := range statusList {