xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-17 16:06:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: safe send channel

Browse Source
This commit is contained in:
CaIon 2024-04-28 16:17:16 +08:00
parent 1c6fd87909
commit 21f32605c8
2 changed files with 18 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
common/go-channel.go
View File

@ -16,7 +16,22 @@ func SafeGoroutine(f func()) {
}() }()
} }
func SafeSend(ch chan bool, value bool) (closed bool) { func SafeSendBool(ch chan bool, value bool) (closed bool) {
defer func() {
// Recover from panic if one occured. A panic would mean the channel was closed.
if recover() != nil {
closed = true
}
}()
// This will panic if the channel is closed.
ch <- value
// If the code reaches here, then the channel was not closed.
return false
}
func SafeSendString(ch chan string, value string) (closed bool) {
defer func() { defer func() {
// Recover from panic if one occured. A panic would mean the channel was closed. // Recover from panic if one occured. A panic would mean the channel was closed.
if recover() != nil { if recover() != nil {

4
relay/channel/openai/relay-openai.go
View File

@ -50,7 +50,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
if data[:6] != "data: " && data[:6] != "[DONE]" { if data[:6] != "data: " && data[:6] != "[DONE]" {
continue continue
} }
dataChan <- data common.SafeSendString(dataChan, data)
data = data[6:] data = data[6:]
if !strings.HasPrefix(data, "[DONE]") { if !strings.HasPrefix(data, "[DONE]") {
streamItems = append(streamItems, data) streamItems = append(streamItems, data)
@ -123,7 +123,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
// wait data out // wait data out
time.Sleep(2 * time.Second) time.Sleep(2 * time.Second)
} }
common.SafeSend(stopChan, true) common.SafeSendBool(stopChan, true)
}() }()
service.SetEventStreamHeaders(c) service.SetEventStreamHeaders(c)
c.Stream(func(w io.Writer) bool { c.Stream(func(w io.Writer) bool {