feat: sync deployed v1.0.0-rc.38-custom.67
This commit is contained in:
cneicy committed
1 parent
d09216bbde
commit
8e6ed0317b
154 files changed
+39540
-20509
No files matched your search
@@ -0,0 +1,247 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/cachex"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/samber/hot"
|
||||
)
|
||||
|
||||
const poolFailoverCooldown = 10 * time.Minute
|
||||
|
||||
type poolCooldown struct {
|
||||
Until int64 `json:"until"`
|
||||
ProbeUntil int64 `json:"probe_until"`
|
||||
ProbeID string `json:"probe_id"`
|
||||
}
|
||||
|
||||
var poolCooldownOnce sync.Once
|
||||
var poolCooldownCache *cachex.HybridCache[poolCooldown]
|
||||
var poolCooldownMu sync.Mutex
|
||||
|
||||
func poolCooldowns() *cachex.HybridCache[poolCooldown] {
|
||||
poolCooldownOnce.Do(func() {
|
||||
poolCooldownCache = cachex.NewHybridCache[poolCooldown](cachex.HybridCacheConfig[poolCooldown]{
|
||||
Namespace: cachex.Namespace("new-api:pool-failover:v3"), Redis: common.RDB,
|
||||
RedisEnabled: func() bool { return common.RedisEnabled && common.RDB != nil }, RedisCodec: cachex.JSONCodec[poolCooldown]{},
|
||||
Memory: func() *hot.HotCache[string, poolCooldown] {
|
||||
return hot.NewHotCache[string, poolCooldown](hot.LRU, 100000).Build()
|
||||
},
|
||||
})
|
||||
})
|
||||
return poolCooldownCache
|
||||
}
|
||||
func poolCooldownKey(c *gin.Context, modelName, group string) string {
|
||||
return common.GenerateHMAC(fmt.Sprintf("%d:%d:%t:%d:%s:%d:%s", c.GetInt("id"), c.GetInt("token_id"), c.GetBool("pool_request_stream"), len(modelName), modelName, len(group), group))
|
||||
}
|
||||
func poolFailures(c *gin.Context) map[string]int {
|
||||
if value, ok := c.Get("pool_failures"); ok {
|
||||
return value.(map[string]int)
|
||||
}
|
||||
value := map[string]int{}
|
||||
c.Set("pool_failures", value)
|
||||
return value
|
||||
}
|
||||
|
||||
// A cooldown may exclude a group only while another eligible route exists.
|
||||
// Inspect alternatives without claiming probes or recursively checking cooldowns.
|
||||
func hasPoolAlternative(c *gin.Context, modelName, current string) bool {
|
||||
for _, entry := range ResolveRequestPool(c, modelName) {
|
||||
group := entry.Group
|
||||
if group == current || poolFailures(c)[group] >= 2 || requestGroupDegraded(c, modelName, group) {
|
||||
continue
|
||||
}
|
||||
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, group))
|
||||
now := time.Now().UnixMilli()
|
||||
if err != nil || (found && (now < value.Until || now < value.ProbeUntil)) {
|
||||
continue
|
||||
}
|
||||
channel, err := model.GetRandomSatisfiedChannel(group, modelName, 0, GetChannelConstraints(c).Filters)
|
||||
if err == nil && channel != nil {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func poolGroupSkipped(c *gin.Context, modelName, group string) bool {
|
||||
if poolFailures(c)[group] >= 2 {
|
||||
return true
|
||||
}
|
||||
if RequestRoutingPolicy(c) != nil && !hasPoolAlternative(c, modelName, group) {
|
||||
return false
|
||||
}
|
||||
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, group))
|
||||
if err != nil || !found {
|
||||
return false
|
||||
}
|
||||
now := time.Now().UnixMilli()
|
||||
skipped := now < value.Until || now < value.ProbeUntil
|
||||
if skipped {
|
||||
wait := int((max(value.Until, value.ProbeUntil) - now + 999) / 1000)
|
||||
previous := c.GetInt("pool_retry_after")
|
||||
if previous == 0 || wait < previous {
|
||||
c.Set("pool_retry_after", wait)
|
||||
}
|
||||
}
|
||||
return skipped
|
||||
}
|
||||
|
||||
// A request rejected after selection or interrupted while streaming must not
|
||||
// leave its recovery reservation blocking subsequent requests for five minutes.
|
||||
func ReleasePoolProbes(c *gin.Context) {
|
||||
raw, ok := c.Get("pool_claimed_groups")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
for group, key := range raw.(map[string]string) {
|
||||
observed, ok := c.Get("pool_probe:" + group)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
trial := observed.(poolCooldown)
|
||||
poolCooldownMu.Lock()
|
||||
current, found, err := poolCooldowns().Get(key)
|
||||
if err == nil && found && current.Until == trial.Until && current.ProbeID == trial.ProbeID {
|
||||
current.ProbeUntil, current.ProbeID = 0, ""
|
||||
_ = poolCooldowns().SetWithTTL(key, current, 0)
|
||||
}
|
||||
poolCooldownMu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// Only the selected candidate can claim a half-open probe. The expired marker
|
||||
// survives the ten-minute cooldown so a failed probe never gets a second try.
|
||||
func claimPoolGroup(c *gin.Context, modelName, group string) bool {
|
||||
if RequestRoutingPolicy(c) != nil && !hasPoolAlternative(c, modelName, group) {
|
||||
return true
|
||||
}
|
||||
poolCooldownMu.Lock()
|
||||
defer poolCooldownMu.Unlock()
|
||||
key := poolCooldownKey(c, modelName, group)
|
||||
value, found, err := poolCooldowns().Get(key)
|
||||
if err != nil || !found {
|
||||
return true
|
||||
}
|
||||
now := time.Now().UnixMilli()
|
||||
if now < value.Until || now < value.ProbeUntil {
|
||||
return false
|
||||
}
|
||||
value.ProbeUntil = now + int64((5*time.Minute)/time.Millisecond)
|
||||
value.ProbeID = fmt.Sprint(time.Now().UnixNano())
|
||||
if err := poolCooldowns().SetWithTTL(key, value, 0); err != nil {
|
||||
return false
|
||||
}
|
||||
c.Set("pool_probe:"+group, value)
|
||||
raw, ok := c.Get("pool_claimed_groups")
|
||||
claimed := map[string]string{}
|
||||
if ok {
|
||||
claimed = raw.(map[string]string)
|
||||
}
|
||||
claimed[group] = key
|
||||
c.Set("pool_claimed_groups", claimed)
|
||||
return true
|
||||
}
|
||||
|
||||
func IsPoolFailoverError(c *gin.Context, err *types.NewAPIError) bool {
|
||||
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
|
||||
return false
|
||||
}
|
||||
if err == nil || RequestRoutingPolicy(c) == nil || GetChannelConstraints(c).SuppressesRetry() || types.IsSkipRetryError(err) {
|
||||
return false
|
||||
}
|
||||
if c.Request != nil && c.Request.Context().Err() != nil {
|
||||
return false
|
||||
}
|
||||
if c.Writer != nil && c.Writer.Written() {
|
||||
return false
|
||||
}
|
||||
if err.StatusCode < 500 || err.StatusCode > 599 {
|
||||
return false
|
||||
}
|
||||
return !operation_setting.IsAlwaysSkipRetryStatusCode(err.StatusCode) && !operation_setting.IsAlwaysSkipRetryCode(err.GetErrorCode()) && operation_setting.ShouldRetryByStatusCode(err.StatusCode)
|
||||
}
|
||||
|
||||
func RecordPoolRetryableFailure(c *gin.Context, modelName, group string) {
|
||||
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
|
||||
return
|
||||
}
|
||||
if RequestRoutingPolicy(c) == nil || group == "" || GetChannelConstraints(c).SuppressesRetry() {
|
||||
return
|
||||
}
|
||||
failures := poolFailures(c)
|
||||
failures[group]++
|
||||
c.Set("pool_retry_group", group)
|
||||
if _, trial := c.Get("pool_probe:" + group); trial {
|
||||
failures[group] = 2
|
||||
}
|
||||
if failures[group] < 2 || !hasPoolAlternative(c, modelName, group) {
|
||||
return
|
||||
}
|
||||
poolCooldownMu.Lock()
|
||||
key := poolCooldownKey(c, modelName, group)
|
||||
previous, _, _ := poolCooldowns().Get(key)
|
||||
value := poolCooldown{Until: max(previous.Until+1, time.Now().Add(poolFailoverCooldown).UnixMilli())}
|
||||
if err := poolCooldowns().SetWithTTL(key, value, 0); err != nil {
|
||||
common.SysError("pool cooldown write failed: " + err.Error())
|
||||
}
|
||||
poolCooldownMu.Unlock()
|
||||
ClearPoolAffinity(c)
|
||||
}
|
||||
|
||||
func RecordPoolSuccess(c *gin.Context, modelName, group string) {
|
||||
raw, trial := c.Get("pool_probe:" + group)
|
||||
if !trial {
|
||||
return
|
||||
}
|
||||
observed := raw.(poolCooldown)
|
||||
poolCooldownMu.Lock()
|
||||
defer poolCooldownMu.Unlock()
|
||||
key := poolCooldownKey(c, modelName, group)
|
||||
current, found, err := poolCooldowns().Get(key)
|
||||
if err == nil && found && current.Until == observed.Until && current.ProbeID == observed.ProbeID {
|
||||
_, _ = poolCooldowns().DeleteMany([]string{key})
|
||||
}
|
||||
}
|
||||
|
||||
func PoolRetryBudget(c *gin.Context, modelName string) int {
|
||||
if _, pinned, _ := GetChannelConstraints(c).ResolvedPin(); pinned {
|
||||
return 0
|
||||
}
|
||||
if RequestRoutingPolicy(c) == nil || GetChannelConstraints(c).SuppressesRetry() {
|
||||
return 0
|
||||
}
|
||||
return max(0, 2*len(ResolveRequestPool(c, modelName))-1)
|
||||
}
|
||||
|
||||
func poolHigherPriorityRecovered(c *gin.Context, modelName, boundGroup string) bool {
|
||||
entries := ResolveRequestPool(c, modelName)
|
||||
priority := -1
|
||||
for _, entry := range entries {
|
||||
if entry.Group == boundGroup {
|
||||
priority = entry.Priority
|
||||
break
|
||||
}
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.Priority >= priority {
|
||||
break
|
||||
}
|
||||
value, found, err := poolCooldowns().Get(poolCooldownKey(c, modelName, entry.Group))
|
||||
if err == nil && found && time.Now().UnixMilli() >= value.Until && !poolGroupSkipped(c, modelName, entry.Group) && !requestGroupDegraded(c, modelName, entry.Group) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func PoolRetryExhausted(c *gin.Context, modelName, group string) bool {
|
||||
return poolFailures(c)[group] >= 2 && !hasPoolAlternative(c, modelName, group)
|
||||
}
|
||||
Reference in new issue
Block a user