Feature/memory analysis

.gitignore

@@ -18,4 +18,5 @@ out.log
 out.log_2
 .env_3
 *.prof
-.env*
+.env*
+.env_prod

common/constants.go

@@ -347,3 +347,8 @@ var ChannelBaseURLs = []string{
 	"",                                          //99
 	"https://ark.cn-beijing.volces.com",         //100 - 豆包离线
 }
+
+func init() {
+	PProfEnabled = false
+	PProfMutex = sync.RWMutex{}
+}

common/pprof.go

@@ -2,10 +2,20 @@ package common
 
 import (
 	"fmt"
-	"github.com/shirou/gopsutil/cpu"
+	"net/http"
 	"os"
 	"runtime/pprof"
+	"sync"
 	"time"
+
+	"github.com/shirou/gopsutil/cpu"
+)
+
+var (
+	PProfEnabled  bool
+	PProfMutex    sync.RWMutex
+	pprofServer   *http.Server
+	serverRunning bool
 )
 
 // Monitor 定时监控cpu使用率，超过阈值输出pprof文件
@@ -42,3 +52,51 @@ func Monitor() {
 		time.Sleep(30 * time.Second)
 	}
 }
+
+// InitPProfServer 初始化并启动 pprof 服务器
+func InitPProfServer() {
+	// 创建 pprof 服务器
+	pprofServer = &http.Server{
+		Addr:    "0.0.0.0:8005",
+		Handler: http.DefaultServeMux,
+	}
+
+	// 启动监控协程
+	go func() {
+		for {
+			PProfMutex.RLock()
+			enabled := PProfEnabled
+			PProfMutex.RUnlock()
+
+			SysLog(fmt.Sprintf("[PPROF] Status check - enabled: %v, serverRunning: %v", enabled, serverRunning))
+
+			if enabled && !serverRunning {
+				// 启动服务器
+				go func() {
+					SysLog("[PPROF] Starting server on :8005")
+					if err := pprofServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+						SysError(fmt.Sprintf("[PPROF] Server error: %v", err))
+					}
+				}()
+				serverRunning = true
+				SysLog("[PPROF] Server marked as running")
+			} else if !enabled && serverRunning {
+				// 关闭服务器
+				SysLog("[PPROF] Stopping server")
+				if err := pprofServer.Close(); err != nil {
+					SysError(fmt.Sprintf("[PPROF] Server close error: %v", err))
+				}
+				// 重新创建服务器实例
+				pprofServer = &http.Server{
+					Addr:    "0.0.0.0:8005",
+					Handler: http.DefaultServeMux,
+				}
+				serverRunning = false
+				SysLog("[PPROF] Server marked as stopped")
+			}
+			time.Sleep(5 * time.Second)
+		}
+	}()
+
+	SysLog("[PPROF] Server initialized and monitoring started")
+}

controller/pprof.go

@@ -0,0 +1,40 @@
+package controller
+
+import (
+	"net/http"
+	"one-api/common"
+
+	"github.com/gin-gonic/gin"
+)
+
+// PProfStatus 获取 pprof 状态
+func PProfStatus(c *gin.Context) {
+	common.PProfMutex.RLock()
+	enabled := common.PProfEnabled
+	common.PProfMutex.RUnlock()
+	c.JSON(http.StatusOK, gin.H{
+		"enabled": enabled,
+	})
+}
+
+// EnablePProf 启用 pprof
+func EnablePProf(c *gin.Context) {
+	common.PProfMutex.Lock()
+	common.PProfEnabled = true
+	common.PProfMutex.Unlock()
+	common.SysLog("pprof enabled via API")
+	c.JSON(http.StatusOK, gin.H{
+		"message": "pprof enabled",
+	})
+}
+
+// DisablePProf 禁用 pprof
+func DisablePProf(c *gin.Context) {
+	common.PProfMutex.Lock()
+	common.PProfEnabled = false
+	common.PProfMutex.Unlock()
+	common.SysLog("pprof disabled via API")
+	c.JSON(http.StatusOK, gin.H{
+		"message": "pprof disabled",
+	})
+}

controller/relay.go

