Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 修复stream handler中因为 c.Stream 提前退出导致 dataChan 阻塞

Browse Source
This commit is contained in:
Xiangyuan Liu 2023-11-28 11:38:07 +08:00
parent 6d2210ad5c
commit 69a1dccf5f
No known key found for this signature in database
GPG Key ID: A2FCA5A37FC0F888
1 changed files with 6 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
controller/relay-openai.go
View File

@ -27,11 +27,15 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
} }
return 0, nil, nil return 0, nil, nil
}) })
dataChan := make(chan string) dataChan := make(chan string, 5)
stopChan := make(chan bool) stopChan := make(chan bool, 2)
defer close(stopChan)
defer close(dataChan)
var wg sync.WaitGroup var wg sync.WaitGroup
go func() { go func() {
wg.Add(1) wg.Add(1)
defer wg.Done()
var streamItems []string var streamItems []string
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
@ -74,7 +78,6 @@ func openaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*O
} }
} }
} }
wg.Done()
stopChan <- true stopChan <- true
}() }()
setEventStreamHeaders(c) setEventStreamHeaders(c)