fix(relay): keep long JSON generation requests alive

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

relay/json_keepalive.go

@@ -0,0 +1,156 @@
+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
+	}
+
+	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
+	}
+
+	headerSnapshot := setJSONKeepaliveHeaders(writer.Header())
+	defer restoreJSONKeepaliveHeaders(writer.Header(), headerSnapshot)
+
+	writer.WriteHeader(http.StatusProcessing)
+	k.written.Store(true)
+	if flusher, ok := writer.(http.Flusher); ok {
+		flusher.Flush()
+	}
+	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
+	}
+	if c == nil || c.Writer == nil {
+		return nil
+	}
+	if unwrapper, ok := c.Writer.(responseWriterUnwrapper); ok {
+		return unwrapper.Unwrap()
+	}
+	return nil
+}

relay/json_keepalive_test.go

@@ -0,0 +1,171 @@
+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
+	informationalHeaders []http.Header
+	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)
+		r.informationalHeaders = append(r.informationalHeaders, cloneTestHeader(r.header))
+		return
+	}
+	r.finalStatus = code
+}
+
+func (r *jsonKeepaliveRecorder) Flush() {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.flushes++
+}
+
+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()
+	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 {
+	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 {
+		snap := rec.snapshot()
+		return keepalive.wasWritten() && len(snap.informational) > 0 && snap.flushes > 0
+	}, time.Second, time.Millisecond)
+	keepalive.stop()
+
+	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) {
+	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)
+
+	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))
+}

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
 	}

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