feat: add healthz and server (#101)

.golangci.yml

@@ -33,6 +33,7 @@ linters:
     - gochecknoinits
     - goerr113
     - gomnd
+    - ireturn
     - nlreturn
     - varnamelen
     - wrapcheck

http/handlers.go

@@ -3,21 +3,31 @@ package http
 import "net/http"
 
 // OK replies to the request with an HTTP 200 ok reply.
+//
+// Deprecated: Use healthz instead.
 func OK(rw http.ResponseWriter, _ *http.Request) { rw.WriteHeader(http.StatusOK) }
 
 // OKHandler returns a simple request handler
 // that replies to each request with a “200 OK” reply.
+//
+// Deprecated: Use healthz instead.
 func OKHandler() http.Handler { return http.HandlerFunc(OK) }
 
 // DefaultHealthPath is the default HTTP path for checking health.
+//
+// Deprecated: Use healthz instead.
 var DefaultHealthPath = "/health"
 
 // Health represents an object that can check its health.
+//
+// Deprecated: Use healthz instead.
 type Health interface {
 	IsHealthy() error
 }
 
 // NewHealthHandler returns a handler for application health checking.
+//
+// Deprecated: Use healthz instead.
 func NewHealthHandler(v ...Health) http.HandlerFunc {
 	return func(rw http.ResponseWriter, r *http.Request) {
 		for _, h := range v {

http/healthz/healthz.go

@@ -0,0 +1,78 @@
+// Package healthz provides HTTP healthz handling.
+package healthz
+
+import (
+	"bytes"
+	"fmt"
+	"net/http"
+)
+
+// HealthChecker represents a named health checker.
+type HealthChecker interface {
+	Name() string
+	Check(*http.Request) error
+}
+
+type healthCheck struct {
+	name  string
+	check func(*http.Request) error
+}
+
+// NamedCheck returns a named health check.
+func NamedCheck(name string, check func(*http.Request) error) HealthChecker {
+	return &healthCheck{
+		name:  name,
+		check: check,
+	}
+}
+
+func (c healthCheck) Name() string { return c.name }
+
+func (c healthCheck) Check(req *http.Request) error { return c.check(req) }
+
+// PingHealth returns true when called.
+var PingHealth HealthChecker = ping{}
+
+type ping struct{}
+
+func (c ping) Name() string { return "ping" }
+
+func (c ping) Check(_ *http.Request) error { return nil }
+
+// Handler returns an HTTP check handler.
+func Handler(name string, errFn func(string), checks ...HealthChecker) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+		var (
+			checkOutput     bytes.Buffer
+			failedChecks    []string
+			failedLogOutput bytes.Buffer
+		)
+		for _, check := range checks {
+			if err := check.Check(req); err != nil {
+				_, _ = fmt.Fprintf(&checkOutput, "- %s failed\n", check.Name())
+				_, _ = fmt.Fprintf(&failedLogOutput, "%s failed: %v\n", check.Name(), err)
+				failedChecks = append(failedChecks, check.Name())
+				continue
+			}
+
+			_, _ = fmt.Fprintf(&checkOutput, "+ %s ok\n", check.Name())
+		}
+
+		if len(failedChecks) > 0 {
+			errFn(failedLogOutput.String())
+			http.Error(rw,
+				fmt.Sprintf("%s%s check failed", checkOutput.String(), name),
+				http.StatusInternalServerError,
+			)
+			return
+		}
+
+		if _, found := req.URL.Query()["verbose"]; !found {
+			_, _ = fmt.Fprint(rw, "ok")
+			return
+		}
+
+		_, _ = checkOutput.WriteTo(rw)
+		_, _ = fmt.Fprintf(rw, "%s check passed", name)
+	})
+}

http/healthz/healthz_test.go

@@ -0,0 +1,65 @@
+package healthz_test
+
+import (
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/hamba/pkg/v2/http/healthz"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHandler(t *testing.T) {
+	goodCheck := healthz.NamedCheck("good", func(*http.Request) error { return nil })
+
+	var gotOutput string
+	h := healthz.Handler("readyz", func(output string) {
+		gotOutput = output
+	}, goodCheck)
+
+	req := httptest.NewRequest(http.MethodGet, "/readyz", nil)
+	rec := httptest.NewRecorder()
+
+	h.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "", gotOutput)
+	assert.Equal(t, `ok`, rec.Body.String())
+}
+func TestHandler_Verbose(t *testing.T) {
+	goodCheck := healthz.NamedCheck("good", func(*http.Request) error { return nil })
+
+	var gotOutput string
+	h := healthz.Handler("readyz", func(output string) {
+		gotOutput = output
+	}, goodCheck)
+
+	req := httptest.NewRequest(http.MethodGet, "/readyz?verbose=1", nil)
+	rec := httptest.NewRecorder()
+
+	h.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusOK, rec.Code)
+	assert.Equal(t, "", gotOutput)
+	assert.Equal(t, "+ good ok\nreadyz check passed", rec.Body.String())
+}
+
+func TestHandler_WithFailingChecks(t *testing.T) {
+	goodCheck := healthz.NamedCheck("good", func(*http.Request) error { return nil })
+	badCheck := healthz.NamedCheck("bad", func(*http.Request) error { return errors.New("test error") })
+
+	var gotOutput string
+	h := healthz.Handler("readyz", func(output string) {
+		gotOutput = output
+	}, goodCheck, badCheck)
+
+	req := httptest.NewRequest(http.MethodGet, "/readyz", nil)
+	rec := httptest.NewRecorder()
+
+	h.ServeHTTP(rec, req)
+
+	assert.Equal(t, http.StatusInternalServerError, rec.Code)
+	assert.Equal(t, "bad failed: test error\n", gotOutput)
+	assert.Equal(t, "+ good ok\n- bad failed\nreadyz check failed\n", rec.Body.String())
+}

