new-api-custom/pkg/perf_metrics/live_status.go

package perfmetrics

import (
	"context"
	"github.com/QuantumNous/new-api/common"
	"strconv"
	"sync"
	"time"
)

var liveSuccess sync.Map

type successTime struct {
	sync.Mutex
	at int64
}

func recordUserSuccess(group string, at int64) {
	if group == "" {
		return
	}
	value, _ := liveSuccess.LoadOrStore(group, &successTime{})
	v := value.(*successTime)
	v.Lock()
	v.at = max(v.at, at)
	v.Unlock()
	if common.RedisEnabled && common.RDB != nil {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		// A delayed writer must not replace a newer successful request.
		_ = common.RDB.Eval(ctx, `local old=tonumber(redis.call('GET',KEYS[1]) or '0'); if tonumber(ARGV[1])>old then redis.call('SET',KEYS[1],ARGV[1],'EX',900) end; return 1`, []string{"status:user-success:" + group}, at).Err()
	}
}

func LatestUserSuccess(ctx context.Context, groups []string) map[string]int64 {
	result := map[string]int64{}
	for _, group := range groups {
		if value, ok := liveSuccess.Load(group); ok {
			v := value.(*successTime)
			v.Lock()
			result[group] = v.at
			v.Unlock()
		}
	}
	if common.RedisEnabled && common.RDB != nil && len(groups) > 0 {
		keys := make([]string, len(groups))
		for i, g := range groups {
			keys[i] = "status:user-success:" + g
		}
		ctx, cancel := context.WithTimeout(ctx, time.Second)
		defer cancel()
		values, err := common.RDB.MGet(ctx, keys...).Result()
		if err == nil {
			for i, value := range values {
				if s, ok := value.(string); ok {
					at, _ := strconv.ParseInt(s, 10, 64)
					result[groups[i]] = max(result[groups[i]], at)
				}
			}
		}
	}
	return result
}