From 811f187d6af15b25cfd910dc9c89e414cc62060a Mon Sep 17 00:00:00 2001 From: StageDog Date: Mon, 1 Sep 2025 12:18:28 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E6=B8=A0=E9=81=93=E7=BA=A7=E5=88=AB?= =?UTF-8?q?=E9=99=90=E5=88=B6=E6=AF=8F=E5=88=86=E9=92=9F=E6=9C=80=E5=A4=A7?= =?UTF-8?q?=E8=AF=B7=E6=B1=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- constant/context_key.go | 1 + middleware/distributor.go | 1 + model/ability.go | 144 +++++++++++++++++++++++---- model/channel.go | 8 ++ model/channel_cache.go | 89 +++++++++++++---- relay/common/relay_info.go | 3 + web/src/pages/Channel/EditChannel.js | 11 ++ 7 files changed, 220 insertions(+), 37 deletions(-) diff --git a/constant/context_key.go b/constant/context_key.go index 4eaf3d007547..8a88a4abfba4 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -31,6 +31,7 @@ const ( ContextKeyChannelIsMultiKey ContextKey = "channel_is_multi_key" ContextKeyChannelMultiKeyIndex ContextKey = "channel_multi_key_index" ContextKeyChannelKey ContextKey = "channel_key" + ContextKeyChannelRateLimit ContextKey = "channel_rate_limit" /* user related keys */ ContextKeyUserId ContextKey = "id" diff --git a/middleware/distributor.go b/middleware/distributor.go index a6889e396465..c15a81781761 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -267,6 +267,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan()) common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping()) common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping()) + common.SetContextKey(c, constant.ContextKeyChannelRateLimit, channel.GetRateLimit()) key, index, newAPIError := channel.GetNextEnabledKey() if newAPIError != nil { diff --git a/model/ability.go b/model/ability.go index f36ff76421a9..c8155f76a880 100644 --- a/model/ability.go +++ b/model/ability.go @@ -6,6 +6,7 @@ import ( "one-api/common" "strings" "sync" + "time" "github.com/samber/lo" "gorm.io/gorm" @@ -102,6 +103,81 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { return channelQuery, nil } +var channelRateLimitStatus sync.Map // 存储每个 Channel 的频率限制状态 +var rateLimitMutex sync.Mutex + +type ChannelRateLimit struct { + Count int64 // 使用次数 + ResetTime time.Time // 上次重置时间 +} + +type ChannelModelKey struct { + ChannelID int +} + +func isRateLimited(channel Channel, channelId int) bool { + if (channel.RateLimit != nil && *channel.RateLimit > 0) { + if _, ok := checkRateLimit(&channel, channelId); !ok { + return true + } + updateRateLimitStatus(channelId) + } + return false +} + + +func checkRateLimit(channel *Channel, channelId int) (*ChannelRateLimit, bool) { + now := time.Now() + key := ChannelModelKey{ChannelID: channelId} + + rateLimitMutex.Lock() + defer rateLimitMutex.Unlock() + + value, exists := channelRateLimitStatus.Load(key) + if !exists { + value = &ChannelRateLimit{ + Count: 1, + ResetTime: now.Add(time.Minute), + } + channelRateLimitStatus.Store(key, value) + return value.(*ChannelRateLimit), true + } + rateLimit := value.(*ChannelRateLimit) + if now.After(rateLimit.ResetTime) { + rateLimit.Count = 1 + rateLimit.ResetTime = now.Add(time.Minute) + return rateLimit, true + } else if int64(*channel.RateLimit) > rateLimit.Count { + rateLimit.Count++ + return rateLimit, true + } + + return rateLimit, false +} + +func updateRateLimitStatus(channelId int) { + now := time.Now() + key := ChannelModelKey{ChannelID: channelId} + + rateLimitMutex.Lock() + defer rateLimitMutex.Unlock() + + val, _ := channelRateLimitStatus.Load(key) + if val == nil { + return + } + + rl := val.(*ChannelRateLimit) + if now.After(rl.ResetTime) { + rl.Count = 1 + rl.ResetTime = now.Add(time.Minute) + } else { + rl.Count++ + } + + channelRateLimitStatus.Store(key, rl) +} + func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) { var abilities []Ability @@ -118,28 +194,62 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, if err != nil { return nil, err } + if len(abilities) <= 0 { + return nil, errors.New("channel not found"); + } + channel := Channel{} - if len(abilities) > 0 { - // Randomly choose one - weightSum := uint(0) - for _, ability_ := range abilities { - weightSum += ability_.Weight + 10 + for len(abilities) > 0 { + selectedIndex, err := getRandomWeightedIndex(abilities) + if err != nil { + return nil, err } - // Randomly choose one - weight := common.GetRandomInt(int(weightSum)) - for _, ability_ := range abilities { - weight -= int(ability_.Weight) + 10 - //log.Printf("weight: %d, ability weight: %d", weight, *ability_.Weight) - if weight <= 0 { - channel.Id = ability_.ChannelId - break + + selectedAbility := abilities[selectedIndex] + channelPtr, err := GetChannelById(selectedAbility.ChannelId, true) + if err != nil { + if err.Error() != "channel not found" { + return nil, err } + abilities = removeAbility(abilities, selectedIndex) + continue } - } else { - return nil, errors.New("channel not found") + + channel = *channelPtr + if isRateLimited(channel, channel.Id) { + abilities = removeAbility(abilities, selectedIndex) + continue + } + + return channelPtr, nil + } + + return nil, errors.New("channel not found") +} + +func getRandomWeightedIndex(abilities []Ability) (int, error) { + weightSum := uint(0) + for _, ability := range abilities { + weightSum += ability.Weight + } + + if weightSum == 0 { + return common.GetRandomInt(len(abilities)), nil } - err = DB.First(&channel, "id = ?", channel.Id).Error - return &channel, err + + randomWeight := common.GetRandomInt(int(weightSum)) + for i, ability := range abilities { + randomWeight -= int(ability.Weight) + if randomWeight <= 0 { + return i, nil + } + } + + return -1, errors.New("unable to select a random weighted index") +} + +func removeAbility(abilities []Ability, index int) []Ability { + return append(abilities[:index], abilities[index+1:]...) } func (channel *Channel) AddAbilities() error { diff --git a/model/channel.go b/model/channel.go index 6277fcda2b32..888948d8363e 100644 --- a/model/channel.go +++ b/model/channel.go @@ -44,6 +44,7 @@ type Channel struct { Tag *string `json:"tag" gorm:"index"` Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置 ParamOverride *string `json:"param_override" gorm:"type:text"` + RateLimit *int `json:"rate_limit" gorm:"default:0"` // add after v0.8.5 ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"` } @@ -397,6 +398,13 @@ func (channel *Channel) GetStatusCodeMapping() string { return *channel.StatusCodeMapping } +func (channel *Channel) GetRateLimit() int { + if channel.RateLimit == nil || *channel.RateLimit <= 0 { + return 0 + } + return *channel.RateLimit +} + func (channel *Channel) Insert() error { var err error err = DB.Create(channel).Error diff --git a/model/channel_cache.go b/model/channel_cache.go index b24512489793..df6bc4203917 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -7,11 +7,13 @@ import ( "one-api/common" "one-api/setting" "sort" + "strconv" "strings" "sync" "time" "github.com/gin-gonic/gin" + "github.com/go-redis/redis/v8" ) var group2model2channels map[string]map[string][]int // enabled channel @@ -130,26 +132,30 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, channelSyncLock.RLock() defer channelSyncLock.RUnlock() - channels := group2model2channels[group][model] + channelIds := group2model2channels[group][model] - if len(channels) == 0 { + validChannels := make([]*Channel, 0) + for _, channelId := range channelIds { + if channel, ok := channelsIDM[channelId]; ok { + if !isRedisLimited(*channel, channelId) { + validChannels = append(validChannels, channel) + } + } else { + return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) + } + } + + if len(validChannels) == 0 { return nil, errors.New("channel not found") } - if len(channels) == 1 { - if channel, ok := channelsIDM[channels[0]]; ok { - return channel, nil - } - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channels[0]) + if len(validChannels) == 1 { + return validChannels[0], nil } uniquePriorities := make(map[int]bool) - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - uniquePriorities[int(channel.GetPriority())] = true - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) - } + for _, channel := range validChannels { + uniquePriorities[int(channel.GetPriority())] = true } var sortedUniquePriorities []int for priority := range uniquePriorities { @@ -164,13 +170,9 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, // get the priority for the given retry number var targetChannels []*Channel - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - if channel.GetPriority() == targetPriority { - targetChannels = append(targetChannels, channel) - } - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) + for _, channel := range validChannels { + if channel.GetPriority() == targetPriority { + targetChannels = append(targetChannels, channel) } } @@ -195,6 +197,53 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, return nil, errors.New("channel not found") } +func isRedisLimited(channel Channel, channelId int) bool { + if channel.RateLimit != nil && *channel.RateLimit > 0 { + if !checkRedisLimit(channel, channelId) { + return true + } + } + return false +} + +func checkRedisLimit(channel Channel, channelId int) bool { + key := fmt.Sprintf("rate_limit:%d", channelId) + + countStr, err := common.RedisGet(key) + if err == redis.Nil { + // Key doesn't exist, set it with expiration + err = common.RedisSet(key, "1", time.Minute) + if err != nil { + common.SysLog(fmt.Sprintf("Error setting rate limit: %v", err)) + return false + } + return true + } else if err != nil { + common.SysLog(fmt.Sprintf("Error checking rate limit: %v", err)) + return false + } + + count, err := strconv.ParseInt(countStr, 10, 64) + if err != nil { + common.SysLog(fmt.Sprintf("Error parsing rate limit count: %v", err)) + return false + } + + if count > int64(*channel.RateLimit) { + return false + } + + // 增加计数 + newCount := strconv.FormatInt(count+1, 10) + err = common.RedisSet(key, newCount, time.Minute) + if err != nil { + common.SysLog(fmt.Sprintf("Error incrementing rate limit: %v", err)) + return false + } + + return true +} + func CacheGetChannel(id int) (*Channel, error) { if !common.MemoryCacheEnabled { return GetChannelById(id, true) diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go index 45fde019cd39..ffef13a34b6c 100644 --- a/relay/common/relay_info.go +++ b/relay/common/relay_info.go @@ -62,6 +62,7 @@ type ResponsesUsageInfo struct { type RelayInfo struct { ChannelType int ChannelId int + ChannelRateLimit int TokenId int TokenKey string UserId int @@ -215,6 +216,7 @@ func GenRelayInfoImage(c *gin.Context) *RelayInfo { func GenRelayInfo(c *gin.Context) *RelayInfo { channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType) channelId := common.GetContextKeyInt(c, constant.ContextKeyChannelId) + channelRateLimit := common.GetContextKeyInt(c, constant.ContextKeyChannelRateLimit) paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride) tokenId := common.GetContextKeyInt(c, constant.ContextKeyTokenId) @@ -235,6 +237,7 @@ func GenRelayInfo(c *gin.Context) *RelayInfo { RequestURLPath: c.Request.URL.String(), ChannelType: channelType, ChannelId: channelId, + ChannelRateLimit: channelRateLimit, TokenId: tokenId, TokenKey: tokenKey, UserId: userId, diff --git a/web/src/pages/Channel/EditChannel.js b/web/src/pages/Channel/EditChannel.js index 0934d8912029..695de60ad164 100644 --- a/web/src/pages/Channel/EditChannel.js +++ b/web/src/pages/Channel/EditChannel.js @@ -102,6 +102,7 @@ const EditChannel = (props) => { groups: ['default'], priority: 0, weight: 0, + rate_limit: 0, tag: '', multi_key_mode: 'random', }; @@ -1371,6 +1372,16 @@ const EditChannel = (props) => { style={{ width: '100%' }} /> + + handleInputChange('rate_limit', value)} + style={{ width: '100%' }} + /> +