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()) + } +}