From 2bd22821f5c7a49a595736f6e430b8f47020cfb2 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 28 Apr 2026 04:38:47 +0000 Subject: [PATCH 1/2] fix(relay): keep long JSON generation requests alive --- relay/image_handler.go | 4 ++ relay/json_keepalive.go | 118 +++++++++++++++++++++++++++++++++++ relay/json_keepalive_test.go | 108 ++++++++++++++++++++++++++++++++ relay/mjproxy_handler.go | 12 ++++ relay/relay_task.go | 4 ++ 5 files changed, 246 insertions(+) create mode 100644 relay/json_keepalive.go create mode 100644 relay/json_keepalive_test.go diff --git a/relay/image_handler.go b/relay/image_handler.go index e986dd897e65..7a6e29955c39 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -85,7 +85,11 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type statusCodeMappingStr := c.GetString("status_code_mapping") + keepalive := startJSONKeepalive(c, jsonKeepaliveInitialDelay, jsonKeepaliveInterval) + defer keepalive.stop() + resp, err := adaptor.DoRequest(c, info, requestBody) + keepalive.stop() if err != nil { return types.NewOpenAIError(err, types.ErrorCodeDoRequestFailed, http.StatusInternalServerError) } diff --git a/relay/json_keepalive.go b/relay/json_keepalive.go new file mode 100644 index 000000000000..4727f260a646 --- /dev/null +++ b/relay/json_keepalive.go @@ -0,0 +1,118 @@ +package relay + +import ( + "net/http" + "sync" + "sync/atomic" + "time" + + "github.com/gin-gonic/gin" +) + +const ( + jsonKeepaliveInitialDelay = 75 * time.Second + jsonKeepaliveInterval = 25 * time.Second +) + +type jsonKeepalive struct { + stopCh chan struct{} + doneCh chan struct{} + stopOnce sync.Once + + written atomic.Bool +} + +func startJSONKeepalive(c *gin.Context, initialDelay, interval time.Duration) *jsonKeepalive { + if c == nil || c.Request == nil || c.Writer == nil || initialDelay <= 0 || interval <= 0 { + return nil + } + + c.Header("Content-Type", "application/json; charset=utf-8") + c.Header("Cache-Control", "no-cache") + c.Header("X-Accel-Buffering", "no") + + keepalive := &jsonKeepalive{ + stopCh: make(chan struct{}), + doneCh: make(chan struct{}), + } + + go keepalive.run(c, initialDelay, interval) + return keepalive +} + +func (k *jsonKeepalive) stop() { + if k == nil { + return + } + k.stopOnce.Do(func() { + close(k.stopCh) + <-k.doneCh + }) +} + +func (k *jsonKeepalive) wasWritten() bool { + return k != nil && k.written.Load() +} + +func (k *jsonKeepalive) run(c *gin.Context, initialDelay, interval time.Duration) { + defer close(k.doneCh) + + timer := time.NewTimer(initialDelay) + defer timer.Stop() + + select { + case <-timer.C: + case <-c.Request.Context().Done(): + return + case <-k.stopCh: + return + } + + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + if !k.write(c) { + return + } + + select { + case <-ticker.C: + case <-c.Request.Context().Done(): + return + case <-k.stopCh: + return + } + } +} + +func (k *jsonKeepalive) write(c *gin.Context) bool { + if c == nil || c.Writer == nil { + return false + } + + writer := jsonKeepaliveResponseWriter(c) + if writer == nil { + return false + } + + writer.WriteHeader(http.StatusProcessing) + k.written.Store(true) + if flusher, ok := writer.(http.Flusher); ok { + flusher.Flush() + } + return true +} + +func jsonKeepaliveResponseWriter(c *gin.Context) http.ResponseWriter { + type responseWriterUnwrapper interface { + Unwrap() http.ResponseWriter + } + if c == nil || c.Writer == nil { + return nil + } + if unwrapper, ok := c.Writer.(responseWriterUnwrapper); ok { + return unwrapper.Unwrap() + } + return nil +} diff --git a/relay/json_keepalive_test.go b/relay/json_keepalive_test.go new file mode 100644 index 000000000000..bc1191bbda5a --- /dev/null +++ b/relay/json_keepalive_test.go @@ -0,0 +1,108 @@ +package relay + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type jsonKeepaliveRecorder struct { + mu sync.Mutex + header http.Header + informational []int + finalStatus int + body []byte + flushes int +} + +func newJSONKeepaliveRecorder() *jsonKeepaliveRecorder { + return &jsonKeepaliveRecorder{header: make(http.Header)} +} + +func (r *jsonKeepaliveRecorder) Header() http.Header { + return r.header +} + +func (r *jsonKeepaliveRecorder) Write(data []byte) (int, error) { + r.mu.Lock() + defer r.mu.Unlock() + if r.finalStatus == 0 { + r.finalStatus = http.StatusOK + } + r.body = append(r.body, data...) + return len(data), nil +} + +func (r *jsonKeepaliveRecorder) WriteHeader(code int) { + r.mu.Lock() + defer r.mu.Unlock() + if code >= 100 && code < 200 { + r.informational = append(r.informational, code) + return + } + r.finalStatus = code +} + +func (r *jsonKeepaliveRecorder) Flush() { + r.mu.Lock() + defer r.mu.Unlock() + r.flushes++ +} + +func (r *jsonKeepaliveRecorder) snapshot() ([]int, int, int, string) { + r.mu.Lock() + defer r.mu.Unlock() + informational := append([]int(nil), r.informational...) + return informational, r.finalStatus, r.flushes, string(r.body) +} + +func newJSONKeepaliveTestContext(rec *jsonKeepaliveRecorder) *gin.Context { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + return c +} + +func TestJSONKeepaliveSendsInformationalProcessing(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, time.Millisecond, time.Millisecond) + require.NotNil(t, keepalive) + require.Eventually(t, func() bool { + informational, _, flushes, _ := rec.snapshot() + return keepalive.wasWritten() && len(informational) > 0 && flushes > 0 + }, time.Second, time.Millisecond) + keepalive.stop() + + informational, finalStatus, flushes, _ := rec.snapshot() + require.Equal(t, http.StatusProcessing, informational[0]) + require.Greater(t, flushes, 0) + require.Zero(t, finalStatus) +} + +func TestJSONKeepalivePreservesFinalJSONStatus(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, time.Millisecond, time.Millisecond) + require.Eventually(t, func() bool { + return keepalive.wasWritten() + }, time.Second, time.Millisecond) + keepalive.stop() + + rec.WriteHeader(http.StatusOK) + _, err := rec.Write([]byte(`{"data":[]}`)) + require.NoError(t, err) + + informational, finalStatus, _, body := rec.snapshot() + require.NotEmpty(t, informational) + require.Equal(t, http.StatusProcessing, informational[0]) + require.Equal(t, http.StatusOK, finalStatus) + require.Equal(t, `{"data":[]}`, body) +} diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index ee48ca64b10b..ffc1adb00457 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -218,7 +218,11 @@ func RelaySwapFace(c *gin.Context, info *relaycommon.RelayInfo) *dto.MidjourneyR requestURL := getMjRequestPath(c.Request.URL.String()) baseURL := c.GetString("base_url") fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL) + keepalive := startJSONKeepalive(c, jsonKeepaliveInitialDelay, jsonKeepaliveInterval) + defer keepalive.stop() + mjResp, _, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) + keepalive.stop() if err != nil { return &mjResp.Response } @@ -301,7 +305,11 @@ func RelayMidjourneyTaskImageSeed(c *gin.Context) *dto.MidjourneyResponse { requestURL := getMjRequestPath(c.Request.URL.String()) fullRequestURL := fmt.Sprintf("%s%s", channel.GetBaseURL(), requestURL) + keepalive := startJSONKeepalive(c, jsonKeepaliveInitialDelay, jsonKeepaliveInterval) + defer keepalive.stop() + midjResponseWithStatus, _, err := service.DoMidjourneyHttpRequest(c, time.Second*30, fullRequestURL) + keepalive.stop() if err != nil { return &midjResponseWithStatus.Response } @@ -523,7 +531,11 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt } } + keepalive := startJSONKeepalive(c, jsonKeepaliveInitialDelay, jsonKeepaliveInterval) + defer keepalive.stop() + midjResponseWithStatus, responseBody, err := service.DoMidjourneyHttpRequest(c, time.Second*60, fullRequestURL) + keepalive.stop() if err != nil { return &midjResponseWithStatus.Response } diff --git a/relay/relay_task.go b/relay/relay_task.go index 098e23828b6c..9d873a2dda67 100644 --- a/relay/relay_task.go +++ b/relay/relay_task.go @@ -217,7 +217,11 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe } // 9. 发送请求 + keepalive := startJSONKeepalive(c, jsonKeepaliveInitialDelay, jsonKeepaliveInterval) + defer keepalive.stop() + resp, err := adaptor.DoRequest(c, info, requestBody) + keepalive.stop() if err != nil { return nil, service.TaskErrorWrapper(err, "do_request_failed", http.StatusInternalServerError) } From d55f048036d889f730b00aff45415d406e3dc9db Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 28 Apr 2026 05:00:40 +0000 Subject: [PATCH 2/2] fix(relay): scope JSON keepalive headers --- relay/json_keepalive.go | 46 ++++++++++++++-- relay/json_keepalive_test.go | 103 ++++++++++++++++++++++++++++------- 2 files changed, 125 insertions(+), 24 deletions(-) diff --git a/relay/json_keepalive.go b/relay/json_keepalive.go index 4727f260a646..ea205cfd3899 100644 --- a/relay/json_keepalive.go +++ b/relay/json_keepalive.go @@ -27,10 +27,6 @@ func startJSONKeepalive(c *gin.Context, initialDelay, interval time.Duration) *j return nil } - c.Header("Content-Type", "application/json; charset=utf-8") - c.Header("Cache-Control", "no-cache") - c.Header("X-Accel-Buffering", "no") - keepalive := &jsonKeepalive{ stopCh: make(chan struct{}), doneCh: make(chan struct{}), @@ -96,6 +92,9 @@ func (k *jsonKeepalive) write(c *gin.Context) bool { return false } + headerSnapshot := setJSONKeepaliveHeaders(writer.Header()) + defer restoreJSONKeepaliveHeaders(writer.Header(), headerSnapshot) + writer.WriteHeader(http.StatusProcessing) k.written.Store(true) if flusher, ok := writer.(http.Flusher); ok { @@ -104,6 +103,45 @@ func (k *jsonKeepalive) write(c *gin.Context) bool { return true } +type jsonKeepaliveHeaderSnapshot struct { + values []string + exists bool +} + +func setJSONKeepaliveHeaders(header http.Header) map[string]jsonKeepaliveHeaderSnapshot { + const ( + contentType = "Content-Type" + cacheControl = "Cache-Control" + accelBuffering = "X-Accel-Buffering" + jsonContentType = "application/json; charset=utf-8" + ) + + keys := []string{contentType, cacheControl, accelBuffering} + snapshot := make(map[string]jsonKeepaliveHeaderSnapshot, len(keys)) + for _, key := range keys { + values, exists := header[key] + snapshot[key] = jsonKeepaliveHeaderSnapshot{ + values: append([]string(nil), values...), + exists: exists, + } + } + + header.Set(contentType, jsonContentType) + header.Set(cacheControl, "no-cache") + header.Set(accelBuffering, "no") + return snapshot +} + +func restoreJSONKeepaliveHeaders(header http.Header, snapshot map[string]jsonKeepaliveHeaderSnapshot) { + for key, previous := range snapshot { + if previous.exists { + header[key] = append([]string(nil), previous.values...) + } else { + delete(header, key) + } + } +} + func jsonKeepaliveResponseWriter(c *gin.Context) http.ResponseWriter { type responseWriterUnwrapper interface { Unwrap() http.ResponseWriter diff --git a/relay/json_keepalive_test.go b/relay/json_keepalive_test.go index bc1191bbda5a..e98d0a35702a 100644 --- a/relay/json_keepalive_test.go +++ b/relay/json_keepalive_test.go @@ -12,12 +12,13 @@ import ( ) type jsonKeepaliveRecorder struct { - mu sync.Mutex - header http.Header - informational []int - finalStatus int - body []byte - flushes int + mu sync.Mutex + header http.Header + informational []int + informationalHeaders []http.Header + finalStatus int + body []byte + flushes int } func newJSONKeepaliveRecorder() *jsonKeepaliveRecorder { @@ -43,6 +44,7 @@ func (r *jsonKeepaliveRecorder) WriteHeader(code int) { defer r.mu.Unlock() if code >= 100 && code < 200 { r.informational = append(r.informational, code) + r.informationalHeaders = append(r.informationalHeaders, cloneTestHeader(r.header)) return } r.finalStatus = code @@ -54,11 +56,34 @@ func (r *jsonKeepaliveRecorder) Flush() { r.flushes++ } -func (r *jsonKeepaliveRecorder) snapshot() ([]int, int, int, string) { +type jsonKeepaliveRecorderSnapshot struct { + informational []int + informationalHeaders []http.Header + finalStatus int + flushes int + body string + header http.Header +} + +func (r *jsonKeepaliveRecorder) snapshot() jsonKeepaliveRecorderSnapshot { r.mu.Lock() defer r.mu.Unlock() - informational := append([]int(nil), r.informational...) - return informational, r.finalStatus, r.flushes, string(r.body) + return jsonKeepaliveRecorderSnapshot{ + informational: append([]int(nil), r.informational...), + informationalHeaders: append([]http.Header(nil), r.informationalHeaders...), + finalStatus: r.finalStatus, + flushes: r.flushes, + body: string(r.body), + header: cloneTestHeader(r.header), + } +} + +func cloneTestHeader(src http.Header) http.Header { + dst := make(http.Header, len(src)) + for key, values := range src { + dst[key] = append([]string(nil), values...) + } + return dst } func newJSONKeepaliveTestContext(rec *jsonKeepaliveRecorder) *gin.Context { @@ -75,15 +100,21 @@ func TestJSONKeepaliveSendsInformationalProcessing(t *testing.T) { keepalive := startJSONKeepalive(c, time.Millisecond, time.Millisecond) require.NotNil(t, keepalive) require.Eventually(t, func() bool { - informational, _, flushes, _ := rec.snapshot() - return keepalive.wasWritten() && len(informational) > 0 && flushes > 0 + snap := rec.snapshot() + return keepalive.wasWritten() && len(snap.informational) > 0 && snap.flushes > 0 }, time.Second, time.Millisecond) keepalive.stop() - informational, finalStatus, flushes, _ := rec.snapshot() - require.Equal(t, http.StatusProcessing, informational[0]) - require.Greater(t, flushes, 0) - require.Zero(t, finalStatus) + snap := rec.snapshot() + require.Equal(t, http.StatusProcessing, snap.informational[0]) + require.Equal(t, "application/json; charset=utf-8", snap.informationalHeaders[0].Get("Content-Type")) + require.Equal(t, "no-cache", snap.informationalHeaders[0].Get("Cache-Control")) + require.Equal(t, "no", snap.informationalHeaders[0].Get("X-Accel-Buffering")) + require.Empty(t, snap.header.Get("Content-Type")) + require.Empty(t, snap.header.Get("Cache-Control")) + require.Empty(t, snap.header.Get("X-Accel-Buffering")) + require.Greater(t, snap.flushes, 0) + require.Zero(t, snap.finalStatus) } func TestJSONKeepalivePreservesFinalJSONStatus(t *testing.T) { @@ -100,9 +131,41 @@ func TestJSONKeepalivePreservesFinalJSONStatus(t *testing.T) { _, err := rec.Write([]byte(`{"data":[]}`)) require.NoError(t, err) - informational, finalStatus, _, body := rec.snapshot() - require.NotEmpty(t, informational) - require.Equal(t, http.StatusProcessing, informational[0]) - require.Equal(t, http.StatusOK, finalStatus) - require.Equal(t, `{"data":[]}`, body) + snap := rec.snapshot() + require.NotEmpty(t, snap.informational) + require.Equal(t, http.StatusProcessing, snap.informational[0]) + require.Equal(t, http.StatusOK, snap.finalStatus) + require.Equal(t, `{"data":[]}`, snap.body) +} + +func TestJSONKeepaliveDoesNotSetHeadersBeforeFirstTick(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, time.Hour, time.Hour) + require.NotNil(t, keepalive) + keepalive.stop() + + snap := rec.snapshot() + require.Empty(t, snap.informational) + require.Empty(t, snap.header.Get("Content-Type")) + require.Empty(t, snap.header.Get("Cache-Control")) + require.Empty(t, snap.header.Get("X-Accel-Buffering")) +} + +func TestJSONKeepaliveRepeatsUntilStoppedThenStaysQuiet(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, time.Millisecond, time.Millisecond) + require.NotNil(t, keepalive) + require.Eventually(t, func() bool { + return len(rec.snapshot().informational) >= 3 + }, time.Second, time.Millisecond) + + keepalive.stop() + countAfterStop := len(rec.snapshot().informational) + time.Sleep(10 * time.Millisecond) + + require.Equal(t, countAfterStop, len(rec.snapshot().informational)) }