Plumb contexts into health checks (#4141)

This commit is contained in:
Milos Gajdos 2023-11-27 20:54:02 +00:00 committed by GitHub
commit d9abc517e8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 179 additions and 120 deletions

View file

@ -1,6 +1,7 @@
package api package api
import ( import (
"context"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"testing" "testing"
@ -59,7 +60,7 @@ func TestPOSTDownHandlerChangeStatus(t *testing.T) {
t.Errorf("Did not get a 200.") 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.") 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.") 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.") t.Errorf("UpHandler didn't remove the error check.")
} }
} }

View file

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

View file

@ -1,25 +1,26 @@
package checks package checks
import ( import (
"context"
"testing" "testing"
) )
func TestFileChecker(t *testing.T) { 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") 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) t.Errorf("NoSuchFileFromMoon was expected as not exists, error:%v", err)
} }
} }
func TestHTTPChecker(t *testing.T) { 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") 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) t.Errorf("Google at Portugal was expected as exists, error:%v", err)
} }
} }

View file

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

View file

@ -1,11 +1,13 @@
package health package health
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"testing" "testing"
"time"
) )
// TestReturns200IfThereAreNoChecks ensures that the result code of the health // TestReturns200IfThereAreNoChecks ensures that the result code of the health
@ -36,7 +38,7 @@ func TestReturns503IfThereAreErrorChecks(t *testing.T) {
} }
// Create a manual error // Create a manual error
Register("some_check", CheckFunc(func() error { Register("some_check", CheckFunc(func(context.Context) error {
return errors.New("This Check did not succeed") return errors.New("This Check did not succeed")
})) }))
@ -105,3 +107,87 @@ func TestHealthHandler(t *testing.T) {
updater.Update(nil) updater.Update(nil)
checkUp(t, "when server is back up") // now we should be back up. checkUp(t, "when server is back up") // now we should be back up.
} }
func TestThresholdStatusUpdater(t *testing.T) {
u := NewThresholdStatusUpdater(3)
assertCheckOK := func() {
t.Helper()
if err := u.Check(context.Background()); err != nil {
t.Errorf("u.Check() = %v; want nil", err)
}
}
assertCheckErr := func(expected string) {
t.Helper()
if err := u.Check(context.Background()); err == nil || err.Error() != expected {
t.Errorf("u.Check() = %v; want %v", err, expected)
}
}
// Updater should report healthy until the threshold is reached.
for i := 1; i <= 3; i++ {
assertCheckOK()
u.Update(fmt.Errorf("fake error %d", i))
}
assertCheckErr("fake error 3")
// The threshold should reset after one successful update.
u.Update(nil)
assertCheckOK()
u.Update(errors.New("first errored update after reset"))
assertCheckOK()
u.Update(nil)
// pollingTerminatedErr should bypass the threshold.
pte := pollingTerminatedErr{Err: errors.New("womp womp")}
u.Update(pte)
assertCheckErr(pte.Error())
}
func TestPoll(t *testing.T) {
type ContextKey struct{}
for _, threshold := range []int{0, 10} {
t.Run(fmt.Sprintf("threshold=%d", threshold), func(t *testing.T) {
ctx, cancel := context.WithCancel(context.WithValue(context.Background(), ContextKey{}, t.Name()))
defer cancel()
checkerCalled := make(chan struct{})
checker := CheckFunc(func(ctx context.Context) error {
if v, ok := ctx.Value(ContextKey{}).(string); !ok || v != t.Name() {
t.Errorf("unexpected context passed into checker: got context with value %q, want %q", v, t.Name())
}
select {
case <-checkerCalled:
default:
close(checkerCalled)
}
return nil
})
updater := NewThresholdStatusUpdater(threshold)
pollReturned := make(chan struct{})
go func() {
Poll(ctx, updater, checker, 1*time.Millisecond)
close(pollReturned)
}()
select {
case <-checkerCalled:
case <-time.After(1 * time.Second):
t.Error("checker has not been polled")
}
cancel()
select {
case <-pollReturned:
case <-time.After(1 * time.Second):
t.Error("poll has not returned after context was canceled")
}
if err := updater.Check(context.Background()); !errors.Is(err, context.Canceled) {
t.Errorf("updater.Check() = %v; want %v", err, context.Canceled)
}
})
}
}

View file

