mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-04-30 15:21:45 +00:00
feat: add comprehensive tests for StreamScannerHandler functionality
- Introduced a new test file for StreamScannerHandler, covering various scenarios including nil inputs, empty bodies, chunk processing, order preservation, and handler failures. - Enhanced error handling and data processing logic in StreamScannerHandler to improve robustness and performance.
This commit is contained in:
@@ -176,10 +176,32 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
|
||||
})
|
||||
}
|
||||
|
||||
dataChan := make(chan string, 10)
|
||||
|
||||
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))
|
||||
}
|
||||
common.SafeSendBool(stopChan, true)
|
||||
}()
|
||||
for data := range dataChan {
|
||||
writeMutex.Lock()
|
||||
success := dataHandler(data)
|
||||
writeMutex.Unlock()
|
||||
if !success {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// Scanner goroutine with improved error handling
|
||||
wg.Add(1)
|
||||
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))
|
||||
@@ -222,22 +244,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
|
||||
if !strings.HasPrefix(data, "[DONE]") {
|
||||
info.SetFirstResponseTime()
|
||||
info.ReceivedResponseCount++
|
||||
// 使用超时机制防止写操作阻塞
|
||||
done := make(chan bool, 1)
|
||||
gopool.Go(func() {
|
||||
writeMutex.Lock()
|
||||
defer writeMutex.Unlock()
|
||||
done <- dataHandler(data)
|
||||
})
|
||||
|
||||
select {
|
||||
case success := <-done:
|
||||
if !success {
|
||||
return
|
||||
}
|
||||
case <-time.After(10 * time.Second):
|
||||
logger.LogError(c, "data handler timeout")
|
||||
return
|
||||
case dataChan <- data:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-stopChan:
|
||||
|
||||
Reference in New Issue
Block a user