diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..974887e791d1 100644 --- a/model/ability.go +++ b/model/ability.go @@ -122,26 +122,20 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha return nil, err } abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model) - channel := Channel{} - if len(abilities) > 0 { - // Randomly choose one - weightSum := uint(0) - for _, ability_ := range abilities { - weightSum += ability_.Weight + 10 - } - // 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 - } - } - } else { + if len(abilities) == 0 { return nil, nil } + + weights := make([]uint, len(abilities)) + for index, ability := range abilities { + weights[index] = ability.Weight + } + selectedIndex, err := selectWeightedIndex(weights) + if err != nil { + return nil, err + } + + channel := Channel{Id: abilities[selectedIndex].ChannelId} err = DB.First(&channel, "id = ?", channel.Id).Error return &channel, err } diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..b0870b65642c 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -3,7 +3,6 @@ package model import ( "errors" "fmt" - "math/rand" "sort" "strings" "sync" @@ -160,13 +159,17 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat targetPriority := int64(sortedUniquePriorities[retry]) // get the priority for the given retry number - var sumWeight = 0 - var targetChannels []*Channel + targetChannels := make([]*Channel, 0, len(channels)) + targetWeights := make([]uint, 0, len(channels)) for _, channelId := range channels { if channel, ok := channelsIDM[channelId]; ok { if channel.GetPriority() == targetPriority { - sumWeight += channel.GetWeight() + weight := uint(0) + if channel.Weight != nil { + weight = *channel.Weight + } targetChannels = append(targetChannels, channel) + targetWeights = append(targetWeights, weight) } } else { return nil, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", channelId) @@ -177,35 +180,11 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority)) } - // smoothing factor and adjustment - smoothingFactor := 1 - smoothingAdjustment := 0 - - if sumWeight == 0 { - // when all channels have weight 0, set sumWeight to the number of channels and set smoothing adjustment to 100 - // each channel's effective weight = 100 - sumWeight = len(targetChannels) * 100 - smoothingAdjustment = 100 - } else if sumWeight/len(targetChannels) < 10 { - // when the average weight is less than 10, set smoothing factor to 100 - smoothingFactor = 100 - } - - // Calculate the total weight of all channels up to endIdx - totalWeight := sumWeight * smoothingFactor - - // Generate a random value in the range [0, totalWeight) - randomWeight := rand.Intn(totalWeight) - - // Find a channel based on its weight - for _, channel := range targetChannels { - randomWeight -= channel.GetWeight()*smoothingFactor + smoothingAdjustment - if randomWeight < 0 { - return channel, nil - } + selectedIndex, err := selectWeightedIndex(targetWeights) + if err != nil { + return nil, err } - // return null if no channel is not found - return nil, errors.New("channel not found") + return targetChannels[selectedIndex], nil } // filterChannelsByRequestPathAndModel restricts candidates by request path and diff --git a/model/channel_weight.go b/model/channel_weight.go new file mode 100644 index 000000000000..b8e9e7042343 --- /dev/null +++ b/model/channel_weight.go @@ -0,0 +1,49 @@ +package model + +import ( + "errors" + "fmt" + "math" + "math/rand/v2" +) + +func selectWeightedIndex(weights []uint) (int, error) { + return selectWeightedIndexWithDraw(weights, rand.Uint64N) +} + +func selectWeightedIndexWithDraw(weights []uint, draw func(uint64) uint64) (int, error) { + if len(weights) == 0 { + return -1, errors.New("cannot select a channel from an empty weight list") + } + + var totalWeight uint64 + for _, weight := range weights { + if uint64(weight) > math.MaxUint64-totalWeight { + return -1, errors.New("channel weight sum overflows uint64") + } + totalWeight += uint64(weight) + } + + useEqualWeights := totalWeight == 0 + if useEqualWeights { + totalWeight = uint64(len(weights)) + } + + randomWeight := draw(totalWeight) + if randomWeight >= totalWeight { + return -1, fmt.Errorf("random channel weight %d is outside [0, %d)", randomWeight, totalWeight) + } + + for index, weight := range weights { + effectiveWeight := uint64(weight) + if useEqualWeights { + effectiveWeight = 1 + } + if randomWeight < effectiveWeight { + return index, nil + } + randomWeight -= effectiveWeight + } + + return -1, errors.New("channel not found in weighted selection") +} diff --git a/model/channel_weight_test.go b/model/channel_weight_test.go new file mode 100644 index 000000000000..61ae79b4f62a --- /dev/null +++ b/model/channel_weight_test.go @@ -0,0 +1,140 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestSelectWeightedIndexUsesRawWeightIntervals(t *testing.T) { + type selection struct { + randomWeight uint64 + expectedIndex int + } + testCases := []struct { + name string + weights []uint + total uint64 + selections []selection + }{ + { + name: "two to one remains two to one", + weights: []uint{2, 1}, + total: 3, + selections: []selection{ + {randomWeight: 0, expectedIndex: 0}, + {randomWeight: 1, expectedIndex: 0}, + {randomWeight: 2, expectedIndex: 1}, + }, + }, + { + name: "all zero weights are equal", + weights: []uint{0, 0, 0}, + total: 3, + selections: []selection{ + {randomWeight: 0, expectedIndex: 0}, + {randomWeight: 1, expectedIndex: 1}, + {randomWeight: 2, expectedIndex: 2}, + }, + }, + { + name: "zero weights get no interval when another weight is positive", + weights: []uint{0, 1, 0}, + total: 1, + selections: []selection{ + {randomWeight: 0, expectedIndex: 1}, + }, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + for _, selection := range testCase.selections { + index, err := selectWeightedIndexWithDraw(testCase.weights, func(total uint64) uint64 { + assert.Equal(t, testCase.total, total) + return selection.randomWeight + }) + require.NoError(t, err) + assert.Equal(t, selection.expectedIndex, index, "random weight %d", selection.randomWeight) + } + }) + } +} + +func TestSelectWeightedIndexRejectsEmptyWeights(t *testing.T) { + _, err := selectWeightedIndexWithDraw(nil, func(uint64) uint64 { + t.Fatal("draw must not be called for an empty weight list") + return 0 + }) + + require.ErrorContains(t, err, "empty weight list") +} + +func TestGetRandomSatisfiedChannelExcludesZeroWeightAcrossCacheModes(t *testing.T) { + const ( + groupName = "weight-selection-test" + modelName = "weight-selection-model" + ) + originalMemoryCacheEnabled := common.MemoryCacheEnabled + positiveWeight := uint(1) + zeroWeight := uint(0) + priority := int64(0) + channels := []*Channel{ + { + Name: "positive-weight", + Key: "test-key-positive", + Status: common.ChannelStatusEnabled, + Weight: &positiveWeight, + Priority: &priority, + Group: groupName, + Models: modelName, + }, + { + Name: "zero-weight", + Key: "test-key-zero", + Status: common.ChannelStatusEnabled, + Weight: &zeroWeight, + Priority: &priority, + Group: groupName, + Models: modelName, + }, + } + t.Cleanup(func() { + defer func() { + common.MemoryCacheEnabled = originalMemoryCacheEnabled + }() + channelIDs := []int{channels[0].Id, channels[1].Id} + require.NoError(t, DB.Where("channel_id IN ?", channelIDs).Delete(&Ability{}).Error) + require.NoError(t, DB.Where("id IN ?", channelIDs).Delete(&Channel{}).Error) + common.MemoryCacheEnabled = true + InitChannelCache() + }) + + for _, channel := range channels { + require.NoError(t, DB.Create(channel).Error) + require.NoError(t, channel.AddAbilities(nil)) + } + + testCases := []struct { + name string + memoryCacheEnabled bool + }{ + {name: "database", memoryCacheEnabled: false}, + {name: "memory cache", memoryCacheEnabled: true}, + } + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + common.MemoryCacheEnabled = testCase.memoryCacheEnabled + if testCase.memoryCacheEnabled { + InitChannelCache() + } + + selected, err := GetRandomSatisfiedChannel(groupName, modelName, 0, "") + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, channels[0].Id, selected.Id) + }) + } +}