Commit 2da94e62 by creamlike1024

fix: 流式请求ping

parent 4febdff3
...@@ -104,40 +104,24 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody ...@@ -104,40 +104,24 @@ func DoWssRequest(a Adaptor, c *gin.Context, info *common.RelayInfo, requestBody
return targetConn, nil return targetConn, nil
} }
func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) { func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.CancelFunc {
var client *http.Client pingerCtx, stopPinger := context.WithCancel(context.Background())
var err error
if proxyURL, ok := info.ChannelSetting["proxy"]; ok {
client, err = service.NewProxyHttpClient(proxyURL.(string))
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
}
} 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 {
pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
var pingerCtx context.Context
pingerCtx, stopPinger = context.WithCancel(c.Request.Context())
// 退出时清理 pingerCtx 防止泄露
defer stopPinger()
pingerWg.Add(1)
gopool.Go(func() { gopool.Go(func() {
defer pingerWg.Done() defer func() {
if common2.DebugEnabled {
println("SSE ping goroutine stopped.")
}
}()
if pingInterval <= 0 { if pingInterval <= 0 {
pingInterval = helper.DefaultPingInterval pingInterval = helper.DefaultPingInterval
} }
ticker := time.NewTicker(pingInterval) ticker := time.NewTicker(pingInterval)
// 退出时清理 ticker
defer ticker.Stop() defer ticker.Stop()
var pingMutex sync.Mutex var pingMutex sync.Mutex
if common2.DebugEnabled { if common2.DebugEnabled {
println("SSE ping goroutine started") println("SSE ping goroutine started")
...@@ -145,39 +129,73 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http ...@@ -145,39 +129,73 @@ func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http
for { for {
select { select {
// 发送 ping 数据
case <-ticker.C: case <-ticker.C:
pingMutex.Lock() if err := sendPingData(c, &pingMutex); err != nil {
err2 := helper.PingData(c)
pingMutex.Unlock()
if err2 != nil {
common2.LogError(c, "SSE ping error: "+err.Error())
return return
} }
if common2.DebugEnabled { // 收到退出信号
println("SSE ping data sent.")
}
case <-pingerCtx.Done(): case <-pingerCtx.Done():
if common2.DebugEnabled { return
println("SSE ping goroutine stopped.") // request 结束
} case <-c.Request.Context().Done():
return 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
}
resp, err := client.Do(req) func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
// request结束后等待 ping goroutine 完成 var client *http.Client
if info.IsStream && pingEnabled { var err error
pingerWg.Wait() if proxyURL, ok := info.ChannelSetting["proxy"]; ok {
client, err = service.NewProxyHttpClient(proxyURL.(string))
if err != nil {
return nil, fmt.Errorf("new proxy http client failed: %w", err)
} }
} else {
client = service.GetHttpClient()
}
if info.IsStream {
helper.SetEventStreamHeaders(c)
// 处理流式请求的 ping 保活
generalSettings := operation_setting.GetGeneralSetting()
if generalSettings.PingIntervalEnabled {
pingInterval := time.Duration(generalSettings.PingIntervalSeconds) * time.Second
stopPinger := startPingKeepAlive(c, pingInterval)
defer stopPinger()
}
}
resp, err := client.Do(req)
if err != nil { if err != nil {
return nil, err return nil, err
} }
if resp == nil { if resp == nil {
return nil, errors.New("resp is nil") return nil, errors.New("resp is nil")
} }
_ = req.Body.Close() _ = req.Body.Close()
_ = c.Request.Body.Close() _ = c.Request.Body.Close()
return resp, nil return resp, nil
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment