feat(monitor): add Prometheus metrics

controller/relay.go

@@ -15,6 +15,7 @@ import (
 	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/middleware"
 	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/monitor"
 	"github.com/QuantumNous/new-api/relay"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
 	relayconstant "github.com/QuantumNous/new-api/relay/constant"
@@ -221,6 +222,17 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 
 		if newAPIError == nil {
 			relayInfo.LastError = nil
+			monitor.RecordRelayRequest(&monitor.RelayMetricsData{
+				ChannelId:      relayInfo.ChannelId,
+				ChannelType:    relayInfo.ChannelType,
+				Model:          relayInfo.OriginModelName,
+				RelayMode:      relayInfo.RelayMode,
+				StatusCode:     200,
+				StartTime:      relayInfo.StartTime,
+				FirstTokenTime: relayInfo.FirstResponseTime,
+				IsStream:       relayInfo.IsStream,
+			})
+			monitor.RecordTokenRequest(relayInfo.UserId, relayInfo.TokenId)
 			return
 		}
 
@@ -232,6 +244,22 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
 			break
 		}
+		monitor.RecordRelayRetry(relayInfo.ChannelId, relayInfo.ChannelType)
+	}
+
+	// Record metrics for failed relay
+	if newAPIError != nil {
+		monitor.RecordRelayRequest(&monitor.RelayMetricsData{
+			ChannelId:   relayInfo.ChannelId,
+			ChannelType: relayInfo.ChannelType,
+			Model:       relayInfo.OriginModelName,
+			RelayMode:   relayInfo.RelayMode,
+			StatusCode:  newAPIError.StatusCode,
+			StartTime:   relayInfo.StartTime,
+			IsStream:    relayInfo.IsStream,
+			ErrorType:   string(newAPIError.GetErrorCode()),
+		})
+		monitor.RecordTokenRequest(relayInfo.UserId, relayInfo.TokenId)
 	}
 
 	useChannel := c.GetStringSlice("use_channel")

go.mod

