oadmin/sub2api
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

fix(流式): 以上游读取判定超时并调大事件缓冲

Browse Source 
- 以读取时间戳判定流式间隔超时，避免下游阻塞误判
- antigravity 流式读取使用 MaxLineSize 配置
- 事件通道缓冲提升到 16

测试: go test ./...
This commit is contained in:
yangjianbo
2026-01-04 20:19:07 +08:00
parent 73ffb58518
commit 7489da49cb
3 changed files with 71 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
backend/internal/service/antigravity_gateway_service.go
View File

@@ -11,6 +11,7 @@ import (
"log" "log"
"net/http" "net/http"
"strings" "strings"
"sync/atomic"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/config"
@@ -680,7 +681,11 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
// 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), defaultMaxLineSize) maxLineSize := defaultMaxLineSize
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
maxLineSize = s.cfg.Gateway.MaxLineSize
}
scanner.Buffer(make([]byte, 64*1024), maxLineSize)
usage := &ClaudeUsage{} usage := &ClaudeUsage{}
var firstTokenMs *int var firstTokenMs *int
@@ -689,7 +694,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
err error err error
} }
// 独立 goroutine 读取上游，避免读取阻塞影响超时处理 // 独立 goroutine 读取上游，避免读取阻塞影响超时处理
events := make(chan scanEvent, 1) events := make(chan scanEvent, 16)
done := make(chan struct{}) done := make(chan struct{})
sendEvent := func(ev scanEvent) bool { sendEvent := func(ev scanEvent) bool {
select { select {
@@ -699,9 +704,12 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
return false return false
} }
} }
var lastReadAt int64
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
go func() { go func() {
defer close(events) defer close(events)
for scanner.Scan() { for scanner.Scan() {
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
if !sendEvent(scanEvent{line: scanner.Text()}) { if !sendEvent(scanEvent{line: scanner.Text()}) {
return return
} }
@@ -717,26 +725,14 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 {
streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
} }
var intervalTimer *time.Timer var intervalTicker *time.Ticker
if streamInterval > 0 { if streamInterval > 0 {
intervalTimer = time.NewTimer(streamInterval) intervalTicker = time.NewTicker(streamInterval)
defer intervalTimer.Stop() defer intervalTicker.Stop()
} }
var intervalCh <-chan time.Time var intervalCh <-chan time.Time
if intervalTimer != nil { if intervalTicker != nil {
intervalCh = intervalTimer.C intervalCh = intervalTicker.C
}
resetInterval := func() {
if intervalTimer == nil {
return
}
if !intervalTimer.Stop() {
select {
case <-intervalTimer.C:
default:
}
}
intervalTimer.Reset(streamInterval)
} }
// 仅发送一次错误事件，避免多次写入导致协议混乱 // 仅发送一次错误事件，避免多次写入导致协议混乱
@@ -758,7 +754,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
} }
if ev.err != nil { if ev.err != nil {
if errors.Is(ev.err, bufio.ErrTooLong) { if errors.Is(ev.err, bufio.ErrTooLong) {
log.Printf("SSE line too long (antigravity): max_size=%d error=%v", defaultMaxLineSize, ev.err) log.Printf("SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err)
sendErrorEvent("response_too_large") sendErrorEvent("response_too_large")
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err
} }
@@ -766,7 +762,6 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
return nil, ev.err return nil, ev.err
} }
resetInterval()
line := ev.line line := ev.line
trimmed := strings.TrimRight(line, "\r\n") trimmed := strings.TrimRight(line, "\r\n")
if strings.HasPrefix(trimmed, "data:") { if strings.HasPrefix(trimmed, "data:") {
@@ -814,6 +809,10 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
flusher.Flush() flusher.Flush()
case <-intervalCh: case <-intervalCh:
lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
if time.Since(lastRead) < streamInterval {
continue
}
log.Printf("Stream data interval timeout (antigravity)") log.Printf("Stream data interval timeout (antigravity)")
sendErrorEvent("stream_timeout") sendErrorEvent("stream_timeout")
return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
@@ -959,7 +958,11 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
var firstTokenMs *int var firstTokenMs *int
// 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), defaultMaxLineSize) maxLineSize := defaultMaxLineSize
if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
maxLineSize = s.cfg.Gateway.MaxLineSize
}
scanner.Buffer(make([]byte, 64*1024), maxLineSize)
// 辅助函数：转换 antigravity.ClaudeUsage 到 service.ClaudeUsage // 辅助函数：转换 antigravity.ClaudeUsage 到 service.ClaudeUsage
convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage { convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage {
@@ -979,7 +982,7 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
err error err error
} }
// 独立 goroutine 读取上游，避免读取阻塞影响超时处理 // 独立 goroutine 读取上游，避免读取阻塞影响超时处理
events := make(chan scanEvent, 1) events := make(chan scanEvent, 16)
done := make(chan struct{}) done := make(chan struct{})
sendEvent := func(ev scanEvent) bool { sendEvent := func(ev scanEvent) bool {
select { select {
@@ -989,9 +992,12 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
return false return false
} }
} }
var lastReadAt int64
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
go func() { go func() {
defer close(events) defer close(events)
for scanner.Scan() { for scanner.Scan() {
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
if !sendEvent(scanEvent{line: scanner.Text()}) { if !sendEvent(scanEvent{line: scanner.Text()}) {
return return
} }
@@ -1006,26 +1012,14 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 {
streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
} }
var intervalTimer *time.Timer var intervalTicker *time.Ticker
if streamInterval > 0 { if streamInterval > 0 {
intervalTimer = time.NewTimer(streamInterval) intervalTicker = time.NewTicker(streamInterval)
defer intervalTimer.Stop() defer intervalTicker.Stop()
} }
var intervalCh <-chan time.Time var intervalCh <-chan time.Time
if intervalTimer != nil { if intervalTicker != nil {
intervalCh = intervalTimer.C intervalCh = intervalTicker.C
}
resetInterval := func() {
if intervalTimer == nil {
return
}
if !intervalTimer.Stop() {
select {
case <-intervalTimer.C:
default:
}
}
intervalTimer.Reset(streamInterval)
} }
// 仅发送一次错误事件，避免多次写入导致协议混乱 // 仅发送一次错误事件，避免多次写入导致协议混乱
@@ -1053,7 +1047,7 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
} }
if ev.err != nil { if ev.err != nil {
if errors.Is(ev.err, bufio.ErrTooLong) { if errors.Is(ev.err, bufio.ErrTooLong) {
log.Printf("SSE line too long (antigravity): max_size=%d error=%v", defaultMaxLineSize, ev.err) log.Printf("SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err)
sendErrorEvent("response_too_large") sendErrorEvent("response_too_large")
return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err
} }
@@ -1061,7 +1055,6 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
return nil, fmt.Errorf("stream read error: %w", ev.err) return nil, fmt.Errorf("stream read error: %w", ev.err)
} }
resetInterval()
line := ev.line line := ev.line
// 处理 SSE 行，转换为 Claude 格式 // 处理 SSE 行，转换为 Claude 格式
claudeEvents := processor.ProcessLine(strings.TrimRight(line, "\r\n")) claudeEvents := processor.ProcessLine(strings.TrimRight(line, "\r\n"))
@@ -1084,6 +1077,10 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
} }
case <-intervalCh: case <-intervalCh:
lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
if time.Since(lastRead) < streamInterval {
continue
}
log.Printf("Stream data interval timeout (antigravity)") log.Printf("Stream data interval timeout (antigravity)")
sendErrorEvent("stream_timeout") sendErrorEvent("stream_timeout")
return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")

40
backend/internal/service/gateway_service.go
View File

@@ -15,6 +15,7 @@ import (
"regexp" "regexp"
"sort" "sort"
"strings" "strings"
"sync/atomic"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/config"
@@ -1460,7 +1461,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
err error err error
} }
// 独立 goroutine 读取上游，避免读取阻塞导致超时/keepalive无法处理 // 独立 goroutine 读取上游，避免读取阻塞导致超时/keepalive无法处理
events := make(chan scanEvent, 1) events := make(chan scanEvent, 16)
done := make(chan struct{}) done := make(chan struct{})
sendEvent := func(ev scanEvent) bool { sendEvent := func(ev scanEvent) bool {
select { select {
@@ -1470,9 +1471,12 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
return false return false
} }
} }
var lastReadAt int64
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
go func() { go func() {
defer close(events) defer close(events)
for scanner.Scan() { for scanner.Scan() {
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
if !sendEvent(scanEvent{line: scanner.Text()}) { if !sendEvent(scanEvent{line: scanner.Text()}) {
return return
} }
@@ -1487,11 +1491,15 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 {
streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
} }
// 仅监控上游数据间隔超时，避免上游挂起占用资源 // 仅监控上游数据间隔超时，避免下游写入阻塞导致误判
var intervalTimer *time.Timer var intervalTicker *time.Ticker
if streamInterval > 0 { if streamInterval > 0 {
intervalTimer = time.NewTimer(streamInterval) intervalTicker = time.NewTicker(streamInterval)
defer intervalTimer.Stop() defer intervalTicker.Stop()
}
var intervalCh <-chan time.Time
if intervalTicker != nil {
intervalCh = intervalTicker.C
} }
// 仅发送一次错误事件，避免多次写入导致协议混乱（写失败时尽力通知客户端） // 仅发送一次错误事件，避免多次写入导致协议混乱（写失败时尽力通知客户端）
@@ -1523,9 +1531,6 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream read error: %w", ev.err) return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream read error: %w", ev.err)
} }
line := ev.line line := ev.line
if intervalTimer != nil {
resetTimer(intervalTimer, streamInterval)
}
if line == "event: error" { if line == "event: error" {
return nil, errors.New("have error in stream") return nil, errors.New("have error in stream")
} }
@@ -1561,12 +1566,11 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
flusher.Flush() flusher.Flush()
} }
case <-func() <-chan time.Time { case <-intervalCh:
if intervalTimer != nil { lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
return intervalTimer.C if time.Since(lastRead) < streamInterval {
continue
} }
return nil
}():
log.Printf("Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) log.Printf("Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval)
sendErrorEvent("stream_timeout") sendErrorEvent("stream_timeout")
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")
@@ -1576,16 +1580,6 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, nil return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, nil
} }
func resetTimer(timer *time.Timer, interval time.Duration) {
if !timer.Stop() {
select {
case <-timer.C:
default:
}
}
timer.Reset(interval)
}
// replaceModelInSSELine 替换SSE数据行中的model字段 // replaceModelInSSELine 替换SSE数据行中的model字段
func (s *GatewayService) replaceModelInSSELine(line, fromModel, toModel string) string { func (s *GatewayService) replaceModelInSSELine(line, fromModel, toModel string) string {
if !sseDataRe.MatchString(line) { if !sseDataRe.MatchString(line) {

25
backend/internal/service/openai_gateway_service.go
View File

@@ -16,6 +16,7 @@ import (
"sort" "sort"
"strconv" "strconv"
"strings" "strings"
"sync/atomic"
"time" "time"
"github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/config"
@@ -786,7 +787,7 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
err error err error
} }
// 独立 goroutine 读取上游，避免读取阻塞影响 keepalive/超时处理 // 独立 goroutine 读取上游，避免读取阻塞影响 keepalive/超时处理
events := make(chan scanEvent, 1) events := make(chan scanEvent, 16)
done := make(chan struct{}) done := make(chan struct{})
sendEvent := func(ev scanEvent) bool { sendEvent := func(ev scanEvent) bool {
select { select {
@@ -796,9 +797,12 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
return false return false
} }
} }
var lastReadAt int64
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
go func() { go func() {
defer close(events) defer close(events)
for scanner.Scan() { for scanner.Scan() {
atomic.StoreInt64(&lastReadAt, time.Now().UnixNano())
if !sendEvent(scanEvent{line: scanner.Text()}) { if !sendEvent(scanEvent{line: scanner.Text()}) {
return return
} }
@@ -813,15 +817,15 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { if s.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 {
streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second streamInterval = time.Duration(s.cfg.Gateway.StreamDataIntervalTimeout) * time.Second
} }
// 仅监控上游数据间隔超时，不被下游 keepalive 影响 // 仅监控上游数据间隔超时，不被下游写入阻塞影响
var intervalTimer *time.Timer var intervalTicker *time.Ticker
if streamInterval > 0 { if streamInterval > 0 {
intervalTimer = time.NewTimer(streamInterval) intervalTicker = time.NewTicker(streamInterval)
defer intervalTimer.Stop() defer intervalTicker.Stop()
} }
var intervalCh <-chan time.Time var intervalCh <-chan time.Time
if intervalTimer != nil { if intervalTicker != nil {
intervalCh = intervalTimer.C intervalCh = intervalTicker.C
} }
keepaliveInterval := time.Duration(0) keepaliveInterval := time.Duration(0)
@@ -872,9 +876,6 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
line := ev.line line := ev.line
lastDataAt = time.Now() lastDataAt = time.Now()
if intervalTimer != nil {
resetTimer(intervalTimer, streamInterval)
}
// Extract data from SSE line (supports both "data: " and "data:" formats) // Extract data from SSE line (supports both "data: " and "data:" formats)
if openaiSSEDataRe.MatchString(line) { if openaiSSEDataRe.MatchString(line) {
@@ -908,6 +909,10 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
} }
case <-intervalCh: case <-intervalCh:
lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt))
if time.Since(lastRead) < streamInterval {
continue
}
log.Printf("Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) log.Printf("Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval)
sendErrorEvent("stream_timeout") sendErrorEvent("stream_timeout")
return &openaiStreamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") return &openaiStreamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout")