diff --git a/relay/channel/api_request.go b/relay/channel/api_request.go index f945a838..f2601600 100644 --- a/relay/channel/api_request.go +++ b/relay/channel/api_request.go @@ -396,10 +396,12 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody return targetConn, nil } -func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.CancelFunc { +func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) (context.CancelFunc, <-chan struct{}) { pingerCtx, stopPinger := context.WithCancel(context.Background()) + done := make(chan struct{}) gopool.Go(func() { + defer close(done) defer func() { // 增加panic恢复处理 if r := recover(); r != nil { @@ -449,36 +451,24 @@ func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.Canc } }) - return stopPinger + return stopPinger, done } func sendPingData(c *gin.Context, mutex *sync.Mutex) error { - // 增加超时控制,防止锁死等待 - done := make(chan error, 1) - go func() { - mutex.Lock() - defer mutex.Unlock() + mutex.Lock() + defer mutex.Unlock() - err := helper.PingData(c) - if err != nil { - logger.LogError(c, "SSE ping error: "+err.Error()) - done <- err - return - } - - logger.LogDebug(c, "SSE ping data sent") - done <- nil - }() - - // 设置发送ping数据的超时时间 - select { - case err := <-done: + // Bound the write so a slow client cannot block this goroutine forever; + // doRequest's defer waits for the pinger to exit before returning. + helper.ExtendWriteDeadline(c) + err := helper.PingData(c) + if err != nil { + logger.LogError(c, "SSE ping error: "+err.Error()) return err - case <-time.After(10 * time.Second): - return errors.New("SSE ping data send timeout") - case <-c.Request.Context().Done(): - return errors.New("request context cancelled during ping") } + + logger.LogDebug(c, "SSE ping data sent") + return nil } func DoRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { @@ -497,17 +487,19 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http } var stopPinger context.CancelFunc + var pingerDone <-chan struct{} if info.IsStream { helper.SetEventStreamHeaders(c) // 处理流式请求的 ping 保活 generalSettings := operation_setting.GetGeneralSetting() if generalSettings.PingIntervalEnabled && !info.DisablePing { pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second - stopPinger = startPingKeepAlive(c, pingInterval) + stopPinger, pingerDone = startPingKeepAlive(c, pingInterval) // 使用defer确保在任何情况下都能停止ping goroutine defer func() { if stopPinger != nil { stopPinger() + <-pingerDone logger.LogDebug(c, "SSE ping goroutine stopped by defer") } }() diff --git a/relay/channel/openai/helper.go b/relay/channel/openai/helper.go index 1a01d06d..d10176b5 100644 --- a/relay/channel/openai/helper.go +++ b/relay/channel/openai/helper.go @@ -206,5 +206,5 @@ func sendResponsesStreamData(c *gin.Context, streamResponse dto.ResponsesStreamR if data == "" { return } - helper.ResponseChunkData(c, streamResponse, data) + _ = helper.ResponseChunkData(c, streamResponse, data) } diff --git a/relay/channel/openai/relay_image.go b/relay/channel/openai/relay_image.go index 436f52cd..1eccbe2c 100644 --- a/relay/channel/openai/relay_image.go +++ b/relay/channel/openai/relay_image.go @@ -136,10 +136,10 @@ func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) { } _ = common.Unmarshal(data, &payload) if eventName := strings.TrimSpace(payload.Type); eventName != "" { - c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", eventName)}) + _ = helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) + return } - c.Render(-1, common.CustomEvent{Data: "data: " + string(data)}) - _ = helper.FlushWriter(c) + _ = helper.StringData(c, string(data)) } // isOpenAIImageStreamErrorEvent detects upstream error chunks by JSON content @@ -269,19 +269,11 @@ func writeOpenaiImageStreamPayload(c *gin.Context, eventName string, payload any return err } if eventName != "" { - if _, err := fmt.Fprintf(c.Writer, "event: %s\n", eventName); err != nil { - return err - } + return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) } - if _, err := fmt.Fprintf(c.Writer, "data: %s\n\n", data); err != nil { - return err - } - return helper.FlushWriter(c) + return helper.StringData(c, string(data)) } func writeOpenaiImageStreamDone(c *gin.Context) error { - if _, err := fmt.Fprint(c.Writer, "data: [DONE]\n\n"); err != nil { - return err - } - return helper.FlushWriter(c) + return helper.StringData(c, "[DONE]") } diff --git a/relay/helper/common.go b/relay/helper/common.go index 17ce79d2..5b118aef 100644 --- a/relay/helper/common.go +++ b/relay/helper/common.go @@ -25,7 +25,7 @@ func FlushWriter(c *gin.Context) (err error) { return nil } - if c.Request != nil && c.Request.Context().Err() != nil { + if requestContextDone(c) { return fmt.Errorf("request context done: %w", c.Request.Context().Err()) } @@ -38,6 +38,10 @@ func FlushWriter(c *gin.Context) (err error) { return nil } +func requestContextDone(c *gin.Context) bool { + return c != nil && c.Request != nil && c.Request.Context().Err() != nil +} + func SetEventStreamHeaders(c *gin.Context) { // 检查是否已经设置过头部 if _, exists := c.Get("event_stream_headers_set"); exists { @@ -55,6 +59,10 @@ func SetEventStreamHeaders(c *gin.Context) { } func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { + if requestContextDone(c) { + return nil + } + jsonData, err := common.Marshal(resp) if err != nil { common.SysError("error marshalling stream response: " + err.Error()) @@ -67,15 +75,23 @@ func ClaudeData(c *gin.Context, resp dto.ClaudeResponse) error { } func ClaudeChunkData(c *gin.Context, resp dto.ClaudeResponse, data string) { + if requestContextDone(c) { + return + } + c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s\n", data)}) _ = FlushWriter(c) } -func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) { +func ResponseChunkData(c *gin.Context, resp dto.ResponsesStreamResponse, data string) error { + if requestContextDone(c) { + return fmt.Errorf("request context done: %w", c.Request.Context().Err()) + } + c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("event: %s\n", resp.Type)}) c.Render(-1, common.CustomEvent{Data: fmt.Sprintf("data: %s", data)}) - _ = FlushWriter(c) + return FlushWriter(c) } func StringData(c *gin.Context, str string) error { @@ -83,7 +99,7 @@ func StringData(c *gin.Context, str string) error { return errors.New("context or writer is nil") } - if c.Request != nil && c.Request.Context().Err() != nil { + if requestContextDone(c) { return fmt.Errorf("request context done: %w", c.Request.Context().Err()) } @@ -96,7 +112,7 @@ func PingData(c *gin.Context) error { return errors.New("context or writer is nil") } - if c.Request != nil && c.Request.Context().Err() != nil { + if requestContextDone(c) { return fmt.Errorf("request context done: %w", c.Request.Context().Err()) } diff --git a/relay/helper/stream_scanner.go b/relay/helper/stream_scanner.go index bc0af1bb..4d200ddd 100644 --- a/relay/helper/stream_scanner.go +++ b/relay/helper/stream_scanner.go @@ -25,6 +25,11 @@ const ( InitialScannerBufferSize = 64 << 10 // 64KB (64*1024) DefaultMaxScannerBufferSize = 128 << 20 // 64MB (64*1024*1024) default SSE buffer size DefaultPingInterval = 10 * time.Second + // streamWriteTimeout bounds a single blocked write to a slow client so the + // unconditional wg.Wait() in cleanup can always finish. Without it, a slow + // but connected client (full TCP buffer, no server WriteTimeout) could hang + // the handler forever. + streamWriteTimeout = 30 * time.Second ) func getScannerBufferSize() int { @@ -40,6 +45,16 @@ func NewStreamScanner(reader io.Reader) *bufio.Scanner { return scanner } +// ExtendWriteDeadline pushes the connection write deadline forward before each +// stream write. Best-effort: writers that don't support deadlines (e.g. +// httptest recorders) are silently ignored. +func ExtendWriteDeadline(c *gin.Context) { + if c == nil || c.Writer == nil { + return + } + _ = http.NewResponseController(c.Writer).SetWriteDeadline(time.Now().Add(streamWriteTimeout)) +} + func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo, dataHandler func(data string, sr *StreamResult)) { if resp == nil || dataHandler == nil { @@ -49,24 +64,27 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon // 无条件新建 StreamStatus info.StreamStatus = relaycommon.NewStreamStatus() - // 确保响应体总是被关闭 - defer func() { - if resp.Body != nil { - resp.Body.Close() - } - }() + ctx, cancel := context.WithCancel(context.Background()) streamingTimeout := time.Duration(constant.StreamingTimeout) * time.Second var ( - stopChan = make(chan bool, 3) // 增加缓冲区避免阻塞 - scanner = NewStreamScanner(resp.Body) - ticker = time.NewTicker(streamingTimeout) - pingTicker *time.Ticker - writeMutex sync.Mutex // Mutex to protect concurrent writes - wg sync.WaitGroup // 用于等待所有 goroutine 退出 + stopChan = make(chan bool, 3) // 增加缓冲区避免阻塞 + scanner = NewStreamScanner(resp.Body) + ticker = time.NewTicker(streamingTimeout) + pingTicker *time.Ticker + writeMutex sync.Mutex // Mutex to protect concurrent writes + wg sync.WaitGroup // 用于等待所有 goroutine 退出 + cleanupOnce sync.Once + stopOnce sync.Once ) + stop := func() { + stopOnce.Do(func() { + close(stopChan) + }) + } + generalSettings := operation_setting.GetGeneralSetting() pingEnabled := generalSettings.PingIntervalEnabled && !info.DisablePing pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second @@ -84,38 +102,28 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon logger.LogDebug(c, "streaming timeout seconds: %d", int64(streamingTimeout.Seconds())) logger.LogDebug(c, "ping interval seconds: %d", int64(pingInterval.Seconds())) - // 改进资源清理,确保所有 goroutine 正确退出 - defer func() { - // 通知所有 goroutine 停止 - common.SafeSendBool(stopChan, true) + cleanup := func() { + cleanupOnce.Do(func() { + cancel() + stop() + if resp.Body != nil { + _ = resp.Body.Close() + } - ticker.Stop() - if pingTicker != nil { - pingTicker.Stop() - } + ticker.Stop() + if pingTicker != nil { + pingTicker.Stop() + } - // 等待所有 goroutine 退出,最多等待5秒 - done := make(chan struct{}) - gopool.Go(func() { wg.Wait() - close(done) }) - - select { - case <-done: - case <-time.After(5 * time.Second): - logger.LogError(c, "timeout waiting for goroutines to exit") - } - - close(stopChan) - }() + } + // Ensure gin.Context is not returned to Gin's pool while any stream goroutine can still use it. + defer cleanup() scanner.Split(bufio.ScanLines) SetEventStreamHeaders(c) - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - ctx = context.WithValue(ctx, "stop_chan", stopChan) // Handle ping data sending with improved error handling @@ -123,13 +131,13 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon wg.Add(1) gopool.Go(func() { defer func() { - wg.Done() if r := recover(); r != nil { logger.LogError(c, fmt.Sprintf("ping goroutine panic: %v", r)) info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("ping panic: %v", r)) - common.SafeSendBool(stopChan, true) + stop() } logger.LogDebug(c, "ping goroutine exited") + wg.Done() }() // 添加超时保护,防止 goroutine 无限运行 @@ -140,31 +148,19 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon for { select { case <-pingTicker.C: - // 使用超时机制防止写操作阻塞 - done := make(chan error, 1) - gopool.Go(func() { + var err error + func() { writeMutex.Lock() defer writeMutex.Unlock() - done <- PingData(c) - }) - - select { - case err := <-done: - if err != nil { - logger.LogError(c, "ping data error: "+err.Error()) - info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, err) - return - } - logger.LogDebug(c, "ping data sent") - case <-time.After(10 * time.Second): - logger.LogError(c, "ping data send timeout") - info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, fmt.Errorf("ping send timeout")) - return - case <-ctx.Done(): - return - case <-stopChan: + ExtendWriteDeadline(c) + err = PingData(c) + }() + if err != nil { + logger.LogError(c, "ping data error: "+err.Error()) + info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPingFail, err) return } + logger.LogDebug(c, "ping data sent") case <-ctx.Done(): return case <-stopChan: @@ -185,19 +181,22 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon wg.Add(1) gopool.Go(func() { defer func() { - wg.Done() if r := recover(); r != nil { logger.LogError(c, fmt.Sprintf("data handler goroutine panic: %v", r)) info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("handler panic: %v", r)) } - common.SafeSendBool(stopChan, true) + stop() + wg.Done() }() sr := newStreamResult(info.StreamStatus) for data := range dataChan { sr.reset() - writeMutex.Lock() - dataHandler(data, sr) - writeMutex.Unlock() + func() { + writeMutex.Lock() + defer writeMutex.Unlock() + ExtendWriteDeadline(c) + dataHandler(data, sr) + }() if sr.IsStopped() { return } @@ -209,13 +208,13 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon common.RelayCtxGo(ctx, func() { defer func() { close(dataChan) - wg.Done() if r := recover(); r != nil { logger.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r)) info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonPanic, fmt.Errorf("scanner panic: %v", r)) } - common.SafeSendBool(stopChan, true) + stop() logger.LogDebug(c, "scanner goroutine exited") + wg.Done() }() for scanner.Scan() { @@ -225,9 +224,6 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon return case <-ctx.Done(): return - case <-c.Request.Context().Done(): - info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, c.Request.Context().Err()) - return default: } @@ -280,9 +276,12 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon case <-stopChan: // EndReason already set by the goroutine that triggered stopChan case <-c.Request.Context().Done(): + // 客户端断开:立即 cleanup 关闭上游 resp.Body,解除 scanner 阻塞并让上游停止生成, + // 避免为已放弃的请求继续消费上游 token。 info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, c.Request.Context().Err()) } + cleanup() if info.StreamStatus.IsNormalEnd() && !info.StreamStatus.HasErrors() { logger.LogInfo(c, fmt.Sprintf("stream ended: %s", info.StreamStatus.Summary())) } else { diff --git a/relay/helper/stream_scanner_test.go b/relay/helper/stream_scanner_test.go index 637b750f..a2119512 100644 --- a/relay/helper/stream_scanner_test.go +++ b/relay/helper/stream_scanner_test.go @@ -2,6 +2,7 @@ package helper import ( "bufio" + "context" "fmt" "io" "net/http" @@ -210,6 +211,78 @@ func TestStreamScannerHandler_DataWithExtraSpaces(t *testing.T) { assert.Equal(t, "{\"trimmed\":true}", got) } +// TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns pins the +// disconnect contract: when the client goes away, the handler must return +// promptly (all goroutines joined, so the gin.Context can never leak into a +// pooled reuse), the upstream body must be closed to stop token generation, +// and no data received after the disconnect may be processed or written. +func TestStreamScannerHandler_ClientCancelAbortsUpstreamAndReturns(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + pr, pw := io.Pipe() + t.Cleanup(func() { + _ = pr.Close() + _ = pw.Close() + }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil).WithContext(ctx) + + resp := &http.Response{Body: pr} + info := &relaycommon.RelayInfo{ + DisablePing: true, + ChannelMeta: &relaycommon.ChannelMeta{}, + } + + var count atomic.Int64 + firstHandled := make(chan struct{}) + done := make(chan struct{}) + go func() { + StreamScannerHandler(c, resp, info, func(data string, sr *StreamResult) { + count.Add(1) + _ = StringData(c, data) + if data == "first" { + close(firstHandled) + } + }) + close(done) + }() + + _, err := fmt.Fprint(pw, "data: first\n") + require.NoError(t, err) + + select { + case <-firstHandled: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for first chunk") + } + + cancel() + + // The handler must return without any further upstream input: cleanup + // closes resp.Body, which unblocks the scanner goroutine. + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("handler did not return after client disconnect") + } + + // Upstream read side must be closed so the provider stops generating + // (and billing) for a request nobody is listening to. + _, err = fmt.Fprint(pw, "data: second\n") + require.ErrorIs(t, err, io.ErrClosedPipe, "upstream body should be closed after client disconnect") + + assert.Equal(t, int64(1), count.Load(), "no chunk after disconnect should be processed") + require.NotNil(t, info.StreamStatus) + assert.Equal(t, relaycommon.StreamEndReasonClientGone, info.StreamStatus.EndReason) + + body := recorder.Body.String() + assert.Contains(t, body, "first") + assert.NotContains(t, body, "second") +} + // ---------- Ping tests ---------- func TestStreamScannerHandler_PingSentDuringSlowUpstream(t *testing.T) {