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

package perfmetrics

import (
	"context"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/QuantumNous/new-api/common"
)

// At is milliseconds so a success and failure in the same second stay ordered.
type UserRequestResult struct {
	At      int64
	Success bool
}

var liveResults sync.Map

type userResultState struct {
	sync.Mutex
	result UserRequestResult
}

func recordUserResult(group string, result UserRequestResult) {
	if group == "" {
		return
	}
	value, _ := liveResults.LoadOrStore(group, &userResultState{})
	v := value.(*userResultState)
	v.Lock()
	if result.At >= v.result.At {
		v.result = result
	}
	v.Unlock()
	if common.RedisEnabled && common.RDB != nil {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		encoded := strconv.FormatInt(result.At, 10) + ":" + strconv.FormatBool(result.Success)
		// Delayed writers cannot replace newer evidence from another instance.
		_ = common.RDB.Eval(ctx, `local old=redis.call('GET',KEYS[1]); local at=old and tonumber(string.match(old,'^(%d+):')) or 0; if tonumber(ARGV[1])>=at then redis.call('SET',KEYS[1],ARGV[2],'EX',900) end; return 1`, []string{"status:user-result:" + group}, result.At, encoded).Err()
	}
}

func LatestUserResults(ctx context.Context, groups []string) map[string]UserRequestResult {
	result := map[string]UserRequestResult{}
	for _, group := range groups {
		if value, ok := liveResults.Load(group); ok {
			v := value.(*userResultState)
			v.Lock()
			result[group] = v.result
			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-result:" + 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 {
				s, ok := value.(string)
				if !ok {
					continue
				}
				atText, successText, ok := strings.Cut(s, ":")
				if !ok {
					continue
				}
				at, atErr := strconv.ParseInt(atText, 10, 64)
				success, successErr := strconv.ParseBool(successText)
				if atErr == nil && successErr == nil && at >= result[groups[i]].At {
					result[groups[i]] = UserRequestResult{At: at, Success: success}
				}
			}
		}
	}
	return result
}