spark-server/internal/bootstrap/health.go

package bootstrap

import (
	"context"
	"sync/atomic"
	"time"

	"spark-server/internal/infrastructure/mysql"
	redisx "spark-server/internal/infrastructure/redis"
	tcp "spark-server/internal/transport/tcp"
)

const (
	depOK      = "ok"
	depSkipped = "skipped"
	depDown    = "down"
	depClosed  = "closed"
)

// ProbeReport is /readyz data. LB should use HTTP 503 when Ready is false.
type ProbeReport struct {
	Status        string `json:"status"`
	Accepting     bool   `json:"accepting"`
	MySQL         string `json:"mysql"`
	Redis         string `json:"redis"`
	ObjectStorage string `json:"object_storage"`
	TaskPool      string `json:"taskpool"`
}

// Probe implements httpx.Health for liveness/readiness.
type Probe struct {
	mysql       *mysql.DB
	redis       *redisx.Client
	objects     pinger
	pool        tcp.TaskPool
	accepting   *atomic.Bool
	requireDeps bool
}

func (p *Probe) Ready(ctx context.Context) (bool, any) {
	if p == nil {
		return true, ProbeReport{Status: depOK, Accepting: true, MySQL: depSkipped, Redis: depSkipped, ObjectStorage: depSkipped, TaskPool: depOK}
	}
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()
	rep := ProbeReport{
		Accepting:     p.accepting != nil && p.accepting.Load(),
		MySQL:         p.checkStore(ctx, p.mysql),
		Redis:         p.checkStore(ctx, p.redis),
		ObjectStorage: p.checkOptional(ctx, p.objects),
		TaskPool:      depOK,
	}
	ready := rep.Accepting
	if p.pool != nil && p.pool.Closed() {
		rep.TaskPool = depClosed
		ready = false
	}
	if rep.MySQL == depDown || rep.Redis == depDown || rep.ObjectStorage == depDown {
		ready = false
	}
	if ready {
		rep.Status = depOK
	} else {
		rep.Status = "not_ready"
	}
	return ready, rep
}

type pinger interface {
	Ready() bool
	Ping(ctx context.Context) error
}

func (p *Probe) checkStore(ctx context.Context, s pinger) string {
	if s == nil || !s.Ready() {
		if p.requireDeps {
			return depDown
		}
		return depSkipped
	}
	if err := s.Ping(ctx); err != nil {
		return depDown
	}
	return depOK
}

func (p *Probe) checkOptional(ctx context.Context, s pinger) string {
	if s == nil || !s.Ready() {
		return depSkipped
	}
	if err := s.Ping(ctx); err != nil {
		return depDown
	}
	return depOK
}