Plumb contexts into health checks (#4141)

health/api/api_test.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -59,7 +60,7 @@ func TestPOSTDownHandlerChangeStatus(t *testing.T) {
 		t.Errorf("Did not get a 200.")
 	}
 
-	if len(health.CheckStatus()) != 1 {
+	if len(health.CheckStatus(context.Background())) != 1 {
 		t.Errorf("DownHandler didn't add an error check.")
 	}
 }
@@ -80,7 +81,7 @@ func TestPOSTUpHandlerChangeStatus(t *testing.T) {
 		t.Errorf("Did not get a 200.")
 	}
 
-	if len(health.CheckStatus()) != 0 {
+	if len(health.CheckStatus(context.Background())) != 0 {
 		t.Errorf("UpHandler didn't remove the error check.")
 	}
 }

health/checks/checks.go

@@ -1,13 +1,13 @@
 package checks
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"net"
 	"net/http"
 	"os"
 	"path/filepath"
-	"strconv"
 	"time"
 
 	"github.com/distribution/distribution/v3/health"
@@ -16,7 +16,7 @@ import (
 // FileChecker checks the existence of a file and returns an error
 // if the file exists.
 func FileChecker(f string) health.Checker {
-	return health.CheckFunc(func() error {
+	return health.CheckFunc(func(context.Context) error {
 		absoluteFilePath, err := filepath.Abs(f)
 		if err != nil {
 			return fmt.Errorf("failed to get absolute path for %q: %v", f, err)
@@ -36,13 +36,13 @@ func FileChecker(f string) health.Checker {
 // HTTPChecker does a HEAD request and verifies that the HTTP status code
 // returned matches statusCode.
 func HTTPChecker(r string, statusCode int, timeout time.Duration, headers http.Header) health.Checker {
-	return health.CheckFunc(func() error {
+	return health.CheckFunc(func(ctx context.Context) error {
 		client := http.Client{
 			Timeout: timeout,
 		}
-		req, err := http.NewRequest(http.MethodHead, r, nil)
+		req, err := http.NewRequestWithContext(ctx, http.MethodHead, r, nil)
 		if err != nil {
-			return errors.New("error creating request: " + r)
+			return fmt.Errorf("%v: error creating request: %w", r, err)
 		}
 		for headerName, headerValues := range headers {
 			for _, headerValue := range headerValues {
@@ -51,22 +51,23 @@ func HTTPChecker(r string, statusCode int, timeout time.Duration, headers http.H
 		}
 		response, err := client.Do(req)
 		if err != nil {
-			return errors.New("error while checking: " + r)
+			return fmt.Errorf("%v: error while checking: %w", r, err)
 		}
 		defer response.Body.Close()
 		if response.StatusCode != statusCode {
-			return errors.New("downstream service returned unexpected status: " + strconv.Itoa(response.StatusCode))
+			return fmt.Errorf("%v: downstream service returned unexpected status: %d", r, response.StatusCode)
 		}
 		return nil
 	})
 }
 
 // TCPChecker attempts to open a TCP connection.
 func TCPChecker(addr string, timeout time.Duration) health.Checker {
-	return health.CheckFunc(func() error {
-		conn, err := net.DialTimeout("tcp", addr, timeout)
+	return health.CheckFunc(func(ctx context.Context) error {
+		d := net.Dialer{Timeout: timeout}
+		conn, err := d.DialContext(ctx, "tcp", addr)
 		if err != nil {
-			return errors.New("connection to " + addr + " failed")
+			return fmt.Errorf("%v: connection failed: %w", addr, err)
 		}
 		conn.Close()
 		return nil

health/checks/checks_test.go

@@ -1,25 +1,26 @@
 package checks
 
 import (
+	"context"
 	"testing"
 )
 
 func TestFileChecker(t *testing.T) {
-	if err := FileChecker("/tmp").Check(); err == nil {
+	if err := FileChecker("/tmp").Check(context.Background()); err == nil {
 		t.Errorf("/tmp was expected as exists")
 	}
 
-	if err := FileChecker("NoSuchFileFromMoon").Check(); err != nil {
+	if err := FileChecker("NoSuchFileFromMoon").Check(context.Background()); err != nil {
 		t.Errorf("NoSuchFileFromMoon was expected as not exists, error:%v", err)
 	}
 }
 
 func TestHTTPChecker(t *testing.T) {
-	if err := HTTPChecker("https://www.google.cybertron", 200, 0, nil).Check(); err == nil {
+	if err := HTTPChecker("https://www.google.cybertron", 200, 0, nil).Check(context.Background()); err == nil {
 		t.Errorf("Google on Cybertron was expected as not exists")
 	}
 
-	if err := HTTPChecker("https://www.google.pt", 200, 0, nil).Check(); err != nil {
+	if err := HTTPChecker("https://www.google.pt", 200, 0, nil).Check(context.Background()); err != nil {
 		t.Errorf("Google at Portugal was expected as exists, error:%v", err)
 	}
 }

health/health.go

@@ -1,7 +1,9 @@
 package health
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 	"sync"
@@ -35,17 +37,17 @@ var DefaultRegistry *Registry
 // Checker is the interface for a Health Checker
 type Checker interface {
 	// Check returns nil if the service is okay.
-	Check() error
+	Check(context.Context) error
 }
 
 // CheckFunc is a convenience type to create functions that implement
 // the Checker interface
-type CheckFunc func() error
+type CheckFunc func(context.Context) error
 
 // Check Implements the Checker interface to allow for any func() error method
 // to be passed as a Checker
-func (cf CheckFunc) Check() error {
-	return cf()
+func (cf CheckFunc) Check(ctx context.Context) error {
+	return cf(ctx)
 }
 
 // Updater implements a health check that is explicitly set.
@@ -66,7 +68,7 @@ type updater struct {
 }
 
 // Check implements the Checker interface
-func (u *updater) Check() error {
+func (u *updater) Check(context.Context) error {
 	u.mu.Lock()
 	defer u.mu.Unlock()
 
@@ -99,11 +101,11 @@ type thresholdUpdater struct {
 }
 
 // Check implements the Checker interface
-func (tu *thresholdUpdater) Check() error {
+func (tu *thresholdUpdater) Check(context.Context) error {
 	tu.mu.Lock()
 	defer tu.mu.Unlock()
 
-	if tu.count >= tu.threshold {
+	if tu.count >= tu.threshold || errors.As(tu.status, new(pollingTerminatedErr)) {
 		return tu.status
 	}
 
@@ -127,47 +129,44 @@ func (tu *thresholdUpdater) Update(status error) {
 
 // NewThresholdStatusUpdater returns a new thresholdUpdater
 func NewThresholdStatusUpdater(t int) Updater {
-	return &thresholdUpdater{threshold: t}
+	if t > 0 {
+		return &thresholdUpdater{threshold: t}
+	}
+	return NewStatusUpdater()
 }
 
-// PeriodicChecker wraps an updater to provide a periodic checker
-func PeriodicChecker(check Checker, period time.Duration) Checker {
-	u := NewStatusUpdater()
-	go func() {
-		t := time.NewTicker(period)
-		defer t.Stop()
-		for {
-			<-t.C
-			u.Update(check.Check())
-		}
-	}()
+type pollingTerminatedErr struct{ Err error }
 
-	return u
+func (e pollingTerminatedErr) Error() string {
+	return fmt.Sprintf("health: check is not polled: %v", e.Err)
 }
 
-// PeriodicThresholdChecker wraps an updater to provide a periodic checker that
-// uses a threshold before it changes status
-func PeriodicThresholdChecker(check Checker, period time.Duration, threshold int) Checker {
-	tu := NewThresholdStatusUpdater(threshold)
-	go func() {
-		t := time.NewTicker(period)
-		defer t.Stop()
-		for {
-			<-t.C
-			tu.Update(check.Check())
+func (e pollingTerminatedErr) Unwrap() error { return e.Err }
+
+// Poll periodically polls the checker c at interval and updates the updater u
+// with the result. The checker is called with ctx as the context. When ctx is
+// done, Poll updates the updater with ctx.Err() and returns.
+func Poll(ctx context.Context, u Updater, c Checker, interval time.Duration) {
+	t := time.NewTicker(interval)
+	defer t.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			u.Update(pollingTerminatedErr{Err: ctx.Err()})
+			return
+		case <-t.C:
+			u.Update(c.Check(ctx))
 		}
-	}()
-
-	return tu
+	}
 }
 
 // CheckStatus returns a map with all the current health check errors
-func (registry *Registry) CheckStatus() map[string]string { // TODO(stevvooe) this needs a proper type
+func (registry *Registry) CheckStatus(ctx context.Context) map[string]string { // TODO(stevvooe) this needs a proper type
 	registry.mu.RLock()
 	defer registry.mu.RUnlock()
 	statusKeys := make(map[string]string)
 	for k, v := range registry.registeredChecks {
-		err := v.Check()
+		err := v.Check(ctx)
 		if err != nil {
 			statusKeys[k] = err.Error()
 		}
@@ -178,8 +177,8 @@ func (registry *Registry) CheckStatus() map[string]string { // TODO(stevvooe) th
 
 // CheckStatus returns a map with all the current health check errors from the
 // default registry.
-func CheckStatus() map[string]string {
-	return DefaultRegistry.CheckStatus()
+func CheckStatus(ctx context.Context) map[string]string {
+	return DefaultRegistry.CheckStatus(ctx)
 }
 
 // Register associates the checker with the provided name.
@@ -203,47 +202,23 @@ func Register(name string, check Checker) {
 }
 
 // RegisterFunc allows the convenience of registering a checker directly from
-// an arbitrary func() error.
-func (registry *Registry) RegisterFunc(name string, check func() error) {
-	registry.Register(name, CheckFunc(check))
+// an arbitrary func(context.Context) error.
+func (registry *Registry) RegisterFunc(name string, check CheckFunc) {
+	registry.Register(name, check)
 }
 
 // RegisterFunc allows the convenience of registering a checker in the default
-// registry directly from an arbitrary func() error.
-func RegisterFunc(name string, check func() error) {
+// registry directly from an arbitrary func(context.Context) error.
+func RegisterFunc(name string, check CheckFunc) {
 	DefaultRegistry.RegisterFunc(name, check)
 }
 
-// RegisterPeriodicFunc allows the convenience of registering a PeriodicChecker
-// from an arbitrary func() error.
-func (registry *Registry) RegisterPeriodicFunc(name string, period time.Duration, check CheckFunc) {
-	registry.Register(name, PeriodicChecker(check, period))
-}
-
-// RegisterPeriodicFunc allows the convenience of registering a PeriodicChecker
-// in the default registry from an arbitrary func() error.
-func RegisterPeriodicFunc(name string, period time.Duration, check CheckFunc) {
-	DefaultRegistry.RegisterPeriodicFunc(name, period, check)
-}
-
-// RegisterPeriodicThresholdFunc allows the convenience of registering a
-// PeriodicChecker from an arbitrary func() error.
-func (registry *Registry) RegisterPeriodicThresholdFunc(name string, period time.Duration, threshold int, check CheckFunc) {
-	registry.Register(name, PeriodicThresholdChecker(check, period, threshold))
-}
-
-// RegisterPeriodicThresholdFunc allows the convenience of registering a
-// PeriodicChecker in the default registry from an arbitrary func() error.
-func RegisterPeriodicThresholdFunc(name string, period time.Duration, threshold int, check CheckFunc) {
-	DefaultRegistry.RegisterPeriodicThresholdFunc(name, period, threshold, check)
-}
-
 // StatusHandler returns a JSON blob with all the currently registered Health Checks
 // and their corresponding status.
 // Returns 503 if any Error status exists, 200 otherwise
 func StatusHandler(w http.ResponseWriter, r *http.Request) {
 	if r.Method == http.MethodGet {
-		checks := CheckStatus()
+		checks := CheckStatus(r.Context())
 		status := http.StatusOK
 
 		// If there is an error, return 503
@@ -263,7 +238,7 @@ func StatusHandler(w http.ResponseWriter, r *http.Request) {
 // disable a web application when the health checks fail.
 func Handler(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		checks := CheckStatus()
+		checks := CheckStatus(r.Context())
 		if len(checks) != 0 {
 			// NOTE(milosgajdos): disable errcheck as the error is
 			// accessible via /debug/health
@@ -282,7 +257,7 @@ func Handler(handler http.Handler) http.Handler {
 func statusResponse(w http.ResponseWriter, r *http.Request, status int, checks map[string]string) {
 	p, err := json.Marshal(checks)
 	if err != nil {
-		dcontext.GetLogger(dcontext.Background()).Errorf("error serializing health status: %v", err)
+		dcontext.GetLogger(r.Context()).Errorf("error serializing health status: %v", err)
 		p, err = json.Marshal(struct {
 			ServerError string `json:"server_error"`
 		}{
@@ -291,7 +266,7 @@ func statusResponse(w http.ResponseWriter, r *http.Request, status int, checks m
 		status = http.StatusInternalServerError
 
 		if err != nil {
-			dcontext.GetLogger(dcontext.Background()).Errorf("error serializing health status failure message: %v", err)
+			dcontext.GetLogger(r.Context()).Errorf("error serializing health status failure message: %v", err)
 			return
 		}
 	}
@@ -300,7 +275,7 @@ func statusResponse(w http.ResponseWriter, r *http.Request, status int, checks m
 	w.Header().Set("Content-Length", fmt.Sprint(len(p)))
 	w.WriteHeader(status)
 	if _, err := w.Write(p); err != nil {
-		dcontext.GetLogger(dcontext.Background()).Errorf("error writing health status response body: %v", err)
+		dcontext.GetLogger(r.Context()).Errorf("error writing health status response body: %v", err)
 	}
 }