github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/utils/liveness/collector.go (about)

     1  package liveness
     2  
     3  import (
     4  	"net/http"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  const (
    10  	// DefaultTolerance is the default time (in seconds) allowed between heartbeats.
    11  	DefaultTolerance = time.Second * 30
    12  
    13  	// ToleranceHeader is the HTTP header name used to override a collector's configured tolerance.
    14  	ToleranceHeader = "X-Liveness-Tolerance"
    15  )
    16  
    17  // CheckCollector produces multiple checks and returns
    18  // live only if all decendant checks are live.
    19  //
    20  // Each child check may only be used by one goroutine,
    21  // the CheckCollector may be used by multiple routines at once to
    22  // produce checks.
    23  type CheckCollector struct {
    24  	lock             sync.RWMutex
    25  	defaultTolerance time.Duration
    26  	checks           []Check
    27  }
    28  
    29  // NewCheckCollector creates a threadsafe Collector.
    30  func NewCheckCollector(tolerance time.Duration) *CheckCollector {
    31  	if tolerance == 0 {
    32  		tolerance = DefaultTolerance
    33  	}
    34  
    35  	return &CheckCollector{
    36  		defaultTolerance: tolerance,
    37  		checks:           make([]Check, 0, 1),
    38  	}
    39  }
    40  
    41  // NewCheck returns a Check which is safe for use on a single routine.
    42  func (c *CheckCollector) NewCheck() Check {
    43  	c.lock.Lock()
    44  	check := &internalCheck{
    45  		defaultTolerance: c.defaultTolerance,
    46  		lastCheckIn:      time.Now(),
    47  	}
    48  	c.checks = append(c.checks, check)
    49  	c.lock.Unlock()
    50  	return check
    51  }
    52  
    53  // Register adds a check to a collector
    54  func (c *CheckCollector) Register(ck Check) {
    55  	c.lock.Lock()
    56  	c.checks = append(c.checks, ck)
    57  	c.lock.Unlock()
    58  }
    59  
    60  func (c *CheckCollector) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    61  	tolerance := c.defaultTolerance
    62  	var err error
    63  
    64  	if toleranceStr := r.Header.Get(ToleranceHeader); toleranceStr != "" {
    65  		tolerance, err = time.ParseDuration(toleranceStr)
    66  		if err != nil {
    67  			http.Error(w, "Invalid tolerace: "+toleranceStr, http.StatusBadRequest)
    68  			return
    69  		}
    70  	}
    71  
    72  	if !c.IsLive(tolerance) {
    73  		w.WriteHeader(http.StatusServiceUnavailable)
    74  		return
    75  	}
    76  
    77  	w.WriteHeader(http.StatusOK)
    78  }
    79  
    80  // IsLive checks if we are still live against the given the tolerace between hearbeats.
    81  //
    82  // If tolerance is 0, the default tolerance is used.
    83  func (c *CheckCollector) IsLive(tolerance time.Duration) bool {
    84  	if tolerance == 0 {
    85  		tolerance = c.defaultTolerance
    86  	}
    87  
    88  	c.lock.RLock()
    89  	defer c.lock.RUnlock()
    90  
    91  	for i := range c.checks {
    92  		if !c.checks[i].IsLive(tolerance) {
    93  			return false
    94  		}
    95  	}
    96  
    97  	return true
    98  }