feat: 渠道级别限制每分钟最大请求数

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"

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 {

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 {

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

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)

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,

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%' }}
                       />
                     </Col>
+                    <Col span={12}>
+                      <Form.InputNumber
+                        field='rate_limit'
+                        label={t('每分钟最大请求数')}
+                        placeholder={t('每分钟最大请求数')}
+                        min={0}
+                        onNumberChange={(value) => handleInputChange('rate_limit', value)}
+                        style={{ width: '100%' }}
+                      />
+                    </Col>
                   </Row>
 
                   <Form.Switch