feat: redis perf and web runtime monitor (#498)

Author: zijiren233

core/common/reqlimit/main.go

@@ -2,15 +2,24 @@ package reqlimit
 
 import (
 	"context"
+	"strconv"
 	"time"
 
 	"github.com/labring/aiproxy/core/common"
+	"github.com/redis/go-redis/v9"
 	log "github.com/sirupsen/logrus"
 )
 
+type ChannelModelRate struct {
+	RPM int64 `json:"rpm"`
+	TPM int64 `json:"tpm"`
+	RPS int64 `json:"rps"`
+	TPS int64 `json:"tps"`
+}
+
 var (
 	memoryGroupModelLimiter = NewInMemoryRecord()
-	redisGroupModelLimiter  = newRedisGroupModelRecord()
+	redisGroupModelLimiter  = newRedisGroupModelRecord(func() *redis.Client { return common.RDB })
 )
 
 func PushGroupModelRequest(
@@ -61,7 +70,9 @@ func GetGroupModelRequest(ctx context.Context, group, model string) (int64, int6
 
 var (
 	memoryGroupModelTokennameLimiter = NewInMemoryRecord()
-	redisGroupModelTokennameLimiter  = newRedisGroupModelTokennameRecord()
+	redisGroupModelTokennameLimiter  = newRedisGroupModelTokennameRecord(
+		func() *redis.Client { return common.RDB },
+	)
 )
 
 func PushGroupModelTokennameRequest(
@@ -120,7 +131,9 @@ func GetGroupModelTokennameRequest(
 
 var (
 	memoryChannelModelRecord = NewInMemoryRecord()
-	redisChannelModelRecord  = newRedisChannelModelRecord()
+	redisChannelModelRecord  = newRedisChannelModelRecord(
+		func() *redis.Client { return common.RDB },
+	)
 )
 
 func PushChannelModelRequest(ctx context.Context, channel, model string) (int64, int64, int64) {
@@ -171,7 +184,9 @@ func GetChannelModelRequest(ctx context.Context, channel, model string) (int64,
 
 var (
 	memoryGroupModelTokensLimiter = NewInMemoryRecord()
-	redisGroupModelTokensLimiter  = newRedisGroupModelTokensRecord()
+	redisGroupModelTokensLimiter  = newRedisGroupModelTokensRecord(
+		func() *redis.Client { return common.RDB },
+	)
 )
 
 func PushGroupModelTokensRequest(
@@ -222,7 +237,9 @@ func GetGroupModelTokensRequest(ctx context.Context, group, model string) (int64
 
 var (
 	memoryGroupModelTokennameTokensLimiter = NewInMemoryRecord()
-	redisGroupModelTokennameTokensLimiter  = newRedisGroupModelTokennameTokensRecord()
+	redisGroupModelTokennameTokensLimiter  = newRedisGroupModelTokennameTokensRecord(
+		func() *redis.Client { return common.RDB },
+	)
 )
 
 func PushGroupModelTokennameTokensRequest(
@@ -289,7 +306,9 @@ func GetGroupModelTokennameTokensRequest(
 
 var (
 	memoryChannelModelTokensRecord = NewInMemoryRecord()
-	redisChannelModelTokensRecord  = newRedisChannelModelTokensRecord()
+	redisChannelModelTokensRecord  = newRedisChannelModelTokensRecord(
+		func() *redis.Client { return common.RDB },
+	)
 )
 
 func PushChannelModelTokensRequest(
@@ -341,3 +360,69 @@ func GetChannelModelTokensRequest(ctx context.Context, channel, model string) (i
 
 	return memoryChannelModelTokensRecord.GetRequest(time.Minute, channel, model)
 }
+
+func GetAllChannelModelRates(ctx context.Context) (map[int64]map[string]ChannelModelRate, error) {
+	requests := make(map[int64]map[string]ChannelModelRate)
+	appendSnapshot := func(snapshot recordSnapshot, assign func(rate *ChannelModelRate)) {
+		if len(snapshot.Keys) != 2 {
+			return
+		}
+
+		channelID, err := strconv.ParseInt(snapshot.Keys[0], 10, 64)
+		if err != nil {
+			return
+		}
+
+		model := snapshot.Keys[1]
+
+		if _, ok := requests[channelID]; !ok {
+			requests[channelID] = make(map[string]ChannelModelRate)
+		}
+
+		rate := requests[channelID][model]
+		assign(&rate)
+		requests[channelID][model] = rate
+	}
+
+	var (
+		requestSnapshots []recordSnapshot
+		tokenSnapshots   []recordSnapshot
+		err              error
+	)
+
+	if common.RedisEnabled {
+		requestSnapshots, err = redisChannelModelRecord.Snapshot(ctx, time.Minute)
+		if err != nil {
+			log.Error("redis snapshot request error: " + err.Error())
+		}
+
+		tokenSnapshots, err = redisChannelModelTokensRecord.Snapshot(ctx, time.Minute)
+		if err != nil {
+			log.Error("redis snapshot token error: " + err.Error())
+		}
+	}
+
+	if len(requestSnapshots) == 0 {
+		requestSnapshots = memoryChannelModelRecord.Snapshot(time.Minute)
+	}
+
+	if len(tokenSnapshots) == 0 {
+		tokenSnapshots = memoryChannelModelTokensRecord.Snapshot(time.Minute)
+	}
+
+	for _, snapshot := range requestSnapshots {
+		appendSnapshot(snapshot, func(rate *ChannelModelRate) {
+			rate.RPM = snapshot.TotalCount
+			rate.RPS = snapshot.SecondCount
+		})
+	}
+
+	for _, snapshot := range tokenSnapshots {
+		appendSnapshot(snapshot, func(rate *ChannelModelRate) {
+			rate.TPM = snapshot.TotalCount
+			rate.TPS = snapshot.SecondCount
+		})
+	}
+
+	return requests, nil
+}

core/common/reqlimit/mem.go

@@ -1,6 +1,8 @@
 package reqlimit
 
 import (
+	"path"
+	"slices"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -14,14 +16,25 @@ type windowCounts struct {
 
 type entry struct {
 	sync.Mutex
-	windows    map[int64]*windowCounts
-	lastAccess atomic.Value
+	windows              map[int64]*windowCounts
+	lastAccess           atomic.Value
+	windowSeconds        int64
+	totalNormal          int64
+	totalOver            int64
+	lastCleanedCutoff    int64
+	aggregateInitialized bool
 }
 
 type InMemoryRecord struct {
 	entries sync.Map
 }
 
+type recordSnapshot struct {
+	Keys        []string
+	TotalCount  int64
+	SecondCount int64
+}
+
 func NewInMemoryRecord() *InMemoryRecord {
 	rl := &InMemoryRecord{
 		entries: sync.Map{},
@@ -45,20 +58,49 @@ func (m *InMemoryRecord) getEntry(keys []string) *entry {
 	return e
 }
 
-func (m *InMemoryRecord) cleanupAndCount(e *entry, cutoff int64) (int64, int64) {
+func (m *InMemoryRecord) rebuildAggregateLocked(e *entry, windowSeconds, cutoff int64) {
 	normalCount := int64(0)
-
 	overCount := int64(0)
+
 	for ts, wc := range e.windows {
 		if ts < cutoff {
 			delete(e.windows, ts)
-		} else {
-			normalCount += wc.normal
-			overCount += wc.over
+			continue
+		}
+
+		normalCount += wc.normal
+		overCount += wc.over
+	}
+
+	e.windowSeconds = windowSeconds
+	e.totalNormal = normalCount
+	e.totalOver = overCount
+	e.lastCleanedCutoff = cutoff
+	e.aggregateInitialized = true
+}
+
+func (m *InMemoryRecord) refreshAggregateLocked(e *entry, nowSecond, windowSeconds int64) {
+	cutoff := nowSecond - windowSeconds
+
+	if !e.aggregateInitialized || e.windowSeconds != windowSeconds {
+		m.rebuildAggregateLocked(e, windowSeconds, cutoff)
+		return
+	}
+
+	for ts := e.lastCleanedCutoff; ts < cutoff; ts++ {
+		wc, ok := e.windows[ts]
+		if !ok {
+			continue
 		}
+
+		e.totalNormal -= wc.normal
+		e.totalOver -= wc.over
+		delete(e.windows, ts)
 	}
 
-	return normalCount, overCount
+	if e.lastCleanedCutoff < cutoff {
+		e.lastCleanedCutoff = cutoff
+	}
 }
 
 func (m *InMemoryRecord) PushRequest(
@@ -76,33 +118,52 @@ func (m *InMemoryRecord) PushRequest(
 	e.lastAccess.Store(now)
 
 	windowStart := now.Unix()
-	cutoff := windowStart - int64(duration.Seconds())
-
-	normalCount, overCount = m.cleanupAndCount(e, cutoff)
+	windowSeconds := int64(duration.Seconds())
+	m.refreshAggregateLocked(e, windowStart, windowSeconds)
 
 	wc, exists := e.windows[windowStart]
 	if !exists {
 		wc = &windowCounts{}
 		e.windows[windowStart] = wc
 	}
 
-	if overed == 0 || normalCount <= overed {
+	if overed == 0 || e.totalNormal <= overed {
 		wc.normal += n
-		normalCount += n
+		e.totalNormal += n
 	} else {
 		wc.over += n
-		overCount += n
+		e.totalOver += n
 	}
 
-	return normalCount, overCount, wc.normal + wc.over
+	return e.totalNormal, e.totalOver, wc.normal + wc.over
 }
 
 func (m *InMemoryRecord) GetRequest(
 	duration time.Duration,
 	keys ...string,
 ) (totalCount, secondCount int64) {
 	nowSecond := time.Now().Unix()
-	cutoff := nowSecond - int64(duration.Seconds())
+	windowSeconds := int64(duration.Seconds())
+
+	if !hasWildcard(keys) {
+		value, ok := m.entries.Load(strings.Join(keys, ":"))
+		if !ok {
+			return 0, 0
+		}
+
+		e, _ := value.(*entry)
+		e.Lock()
+		m.refreshAggregateLocked(e, nowSecond, windowSeconds)
+		nowWindow := e.windows[nowSecond]
+		totalCount = e.totalNormal + e.totalOver
+		e.Unlock()
+
+		if nowWindow != nil {
+			secondCount = nowWindow.normal + nowWindow.over
+		}
+
+		return totalCount, secondCount
+	}
 
 	m.entries.Range(func(key, value any) bool {
 		k, _ := key.(string)
@@ -111,11 +172,12 @@ func (m *InMemoryRecord) GetRequest(
 		if matchKeys(keys, currentKeys) {
 			e, _ := value.(*entry)
 			e.Lock()
-			normalCount, overCount := m.cleanupAndCount(e, cutoff)
+			m.refreshAggregateLocked(e, nowSecond, windowSeconds)
 			nowWindow := e.windows[nowSecond]
+			entryTotalCount := e.totalNormal + e.totalOver
 			e.Unlock()
 
-			totalCount += normalCount + overCount
+			totalCount += entryTotalCount
 
 			if nowWindow != nil {
 				secondCount += nowWindow.normal + nowWindow.over
@@ -151,6 +213,37 @@ func (m *InMemoryRecord) cleanupInactiveEntries(interval, maxInactivity time.Dur
 	}
 }
 
+func (m *InMemoryRecord) Snapshot(duration time.Duration) []recordSnapshot {
+	nowSecond := time.Now().Unix()
+	windowSeconds := int64(duration.Seconds())
+	snapshots := make([]recordSnapshot, 0)
+
+	m.entries.Range(func(key, value any) bool {
+		k, _ := key.(string)
+		e, _ := value.(*entry)
+		e.Lock()
+		m.refreshAggregateLocked(e, nowSecond, windowSeconds)
+		nowWindow := e.windows[nowSecond]
+		totalCount := e.totalNormal + e.totalOver
+		e.Unlock()
+
+		secondCount := int64(0)
+		if nowWindow != nil {
+			secondCount = nowWindow.normal + nowWindow.over
+		}
+
+		snapshots = append(snapshots, recordSnapshot{
+			Keys:        parseKeys(k),
+			TotalCount:  totalCount,
+			SecondCount: secondCount,
+		})
+
+		return true
+	})
+
+	return snapshots
+}
+
 func parseKeys(key string) []string {
 	return strings.Split(key, ":")
 }
@@ -161,10 +254,27 @@ func matchKeys(pattern, keys []string) bool {
 	}
 
 	for i, p := range pattern {
-		if p != "*" && p != keys[i] {
+		if isGlobPattern(p) {
+			matched, err := path.Match(p, keys[i])
+			if err != nil || !matched {
+				return false
+			}
+
+			continue
+		}
+
+		if p != keys[i] {
 			return false
 		}
 	}
 
 	return true
 }
+
+func hasWildcard(keys []string) bool {
+	return slices.ContainsFunc(keys, isGlobPattern)
+}
+
+func isGlobPattern(key string) bool {
+	return strings.ContainsAny(key, "*?[")
+}