90CNServer/server/internal/platform/health/monitor.go

package health

import (
	"context"
	"strings"
	"sync"
	"time"

	healthpb "google.golang.org/grpc/health/grpc_health_v1"
)

const (
	defMonitorEvery   = 5 * time.Second
	defMonitorTimeout = 2 * time.Second
)

type ProbeTarget struct {
	Name    string
	Service string
	Client  healthpb.HealthClient
}

type ProbeResult struct {
	Target    string    `json:"target"`
	Service   string    `json:"service,omitempty"`
	State     State     `json:"state"`
	Error     string    `json:"error,omitempty"`
	CheckedAt time.Time `json:"checked_at"`
}

type MonitorOptions struct {
	Interval time.Duration
	Timeout  time.Duration
	Targets  []ProbeTarget
	Notify   func(ProbeResult)
	Now      func() time.Time
}

type Monitor struct {
	interval time.Duration
	timeout  time.Duration
	targets  []ProbeTarget
	notify   func(ProbeResult)
	now      func() time.Time

	mu   sync.Mutex
	last map[string]State
}

func NewMonitor(options MonitorOptions) *Monitor {
	interval := options.Interval
	if interval <= 0 {
		interval = defMonitorEvery
	}
	timeout := options.Timeout
	if timeout <= 0 {
		timeout = defMonitorTimeout
	}
	now := options.Now
	if now == nil {
		now = func() time.Time { return time.Now().UTC() }
	}
	targets := make([]ProbeTarget, 0, len(options.Targets))
	for _, target := range options.Targets {
		target.Name = strings.TrimSpace(target.Name)
		target.Service = strings.TrimSpace(target.Service)
		if target.Name == "" {
			target.Name = target.Service
		}
		if target.Name == "" {
			continue
		}
		targets = append(targets, target)
	}
	return &Monitor{
		interval: interval,
		timeout:  timeout,
		targets:  targets,
		notify:   options.Notify,
		now:      now,
		last:     make(map[string]State),
	}
}

func (m *Monitor) CheckOnce(ctx context.Context) []ProbeResult {
	if m == nil {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	results := make([]ProbeResult, 0, len(m.targets))
	for _, target := range m.targets {
		result := m.probe(ctx, target)
		results = append(results, result)
		m.notifyIfChanged(result)
	}
	return results
}

func (m *Monitor) Run(ctx context.Context) error {
	if m == nil {
		return nil
	}
	if ctx == nil {
		ctx = context.Background()
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	m.CheckOnce(ctx)
	ticker := time.NewTicker(m.interval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
			m.CheckOnce(ctx)
		}
	}
}

func (m *Monitor) probe(ctx context.Context, target ProbeTarget) ProbeResult {
	probeCtx := ctx
	cancel := func() {}
	if m.timeout > 0 {
		probeCtx, cancel = context.WithTimeout(ctx, m.timeout)
	}
	defer cancel()
	state, err := ProbeGRPC(probeCtx, target.Client, target.Service)
	result := ProbeResult{
		Target:    target.Name,
		Service:   target.Service,
		State:     state,
		CheckedAt: m.now().UTC(),
	}
	if err != nil {
		result.Error = err.Error()
	}
	return result
}

func (m *Monitor) notifyIfChanged(result ProbeResult) {
	if m.notify == nil {
		return
	}
	key := result.Target + "\x00" + result.Service
	m.mu.Lock()
	last, ok := m.last[key]
	if ok && last == result.State {
		m.mu.Unlock()
		return
	}
	m.last[key] = result.State
	m.mu.Unlock()
	m.notify(result)
}