From 443f1cbc2df8ac8e2dfe71ff757e7d7ce411f53c Mon Sep 17 00:00:00 2001 From: yetone Date: Wed, 7 Oct 2026 03:58:27 +0800 Subject: [PATCH] gateway: a held stream's first keepalive reads only magpie's notes until the vendor's reply has its status, so it never races the try writing the reply's headers go test -race ./internal/gateway failed in TestQuietStreamsKeptAliveEveryProtocol (chat held) with six data races: keepAlive, on watch's goroutine, sent the held stream's 200 with the try's headers by iterating holdWriter.header, while the try's goroutine wrote that map with no lock: noteMember (#822's X-Magpie-Provider/Model) and passthrough copying the vendor's headers in. In production, a vendor that answers only after keepHeldAfter does the same, and outside -race that can end the process with "concurrent map iteration and map write". Before the try has its status (WriteHeader, under mu), keepAlive now sends only magpie's own notes, which noteMember sets under mu (holdWriter.note); the vendor has said nothing yet, so nothing of its is lost. It also clones what it copies, so the agent's header shares no slice with the try's. TestKeepAliveBeforeTheVendorsHeaders (a try writing its headers while the first comment goes out) reports the race on the old code (2 of 3 runs); -race -count=20 passes, and so does the whole package under -race. --- internal/gateway/fallback.go | 27 ++++++++++++++-- internal/gateway/member.go | 12 ++++++-- internal/gateway/quiet_keepalive_test.go | 39 ++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/internal/gateway/fallback.go b/internal/gateway/fallback.go index be1e6f5c..04069242 100644 --- a/internal/gateway/fallback.go +++ b/internal/gateway/fallback.go @@ -987,6 +987,10 @@ type holdWriter struct { // reasoning keeps the agent from its idle timeout with SSE comments // meanwhile (#751) alive *keptAlive + // notes are the headers magpie says of the try (noteMember), set + // under mu: all keepAlive sends of the try's before it has its status, + // as until then the try is still writing header, with no lock + notes http.Header first firstToken // when its first content and text came (#196) @@ -1026,6 +1030,18 @@ func newHoldWriter(w http.ResponseWriter, hold bool) *holdWriter { func (h *holdWriter) Header() http.Header { return h.header } +// note sets a header magpie says of the try, which keepAlive may send +// from watch's goroutine before the try has its status. +func (h *holdWriter) note(k, v string) { + h.mu.Lock() + defer h.mu.Unlock() + h.header.Set(k, v) + if h.notes == nil { + h.notes = http.Header{} + } + h.notes.Set(k, v) +} + func (h *holdWriter) WriteHeader(code int) { h.mu.Lock() defer h.mu.Unlock() @@ -1244,10 +1260,17 @@ func (h *holdWriter) keepAlive() { return } if !a.sent { + // the try's header is whole once it has its status (WriteHeader, + // under mu); before, the vendor's are still being copied into it + // from the try's goroutine, and only magpie's notes are read + src := h.header + if h.status == 0 { + src = h.notes + } dst := h.w.Header() - for k, v := range h.header { + for k, v := range src { if k != resetsHeader && k != refusedHeader && k != "Content-Length" { - dst[k] = v + dst[k] = slices.Clone(v) } } dst.Set("Content-Type", "text/event-stream; charset=utf-8") diff --git a/internal/gateway/member.go b/internal/gateway/member.go index bd46f4ef..6b728d2b 100644 --- a/internal/gateway/member.go +++ b/internal/gateway/member.go @@ -43,9 +43,15 @@ type memberKey struct{} // and tells the request's memberWriter, if it has one. func noteMember(w http.ResponseWriter, r *http.Request, p provider.Provider, model string) { id := p.ID + "/" + model - h := w.Header() - h.Set(providerHeader, p.ID) - h.Set(modelHeader, id) + if hw, ok := w.(*holdWriter); ok { + // watch's keepAlive may send them meanwhile + hw.note(providerHeader, p.ID) + hw.note(modelHeader, id) + } else { + h := w.Header() + h.Set(providerHeader, p.ID) + h.Set(modelHeader, id) + } if mw, _ := r.Context().Value(memberKey{}).(*memberWriter); mw != nil { mw.setMember(id) } diff --git a/internal/gateway/quiet_keepalive_test.go b/internal/gateway/quiet_keepalive_test.go index 521f68f1..89117360 100644 --- a/internal/gateway/quiet_keepalive_test.go +++ b/internal/gateway/quiet_keepalive_test.go @@ -5,6 +5,7 @@ import ( "io" "net/http" "net/http/httptest" + "strconv" "strings" "sync" "sync/atomic" @@ -270,3 +271,41 @@ func TestQuietStreamsKeptAliveEveryProtocol(t *testing.T) { } } } + +// A vendor slow to answer at all gets its stream's agent kept alive from +// watch's goroutine while the try, once the vendor answers, copies the +// vendor's headers in: keepAlive reads only magpie's notes until the try +// has its status, so the two never touch one map (a race, and in a build +// without -race "concurrent map iteration and map write"). +func TestKeepAliveBeforeTheVendorsHeaders(t *testing.T) { + quietFast(t) + rec := httptest.NewRecorder() + h := newHoldWriter(rec, true) + h.streams, h.alive = true, &keptAlive{proto: provider.Responses} + r := httptest.NewRequest("POST", "/v1/responses", nil) + noteMember(h, r, provider.Provider{ID: "q"}, "m") + end := h.watch() + // the vendor's headers come in as passthrough copies them, the first + // comment going out meanwhile + deadline := time.Now().Add(5 * time.Second) + for i := 0; ; i++ { + h.Header().Set("X-Vendor", strconv.Itoa(i)) + h.mu.Lock() + sent := h.alive.sent + h.mu.Unlock() + if sent && i > 100 || time.Now().After(deadline) { + break + } + time.Sleep(100 * time.Microsecond) + } + end() + if !h.alive.sent { + t.Fatal("no comment went out") + } + if got := rec.Header().Get(providerHeader); got != "q" { + t.Fatalf("the stream's headers before its reply lack magpie's notes: %q", got) + } + if !strings.Contains(rec.Body.String(), ": keepalive") { + t.Fatalf("body: %q", rec.Body.String()) + } +}