easyai-ai-gateway/apps/api/internal/runner/route_prober.go

package runner

import (
	"context"
	"errors"
	"log/slog"
	"math"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type routeProbeState struct {
	next                time.Time
	lastRequestObserved time.Time
}

func (s *Service) StartRouteHealthProber(ctx context.Context) {
	if !s.cfg.RunsRouteProber() {
		return
	}
	go s.runRouteHealthProber(ctx)
}

func (s *Service) runRouteHealthProber(ctx context.Context) {
	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()
	states := make(map[string]routeProbeState)
	probeConcurrency := s.cfg.RouteProbeConcurrency
	if probeConcurrency <= 0 {
		probeConcurrency = 8
	}
	probeSlots := make(chan struct{}, probeConcurrency)
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			targets, err := s.coordinationStore.ListRouteProbeTargets(ctx)
			if err != nil {
				s.logRouteProbeFailure("list_targets", err)
				continue
			}
			requested, requestErr := s.coordinationStore.ListRequestedRouteProbes(ctx)
			if requestErr != nil {
				s.logRouteProbeFailure("list_requests", requestErr)
				continue
			}
			now := time.Now()
			for _, target := range targets {
				profile, profileErr := routeProfileForCandidate(target.Candidate)
				if profileErr != nil {
					continue
				}
				state := states[profile.Key]
				requestedAt := requested[profile.Key]
				requestDue := !requestedAt.IsZero() && requestedAt.After(state.lastRequestObserved)
				if !requestDue && now.Before(state.next) {
					continue
				}
				interval := time.Duration(s.cfg.RouteProbeColdIntervalSeconds) * time.Second
				if target.Hot {
					interval = time.Duration(s.cfg.RouteProbeHotIntervalSeconds) * time.Second
				}
				select {
				case probeSlots <- struct{}{}:
					states[profile.Key] = routeProbeState{next: now.Add(interval), lastRequestObserved: requestedAt}
				case <-ctx.Done():
					return
				default:
					continue
				}
				go func(target store.RouteProbeTarget, profile executionpool.RouteProfile) {
					defer func() { <-probeSlots }()
					s.probeRouteTarget(ctx, target, profile)
				}(target, profile)
			}
		}
	}
}

func (s *Service) probeRouteTarget(ctx context.Context, target store.RouteProbeTarget, profile executionpool.RouteProfile) {
	leaseToken, locked, err := s.coordinationStore.TryClaimRouteProbe(ctx, s.cfg.ExecutionPoolID, profile.Key, time.Duration(s.cfg.RouteProbeTimeoutMS)*time.Millisecond+5*time.Second)
	if err != nil || !locked {
		return
	}
	defer func() {
		releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := s.coordinationStore.ReleaseRouteProbe(releaseCtx, s.cfg.ExecutionPoolID, profile.Key, leaseToken); err != nil {
			s.logRouteProbeFailure("release_lease", err)
		}
	}()
	client, err := s.httpClientForCandidate(target.Candidate, false, "")
	if err != nil {
		s.logRouteProbeFailure("prepare_client", err)
		return
	}
	probe := executionpool.HTTPProbe{Client: client}
	result := probe.Probe(ctx, executionpool.ProbeTarget{
		RouteProfile: profile,
		URL:          routeEndpointForCandidate(target.Candidate),
		Timeout:      time.Duration(s.cfg.RouteProbeTimeoutMS) * time.Millisecond,
	})
	previous := executionpool.RouteHealth{}
	if history, historyErr := s.coordinationStore.ListRouteHealth(ctx, profile.Key, time.Now()); historyErr == nil {
		for _, item := range history {
			if item.PoolID == s.cfg.ExecutionPoolID {
				previous = item
				break
			}
		}
	}
	health := nextRouteHealth(previous, s.cfg.ExecutionPoolID, profile.Key, result)
	s.observeRouteProbe(string(health.State))
	if err := s.coordinationStore.RecordRouteHealth(context.WithoutCancel(ctx), health); err != nil {
		s.logRouteProbeFailure("record", err)
	}
}

func (s *Service) observeRouteProbe(state string) {
	observer, ok := s.billingMetrics.(interface{ ObserveRouteProbe(string) })
	if ok {
		observer.ObserveRouteProbe(state)
	}
}

func nextRouteHealth(previous executionpool.RouteHealth, poolID, profileKey string, result executionpool.ProbeResult) executionpool.RouteHealth {
	health := previous
	health.PoolID = poolID
	health.RouteProfileKey = profileKey
	health.SampleCount++
	health.SampledAt = result.SampledAt
	health.ExpiresAt = result.SampledAt.Add(45 * time.Second)
	observation := 0.0
	if result.Reachable {
		observation = 1
		health.ConsecutiveSuccesses++
		health.ConsecutiveFailures = 0
		if health.ConsecutiveSuccesses >= 2 || previous.State == executionpool.RouteHealthy {
			health.State = executionpool.RouteHealthy
		} else {
			health.State = executionpool.RouteDegraded
		}
	} else {
		health.ConsecutiveFailures++
		health.ConsecutiveSuccesses = 0
		switch {
		case health.ConsecutiveFailures >= 3:
			health.State = executionpool.RouteUnreachable
		case previous.State == executionpool.RouteHealthy || previous.State == executionpool.RouteDegraded:
			health.State = executionpool.RouteDegraded
		default:
			health.State = executionpool.RouteUnknown
		}
	}
	if previous.SampleCount == 0 {
		health.SuccessRate = observation
	} else {
		health.SuccessRate = 0.8*previous.SuccessRate + 0.2*observation
	}
	connectTLS := result.TCPDuration + result.TLSDuration
	if connectTLS > 0 {
		health.ConnectTLSP95 = ewmaDuration(previous.ConnectTLSP95, connectTLS)
	}
	if result.FirstByte > 0 {
		if previous.FirstByteP95 > 0 {
			delta := math.Abs(float64(result.FirstByte - previous.FirstByteP95))
			health.JitterP95 = ewmaDuration(previous.JitterP95, time.Duration(delta))
		}
		health.FirstByteP95 = ewmaDuration(previous.FirstByteP95, result.FirstByte)
	}
	return health
}

func ewmaDuration(previous, current time.Duration) time.Duration {
	if previous <= 0 {
		return current
	}
	return time.Duration(0.8*float64(previous) + 0.2*float64(current))
}

func (s *Service) logRouteProbeFailure(stage string, err error) {
	if err == nil || errors.Is(err, context.Canceled) || s.logger == nil {
		return
	}
	s.logger.Log(context.Background(), slog.LevelWarn, "route health probe failed", "stage", stage, "error_category", "route_probe_failed")
}