@ -348,22 +348,20 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
interval = defaultCheckInterval interval = defaultCheckInterval
} }
storageDriverCheck := func() error { storageDriverCheck := health.CheckFunc(func(ctx context.Context) error {
_, err := app.driver.Stat(app, "/") // "/" should always exist _, err := app.driver.Stat(ctx, "/") // "/" should always exist
if _, ok := err.(storagedriver.PathNotFoundError); ok { if _, ok := err.(storagedriver.PathNotFoundError); ok {
err = nil // pass this through, backend is responding, but this path doesn't exist. err = nil // pass this through, backend is responding, but this path doesn't exist.
} }
if err != nil { if err != nil {
dcontext.GetLogger(app).Errorf("storage driver health check: %v", err) dcontext.GetLogger(ctx).Errorf("storage driver health check: %v", err)
} }
return err return err
} })
if app.Config.Health.StorageDriver.Threshold != 0 { updater := health.NewThresholdStatusUpdater(app.Config.Health.StorageDriver.Threshold)
healthRegistry.RegisterPeriodicThresholdFunc("storagedriver_"+app.Config.Storage.Type(), interval, app.Config.Health.StorageDriver.Threshold, storageDriverCheck) healthRegistry.Register("storagedriver_"+app.Config.Storage.Type(), updater)
} else { go health.Poll(app, updater, storageDriverCheck, interval)
healthRegistry.RegisterPeriodicFunc("storagedriver_"+app.Config.Storage.Type(), interval, storageDriverCheck)
}
} }
for _, fileChecker := range app.Config.Health.FileCheckers { for _, fileChecker := range app.Config.Health.FileCheckers {
@ -372,7 +370,9 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
interval = defaultCheckInterval interval = defaultCheckInterval
} }
dcontext.GetLogger(app).Infof("configuring file health check path=%s, interval=%d", fileChecker.File, interval/time.Second) dcontext.GetLogger(app).Infof("configuring file health check path=%s, interval=%d", fileChecker.File, interval/time.Second)
healthRegistry.Register(fileChecker.File, health.PeriodicChecker(checks.FileChecker(fileChecker.File), interval)) u := health.NewStatusUpdater()
healthRegistry.Register(fileChecker.File, u)
go health.Poll(app, u, checks.FileChecker(fileChecker.File), interval)
} }
for _, httpChecker := range app.Config.Health.HTTPCheckers { for _, httpChecker := range app.Config.Health.HTTPCheckers {
@ -388,13 +388,10 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
checker := checks.HTTPChecker(httpChecker.URI, statusCode, httpChecker.Timeout, httpChecker.Headers) checker := checks.HTTPChecker(httpChecker.URI, statusCode, httpChecker.Timeout, httpChecker.Headers)
if httpChecker.Threshold != 0 {
dcontext.GetLogger(app).Infof("configuring HTTP health check uri=%s, interval=%d, threshold=%d", httpChecker.URI, interval/time.Second, httpChecker.Threshold) dcontext.GetLogger(app).Infof("configuring HTTP health check uri=%s, interval=%d, threshold=%d", httpChecker.URI, interval/time.Second, httpChecker.Threshold)
healthRegistry.Register(httpChecker.URI, health.PeriodicThresholdChecker(checker, interval, httpChecker.Threshold)) updater := health.NewThresholdStatusUpdater(httpChecker.Threshold)
} else { healthRegistry.Register(httpChecker.URI, updater)
dcontext.GetLogger(app).Infof("configuring HTTP health check uri=%s, interval=%d", httpChecker.URI, interval/time.Second) go health.Poll(app, updater, checker, interval)
healthRegistry.Register(httpChecker.URI, health.PeriodicChecker(checker, interval))
}
} }
for _, tcpChecker := range app.Config.Health.TCPCheckers { for _, tcpChecker := range app.Config.Health.TCPCheckers {
@ -405,13 +402,10 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
checker := checks.TCPChecker(tcpChecker.Addr, tcpChecker.Timeout) checker := checks.TCPChecker(tcpChecker.Addr, tcpChecker.Timeout)
if tcpChecker.Threshold != 0 {
dcontext.GetLogger(app).Infof("configuring TCP health check addr=%s, interval=%d, threshold=%d", tcpChecker.Addr, interval/time.Second, tcpChecker.Threshold) dcontext.GetLogger(app).Infof("configuring TCP health check addr=%s, interval=%d, threshold=%d", tcpChecker.Addr, interval/time.Second, tcpChecker.Threshold)
healthRegistry.Register(tcpChecker.Addr, health.PeriodicThresholdChecker(checker, interval, tcpChecker.Threshold)) updater := health.NewThresholdStatusUpdater(tcpChecker.Threshold)
} else { healthRegistry.Register(tcpChecker.Addr, updater)
dcontext.GetLogger(app).Infof("configuring TCP health check addr=%s, interval=%d", tcpChecker.Addr, interval/time.Second) go health.Poll(app, updater, checker, interval)
healthRegistry.Register(tcpChecker.Addr, health.PeriodicChecker(checker, interval))
}
} }
} }

