diff --git a/core/observe/statserver.go b/core/observe/statserver.go index 5962d959a..fbd29324b 100644 --- a/core/observe/statserver.go +++ b/core/observe/statserver.go @@ -8,6 +8,7 @@ import ( "net" "net/http" "net/http/pprof" + "runtime" "strings" "time" @@ -107,6 +108,13 @@ func NewStatServer(addr string, configProvider ConfigProvider, statsProvider Sta mux.Handle("/tls-bridge/unread", o.tlsBridgeUnread) } if o.pprof { + // /debug/pprof/block and /mutex stay empty without these: the handlers + // exist, but the runtime records nothing unless asked. Both cost only + // when the event they measure happens — a blocked goroutine or a contended + // mutex — so they ride the same gate as the handlers rather than running + // on deployments where pprof is off. + runtime.SetBlockProfileRate(1) + runtime.SetMutexProfileFraction(1) pprofGuard := func(h http.HandlerFunc) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { if !loopbackHost(r.Host) { diff --git a/core/observe/statserver_test.go b/core/observe/statserver_test.go index a83115af7..d65922400 100644 --- a/core/observe/statserver_test.go +++ b/core/observe/statserver_test.go @@ -6,8 +6,10 @@ import ( "net" "net/http" "net/http/httptest" + "runtime" "strings" "testing" + "time" "github.com/rossoctl/cortex/core/auth" "github.com/rossoctl/cortex/core/config" @@ -241,6 +243,33 @@ func TestStatServerPprofServesProfiles(t *testing.T) { } } +// /debug/pprof/block and /mutex answer 200 but stay empty forever unless the +// runtime is told to record contention — the handlers exist either way. WithPprof +// turns the recorders on with the handlers, so the endpoints can show something. +func TestStatServerPprofEnablesContentionProfiles(t *testing.T) { + handler := serveMuxOpts(newTestConfig(), auth.NewStats(), WithPprof(true)) + + // Mutex has a getter; block does not, so block is proven by recording one + // event: a channel receive that waits is a block, and at rate 0 it would + // never reach the profile. + ch := make(chan struct{}) + go func() { time.Sleep(100 * time.Millisecond); close(ch) }() + <-ch + + req := httptest.NewRequest(http.MethodGet, "/debug/pprof/block?debug=1", nil) + req.Host = "localhost:47802" // a loopback Host, as a local client sends + w := httptest.NewRecorder() + handler.ServeHTTP(w, req) + // The blocked receive's stack names this test; an empty profile carries only + // the header lines. + if !strings.Contains(w.Body.String(), "TestStatServerPprofEnablesContentionProfiles") { + t.Error("a 100ms blocked receive is absent from /debug/pprof/block: the recorder is off") + } + if prev := runtime.SetMutexProfileFraction(1); prev == 0 { + t.Error("WithPprof left mutex profiling off: /debug/pprof/mutex would stay empty") + } +} + // A page can rebind a DNS name to 127.0.0.1 and have a browser read a loopback // port, so the profiles refuse any Host that is not exactly loopback — the same // defense as the session API's erase guards. diff --git a/core/pipeline/model.go b/core/pipeline/model.go index 63e3bd19f..886f9861a 100644 --- a/core/pipeline/model.go +++ b/core/pipeline/model.go @@ -5,6 +5,7 @@ import ( "fmt" "log/slog" "strings" + "unsafe" "github.com/tidwall/gjson" "github.com/tidwall/sjson" @@ -115,7 +116,12 @@ func requestModel(body []byte) (string, error) { } var model gjson.Result exact, anyCase := 0, 0 - gjson.ParseBytes(body).ForEach(func(key, value gjson.Result) bool { + // Zero-copy: ParseBytes is Parse(string(body)) — a full copy of every routed + // request just to scan its top-level keys. The alias is safe because gjson + // never writes the string it parses, the body is not written concurrently on + // the request path, and nothing that outlives the scan points into it: the + // only string that escapes, below, is cloned. + gjson.Parse(unsafe.String(unsafe.SliceData(body), len(body))).ForEach(func(key, value gjson.Result) bool { if key.Type != gjson.String || !strings.EqualFold(key.String(), "model") { return true } diff --git a/core/pipeline/model_test.go b/core/pipeline/model_test.go index 00f319934..5078af603 100644 --- a/core/pipeline/model_test.go +++ b/core/pipeline/model_test.go @@ -282,3 +282,27 @@ func TestRequestModelDoesNotPinTheBody(t *testing.T) { t.Errorf("name = %q", names[0]) } } + +// TestRequestModelDoesNotCopyTheBody is the allocation regression for the same +// scan: ParseBytes is Parse(string(body)), so reading the top-level keys used to +// copy the whole body once per routed request — ~1GB/hour of churn on a live +// proxy, third behind tool-prune's copies. The scan now runs over an alias of +// the body, so one call on a 4MiB body must allocate well under the body's size. +func TestRequestModelDoesNotCopyTheBody(t *testing.T) { + pad := strings.Repeat("x", 4<<20) + body := []byte(`{"model":"claude-opus-5-5","pad":"` + pad + `"}`) + if _, err := requestModel(body); err != nil { // warm: page in gjson's code + t.Fatalf("requestModel: %v", err) + } + runtime.GC() + var before runtime.MemStats + runtime.ReadMemStats(&before) + if _, err := requestModel(body); err != nil { + t.Fatalf("requestModel: %v", err) + } + var after runtime.MemStats + runtime.ReadMemStats(&after) + if got := after.TotalAlloc - before.TotalAlloc; got > 64<<10 { + t.Errorf("requestModel on a %dMiB body allocated %dKiB; the scan copies the body", len(body)>>20, got>>10) + } +}