255 lines
8.2 KiB
Go
255 lines
8.2 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"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 == http.StatusForbidden {
|
|
// Only upstream refusals participate in pool failover. Local permission
|
|
// failures must not gain the pool's retry budget or override strict affinity.
|
|
if err.GetErrorType() == types.ErrorTypeNewAPIError {
|
|
return false
|
|
}
|
|
} else 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)
|
|
}
|