Merge pull request #1134 from QuantumNous/fix_ping_keepalive

fix: 流式请求ping

relay/channel/api_request.go

@@ -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) {
-	var client *http.Client
-	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)
+func startPingKeepAlive(c *gin.Context, pingInterval time.Duration) context.CancelFunc {
+	pingerCtx, stopPinger := context.WithCancel(context.Background())
 
 
-		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
 
 
 		for {
 		for {
 			select {
 			select {
+			// 发送 ping 数据
 			case <-ticker.C:
 			case <-ticker.C:
-						pingMutex.Lock()
-						err2 := helper.PingData(c)
-						pingMutex.Unlock()
-						if err2 != nil {
-							common2.LogError(c, "SSE ping error: "+err.Error())
+				if err := sendPingData(c, &pingMutex); err != nil {
 					return
 					return
 				}
 				}
-						if common2.DebugEnabled {
-							println("SSE ping data sent.")
-						}
+			// 收到退出信号
 			case <-pingerCtx.Done():
 			case <-pingerCtx.Done():
-						if common2.DebugEnabled {
-							println("SSE ping goroutine stopped.")
-						}
+				return
+			// 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)
-	// request结束后等待 ping goroutine 完成
-	if info.IsStream && pingEnabled {
-		pingerWg.Wait()
+func doRequest(c *gin.Context, req *http.Request, info *common.RelayInfo) (*http.Response, error) {
+	var client *http.Client
+	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()
+	}
+
+	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