diff --git a/controller/channel_affinity_cache.go b/controller/channel_affinity_cache.go index a72b04b8b9d6..ae764d1b7107 100644 --- a/controller/channel_affinity_cache.go +++ b/controller/channel_affinity_cache.go @@ -2,6 +2,7 @@ package controller import ( "net/http" + "strconv" "strings" "github.com/QuantumNous/new-api/service" @@ -59,6 +60,58 @@ func ClearChannelAffinityCache(c *gin.Context) { }) } +func GetChannelAffinityExclusiveCacheStats(c *gin.Context) { + stats := service.GetChannelAffinityExclusiveCacheStats() + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": stats, + }) +} + +func ClearChannelAffinityExclusiveCache(c *gin.Context) { + all := strings.TrimSpace(c.Query("all")) + channelIDStr := strings.TrimSpace(c.Query("channel_id")) + + if all == "true" { + deleted := service.ClearChannelAffinityExclusiveCacheAll() + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": gin.H{ + "deleted": deleted, + }, + }) + return + } + + if channelIDStr == "" { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "缺少参数:channel_id,或使用 all=true 清空全部", + }) + return + } + + channelID, err := strconv.Atoi(channelIDStr) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "success": false, + "message": "channel_id 必须是数字", + }) + return + } + + ok := service.ClearChannelAffinityExclusiveCacheByChannelID(channelID) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": gin.H{ + "deleted": ok, + }, + }) +} + func GetChannelAffinityUsageCacheStats(c *gin.Context) { ruleName := strings.TrimSpace(c.Query("rule_name")) usingGroup := strings.TrimSpace(c.Query("using_group")) diff --git a/dto/channel_settings.go b/dto/channel_settings.go index 8d7466d25966..7ac6052c88e6 100644 --- a/dto/channel_settings.go +++ b/dto/channel_settings.go @@ -40,6 +40,7 @@ type ChannelOtherSettings struct { UpstreamModelUpdateLastDetectedModels []string `json:"upstream_model_update_last_detected_models,omitempty"` // 上次检测到的可加入模型 UpstreamModelUpdateLastRemovedModels []string `json:"upstream_model_update_last_removed_models,omitempty"` // 上次检测到的可删除模型 UpstreamModelUpdateIgnoredModels []string `json:"upstream_model_update_ignored_models,omitempty"` // 手动忽略的模型 + AffinityExclusive bool `json:"affinity_exclusive,omitempty"` // 亲和性独占:开启后仅允许一个令牌通过亲和性绑定使用此渠道 } func (s *ChannelOtherSettings) IsOpenRouterEnterprise() bool { diff --git a/i18n/keys.go b/i18n/keys.go index 4d98540a77ce..33bdbeea5842 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -290,18 +290,19 @@ const ( // Distributor related messages const ( - MsgDistributorInvalidRequest = "distributor.invalid_request" - MsgDistributorInvalidChannelId = "distributor.invalid_channel_id" - MsgDistributorChannelDisabled = "distributor.channel_disabled" - MsgDistributorTokenNoModelAccess = "distributor.token_no_model_access" - MsgDistributorTokenModelForbidden = "distributor.token_model_forbidden" - MsgDistributorModelNameRequired = "distributor.model_name_required" - MsgDistributorInvalidPlayground = "distributor.invalid_playground_request" - MsgDistributorGroupAccessDenied = "distributor.group_access_denied" - MsgDistributorGetChannelFailed = "distributor.get_channel_failed" - MsgDistributorNoAvailableChannel = "distributor.no_available_channel" - MsgDistributorInvalidMidjourney = "distributor.invalid_midjourney_request" - MsgDistributorInvalidParseModel = "distributor.invalid_request_parse_model" + MsgDistributorInvalidRequest = "distributor.invalid_request" + MsgDistributorInvalidChannelId = "distributor.invalid_channel_id" + MsgDistributorChannelDisabled = "distributor.channel_disabled" + MsgDistributorChannelExclusiveLocked = "distributor.channel_exclusive_locked" + MsgDistributorTokenNoModelAccess = "distributor.token_no_model_access" + MsgDistributorTokenModelForbidden = "distributor.token_model_forbidden" + MsgDistributorModelNameRequired = "distributor.model_name_required" + MsgDistributorInvalidPlayground = "distributor.invalid_playground_request" + MsgDistributorGroupAccessDenied = "distributor.group_access_denied" + MsgDistributorGetChannelFailed = "distributor.get_channel_failed" + MsgDistributorNoAvailableChannel = "distributor.no_available_channel" + MsgDistributorInvalidMidjourney = "distributor.invalid_midjourney_request" + MsgDistributorInvalidParseModel = "distributor.invalid_request_parse_model" ) // Custom OAuth provider related messages diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 54dbf9181b8b..a9add0d260a7 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -245,6 +245,7 @@ common.invalid_input: "Invalid input" distributor.invalid_request: "Invalid request: {{.Error}}" distributor.invalid_channel_id: "Invalid channel ID" distributor.channel_disabled: "This channel has been disabled" +distributor.channel_exclusive_locked: "This channel is exclusively occupied by another token" distributor.token_no_model_access: "This token has no access to any models" distributor.token_model_forbidden: "This token has no access to model {{.Model}}" distributor.model_name_required: "Model name not specified, model name cannot be empty" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index 4e0b5cd15d3a..aedc9bfdf725 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -246,6 +246,7 @@ common.invalid_input: "输入不合法" distributor.invalid_request: "无效的请求,{{.Error}}" distributor.invalid_channel_id: "无效的渠道 Id" distributor.channel_disabled: "该渠道已被禁用" +distributor.channel_exclusive_locked: "该渠道已被其他令牌独占" distributor.token_no_model_access: "该令牌无权访问任何模型" distributor.token_model_forbidden: "该令牌无权访问模型 {{.Model}}" distributor.model_name_required: "未指定模型名称,模型名称不能为空" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index dcdd331b39a3..4e269d285d41 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -246,6 +246,7 @@ common.invalid_input: "輸入不合法" distributor.invalid_request: "無效的請求,{{.Error}}" distributor.invalid_channel_id: "無效的管道 Id" distributor.channel_disabled: "該管道已被禁用" +distributor.channel_exclusive_locked: "該管道已被其他令牌獨佔" distributor.token_no_model_access: "該令牌無權存取任何模型" distributor.token_model_forbidden: "該令牌無權存取模型 {{.Model}}" distributor.model_name_required: "未指定模型名稱,模型名稱不能為空" diff --git a/middleware/distributor.go b/middleware/distributor.go index d626941456c7..b5c4dcd5e370 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -51,6 +51,13 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } + if channel.GetOtherSettings().AffinityExclusive { + tokenID := c.GetInt(string(constant.ContextKeyTokenId)) + if !service.IsChannelAffinityExclusiveAvailable(channel.Id, tokenID) { + abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelExclusiveLocked)) + return + } + } } else { // Select a channel for the user // check token model mapping @@ -107,22 +114,34 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled)) return } - } else if usingGroup == "auto" { - userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) - for _, g := range autoGroups { - if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { - selectGroup = g - common.SetContextKey(c, constant.ContextKeyAutoGroup, g) + } else { + // 亲和性独占:检查该渠道是否被其他令牌锁定 + usePreferred := true + if preferred.GetOtherSettings().AffinityExclusive { + tokenID := c.GetInt(string(constant.ContextKeyTokenId)) + if !service.IsChannelAffinityExclusiveAvailable(preferred.Id, tokenID) { + usePreferred = false // 被其他令牌锁定,回退到随机选择 + } + } + if usePreferred { + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) + autoGroups := service.GetUserAutoGroup(userGroup) + for _, g := range autoGroups { + if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { + selectGroup = g + common.SetContextKey(c, constant.ContextKeyAutoGroup, g) + channel = preferred + service.MarkChannelAffinityUsed(c, g, preferred.Id) + break + } + } + } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { channel = preferred - service.MarkChannelAffinityUsed(c, g, preferred.Id) - break + selectGroup = usingGroup + service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) } } - } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { - channel = preferred - selectGroup = usingGroup - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) } } } diff --git a/model/channel_cache.go b/model/channel_cache.go index c9c503576038..6490a09cbb35 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -93,10 +93,113 @@ func SyncChannelCache(frequency int) { } } -func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) { +// ChannelFilter is a predicate that returns true if the channel should be included in selection. +type ChannelFilter func(*Channel) bool + +func passChannelFilters(channel *Channel, filters []ChannelFilter) bool { + for _, f := range filters { + if f != nil && !f(channel) { + return false + } + } + return true +} + +func getRandomSatisfiedChannelFromDB(group string, modelName string, retry int, filters []ChannelFilter) (*Channel, bool, error) { + var abilities []Ability + if err := DB.Where(commonGroupCol+" = ? and model = ? and enabled = ?", group, modelName, true). + Order("priority DESC"). + Order("weight DESC"). + Find(&abilities).Error; err != nil { + return nil, false, err + } + if len(abilities) == 0 { + return nil, false, nil + } + + abilitiesByPriority := make(map[int64][]Ability) + priorityChannels := make(map[int64][]*Channel) + channelIDs := make([]int, 0, len(abilities)) + for _, ability := range abilities { + channelIDs = append(channelIDs, ability.ChannelId) + } + + channels, err := GetChannelsByIds(channelIDs) + if err != nil { + return nil, true, err + } + channelByID := make(map[int]*Channel, len(channels)) + for _, channel := range channels { + channelByID[channel.Id] = channel + } + + availablePriorities := make([]int64, 0) + seenPriorities := make(map[int64]bool) + for _, ability := range abilities { + priority := int64(0) + if ability.Priority != nil { + priority = *ability.Priority + } + channel, ok := channelByID[ability.ChannelId] + if !ok { + return nil, true, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", ability.ChannelId) + } + if !passChannelFilters(channel, filters) { + continue + } + abilitiesByPriority[priority] = append(abilitiesByPriority[priority], ability) + priorityChannels[priority] = append(priorityChannels[priority], channel) + if !seenPriorities[priority] { + availablePriorities = append(availablePriorities, priority) + seenPriorities[priority] = true + } + } + + if len(availablePriorities) == 0 { + return nil, true, nil + } + sort.Slice(availablePriorities, func(i, j int) bool { + return availablePriorities[i] > availablePriorities[j] + }) + + if retry >= len(availablePriorities) { + retry = len(availablePriorities) - 1 + } + targetPriority := availablePriorities[retry] + targetAbilities := abilitiesByPriority[targetPriority] + targetChannels := priorityChannels[targetPriority] + if len(targetAbilities) == 0 || len(targetChannels) == 0 { + return nil, true, nil + } + + weightSum := uint(0) + for _, ability := range targetAbilities { + weightSum += ability.Weight + 10 + } + weight := common.GetRandomInt(int(weightSum)) + for idx, ability := range targetAbilities { + weight -= int(ability.Weight) + 10 + if weight <= 0 { + return targetChannels[idx], true, nil + } + } + return targetChannels[len(targetChannels)-1], true, nil +} + +func GetRandomSatisfiedChannel(group string, model string, retry int, filters ...ChannelFilter) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry) + channel, modelMatched, err := getRandomSatisfiedChannelFromDB(group, model, retry, filters) + if err != nil || channel != nil || modelMatched { + return channel, err + } + + normalizedModel := ratio_setting.FormatMatchingModelName(model) + if normalizedModel == model { + return nil, nil + } + fallbackChannel, _, fallbackErr := getRandomSatisfiedChannelFromDB(group, normalizedModel, retry, filters) + return fallbackChannel, fallbackErr } channelSyncLock.RLock() @@ -117,6 +220,9 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, if len(channels) == 1 { if channel, ok := channelsIDM[channels[0]]; ok { + if !passChannelFilters(channel, filters) { + return nil, nil + } return channel, nil } return nil, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", channels[0]) @@ -125,11 +231,17 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, uniquePriorities := make(map[int]bool) for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { + if !passChannelFilters(channel, filters) { + continue + } uniquePriorities[int(channel.GetPriority())] = true } else { return nil, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", channelId) } } + if len(uniquePriorities) == 0 { + return nil, nil + } var sortedUniquePriorities []int for priority := range uniquePriorities { sortedUniquePriorities = append(sortedUniquePriorities, priority) @@ -147,6 +259,9 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { if channel.GetPriority() == targetPriority { + if !passChannelFilters(channel, filters) { + continue + } sumWeight += channel.GetWeight() targetChannels = append(targetChannels, channel) } @@ -156,7 +271,7 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, } if len(targetChannels) == 0 { - return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority)) + return nil, nil } // smoothing factor and adjustment diff --git a/model/channel_cache_test.go b/model/channel_cache_test.go new file mode 100644 index 000000000000..79c0496fd49e --- /dev/null +++ b/model/channel_cache_test.go @@ -0,0 +1,50 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/require" +) + +func TestGetRandomSatisfiedChannelSkipsFilteredHighPriority(t *testing.T) { + prevMemoryCacheEnabled := common.MemoryCacheEnabled + + channelSyncLock.Lock() + prevGroup2Model2Channels := group2model2channels + prevChannelsIDM := channelsIDM + channelSyncLock.Unlock() + + common.MemoryCacheEnabled = true + + highPriority := int64(10) + lowPriority := int64(5) + weight := uint(1) + + channelSyncLock.Lock() + group2model2channels = map[string]map[string][]int{ + "default": { + "gpt-test": {1, 2}, + }, + } + channelsIDM = map[int]*Channel{ + 1: {Id: 1, Priority: &highPriority, Weight: &weight}, + 2: {Id: 2, Priority: &lowPriority, Weight: &weight}, + } + channelSyncLock.Unlock() + + t.Cleanup(func() { + common.MemoryCacheEnabled = prevMemoryCacheEnabled + channelSyncLock.Lock() + group2model2channels = prevGroup2Model2Channels + channelsIDM = prevChannelsIDM + channelSyncLock.Unlock() + }) + + channel, err := GetRandomSatisfiedChannel("default", "gpt-test", 0, func(channel *Channel) bool { + return channel.Id != 1 + }) + require.NoError(t, err) + require.NotNil(t, channel) + require.Equal(t, 2, channel.Id) +} diff --git a/router/api-router.go b/router/api-router.go index 35d113768be7..62a1b998585e 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -172,6 +172,8 @@ func SetApiRouter(router *gin.Engine) { optionRoute.PUT("/", controller.UpdateOption) optionRoute.GET("/channel_affinity_cache", controller.GetChannelAffinityCacheStats) optionRoute.DELETE("/channel_affinity_cache", controller.ClearChannelAffinityCache) + optionRoute.GET("/channel_affinity_exclusive_cache", controller.GetChannelAffinityExclusiveCacheStats) + optionRoute.DELETE("/channel_affinity_exclusive_cache", controller.ClearChannelAffinityExclusiveCache) optionRoute.POST("/rest_model_ratio", controller.ResetModelRatio) optionRoute.POST("/migrate_console_setting", controller.MigrateConsoleSetting) // 用于迁移检测的旧键,下个版本会删除 } diff --git a/service/channel_affinity.go b/service/channel_affinity.go index 9f89585fac03..f59b9e039e3b 100644 --- a/service/channel_affinity.go +++ b/service/channel_affinity.go @@ -1,6 +1,7 @@ package service import ( + "context" "fmt" "hash/fnv" "regexp" @@ -10,10 +11,13 @@ import ( "time" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/cachex" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/QuantumNous/new-api/types" + "github.com/go-redis/redis/v8" "github.com/gin-gonic/gin" "github.com/samber/hot" "github.com/tidwall/gjson" @@ -28,6 +32,8 @@ const ( channelAffinityCacheNamespace = "new-api:channel_affinity:v1" channelAffinityUsageCacheStatsNamespace = "new-api:channel_affinity_usage_cache_stats:v1" + channelAffinityExclusiveNamespace = "new-api:channel_affinity_exclusive:v1" + channelAffinityExclusiveBindingNamespace = "new-api:channel_affinity_exclusive_binding:v1" ) var ( @@ -37,9 +43,22 @@ var ( channelAffinityUsageCacheStatsOnce sync.Once channelAffinityUsageCacheStatsCache *cachex.HybridCache[ChannelAffinityUsageCacheCounters] + channelAffinityExclusiveCacheOnce sync.Once + channelAffinityExclusiveCache *cachex.HybridCache[int] // channelID -> tokenID + + channelAffinityExclusiveBindingCacheOnce sync.Once + channelAffinityExclusiveBindingCache *cachex.HybridCache[ChannelAffinityExclusiveBinding] + + channelAffinityExclusiveMemoryLock sync.Mutex + channelAffinityRegexCache sync.Map // map[string]*regexp.Regexp ) +type ChannelAffinityExclusiveBinding struct { + ChannelID int `json:"channel_id"` + TokenID int `json:"token_id"` +} + type channelAffinityMeta struct { CacheKey string TTLSeconds int @@ -108,6 +127,410 @@ func getChannelAffinityCache() *cachex.HybridCache[int] { return channelAffinityCache } +func getChannelAffinityExclusiveCache() *cachex.HybridCache[int] { + channelAffinityExclusiveCacheOnce.Do(func() { + setting := operation_setting.GetChannelAffinitySetting() + defaultTTLSeconds := setting.DefaultTTLSeconds + if defaultTTLSeconds <= 0 { + defaultTTLSeconds = 3600 + } + + channelAffinityExclusiveCache = cachex.NewHybridCache[int](cachex.HybridCacheConfig[int]{ + Namespace: cachex.Namespace(channelAffinityExclusiveNamespace), + Redis: common.RDB, + RedisEnabled: func() bool { + return common.RedisEnabled && common.RDB != nil + }, + RedisCodec: cachex.IntCodec{}, + Memory: func() *hot.HotCache[string, int] { + return hot.NewHotCache[string, int](hot.LRU, 10_000). + WithTTL(time.Duration(defaultTTLSeconds) * time.Second). + WithJanitor(). + Build() + }, + }) + }) + return channelAffinityExclusiveCache +} + +func getChannelAffinityExclusiveBindingCache() *cachex.HybridCache[ChannelAffinityExclusiveBinding] { + channelAffinityExclusiveBindingCacheOnce.Do(func() { + setting := operation_setting.GetChannelAffinitySetting() + capacity := setting.MaxEntries + if capacity <= 0 { + capacity = 100_000 + } + defaultTTLSeconds := setting.DefaultTTLSeconds + if defaultTTLSeconds <= 0 { + defaultTTLSeconds = 3600 + } + + channelAffinityExclusiveBindingCache = cachex.NewHybridCache[ChannelAffinityExclusiveBinding](cachex.HybridCacheConfig[ChannelAffinityExclusiveBinding]{ + Namespace: cachex.Namespace(channelAffinityExclusiveBindingNamespace), + Redis: common.RDB, + RedisEnabled: func() bool { + return common.RedisEnabled && common.RDB != nil + }, + RedisCodec: cachex.JSONCodec[ChannelAffinityExclusiveBinding]{}, + Memory: func() *hot.HotCache[string, ChannelAffinityExclusiveBinding] { + return hot.NewHotCache[string, ChannelAffinityExclusiveBinding](hot.LRU, capacity). + WithTTL(time.Duration(defaultTTLSeconds) * time.Second). + WithJanitor(). + Build() + }, + }) + }) + return channelAffinityExclusiveBindingCache +} + +func normalizeChannelAffinityCacheKey(cacheKey string) string { + cacheKey = strings.TrimSpace(cacheKey) + return strings.TrimPrefix(cacheKey, channelAffinityCacheNamespace+":") +} + +func channelAffinityExclusiveRedisEnabled() bool { + return common.RedisEnabled && common.RDB != nil +} + +// TrySetChannelAffinityExclusiveLock attempts to set an exclusive lock for the given channel. +// Returns true if the lock was set (no existing lock or same token holds it). +// Returns false if the channel is locked by a different token. +func TrySetChannelAffinityExclusiveLock(channelID, tokenID int, ttl time.Duration) bool { + if channelID <= 0 || tokenID <= 0 { + return false + } + cache := getChannelAffinityExclusiveCache() + key := strconv.Itoa(channelID) + if ttl <= 0 { + ttl = time.Second + } + + if channelAffinityExclusiveRedisEnabled() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + fullKey := cache.FullKey(key) + script := ` +local current = redis.call("GET", KEYS[1]) +if not current then + redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2]) + return 1 +end +if current == ARGV[1] then + redis.call("SET", KEYS[1], ARGV[1], "PX", ARGV[2]) + return 1 +end +return 0 +` + res, err := common.RDB.Eval(ctx, script, []string{fullKey}, strconv.Itoa(tokenID), ttl.Milliseconds()).Int64() + if err != nil && err != redis.Nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock set failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + return res == 1 + } + + channelAffinityExclusiveMemoryLock.Lock() + defer channelAffinityExclusiveMemoryLock.Unlock() + + existingTokenID, found, err := cache.Get(key) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock get failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + if found && existingTokenID != tokenID { + return false + } + + if err := cache.SetWithTTL(key, tokenID, ttl); err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock set failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + return true +} + +func deleteChannelAffinityExclusiveLockIfOwned(channelID, tokenID int) bool { + if channelID <= 0 || tokenID <= 0 { + return false + } + cache := getChannelAffinityExclusiveCache() + key := strconv.Itoa(channelID) + + if channelAffinityExclusiveRedisEnabled() { + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + + fullKey := cache.FullKey(key) + script := ` +local current = redis.call("GET", KEYS[1]) +if current == ARGV[1] then + return redis.call("DEL", KEYS[1]) +end +return 0 +` + res, err := common.RDB.Eval(ctx, script, []string{fullKey}, strconv.Itoa(tokenID)).Int64() + if err != nil && err != redis.Nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock delete failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + return res > 0 + } + + channelAffinityExclusiveMemoryLock.Lock() + defer channelAffinityExclusiveMemoryLock.Unlock() + + existingTokenID, found, err := cache.Get(key) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock get failed on delete: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + if !found || existingTokenID != tokenID { + return false + } + result, err := cache.DeleteMany([]string{key}) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock delete failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return false + } + return result[cache.FullKey(key)] +} + +// IsChannelAffinityExclusiveAvailable checks if a channel is available for the given token. +// Returns true if: no exclusive lock exists, or the lock is held by the same token. +// When currentTokenID <= 0, the channel is only available when no lock exists. +func IsChannelAffinityExclusiveAvailable(channelID, currentTokenID int) bool { + cache := getChannelAffinityExclusiveCache() + key := strconv.Itoa(channelID) + lockedTokenID, found, err := cache.Get(key) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock get failed: channel=%d, err=%v", channelID, err)) + return false + } + if !found { + return true + } + return currentTokenID > 0 && lockedTokenID == currentTokenID +} + +// GetChannelAffinityExclusiveLockHolder returns the token ID holding the exclusive lock, if any. +func GetChannelAffinityExclusiveLockHolder(channelID int) (int, bool) { + cache := getChannelAffinityExclusiveCache() + key := strconv.Itoa(channelID) + tokenID, found, err := cache.Get(key) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive lock get failed: channel=%d, err=%v", channelID, err)) + return 0, false + } + return tokenID, found +} + +func getChannelAffinityExclusiveBinding(cacheKey string) (ChannelAffinityExclusiveBinding, bool) { + cache := getChannelAffinityExclusiveBindingCache() + rawKey := normalizeChannelAffinityCacheKey(cacheKey) + if rawKey == "" { + return ChannelAffinityExclusiveBinding{}, false + } + binding, found, err := cache.Get(rawKey) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding get failed: affinity_key=%s, err=%v", rawKey, err)) + return ChannelAffinityExclusiveBinding{}, false + } + return binding, found +} + +func setChannelAffinityExclusiveBinding(cacheKey string, binding ChannelAffinityExclusiveBinding, ttl time.Duration) error { + cache := getChannelAffinityExclusiveBindingCache() + rawKey := normalizeChannelAffinityCacheKey(cacheKey) + if rawKey == "" { + return nil + } + return cache.SetWithTTL(rawKey, binding, ttl) +} + +func hasOtherChannelAffinityExclusiveBinding(channelID, tokenID int) bool { + if channelID <= 0 || tokenID <= 0 { + return false + } + cache := getChannelAffinityExclusiveBindingCache() + keys, err := cache.Keys() + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding keys failed: channel=%d, token=%d, err=%v", channelID, tokenID, err)) + return true + } + prefix := channelAffinityExclusiveBindingNamespace + ":" + for _, key := range keys { + rawKey := strings.TrimPrefix(key, prefix) + binding, found, err := cache.Get(rawKey) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding get failed: key=%s, err=%v", rawKey, err)) + return true + } + if found && binding.ChannelID == channelID && binding.TokenID == tokenID { + return true + } + } + return false +} + +func maybeReleaseChannelAffinityExclusiveLock(channelID, tokenID int) { + if channelID <= 0 || tokenID <= 0 { + return + } + if hasOtherChannelAffinityExclusiveBinding(channelID, tokenID) { + return + } + _ = deleteChannelAffinityExclusiveLockIfOwned(channelID, tokenID) +} + +func clearChannelAffinityExclusiveBindingCacheAll() int { + cache := getChannelAffinityExclusiveBindingCache() + keys, err := cache.Keys() + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache list keys failed: err=%v", err)) + return 0 + } + if len(keys) > 0 { + if _, err := cache.DeleteMany(keys); err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache delete many failed: err=%v", err)) + } + } + return len(keys) +} + +func cleanupChannelAffinityExclusiveBindingsByAffinityKeys(cacheKeys []string) { + if len(cacheKeys) == 0 { + return + } + cache := getChannelAffinityExclusiveBindingCache() + released := make(map[string]ChannelAffinityExclusiveBinding, len(cacheKeys)) + rawKeys := make([]string, 0, len(cacheKeys)) + + for _, cacheKey := range cacheKeys { + rawKey := normalizeChannelAffinityCacheKey(cacheKey) + if rawKey == "" { + continue + } + if binding, found := getChannelAffinityExclusiveBinding(rawKey); found { + released[fmt.Sprintf("%d:%d", binding.ChannelID, binding.TokenID)] = binding + } + rawKeys = append(rawKeys, rawKey) + } + + if len(rawKeys) > 0 { + if _, err := cache.DeleteMany(rawKeys); err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache delete many failed: err=%v", err)) + return + } + } + + for _, binding := range released { + maybeReleaseChannelAffinityExclusiveLock(binding.ChannelID, binding.TokenID) + } +} + +func clearChannelAffinityExclusiveBindingCacheByChannelID(channelID int) int { + if channelID <= 0 { + return 0 + } + cache := getChannelAffinityExclusiveBindingCache() + keys, err := cache.Keys() + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache list keys failed: channel=%d, err=%v", channelID, err)) + return 0 + } + prefix := channelAffinityExclusiveBindingNamespace + ":" + targetKeys := make([]string, 0) + for _, key := range keys { + rawKey := strings.TrimPrefix(key, prefix) + binding, found, err := cache.Get(rawKey) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache get failed: key=%s, err=%v", rawKey, err)) + continue + } + if found && binding.ChannelID == channelID { + targetKeys = append(targetKeys, rawKey) + } + } + if len(targetKeys) == 0 { + return 0 + } + result, err := cache.DeleteMany(targetKeys) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding cache delete many failed: channel=%d, err=%v", channelID, err)) + return 0 + } + deleted := 0 + for _, ok := range result { + if ok { + deleted++ + } + } + return deleted +} + +type ChannelAffinityExclusiveCacheStats struct { + Total int `json:"total"` + Entries map[int]int `json:"entries"` // channelID -> tokenID +} + +func GetChannelAffinityExclusiveCacheStats() ChannelAffinityExclusiveCacheStats { + cache := getChannelAffinityExclusiveCache() + keys, err := cache.Keys() + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive cache list keys failed: err=%v", err)) + return ChannelAffinityExclusiveCacheStats{Total: 0, Entries: map[int]int{}} + } + entries := make(map[int]int, len(keys)) + for _, k := range keys { + suffix := strings.TrimPrefix(k, channelAffinityExclusiveNamespace+":") + channelID, err := strconv.Atoi(suffix) + if err != nil { + continue + } + tokenID, found, _ := cache.Get(suffix) + if found { + entries[channelID] = tokenID + } + } + return ChannelAffinityExclusiveCacheStats{ + Total: len(entries), + Entries: entries, + } +} + +func ClearChannelAffinityExclusiveCacheAll() int { + cache := getChannelAffinityExclusiveCache() + keys, err := cache.Keys() + if err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive cache list keys failed: err=%v", err)) + clearChannelAffinityExclusiveBindingCacheAll() + return 0 + } + if len(keys) > 0 { + if _, err := cache.DeleteMany(keys); err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive cache delete many failed: err=%v", err)) + } + } + clearChannelAffinityExclusiveBindingCacheAll() + return len(keys) +} + +func ClearChannelAffinityExclusiveCacheByChannelID(channelID int) bool { + cache := getChannelAffinityExclusiveCache() + key := strconv.Itoa(channelID) + bindingDeleted := clearChannelAffinityExclusiveBindingCacheByChannelID(channelID) + result, err := cache.DeleteMany([]string{key}) + if err != nil { + return bindingDeleted > 0 + } + for _, deleted := range result { + if deleted { + return true + } + } + return bindingDeleted > 0 +} + func GetChannelAffinityCacheStats() ChannelAffinityCacheStats { setting := operation_setting.GetChannelAffinitySetting() if setting == nil { @@ -197,6 +620,7 @@ func ClearChannelAffinityCacheAll() int { common.SysError(fmt.Sprintf("channel affinity cache delete many failed: err=%v", err)) } } + ClearChannelAffinityExclusiveCacheAll() return len(keys) } @@ -228,10 +652,37 @@ func ClearChannelAffinityCacheByRuleName(ruleName string) (int, error) { } cache := getChannelAffinityCache() - deleted, err := cache.DeleteByPrefix(ruleName) + keys, err := cache.Keys() + if err != nil { + return 0, err + } + + prefix := channelAffinityCacheNamespace + ":" + ruleName + ":" + matchedFullKeys := make([]string, 0) + matchedAffinityKeys := make([]string, 0) + for _, key := range keys { + if !strings.HasPrefix(key, prefix) { + continue + } + matchedFullKeys = append(matchedFullKeys, key) + matchedAffinityKeys = append(matchedAffinityKeys, normalizeChannelAffinityCacheKey(key)) + } + if len(matchedFullKeys) == 0 { + return 0, nil + } + + result, err := cache.DeleteMany(matchedFullKeys) if err != nil { return 0, err } + cleanupChannelAffinityExclusiveBindingsByAffinityKeys(matchedAffinityKeys) + + deleted := 0 + for _, ok := range result { + if ok { + deleted++ + } + } return deleted, nil } @@ -683,9 +1134,51 @@ func RecordChannelAffinity(c *gin.Context, channelID int) { if ttlSeconds <= 0 { ttlSeconds = 3600 } + ttl := time.Duration(ttlSeconds) * time.Second + affinityCacheKey := normalizeChannelAffinityCacheKey(cacheKey) + if affinityCacheKey == "" { + return + } + oldBinding, oldBindingFound := getChannelAffinityExclusiveBinding(affinityCacheKey) + + tokenID := c.GetInt(string(constant.ContextKeyTokenId)) + shouldBindExclusive := false + newBinding := ChannelAffinityExclusiveBinding{} + ch, err := model.CacheGetChannel(channelID) + if err == nil && ch != nil && ch.GetOtherSettings().AffinityExclusive && tokenID > 0 { + shouldBindExclusive = true + newBinding = ChannelAffinityExclusiveBinding{ + ChannelID: channelID, + TokenID: tokenID, + } + if !TrySetChannelAffinityExclusiveLock(channelID, tokenID, ttl) { + common.SysLog(fmt.Sprintf("channel affinity exclusive lock skipped: channel=%d, token=%d", channelID, tokenID)) + return + } + } + cache := getChannelAffinityCache() - if err := cache.SetWithTTL(cacheKey, channelID, time.Duration(ttlSeconds)*time.Second); err != nil { + if err := cache.SetWithTTL(cacheKey, channelID, ttl); err != nil { common.SysError(fmt.Sprintf("channel affinity cache set failed: key=%s, err=%v", cacheKey, err)) + if shouldBindExclusive && (!oldBindingFound || oldBinding != newBinding) { + _ = deleteChannelAffinityExclusiveLockIfOwned(channelID, tokenID) + } + return + } + + if shouldBindExclusive { + if err := setChannelAffinityExclusiveBinding(affinityCacheKey, newBinding, ttl); err != nil { + common.SysError(fmt.Sprintf("channel affinity exclusive binding set failed: affinity_key=%s, channel=%d, token=%d, err=%v", affinityCacheKey, channelID, tokenID, err)) + return + } + if oldBindingFound && oldBinding != newBinding { + maybeReleaseChannelAffinityExclusiveLock(oldBinding.ChannelID, oldBinding.TokenID) + } + return + } + + if oldBindingFound { + cleanupChannelAffinityExclusiveBindingsByAffinityKeys([]string{affinityCacheKey}) } } diff --git a/service/channel_affinity_exclusive_test.go b/service/channel_affinity_exclusive_test.go new file mode 100644 index 000000000000..f2d61b9cf183 --- /dev/null +++ b/service/channel_affinity_exclusive_test.go @@ -0,0 +1,86 @@ +package service + +import ( + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/stretchr/testify/require" +) + +func TestChannelAffinityExclusiveLockRejectsOtherToken(t *testing.T) { + prevRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + t.Cleanup(func() { + common.RedisEnabled = prevRedisEnabled + ClearChannelAffinityExclusiveCacheAll() + }) + + ClearChannelAffinityExclusiveCacheAll() + + require.True(t, TrySetChannelAffinityExclusiveLock(101, 1001, time.Minute)) + require.True(t, IsChannelAffinityExclusiveAvailable(101, 1001)) + require.False(t, IsChannelAffinityExclusiveAvailable(101, 1002)) + require.False(t, IsChannelAffinityExclusiveAvailable(101, 0)) + require.False(t, TrySetChannelAffinityExclusiveLock(101, 1002, time.Minute)) + + holder, found := GetChannelAffinityExclusiveLockHolder(101) + require.True(t, found) + require.Equal(t, 1001, holder) +} + +func TestCleanupChannelAffinityExclusiveBindingsByAffinityKeysKeepsLockUntilLastBindingRemoved(t *testing.T) { + prevRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + t.Cleanup(func() { + common.RedisEnabled = prevRedisEnabled + ClearChannelAffinityExclusiveCacheAll() + }) + + ClearChannelAffinityExclusiveCacheAll() + + ttl := time.Minute + require.True(t, TrySetChannelAffinityExclusiveLock(201, 2001, ttl)) + require.NoError(t, setChannelAffinityExclusiveBinding("rule-a:key-1", ChannelAffinityExclusiveBinding{ChannelID: 201, TokenID: 2001}, ttl)) + require.NoError(t, setChannelAffinityExclusiveBinding("rule-a:key-2", ChannelAffinityExclusiveBinding{ChannelID: 201, TokenID: 2001}, ttl)) + + cleanupChannelAffinityExclusiveBindingsByAffinityKeys([]string{"rule-a:key-1"}) + require.False(t, IsChannelAffinityExclusiveAvailable(201, 3001)) + + cleanupChannelAffinityExclusiveBindingsByAffinityKeys([]string{"rule-a:key-2"}) + require.True(t, IsChannelAffinityExclusiveAvailable(201, 3001)) +} + +func TestClearChannelAffinityCacheByRuleNameAlsoClearsExclusiveBinding(t *testing.T) { + prevRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + t.Cleanup(func() { + common.RedisEnabled = prevRedisEnabled + ClearChannelAffinityExclusiveCacheAll() + _ = ClearChannelAffinityCacheAll() + }) + + ClearChannelAffinityExclusiveCacheAll() + _ = ClearChannelAffinityCacheAll() + + var matchedRule *operation_setting.ChannelAffinityRule + for i := range operation_setting.GetChannelAffinitySetting().Rules { + rule := &operation_setting.GetChannelAffinitySetting().Rules[i] + if rule.IncludeRuleName { + matchedRule = rule + break + } + } + require.NotNil(t, matchedRule) + + cacheKey := buildChannelAffinityCacheKeySuffix(*matchedRule, "default", "exclusive-test") + require.NoError(t, getChannelAffinityCache().SetWithTTL(cacheKey, 301, time.Minute)) + require.True(t, TrySetChannelAffinityExclusiveLock(301, 3001, time.Minute)) + require.NoError(t, setChannelAffinityExclusiveBinding(cacheKey, ChannelAffinityExclusiveBinding{ChannelID: 301, TokenID: 3001}, time.Minute)) + + deleted, err := ClearChannelAffinityCacheByRuleName(matchedRule.Name) + require.NoError(t, err) + require.Equal(t, 1, deleted) + require.True(t, IsChannelAffinityExclusiveAvailable(301, 4001)) +} diff --git a/service/channel_select.go b/service/channel_select.go index a3710ef8cec3..46110c42af51 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -86,6 +86,19 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, selectGroup := param.TokenGroup userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) + // 构建亲和性独占过滤器 + var filters []model.ChannelFilter + currentTokenID := 0 + if param.Ctx != nil { + currentTokenID = param.Ctx.GetInt(string(constant.ContextKeyTokenId)) + } + filters = append(filters, func(ch *model.Channel) bool { + if !ch.GetOtherSettings().AffinityExclusive { + return true + } + return IsChannelAffinityExclusiveAvailable(ch.Id, currentTokenID) + }) + if param.TokenGroup == "auto" { if len(setting.GetAutoGroups()) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") @@ -115,7 +128,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry) + channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry, filters...) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道,尝试下一个分组 @@ -153,7 +166,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry()) + channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry(), filters...) if err != nil { return nil, param.TokenGroup, err } diff --git a/web/src/components/table/channels/modals/EditChannelModal.jsx b/web/src/components/table/channels/modals/EditChannelModal.jsx index 899e290594b8..e3c635804f16 100644 --- a/web/src/components/table/channels/modals/EditChannelModal.jsx +++ b/web/src/components/table/channels/modals/EditChannelModal.jsx @@ -214,6 +214,7 @@ const EditChannelModal = (props) => { upstream_model_update_last_check_time: 0, upstream_model_update_last_detected_models: [], upstream_model_update_ignored_models: '', + affinity_exclusive: false, }; const [batch, setBatch] = useState(false); const [multiToSingle, setMultiToSingle] = useState(false); @@ -891,6 +892,7 @@ const EditChannelModal = (props) => { data.allow_inference_geo = parsedSettings.allow_inference_geo || false; data.claude_beta_query = parsedSettings.claude_beta_query || false; + data.affinity_exclusive = parsedSettings.affinity_exclusive || false; data.upstream_model_update_check_enabled = parsedSettings.upstream_model_update_check_enabled === true; data.upstream_model_update_auto_sync_enabled = @@ -920,6 +922,7 @@ const EditChannelModal = (props) => { data.allow_include_obfuscation = false; data.allow_inference_geo = false; data.claude_beta_query = false; + data.affinity_exclusive = false; data.upstream_model_update_check_enabled = false; data.upstream_model_update_auto_sync_enabled = false; data.upstream_model_update_last_check_time = 0; @@ -937,6 +940,7 @@ const EditChannelModal = (props) => { data.allow_include_obfuscation = false; data.allow_inference_geo = false; data.claude_beta_query = false; + data.affinity_exclusive = false; data.upstream_model_update_check_enabled = false; data.upstream_model_update_auto_sync_enabled = false; data.upstream_model_update_last_check_time = 0; @@ -1015,6 +1019,7 @@ const EditChannelModal = (props) => { data.pass_through_body_enabled || data.force_format || data.claude_beta_query || + data.affinity_exclusive || data.system_prompt_override; if (hasAdvancedValues) { setAdvancedSettingsOpen(true); @@ -1780,6 +1785,8 @@ const EditChannelModal = (props) => { } } + settings.affinity_exclusive = localInputs.affinity_exclusive === true; + settings.upstream_model_update_check_enabled = localInputs.upstream_model_update_check_enabled === true; settings.upstream_model_update_auto_sync_enabled = @@ -1824,6 +1831,7 @@ const EditChannelModal = (props) => { delete localInputs.allow_include_obfuscation; delete localInputs.allow_inference_geo; delete localInputs.claude_beta_query; + delete localInputs.affinity_exclusive; delete localInputs.upstream_model_update_check_enabled; delete localInputs.upstream_model_update_auto_sync_enabled; delete localInputs.upstream_model_update_last_check_time; @@ -2505,6 +2513,7 @@ const EditChannelModal = (props) => {