cneicy/new-api-custom
1
0 Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

fix: follow live request results and prioritize usable groups

Browse Source
This commit is contained in:
cneicy committed 2026-09-26 14:49:12 +08:00
1 parent 5152a4080e
commit d09216bbde
12 files changed
+257 -86
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pkg/perf_metrics/live_status.go
+41 -21
View File
@@ -2,59 +2,79 @@ package perfmetrics
import (
"context"
"github.com/QuantumNous/new-api/common"
"strconv"
"strings"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
)
var liveSuccess sync.Map
type successTime struct {
sync.Mutex
at int64
// At is milliseconds so a success and failure in the same second stay ordered.
type UserRequestResult struct {
At int64
Success bool
}
func recordUserSuccess(group string, at int64) {
var liveResults sync.Map
type userResultState struct {
sync.Mutex
result UserRequestResult
}
func recordUserResult(group string, result UserRequestResult) {
if group == "" {
return
}
value, _ := liveSuccess.LoadOrStore(group, &successTime{})
v := value.(*successTime)
value, _ := liveResults.LoadOrStore(group, &userResultState{})
v := value.(*userResultState)
v.Lock()
v.at = max(v.at, at)
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()
// 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()
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 LatestUserSuccess(ctx context.Context, groups []string) map[string]int64 {
result := map[string]int64{}
func LatestUserResults(ctx context.Context, groups []string) map[string]UserRequestResult {
result := map[string]UserRequestResult{}
for _, group := range groups {
if value, ok := liveSuccess.Load(group); ok {
v := value.(*successTime)
if value, ok := liveResults.Load(group); ok {
v := value.(*userResultState)
v.Lock()
result[group] = v.at
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-success:" + g
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 {
if s, ok := value.(string); ok {
at, _ := strconv.ParseInt(s, 10, 64)
result[groups[i]] = max(result[groups[i]], at)
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}
}
}
}
pkg/perf_metrics/live_status_test.go
+68 -20
View File
@@ -2,35 +2,83 @@ package perfmetrics
import (
"context"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/stretchr/testify/assert"
"testing"
"time"
"github.com/QuantumNous/new-api/common"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/types"
"github.com/alicebob/miniredis/v2"
"github.com/go-redis/redis/v8"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestLiveUserSuccess(t *testing.T) {
t.Cleanup(func() {
for _, name := range []string{"user", "probe", "internal", "failed"} {
liveSuccess.Delete(name)
}
})
func TestLiveUserResults(t *testing.T) {
oldRedis := common.RedisEnabled
common.RedisEnabled = false
t.Cleanup(func() { common.RedisEnabled = oldRedis })
canceled, cancel := context.WithCancel(context.Background())
cancel()
for _, tc := range []struct {
name string
test bool
token int
fail bool
want bool
name string
test bool
token int
ctx context.Context
err *types.NewAPIError
want bool
success bool
}{
{"user", false, 1, false, true}, {"probe", true, 1, false, false}, {"internal", false, 0, false, false}, {"failed", false, 1, true, false},
{"success", false, 1, context.Background(), nil, true, true},
{"upstream 503", false, 1, context.Background(), types.InitOpenAIError("server_error", 503), true, false},
{"probe", true, 1, context.Background(), types.InitOpenAIError("server_error", 503), false, false},
{"internal", false, 0, context.Background(), nil, false, false},
{"canceled", false, 1, canceled, types.InitOpenAIError("server_error", 503), false, false},
{"bad prompt", false, 1, context.Background(), types.InitOpenAIError("context_length_exceeded", 400), false, false},
} {
t.Run(tc.name, func(t *testing.T) {
var apiErr *types.NewAPIError
if tc.fail {
apiErr = types.InitOpenAIError("server_error", 500)
t.Cleanup(func() { liveResults.Delete(tc.name) })
RecordRelayResult(tc.ctx, &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, tc.err)
result := LatestUserResults(context.Background(), []string{tc.name})[tc.name]
assert.Equal(t, tc.want, result.At > 0)
if tc.want {
assert.Equal(t, tc.success, result.Success)
}
RecordRelayResult(context.Background(), &relaycommon.RelayInfo{UsingGroup: tc.name, TokenId: tc.token, IsChannelTest: tc.test, StartTime: time.Now()}, apiErr)
assert.Equal(t, tc.want, LatestUserSuccess(context.Background(), []string{tc.name})[tc.name] > 0)
})
}
}
func TestUserResultsAreOrderedAndSurviveRestart(t *testing.T) {
server := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: server.Addr()})
oldRedis, oldClient := common.RedisEnabled, common.RDB
common.RedisEnabled, common.RDB = true, client
t.Cleanup(func() {
common.RedisEnabled, common.RDB = oldRedis, oldClient
liveResults.Delete("group")
liveResults.Delete("other")
client.Close()
})
ctx := context.Background()
check := func(want UserRequestResult) {
assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
liveResults.Delete("group")
assert.Equal(t, want, LatestUserResults(ctx, []string{"group"})["group"])
}
recordUserResult("other", UserRequestResult{At: 2000, Success: true})
recordUserResult("group", UserRequestResult{At: 2000, Success: true})
check(UserRequestResult{At: 2000, Success: true})
recordUserResult("group", UserRequestResult{At: 2001, Success: false})
check(UserRequestResult{At: 2001, Success: false})
recordUserResult("group", UserRequestResult{At: 2000, Success: true})
check(UserRequestResult{At: 2001, Success: false})
recordUserResult("group", UserRequestResult{At: 2002, Success: true})
recordUserResult("group", UserRequestResult{At: 2001, Success: false})
check(UserRequestResult{At: 2002, Success: true})
assert.True(t, LatestUserResults(ctx, []string{"other"})["other"].Success)
ttl, err := client.TTL(ctx, "status:user-result:group").Result()
require.NoError(t, err)
assert.Equal(t, 15*time.Minute, ttl)
server.FastForward(15 * time.Minute)
assert.Zero(t, LatestUserResults(ctx, []string{"group"})["group"].At)
}
pkg/perf_metrics/metrics.go
+2 -2
View File
@@ -38,8 +38,8 @@ func RecordRelayResult(ctx context.Context, info *relaycommon.RelayInfo, apiErr
return
}
now := time.Now()
if outcome == OutcomeSuccess && !info.IsChannelTest && info.TokenId > 0 {
recordUserSuccess(info.UsingGroup, now.Unix())
if !info.IsChannelTest && info.TokenId > 0 {
recordUserResult(info.UsingGroup, UserRequestResult{At: now.UnixMilli(), Success: outcome == OutcomeSuccess})
}
hasTtft := info.IsStream && info.HasSendResponse()
ttftMs := int64(0)