diff --git a/README.md b/README.md index ea494b4..b4f2995 100644 --- a/README.md +++ b/README.md @@ -383,7 +383,8 @@ config accepts them as YAML, which decodes to the same JSON shape. | `requireModelPrefix` | `bool` | `false` | When `true`, a model resolves only as `provider/model` or as a `modelAliases` id. A bare id returns 404 `unknown model` and `GET /v1/models` lists no bare ids. `false` keeps the sorted-first bare-id collision winner. | | `stripClientAnthropicBeta` | `bool` | `false` | When `true`, `POST /v1/messages` does not forward the client's `anthropic-beta` header to the provider, as native passthrough already does. `false` forwards it, so a client can enable beta features on the operator's own Anthropic account. | | `modelMeta` | `map[string]*ModelMetaConfig` | `{}` | Per-model context-window and cost overrides, keyed by exact `provider/model` or a bare model id — see [Model metadata](#model-metadata). Omitted or empty means no overrides; every model's metadata still resolves through discovery and the built-in table. | -| `maxConcurrency` | `int` | `0` | When above `0`, every response carries `X-Concurrency-Limit: `, so clients that honour it (omps) keep at most `n` requests in flight. Set it to what your upstreams serve at once. The gateway only advertises the value; it does not queue or reject on it. `0` sends no header. Negative values are rejected at startup. | +| `maxConcurrency` | `int` | `0` | When above `0`, every response carries `RateLimit-Policy: "concurrency";q=;qu="concurrent-requests"` (IETF draft `draft-ietf-httpapi-ratelimit-headers`, an Internet-Draft, not a final standard), so clients that read it can keep at most `n` requests in flight. Set it to what your upstreams serve at once. The gateway only advertises the value; it does not queue or reject on it. It always wins over a learned value. `0` sends no header. Negative values are rejected at startup. | +| `concurrencyAutodetect` | `bool` | `false` | When `true`, the gateway learns per provider how many requests the upstream serves at once. A 429, 503, 504 or upstream timeout sets the learned value to half the in-flight requests (minimum 1) and never raises it. Every 4 consecutive successes add 1. At 64 the value counts as unknown again. Responses served via that provider advertise the smaller of `maxConcurrency` and the learned value in `RateLimit-Policy`; with neither known, no header is sent. Learned values live in memory and reset on restart. | | `passthroughUnknown` | `bool` | `false` | `false`: a request matching none of the plugin's routes gets a 404 JSON envelope. `true`: it falls through to the router's own backing service. **Does not cover the three media routes** — see the upgrade note below. | | `maxRequestBytes` | `int` | `0` (4 MiB) | Raises the unified chat/embeddings request-body cap. `0` or negative keeps 4 MiB; a positive value is clamped to 64 MiB. Pair with `maxInFlightBodyRequests` and the pod memory limit: an adversarial body decodes to ~12x live heap. | | `maxInFlightBodyRequests` | `int` | `0` (self-tuned) | Caps how many unified/media requests may concurrently be inside their own request-body read+decode step — see [Body-admission limit](#body-admission-limit). `0` self-tunes from the host's own `GOMAXPROCS`; a positive value pins an exact cap instead, clamped to a maximum of 10,000. | diff --git a/concurrency.go b/concurrency.go new file mode 100644 index 0000000..d6c5090 --- /dev/null +++ b/concurrency.go @@ -0,0 +1,165 @@ +package traefikllmgateway + +import ( + "context" + "net/http" + "strconv" + "sync" +) + +// rateLimitPolicyHeader carries the concurrency limit as an RFC 8941 +// structured field, per draft-ietf-httpapi-ratelimit-headers. +const rateLimitPolicyHeader = "RateLimit-Policy" + +const ( + concurrencyPolicyName = "concurrency" + concurrencyUnit = "concurrent-requests" + // successesPerRegrow is how many consecutive upstream successes raise + // the learned limit by one. + successesPerRegrow = 4 + // learnedCeiling is the learned value at which the limit counts as + // unknown again. + learnedCeiling = 64 +) + +// concurrencyPolicyValue renders n as a RateLimit-Policy list member. +func concurrencyPolicyValue(n int) string { + return `"` + concurrencyPolicyName + `";q=` + strconv.Itoa(n) + `;qu="` + concurrencyUnit + `"` +} + +// providerConcurrency is one provider's learned state. +type providerConcurrency struct { + inFlight int + learned int // 0 means unknown + successes int +} + +// concurrencyLearner learns, per configured provider, how many requests the +// upstream serves at once. Keys are configured provider names only, so +// memory stays bounded. +type concurrencyLearner struct { + byProvider map[string]*providerConcurrency + mu sync.Mutex + explicit int +} + +func newConcurrencyLearner(explicit int) *concurrencyLearner { + return &concurrencyLearner{byProvider: map[string]*providerConcurrency{}, explicit: explicit} +} + +func (c *concurrencyLearner) state(provider string) *providerConcurrency { + s := c.byProvider[provider] + if s == nil { + s = &providerConcurrency{} + c.byProvider[provider] = s + } + return s +} + +// begin counts one in-flight upstream request; the returned func ends it. +func (c *concurrencyLearner) begin(provider string) func() { + c.mu.Lock() + c.state(provider).inFlight++ + c.mu.Unlock() + var once sync.Once + return func() { + once.Do(func() { + c.mu.Lock() + c.state(provider).inFlight-- + c.mu.Unlock() + }) + } +} + +// observe feeds one upstream attempt outcome into the provider's state. +func (c *concurrencyLearner) observe(provider string, resp *http.Response, err error) { + overload := false + success := false + switch { + case resp != nil: + code := resp.StatusCode + overload = code == http.StatusTooManyRequests || code == http.StatusServiceUnavailable || code == http.StatusGatewayTimeout + success = code < http.StatusBadRequest + case err != nil: + overload = isDeadlineExceeded(err) + } + if !overload && !success { + return + } + c.mu.Lock() + defer c.mu.Unlock() + s := c.state(provider) + if overload { + s.successes = 0 + next := s.inFlight / 2 + if next < 1 { + next = 1 + } + if s.learned == 0 || next < s.learned { + s.learned = next + } + return + } + if s.learned == 0 { + return + } + s.successes++ + if s.successes < successesPerRegrow { + return + } + s.successes = 0 + s.learned++ + if s.learned >= learnedCeiling { + s.learned = 0 + } +} + +// limit returns the effective advertised limit for provider: the smaller of +// the explicit and learned values, or 0 when neither is known. +func (c *concurrencyLearner) limit(provider string) int { + c.mu.Lock() + learned := 0 + if s := c.byProvider[provider]; s != nil { + learned = s.learned + } + c.mu.Unlock() + switch { + case c.explicit > 0 && (learned == 0 || c.explicit < learned): + return c.explicit + default: + return learned + } +} + +// setConcurrencyHeader sets or clears the RateLimit-Policy header on h. +func setConcurrencyHeader(h http.Header, n int) { + if n > 0 { + h.Set(rateLimitPolicyHeader, concurrencyPolicyValue(n)) + return + } + h.Del(rateLimitPolicyHeader) +} + +// trackUpstream wires per-provider accounting into ctx: the limiter's +// attempt record and, when concurrencyAutodetect is on, in-flight tracking, +// learning and the RateLimit-Policy header on w. The returned func must run +// when the upstream exchange ends. +func (g *Gateway) trackUpstream(ctx context.Context, w http.ResponseWriter, provider, model string) (context.Context, func()) { + c := g.concurrency + if c == nil { + return withAttemptRecorder(ctx, func(resp *http.Response, attemptErr error) { + g.limiter.recordProviderAttempt(provider, model, resp, attemptErr) + }), func() {} + } + end := c.begin(provider) + setConcurrencyHeader(w.Header(), c.limit(provider)) + return withAttemptRecorder(ctx, func(resp *http.Response, attemptErr error) { + g.limiter.recordProviderAttempt(provider, model, resp, attemptErr) + c.observe(provider, resp, attemptErr) + // The watchdog also calls this from a timer goroutine with a nil + // resp; only a real response means the request goroutine still owns w. + if resp != nil { + setConcurrencyHeader(w.Header(), c.limit(provider)) + } + }), end +} diff --git a/concurrency_header_test.go b/concurrency_header_test.go index 131ea47..dbee580 100644 --- a/concurrency_header_test.go +++ b/concurrency_header_test.go @@ -4,6 +4,7 @@ import ( "context" "net/http" "net/http/httptest" + "sync/atomic" "testing" ) @@ -26,7 +27,7 @@ func TestMaxConcurrency_AdvertisedOnResponses(t *testing.T) { header string max int }{ - {"configured limit is advertised", "3", 3}, + {"configured limit is advertised", `"concurrency";q=3;qu="concurrent-requests"`, 3}, {"zero sends no header", "", 0}, } for _, tc := range tests { @@ -40,8 +41,11 @@ func TestMaxConcurrency_AdvertisedOnResponses(t *testing.T) { req.Header.Set("Authorization", "Bearer sk-user") rec := httptest.NewRecorder() h.ServeHTTP(rec, req) - if got := rec.Header().Get("X-Concurrency-Limit"); got != tc.header { - t.Fatalf("X-Concurrency-Limit = %q, want %q", got, tc.header) + if got := rec.Header().Get("RateLimit-Policy"); got != tc.header { + t.Fatalf("RateLimit-Policy = %q, want %q", got, tc.header) + } + if got := rec.Header().Get("X-Concurrency-Limit"); got != "" { + t.Fatalf("legacy X-Concurrency-Limit = %q, want none", got) } }) } @@ -53,3 +57,160 @@ func TestMaxConcurrency_NegativeRejected(t *testing.T) { t.Fatal("New accepted a negative maxConcurrency") } } + +func TestConcurrencyLearner_Learning(t *testing.T) { + t.Parallel() + resp := func(code int) *http.Response { return &http.Response{StatusCode: code} } + tests := []struct { + err error + name string + events []*http.Response + explicit int + inFlight int + want int + }{ + {nil, "503 halves in-flight", []*http.Response{resp(503)}, 0, 8, 4}, + {nil, "429 halves in-flight", []*http.Response{resp(429)}, 0, 7, 3}, + {nil, "504 halves in-flight", []*http.Response{resp(504)}, 0, 6, 3}, + {context.DeadlineExceeded, "timeout halves in-flight", []*http.Response{nil}, 0, 10, 5}, + {nil, "floor of one", []*http.Response{resp(503)}, 0, 1, 1}, + {nil, "never raised on overload", []*http.Response{resp(503), resp(503)}, 0, 8, 4}, + {nil, "success without learning stays unknown", []*http.Response{resp(200), resp(200), resp(200), resp(200)}, 0, 8, 0}, + {nil, "four successes add one", []*http.Response{resp(503), resp(200), resp(200), resp(200), resp(200)}, 0, 8, 5}, + {nil, "three successes add nothing", []*http.Response{resp(503), resp(200), resp(200), resp(200)}, 0, 8, 4}, + {nil, "overload resets the streak", []*http.Response{resp(503), resp(200), resp(200), resp(200), resp(503), resp(200)}, 0, 8, 4}, + {nil, "client error is neutral", []*http.Response{resp(503), resp(404), resp(404), resp(404), resp(404)}, 0, 8, 4}, + {nil, "explicit caps learned", []*http.Response{resp(503)}, 2, 8, 2}, + {nil, "learned below explicit wins", []*http.Response{resp(503)}, 6, 4, 2}, + {nil, "explicit alone", []*http.Response{resp(200)}, 3, 4, 3}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + c := newConcurrencyLearner(tc.explicit) + for i := 0; i < tc.inFlight; i++ { + c.begin("p") + } + for _, r := range tc.events { + c.observe("p", r, tc.err) + } + if got := c.limit("p"); got != tc.want { + t.Fatalf("limit = %d, want %d", got, tc.want) + } + }) + } +} + +func TestConcurrencyLearner_ClearsAtCeiling(t *testing.T) { + t.Parallel() + c := newConcurrencyLearner(0) + for i := 0; i < 2*(learnedCeiling-1); i++ { + c.begin("p") + } + c.observe("p", &http.Response{StatusCode: 503}, nil) + if got := c.limit("p"); got != learnedCeiling-1 { + t.Fatalf("limit = %d, want %d", got, learnedCeiling-1) + } + for i := 0; i < successesPerRegrow; i++ { + c.observe("p", &http.Response{StatusCode: 200}, nil) + } + if got := c.limit("p"); got != 0 { + t.Fatalf("limit = %d after reaching ceiling, want 0 (unknown)", got) + } +} + +func TestConcurrencyLearner_EndReleasesInFlight(t *testing.T) { + t.Parallel() + c := newConcurrencyLearner(0) + end := c.begin("p") + c.begin("p") + end() + end() + c.observe("p", &http.Response{StatusCode: 503}, nil) + if got := c.limit("p"); got != 1 { + t.Fatalf("limit = %d, want 1 (one request still in flight, end is idempotent)", got) + } +} + +func newAutodetectGateway(t *testing.T, upstream string, autodetect bool, max int) http.Handler { + t.Helper() + cfg := CreateConfig() + cfg.ConcurrencyAutodetect = autodetect + cfg.MaxConcurrency = max + cfg.Providers = map[string]*ProviderConfig{ + "openai": {Type: "openai", BaseURL: upstream, APIKey: "sk-up", Models: []string{"gpt-test"}}, + } + cfg.Groups = map[string]*GroupConfig{"default": {}} + cfg.Users = &UsersConfig{Inline: []*UserConfig{{Name: "alice", Group: "default", APIKey: "sk-alice", Limits: &LimitsConfig{}}}} + h, err := New(context.Background(), http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}), cfg, "llmgw") + if err != nil { + t.Fatalf("New: %v", err) + } + return h +} + +func chatPolicyHeader(t *testing.T, h http.Handler) (int, string) { + t.Helper() + 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)) + return rec.Code, rec.Header().Get("RateLimit-Policy") +} + +func TestConcurrencyAutodetect_EndToEnd(t *testing.T) { + t.Parallel() + var status atomic.Int64 + status.Store(http.StatusOK) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(int(status.Load())) + _, _ = w.Write([]byte(`{"id":"x","object":"chat.completion","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)) + })) + defer srv.Close() + + t.Run("overload then regrow", func(t *testing.T) { + h := newAutodetectGateway(t, srv.URL, true, 0) + status.Store(http.StatusOK) + if _, got := chatPolicyHeader(t, h); got != "" { + t.Fatalf("header before any overload = %q, want none", got) + } + for _, code := range []int64{http.StatusServiceUnavailable, http.StatusTooManyRequests, http.StatusGatewayTimeout} { + status.Store(code) + gotCode, got := chatPolicyHeader(t, h) + if gotCode != int(code) { + t.Fatalf("status = %d, want %d passed through", gotCode, code) + } + if want := `"concurrency";q=1;qu="concurrent-requests"`; got != want { + t.Fatalf("header after %d = %q, want %q", code, got, want) + } + } + status.Store(http.StatusOK) + var got string + for i := 0; i < successesPerRegrow; i++ { + _, got = chatPolicyHeader(t, h) + } + if want := `"concurrency";q=2;qu="concurrent-requests"`; got != want { + t.Fatalf("header after regrow = %q, want %q", got, want) + } + }) + + t.Run("explicit caps learned", func(t *testing.T) { + h := newAutodetectGateway(t, srv.URL, true, 3) + status.Store(http.StatusOK) + if _, got := chatPolicyHeader(t, h); got != `"concurrency";q=3;qu="concurrent-requests"` { + t.Fatalf("header = %q, want explicit 3", got) + } + status.Store(http.StatusServiceUnavailable) + if _, got := chatPolicyHeader(t, h); got != `"concurrency";q=1;qu="concurrent-requests"` { + t.Fatalf("header after overload = %q, want learned 1", got) + } + }) + + t.Run("autodetect off never learns", func(t *testing.T) { + h := newAutodetectGateway(t, srv.URL, false, 0) + status.Store(http.StatusServiceUnavailable) + if _, got := chatPolicyHeader(t, h); got != "" { + t.Fatalf("header = %q, want none with autodetect off", got) + } + }) +} diff --git a/llmgateway.go b/llmgateway.go index e493741..c34e0cf 100644 --- a/llmgateway.go +++ b/llmgateway.go @@ -12,7 +12,6 @@ import ( "net/url" "os" "runtime" - "strconv" "strings" "sync" "time" @@ -184,12 +183,15 @@ type Config struct { // resolved (404 unknown model) and is not listed by GET /v1/models. It // removes the sorted-first bare-id collision winner. Defaults to false. RequireModelPrefix bool `json:"requireModelPrefix,omitempty"` - // MaxConcurrency, when > 0, is advertised on every response as - // X-Concurrency-Limit so clients that honour it (omps) cap their own - // in-flight requests to this gateway. It documents what the upstreams - // can serve at once; the gateway does not queue or reject on it. 0 - // (the default) sends no header. Negative values are rejected at - // construction. + // ConcurrencyAutodetect, when true, learns each provider's concurrency + // from overload answers (429/503/504/timeout) and advertises + // min(MaxConcurrency, learned). Defaults to false. + ConcurrencyAutodetect bool `json:"concurrencyAutodetect,omitempty"` + // MaxConcurrency, when > 0, is advertised on every response as a + // RateLimit-Policy concurrent-requests quota + // (draft-ietf-httpapi-ratelimit-headers). The gateway does not queue or + // reject on it. It wins over any learned value. 0 (the default) sends + // nothing. Negative values are rejected at construction. MaxConcurrency int `json:"maxConcurrency,omitempty"` } @@ -789,6 +791,7 @@ type Gateway struct { teeBudget *teeBudget fanoutGate chan struct{} failoverHealth *requestHealthTracker + concurrency *concurrencyLearner targetHealth *targetHealthTracker limiter *limiter registry *modelRegistry @@ -971,6 +974,9 @@ func newGateway(ctx context.Context, next http.Handler, config *Config, name str } g.failover = failoverCfg g.failoverHealth = newRequestHealthTracker() + if config.ConcurrencyAutodetect { + g.concurrency = newConcurrencyLearner(config.MaxConcurrency) + } // Security audit run-1: nothing validates the PRODUCT of the two // retry axes, which compose multiplicatively at runtime. Surface it @@ -1358,14 +1364,11 @@ func attachUsersFile(auth *authStore, path string, log gatewayLogger) error { } // ServeHTTP is the internal router. It grows in later tasks. -// concurrencyLimitHeader advertises Config.MaxConcurrency to clients. -const concurrencyLimitHeader = "X-Concurrency-Limit" - func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) { sw := &statusTrackingWriter{ResponseWriter: w} defer recoverPanic(sw, g) if g.cfg.MaxConcurrency > 0 { - sw.Header().Set(concurrencyLimitHeader, strconv.Itoa(g.cfg.MaxConcurrency)) + sw.Header().Set(rateLimitPolicyHeader, concurrencyPolicyValue(g.cfg.MaxConcurrency)) } // pruneRejectionsAfterReload only runs when maybeReload just performed // a REAL reload (users_file.go's own doc comment explains why it must diff --git a/routes_media.go b/routes_media.go index 7e78739..142b3ad 100644 --- a/routes_media.go +++ b/routes_media.go @@ -124,9 +124,8 @@ func (g *Gateway) handleImagesGenerations(w http.ResponseWriter, r *http.Request // Feature A (v0.22): see routes_unified.go's runUnified for the // identical pattern this mirrors. - ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) { - g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr) - }) + ctx, endUpstream := g.trackUpstream(r.Context(), sw, adapter.name(), upstreamModel) + defer endUpstream() // feat: instrument upstream latency — gated on metricsEnabled OR // admin.stats.latency, same reasoning as runUnified's identical wiring // (routes_unified.go). @@ -200,9 +199,8 @@ func (g *Gateway) handleAudioSpeech(w http.ResponseWriter, r *http.Request, u *u // Feature A (v0.22): see routes_unified.go's runUnified for the // identical pattern this mirrors. - ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) { - g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr) - }) + ctx, endUpstream := g.trackUpstream(r.Context(), sw, adapter.name(), upstreamModel) + defer endUpstream() // feat: instrument upstream latency — gated on metricsEnabled OR // admin.stats.latency, same reasoning as runUnified's identical wiring // (routes_unified.go). @@ -315,9 +313,8 @@ func (g *Gateway) handleAudioTranscriptions(w http.ResponseWriter, r *http.Reque // Feature A (v0.22): see routes_unified.go's runUnified for the // identical pattern this mirrors. - ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) { - g.limiter.recordProviderAttempt(adapter.name(), upstreamModel, resp, attemptErr) - }) + ctx, endUpstream := g.trackUpstream(r.Context(), sw, adapter.name(), upstreamModel) + defer endUpstream() // feat: instrument upstream latency — gated on metricsEnabled OR // admin.stats.latency, same reasoning as runUnified's identical wiring // (routes_unified.go). diff --git a/routes_passthrough.go b/routes_passthrough.go index cdbfb05..a5a5198 100644 --- a/routes_passthrough.go +++ b/routes_passthrough.go @@ -1297,9 +1297,9 @@ func (g *Gateway) handlePassthrough(w http.ResponseWriter, r *http.Request, u *u // (extractPassthroughUsage, below), read only after this attempt // already resolved, so there is no model to attribute it to yet. See // recordProviderAttempt's own doc comment (limits.go). - r = r.WithContext(withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) { - g.limiter.recordProviderAttempt(providerName, "", resp, attemptErr) - })) + trackedCtx, endUpstream := g.trackUpstream(r.Context(), w, providerName, "") + defer endUpstream() + r = r.WithContext(trackedCtx) // feat: instrument upstream latency — gated on metricsEnabled OR // admin.stats.latency (admin-redesign WP-A), same reasoning as // runUnified's identical wiring (routes_unified.go). No model here diff --git a/routes_unified.go b/routes_unified.go index 65591d1..093dbe4 100644 --- a/routes_unified.go +++ b/routes_unified.go @@ -587,9 +587,7 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop return } - ctx := withAttemptRecorder(r.Context(), func(resp *http.Response, attemptErr error) { - g.limiter.recordProviderAttempt(providerName, upstreamModel, resp, attemptErr) - }) + ctx, endUpstream := g.trackUpstream(r.Context(), respWriter, providerName, upstreamModel) // feat: instrument upstream latency — gated on metricsEnabled OR // admin.stats.latency (admin-redesign WP-A), so a deployment with // BOTH off wires nothing here at all (not even the closure @@ -615,6 +613,7 @@ func (g *Gateway) runMeteredCall(sw *statusTrackingWriter, r *http.Request, scop }) } result, callErr := call(cand.adapter, ctx, respWriter, attemptReq) + endUpstream() // Request-path health, recorded ONCE per candidate per logical // request (F3/F6 fix, adversarial-review) — see