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.
This commit is contained in:
yetone committed 2026-10-07 03:58:27 +08:00
1 parent b234507842
commit 443f1cbc2d
3 files changed
+73 -5

No files matched your search

+25 -2
View File
@@ -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")
+9 -3
View File
@@ -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)
}
+39
View File
@@ -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())
}
}