http/server.go

@@ -4,10 +4,18 @@ import (
 	"context"
 	"crypto/tls"
 	"errors"
+	"fmt"
 	"net"
 	"net/http"
+	"strings"
+	"sync"
 	"time"
 
+	"github.com/hamba/logger/v2"
+	lctx "github.com/hamba/logger/v2/ctx"
+	"github.com/hamba/pkg/v2/http/healthz"
+	"github.com/hamba/pkg/v2/http/middleware"
+	"github.com/hamba/statter/v2"
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
 )
@@ -42,7 +50,6 @@ func WithH2C() SrvOptFunc {
 		h2s := &http2.Server{
 			IdleTimeout: 120 * time.Second,
 		}
-
 		srv.Handler = h2c.NewHandler(srv.Handler, h2s)
 	}
 }
@@ -53,7 +60,7 @@ type Server struct {
 	srv *http.Server
 }
 
-// NewServer returns a server.
+// NewServer returns a server with the base context ctx.
 func NewServer(ctx context.Context, addr string, h http.Handler, opts ...SrvOptFunc) *Server {
 	srv := &http.Server{
 		BaseContext: func(_ net.Listener) context.Context {
@@ -97,3 +104,158 @@ func (s *Server) Shutdown(timeout time.Duration) error {
 func (s *Server) Close() error {
 	return s.srv.Close()
 }
+
+// HealthServerConfig configures a HealthServer.
+type HealthServerConfig struct {
+	Addr    string
+	Handler http.Handler
+	Stats   *statter.Statter
+	Log     *logger.Logger
+}
+
+// HealthServer is an HTTP server with healthz capabilities.
+type HealthServer struct {
+	srv *Server
+
+	shudownCh chan struct{}
+
+	readyzMu        sync.Mutex
+	readyzInstalled bool
+	readyzChecks    []healthz.HealthChecker
+
+	livezMu        sync.Mutex
+	livezInstalled bool
+	livezChecks    []healthz.HealthChecker
+
+	stats *statter.Statter
+	log   *logger.Logger
+}
+
+// NewHealthServer returns an HTTP server with healthz capabilities.
+func NewHealthServer(ctx context.Context, cfg HealthServerConfig, opts ...SrvOptFunc) *HealthServer {
+	srv := NewServer(ctx, cfg.Addr, cfg.Handler, opts...)
+
+	return &HealthServer{
+		srv:       srv,
+		shudownCh: make(chan struct{}),
+		stats:     cfg.Stats,
+		log:       cfg.Log,
+	}
+}
+
+// AddHealthzChecks adds health checks to both readyz and livez.
+func (s *HealthServer) AddHealthzChecks(checks ...healthz.HealthChecker) error {
+	if err := s.AddReadyzChecks(checks...); err != nil {
+		return err
+	}
+	return s.AddLivezChecks(checks...)
+}
+
+// AddReadyzChecks adds health checks to readyz.
+func (s *HealthServer) AddReadyzChecks(checks ...healthz.HealthChecker) error {
+	s.readyzMu.Lock()
+	defer s.readyzMu.Unlock()
+	if s.readyzInstalled {
+		return errors.New("could not add checks as readyz has already been installed")
+	}
+	s.readyzChecks = append(s.readyzChecks, checks...)
+	return nil
+}
+
+// AddLivezChecks adds health checks to livez.
+func (s *HealthServer) AddLivezChecks(checks ...healthz.HealthChecker) error {
+	s.livezMu.Lock()
+	defer s.livezMu.Unlock()
+	if s.livezInstalled {
+		return errors.New("could not add checks as livez has already been installed")
+	}
+	s.livezChecks = append(s.livezChecks, checks...)
+	return nil
+}
+
+// Serve installs the health checks and starts the server in a non-blocking way.
+func (s *HealthServer) Serve(errFn func(error)) {
+	s.installChecks()
+
+	s.srv.Serve(errFn)
+}
+
+func (s *HealthServer) installChecks() {
+	mux := http.NewServeMux()
+	s.installLivezChecks(mux)
+
+	// When shutdown is started, the readyz check should start failing.
+	if err := s.AddReadyzChecks(shutdownCheck{ch: s.shudownCh}); err != nil {
+		s.log.Error("Could not install readyz shutdown check", lctx.Err(err))
+	}
+	s.installReadyzChecks(mux)
+
+	mux.Handle("/", s.srv.srv.Handler)
+	s.srv.srv.Handler = mux
+}
+
+func (s *HealthServer) installReadyzChecks(mux *http.ServeMux) {
+	s.readyzMu.Lock()
+	defer s.readyzMu.Unlock()
+	s.readyzInstalled = true
+	s.installCheckers(mux, "/readyz", s.readyzChecks)
+}
+
+func (s *HealthServer) installLivezChecks(mux *http.ServeMux) {
+	s.livezMu.Lock()
+	defer s.livezMu.Unlock()
+	s.livezInstalled = true
+	s.installCheckers(mux, "/livez", s.livezChecks)
+}
+
+func (s *HealthServer) installCheckers(mux *http.ServeMux, path string, checks []healthz.HealthChecker) {
+	if len(checks) == 0 {
+		checks = []healthz.HealthChecker{healthz.PingHealth}
+	}
+
+	s.log.Info("Installing health checkers",
+		lctx.Str("path", path),
+		lctx.Str("checks", strings.Join(checkNames(checks), ",")),
+	)
+
+	name := strings.TrimPrefix(path, "/")
+	h := healthz.Handler(name, func(output string) {
+		s.log.Info(fmt.Sprintf("%s check failed\n%s", name, output))
+	}, checks...)
+	mux.Handle(path, middleware.WithStats(name, s.stats, h))
+}
+
+// Shutdown attempts to close all server connections.
+func (s *HealthServer) Shutdown(timeout time.Duration) error {
+	close(s.shudownCh)
+
+	return s.srv.Shutdown(timeout)
+}
+
+// Close closes the server.
+func (s *HealthServer) Close() error {
+	return s.srv.Close()
+}
+
+func checkNames(checks []healthz.HealthChecker) []string {
+	names := make([]string, len(checks))
+	for i, check := range checks {
+		names[i] = check.Name()
+	}
+	return names
+}
+
+type shutdownCheck struct {
+	ch <-chan struct{}
+}
+
+func (s shutdownCheck) Name() string { return "shutdown" }
+
+func (s shutdownCheck) Check(*http.Request) error {
+	select {
+	case <-s.ch:
+		return errors.New("server is shutting down")
+	default:
+		return nil
+	}
+}