feat: 算力引擎 Go 后端核心（new-api）
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包：relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
@@ -0,0 +1,310 @@
|
||||
package helper
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
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 {
|
||||
if constant.StreamScannerMaxBufferMB > 0 {
|
||||
return constant.StreamScannerMaxBufferMB << 20
|
||||
}
|
||||
return DefaultMaxScannerBufferSize
|
||||
}
|
||||
|
||||
func NewStreamScanner(reader io.Reader) *bufio.Scanner {
|
||||
scanner := bufio.NewScanner(reader)
|
||||
scanner.Buffer(make([]byte, InitialScannerBufferSize), getScannerBufferSize())
|
||||
return scanner
|
||||
}
|
||||
|
||||
func copyCodexSSEHeaders(c *gin.Context, resp *http.Response) {
|
||||
if c == nil || c.Writer == nil || resp == nil {
|
||||
return
|
||||
}
|
||||
// codex
|
||||
for _, name := range []string{"X-Reasoning-Included", "X-Codex-Turn-State"} {
|
||||
values := resp.Header.Values(name)
|
||||
if !service.ShouldCopyUpstreamHeader(c, name, values) {
|
||||
continue
|
||||
}
|
||||
for _, value := range values {
|
||||
if value != "" {
|
||||
c.Writer.Header().Add(name, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
return
|
||||
}
|
||||
|
||||
// 无条件新建 StreamStatus
|
||||
info.StreamStatus = relaycommon.NewStreamStatus()
|
||||
|
||||
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 退出
|
||||
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
|
||||
if pingInterval <= 0 {
|
||||
pingInterval = DefaultPingInterval
|
||||
}
|
||||
|
||||
if pingEnabled {
|
||||
pingTicker = time.NewTicker(pingInterval)
|
||||
}
|
||||
|
||||
logger.LogDebug(c, "relay timeout seconds: %d", common.RelayTimeout)
|
||||
logger.LogDebug(c, "relay max idle conns: %d", common.RelayMaxIdleConns)
|
||||
logger.LogDebug(c, "relay max idle conns per host: %d", common.RelayMaxIdleConnsPerHost)
|
||||
logger.LogDebug(c, "streaming timeout seconds: %d", int64(streamingTimeout.Seconds()))
|
||||
logger.LogDebug(c, "ping interval seconds: %d", int64(pingInterval.Seconds()))
|
||||
|
||||
cleanup := func() {
|
||||
cleanupOnce.Do(func() {
|
||||
cancel()
|
||||
stop()
|
||||
if resp.Body != nil {
|
||||
_ = resp.Body.Close()
|
||||
}
|
||||
|
||||
ticker.Stop()
|
||||
if pingTicker != nil {
|
||||
pingTicker.Stop()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
})
|
||||
}
|
||||
// Ensure gin.Context is not returned to Gin's pool while any stream goroutine can still use it.
|
||||
defer cleanup()
|
||||
|
||||
scanner.Split(bufio.ScanLines)
|
||||
copyCodexSSEHeaders(c, resp)
|
||||
SetEventStreamHeaders(c)
|
||||
|
||||
ctx = context.WithValue(ctx, "stop_chan", stopChan)
|
||||
|
||||
// Handle ping data sending with improved error handling
|
||||
if pingEnabled && pingTicker != nil {
|
||||
wg.Add(1)
|
||||
gopool.Go(func() {
|
||||
defer func() {
|
||||
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))
|
||||
stop()
|
||||
}
|
||||
logger.LogDebug(c, "ping goroutine exited")
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// 添加超时保护，防止 goroutine 无限运行
|
||||
maxPingDuration := 30 * time.Minute // 最大 ping 持续时间
|
||||
pingTimeout := time.NewTimer(maxPingDuration)
|
||||
defer pingTimeout.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-pingTicker.C:
|
||||
var err error
|
||||
func() {
|
||||
writeMutex.Lock()
|
||||
defer writeMutex.Unlock()
|
||||
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:
|
||||
return
|
||||
case <-c.Request.Context().Done():
|
||||
// 监听客户端断开连接
|
||||
return
|
||||
case <-pingTimeout.C:
|
||||
logger.LogError(c, "ping goroutine max duration reached")
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
dataChan := make(chan string, 10)
|
||||
|
||||
wg.Add(1)
|
||||
gopool.Go(func() {
|
||||
defer func() {
|
||||
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))
|
||||
}
|
||||
stop()
|
||||
wg.Done()
|
||||
}()
|
||||
sr := newStreamResult(info.StreamStatus)
|
||||
for data := range dataChan {
|
||||
sr.reset()
|
||||
func() {
|
||||
writeMutex.Lock()
|
||||
defer writeMutex.Unlock()
|
||||
ExtendWriteDeadline(c)
|
||||
dataHandler(data, sr)
|
||||
}()
|
||||
if sr.IsStopped() {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// Scanner goroutine with improved error handling
|
||||
wg.Add(1)
|
||||
common.RelayCtxGo(ctx, func() {
|
||||
defer func() {
|
||||
close(dataChan)
|
||||
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))
|
||||
}
|
||||
stop()
|
||||
logger.LogDebug(c, "scanner goroutine exited")
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
for scanner.Scan() {
|
||||
// 检查是否需要停止
|
||||
select {
|
||||
case <-stopChan:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
ticker.Reset(streamingTimeout)
|
||||
data := scanner.Text()
|
||||
logger.LogDebug(c, "stream scanner data: %s", data)
|
||||
|
||||
if len(data) < 6 {
|
||||
continue
|
||||
}
|
||||
if data[:5] != "data:" && data[:6] != "[DONE]" {
|
||||
continue
|
||||
}
|
||||
data = data[5:]
|
||||
data = strings.TrimSpace(data)
|
||||
if data == "" {
|
||||
continue
|
||||
}
|
||||
if !strings.HasPrefix(data, "[DONE]") {
|
||||
info.SetFirstResponseTime()
|
||||
info.ReceivedResponseCount++
|
||||
|
||||
select {
|
||||
case dataChan <- data:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-stopChan:
|
||||
return
|
||||
}
|
||||
} else {
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonDone, nil)
|
||||
logger.LogDebug(c, "received [DONE], stopping scanner")
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := scanner.Err(); err != nil {
|
||||
if err != io.EOF {
|
||||
logger.LogError(c, "scanner error: "+err.Error())
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonScannerErr, err)
|
||||
}
|
||||
}
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonEOF, nil)
|
||||
})
|
||||
|
||||
// 主循环等待完成或超时
|
||||
select {
|
||||
case <-ticker.C:
|
||||
info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonTimeout, nil)
|
||||
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 {
|
||||
logger.LogError(c, fmt.Sprintf("stream ended: %s, received=%d", info.StreamStatus.Summary(), info.ReceivedResponseCount))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user