Watch signals from internal server

Signed-off-by: Jesus Vazquez <jesus.vazquez@grafana.com>

pkg/appcommon/harness.go

@@ -153,9 +153,12 @@ func New(cfg Config, reg prometheus.Registerer, metricPrefix string, tracer open
 	}
 	app.Server = srv
 
+	signalHandler := stopsignal.NewSignalHandler(cfg.InternalServerConfig.ServerGracefulShutdownTimeout, logger)
+	cfg.InternalServerConfig.ReadinessProvider = signalHandler
+
 	app.Group.Add(app.Server.Handler())
 	app.Group.Add(internalserver.Handler(logger, cfg.InternalServerConfig))
-	app.Group.Add(stopsignal.Handler(logger, syscall.SIGTERM, syscall.SIGINT))
+	app.Group.Add(signalHandler.Handler(syscall.SIGTERM, syscall.SIGINT))
 
 	if err := registerVersionMetrics(reg, cfg.ServiceName, metricPrefix); err != nil {
 		return app, err

pkg/internalserver/internalserver_handler.go

@@ -24,6 +24,8 @@ type Config struct {
 	HTTPListenAddress             string        `yaml:"http_listen_address"`
 	HTTPListenPort                int           `yaml:"http_listen_port"`
 	ServerGracefulShutdownTimeout time.Duration `yaml:"graceful_shutdown_timeout"`
+
+	ReadinessProvider ReadinessProvider `yaml:"-"`
 }
 
 // RegisterFlags adds the flags required to config this to the given FlagSet
@@ -47,8 +49,7 @@ func Handler(logger log.Logger, cfg Config) (run func() error, stop func(error))
 	mux := http.NewServeMux()
 	mux.Handle("/metrics", promhttp.Handler())
 
-	signalHandler := newSignalHandler(cfg.ServerGracefulShutdownTimeout, logger)
-	mux.Handle("/healthz", http.HandlerFunc(NewReadinessHandler(signalHandler, logger)))
+	mux.Handle("/healthz", http.HandlerFunc(NewReadinessHandler(cfg.ReadinessProvider, logger)))
 
 	// Pprof.
 	mux.HandleFunc("/debug/pprof/", pprof.Index)

pkg/stopsignal/signal_handler.go

@@ -0,0 +1,64 @@
+package stopsignal
+
+import (
+	"os"
+	"os/signal"
+	"sync/atomic"
+	"time"
+
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+)
+
+func NewSignalHandler(shutdownDelay time.Duration, logger log.Logger) *SignalHandler {
+	return &SignalHandler{
+		quit:          make(chan struct{}),
+		shutdownDelay: shutdownDelay,
+		logger:        logger,
+	}
+}
+
+type SignalHandler struct {
+	quit          chan struct{}
+	ready         atomic.Bool
+	shutdownDelay time.Duration
+	logger        log.Logger
+}
+
+func (sh *SignalHandler) Handler(signals ...os.Signal) (run func() error, stop func(error)) {
+	sh.ready.Store(true)
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, signals...)
+
+	return func() error {
+			level.Info(sh.logger).Log("msg", "Waiting for stop signal...")
+			select {
+			case <-sh.quit:
+				return nil
+
+			case <-sigs:
+				level.Info(sh.logger).Log("msg", "=== received SIGINT/SIGTERM ===", "sleep", sh.shutdownDelay)
+
+				// Not ready anymore.
+				sh.ready.Store(false)
+				if sh.shutdownDelay > 0 {
+					time.Sleep(sh.shutdownDelay)
+				}
+
+				level.Info(sh.logger).Log("msg", "shutting down")
+				return nil
+			}
+		},
+		func(_ error) {
+			sh.Stop()
+		}
+}
+
+func (sh *SignalHandler) Stop() {
+	close(sh.quit)
+}
+
+func (sh *SignalHandler) Ready() bool {
+	return sh.ready.Load()
+}