From 5e81e19bc81e88d5df15a04f6a6268886127e002 Mon Sep 17 00:00:00 2001 From: JustSong Date: Wed, 27 Mar 2024 19:09:27 +0800 Subject: [PATCH] fix: fix SQL channel selection algo (#1197) --- model/ability.go | 12 +++++++++--- model/cache.go | 2 +- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/model/ability.go b/model/ability.go index 7127abc395..48b856a2f3 100644 --- a/model/ability.go +++ b/model/ability.go @@ -2,6 +2,7 @@ package model import ( "github.com/songquanpeng/one-api/common" + "gorm.io/gorm" "strings" ) @@ -13,7 +14,7 @@ type Ability struct { Priority *int64 `json:"priority" gorm:"bigint;default:0;index"` } -func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) { +func GetRandomSatisfiedChannel(group string, model string, ignoreFirstPriority bool) (*Channel, error) { ability := Ability{} groupCol := "`group`" trueVal := "1" @@ -23,8 +24,13 @@ func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) { } var err error = nil - maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model) - channelQuery := DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal+" and priority = (?)", group, model, maxPrioritySubQuery) + var channelQuery *gorm.DB + if ignoreFirstPriority { + channelQuery = DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model) + } else { + maxPrioritySubQuery := DB.Model(&Ability{}).Select("MAX(priority)").Where(groupCol+" = ? and model = ? and enabled = "+trueVal, group, model) + channelQuery = DB.Where(groupCol+" = ? and model = ? and enabled = "+trueVal+" and priority = (?)", group, model, maxPrioritySubQuery) + } if common.UsingSQLite || common.UsingPostgreSQL { err = channelQuery.Order("RANDOM()").First(&ability).Error } else { diff --git a/model/cache.go b/model/cache.go index dd20d857b8..244fe6acf7 100644 --- a/model/cache.go +++ b/model/cache.go @@ -205,7 +205,7 @@ func SyncChannelCache(frequency int) { func CacheGetRandomSatisfiedChannel(group string, model string, ignoreFirstPriority bool) (*Channel, error) { if !config.MemoryCacheEnabled { - return GetRandomSatisfiedChannel(group, model) + return GetRandomSatisfiedChannel(group, model, ignoreFirstPriority) } channelSyncLock.RLock() defer channelSyncLock.RUnlock()