@@ -35,6 +35,7 @@ require (
 	github.com/nicksnyder/go-i18n/v2 v2.6.1
 	github.com/pkg/errors v0.9.1
 	github.com/pquerna/otp v1.5.0
+	github.com/prometheus/client_golang v1.22.0
 	github.com/samber/hot v0.11.0
 	github.com/samber/lo v1.52.0
 	github.com/shirou/gopsutil v3.21.11+incompatible
@@ -115,7 +116,6 @@ require (
 	github.com/ncruces/go-strftime v0.1.9 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
-	github.com/prometheus/client_golang v1.22.0 // indirect
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.62.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect

go.sum

@@ -186,6 +186,8 @@ github.com/kr/pty v1.1.8/go.mod h1:O1sed60cT9XZ5uDucP5qwvh+TE3NnUj51EiZO/lmSfw=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
+github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
+github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
 github.com/leodido/go-urn v1.2.0/go.mod h1:+8+nEpDfqqsY+g338gtMEUOtuK+4dEMhiQEgxpxOKII=
 github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY=
 github.com/leodido/go-urn v1.4.0 h1:WT9HwE9SGECu3lg4d/dIA+jxlljEa1/ffXKmRjqdmIQ=

main.go

@@ -18,6 +18,7 @@ import (
 	"github.com/QuantumNous/new-api/logger"
 	"github.com/QuantumNous/new-api/middleware"
 	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/monitor"
 	"github.com/QuantumNous/new-api/oauth"
 	"github.com/QuantumNous/new-api/relay"
 	"github.com/QuantumNous/new-api/router"
@@ -165,6 +166,7 @@ func main() {
 	// This will cause SSE not to work!!!
 	//server.Use(gzip.Gzip(gzip.DefaultCompression))
 	server.Use(middleware.RequestId())
+	server.Use(monitor.PrometheusMiddleware())
 	server.Use(middleware.PoweredBy())
 	server.Use(middleware.I18n())
 	middleware.SetUpLogger(server)
@@ -252,6 +254,12 @@ func InitResources() error {
 	// 加载环境变量
 	common.InitEnv()
 
+	// Initialize Prometheus metrics
+	if os.Getenv("PROMETHEUS_ENABLED") != "false" {
+		monitor.InitMetrics()
+		go monitor.StartMetricsServer()
+	}
+
 	logger.SetupLogger()
 
 	// Initialize model settings
@@ -291,6 +299,11 @@ func InitResources() error {
 		return err
 	}
 
+	// Register Prometheus metrics hook for Redis
+	if common.RedisEnabled {
+		monitor.RegisterRedisHook(common.RDB)
+	}
+
 	// 启动系统监控
 	common.StartSystemMonitor()
 

model/main.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/monitor"
 
 	"github.com/glebarez/sqlite"
 	"gorm.io/driver/mysql"
@@ -195,6 +196,10 @@ func InitDB() (err error) {
 		sqlDB.SetMaxOpenConns(common.GetEnvOrDefault("SQL_MAX_OPEN_CONNS", 1000))
 		sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetEnvOrDefault("SQL_MAX_LIFETIME", 60)))
 
+		// Register Prometheus DB metrics
+		monitor.RegisterDBCallbacks(DB)
+		monitor.StartDBPoolCollector(sqlDB)
+
 		if !common.IsMasterNode {
 			return nil
 		}

monitor/db.go

@@ -0,0 +1,100 @@
+package monitor
+
+import (
+	"database/sql"
+	"time"
+
+	"gorm.io/gorm"
+)
+
+const dbStartTimeKey = "prom_start_time"
+
+// RegisterDBCallbacks adds Prometheus instrumentation callbacks to a GORM DB.
+func RegisterDBCallbacks(db *gorm.DB) {
+	if db == nil || !Enabled() {
+		return
+	}
+
+	registerCallback(db, "create")
+	registerCallback(db, "query")
+	registerCallback(db, "update")
+	registerCallback(db, "delete")
+	registerCallback(db, "raw")
+}
+
+func registerCallback(db *gorm.DB, operation string) {
+	beforeName := "prometheus:before_" + operation
+	afterName := "prometheus:after_" + operation
+
+	before := func(db *gorm.DB) {
+		db.Set(dbStartTimeKey, time.Now())
+	}
+	after := func(db *gorm.DB) {
+		recordDBMetrics(db, operation)
+	}
+
+	switch operation {
+	case "create":
+		_ = db.Callback().Create().Before("*").Register(beforeName, before)
+		_ = db.Callback().Create().After("*").Register(afterName, after)
+	case "query":
+		_ = db.Callback().Query().Before("*").Register(beforeName, before)
+		_ = db.Callback().Query().After("*").Register(afterName, after)
+	case "update":
+		_ = db.Callback().Update().Before("*").Register(beforeName, before)
+		_ = db.Callback().Update().After("*").Register(afterName, after)
+	case "delete":
+		_ = db.Callback().Delete().Before("*").Register(beforeName, before)
+		_ = db.Callback().Delete().After("*").Register(afterName, after)
+	case "raw":
+		_ = db.Callback().Raw().Before("*").Register(beforeName, before)
+		_ = db.Callback().Raw().After("*").Register(afterName, after)
+	}
+}
+
+func recordDBMetrics(db *gorm.DB, operation string) {
+	if !Enabled() {
+		return
+	}
+	status := "success"
+	if db.Error != nil {
+		status = "error"
+	}
+	DBQueriesTotal.WithLabelValues(operation, status).Inc()
+
+	if v, ok := db.Get(dbStartTimeKey); ok {
+		if start, ok := v.(time.Time); ok {
+			DBQueryDuration.WithLabelValues(operation).Observe(time.Since(start).Seconds())
+		}
+	}
+}
+
+// StartDBPoolCollector starts a background goroutine that periodically collects DB pool stats.
+func StartDBPoolCollector(sqlDB *sql.DB) {
+	if sqlDB == nil || !Enabled() {
+		return
+	}
+
+	go func() {
+		var lastWaitCount int64
+		var lastWaitDuration time.Duration
+
+		ticker := time.NewTicker(15 * time.Second)
+		defer ticker.Stop()
+		for range ticker.C {
+			stats := sqlDB.Stats()
+			DBOpenConnections.WithLabelValues("open").Set(float64(stats.OpenConnections))
+			DBOpenConnections.WithLabelValues("in_use").Set(float64(stats.InUse))
+			DBOpenConnections.WithLabelValues("idle").Set(float64(stats.Idle))
+
+			if delta := stats.WaitCount - lastWaitCount; delta > 0 {
+				DBWaitTotal.Add(float64(delta))
+				lastWaitCount = stats.WaitCount
+			}
+			if delta := stats.WaitDuration - lastWaitDuration; delta > 0 {
+				DBWaitDuration.Add(delta.Seconds())
+				lastWaitDuration = stats.WaitDuration
+			}
+		}
+	}()
+}