Browse Source

✨ feat(middleware): add HTTP statistics middleware

CaIon 9 months ago
parent
commit
2509f644bc
3 changed files with 49 additions and 2 deletions
  1. 7 2
      controller/misc.go
  2. 41 0
      middleware/stats.go
  3. 1 0
      router/relay-router.go

+ 7 - 2
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": "",

+ 41 - 0
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),
+	}
+} 

+ 1 - 0
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())