fix(channel): unify weighted selection across cache modes

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
 }

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

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")
+}

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)
+		})
+	}
+}