From 7e48d80133b35561f79137546fd1b5d530cb5db1 Mon Sep 17 00:00:00 2001
From: Nicholas Wiersma <nick@wiersma.co.za>
Date: Wed, 15 Nov 2023 14:14:21 +0200
Subject: [PATCH] feat: support pre health checks (#104)

---
 http/server.go      | 34 +++++++++++++++++++++++++------
 http/server_test.go | 49 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/http/server.go b/http/server.go
index 72f50f6..3315343 100644
--- a/http/server.go
+++ b/http/server.go
@@ -121,8 +121,28 @@ func (s *Server) Close() error {
 type HealthServerConfig struct {
 	Addr    string
 	Handler http.Handler
-	Stats   *statter.Statter
-	Log     *logger.Logger
+
+	ReadyzChecks []healthz.HealthChecker
+	LivezChecks  []healthz.HealthChecker
+
+	Stats *statter.Statter
+	Log   *logger.Logger
+}
+
+// AddHealthzChecks adds the given checks to the config.
+func (c *HealthServerConfig) AddHealthzChecks(checks ...healthz.HealthChecker) {
+	c.AddReadyzChecks(checks...)
+	c.AddLivezChecks(checks...)
+}
+
+// AddReadyzChecks adds the given checks to the config.
+func (c *HealthServerConfig) AddReadyzChecks(checks ...healthz.HealthChecker) {
+	c.ReadyzChecks = append(c.ReadyzChecks, checks...)
+}
+
+// AddLivezChecks adds the given checks to the config.
+func (c *HealthServerConfig) AddLivezChecks(checks ...healthz.HealthChecker) {
+	c.LivezChecks = append(c.LivezChecks, checks...)
 }
 
 // HealthServer is an HTTP server with healthz capabilities.
@@ -148,10 +168,12 @@ func NewHealthServer(ctx context.Context, cfg HealthServerConfig, opts ...SrvOpt
 	srv := NewServer(ctx, cfg.Addr, cfg.Handler, opts...)
 
 	return &HealthServer{
-		srv:       srv,
-		shudownCh: make(chan struct{}),
-		stats:     cfg.Stats,
-		log:       cfg.Log,
+		srv:          srv,
+		shudownCh:    make(chan struct{}),
+		readyzChecks: cfg.ReadyzChecks,
+		livezChecks:  cfg.LivezChecks,
+		stats:        cfg.Stats,
+		log:          cfg.Log,
 	}
 }
 
diff --git a/http/server_test.go b/http/server_test.go
index d082bad..d42baed 100644
--- a/http/server_test.go
+++ b/http/server_test.go
@@ -196,6 +196,55 @@ func TestHealthServer(t *testing.T) {
 	assert.Equal(t, "+ test ok\nlivez check passed", body)
 }
 
+func TestHealthServer_WithPreChecks(t *testing.T) {
+	stats := statter.New(statter.DiscardReporter, time.Minute)
+	log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Error)
+
+	lnCh := make(chan net.Listener, 1)
+	setTestHookServerServe(func(ln net.Listener) {
+		lnCh <- ln
+	})
+	t.Cleanup(func() { setTestHookServerServe(nil) })
+
+	check := healthz.NamedCheck("test", func(*http.Request) error {
+		return nil
+	})
+
+	cfg := &HealthServerConfig{
+		Addr:    "localhost:0",
+		Handler: http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}),
+
+		Stats: stats,
+		Log:   log,
+	}
+	cfg.AddHealthzChecks(check)
+
+	srv := NewHealthServer(context.Background(), *cfg)
+	srv.Serve(func(err error) {
+		require.NoError(t, err)
+	})
+	t.Cleanup(func() {
+		_ = srv.Close()
+	})
+
+	var ln net.Listener
+	select {
+	case <-time.After(30 * time.Second):
+		require.Fail(t, "Timed out waiting for server listener")
+	case ln = <-lnCh:
+	}
+
+	url := "http://" + ln.Addr().String() + "/readyz?verbose=1"
+	statusCode, body := requireDoRequest(t, url)
+	assert.Equal(t, statusCode, http.StatusOK)
+	assert.Equal(t, "+ test ok\n+ shutdown ok\nreadyz check passed", body)
+
+	url = "http://" + ln.Addr().String() + "/livez?verbose=1"
+	statusCode, body = requireDoRequest(t, url)
+	assert.Equal(t, statusCode, http.StatusOK)
+	assert.Equal(t, "+ test ok\nlivez check passed", body)
+}
+
 func TestHealthServer_ShutdownCausesReadyzCheckToFail(t *testing.T) {
 	stats := statter.New(statter.DiscardReporter, time.Minute)
 	log := logger.New(io.Discard, logger.LogfmtFormat(), logger.Error)