new-api-custom/service/pool_failover.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)
}