View file

@ -5,6 +5,7 @@ import (
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"os" "os"
"strings"
"testing" "testing"
"time" "time"
@ -48,7 +49,7 @@ func TestFileHealthCheck(t *testing.T) {
// Wait for health check to happen // Wait for health check to happen
<-time.After(2 * interval) <-time.After(2 * interval)
status := healthRegistry.CheckStatus() status := healthRegistry.CheckStatus(ctx)
if len(status) != 1 { if len(status) != 1 {
t.Fatal("expected 1 item in health check results") t.Fatal("expected 1 item in health check results")
} }
@ -59,7 +60,7 @@ func TestFileHealthCheck(t *testing.T) {
os.Remove(tmpfile.Name()) os.Remove(tmpfile.Name())
<-time.After(2 * interval) <-time.After(2 * interval)
if len(healthRegistry.CheckStatus()) != 0 { if len(healthRegistry.CheckStatus(ctx)) != 0 {
t.Fatal("expected 0 items in health check results") t.Fatal("expected 0 items in health check results")
} }
} }
@ -112,7 +113,7 @@ func TestTCPHealthCheck(t *testing.T) {
// Wait for health check to happen // Wait for health check to happen
<-time.After(2 * interval) <-time.After(2 * interval)
if len(healthRegistry.CheckStatus()) != 0 { if len(healthRegistry.CheckStatus(ctx)) != 0 {
t.Fatal("expected 0 items in health check results") t.Fatal("expected 0 items in health check results")
} }
@ -120,11 +121,11 @@ func TestTCPHealthCheck(t *testing.T) {
<-time.After(2 * interval) <-time.After(2 * interval)
// Health check should now fail // Health check should now fail
status := healthRegistry.CheckStatus() status := healthRegistry.CheckStatus(ctx)
if len(status) != 1 { if len(status) != 1 {
t.Fatal("expected 1 item in health check results") t.Fatal("expected 1 item in health check results")
} }
if status[addrStr] != "connection to "+addrStr+" failed" { if !strings.Contains(status[addrStr], "connection failed") {
t.Fatal(`did not get "connection failed" result for health check`) t.Fatal(`did not get "connection failed" result for health check`)
} }
} }
@ -174,7 +175,7 @@ func TestHTTPHealthCheck(t *testing.T) {
for i := 0; ; i++ { for i := 0; ; i++ {
<-time.After(interval) <-time.After(interval)
status := healthRegistry.CheckStatus() status := healthRegistry.CheckStatus(ctx)
if i < threshold-1 { if i < threshold-1 {
// definitely shouldn't have hit the threshold yet // definitely shouldn't have hit the threshold yet
@ -191,7 +192,7 @@ func TestHTTPHealthCheck(t *testing.T) {
if len(status) != 1 { if len(status) != 1 {
t.Fatal("expected 1 item in health check results") t.Fatal("expected 1 item in health check results")
} }
if status[checkedServer.URL] != "downstream service returned unexpected status: 500" { if !strings.Contains(status[checkedServer.URL], "downstream service returned unexpected status: 500") {
t.Fatal("did not get expected result for health check") t.Fatal("did not get expected result for health check")
} }
@ -203,7 +204,7 @@ func TestHTTPHealthCheck(t *testing.T) {
<-time.After(2 * interval) <-time.After(2 * interval)
if len(healthRegistry.CheckStatus()) != 0 { if len(healthRegistry.CheckStatus(ctx)) != 0 {
t.Fatal("expected 0 items in health check results") t.Fatal("expected 0 items in health check results")
} }
} }