fix: 流式请求ping

This commit is contained in:
creamlike1024
2025-05-31 22:13:17 +08:00
parent 9c4d3a6359
commit c51a30b862

View File

@@ -104,6 +104,65 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
return targetConn, nil
}
func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.CancelFunc {
pingerCtx, stopPinger := context.WithCancel(context.Background())
gopool.Go(func() {
defer func() {
if common2.DebugEnabled {
println("SSE ping goroutine stopped.")
}
}()
if pingInterval <= 0 {
pingInterval = helper.DefaultPingInterval
}
ticker := time.NewTicker(pingInterval)
// 退出时清理 ticker
defer ticker.Stop()
var pingMutex sync.Mutex
if common2.DebugEnabled {
println("SSE ping goroutine started")
}
for {
select {
// 发送 ping 数据
case <-ticker.C:
if err := sendPingData(c, &pingMutex); err != nil {
return
}
// 收到退出信号
case <-pingerCtx.Done():
return
// request 结束
case <-c.Request.Context().Done():
return
}
}
})
return stopPinger
}
func sendPingData(c *gin.Context, mutex *sync.Mutex) error {
mutex.Lock()
defer mutex.Unlock()
err := helper.PingData(c)
if err != nil {
common2.LogError(c, "SSE ping error: "+err.Error())
return err
}
if common2.DebugEnabled {
println("SSE ping data sent.")
}
return nil
}
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
var client *http.Client
var err error
@@ -115,69 +174,28 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
} else {
client = service.GetHttpClient()
}
// 流式请求 ping 保活
var stopPinger func()
generalSettings := operation_setting.GetGeneralSetting()
pingEnabled := generalSettings.PingIntervalEnabled
var pingerWg sync.WaitGroup
if info.IsStream {
helper.SetEventStreamHeaders(c)
if pingEnabled {
// 处理流式请求的 ping 保活
generalSettings := operation_setting.GetGeneralSetting()
if generalSettings.PingIntervalEnabled {
pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
var pingerCtx context.Context
pingerCtx, stopPinger = context.WithCancel(c.Request.Context())
// 退出时清理 pingerCtx 防止泄露
stopPinger := startPingKeepAlive(c, pingInterval)
defer stopPinger()
pingerWg.Add(1)
gopool.Go(func() {
defer pingerWg.Done()
if pingInterval <= 0 {
pingInterval = helper.DefaultPingInterval
}
ticker := time.NewTicker(pingInterval)
defer ticker.Stop()
var pingMutex sync.Mutex
if common2.DebugEnabled {
println("SSE ping goroutine started")
}
for {
select {
case <-ticker.C:
pingMutex.Lock()
err2 := helper.PingData(c)
pingMutex.Unlock()
if err2 != nil {
common2.LogError(c, "SSE ping error: "+err.Error())
return
}
if common2.DebugEnabled {
println("SSE ping data sent.")
}
case <-pingerCtx.Done():
if common2.DebugEnabled {
println("SSE ping goroutine stopped.")
}
return
}
}
})
}
}
resp, err := client.Do(req)
// request结束后等待 ping goroutine 完成
if info.IsStream && pingEnabled {
pingerWg.Wait()
}
if err != nil {
return nil, err
}
if resp == nil {
return nil, errors.New("resp is nil")
}
_ = req.Body.Close()
_ = c.Request.Body.Close()
return resp, nil