fix: fix stop signal calling stack

fix: fix stop signal calling stack
Co-Authored-By: Minghan Zhang <112773885+zmh-program@users.noreply.github.com>
This commit is contained in:
Deng Junhai 2024-06-28 19:51:32 +08:00
parent b0a684cada
commit 3dffa7fdab

View File

@ -10,6 +10,7 @@ import (
"chat/globals" "chat/globals"
"chat/manager/conversation" "chat/manager/conversation"
"chat/utils" "chat/utils"
"time"
"database/sql" "database/sql"
"errors" "errors"
@ -51,7 +52,9 @@ type partialChunk struct {
func createStopSignal(conn *Connection) chan bool { func createStopSignal(conn *Connection) chan bool {
stopSignal := make(chan bool, 1) stopSignal := make(chan bool, 1)
go func(conn *Connection, stopSignal chan bool) { go func(conn *Connection, stopSignal chan bool) {
ticker := time.NewTicker(100 * time.Millisecond)
defer func() { defer func() {
ticker.Stop()
if r := recover(); r != nil && !strings.Contains(fmt.Sprintf("%s", r), "closed channel") { if r := recover(); r != nil && !strings.Contains(fmt.Sprintf("%s", r), "closed channel") {
stack := debug.Stack() stack := debug.Stack()
globals.Warn(fmt.Sprintf("caught panic from stop signal: %s\n%s", r, stack)) globals.Warn(fmt.Sprintf("caught panic from stop signal: %s\n%s", r, stack))
@ -59,11 +62,16 @@ func createStopSignal(conn *Connection) chan bool {
}() }()
for { for {
if conn.PeekStop() != nil { select {
stopSignal <- true case <-ticker.C:
state := conn.PeekStop() != nil // check the stop state
stopSignal <- state
if state {
break break
} }
} }
}
}(conn, stopSignal) }(conn, stopSignal)
return stopSignal return stopSignal
@ -169,7 +177,9 @@ func createChatTask(
return hit, sendPackError return hit, sendPackError
} }
case <-stopSignal: case signal := <-stopSignal:
// if stop signal is received
if signal {
globals.Info(fmt.Sprintf("client stopped the chat request (model: %s, client: %s)", model, conn.GetCtx().ClientIP())) globals.Info(fmt.Sprintf("client stopped the chat request (model: %s, client: %s)", model, conn.GetCtx().ClientIP()))
_ = conn.SendClient(globals.ChatSegmentResponse{ _ = conn.SendClient(globals.ChatSegmentResponse{
Quota: buffer.GetQuota(), Quota: buffer.GetQuota(),
@ -181,6 +191,7 @@ func createChatTask(
return return
} }
} }
}
} }
func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conversation, restart bool) string { func ChatHandler(conn *Connection, user *auth.User, instance *conversation.Conversation, restart bool) string {