diff --git a/controller/misc.go b/controller/misc.go index be76cab5..d95ced5b 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -6,6 +6,7 @@ import ( "net/http" "one-api/common" "one-api/constant" + "one-api/middleware" "one-api/model" "one-api/setting" "one-api/setting/operation_setting" @@ -24,14 +25,18 @@ func TestStatus(c *gin.Context) { }) return } + // 获取HTTP统计信息 + httpStats := middleware.GetStats() c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "Server is running", + "success": true, + "message": "Server is running", + "http_stats": httpStats, }) return } func GetStatus(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/middleware/stats.go b/middleware/stats.go new file mode 100644 index 00000000..1c97983f --- /dev/null +++ b/middleware/stats.go @@ -0,0 +1,41 @@ +package middleware + +import ( + "sync/atomic" + + "github.com/gin-gonic/gin" +) + +// HTTPStats 存储HTTP统计信息 +type HTTPStats struct { + activeConnections int64 +} + +var globalStats = &HTTPStats{} + +// StatsMiddleware 统计中间件 +func StatsMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + // 增加活跃连接数 + atomic.AddInt64(&globalStats.activeConnections, 1) + + // 确保在请求结束时减少连接数 + defer func() { + atomic.AddInt64(&globalStats.activeConnections, -1) + }() + + c.Next() + } +} + +// StatsInfo 统计信息结构 +type StatsInfo struct { + ActiveConnections int64 `json:"active_connections"` +} + +// GetStats 获取统计信息 +func GetStats() StatsInfo { + return StatsInfo{ + ActiveConnections: atomic.LoadInt64(&globalStats.activeConnections), + } +} \ No newline at end of file diff --git a/router/relay-router.go b/router/relay-router.go index 1115a491..aa7f27a8 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -11,6 +11,7 @@ import ( func SetRelayRouter(router *gin.Engine) { router.Use(middleware.CORS()) router.Use(middleware.DecompressRequestMiddleware()) + router.Use(middleware.StatsMiddleware()) // https://platform.openai.com/docs/api-reference/introduction modelsRouter := router.Group("/v1/models") modelsRouter.Use(middleware.TokenAuth())