From ed4872f37fb1d8b6b79936e58f25b90c3f2a0472 Mon Sep 17 00:00:00 2001
From: Jesus Vazquez <jesus.vazquez@grafana.com>
Date: Fri, 15 Sep 2023 12:27:11 +0200
Subject: [PATCH] Update internalserver healthz http handler

Signed-off-by: Jesus Vazquez <jesus.vazquez@grafana.com>
---
 pkg/internalserver/internalserver_handler.go |  6 +-
 pkg/internalserver/ready_handler.go          | 38 ++++++++++++
 pkg/internalserver/signal_handler.go         | 62 ++++++++++++++++++++
 3 files changed, 102 insertions(+), 4 deletions(-)
 create mode 100644 pkg/internalserver/ready_handler.go
 create mode 100644 pkg/internalserver/signal_handler.go

diff --git a/pkg/internalserver/internalserver_handler.go b/pkg/internalserver/internalserver_handler.go
index 967aa74..27cc2e7 100644
--- a/pkg/internalserver/internalserver_handler.go
+++ b/pkg/internalserver/internalserver_handler.go
@@ -47,10 +47,8 @@ func Handler(logger log.Logger, cfg Config) (run func() error, stop func(error))
 	mux := http.NewServeMux()
 	mux.Handle("/metrics", promhttp.Handler())
 
-	mux.Handle("/healthz", http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-		rw.WriteHeader(http.StatusOK)
-		_, _ = rw.Write([]byte("ok"))
-	}))
+	signalHandler := newSignalHandler(cfg.ServerGracefulShutdownTimeout, logger)
+	mux.Handle("/healthz", http.HandlerFunc(NewReadinessHandler(signalHandler, logger)))
 
 	// Pprof.
 	mux.HandleFunc("/debug/pprof/", pprof.Index)
diff --git a/pkg/internalserver/ready_handler.go b/pkg/internalserver/ready_handler.go
new file mode 100644
index 0000000..cecef03
--- /dev/null
+++ b/pkg/internalserver/ready_handler.go
@@ -0,0 +1,38 @@
+package internalserver
+
+import (
+	"net/http"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+type ReadinessProvider interface {
+	Ready() bool
+}
+
+// NewReadinessHandler returns an endpoint that returns a simple 200 to denote
+// the web server is active
+func NewReadinessHandler(ready ReadinessProvider, logger log.Logger) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		var err error
+
+		if ready.Ready() {
+			w.WriteHeader(http.StatusOK)
+			_, err = w.Write([]byte("OK"))
+		} else {
+			w.WriteHeader(http.StatusInternalServerError)
+			_, err = w.Write([]byte("Not ready"))
+		}
+
+		if err != nil {
+			level.Error(logger).Log("msg", "ready endpoint error", "err", err)
+		}
+	}
+}
+
+type AlwaysReady struct{}
+
+func (AlwaysReady) Ready() bool {
+	return true
+}
diff --git a/pkg/internalserver/signal_handler.go b/pkg/internalserver/signal_handler.go
new file mode 100644
index 0000000..c942b02
--- /dev/null
+++ b/pkg/internalserver/signal_handler.go
@@ -0,0 +1,62 @@
+package internalserver
+
+import (
+	"os"
+	"os/signal"
+	"sync/atomic"
+	"syscall"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+func newSignalHandler(shutdownDelay time.Duration, logger log.Logger) *serverSignalHandler {
+	return &serverSignalHandler{
+		quit:          make(chan struct{}),
+		shutdownDelay: shutdownDelay,
+		logger:        logger,
+	}
+}
+
+type serverSignalHandler struct {
+	quit          chan struct{}
+	ready         atomic.Bool
+	shutdownDelay time.Duration
+	logger        log.Logger
+}
+
+func (dh *serverSignalHandler) Loop() {
+	dh.ready.Store(true)
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+	defer signal.Stop(sigs)
+
+	for {
+		select {
+		case <-dh.quit:
+			return
+
+		case <-sigs:
+			level.Info(dh.logger).Log("msg", "=== received SIGINT/SIGTERM ===", "sleep", dh.shutdownDelay)
+
+			// Not ready anymore.
+			dh.ready.Store(false)
+			if dh.shutdownDelay > 0 {
+				time.Sleep(dh.shutdownDelay)
+			}
+
+			level.Info(dh.logger).Log("msg", "shutting down")
+			return
+		}
+	}
+}
+
+func (dh *serverSignalHandler) Stop() {
+	close(dh.quit)
+}
+
+func (dh *serverSignalHandler) Ready() bool {
+	return dh.ready.Load()
+}