120 lines
2.1 KiB
Go
120 lines
2.1 KiB
Go
package probe
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type (
|
|
ReadinessFunc func() Readiness
|
|
LivenessFunc func(context.Context) Liveness
|
|
|
|
ReadinessAggregateFuncs []ReadinessFunc
|
|
)
|
|
|
|
func (fs ReadinessAggregateFuncs) check() (a ReadinessAggregate) {
|
|
a = make(ReadinessAggregate, 0, len(fs))
|
|
for _, f := range fs {
|
|
a = append(a, f())
|
|
}
|
|
|
|
return a
|
|
}
|
|
|
|
type Reporter interface {
|
|
ReportReadiness() Readiness
|
|
ReportLiveness() Liveness
|
|
|
|
RegisterReadiness(ReadinessFunc)
|
|
RegisterLiveness(LivenessFunc)
|
|
}
|
|
|
|
func NewReporter(livenessTimeout time.Duration) *reporter {
|
|
return &reporter{
|
|
livenessTimeout: livenessTimeout,
|
|
}
|
|
}
|
|
|
|
type reporter struct {
|
|
readinessComponents ReadinessAggregateFuncs
|
|
livenessComponents []LivenessFunc
|
|
|
|
livenessTimeout time.Duration
|
|
|
|
mu sync.Mutex
|
|
livemu sync.Mutex
|
|
readmu sync.Mutex
|
|
}
|
|
|
|
func (r *reporter) ReportReadiness() Readiness {
|
|
r.readmu.Lock()
|
|
defer r.readmu.Unlock()
|
|
|
|
return r.readinessComponents.check().Status()
|
|
}
|
|
|
|
func (r *reporter) ReportLiveness() (out Liveness) {
|
|
r.livemu.Lock()
|
|
defer r.livemu.Unlock()
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), r.livenessTimeout)
|
|
defer cancel()
|
|
|
|
for _, f := range r.livenessComponents {
|
|
status := f(ctx)
|
|
if status == LivenessTimeout {
|
|
return status
|
|
}
|
|
}
|
|
|
|
return LivenessOk
|
|
}
|
|
|
|
func (r *reporter) RegisterReadiness(f ReadinessFunc) {
|
|
r.readmu.Lock()
|
|
defer r.readmu.Unlock()
|
|
|
|
r.readinessComponents = append(r.readinessComponents, f)
|
|
}
|
|
|
|
func (r *reporter) RegisterLiveness(f LivenessFunc) {
|
|
r.livemu.Lock()
|
|
defer r.livemu.Unlock()
|
|
|
|
r.livenessComponents = append(r.livenessComponents, f)
|
|
}
|
|
|
|
func SimpleReadinessSwitcher() (f ReadinessFunc, toggle func(newStatus Readiness)) {
|
|
var status atomic.Int32
|
|
|
|
f = func() Readiness {
|
|
return Readiness(status.Load())
|
|
}
|
|
|
|
toggle = func(newStatus Readiness) {
|
|
status.Store(int32(newStatus))
|
|
}
|
|
|
|
return f, toggle
|
|
}
|
|
|
|
func SimpleLivenessSwitcher() (f LivenessFunc, toggle func()) {
|
|
var liveness atomic.Bool
|
|
|
|
f = func(context.Context) Liveness {
|
|
if liveness.Load() {
|
|
return LivenessOk
|
|
}
|
|
|
|
return LivenessUnknown
|
|
}
|
|
|
|
toggle = func() {
|
|
liveness.Store(true)
|
|
}
|
|
|
|
return f, toggle
|
|
}
|