Skip to content

Commit 176f764

Browse files
committed
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.
1 parent 89c0b79 commit 176f764

File tree

2 files changed

+544
-14
lines changed

2 files changed

+544
-14
lines changed

relay/helper/stream_scanner.go

Lines changed: 23 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -176,10 +176,32 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
176176
})
177177
}
178178

179+
dataChan := make(chan string, 10)
180+
181+
wg.Add(1)
182+
gopool.Go(func() {
183+
defer func() {
184+
wg.Done()
185+
if r := recover(); r != nil {
186+
logger.LogError(c, fmt.Sprintf("data handler goroutine panic: %v", r))
187+
}
188+
common.SafeSendBool(stopChan, true)
189+
}()
190+
for data := range dataChan {
191+
writeMutex.Lock()
192+
success := dataHandler(data)
193+
writeMutex.Unlock()
194+
if !success {
195+
return
196+
}
197+
}
198+
})
199+
179200
// Scanner goroutine with improved error handling
180201
wg.Add(1)
181202
common.RelayCtxGo(ctx, func() {
182203
defer func() {
204+
close(dataChan)
183205
wg.Done()
184206
if r := recover(); r != nil {
185207
logger.LogError(c, fmt.Sprintf("scanner goroutine panic: %v", r))
@@ -222,22 +244,9 @@ func StreamScannerHandler(c *gin.Context, resp *http.Response, info *relaycommon
222244
if !strings.HasPrefix(data, "[DONE]") {
223245
info.SetFirstResponseTime()
224246
info.ReceivedResponseCount++
225-
// 使用超时机制防止写操作阻塞
226-
done := make(chan bool, 1)
227-
gopool.Go(func() {
228-
writeMutex.Lock()
229-
defer writeMutex.Unlock()
230-
done <- dataHandler(data)
231-
})
232247

233248
select {
234-
case success := <-done:
235-
if !success {
236-
return
237-
}
238-
case <-time.After(10 * time.Second):
239-
logger.LogError(c, "data handler timeout")
240-
return
249+
case dataChan <- data:
241250
case <-ctx.Done():
242251
return
243252
case <-stopChan:

0 commit comments

Comments
 (0)