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

package runner

import (
	"context"
	"crypto/tls"
	"net/http"
	"net/http/httptrace"
	"sort"
	"sync"
	"time"

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

type passiveRouteObserver struct {
	mu              sync.Mutex
	samples         int
	successes       int
	connectTLS      []time.Duration
	uploadBytesRate []float64
}

type passiveRouteRoundTripper struct {
	base     http.RoundTripper
	observer *passiveRouteObserver
}

func (observer *passiveRouteObserver) wrap(client *http.Client) *http.Client {
	if client == nil {
		client = http.DefaultClient
	}
	cloned := *client
	base := cloned.Transport
	if base == nil {
		base = http.DefaultTransport
	}
	cloned.Transport = passiveRouteRoundTripper{base: base, observer: observer}
	return &cloned
}

func (transport passiveRouteRoundTripper) RoundTrip(request *http.Request) (*http.Response, error) {
	startedAt := time.Now()
	var connectStartedAt time.Time
	var connectDuration time.Duration
	var tlsStartedAt time.Time
	var tlsDuration time.Duration
	var wroteRequestAt time.Time
	trace := &httptrace.ClientTrace{
		ConnectStart: func(_, _ string) { connectStartedAt = time.Now() },
		ConnectDone: func(_, _ string, _ error) {
			if !connectStartedAt.IsZero() {
				connectDuration = time.Since(connectStartedAt)
			}
		},
		TLSHandshakeStart: func() { tlsStartedAt = time.Now() },
		TLSHandshakeDone: func(tls.ConnectionState, error) {
			if !tlsStartedAt.IsZero() {
				tlsDuration = time.Since(tlsStartedAt)
			}
		},
		WroteRequest: func(httptrace.WroteRequestInfo) { wroteRequestAt = time.Now() },
	}
	traced := request.Clone(httptrace.WithClientTrace(request.Context(), trace))
	response, err := transport.base.RoundTrip(traced)
	transport.observer.record(
		err == nil,
		connectDuration+tlsDuration,
		request.ContentLength,
		startedAt,
		wroteRequestAt,
	)
	return response, err
}

func (observer *passiveRouteObserver) record(success bool, connectTLS time.Duration, contentLength int64, startedAt, wroteRequestAt time.Time) {
	observer.mu.Lock()
	defer observer.mu.Unlock()
	observer.samples++
	if success {
		observer.successes++
	}
	if connectTLS > 0 {
		observer.connectTLS = append(observer.connectTLS, connectTLS)
	}
	if contentLength > 0 && !wroteRequestAt.IsZero() && wroteRequestAt.After(startedAt) {
		observer.uploadBytesRate = append(observer.uploadBytesRate, float64(contentLength)/wroteRequestAt.Sub(startedAt).Seconds())
	}
}

func (observer *passiveRouteObserver) snapshot(poolID, routeProfileKey string) executionpool.RouteObservation {
	observer.mu.Lock()
	defer observer.mu.Unlock()
	return executionpool.RouteObservation{
		PoolID: poolID, RouteProfileKey: routeProfileKey,
		SampleCount: observer.samples, SuccessCount: observer.successes,
		ConnectTLSP95:        durationP95(observer.connectTLS),
		UploadBytesPerSecond: floatP95(observer.uploadBytesRate),
	}
}

func durationP95(values []time.Duration) time.Duration {
	if len(values) == 0 {
		return 0
	}
	copyOfValues := append([]time.Duration(nil), values...)
	sort.Slice(copyOfValues, func(i, j int) bool { return copyOfValues[i] < copyOfValues[j] })
	return copyOfValues[(len(copyOfValues)*95+99)/100-1]
}

func floatP95(values []float64) float64 {
	if len(values) == 0 {
		return 0
	}
	copyOfValues := append([]float64(nil), values...)
	sort.Float64s(copyOfValues)
	return copyOfValues[(len(copyOfValues)*95+99)/100-1]
}

func (s *Service) recordPassiveRouteObservation(task store.GatewayTask, observer *passiveRouteObserver) {
	if observer == nil || !s.routingEnabled() || task.RouteProfileKey == "" {
		return
	}
	poolID := task.AssignedPoolID
	if poolID == "" && s.cfg.RunsAsyncExecutionWorker() {
		poolID = s.cfg.ExecutionPoolID
	}
	observation := observer.snapshot(poolID, task.RouteProfileKey)
	if observation.PoolID == "" || observation.SampleCount == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
	defer cancel()
	if err := s.coordinationStore.RecordRouteObservation(ctx, observation); err != nil && s.logger != nil {
		s.logger.Warn("record passive route observation failed", "error_category", "route_observation_failed")
	}
}