add healthcheck service

internal/probe/reporter.go

@@ -0,0 +1,119 @@
+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
+}