diff --git a/.gitignore b/.gitignore index 6e3b5b7ab521..ebc614dd0cb9 100644 --- a/.gitignore +++ b/.gitignore @@ -18,4 +18,5 @@ out.log out.log_2 .env_3 *.prof -.env* \ No newline at end of file +.env* +.env_prod diff --git a/common/constants.go b/common/constants.go index e01593b0cd54..1bac0c74e2e7 100644 --- a/common/constants.go +++ b/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{} +} diff --git a/common/pprof.go b/common/pprof.go index 4bec30f1551e..3f2539057cf3 100644 --- a/common/pprof.go +++ b/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") +} diff --git a/controller/pprof.go b/controller/pprof.go new file mode 100644 index 000000000000..82f8607d8aff --- /dev/null +++ b/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", + }) +} diff --git a/controller/relay.go b/controller/relay.go index 1d78978487e7..c847e91ff084 100644 --- a/controller/relay.go +++ b/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 { // 超时不重试 diff --git a/main.go b/main.go index ecd080a5b2f0..fbbc02ab5722 100644 --- a/main.go +++ b/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() diff --git a/relay/relay-text.go b/relay/relay-text.go index cb3930dad4a1..b2ff9de69bc9 100644 --- a/relay/relay-text.go +++ b/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 diff --git a/router/api-router.go b/router/api-router.go index cdcc4d2e14e1..afb252c77ed6 100644 --- a/router/api-router.go +++ b/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()) {