mirror of
https://github.com/lukaszraczylo/traefik-llmgateway.git
synced 2026-10-12 04:54:45 +00:00
feat(metrics): instrument upstream latency
Add time-to-first-byte and total-duration histograms for every upstream call, observation only. No throttling, priorities, or control logic reads this data. The operator wants strain-based throttling later and needs real traffic data first. Extend watchdogBody (timeout.go) rather than add a new interception layer: it already wraps every upstream response body for provider adapters, native passthrough, and the MCP/A2A target proxy, and already observes the first successful read. armLatency arms a watchdogBody after construction, so newWatchdogBody keeps its existing signature and behavior; its own pinned tests still pass unchanged. A timestamp plus a few atomics are the only added cost; Read and Close each pay one nil check when latency collection is disabled. Expose llmgateway_upstream_ttfb_seconds and llmgateway_upstream_duration_seconds on the existing Prometheus endpoint (metrics.go). Both are in-process, per-replica histograms, never store-backed, so sum() is always the correct cross-replica aggregation. Buckets run 0.1s to 600s, chosen to resolve fast TTFB and short completions finely, cover the bulk of real chat completions, and bracket the five-minute default request timeout. Labels are provider and stream always, model only when Config.Metrics.ModelLabel is set, matching the existing per-model opt-in convention. Add a compact per-provider, per-stream average summary to GET /admin/api/overview, sourced from the same in-process store, without a full histogram or a per-model breakdown. Every write path is gated on metricsEnabled(cfg), so a deployment with metrics disabled collects nothing and the histogram families never render. No behavior change to request handling: same responses, status codes, headers, and accounting. Fix a Yaegi-only bug this feature's own context-threading change exposed: two distinct empty struct types used as context.WithValue keys compare equal to each other under the interpreter's interface equality, though never under compiled Go. Threading a second recorder (latencyRecorder) alongside the existing attemptRecorder made this observable for the first time, since compiled tests never see it and the plugin previously carried only one such key. Give latencyRecorderCtxKey a field so its size differs from attemptRecorderCtxKey's, confirmed against a minimal Yaegi reproduction before and after the fix. Fix a second bug, reproducible under compiled Go too: io.Reader may legally return its final data together with io.EOF in one call, and Go's real net/http body reader does this for small, fully-buffered responses. The original TTFB capture lived inside the err == nil branch, so it never fired for a typical small non-streaming completion. Move the check ahead of the error branch so it fires on any n > 0, leaving the pinned timer-reset and Close semantics untouched. Tests mutate the real source per case and confirm each one fails first, recorded in each test's own doc comment.
This commit is contained in:
@@ -286,9 +286,26 @@ type adminProviderView struct {
|
||||
// marshals as "{}" for a provider with no known models yet, not
|
||||
// omitted.
|
||||
ModelRates map[string]adminModelRateView `json:"modelRates"`
|
||||
Type string `json:"type"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
LastErr string `json:"lastErr,omitempty"`
|
||||
// Latency is a compact per-stream ("streaming"/"non-streaming")
|
||||
// average TTFB/duration summary for this provider (feat: instrument
|
||||
// upstream latency), sourced from the SAME in-process g.latency
|
||||
// accumulator /metrics reads (metrics.go). Averages only — never a
|
||||
// full histogram or a per-model breakdown — to keep this already-
|
||||
// ~113KB, 5s-polled response from growing meaningfully; see
|
||||
// buildAdminLatencyViews' own doc comment. nil (omitted) for a
|
||||
// provider with no observations yet, or for a deployment with metrics
|
||||
// collection gated off entirely. Grouped here with the other map
|
||||
// fields above, not declared near AttemptsDay/ModelCount below,
|
||||
// because it is a pointer-shaped field (a map header) and this
|
||||
// struct's fields are ordered pointer-shaped-first, then plain
|
||||
// scalars (int64/int/bool) last — fieldalignment (govet, enabled via
|
||||
// this repo's global golangci-lint config) flags a pointer field
|
||||
// declared after the scalar block as costing extra GC pointer-scan
|
||||
// bytes. See timeout.go's watchdogBody for the identical convention.
|
||||
Latency map[string]adminLatencyView `json:"latency,omitempty"`
|
||||
Type string `json:"type"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
LastErr string `json:"lastErr,omitempty"`
|
||||
// HealthState is this provider's discovery circuit breaker state
|
||||
// (feat/provider-health): "closed" (normal), "open" (backing off
|
||||
// after repeated discovery failures — maybeRefresh skips it until
|
||||
@@ -347,6 +364,59 @@ type adminModelRateView struct {
|
||||
FailuresDay int64 `json:"failuresDay"`
|
||||
}
|
||||
|
||||
// adminLatencyView is one provider's one stream-state's compact latency
|
||||
// summary (feat: instrument upstream latency) — adminProviderView.
|
||||
// Latency's value type. Plain averages (sum/count from the same
|
||||
// latencySnapshot /metrics reads, metrics.go), never a percentile: no
|
||||
// interpolation-correctness question to get wrong, and "compact" (the
|
||||
// task brief's own word for this admin surface) rules out shipping a
|
||||
// full bucket set here anyway — that detail belongs on /metrics.
|
||||
// AvgTTFBMs is 0 (omitted) when no successful read was ever observed
|
||||
// for this stream state (latencySnapshot.ttfbCount == 0) — never
|
||||
// fabricated from AvgDurationMs, which is a materially different
|
||||
// number for a genuine stream (see llmgateway_upstream_ttfb_seconds'
|
||||
// own HELP text, metrics.go, for why the two must never be conflated).
|
||||
type adminLatencyView struct {
|
||||
AvgTTFBMs float64 `json:"avgTtfbMs,omitempty"`
|
||||
AvgDurationMs float64 `json:"avgDurationMs,omitempty"`
|
||||
Count int64 `json:"count,omitempty"`
|
||||
}
|
||||
|
||||
// buildAdminLatencyViews groups snaps (g.latency.snapshot(), metrics.go)
|
||||
// by provider, discarding the opt-in model dimension entirely (task
|
||||
// brief: "compact... do not bloat" — GET /admin/api/overview already
|
||||
// sits around 113 KB and is polled every 5s; a full per-model latency
|
||||
// breakdown belongs on /metrics, not here). The outer map is keyed by
|
||||
// provider name, the inner by "streaming"/"non-streaming"; a provider
|
||||
// with no observations at all is simply absent from the outer map, so
|
||||
// adminOverviewResponse's own lookup (buildAdminOverview, below) yields
|
||||
// nil for it — adminProviderView.Latency's own documented "nil means no
|
||||
// observations yet".
|
||||
func buildAdminLatencyViews(snaps []latencySnapshot) map[string]map[string]adminLatencyView {
|
||||
out := make(map[string]map[string]adminLatencyView)
|
||||
for _, s := range snaps {
|
||||
if s.key.model != "" || s.durationCount == 0 {
|
||||
continue
|
||||
}
|
||||
streamKey := "non-streaming"
|
||||
if s.key.streaming {
|
||||
streamKey = "streaming"
|
||||
}
|
||||
v := adminLatencyView{
|
||||
AvgDurationMs: s.durationSum / float64(s.durationCount) * 1000,
|
||||
Count: s.durationCount,
|
||||
}
|
||||
if s.ttfbCount > 0 {
|
||||
v.AvgTTFBMs = s.ttfbSum / float64(s.ttfbCount) * 1000
|
||||
}
|
||||
if out[s.key.provider] == nil {
|
||||
out[s.key.provider] = make(map[string]adminLatencyView)
|
||||
}
|
||||
out[s.key.provider][streamKey] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// adminModelMetaView is one upstream model's resolved metadata (feature
|
||||
// v0.23: context window, per-token cost) within its provider —
|
||||
// adminProviderView.ModelMeta's value type. Every field is a pointer,
|
||||
@@ -485,6 +555,11 @@ func (g *Gateway) buildAdminOverview() adminOverviewResponse {
|
||||
providerLevelCounters := allCounters[:len(snaps)]
|
||||
modelCounters := allCounters[len(snaps):]
|
||||
|
||||
// feat: instrument upstream latency — one g.latency.snapshot() call
|
||||
// for the whole response, mirroring providerUsage's own one-batched-
|
||||
// read-for-every-provider shape above, not a per-provider read.
|
||||
latencyViews := buildAdminLatencyViews(g.latency.snapshot())
|
||||
|
||||
providers := make([]adminProviderView, len(snaps))
|
||||
mi := 0
|
||||
for i, s := range snaps {
|
||||
@@ -517,6 +592,7 @@ func (g *Gateway) buildAdminOverview() adminOverviewResponse {
|
||||
FailuresMinute: pc.failuresMinute,
|
||||
ModelRates: modelRates,
|
||||
ModelMeta: modelMeta,
|
||||
Latency: latencyViews[s.name],
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
@@ -551,6 +552,85 @@ func TestAdminOverview_ProviderRates_EndToEndAfterTraffic(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdminOverview_LatencySummary_AvgReflectsSeededObservations proves
|
||||
// buildAdminLatencyViews' own math (admin.go, feat: instrument upstream
|
||||
// latency): plain averages (sum/count), split by stream state, and that
|
||||
// the opt-in per-model dimension never leaks into this compact admin
|
||||
// surface (the task brief's own "compact... do not bloat" requirement).
|
||||
//
|
||||
// alpha's observation is seeded directly via gw.latency.record under a
|
||||
// MODEL-keyed latencyKey only — bypassing recordLatency's own
|
||||
// dual-write (which always also writes the bare provider+stream key,
|
||||
// tested separately by metrics_test.go's TestMetrics_ModelLabel_*
|
||||
// family) — specifically so no base (model="") entry exists for
|
||||
// alpha's "streaming" state at all. That isolates buildAdminLatencyViews'
|
||||
// own model-exclusion filter: if it ever stopped skipping model-keyed
|
||||
// entries, THIS is the one that would leak through as a phantom row,
|
||||
// since there is no correctly-shaped base entry that could coincidentally
|
||||
// produce the same numbers instead (unlike a same-provider dual-write,
|
||||
// where the base and model entries are byte-identical and an overwrite
|
||||
// would go unnoticed).
|
||||
//
|
||||
// MUTATION VERIFIED: removing the `if s.key.model != "" ... continue`
|
||||
// guard from buildAdminLatencyViews (admin.go) made this test fail —
|
||||
// alpha.Latency gained a "streaming" entry (Count=1, AvgDurationMs=5000)
|
||||
// that must never exist, since alpha never received any base-key
|
||||
// observation. Reverted before committing.
|
||||
func TestAdminOverview_LatencySummary_AvgReflectsSeededObservations(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newAdminTestConfig()
|
||||
h, gw := newAdminGatewayHandle(t, cfg)
|
||||
|
||||
// zeta, non-streaming: two base-key observations, 100ms and 300ms —
|
||||
// average 200ms.
|
||||
gw.recordLatency("zeta", "", latencySample{duration: 100 * time.Millisecond, ttfb: 100 * time.Millisecond, hasTTFB: true, streaming: false})
|
||||
gw.recordLatency("zeta", "", latencySample{duration: 300 * time.Millisecond, ttfb: 300 * time.Millisecond, hasTTFB: true, streaming: false})
|
||||
|
||||
// alpha, streaming: model-keyed only (see doc comment above).
|
||||
gw.latency.record(latencyKey{provider: "alpha", streaming: true, model: "a-model-1"},
|
||||
latencySample{duration: 5 * time.Second, ttfb: time.Second, hasTTFB: true, streaming: true})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, adminRequest(http.MethodGet, adminOverviewPath, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var got adminOverviewResponse
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &got); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
var zeta, alpha *adminProviderView
|
||||
for i := range got.Providers {
|
||||
switch got.Providers[i].Name {
|
||||
case "zeta":
|
||||
zeta = &got.Providers[i]
|
||||
case "alpha":
|
||||
alpha = &got.Providers[i]
|
||||
}
|
||||
}
|
||||
if zeta == nil || alpha == nil {
|
||||
t.Fatalf("providers = %+v, want entries named zeta and alpha", got.Providers)
|
||||
}
|
||||
|
||||
nonStream, ok := zeta.Latency["non-streaming"]
|
||||
if !ok {
|
||||
t.Fatalf("zeta.Latency missing \"non-streaming\" entry: %+v", zeta.Latency)
|
||||
}
|
||||
if nonStream.Count != 2 {
|
||||
t.Errorf("non-streaming Count = %d, want 2", nonStream.Count)
|
||||
}
|
||||
if math.Abs(nonStream.AvgDurationMs-200) > 0.01 {
|
||||
t.Errorf("non-streaming AvgDurationMs = %v, want 200 ((100+300)/2)", nonStream.AvgDurationMs)
|
||||
}
|
||||
if math.Abs(nonStream.AvgTTFBMs-200) > 0.01 {
|
||||
t.Errorf("non-streaming AvgTTFBMs = %v, want 200", nonStream.AvgTTFBMs)
|
||||
}
|
||||
|
||||
if _, ok := alpha.Latency["streaming"]; ok {
|
||||
t.Errorf(`alpha.Latency has a "streaming" entry despite the only observation being model-keyed (never a base provider+stream one): %+v — the opt-in per-model dimension must never leak into this compact admin summary`, alpha.Latency)
|
||||
}
|
||||
}
|
||||
|
||||
// --- usage: math against seeded counters ---
|
||||
|
||||
func TestAdminUsage_MathAgainstSeededCounters(t *testing.T) {
|
||||
|
||||
+343
@@ -0,0 +1,343 @@
|
||||
package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// delayedReadCloser is a minimal io.ReadCloser test double for
|
||||
// watchdogBody's own latency unit tests below: its FIRST Read call
|
||||
// sleeps firstReadDelay before delivering one byte; its SECOND sleeps
|
||||
// afterFirstRead before returning io.EOF; every call after that is a
|
||||
// no-op io.EOF. This lets a test control exactly how much wall-clock
|
||||
// time falls BEFORE the first successful read (what TTFB must measure)
|
||||
// versus AFTER it (what must count toward duration but never toward
|
||||
// TTFB) without a real network connection.
|
||||
type delayedReadCloser struct {
|
||||
firstReadDelay time.Duration
|
||||
afterFirstRead time.Duration
|
||||
read int
|
||||
}
|
||||
|
||||
func (d *delayedReadCloser) Read(p []byte) (int, error) {
|
||||
d.read++
|
||||
switch d.read {
|
||||
case 1:
|
||||
time.Sleep(d.firstReadDelay)
|
||||
p[0] = 'x'
|
||||
return 1, nil
|
||||
default:
|
||||
time.Sleep(d.afterFirstRead)
|
||||
return 0, io.EOF
|
||||
}
|
||||
}
|
||||
|
||||
func (d *delayedReadCloser) Close() error { return nil }
|
||||
|
||||
// TestChatCompletion_NonStreaming_SmallBufferedBody_HasTTFB proves TTFB
|
||||
// is still captured when the underlying reader combines its final data
|
||||
// with io.EOF in ONE Read call — io.Reader's own contract explicitly
|
||||
// permits this ("an instance ... may return either err == EOF or err ==
|
||||
// nil"), and Go's real net/http body reader does exactly this for a
|
||||
// small, fully-buffered response: measured directly against this test's
|
||||
// own httptest server, a ~200-byte JSON completion arrives as a single
|
||||
// Read call reporting (n=200, err=io.EOF) together, not as a data call
|
||||
// followed by a separate EOF call. This is the common case for a
|
||||
// typical, fully-buffered non-streaming completion — most real traffic
|
||||
// — so getting it wrong here means hasTTFB is false for most requests,
|
||||
// not an edge case.
|
||||
//
|
||||
// MUTATION VERIFIED: moving the TTFB-capture block (timeout.go's Read)
|
||||
// back inside the `if err != nil { ... return ... }` branch's implicit
|
||||
// else (i.e., gating it on err == nil, matching the pre-fix shape) made
|
||||
// this test fail — samples[0].hasTTFB read false, since this test's own
|
||||
// server response arrives in one combined (n>0, io.EOF) call. Reverted
|
||||
// before committing.
|
||||
func TestChatCompletion_NonStreaming_SmallBufferedBody_HasTTFB(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a := newTimeoutTestAdapter(srv.URL, 5*time.Second)
|
||||
|
||||
var mu sync.Mutex
|
||||
var samples []latencySample
|
||||
ctx := withLatencyRecorder(context.Background(), func(s latencySample) {
|
||||
mu.Lock()
|
||||
samples = append(samples, s)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
if _, err := a.chatCompletion(ctx, rec, map[string]any{"model": "gpt-5", "messages": []any{}}); err != nil {
|
||||
t.Fatalf("chatCompletion: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(samples) != 1 {
|
||||
t.Fatalf("latencyRecorder invoked %d times, want 1: %+v", len(samples), samples)
|
||||
}
|
||||
if samples[0].streaming {
|
||||
t.Error("streaming = true, want false (Content-Type: application/json)")
|
||||
}
|
||||
if !samples[0].hasTTFB {
|
||||
t.Fatal("hasTTFB = false, want true — a small, fully-buffered non-streaming response must still report a TTFB, even when its data arrives combined with io.EOF in one Read call")
|
||||
}
|
||||
if samples[0].ttfb <= 0 {
|
||||
t.Errorf("ttfb = %s, want > 0", samples[0].ttfb)
|
||||
}
|
||||
if samples[0].ttfb > samples[0].duration {
|
||||
t.Errorf("ttfb (%s) > duration (%s) — ttfb must never exceed the total duration it is a prefix of", samples[0].ttfb, samples[0].duration)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWatchdogBody_ArmLatency_TTFBAtFirstRead_NotConstructionOrClose
|
||||
// proves the task brief's core timing requirement directly, at the
|
||||
// watchdogBody level: TTFB is captured strictly at the moment of the
|
||||
// FIRST successful (n>0) Read, never at armLatency's own construction-
|
||||
// time "start" and never re-derived at Close. firstReadDelay elapses
|
||||
// BEFORE the first Read returns; afterFirstRead elapses AFTER it,
|
||||
// before Close. A correct capture reports ttfb in
|
||||
// [firstReadDelay, firstReadDelay+afterFirstRead) and duration >=
|
||||
// firstReadDelay+afterFirstRead.
|
||||
//
|
||||
// MUTATION VERIFIED: changing Close (timeout.go) to unconditionally set
|
||||
// `sample.ttfb = time.Since(wb.latencyStart); sample.hasTTFB = true`
|
||||
// (ignoring the Read-captured wb.ttfb/ttfbRecorded entirely) made this
|
||||
// test fail: got.ttfb read ~180ms (firstReadDelay+afterFirstRead, i.e.
|
||||
// Close time), failing the "< firstReadDelay+afterFirstRead" assertion
|
||||
// below. Reverted before committing.
|
||||
func TestWatchdogBody_ArmLatency_TTFBAtFirstRead_NotConstructionOrClose(t *testing.T) {
|
||||
const (
|
||||
firstReadDelay = 60 * time.Millisecond
|
||||
afterFirstRead = 120 * time.Millisecond
|
||||
)
|
||||
rc := &delayedReadCloser{firstReadDelay: firstReadDelay, afterFirstRead: afterFirstRead}
|
||||
_, cancel := context.WithCancel(context.Background())
|
||||
wb := newWatchdogBody(rc, cancel, 5*time.Second, "p1", nil)
|
||||
|
||||
var mu sync.Mutex
|
||||
var got *latencySample
|
||||
wb.armLatency(time.Now(), false, func(s latencySample) {
|
||||
mu.Lock()
|
||||
got = &s
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
buf := make([]byte, 16)
|
||||
if _, err := wb.Read(buf); err != nil {
|
||||
t.Fatalf("first Read: %v", err)
|
||||
}
|
||||
if _, err := wb.Read(buf); err != io.EOF {
|
||||
t.Fatalf("second Read: err = %v, want io.EOF", err)
|
||||
}
|
||||
if err := wb.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if got == nil {
|
||||
t.Fatal("latencyRecorder never invoked")
|
||||
}
|
||||
if !got.hasTTFB {
|
||||
t.Fatal("hasTTFB = false, want true — one successful read did happen")
|
||||
}
|
||||
if got.ttfb < firstReadDelay {
|
||||
t.Errorf("ttfb = %s, want >= %s (the first read's own delay)", got.ttfb, firstReadDelay)
|
||||
}
|
||||
if got.ttfb >= firstReadDelay+afterFirstRead {
|
||||
t.Errorf("ttfb = %s, want < %s — ttfb must not include time spent waiting on the SECOND read, proving it was captured at the FIRST read, not re-derived at Close", got.ttfb, firstReadDelay+afterFirstRead)
|
||||
}
|
||||
if got.duration < firstReadDelay+afterFirstRead {
|
||||
t.Errorf("duration = %s, want >= %s (the whole body's lifetime, both reads)", got.duration, firstReadDelay+afterFirstRead)
|
||||
}
|
||||
}
|
||||
|
||||
// TestWatchdogBody_ArmLatency_NeverArmed_NoSampleEmitted proves the
|
||||
// "no measurable hot-path cost, nothing collected" contract for a
|
||||
// watchdogBody that is never armed at all (the disabled-metrics path,
|
||||
// end to end: upstreamBytes/proxyUpstream only call armLatency when a
|
||||
// non-nil latencyRecorder was actually looked up from context) — Read
|
||||
// and Close must both behave exactly as they did before this feature
|
||||
// existed, and nothing resembling a latencySample must ever be built or
|
||||
// reported.
|
||||
//
|
||||
// MUTATION VERIFIED: changing Close (timeout.go) to call wb.onLatency
|
||||
// unconditionally (dropping the `wb.onLatency != nil` guard) made this
|
||||
// test panic with a nil-function-call ("runtime error: invalid memory
|
||||
// address or nil pointer dereference"), which go test reports as a
|
||||
// failure — confirming the guard is what this test actually exercises.
|
||||
// Reverted before committing.
|
||||
func TestWatchdogBody_ArmLatency_NeverArmed_NoSampleEmitted(t *testing.T) {
|
||||
rc := &delayedReadCloser{firstReadDelay: time.Millisecond, afterFirstRead: time.Millisecond}
|
||||
_, cancel := context.WithCancel(context.Background())
|
||||
wb := newWatchdogBody(rc, cancel, 5*time.Second, "p1", nil)
|
||||
// armLatency deliberately never called.
|
||||
|
||||
buf := make([]byte, 16)
|
||||
if _, err := wb.Read(buf); err != nil {
|
||||
t.Fatalf("first Read: %v", err)
|
||||
}
|
||||
if _, err := wb.Read(buf); err != io.EOF {
|
||||
t.Fatalf("second Read: err = %v, want io.EOF", err)
|
||||
}
|
||||
if err := wb.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
if wb.onLatency != nil {
|
||||
t.Error("wb.onLatency is non-nil despite armLatency never being called")
|
||||
}
|
||||
}
|
||||
|
||||
// TestChatCompletion_Stream_LongSteadyStream_ReportsSmallTTFBLargeDuration
|
||||
// is the end-to-end proof for the exact case this whole feature exists
|
||||
// to distinguish (task brief, "The signal, and the trap"): a stream
|
||||
// whose total wall-clock duration is large, but whose FIRST chunk
|
||||
// arrives quickly and every later chunk keeps the connection making
|
||||
// progress, must report a SMALL time-to-first-byte and a LARGE total
|
||||
// duration — never a duration that collapses to roughly the same value
|
||||
// as TTFB (which is what a non-streaming, fully-buffered response looks
|
||||
// like instead).
|
||||
//
|
||||
// MUTATION VERIFIED: deleting the `if n > 0 && wb.onLatency != nil && ...
|
||||
// { ... }` TTFB-capture block from watchdogBody.Read (timeout.go)
|
||||
// entirely made this test fail: ttfbRecorded never flips, hasTTFB stays
|
||||
// false on every sample, and the `!s.hasTTFB` assertion below fails
|
||||
// immediately. Reverted before committing.
|
||||
func TestChatCompletion_Stream_LongSteadyStream_ReportsSmallTTFBLargeDuration(t *testing.T) {
|
||||
const (
|
||||
chunkGap = 40 * time.Millisecond
|
||||
numChunks = 5 // 5*40ms = 200ms total: a large duration by this test's own standard
|
||||
)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fl := w.(http.Flusher)
|
||||
for i := 0; i < numChunks; i++ {
|
||||
time.Sleep(chunkGap)
|
||||
_, _ = fmt.Fprintf(w, "data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"tok%d\"}}]}\n\n", i)
|
||||
fl.Flush()
|
||||
}
|
||||
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
fl.Flush()
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
a := newTimeoutTestAdapter(srv.URL, 5*time.Second) // generous: the idle watchdog must never fire
|
||||
|
||||
var mu sync.Mutex
|
||||
var samples []latencySample
|
||||
ctx := withLatencyRecorder(context.Background(), func(s latencySample) {
|
||||
mu.Lock()
|
||||
samples = append(samples, s)
|
||||
mu.Unlock()
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
if _, err := a.chatCompletion(ctx, rec, map[string]any{"model": "gpt-5", "messages": []any{}, "stream": true}); err != nil {
|
||||
t.Fatalf("chatCompletion: %v", err)
|
||||
}
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(samples) != 1 {
|
||||
t.Fatalf("latencyRecorder invoked %d times, want 1: %+v", len(samples), samples)
|
||||
}
|
||||
s := samples[0]
|
||||
if !s.streaming {
|
||||
t.Error("streaming = false, want true (Content-Type: text/event-stream)")
|
||||
}
|
||||
if !s.hasTTFB {
|
||||
t.Fatal("hasTTFB = false, want true — the stream's first chunk was read successfully")
|
||||
}
|
||||
if s.ttfb >= 2*chunkGap {
|
||||
t.Errorf("ttfb = %s, want well under %s (roughly one chunkGap, not the whole stream)", s.ttfb, 2*chunkGap)
|
||||
}
|
||||
wantMinDuration := time.Duration(numChunks) * chunkGap
|
||||
if s.duration < wantMinDuration {
|
||||
t.Errorf("duration = %s, want at least %s (the whole stream's wall-clock time)", s.duration, wantMinDuration)
|
||||
}
|
||||
if s.duration <= 2*s.ttfb {
|
||||
t.Errorf("duration (%s) is not meaningfully larger than ttfb (%s) — a long, steadily-emitting stream must report a small TTFB and a large duration, the whole point of this feature", s.duration, s.ttfb)
|
||||
}
|
||||
}
|
||||
|
||||
// TestProxyUpstream_Latency_ArmedFromContext_StreamingDerivedFromContentType
|
||||
// proves proxyUpstream (routes_passthrough.go, native passthrough's
|
||||
// shared core) wires latency the same way upstreamBytes (providers.go)
|
||||
// does: armed with whatever latencyRecorder its caller's request context
|
||||
// carries, with streaming derived from the upstream response's own
|
||||
// Content-Type — not from anything passthrough itself would have no way
|
||||
// to know (it forwards an opaque body, never decoding a "stream"
|
||||
// request field the way an adapter's chatCompletion does). The
|
||||
// recorder is injected directly onto the request's context here (cfg.
|
||||
// Metrics is left unset, so handlePassthrough's own metricsEnabled-
|
||||
// gated wiring never runs and never overwrites it) — this test is about
|
||||
// proxyUpstream's OWN arming/derivation logic, not about the
|
||||
// metricsEnabled gate itself (covered separately).
|
||||
//
|
||||
// MUTATION VERIFIED: removing the `if latRec != nil { wb.armLatency(...)
|
||||
// }` block from proxyUpstream (routes_passthrough.go) made this test
|
||||
// fail — the spy recorder was never invoked, so `len(samples) != 1`
|
||||
// failed with 0. Reverted before committing.
|
||||
func TestProxyUpstream_Latency_ArmedFromContext_StreamingDerivedFromContentType(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fl := w.(http.Flusher)
|
||||
_, _ = fmt.Fprint(w, "data: {\"hello\":\"world\"}\n\n")
|
||||
fl.Flush()
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := CreateConfig()
|
||||
cfg.Providers = map[string]*ProviderConfig{
|
||||
"openai": {Type: "openai", BaseURL: srv.URL, APIKey: "sk-up"},
|
||||
}
|
||||
cfg.Groups = map[string]*GroupConfig{"default": {}}
|
||||
cfg.Users = &UsersConfig{Inline: []*UserConfig{
|
||||
{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}},
|
||||
}}
|
||||
next := http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})
|
||||
h, err := New(context.Background(), next, cfg, "llmgw")
|
||||
if err != nil {
|
||||
t.Fatalf("New: %v", err)
|
||||
}
|
||||
gw := h.(*Gateway)
|
||||
gw.limiter.spawn = func(f func()) { f() }
|
||||
|
||||
var mu sync.Mutex
|
||||
var samples []latencySample
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/native-endpoint", nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-alice")
|
||||
req = req.WithContext(withAttemptRecorder(req.Context(), func(*http.Response, error) {}))
|
||||
req = req.WithContext(withLatencyRecorder(req.Context(), func(s latencySample) {
|
||||
mu.Lock()
|
||||
samples = append(samples, s)
|
||||
mu.Unlock()
|
||||
}))
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, req)
|
||||
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if len(samples) != 1 {
|
||||
t.Fatalf("latencyRecorder invoked %d times, want 1", len(samples))
|
||||
}
|
||||
if !samples[0].streaming {
|
||||
t.Error("streaming = false, want true (Content-Type: text/event-stream)")
|
||||
}
|
||||
}
|
||||
+58
-32
@@ -537,24 +537,16 @@ const maxExplicitBodyAdmissionCap = 10_000
|
||||
// zero-waste sequence, then hand-applied here so every field keeps its
|
||||
// original doc comment) — see providerState's own doc comment
|
||||
// (registry.go) for the general convention.
|
||||
// Field order below is fieldalignment-derived (govet, enabled via this
|
||||
// repo's global golangci-lint config), not declaration-convenience
|
||||
// order: pointer-shaped fields (interfaces, pointers, maps, chans,
|
||||
// strings, slices) are grouped first, then plain scalars (logGate/
|
||||
// failoverConfig/time.Duration) last, so the GC's pointer-scan span
|
||||
// stays as short as possible. See adminProviderView's identical
|
||||
// convention (admin.go) for the same reasoning spelled out once; this
|
||||
// struct is the other place it applies.
|
||||
type Gateway struct {
|
||||
next http.Handler
|
||||
// targetClient is the shared, connection-pooled *http.Client the
|
||||
// MCP/A2A target proxy (mcp_a2a.go) issues every upstream request
|
||||
// through — built once via newAdapterHTTPClient, the same constructor
|
||||
// each provider adapter uses for its own client.
|
||||
targetClient *http.Client
|
||||
auth *authStore
|
||||
limiter *limiter
|
||||
registry *modelRegistry
|
||||
adapters map[string]providerAdapter
|
||||
cfg *Config
|
||||
// cache is nil whenever response caching is not configured or not
|
||||
// usable (cfg.Cache.Enabled is false, or true with no config.Redis —
|
||||
// see buildResponseCache, cache.go). Every call site checks for nil
|
||||
// before using it, rather than responseCache having its own
|
||||
// always-disabled zero value.
|
||||
cache *responseCache
|
||||
// redisClient is the same instance newGateway hands to both the
|
||||
// limiter's redisStore and the response cache (its own doc comment,
|
||||
// below, explains why it's built once and shared) — kept here too,
|
||||
@@ -562,13 +554,6 @@ type Gateway struct {
|
||||
// pooled connections through. nil when config.Redis is absent, same
|
||||
// as buildRedisClient's own nil-for-unconfigured contract.
|
||||
redisClient *respClient
|
||||
// bodyAdmission is the buffered-channel semaphore acquireBodyAdmission
|
||||
// (routes_unified.go) claims from and releases: security review
|
||||
// finding 1b, 2026-08-22. Sized once, here, by newGateway (see
|
||||
// defaultBodyAdmissionCap/Config.MaxInFlightBodyRequests) — never
|
||||
// resized afterward, matching a Go channel's own fixed-capacity
|
||||
// contract.
|
||||
bodyAdmission chan struct{}
|
||||
// failoverHealth is feat/failover's own per-pod, in-memory
|
||||
// request-path health signal (failover.go), built once, here, by
|
||||
// newGateway. Every method on it is nil-receiver-safe, so a Gateway
|
||||
@@ -576,12 +561,43 @@ type Gateway struct {
|
||||
// degrades to "failover never skips a candidate for request health"
|
||||
// rather than a nil-pointer panic.
|
||||
failoverHealth *requestHealthTracker
|
||||
// metricsNets is Config.Metrics.AllowedCIDRs, parsed once at
|
||||
// construction (parseMetricsCIDRs, metrics.go) so the request path
|
||||
// never re-parses a CIDR string per scrape. nil when Metrics is
|
||||
// unconfigured or carries no AllowedCIDRs entries.
|
||||
metricsNets []*net.IPNet
|
||||
name string
|
||||
limiter *limiter
|
||||
registry *modelRegistry
|
||||
adapters map[string]providerAdapter
|
||||
cfg *Config
|
||||
auth *authStore
|
||||
// cache is nil whenever response caching is not configured or not
|
||||
// usable (cfg.Cache.Enabled is false, or true with no config.Redis —
|
||||
// see buildResponseCache, cache.go). Every call site checks for nil
|
||||
// before using it, rather than responseCache having its own
|
||||
// always-disabled zero value.
|
||||
cache *responseCache
|
||||
// bodyAdmission is the buffered-channel semaphore acquireBodyAdmission
|
||||
// (routes_unified.go) claims from and releases: security review
|
||||
// finding 1b, 2026-08-22. Sized once, here, by newGateway (see
|
||||
// defaultBodyAdmissionCap/Config.MaxInFlightBodyRequests) — never
|
||||
// resized afterward, matching a Go channel's own fixed-capacity
|
||||
// contract.
|
||||
bodyAdmission chan struct{}
|
||||
// targetClient is the shared, connection-pooled *http.Client the
|
||||
// MCP/A2A target proxy (mcp_a2a.go) issues every upstream request
|
||||
// through — built once via newAdapterHTTPClient, the same constructor
|
||||
// each provider adapter uses for its own client.
|
||||
targetClient *http.Client
|
||||
// latency is the in-process upstream-latency accumulator (feat:
|
||||
// instrument upstream latency, metrics.go's latencyStore) —
|
||||
// recordLatency's only write target, and writeLatencyMetrics/
|
||||
// buildAdminOverview's (admin.go) own read source. Always
|
||||
// constructed, here, by newGateway, even when Config.Metrics is nil
|
||||
// or disabled: g.latency simply never receives anything to
|
||||
// accumulate in that case (every withLatencyRecorder call site gates
|
||||
// its own wiring on metricsEnabled(g.cfg), not this field's
|
||||
// existence) — every method on it is also nil-receiver-safe
|
||||
// (latencyStore's own doc comment), matching failoverHealth's
|
||||
// identical "safe even off a bare &Gateway{} literal" convention
|
||||
// above.
|
||||
latency *latencyStore
|
||||
name string
|
||||
// failoverLogGate rate-limits runMeteredCall's generic "failing over"
|
||||
// log line (routes_unified.go) to once per storeErrorLogEvery
|
||||
// (adversarial-review fix, F10) — reuses auth.go's own logGate type,
|
||||
@@ -590,6 +606,14 @@ type Gateway struct {
|
||||
// on a fresh Gateway always logs. Deliberately does NOT gate the 404
|
||||
// loud-log line (operator ruling: that one must always log).
|
||||
failoverLogGate logGate
|
||||
// metricsNets is Config.Metrics.AllowedCIDRs, parsed once at
|
||||
// construction (parseMetricsCIDRs, metrics.go) so the request path
|
||||
// never re-parses a CIDR string per scrape. nil when Metrics is
|
||||
// unconfigured or carries no AllowedCIDRs entries.
|
||||
metricsNets []*net.IPNet
|
||||
// failover is Config.Failover, validated and resolved once by
|
||||
// newGateway (validateFailoverConfig, failover.go).
|
||||
failover failoverConfig
|
||||
// targetTimeout is targetClient's own resolved request timeout —
|
||||
// Config.RequestTimeout only; an MCP/A2A target has no per-target
|
||||
// override the way a provider does. proxyUpstream's own callers
|
||||
@@ -597,9 +621,6 @@ type Gateway struct {
|
||||
// idle-progress body watchdog (timeout.go) matches the
|
||||
// ResponseHeaderTimeout already set on targetClient's Transport.
|
||||
targetTimeout time.Duration
|
||||
// failover is Config.Failover, validated and resolved once by
|
||||
// newGateway (validateFailoverConfig, failover.go).
|
||||
failover failoverConfig
|
||||
}
|
||||
|
||||
// telemetryStartupOnce keeps the anonymous "plugin loaded" ping to one per
|
||||
@@ -718,6 +739,11 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str
|
||||
}
|
||||
g.failover = failoverCfg
|
||||
g.failoverHealth = newRequestHealthTracker()
|
||||
// feat: instrument upstream latency — always constructed, regardless
|
||||
// of whether Config.Metrics is nil/disabled; see g.latency's own doc
|
||||
// comment above for why an always-present, usually-empty store is the
|
||||
// right default rather than a conditional nil.
|
||||
g.latency = newLatencyStore()
|
||||
|
||||
// A malformed Config.Metrics.AllowedCIDRs entry is a constructor
|
||||
// error (validate-at-construction), not a silently-ignored allowlist
|
||||
|
||||
+345
@@ -8,6 +8,8 @@ import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
@@ -361,6 +363,56 @@ func (m *metricWriter) sampleFloat(name string, labels []metricLabel, v float64)
|
||||
m.sample(name, labels, strconv.FormatFloat(v, 'g', -1, 64))
|
||||
}
|
||||
|
||||
// histogram writes one Prometheus histogram family's samples for one
|
||||
// label set (feat: instrument upstream latency): cumulative `name+
|
||||
// "_bucket"` series in ascending `le` order, ending with `le="+Inf"`,
|
||||
// followed by `name+"_sum"` then `name+"_count"` — the exact shape a
|
||||
// real scrape parser requires (see escapeLabelValue's own doc comment
|
||||
// for what getting exposition format wrong costs in general: this is
|
||||
// not "one bad metric gets dropped", it is "the whole scrape fails to
|
||||
// parse and the target reports DOWN"). Getting a histogram wrong is
|
||||
// worse than getting a label escape wrong: a non-cumulative or
|
||||
// out-of-order bucket set, or a missing +Inf terminator, are ALSO
|
||||
// spec violations a strict parser rejects outright.
|
||||
//
|
||||
// buckets is DENSE, one entry per latencyBucketBounds index — index i
|
||||
// counts every observation in (latencyBucketBounds[i-1],
|
||||
// latencyBucketBounds[i]] — converted to the CUMULATIVE form the
|
||||
// exposition format requires here, once, at render time
|
||||
// (observeLatencyBucket, this file's own accumulation-side counterpart,
|
||||
// only ever needs to bump one dense bucket per observation, never
|
||||
// rewrite every bucket at or above its own). overflow is the +Inf
|
||||
// bucket's own dense count: every observation past the largest
|
||||
// configured bound. _count is derived from the SAME cumulative running
|
||||
// total the last +Inf bucket line already computed, rather than a
|
||||
// separately-tracked counter passed in alongside sum — by construction,
|
||||
// these can never diverge from what the bucket lines above them show,
|
||||
// closing off exactly the "buckets and _count silently disagree" failure
|
||||
// mode a real parser would also reject.
|
||||
//
|
||||
// labels is never mutated or its backing array reused across calls: a
|
||||
// fresh []metricLabel is built per bucket line instead of append(labels,
|
||||
// ...), so this is safe regardless of what spare capacity the caller's
|
||||
// own labels slice happens to have.
|
||||
func (m *metricWriter) histogram(name string, labels []metricLabel, buckets []int64, overflow int64, sum float64) {
|
||||
cumulative := int64(0)
|
||||
for i, bound := range latencyBucketBounds {
|
||||
cumulative += buckets[i]
|
||||
bucketLabels := make([]metricLabel, len(labels)+1)
|
||||
copy(bucketLabels, labels)
|
||||
bucketLabels[len(labels)] = metricLabel{"le", strconv.FormatFloat(bound, 'g', -1, 64)}
|
||||
m.sampleInt(name+"_bucket", bucketLabels, cumulative)
|
||||
}
|
||||
cumulative += overflow
|
||||
infLabels := make([]metricLabel, len(labels)+1)
|
||||
copy(infLabels, labels)
|
||||
infLabels[len(labels)] = metricLabel{"le", "+Inf"}
|
||||
m.sampleInt(name+"_bucket", infLabels, cumulative)
|
||||
|
||||
m.sampleFloat(name+"_sum", labels, sum)
|
||||
m.sampleInt(name+"_count", labels, cumulative)
|
||||
}
|
||||
|
||||
// aggregationNoteStoreBacked is appended to every store-backed family's
|
||||
// HELP text below (item 1, adversarial verification 2026-08-23,
|
||||
// reproduced against a 3-replica Traefik deployment): windowKey
|
||||
@@ -415,6 +467,7 @@ func (g *Gateway) renderMetrics() []byte {
|
||||
var m metricWriter
|
||||
g.writeUsageMetrics(&m)
|
||||
g.writeProviderMetrics(&m)
|
||||
g.writeLatencyMetrics(&m)
|
||||
g.writeRejectionMetrics(&m)
|
||||
g.writeStoreHealthMetrics(&m)
|
||||
return m.buf.Bytes()
|
||||
@@ -616,6 +669,298 @@ func (g *Gateway) writeProviderMetrics(m *metricWriter) {
|
||||
}
|
||||
}
|
||||
|
||||
// --- upstream latency (feat: instrument upstream latency) ---
|
||||
//
|
||||
// The task: "priority-based throttling that engages when providers are
|
||||
// under strain" needs real data on what strain looks like first — this
|
||||
// section adds the measurement and nothing else. No throttling, no
|
||||
// priorities, no control logic reads any of it; it exists purely to be
|
||||
// scraped and looked at.
|
||||
|
||||
// latencyBucketBounds are the shared histogram bucket upper bounds, in
|
||||
// seconds, for both llmgateway_upstream_ttfb_seconds and
|
||||
// llmgateway_upstream_duration_seconds below.
|
||||
//
|
||||
// Bucket choice, reasoned from how upstream LLM latency actually
|
||||
// distributes: total upstream duration is roughly TTFB + output_tokens/
|
||||
// tokens_per_sec (task brief), so the useful range spans sub-second
|
||||
// (a fast TTFB, or a very short completion) through several minutes (a
|
||||
// long generation, or a request approaching defaultRequestTimeout,
|
||||
// timeout.go — 5 minutes, the point this gateway aborts an idle
|
||||
// upstream outright, regardless of Config.RequestTimeout overrides).
|
||||
//
|
||||
// - 0.1/0.25/0.5/1s: resolves a fast TTFB and short completions at
|
||||
// useful granularity — most of a healthy provider's TTFB population
|
||||
// lives in this range.
|
||||
// - 2/5/10/15/30/60s: the bulk of real chat completions, where most
|
||||
// observations should land; finer-grained here than at the extremes
|
||||
// because this is where an operator needs to tell "a bit slower"
|
||||
// from "a lot slower".
|
||||
// - 120/300s: brackets defaultRequestTimeout itself (300s = 5
|
||||
// minutes), so mass building up in this range is visibly "close to
|
||||
// where this gateway would abort the request outright", not just
|
||||
// "slow".
|
||||
// - 600s (10 minutes): headroom for a deployment that configures a
|
||||
// longer RequestTimeout override; anything past it falls into the
|
||||
// mandatory +Inf bucket.
|
||||
//
|
||||
// Both histograms share this one bound set deliberately, not two
|
||||
// independently tuned ones: it lets TTFB and duration be read off the
|
||||
// same axis on one dashboard, and one set is one less thing to keep in
|
||||
// sync as this feature evolves.
|
||||
var latencyBucketBounds = []float64{0.1, 0.25, 0.5, 1, 2, 5, 10, 15, 30, 60, 120, 300, 600}
|
||||
|
||||
// aggregationNoteLatency is appended to both upstream-latency histogram
|
||||
// families' HELP text: g.latency (latencyStore, below) is in-process
|
||||
// only, like limiter.rejections (aggregationNotePerProcess above) —
|
||||
// NEVER written to or read from Redis regardless of Config.Redis, unlike
|
||||
// the store-backed request/token/cost/provider-attempt families
|
||||
// elsewhere in this file, which flip between sum() and max() depending
|
||||
// on whether Redis is configured (aggregationNoteStoreBacked's own doc
|
||||
// comment). Each replica reports only the upstream traffic it personally
|
||||
// handled, so sum() (or, for a real percentile, Prometheus's own
|
||||
// histogram_quantile over a sum()-aggregated set) is always the correct
|
||||
// cross-replica aggregation here, unconditionally.
|
||||
const aggregationNoteLatency = " AGGREGATION ACROSS REPLICAS: this histogram lives only in this process's memory, never in Redis, regardless of Config.Redis — each replica reports only the upstream traffic it personally handled, so aggregate with sum() (histogram_quantile over a sum()-aggregated set, for a real percentile), which is always correct for it, unlike the store-backed families above that flip between sum() and max() depending on whether Redis is configured."
|
||||
|
||||
// latencySample is what one completed upstream watchdogBody (timeout.go)
|
||||
// hands its latencyRecorder exactly once, from Close: duration always
|
||||
// (the whole body's lifetime, start to Close); ttfb only when hasTTFB —
|
||||
// a body that never delivered one successful byte (e.g. the watchdog
|
||||
// fired before any read succeeded) has nothing meaningful to report
|
||||
// there, and must not fabricate a zero. streaming is the distinguishing
|
||||
// label the task brief requires, derived once by the caller from the
|
||||
// response's own Content-Type (isEventStreamResponse, timeout.go).
|
||||
type latencySample struct {
|
||||
duration time.Duration
|
||||
ttfb time.Duration
|
||||
hasTTFB bool
|
||||
streaming bool
|
||||
}
|
||||
|
||||
// latencyKey identifies one label set's histogram pair within
|
||||
// latencyStore: provider and streaming are always populated (the task
|
||||
// brief's "always" cardinality); model is empty unless Config.Metrics.
|
||||
// ModelLabel opted into the per-model breakdown AND the caller actually
|
||||
// knew the upstream model at record time (recordLatency below —
|
||||
// native passthrough never does, mirroring recordProviderAttempt's own
|
||||
// identical limitation, limits.go).
|
||||
type latencyKey struct {
|
||||
provider string
|
||||
model string
|
||||
streaming bool
|
||||
}
|
||||
|
||||
// latencyHistogram accumulates one latencyKey's two histograms (TTFB,
|
||||
// duration): DENSE per-bucket occurrence counts (observeLatencyBucket
|
||||
// below converts to the exposition format's required CUMULATIVE form
|
||||
// only at render time, metricWriter.histogram above), each with its own
|
||||
// running sum/count, plus overflow — the +Inf bucket's own dense count,
|
||||
// every observation past the largest configured bound.
|
||||
type latencyHistogram struct {
|
||||
ttfbBuckets []int64
|
||||
durationBuckets []int64
|
||||
ttfbSum float64
|
||||
durationSum float64
|
||||
ttfbCount int64
|
||||
ttfbOverflow int64
|
||||
durationCount int64
|
||||
durationOverflow int64
|
||||
}
|
||||
|
||||
// newLatencyHistogram allocates one latencyHistogram with dense bucket
|
||||
// slices sized to latencyBucketBounds, freshly zeroed — called at most
|
||||
// once per DISTINCT latencyKey ever observed (latencyStore.record
|
||||
// below), never per observation.
|
||||
func newLatencyHistogram() *latencyHistogram {
|
||||
return &latencyHistogram{
|
||||
ttfbBuckets: make([]int64, len(latencyBucketBounds)),
|
||||
durationBuckets: make([]int64, len(latencyBucketBounds)),
|
||||
}
|
||||
}
|
||||
|
||||
// observeLatencyBucket bumps the correct dense bucket (or overflow, for
|
||||
// a value past every configured bound) for one observation, and updates
|
||||
// the running sum/count alongside it — shared by latencyHistogram's own
|
||||
// two observations (TTFB, duration) so the bucket-selection rule can
|
||||
// never drift between them. latencyBucketBounds is ascending, so the
|
||||
// first bound the observation is <= is its bucket, matching Prometheus's
|
||||
// own "le" (less-than-or-equal) histogram semantics exactly.
|
||||
func observeLatencyBucket(buckets []int64, overflow *int64, sum *float64, count *int64, d time.Duration) {
|
||||
v := d.Seconds()
|
||||
*sum += v
|
||||
*count++
|
||||
for i, bound := range latencyBucketBounds {
|
||||
if v <= bound {
|
||||
buckets[i]++
|
||||
return
|
||||
}
|
||||
}
|
||||
*overflow++
|
||||
}
|
||||
|
||||
// latencyStore is the Gateway's in-process, per-replica upstream-latency
|
||||
// accumulator (feat: instrument upstream latency) — see
|
||||
// aggregationNoteLatency above for why sum(), never max(), is always the
|
||||
// correct cross-replica aggregation for it.
|
||||
//
|
||||
// Every method is nil-receiver-safe, mirroring requestHealthTracker's
|
||||
// own convention (failoverHealth, llmgateway.go): a Gateway assembled
|
||||
// directly as a bare &Gateway{} literal (bypassing newGateway, as a few
|
||||
// older tests do) degrades to "no latency observed", rather than a
|
||||
// nil-pointer panic, even though newGateway itself always constructs a
|
||||
// real one.
|
||||
type latencyStore struct {
|
||||
// data declared before mu (fieldalignment, govet): a map header is
|
||||
// fully pointer-shaped, while sync.Mutex is a plain int32+uint32 pair
|
||||
// with no pointer at all — grouping the pointer-shaped field first
|
||||
// keeps the GC's pointer-scan span as short as possible, the same
|
||||
// convention Gateway and adminProviderView follow (llmgateway.go,
|
||||
// admin.go).
|
||||
data map[latencyKey]*latencyHistogram
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// newLatencyStore returns an empty, ready-to-use latencyStore.
|
||||
func newLatencyStore() *latencyStore {
|
||||
return &latencyStore{data: make(map[latencyKey]*latencyHistogram)}
|
||||
}
|
||||
|
||||
// record accumulates sample into key's histogram, allocating a fresh
|
||||
// latencyHistogram the first time key is seen. Called at most twice per
|
||||
// completed upstream body — recordLatency below's own dual-write, once
|
||||
// for the always-on provider+streaming key and once more for the opt-in
|
||||
// provider+streaming+model key — nowhere near the hot per-Read path
|
||||
// watchdogBody.Read/Close themselves stay off (armLatency's and Read's
|
||||
// own doc comments, timeout.go, make that constraint explicit): this
|
||||
// lock is paid once per REQUEST, never once per chunk of a stream.
|
||||
func (s *latencyStore) record(key latencyKey, sample latencySample) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
h := s.data[key]
|
||||
if h == nil {
|
||||
h = newLatencyHistogram()
|
||||
s.data[key] = h
|
||||
}
|
||||
if sample.hasTTFB {
|
||||
observeLatencyBucket(h.ttfbBuckets, &h.ttfbOverflow, &h.ttfbSum, &h.ttfbCount, sample.ttfb)
|
||||
}
|
||||
observeLatencyBucket(h.durationBuckets, &h.durationOverflow, &h.durationSum, &h.durationCount, sample.duration)
|
||||
}
|
||||
|
||||
// latencySnapshot is one latencyKey's fully-copied histogram state, safe
|
||||
// to read after latencyStore.snapshot returns without holding its lock.
|
||||
type latencySnapshot struct {
|
||||
key latencyKey
|
||||
ttfbBuckets []int64
|
||||
durationBuckets []int64
|
||||
ttfbSum float64
|
||||
durationSum float64
|
||||
ttfbCount int64
|
||||
ttfbOverflow int64
|
||||
durationCount int64
|
||||
durationOverflow int64
|
||||
}
|
||||
|
||||
// snapshot returns a fully-copied view of every latencyKey s currently
|
||||
// holds, safe to read without s's lock — writeLatencyMetrics and
|
||||
// buildAdminLatencyViews (admin.go) both read through this rather than
|
||||
// s.data directly.
|
||||
func (s *latencyStore) snapshot() []latencySnapshot {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]latencySnapshot, 0, len(s.data))
|
||||
for k, h := range s.data {
|
||||
out = append(out, latencySnapshot{
|
||||
key: k,
|
||||
ttfbBuckets: append([]int64(nil), h.ttfbBuckets...),
|
||||
durationBuckets: append([]int64(nil), h.durationBuckets...),
|
||||
ttfbSum: h.ttfbSum,
|
||||
durationSum: h.durationSum,
|
||||
ttfbCount: h.ttfbCount,
|
||||
ttfbOverflow: h.ttfbOverflow,
|
||||
durationCount: h.durationCount,
|
||||
durationOverflow: h.durationOverflow,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// recordLatency accumulates sample into g.latency: always under
|
||||
// (provider, streaming) — the always-on dimension the task brief
|
||||
// requires — and additionally under (provider, streaming, model) when
|
||||
// both Config.Metrics.ModelLabel is enabled and model is known,
|
||||
// mirroring limiter.recordProviderAttempt's own dual-write rule (limits.
|
||||
// go) for the identical opt-in-cardinality reason.
|
||||
//
|
||||
// This is the ONLY write path into g.latency, and it is reached only
|
||||
// when a latencyRecorder was ever wired onto a request's context in the
|
||||
// first place — every withLatencyRecorder call site (routes_unified.go,
|
||||
// routes_media.go, routes_passthrough.go) gates that wiring on
|
||||
// metricsEnabled(g.cfg) — so a deployment with metrics disabled calls
|
||||
// this exactly zero times, and g.latency's map simply never grows.
|
||||
func (g *Gateway) recordLatency(provider, model string, sample latencySample) {
|
||||
g.latency.record(latencyKey{provider: provider, streaming: sample.streaming}, sample)
|
||||
if model != "" && g.cfg.Metrics != nil && g.cfg.Metrics.ModelLabel {
|
||||
g.latency.record(latencyKey{provider: provider, streaming: sample.streaming, model: model}, sample)
|
||||
}
|
||||
}
|
||||
|
||||
// writeLatencyMetrics emits llmgateway_upstream_ttfb_seconds and
|
||||
// llmgateway_upstream_duration_seconds — this feature's own two
|
||||
// histograms. Absent entirely (no # HELP/TYPE lines even emitted) when
|
||||
// g.latency has recorded nothing yet — a fresh process, or a deployment
|
||||
// with metrics collection gated off at every wiring site — mirroring how
|
||||
// a brand-new histogram with zero observations would look on any real
|
||||
// Prometheus exporter too: no series until the first observation, not a
|
||||
// family of all-zero ones.
|
||||
//
|
||||
// snaps is sorted (provider, then streaming, then model) purely for
|
||||
// deterministic output across repeated scrapes of identical underlying
|
||||
// state — Prometheus does not require sample ordering, but a stable
|
||||
// scrape body makes diffing two scrapes by eye meaningful, matching
|
||||
// writeRejectionMetrics' own sort.Slice above.
|
||||
func (g *Gateway) writeLatencyMetrics(m *metricWriter) {
|
||||
snaps := g.latency.snapshot()
|
||||
if len(snaps) == 0 {
|
||||
return
|
||||
}
|
||||
sort.Slice(snaps, func(i, j int) bool {
|
||||
a, b := snaps[i].key, snaps[j].key
|
||||
if a.provider != b.provider {
|
||||
return a.provider < b.provider
|
||||
}
|
||||
if a.streaming != b.streaming {
|
||||
return !a.streaming // "false" (non-streaming) sorts before "true"
|
||||
}
|
||||
return a.model < b.model
|
||||
})
|
||||
|
||||
m.family("llmgateway_upstream_ttfb_seconds", "histogram",
|
||||
"Time to first successful byte read from an upstream response body, in seconds, measured from just before the request was sent (watchdogBody, timeout.go). Meaningful as a distinct \"how fast did the provider start responding\" signal only for stream=\"true\": a non-streaming provider buffers its whole completion before sending anything, so its own TTFB is approximately equal to its own llmgateway_upstream_duration_seconds and carries the same output-length contamination duration does — always split by the stream label before comparing across requests of different lengths, never averaged across it."+aggregationNoteLatency)
|
||||
m.family("llmgateway_upstream_duration_seconds", "histogram",
|
||||
"Total upstream response body duration, in seconds, from just before the request was sent to the last successful read or Close (watchdogBody, timeout.go) — covers the whole body, not just headers. Roughly TTFB + output_tokens/tokens_per_sec: a long generation legitimately takes longer than a short one on an equally healthy provider, so a raw average across requests of very different output length is not by itself a \"provider degraded\" signal — llmgateway_tokens_total's completion direction (this file) gives the token side of that ratio for the non-streaming population, where duration approximates total generation time."+aggregationNoteLatency)
|
||||
|
||||
for _, s := range snaps {
|
||||
labels := []metricLabel{{"provider", s.key.provider}, {"stream", strconv.FormatBool(s.key.streaming)}}
|
||||
if s.key.model != "" {
|
||||
labels = append(labels, metricLabel{"model", s.key.model})
|
||||
}
|
||||
if s.ttfbCount > 0 {
|
||||
m.histogram("llmgateway_upstream_ttfb_seconds", labels, s.ttfbBuckets, s.ttfbOverflow, s.ttfbSum)
|
||||
}
|
||||
if s.durationCount > 0 {
|
||||
m.histogram("llmgateway_upstream_duration_seconds", labels, s.durationBuckets, s.durationOverflow, s.durationSum)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// writeRejectionMetrics emits llmgateway_rate_limit_rejections_total from
|
||||
// limiter.rejections (limits.go) — the one metric family in this file
|
||||
// with no existing counterStore-backed accounting to read: checkAndCount
|
||||
|
||||
+255
-1
@@ -2,6 +2,9 @@ package traefikllmgateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
@@ -430,13 +433,32 @@ func TestMetrics_OutputParsesAsValidExposition(t *testing.T) {
|
||||
if !validMetricName(s.name) {
|
||||
t.Errorf("metric name %q does not match [a-zA-Z_:][a-zA-Z0-9_:]*", s.name)
|
||||
}
|
||||
kind, ok := types[s.name]
|
||||
// familyName strips a histogram's own _bucket/_sum/_count sample
|
||||
// suffix before the # TYPE lookup: a real Prometheus exposition
|
||||
// document carries exactly ONE "# TYPE <base> histogram" line per
|
||||
// family (metricWriter.family, metrics.go's histogram method),
|
||||
// covering all three suffixed sample names — that is the format's
|
||||
// actual rule (promparse.go's own histogram handling), not a gap
|
||||
// in this test. Every non-histogram family's sample name already
|
||||
// equals its own base name, so this is a no-op for them.
|
||||
familyName := s.name
|
||||
for _, suffix := range []string{"_bucket", "_sum", "_count"} {
|
||||
if trimmed, ok := strings.CutSuffix(s.name, suffix); ok {
|
||||
if k, ok := types[trimmed]; ok && k == "histogram" {
|
||||
familyName = trimmed
|
||||
}
|
||||
}
|
||||
}
|
||||
kind, ok := types[familyName]
|
||||
if !ok {
|
||||
t.Errorf("sample %q has no # TYPE line", s.name)
|
||||
}
|
||||
if kind == "counter" && !strings.HasSuffix(s.name, "_total") {
|
||||
t.Errorf("counter %q does not end in _total", s.name)
|
||||
}
|
||||
if kind == "histogram" && s.name != familyName+"_bucket" && s.name != familyName+"_sum" && s.name != familyName+"_count" {
|
||||
t.Errorf("histogram sample %q does not end in _bucket/_sum/_count", s.name)
|
||||
}
|
||||
|
||||
labelKeys := make([]string, 0, len(s.labels))
|
||||
for k := range s.labels {
|
||||
@@ -469,6 +491,8 @@ func TestMetrics_OutputParsesAsValidExposition(t *testing.T) {
|
||||
"llmgateway_provider_model_failures_total",
|
||||
"llmgateway_rate_limit_rejections_total",
|
||||
"llmgateway_limit_store_up",
|
||||
"llmgateway_upstream_ttfb_seconds",
|
||||
"llmgateway_upstream_duration_seconds",
|
||||
}
|
||||
for _, name := range wantFamilies {
|
||||
if _, ok := types[name]; !ok {
|
||||
@@ -1121,3 +1145,233 @@ func TestMetrics_BudgetConsumedRatio_ReflectsUsageVsLimit(t *testing.T) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- upstream latency (feat: instrument upstream latency) ---
|
||||
|
||||
// TestMetrics_LatencyHistogram_BucketsCumulativeOrdered_PlusInfPresent_SumCountConsistent
|
||||
// is the exposition-format correctness test the task brief calls out by
|
||||
// name: a malformed histogram does not drop one metric, it fails the
|
||||
// whole scrape and the target reports DOWN — so this parses the real
|
||||
// rendered document (not a golden string) and checks the format's own
|
||||
// required shape by hand: buckets present for every configured bound
|
||||
// plus +Inf, strictly ascending `le`, cumulative (non-decreasing)
|
||||
// counts, +Inf last, and `_count` exactly equal to the final cumulative
|
||||
// bucket.
|
||||
//
|
||||
// Observations are seeded directly via gw.recordLatency (bypassing
|
||||
// watchdogBody's own real-clock capture, already covered by
|
||||
// latency_test.go's TestWatchdogBody_ArmLatency_TTFBAtFirstRead_
|
||||
// NotConstructionOrClose) so each one lands in a KNOWN bucket
|
||||
// deterministically, including one past every configured bound (the
|
||||
// overflow/+Inf case) — a real network delay could not reliably hit
|
||||
// that case without an unreasonably slow test.
|
||||
//
|
||||
// MUTATION VERIFIED: changing metricWriter.histogram (metrics.go) to
|
||||
// emit the DENSE per-bucket count directly (dropping the running
|
||||
// `cumulative +=` accumulation, i.e. `m.sampleInt(name+"_bucket",
|
||||
// bucketLabels, buckets[i])` instead of `cumulative`) made this test
|
||||
// fail: consecutive bucket values were no longer non-decreasing (a
|
||||
// dense count can drop back to 0 after a nonzero one), failing the
|
||||
// "must be cumulative" assertion below. Reverted before committing.
|
||||
func TestMetrics_LatencyHistogram_BucketsCumulativeOrdered_PlusInfPresent_SumCountConsistent(t *testing.T) {
|
||||
t.Parallel()
|
||||
cfg := newMetricsTestConfig()
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
durations := []time.Duration{
|
||||
50 * time.Millisecond, // <= 0.1 bucket
|
||||
400 * time.Millisecond, // <= 0.5 bucket
|
||||
3 * time.Second, // <= 5 bucket
|
||||
20 * time.Second, // <= 30 bucket
|
||||
1000 * time.Second, // past every bound: overflow / +Inf
|
||||
}
|
||||
for _, d := range durations {
|
||||
gw.recordLatency("openai", "", latencySample{duration: d, ttfb: d / 10, hasTTFB: true, streaming: false})
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
|
||||
type bucketLine struct {
|
||||
le float64
|
||||
value int64
|
||||
}
|
||||
var buckets []bucketLine
|
||||
var sumVal, countVal string
|
||||
for _, s := range samples {
|
||||
if s.labels["provider"] != "openai" || s.labels["stream"] != "false" {
|
||||
continue
|
||||
}
|
||||
switch s.name {
|
||||
case "llmgateway_upstream_duration_seconds_bucket":
|
||||
le := math.Inf(1)
|
||||
if s.labels["le"] != "+Inf" {
|
||||
var err error
|
||||
le, err = strconv.ParseFloat(s.labels["le"], 64)
|
||||
if err != nil {
|
||||
t.Fatalf("le=%q does not parse: %v", s.labels["le"], err)
|
||||
}
|
||||
}
|
||||
v, err := strconv.ParseInt(s.value, 10, 64)
|
||||
if err != nil {
|
||||
t.Fatalf("bucket value %q does not parse: %v", s.value, err)
|
||||
}
|
||||
buckets = append(buckets, bucketLine{le: le, value: v})
|
||||
case "llmgateway_upstream_duration_seconds_sum":
|
||||
sumVal = s.value
|
||||
case "llmgateway_upstream_duration_seconds_count":
|
||||
countVal = s.value
|
||||
}
|
||||
}
|
||||
|
||||
if len(buckets) != len(latencyBucketBounds)+1 {
|
||||
t.Fatalf("got %d bucket lines, want %d (%d configured bounds + Inf)", len(buckets), len(latencyBucketBounds)+1, len(latencyBucketBounds))
|
||||
}
|
||||
for i := 1; i < len(buckets); i++ {
|
||||
if buckets[i].le <= buckets[i-1].le {
|
||||
t.Errorf("bucket %d le=%v is not strictly greater than bucket %d le=%v — buckets must be strictly ascending", i, buckets[i].le, i-1, buckets[i-1].le)
|
||||
}
|
||||
if buckets[i].value < buckets[i-1].value {
|
||||
t.Errorf("bucket %d value=%d < bucket %d value=%d — buckets must be cumulative (non-decreasing)", i, buckets[i].value, i-1, buckets[i-1].value)
|
||||
}
|
||||
}
|
||||
if !math.IsInf(buckets[len(buckets)-1].le, 1) {
|
||||
t.Errorf("last bucket le=%v, want +Inf", buckets[len(buckets)-1].le)
|
||||
}
|
||||
if sumVal == "" {
|
||||
t.Fatal("no _sum sample found for this label set")
|
||||
}
|
||||
if countVal == "" {
|
||||
t.Fatal("no _count sample found for this label set")
|
||||
}
|
||||
wantCount := int64(len(durations))
|
||||
if buckets[len(buckets)-1].value != wantCount {
|
||||
t.Errorf("final cumulative (+Inf) bucket = %d, want %d (one per seeded observation, including the overflow one)", buckets[len(buckets)-1].value, wantCount)
|
||||
}
|
||||
if countVal != strconv.FormatInt(buckets[len(buckets)-1].value, 10) {
|
||||
t.Errorf("_count = %s, want %d — must exactly equal the final +Inf cumulative bucket, never drift from it", countVal, buckets[len(buckets)-1].value)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMetrics_LatencyHistogram_StreamingAndNonStreaming_DistinctLabelSets
|
||||
// drives one streaming and one non-streaming request through the real
|
||||
// unified route (not a direct seed) and asserts both land in distinct
|
||||
// stream="true"/stream="false" series for the same provider — proving
|
||||
// the end-to-end wiring (isEventStreamResponse's derivation at
|
||||
// newWatchdogBody's call sites, through to recordLatency's dual key)
|
||||
// actually separates them, not just that the renderer could in
|
||||
// principle.
|
||||
//
|
||||
// MUTATION VERIFIED: hardcoding isEventStreamResponse (timeout.go) to
|
||||
// always return false made this test fail — no
|
||||
// llmgateway_upstream_duration_seconds_count{...,stream="true"} sample
|
||||
// was ever emitted, since every response (including the real SSE one)
|
||||
// was misclassified as non-streaming. Reverted before committing.
|
||||
func TestMetrics_LatencyHistogram_StreamingAndNonStreaming_DistinctLabelSets(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
if streaming, _ := body["stream"].(bool); streaming {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
fl := w.(http.Flusher)
|
||||
_, _ = fmt.Fprint(w, "data: {\"id\":\"c\",\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n")
|
||||
fl.Flush()
|
||||
_, _ = fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
fl.Flush()
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Providers["openai"].BaseURL = srv.URL
|
||||
h, _ := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
nonStreamBody := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
streamBody := map[string]any{"model": "gpt-test", "stream": true, "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
|
||||
rec1 := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec1, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", nonStreamBody))
|
||||
if rec1.Code != http.StatusOK {
|
||||
t.Fatalf("non-streaming request status = %d, want 200, body=%s", rec1.Code, rec1.Body.String())
|
||||
}
|
||||
rec2 := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec2, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", streamBody))
|
||||
if rec2.Code != http.StatusOK {
|
||||
t.Fatalf("streaming request status = %d, want 200, body=%s", rec2.Code, rec2.Body.String())
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, metricsRequest(http.MethodGet, metricsPathDefault, "sk-admin1"))
|
||||
_, samples := parsePrometheusText(t, rec.Body.Bytes())
|
||||
|
||||
var sawStreamTrue, sawStreamFalse bool
|
||||
for _, s := range samples {
|
||||
if s.name != "llmgateway_upstream_duration_seconds_count" || s.labels["provider"] != "openai" {
|
||||
continue
|
||||
}
|
||||
switch s.labels["stream"] {
|
||||
case "true":
|
||||
sawStreamTrue = true
|
||||
case "false":
|
||||
sawStreamFalse = true
|
||||
}
|
||||
}
|
||||
if !sawStreamTrue {
|
||||
t.Error(`no llmgateway_upstream_duration_seconds_count{provider="openai",stream="true"} sample — the streaming request must land in the stream="true" label set`)
|
||||
}
|
||||
if !sawStreamFalse {
|
||||
t.Error(`no llmgateway_upstream_duration_seconds_count{provider="openai",stream="false"} sample — the non-streaming request must land in the stream="false" label set`)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecordLatency_MetricsDisabled_NoObservationCollected proves the
|
||||
// "no collection cost, no series emitted" requirement at the
|
||||
// Gateway/collection level (not merely "the renderer skips an empty
|
||||
// store", which would be true even if collection were silently still
|
||||
// happening): with Config.Metrics.Enabled false, driving a real,
|
||||
// successful request through the unified route must leave g.latency
|
||||
// completely empty — every withLatencyRecorder call site
|
||||
// (routes_unified.go, routes_media.go, routes_passthrough.go) gates its
|
||||
// own wiring on metricsEnabled(g.cfg), so no latencyRecorder is ever
|
||||
// looked up from context in the first place, and armLatency is never
|
||||
// called at all.
|
||||
//
|
||||
// MUTATION VERIFIED: removing the `if metricsEnabled(g.cfg)` guard
|
||||
// around runUnified's own withLatencyRecorder wiring (routes_unified.go)
|
||||
// made this test fail — gw.latency.snapshot() returned 1 entry instead
|
||||
// of 0. Reverted before committing.
|
||||
func TestRecordLatency_MetricsDisabled_NoObservationCollected(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(`{"id":"c1","object":"chat.completion","model":"gpt-test","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
cfg := newMetricsTestConfig()
|
||||
cfg.Metrics.Enabled = false // metrics collection disabled, but Metrics itself stays non-nil
|
||||
cfg.Providers["openai"].BaseURL = srv.URL
|
||||
h, gw := newMetricsGatewayHandle(t, cfg)
|
||||
|
||||
body := map[string]any{"model": "gpt-test", "messages": []any{map[string]any{"role": "user", "content": "hi"}}}
|
||||
rec := httptest.NewRecorder()
|
||||
h.ServeHTTP(rec, newUnifiedRequest(t, http.MethodPost, "/v1/chat/completions", "sk-alice", body))
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200, body=%s", rec.Code, rec.Body.String())
|
||||
}
|
||||
|
||||
if snaps := gw.latency.snapshot(); len(snaps) != 0 {
|
||||
t.Errorf("g.latency has %d entries after a request with metrics disabled, want 0 — no latency observation should ever be collected when metrics is off", len(snaps))
|
||||
}
|
||||
}
|
||||
|
||||
+93
-1
@@ -249,6 +249,73 @@ func attemptRecorderFromContext(ctx context.Context) attemptRecorder {
|
||||
return rec
|
||||
}
|
||||
|
||||
// latencyRecorder is invoked exactly once per completed upstream body —
|
||||
// from watchdogBody.Close (timeout.go), never per Read — with the
|
||||
// latencySample (metrics.go) that body's own watchdog measured: time to
|
||||
// first byte, total duration, and whether the response was a stream
|
||||
// (feat: instrument upstream latency, observation-only per the task
|
||||
// brief — no throttling or control logic reads this).
|
||||
//
|
||||
// Threaded via context exactly like attemptRecorder above, for the same
|
||||
// reason: upstreamBytes/proxyUpstream have no other way to reach
|
||||
// Gateway-level state (g.latency, metrics.go). A call site that never
|
||||
// wires one via withLatencyRecorder gets nil back from
|
||||
// latencyRecorderFromContext, so registry.go's discovery listModels
|
||||
// calls and the MCP/A2A target proxy (handleTargetProxy, mcp_a2a.go —
|
||||
// neither ever wraps its context this way, mirroring
|
||||
// withAttemptRecorder's own identical exclusion) never observe latency
|
||||
// either — this feature accounts the same traffic Feature A does, not a
|
||||
// byte more.
|
||||
type latencyRecorder func(sample latencySample)
|
||||
|
||||
// latencyRecorderCtxKey is the unexported context.Value key
|
||||
// withLatencyRecorder/latencyRecorderFromContext share — its own type,
|
||||
// never a string, so no other package's context.WithValue call can
|
||||
// collide with it, mirroring attemptRecorderCtxKey above.
|
||||
//
|
||||
// Deliberately NOT a bare `struct{}` like attemptRecorderCtxKey (yaegi-
|
||||
// check finding, this feature): under the REAL Yaegi interpreter, two
|
||||
// DISTINCT zero-sized struct types used as `any`-typed context.WithValue
|
||||
// keys compare EQUAL to each other via `==` — an interpreter-only
|
||||
// divergence from compiled Go's own interface-equality semantics, where
|
||||
// two values of different concrete types are never equal regardless of
|
||||
// value. With attemptRecorderCtxKey ALSO a bare struct{}, a context
|
||||
// carrying both (routes_unified.go's withAttemptRecorder then
|
||||
// withLatencyRecorder) let ctx.Value(attemptRecorderCtxKey{}) wrongly
|
||||
// match the LATER latencyRecorderCtxKey{} entry instead, silently
|
||||
// returning nil once type-asserted to attemptRecorder — retryPolicy.do
|
||||
// (retry.go) then never called rec(resp, err) at all, and
|
||||
// limiter.recordProviderAttempt (Feature A) went permanently dark for
|
||||
// every metered request, with no error anywhere: exerciseAttemptAccounting
|
||||
// (tools/yaegi-check/main.go) is what caught it — a compiled `go test`
|
||||
// run never sees this divergence, since compiled Go's interface equality
|
||||
// is correct. The `_ byte` field costs nothing at runtime and makes this
|
||||
// type's size differ from attemptRecorderCtxKey's, which is enough to
|
||||
// keep the two apart under the interpreter too.
|
||||
type latencyRecorderCtxKey struct{ _ byte }
|
||||
|
||||
// withLatencyRecorder returns a context carrying rec. A nil rec (the
|
||||
// common case: metrics.go's own callers only ever pass a non-nil one
|
||||
// when metricsEnabled(g.cfg) was true at the point of the call) returns
|
||||
// ctx unchanged, so a caller may call this unconditionally, exactly
|
||||
// mirroring withAttemptRecorder's own nil handling above.
|
||||
func withLatencyRecorder(ctx context.Context, rec latencyRecorder) context.Context {
|
||||
if rec == nil {
|
||||
return ctx
|
||||
}
|
||||
return context.WithValue(ctx, latencyRecorderCtxKey{}, rec)
|
||||
}
|
||||
|
||||
// latencyRecorderFromContext returns the latencyRecorder ctx carries via
|
||||
// withLatencyRecorder, or nil when none was set — the common case for
|
||||
// any call path this feature does not account, or for any deployment
|
||||
// with metrics collection disabled (see latencyRecorder's own doc
|
||||
// comment).
|
||||
func latencyRecorderFromContext(ctx context.Context) latencyRecorder {
|
||||
rec, _ := ctx.Value(latencyRecorderCtxKey{}).(latencyRecorder)
|
||||
return rec
|
||||
}
|
||||
|
||||
// upstreamJSON issues an HTTP request to url: body, when non-nil, is
|
||||
// marshaled as the JSON request body; hdr's values are added to the
|
||||
// request (a caller builds this from its adapter's auth and content-type
|
||||
@@ -320,6 +387,15 @@ func upstreamBytes(ctx context.Context, client *http.Client, method, url string,
|
||||
// provider-health accounting, since nothing else ever called rec
|
||||
// again for that attempt.
|
||||
rec := attemptRecorderFromContext(ctx)
|
||||
// latRec: same one-lookup-outside-the-retry-loop shape as rec above,
|
||||
// and the same reasoning applies — every attempt shares ctx, so a
|
||||
// latencyRecorder wired via withLatencyRecorder (nil when metrics
|
||||
// collection is disabled, latencyRecorder's own doc comment) is
|
||||
// identical across attempts. A retried request still gets its own
|
||||
// independent latency observation per attempt (below, start is
|
||||
// captured fresh inside the closure), matching how watchdogBody
|
||||
// itself is rebuilt fresh per attempt.
|
||||
latRec := latencyRecorderFromContext(ctx)
|
||||
|
||||
call := func() (*http.Response, error) {
|
||||
var r io.Reader
|
||||
@@ -361,12 +437,28 @@ func upstreamBytes(ctx context.Context, client *http.Client, method, url string,
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
|
||||
// start: captured just before the request is actually sent, so a
|
||||
// non-nil latRec's eventual sample (armLatency below) measures from
|
||||
// here — matching the task brief's own "from just before the
|
||||
// upstream request is sent" definition, not from whatever moment
|
||||
// newWatchdogBody happens to run at (which is already after
|
||||
// headers arrived).
|
||||
start := time.Now()
|
||||
resp, err := client.Do(req) //nolint:bodyclose,gosec // caller closes resp.Body; upstreamBytes hands the response, not its lifecycle, back — same operator-configured URL as above
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("%w: %w", errUpstream, err)
|
||||
}
|
||||
resp.Body = newWatchdogBody(resp.Body, cancel, timeout, fmt.Sprintf("provider %q", providerName), rec)
|
||||
wb := newWatchdogBody(resp.Body, cancel, timeout, fmt.Sprintf("provider %q", providerName), rec)
|
||||
if latRec != nil {
|
||||
// streaming: derived from resp's own Content-Type, the SAME
|
||||
// convention each adapter's forwardStream/forwardJSON split
|
||||
// already applies (isEventStreamResponse, timeout.go) — decided
|
||||
// here because it is already knowable, ahead of whichever of
|
||||
// those two the caller picks moments from now.
|
||||
wb.armLatency(start, isEventStreamResponse(resp), latRec)
|
||||
}
|
||||
resp.Body = wb
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -119,6 +119,13 @@ func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request
|
||||
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
|
||||
})
|
||||
// feat: instrument upstream latency — gated on metricsEnabled, same
|
||||
// reasoning as runUnified's identical wiring (routes_unified.go).
|
||||
if metricsEnabled(g.cfg) {
|
||||
ctx = withLatencyRecorder(ctx, func(sample latencySample) {
|
||||
g.recordLatency(adapter.name(), upstreamModel, sample)
|
||||
})
|
||||
}
|
||||
if _, err := adapter.imagesGeneration(ctx, sw, req); err != nil {
|
||||
g.handleAdapterError(sw, err, adapter.name())
|
||||
}
|
||||
@@ -177,6 +184,13 @@ func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *u
|
||||
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
|
||||
})
|
||||
// feat: instrument upstream latency — gated on metricsEnabled, same
|
||||
// reasoning as runUnified's identical wiring (routes_unified.go).
|
||||
if metricsEnabled(g.cfg) {
|
||||
ctx = withLatencyRecorder(ctx, func(sample latencySample) {
|
||||
g.recordLatency(adapter.name(), upstreamModel, sample)
|
||||
})
|
||||
}
|
||||
if _, err := adapter.audioSpeech(ctx, sw, body, "application/json"); err != nil {
|
||||
g.handleAdapterError(sw, err, adapter.name())
|
||||
}
|
||||
@@ -274,6 +288,13 @@ func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Reque
|
||||
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr)
|
||||
})
|
||||
// feat: instrument upstream latency — gated on metricsEnabled, same
|
||||
// reasoning as runUnified's identical wiring (routes_unified.go).
|
||||
if metricsEnabled(g.cfg) {
|
||||
ctx = withLatencyRecorder(ctx, func(sample latencySample) {
|
||||
g.recordLatency(adapter.name(), upstreamModel, sample)
|
||||
})
|
||||
}
|
||||
if _, err := adapter.audioTranscription(ctx, sw, uploadBody, uploadContentType); err != nil {
|
||||
g.handleAdapterError(sw, err, adapter.name())
|
||||
}
|
||||
|
||||
+28
-1
@@ -862,6 +862,16 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u
|
||||
r = r.WithContext(withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(providerName, "", resp, attemptErr)
|
||||
}))
|
||||
// feat: instrument upstream latency — gated on metricsEnabled, same
|
||||
// reasoning as runUnified's identical wiring (routes_unified.go). No
|
||||
// model here either, for the identical reason the attemptRecorder
|
||||
// above passes "": the upstream model lives in the response body,
|
||||
// read only after this attempt already resolved (comment above).
|
||||
if metricsEnabled(g.cfg) {
|
||||
r = r.WithContext(withLatencyRecorder(r.Context(), func(sample latencySample) {
|
||||
g.recordLatency(providerName, "", sample)
|
||||
}))
|
||||
}
|
||||
|
||||
result, ok := g.proxyUpstream(w, r, upstreamURL, adapter.httpClient(), adapter.injectAuth, providerCredentialRetargetHeaders, true, "passthrough (provider "+providerName+")", adapter.requestTimeout())
|
||||
if !ok || !result.isJSON {
|
||||
@@ -1007,6 +1017,12 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
|
||||
injectAuth(upstreamReq)
|
||||
}
|
||||
|
||||
// start: captured just before the request is actually sent, matching
|
||||
// upstreamBytes' own identical capture (providers.go) — see
|
||||
// watchdogBody.armLatency's doc comment (timeout.go) for why this,
|
||||
// not whatever moment newWatchdogBody itself runs at, is what "just
|
||||
// before the upstream request is sent" (task brief) means here.
|
||||
start := time.Now()
|
||||
resp, err := client.Do(upstreamReq) //nolint:gosec // same upstreamReq built above; operator-fixed host, traversal-checked, see its construction comment
|
||||
// Feature A (v0.22): proxyUpstream makes exactly one attempt (no
|
||||
// retry.go policy wraps this path), so this fires once per call,
|
||||
@@ -1020,6 +1036,13 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
|
||||
// accounting a build/send failure already does (coordinator
|
||||
// adversarial review, 2026-08-23, finding F4).
|
||||
rec := attemptRecorderFromContext(r.Context())
|
||||
// latRec: nil for the MCP/A2A target proxy (handleTargetProxy,
|
||||
// mcp_a2a.go, never wraps r's context this way — the SAME exclusion
|
||||
// rec's own doc comment above already documents for attemptRecorder),
|
||||
// non-nil for native passthrough whenever metrics collection is
|
||||
// enabled (handlePassthrough's own withLatencyRecorder wiring, above
|
||||
// in this file).
|
||||
latRec := latencyRecorderFromContext(r.Context())
|
||||
if rec != nil {
|
||||
rec(resp, err)
|
||||
}
|
||||
@@ -1037,7 +1060,11 @@ func (g *Gateway) proxyUpstream(w http.ResponseWriter, r *http.Request, upstream
|
||||
// calls here (mcp_a2a.go) pass "mcp target (name ...)"/"a2a target
|
||||
// (name ...)" — an MCP/A2A target is not a provider, and the watchdog
|
||||
// error text must not claim it is (finding F9).
|
||||
resp.Body = newWatchdogBody(resp.Body, cancel, timeout, logPrefix, rec)
|
||||
wb := newWatchdogBody(resp.Body, cancel, timeout, logPrefix, rec)
|
||||
if latRec != nil {
|
||||
wb.armLatency(start, isEventStreamResponse(resp), latRec)
|
||||
}
|
||||
resp.Body = wb
|
||||
defer resp.Body.Close() //nolint:errcheck // read-side close; nothing actionable on failure
|
||||
|
||||
copyHeadersExcept(w.Header(), resp.Header, hopByHopHeaders, dangerousResponseHeaders)
|
||||
|
||||
@@ -494,6 +494,15 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop
|
||||
ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) {
|
||||
g.limiter.recordProviderAttempt(providerName, upstreamModel, resp, attemptErr)
|
||||
})
|
||||
// feat: instrument upstream latency — gated on metricsEnabled so a
|
||||
// deployment with metrics off wires nothing here at all (not even
|
||||
// the closure allocation), one level up from watchdogBody's own
|
||||
// per-Read nil check (timeout.go).
|
||||
if metricsEnabled(g.cfg) {
|
||||
ctx = withLatencyRecorder(ctx, func(sample latencySample) {
|
||||
g.recordLatency(providerName, upstreamModel, sample)
|
||||
})
|
||||
}
|
||||
result, callErr := call(cand.adapter, ctx, respWriter, attemptReq)
|
||||
|
||||
// Request-path health, recorded ONCE per candidate per logical
|
||||
|
||||
+128
-7
@@ -5,6 +5,8 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
@@ -148,14 +150,37 @@ func resolveProviderTimeout(providerName, providerRaw string, globalTimeout time
|
||||
// timedOut is accessed via sync/atomic's function API (atomic.StoreInt32/
|
||||
// LoadInt32), not the newer atomic.Bool type, per this plugin's
|
||||
// yaegi-interpreted-stdlib guardrail.
|
||||
//
|
||||
// onLatency/latencyStart/ttfb/streaming/ttfbRecorded/latencyRecorded
|
||||
// (feat: instrument upstream latency, task brief) are this same
|
||||
// watchdogBody's SECOND, independent observation: onLatency stays nil
|
||||
// unless armLatency (below) is called with a non-nil recorder, and every
|
||||
// access to the group below is gated on that one nil check first — a
|
||||
// disabled deployment (or any caller that never observes latency at
|
||||
// all, e.g. registry.go's discovery calls) pays one pointer comparison
|
||||
// per Read and nothing else. See armLatency's own doc comment for why
|
||||
// this is a separate method rather than additional newWatchdogBody
|
||||
// parameters, and Read/Close below for where each field is written.
|
||||
type watchdogBody struct {
|
||||
rc io.ReadCloser
|
||||
cancel context.CancelFunc
|
||||
timer *time.Timer
|
||||
onTimeout attemptRecorder
|
||||
label string
|
||||
timeout time.Duration
|
||||
timedOut int32
|
||||
// latencyStart is declared first (fieldalignment, govet — enabled via
|
||||
// this repo's global golangci-lint config): time.Time embeds a
|
||||
// *time.Location, making it pointer-shaped, and every pointer-shaped
|
||||
// field in this struct is grouped at the front, then plain scalars
|
||||
// (time.Duration/int32/bool) last — see adminProviderView's identical
|
||||
// convention (admin.go) for the same reasoning spelled out once.
|
||||
latencyStart time.Time
|
||||
rc io.ReadCloser
|
||||
cancel context.CancelFunc
|
||||
timer *time.Timer
|
||||
onTimeout attemptRecorder
|
||||
onLatency latencyRecorder
|
||||
label string
|
||||
timeout time.Duration
|
||||
ttfb time.Duration
|
||||
timedOut int32
|
||||
ttfbRecorded int32
|
||||
latencyRecorded int32
|
||||
streaming bool
|
||||
}
|
||||
|
||||
// newWatchdogBody returns a watchdogBody wrapping rc, arming its watchdog
|
||||
@@ -181,6 +206,54 @@ func newWatchdogBody(rc io.ReadCloser, cancel context.CancelFunc, timeout time.D
|
||||
return wb
|
||||
}
|
||||
|
||||
// armLatency wires wb to observe upstream latency (feat: instrument
|
||||
// upstream latency, task brief): rec is invoked exactly once, from
|
||||
// Close below, with the latencySample (metrics.go) this watchdogBody
|
||||
// measured. rec == nil disables all of it — Read/Close each take a
|
||||
// single, cost-free nil check on wb.onLatency and do nothing further,
|
||||
// so calling armLatency at all (even with rec nil) adds no measurable
|
||||
// cost, and a caller may call it unconditionally.
|
||||
//
|
||||
// This is a separate method, not additional newWatchdogBody parameters:
|
||||
// newWatchdogBody's existing signature and behavior are pinned by
|
||||
// timeout_test.go's direct, positional 5-argument calls, which must
|
||||
// keep passing unchanged (task constraint) — changing that signature
|
||||
// would have broken every one of them for no benefit, since every real
|
||||
// call site (providers.go, routes_passthrough.go) already holds a
|
||||
// *watchdogBody reference to call this on immediately afterward.
|
||||
//
|
||||
// start is the moment just BEFORE the request was sent (the caller's
|
||||
// own time.Now(), captured ahead of client.Do) — deliberately earlier
|
||||
// than "now" here: newWatchdogBody itself only ever runs once the
|
||||
// response's headers have already arrived, so using the caller's
|
||||
// earlier timestamp is what makes both TTFB and total duration include
|
||||
// the network round trip for headers too, matching the task brief's own
|
||||
// definition ("from just before the upstream request is sent").
|
||||
//
|
||||
// streaming labels the eventual sample: the caller derives it from the
|
||||
// response's own Content-Type via isEventStreamResponse (below) —
|
||||
// deliberately not from the client's requested "stream" flag alone,
|
||||
// since a provider that ignores it and answers with one buffered JSON
|
||||
// body is, for latency purposes, indistinguishable from a genuinely
|
||||
// non-streaming call: its TTFB IS its own total duration either way.
|
||||
func (wb *watchdogBody) armLatency(start time.Time, streaming bool, rec latencyRecorder) {
|
||||
wb.onLatency = rec
|
||||
wb.latencyStart = start
|
||||
wb.streaming = streaming
|
||||
}
|
||||
|
||||
// isEventStreamResponse reports whether resp's Content-Type indicates a
|
||||
// text/event-stream body — the SAME convention each provider adapter's
|
||||
// own forwardStream/forwardJSON split already applies (provider_openai.
|
||||
// go, provider_anthropic.go, provider_gemini.go), duplicated here rather
|
||||
// than imported from any one adapter file: newWatchdogBody has two call
|
||||
// sites (providers.go, routes_passthrough.go), and routes_passthrough.go
|
||||
// (native passthrough, the MCP/A2A target proxy) has no adapter-specific
|
||||
// streaming decision of its own to share this with.
|
||||
func isEventStreamResponse(resp *http.Response) bool {
|
||||
return strings.Contains(strings.ToLower(resp.Header.Get("Content-Type")), "text/event-stream")
|
||||
}
|
||||
|
||||
// fire runs once, in its own goroutine (time.AfterFunc's own contract —
|
||||
// no goroutine is ever spawned for a request whose timer is Stopped
|
||||
// before it fires, which is the common, non-timed-out case), the moment
|
||||
@@ -227,8 +300,32 @@ func (wb *watchdogBody) timeoutError() error {
|
||||
// actually produced. Any other error — a real client disconnect, a clean
|
||||
// EOF, a genuine network failure that has nothing to do with this
|
||||
// watchdog — passes through unchanged.
|
||||
//
|
||||
// The TTFB-capture block below (feat: instrument upstream latency) is
|
||||
// purely additive and never changes what this method returns or the
|
||||
// timer-reset logic's own behavior — it is checked FIRST, ahead of the
|
||||
// `if err != nil` branch, and deliberately does NOT require err == nil
|
||||
// the way the timer-reset block below still does: io.Reader's contract
|
||||
// explicitly permits a reader to return its final data together with
|
||||
// io.EOF in the SAME call (io.Reader's own doc: "an instance ... may
|
||||
// return either err == EOF or err == nil"), and Go's real net/http body
|
||||
// reader does exactly that for a small, fully-buffered response —
|
||||
// measured directly against a real httptest server here, not assumed:
|
||||
// a ~200-byte JSON completion arrives as ONE Read call reporting
|
||||
// (n=200, err=io.EOF) together. Gating the capture on err == nil (as
|
||||
// timer-reset still correctly does, per finding F6 below) would then
|
||||
// skip it for every such response — the common case for a typical,
|
||||
// fully-buffered non-streaming completion — leaving hasTTFB false for
|
||||
// most real traffic. Checking `n > 0` alone, before the error branch,
|
||||
// fixes that without touching the timer-reset semantics those tests
|
||||
// pin: a self-contained top-of-function check that only ever WRITES
|
||||
// wb.ttfb/wb.ttfbRecorded, never wb.timer or any error path.
|
||||
func (wb *watchdogBody) Read(p []byte) (int, error) {
|
||||
n, err := wb.rc.Read(p)
|
||||
if n > 0 && wb.onLatency != nil && atomic.LoadInt32(&wb.ttfbRecorded) == 0 {
|
||||
wb.ttfb = time.Since(wb.latencyStart)
|
||||
atomic.StoreInt32(&wb.ttfbRecorded, 1)
|
||||
}
|
||||
if err != nil {
|
||||
if atomic.LoadInt32(&wb.timedOut) == 1 {
|
||||
return n, wb.timeoutError()
|
||||
@@ -248,8 +345,32 @@ func (wb *watchdogBody) Read(p []byte) (int, error) {
|
||||
// gave up early, or because the watchdog itself already fired — instead
|
||||
// of being left to be reclaimed only once its parent context ends. Then
|
||||
// closes the wrapped body.
|
||||
//
|
||||
// The `if wb.onLatency != nil` block below (feat: instrument upstream
|
||||
// latency) is purely additive, runs after the pre-existing timer.Stop/
|
||||
// cancel/rc.Close sequence above is already decided, and never changes
|
||||
// this method's return value: it reports exactly one latencySample
|
||||
// (metrics.go) — duration always (time.Since(wb.latencyStart), this
|
||||
// body's whole lifetime), ttfb only when a successful read actually
|
||||
// happened first (hasTTFB, guarded by ttfbRecorded — Read's own doc
|
||||
// comment above). latencyRecorded is a CompareAndSwap, not a plain
|
||||
// check, so a caller that closes this watchdogBody more than once (not
|
||||
// expected by any current call site, but not this method's job to
|
||||
// forbid either) still reports the sample exactly once rather than
|
||||
// double-counting it.
|
||||
func (wb *watchdogBody) Close() error {
|
||||
wb.timer.Stop()
|
||||
wb.cancel()
|
||||
if wb.onLatency != nil && atomic.CompareAndSwapInt32(&wb.latencyRecorded, 0, 1) {
|
||||
sample := latencySample{
|
||||
duration: time.Since(wb.latencyStart),
|
||||
streaming: wb.streaming,
|
||||
}
|
||||
if atomic.LoadInt32(&wb.ttfbRecorded) == 1 {
|
||||
sample.ttfb = wb.ttfb
|
||||
sample.hasTTFB = true
|
||||
}
|
||||
wb.onLatency(sample)
|
||||
}
|
||||
return wb.rc.Close()
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user