@@ -351,6 +351,16 @@ func shouldRetry(c *gin.Context, openaiErr *dto.OpenAIErrorWithStatusCode, retry
 	if _, ok := c.Get("specific_channel_id"); ok {
 		return false
 	}
+	if strings.Contains(openaiErr.Error.Message, "deadline exceeded") ||
+		strings.Contains(openaiErr.Error.Message, "request canceled") ||
+		strings.Contains(openaiErr.Error.Message, "copy_response_body_failed") {
+		common.LogInfo(c, fmt.Sprintf("客户端请求下游超时，不再重试 : %s", openaiErr.Error.Message))
+		return false
+	}
+	if openaiErr.Error.Code == "copy_response_body_failed" {
+		common.LogInfo(c, fmt.Sprintf("客户端连接断开，不再重试 : %s", openaiErr.Error.Message))
+		return false
+	}
 	if openaiErr.StatusCode == http.StatusTooManyRequests {
 		return true
 	}
@@ -377,10 +387,6 @@ func shouldRetry(c *gin.Context, openaiErr *dto.OpenAIErrorWithStatusCode, retry
 	if openaiErr.StatusCode == 307 {
 		return true
 	}
-	if strings.Contains(openaiErr.Error.Message, "deadline exceeded") || strings.Contains(openaiErr.Error.Message, "request canceled") || strings.Contains(openaiErr.Error.Message, "copy_response_body_failed") {
-		common.LogInfo(c, fmt.Sprintf("客户端请求下游超时，不再重试 : %s", openaiErr.Error.Message))
-		return false
-	}
 
 	if openaiErr.StatusCode/100 == 5 {
 		// 超时不重试

main.go

@@ -221,12 +221,9 @@ func main() {
 	}
 
 	if os.Getenv("ENABLE_PPROF") == "true" {
-		gopool.Go(func() {
-			log.Println(http.ListenAndServe("0.0.0.0:8005", nil))
-		})
-		go common.Monitor()
-		common.SysLog("pprof enabled")
+		common.PProfEnabled = true
 	}
+	common.InitPProfServer()
 
 	service.InitTokenEncoders()
 

relay/relay-text.go

@@ -249,6 +249,12 @@ func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *d
 
 	if resp != nil {
 		httpResp = resp.(*http.Response)
+		// 确保响应体被正确关闭
+		defer func() {
+			if httpResp != nil && httpResp.Body != nil {
+				httpResp.Body.Close()
+			}
+		}()
 		// // 直接设置到 gin 的响应头
 		// c.Writer.Header().Set("X-Origin-User-ID", strconv.Itoa(relayInfo.UserId))
 		// c.Writer.Header().Set("X-Origin-Channel-ID", strconv.Itoa(relayInfo.ChannelId))
@@ -271,6 +277,10 @@ func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *d
 			service.ResetStatusCode(openaiErr, statusCodeMappingStr)
 			return openaiErr
 		}
+	} else {
+		// 如果resp为nil，返回错误
+		funcErr = service.OpenAIErrorWrapperLocal(fmt.Errorf("no response received"), "no_response", http.StatusInternalServerError)
+		return funcErr
 	}
 
 	// 读取响应体并创建副本
@@ -281,7 +291,14 @@ func TextHelper(c *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest *d
 	}
 	// 为adaptor创建一个新的响应体
 	httpResp.Body = io.NopCloser(bytes.NewBuffer(responseBodyBytes))
-	common.LogInfo(c, fmt.Sprintf("response body: %s", string(responseBodyBytes)))
+
+	// 限制日志大小，避免内存泄漏
+	if len(responseBodyBytes) <= 2048 {
+		common.LogInfo(c, fmt.Sprintf("response body: %s", string(responseBodyBytes)))
+	} else {
+		common.LogInfo(c, fmt.Sprintf("response body too large (size: %d bytes), skipping print", len(responseBodyBytes)))
+	}
+
 	usage, openaiErr := adaptor.DoResponse(c, httpResp, relayInfo)
 	if openaiErr != nil {
 		funcErr = openaiErr

router/api-router.go

@@ -80,6 +80,16 @@ func SetApiRouter(router *gin.Engine) {
 			optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio)
 			optionRoute.POST("/request_log", controller.ToggleRequestLog)
 		}
+
+		// 添加 pprof 控制路由
+		pprofRoute := apiRouter.Group("/pprof")
+		pprofRoute.Use(middleware.RootAuth())
+		{
+			pprofRoute.GET("/status", controller.PProfStatus)
+			pprofRoute.POST("/enable", controller.EnablePProf)
+			pprofRoute.POST("/disable", controller.DisablePProf)
+		}
+
 		channelRoute := apiRouter.Group("/channel")
 		channelRoute.Use(middleware.AdminAuth())
 		{