health.go

Documentation: github.com/docker/distribution/health

     1  package health
     2  
     3  import (
     4  	"encoding/json"
     5  	"fmt"
     6  	"net/http"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/docker/distribution/context"
    11  	"github.com/docker/distribution/registry/api/errcode"
    12  )
    13  
    14  // A Registry is a collection of checks. Most applications will use the global
    15  // registry defined in DefaultRegistry. However, unit tests may need to create
    16  // separate registries to isolate themselves from other tests.
    17  type Registry struct {
    18  	mu               sync.RWMutex
    19  	registeredChecks map[string]Checker
    20  }
    21  
    22  // NewRegistry creates a new registry. This isn't necessary for normal use of
    23  // the package, but may be useful for unit tests so individual tests have their
    24  // own set of checks.
    25  func NewRegistry() *Registry {
    26  	return &Registry{
    27  		registeredChecks: make(map[string]Checker),
    28  	}
    29  }
    30  
    31  // DefaultRegistry is the default registry where checks are registered. It is
    32  // the registry used by the HTTP handler.
    33  var DefaultRegistry *Registry
    34  
    35  // Checker is the interface for a Health Checker
    36  type Checker interface {
    37  	// Check returns nil if the service is okay.
    38  	Check() error
    39  }
    40  
    41  // CheckFunc is a convenience type to create functions that implement
    42  // the Checker interface
    43  type CheckFunc func() error
    44  
    45  // Check Implements the Checker interface to allow for any func() error method
    46  // to be passed as a Checker
    47  func (cf CheckFunc) Check() error {
    48  	return cf()
    49  }
    50  
    51  // Updater implements a health check that is explicitly set.
    52  type Updater interface {
    53  	Checker
    54  
    55  	// Update updates the current status of the health check.
    56  	Update(status error)
    57  }
    58  
    59  // updater implements Checker and Updater, providing an asynchronous Update
    60  // method.
    61  // This allows us to have a Checker that returns the Check() call immediately
    62  // not blocking on a potentially expensive check.
    63  type updater struct {
    64  	mu     sync.Mutex
    65  	status error
    66  }
    67  
    68  // Check implements the Checker interface
    69  func (u *updater) Check() error {
    70  	u.mu.Lock()
    71  	defer u.mu.Unlock()
    72  
    73  	return u.status
    74  }
    75  
    76  // Update implements the Updater interface, allowing asynchronous access to
    77  // the status of a Checker.
    78  func (u *updater) Update(status error) {
    79  	u.mu.Lock()
    80  	defer u.mu.Unlock()
    81  
    82  	u.status = status
    83  }
    84  
    85  // NewStatusUpdater returns a new updater
    86  func NewStatusUpdater() Updater {
    87  	return &updater{}
    88  }
    89  
    90  // thresholdUpdater implements Checker and Updater, providing an asynchronous Update
    91  // method.
    92  // This allows us to have a Checker that returns the Check() call immediately
    93  // not blocking on a potentially expensive check.
    94  type thresholdUpdater struct {
    95  	mu        sync.Mutex
    96  	status    error
    97  	threshold int
    98  	count     int
    99  }
   100  
   101  // Check implements the Checker interface
   102  func (tu *thresholdUpdater) Check() error {
   103  	tu.mu.Lock()
   104  	defer tu.mu.Unlock()
   105  
   106  	if tu.count >= tu.threshold {
   107  		return tu.status
   108  	}
   109  
   110  	return nil
   111  }
   112  
   113  // thresholdUpdater implements the Updater interface, allowing asynchronous
   114  // access to the status of a Checker.
   115  func (tu *thresholdUpdater) Update(status error) {
   116  	tu.mu.Lock()
   117  	defer tu.mu.Unlock()
   118  
   119  	if status == nil {
   120  		tu.count = 0
   121  	} else if tu.count < tu.threshold {
   122  		tu.count++
   123  	}
   124  
   125  	tu.status = status
   126  }
   127  
   128  // NewThresholdStatusUpdater returns a new thresholdUpdater
   129  func NewThresholdStatusUpdater(t int) Updater {
   130  	return &thresholdUpdater{threshold: t}
   131  }
   132  
   133  // PeriodicChecker wraps an updater to provide a periodic checker
   134  func PeriodicChecker(check Checker, period time.Duration) Checker {
   135  	u := NewStatusUpdater()
   136  	go func() {
   137  		t := time.NewTicker(period)
   138  		for {
   139  			<-t.C
   140  			u.Update(check.Check())
   141  		}
   142  	}()
   143  
   144  	return u
   145  }
   146  
   147  // PeriodicThresholdChecker wraps an updater to provide a periodic checker that
   148  // uses a threshold before it changes status
   149  func PeriodicThresholdChecker(check Checker, period time.Duration, threshold int) Checker {
   150  	tu := NewThresholdStatusUpdater(threshold)
   151  	go func() {
   152  		t := time.NewTicker(period)
   153  		for {
   154  			<-t.C
   155  			tu.Update(check.Check())
   156  		}
   157  	}()
   158  
   159  	return tu
   160  }
   161  
   162  // CheckStatus returns a map with all the current health check errors
   163  func (registry *Registry) CheckStatus() map[string]string { // TODO(stevvooe) this needs a proper type
   164  	registry.mu.RLock()
   165  	defer registry.mu.RUnlock()
   166  	statusKeys := make(map[string]string)
   167  	for k, v := range registry.registeredChecks {
   168  		err := v.Check()
   169  		if err != nil {
   170  			statusKeys[k] = err.Error()
   171  		}
   172  	}
   173  
   174  	return statusKeys
   175  }
   176  
   177  // CheckStatus returns a map with all the current health check errors from the
   178  // default registry.
   179  func CheckStatus() map[string]string {
   180  	return DefaultRegistry.CheckStatus()
   181  }
   182  
   183  // Register associates the checker with the provided name.
   184  func (registry *Registry) Register(name string, check Checker) {
   185  	if registry == nil {
   186  		registry = DefaultRegistry
   187  	}
   188  	registry.mu.Lock()
   189  	defer registry.mu.Unlock()
   190  	_, ok := registry.registeredChecks[name]
   191  	if ok {
   192  		panic("Check already exists: " + name)
   193  	}
   194  	registry.registeredChecks[name] = check
   195  }
   196  
   197  // Register associates the checker with the provided name in the default
   198  // registry.
   199  func Register(name string, check Checker) {
   200  	DefaultRegistry.Register(name, check)
   201  }
   202  
   203  // RegisterFunc allows the convenience of registering a checker directly from
   204  // an arbitrary func() error.
   205  func (registry *Registry) RegisterFunc(name string, check func() error) {
   206  	registry.Register(name, CheckFunc(check))
   207  }
   208  
   209  // RegisterFunc allows the convenience of registering a checker in the default
   210  // registry directly from an arbitrary func() error.
   211  func RegisterFunc(name string, check func() error) {
   212  	DefaultRegistry.RegisterFunc(name, check)
   213  }
   214  
   215  // RegisterPeriodicFunc allows the convenience of registering a PeriodicChecker
   216  // from an arbitrary func() error.
   217  func (registry *Registry) RegisterPeriodicFunc(name string, period time.Duration, check CheckFunc) {
   218  	registry.Register(name, PeriodicChecker(check, period))
   219  }
   220  
   221  // RegisterPeriodicFunc allows the convenience of registering a PeriodicChecker
   222  // in the default registry from an arbitrary func() error.
   223  func RegisterPeriodicFunc(name string, period time.Duration, check CheckFunc) {
   224  	DefaultRegistry.RegisterPeriodicFunc(name, period, check)
   225  }
   226  
   227  // RegisterPeriodicThresholdFunc allows the convenience of registering a
   228  // PeriodicChecker from an arbitrary func() error.
   229  func (registry *Registry) RegisterPeriodicThresholdFunc(name string, period time.Duration, threshold int, check CheckFunc) {
   230  	registry.Register(name, PeriodicThresholdChecker(check, period, threshold))
   231  }
   232  
   233  // RegisterPeriodicThresholdFunc allows the convenience of registering a
   234  // PeriodicChecker in the default registry from an arbitrary func() error.
   235  func RegisterPeriodicThresholdFunc(name string, period time.Duration, threshold int, check CheckFunc) {
   236  	DefaultRegistry.RegisterPeriodicThresholdFunc(name, period, threshold, check)
   237  }
   238  
   239  // StatusHandler returns a JSON blob with all the currently registered Health Checks
   240  // and their corresponding status.
   241  // Returns 503 if any Error status exists, 200 otherwise
   242  func StatusHandler(w http.ResponseWriter, r *http.Request) {
   243  	if r.Method == "GET" {
   244  		checks := CheckStatus()
   245  		status := http.StatusOK
   246  
   247  		// If there is an error, return 503
   248  		if len(checks) != 0 {
   249  			status = http.StatusServiceUnavailable
   250  		}
   251  
   252  		statusResponse(w, r, status, checks)
   253  	} else {
   254  		http.NotFound(w, r)
   255  	}
   256  }
   257  
   258  // Handler returns a handler that will return 503 response code if the health
   259  // checks have failed. If everything is okay with the health checks, the
   260  // handler will pass through to the provided handler. Use this handler to
   261  // disable a web application when the health checks fail.
   262  func Handler(handler http.Handler) http.Handler {
   263  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   264  		checks := CheckStatus()
   265  		if len(checks) != 0 {
   266  			errcode.ServeJSON(w, errcode.ErrorCodeUnavailable.
   267  				WithDetail("health check failed: please see /debug/health"))
   268  			return
   269  		}
   270  
   271  		handler.ServeHTTP(w, r) // pass through
   272  	})
   273  }
   274  
   275  // statusResponse completes the request with a response describing the health
   276  // of the service.
   277  func statusResponse(w http.ResponseWriter, r *http.Request, status int, checks map[string]string) {
   278  	p, err := json.Marshal(checks)
   279  	if err != nil {
   280  		context.GetLogger(context.Background()).Errorf("error serializing health status: %v", err)
   281  		p, err = json.Marshal(struct {
   282  			ServerError string `json:"server_error"`
   283  		}{
   284  			ServerError: "Could not parse error message",
   285  		})
   286  		status = http.StatusInternalServerError
   287  
   288  		if err != nil {
   289  			context.GetLogger(context.Background()).Errorf("error serializing health status failure message: %v", err)
   290  			return
   291  		}
   292  	}
   293  
   294  	w.Header().Set("Content-Type", "application/json; charset=utf-8")
   295  	w.Header().Set("Content-Length", fmt.Sprint(len(p)))
   296  	w.WriteHeader(status)
   297  	if _, err := w.Write(p); err != nil {
   298  		context.GetLogger(context.Background()).Errorf("error writing health status response body: %v", err)
   299  	}
   300  }
   301  
   302  // Registers global /debug/health api endpoint, creates default registry
   303  func init() {
   304  	DefaultRegistry = NewRegistry()
   305  	http.HandleFunc("/debug/health", StatusHandler)
   306  }
   307
View as plain text