597 lines
22 KiB
Go
597 lines
22 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"sort"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/dto"
|
|
"github.com/QuantumNous/new-api/i18n"
|
|
"github.com/QuantumNous/new-api/logger"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
|
"github.com/QuantumNous/new-api/relaykit/types"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func GetChannelConstraints(c *gin.Context) *dto.ChannelConstraints {
|
|
if c == nil {
|
|
return &dto.ChannelConstraints{}
|
|
}
|
|
if existing, ok := common.GetContextKeyType[*dto.ChannelConstraints](c, constant.ContextKeyChannelConstraints); ok && existing != nil {
|
|
return existing
|
|
}
|
|
constraints := &dto.ChannelConstraints{}
|
|
common.SetContextKey(c, constant.ContextKeyChannelConstraints, constraints)
|
|
return constraints
|
|
}
|
|
|
|
func AppendTaskPluginIdentityFilter(c *gin.Context, pluginKey string) {
|
|
if c == nil {
|
|
return
|
|
}
|
|
channelTypes, pluginKeys := pinnedTaskPluginIdentities(c, pluginKey)
|
|
GetChannelConstraints(c).AddFilter(dto.ChannelFilter{
|
|
Kind: dto.FilterTaskPluginIdentity,
|
|
TaskPluginKey: pluginKey,
|
|
TaskPluginChannelTypes: channelTypes,
|
|
TaskPluginKeys: pluginKeys,
|
|
})
|
|
}
|
|
|
|
type RetryParam struct {
|
|
Ctx *gin.Context
|
|
TokenGroup string
|
|
ModelName string
|
|
RequestPath string
|
|
Retry *int
|
|
resetNextTry bool
|
|
}
|
|
|
|
func (p *RetryParam) GetRetry() int {
|
|
if p.Retry == nil {
|
|
return 0
|
|
}
|
|
return *p.Retry
|
|
}
|
|
|
|
func (p *RetryParam) SetRetry(retry int) {
|
|
p.Retry = &retry
|
|
}
|
|
|
|
func (p *RetryParam) IncreaseRetry() {
|
|
if p.resetNextTry {
|
|
p.resetNextTry = false
|
|
return
|
|
}
|
|
if p.Retry == nil {
|
|
p.Retry = new(int)
|
|
}
|
|
*p.Retry++
|
|
}
|
|
|
|
func (p *RetryParam) ResetRetryNextTry() {
|
|
p.resetNextTry = true
|
|
}
|
|
|
|
// groupPoolTier is a set of groups sharing one priority. Tiers are tried in
|
|
// ascending priority order, and the groups inside a tier are load-balanced.
|
|
type groupPoolTier struct {
|
|
priority int
|
|
groups []string
|
|
}
|
|
|
|
var groupPoolRoundRobinCounters sync.Map
|
|
|
|
// buildGroupPoolTiers resolves the token's routing pool into ordered tiers.
|
|
// Entries the user group cannot select are dropped, and the token's own group is
|
|
// placed ahead of the configured pool when it is missing from it.
|
|
func buildGroupPoolTiers(c *gin.Context, primary, userGroup string) []groupPoolTier {
|
|
value, exists := common.GetContextKey(c, constant.ContextKeyTokenGroupPool)
|
|
if !exists {
|
|
return nil
|
|
}
|
|
entries, ok := value.([]model.TokenGroupEntry)
|
|
if !ok || len(entries) == 0 {
|
|
return nil
|
|
}
|
|
maxCount := setting.GetMaxTokenAutoGroups()
|
|
ordered := make([]model.TokenGroupEntry, 0, len(entries)+1)
|
|
seen := make(map[string]struct{}, len(entries)+1)
|
|
hasPrimary := false
|
|
for _, entry := range entries {
|
|
if len(ordered) >= maxCount {
|
|
break
|
|
}
|
|
if !IsUserSelectableGroup(userGroup, entry.Group) {
|
|
continue
|
|
}
|
|
if _, duplicate := seen[entry.Group]; duplicate {
|
|
continue
|
|
}
|
|
seen[entry.Group] = struct{}{}
|
|
if entry.Group == primary {
|
|
hasPrimary = true
|
|
}
|
|
ordered = append(ordered, entry)
|
|
}
|
|
if len(ordered) == 0 {
|
|
return nil
|
|
}
|
|
if !hasPrimary && primary != "" && primary != "auto" {
|
|
minPriority := ordered[0].Priority
|
|
for _, entry := range ordered {
|
|
if entry.Priority < minPriority {
|
|
minPriority = entry.Priority
|
|
}
|
|
}
|
|
ordered = append([]model.TokenGroupEntry{{Group: primary, Priority: minPriority - 1}}, ordered...)
|
|
}
|
|
sort.SliceStable(ordered, func(i, j int) bool { return ordered[i].Priority < ordered[j].Priority })
|
|
tiers := make([]groupPoolTier, 0, len(ordered))
|
|
for _, entry := range ordered {
|
|
if len(tiers) > 0 && tiers[len(tiers)-1].priority == entry.Priority {
|
|
tiers[len(tiers)-1].groups = append(tiers[len(tiers)-1].groups, entry.Group)
|
|
continue
|
|
}
|
|
tiers = append(tiers, groupPoolTier{priority: entry.Priority, groups: []string{entry.Group}})
|
|
}
|
|
return tiers
|
|
}
|
|
|
|
func failedRequestGroups(c *gin.Context) map[string]bool {
|
|
failed := make(map[string]bool)
|
|
groups, ok := common.GetContextKeyType[[]string](c, constant.ContextKeyTokenFailedGroups)
|
|
if !ok {
|
|
return failed
|
|
}
|
|
for _, group := range groups {
|
|
failed[group] = true
|
|
}
|
|
return failed
|
|
}
|
|
|
|
// MarkRequestGroupFailed records that a group already failed during the current
|
|
// request so the next selection switches away from it instead of retrying it.
|
|
func MarkRequestGroupFailed(c *gin.Context, group string) {
|
|
if c == nil || group == "" {
|
|
return
|
|
}
|
|
groups, _ := common.GetContextKeyType[[]string](c, constant.ContextKeyTokenFailedGroups)
|
|
for _, existing := range groups {
|
|
if existing == group {
|
|
return
|
|
}
|
|
}
|
|
common.SetContextKey(c, constant.ContextKeyTokenFailedGroups, append(groups, group))
|
|
}
|
|
|
|
func nextTierIndex(key string, size int) int {
|
|
if size <= 1 {
|
|
return 0
|
|
}
|
|
value, _ := groupPoolRoundRobinCounters.LoadOrStore(key, new(uint64))
|
|
counter, ok := value.(*uint64)
|
|
if !ok {
|
|
return 0
|
|
}
|
|
return int(atomic.AddUint64(counter, 1)-1) % size
|
|
}
|
|
|
|
func selectFromGroupPool(param *RetryParam, tiers []groupPoolTier, filters []dto.ChannelFilter) (*model.Channel, string) {
|
|
failed := failedRequestGroups(param.Ctx)
|
|
// Each retry advances one priority tier, so a retryable failure reaches the
|
|
// next available group instead of repeating the group that just failed.
|
|
start := param.GetRetry()
|
|
if start < 0 {
|
|
start = 0
|
|
}
|
|
if start > len(tiers)-1 {
|
|
start = len(tiers) - 1
|
|
}
|
|
for _, tier := range tiers[start:] {
|
|
eligible := make([]string, 0, len(tier.groups))
|
|
for _, group := range tier.groups {
|
|
if failed[group] || !IsGroupAvailable(group) {
|
|
continue
|
|
}
|
|
eligible = append(eligible, group)
|
|
}
|
|
if len(eligible) == 0 {
|
|
continue
|
|
}
|
|
start := nextTierIndex(tier.groups[0], len(eligible))
|
|
for offset := 0; offset < len(eligible); offset++ {
|
|
group := eligible[(start+offset)%len(eligible)]
|
|
if channel, _ := selectWithTrace(param, group, 0); channel != nil {
|
|
logger.LogInfo(param.Ctx, fmt.Sprintf("group pool selected: group=%s channel=%d", group, channel.Id))
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
|
|
return channel, group
|
|
}
|
|
}
|
|
}
|
|
// Availability is a heuristic, so when every tier looks degraded fall back to
|
|
// the highest-priority group this request has not tried yet rather than
|
|
// failing without an attempt.
|
|
for _, tier := range tiers {
|
|
for _, group := range tier.groups {
|
|
if failed[group] {
|
|
continue
|
|
}
|
|
if channel, _ := selectWithTrace(param, group, 0); channel != nil {
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
|
|
return channel, group
|
|
}
|
|
}
|
|
}
|
|
return nil, param.TokenGroup
|
|
}
|
|
|
|
// CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements.
|
|
// 尝试获取一个满足要求的随机渠道。
|
|
//
|
|
// For ordered Auto groups or a fixed token group with a backup group:
|
|
// 对于有序 Auto 分组，或配置了备用分组的固定令牌分组：
|
|
//
|
|
// - Each group will exhaust all its priorities before moving to the next group.
|
|
// 每个分组会用完所有优先级后才会切换到下一个分组。
|
|
//
|
|
// - Uses ContextKeyAutoGroupIndex to track current group index.
|
|
// 使用 ContextKeyAutoGroupIndex 跟踪当前分组索引。
|
|
//
|
|
// - Uses ContextKeyAutoGroupRetryIndex to track the global Retry count when current group started.
|
|
// 使用 ContextKeyAutoGroupRetryIndex 跟踪当前分组开始时的全局重试次数。
|
|
//
|
|
// - priorityRetry = Retry - startRetryIndex, represents the priority level within current group.
|
|
// priorityRetry = Retry - startRetryIndex，表示当前分组内的优先级级别。
|
|
//
|
|
// - When GetRandomSatisfiedChannel returns nil (priorities exhausted), moves to next group.
|
|
// 当 GetRandomSatisfiedChannel 返回 nil（优先级用完）时，切换到下一个分组。
|
|
//
|
|
// Example flow (2 groups, each with 2 priorities, RetryTimes=3):
|
|
// 示例流程（2个分组，每个有2个优先级，RetryTimes=3）：
|
|
//
|
|
// Retry=0: GroupA, priority0 (startRetryIndex=0, priorityRetry=0)
|
|
// 分组A, 优先级0
|
|
//
|
|
// Retry=1: GroupA, priority1 (startRetryIndex=0, priorityRetry=1)
|
|
// 分组A, 优先级1
|
|
//
|
|
// Retry=2: GroupA exhausted → GroupB, priority0 (startRetryIndex=2, priorityRetry=0)
|
|
// 分组A用完 → 分组B, 优先级0
|
|
//
|
|
// Retry=3: GroupB, priority1 (startRetryIndex=2, priorityRetry=1)
|
|
// 分组B, 优先级1
|
|
func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, error) {
|
|
if RequestRoutingPolicy(param.Ctx) != nil {
|
|
return selectTokenRoutingPool(param)
|
|
}
|
|
if group := imageGenerationRequestGroup(param.Ctx, param.ModelName); group != "" {
|
|
channel, err := selectWithTrace(param, group, param.GetRetry())
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, group)
|
|
return channel, group, err
|
|
}
|
|
var channel *model.Channel
|
|
var err error
|
|
selectGroup := param.TokenGroup
|
|
userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup)
|
|
filters := GetChannelConstraints(param.Ctx).Filters
|
|
crossGroupRetry := common.GetContextKeyBool(param.Ctx, constant.ContextKeyTokenCrossGroupRetry)
|
|
orderedGroups := make([]string, 0)
|
|
tiers := make([]groupPoolTier, 0)
|
|
|
|
if param.TokenGroup == "auto" {
|
|
orderedGroups = GetRequestAutoGroups(param.Ctx, userGroup)
|
|
if len(orderedGroups) == 0 {
|
|
return nil, selectGroup, errors.New("auto groups is not enabled")
|
|
}
|
|
} else if crossGroupRetry {
|
|
if _, exists := common.GetContextKey(param.Ctx, constant.ContextKeyTokenGroupPool); exists {
|
|
tiers = buildGroupPoolTiers(param.Ctx, param.TokenGroup, userGroup)
|
|
}
|
|
if len(tiers) == 0 {
|
|
if value, exists := common.GetContextKey(param.Ctx, constant.ContextKeyTokenAutoGroups); exists {
|
|
if backupGroups, ok := value.([]string); ok {
|
|
orderedGroups = append(orderedGroups, param.TokenGroup)
|
|
for _, group := range FilterUserTokenAutoGroups(userGroup, backupGroups) {
|
|
if group != param.TokenGroup {
|
|
orderedGroups = append(orderedGroups, group)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(tiers) > 0 {
|
|
logger.LogInfo(param.Ctx, fmt.Sprintf("group pool selection: tokenGroup=%s tiers=%d retry=%d", param.TokenGroup, len(tiers), param.GetRetry()))
|
|
channel, selectGroup = selectFromGroupPool(param, tiers, filters)
|
|
return channel, selectGroup, nil
|
|
}
|
|
|
|
if len(orderedGroups) > 0 {
|
|
logger.LogInfo(param.Ctx, fmt.Sprintf("legacy group selection: groups=%v retry=%d", orderedGroups, param.GetRetry()))
|
|
// startGroupIndex: the group index to start searching from
|
|
// startGroupIndex: 开始搜索的分组索引
|
|
startGroupIndex := 0
|
|
|
|
if lastGroupIndex, exists := common.GetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex); exists {
|
|
if idx, ok := lastGroupIndex.(int); ok {
|
|
startGroupIndex = idx
|
|
}
|
|
}
|
|
|
|
for i := startGroupIndex; i < len(orderedGroups); i++ {
|
|
autoGroup := orderedGroups[i]
|
|
// Calculate priorityRetry for current group
|
|
// 计算当前分组的 priorityRetry
|
|
priorityRetry := param.GetRetry()
|
|
// If moved to a new group, reset priorityRetry and update startRetryIndex
|
|
// 如果切换到新分组，重置 priorityRetry 并更新 startRetryIndex
|
|
if i > startGroupIndex {
|
|
priorityRetry = 0
|
|
}
|
|
logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry)
|
|
|
|
channel, _ = selectWithTrace(param, autoGroup, priorityRetry)
|
|
if channel == nil {
|
|
// Current group has no available channel for this model, try next group
|
|
// 当前分组没有该模型的可用渠道，尝试下一个分组
|
|
logger.LogDebug(param.Ctx, "No available channel in group %s for model %s at priorityRetry %d, trying next group", autoGroup, param.ModelName, priorityRetry)
|
|
// 重置状态以尝试下一个分组
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1)
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupRetryIndex, 0)
|
|
// Reset retry counter so outer loop can continue for next group
|
|
// 重置重试计数器，以便外层循环可以为下一个分组继续
|
|
param.SetRetry(0)
|
|
continue
|
|
}
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroup, autoGroup)
|
|
selectGroup = autoGroup
|
|
logger.LogDebug(param.Ctx, "Auto selected group: %s", autoGroup)
|
|
|
|
// Prepare state for next retry
|
|
// 为下一次重试准备状态
|
|
if crossGroupRetry && priorityRetry >= common.RetryTimes {
|
|
// Current group has exhausted all retries, prepare to switch to next group
|
|
// This request still uses current group, but next retry will use next group
|
|
// 当前分组已用完所有重试次数，准备切换到下一个分组
|
|
// 本次请求仍使用当前分组，但下次重试将使用下一个分组
|
|
logger.LogDebug(param.Ctx, "Current group %s retries exhausted (priorityRetry=%d >= RetryTimes=%d), preparing switch to next group for next retry", autoGroup, priorityRetry, common.RetryTimes)
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i+1)
|
|
// Reset retry counter so outer loop can continue for next group
|
|
// 重置重试计数器，以便外层循环可以为下一个分组继续
|
|
param.SetRetry(0)
|
|
param.ResetRetryNextTry()
|
|
} else {
|
|
// Stay in current group, save current state
|
|
// 保持在当前分组，保存当前状态
|
|
common.SetContextKey(param.Ctx, constant.ContextKeyAutoGroupIndex, i)
|
|
}
|
|
break
|
|
}
|
|
} else {
|
|
channel, err = selectWithTrace(param, param.TokenGroup, param.GetRetry())
|
|
if err != nil {
|
|
return nil, param.TokenGroup, err
|
|
}
|
|
}
|
|
return channel, selectGroup, nil
|
|
}
|
|
|
|
func pinnedTaskPluginIdentities(c *gin.Context, expected string) ([]int, []string) {
|
|
if c == nil || expected == "" {
|
|
return nil, nil
|
|
}
|
|
if value, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint); exists {
|
|
pinned, ok := value.(jsplugin.PinnedEndpoint)
|
|
if ok && pinned.Generation != nil && len(pinned.Candidates) > 1 {
|
|
expectedFound := false
|
|
channelTypes := make([]int, 0, len(pinned.Candidates))
|
|
pluginKeys := make([]string, 0, len(pinned.Candidates))
|
|
seen := make(map[int]struct{}, len(pinned.Candidates))
|
|
for _, candidate := range pinned.Candidates {
|
|
if candidate.Plugin == nil {
|
|
continue
|
|
}
|
|
if candidate.Plugin.Meta.Key == expected {
|
|
expectedFound = true
|
|
}
|
|
pluginKeys = append(pluginKeys, candidate.Plugin.Meta.Key)
|
|
for _, channelType := range candidate.Plugin.Meta.ChannelTypes {
|
|
if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
|
continue
|
|
}
|
|
if _, duplicate := seen[channelType]; duplicate {
|
|
continue
|
|
}
|
|
if plugin, indexed := pinned.Generation.GetByChannelType(channelType); indexed && plugin == candidate.Plugin {
|
|
seen[channelType] = struct{}{}
|
|
channelTypes = append(channelTypes, channelType)
|
|
}
|
|
}
|
|
}
|
|
if expectedFound {
|
|
return channelTypes, pluginKeys
|
|
}
|
|
}
|
|
}
|
|
value, exists := c.Get(jsplugin.ContextKeyPinnedPlugin)
|
|
pinned, ok := value.(jsplugin.PinnedPlugin)
|
|
if !exists || !ok || pinned.Generation == nil || pinned.Plugin == nil || pinned.Plugin.Meta.Key != expected {
|
|
return nil, nil
|
|
}
|
|
channelTypes := make([]int, 0, len(pinned.Plugin.Meta.ChannelTypes))
|
|
for _, channelType := range pinned.Plugin.Meta.ChannelTypes {
|
|
if channelType == 0 || channelType == constant.ChannelTypeTaskPlugin {
|
|
continue
|
|
}
|
|
channelTypes = append(channelTypes, channelType)
|
|
}
|
|
return channelTypes, []string{expected}
|
|
}
|
|
|
|
// ChannelSelectError explains why SelectChannelForRequest found no channel.
|
|
// Callers render it for their transport: the HTTP distributor localizes
|
|
// MessageID with its own helpers and the Responses WebSocket relay wraps it in
|
|
// a NewAPIError. Message is set instead of MessageID when the text is a fixed
|
|
// error code that clients match on.
|
|
type ChannelSelectError struct {
|
|
StatusCode int
|
|
Code types.ErrorCode
|
|
MessageID string
|
|
Params map[string]any
|
|
Message string
|
|
// FilterKind and Channel identify a candidate rejected by request filters.
|
|
FilterKind dto.ChannelFilterKind
|
|
Channel *model.Channel
|
|
// NoAvailableChannel marks the "no channel for this group and model"
|
|
// outcome so the distributor can name the claiming task plugin.
|
|
NoAvailableChannel bool
|
|
}
|
|
|
|
// SelectChannelForRequest resolves the channel for one attempt with the rules
|
|
// shared by the HTTP distributor and the Responses WebSocket relay: a pinned
|
|
// channel wins, then session affinity (first attempt only), then a random
|
|
// eligible channel; every candidate must satisfy the request's channel
|
|
// filters. The group the channel was chosen from is returned for auto-group
|
|
// callers. The caller still applies SetupContextForSelectedChannel.
|
|
func SelectChannelForRequest(c *gin.Context, modelName string, retry *RetryParam) (*model.Channel, string, *ChannelSelectError) {
|
|
c.Set("original_model", modelName)
|
|
constraints := GetChannelConstraints(c)
|
|
if pin, found, overridden := constraints.ResolvedPin(); found {
|
|
for _, lost := range overridden {
|
|
logger.LogWarn(c, fmt.Sprintf(
|
|
"channel pin overridden: winning_source=%s winning_channel_id=%d overridden_source=%s overridden_channel_id=%d",
|
|
pin.Source, pin.ChannelId, lost.Source, lost.ChannelId,
|
|
))
|
|
}
|
|
channel, err := model.CacheGetChannel(pin.ChannelId)
|
|
if err != nil {
|
|
recordDirectSelection(c, nil, pin.ChannelId, retry.TokenGroup, modelName, "pinned_channel_missing")
|
|
return nil, "", pinnedChannelUnavailable(pin, http.StatusBadRequest, i18n.MsgDistributorInvalidChannelId)
|
|
}
|
|
if channel.Status != common.ChannelStatusEnabled {
|
|
recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "channel_disabled")
|
|
return nil, "", pinnedChannelUnavailable(pin, http.StatusForbidden, i18n.MsgDistributorChannelDisabled)
|
|
}
|
|
if ok, kind := model.ChannelSatisfiesFilters(channel, modelName, constraints.Filters); !ok {
|
|
recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "excluded_"+string(kind))
|
|
return nil, "", &ChannelSelectError{
|
|
StatusCode: http.StatusBadRequest, Code: types.ErrorCode(kind), MessageID: i18n.MsgDistributorNoAvailableChannel,
|
|
Params: map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelName},
|
|
FilterKind: kind, Channel: channel,
|
|
}
|
|
}
|
|
recordDirectSelection(c, channel, channel.Id, retry.TokenGroup, modelName, "pinned_channel")
|
|
return channel, "", nil
|
|
}
|
|
|
|
usingGroup := retry.TokenGroup
|
|
if RequestRoutingPolicy(c) != nil {
|
|
if retry.GetRetry() == 0 {
|
|
if channel, group := PreferredPoolChannel(c, modelName); channel != nil {
|
|
return channel, group, nil
|
|
}
|
|
}
|
|
channel, group, err := selectTokenRoutingPool(retry)
|
|
if err != nil || channel == nil {
|
|
return nil, group, &ChannelSelectError{StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, Message: "No eligible channel in the group pool", NoAvailableChannel: true}
|
|
}
|
|
return channel, group, nil
|
|
}
|
|
var channel *model.Channel
|
|
var selectGroup string
|
|
if retry.GetRetry() == 0 {
|
|
if preferredChannelID, found := GetPreferredChannelByAffinity(c, modelName, usingGroup); found {
|
|
affinityUsable := false
|
|
preferred, err := model.CacheGetChannel(preferredChannelID)
|
|
affinitySatisfied := false
|
|
if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
|
|
affinitySatisfied, _ = model.ChannelSatisfiesFilters(preferred, modelName, constraints.Filters)
|
|
}
|
|
if affinitySatisfied {
|
|
if usingGroup == "auto" {
|
|
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
|
|
for _, g := range GetRequestAutoGroups(c, userGroup) {
|
|
if model.IsChannelEnabledForGroupModel(g, modelName, preferred.Id) {
|
|
selectGroup = g
|
|
common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
|
|
channel = preferred
|
|
affinityUsable = true
|
|
MarkChannelAffinityUsed(c, g, preferred.Id)
|
|
break
|
|
}
|
|
}
|
|
} else if model.IsChannelEnabledForGroupModel(usingGroup, modelName, preferred.Id) {
|
|
channel = preferred
|
|
selectGroup = usingGroup
|
|
affinityUsable = true
|
|
MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
|
|
}
|
|
}
|
|
reason := "affinity_unavailable"
|
|
if affinityUsable {
|
|
reason = "affinity_hit"
|
|
}
|
|
recordDirectSelection(c, preferred, preferredChannelID, usingGroup, modelName, reason)
|
|
if !affinityUsable && !ShouldKeepChannelAffinityOnChannelDisabled() {
|
|
ClearCurrentChannelAffinityCache(c)
|
|
}
|
|
if !affinityUsable && RequestPolicy(c).SessionMode == "strict" {
|
|
return nil, "", &ChannelSelectError{StatusCode: http.StatusServiceUnavailable, Message: "strict_session_binding_unavailable"}
|
|
}
|
|
}
|
|
}
|
|
|
|
if channel == nil {
|
|
var err error
|
|
channel, selectGroup, err = CacheGetRandomSatisfiedChannel(retry)
|
|
if err != nil {
|
|
showGroup := usingGroup
|
|
if usingGroup == "auto" {
|
|
showGroup = fmt.Sprintf("auto(%s)", selectGroup)
|
|
}
|
|
return nil, selectGroup, &ChannelSelectError{
|
|
StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorGetChannelFailed,
|
|
Params: map[string]any{"Group": showGroup, "Model": modelName, "Error": err.Error()},
|
|
}
|
|
}
|
|
if channel == nil {
|
|
return nil, selectGroup, &ChannelSelectError{
|
|
StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorNoAvailableChannel,
|
|
Params: map[string]any{"Group": usingGroup, "Model": modelName}, NoAvailableChannel: true,
|
|
}
|
|
}
|
|
}
|
|
if ok, kind := model.ChannelSatisfiesFilters(channel, modelName, constraints.Filters); !ok {
|
|
return nil, selectGroup, &ChannelSelectError{
|
|
StatusCode: http.StatusServiceUnavailable, Code: types.ErrorCodeModelNotFound, MessageID: i18n.MsgDistributorNoAvailableChannel,
|
|
Params: map[string]any{"Group": common.GetContextKeyString(c, constant.ContextKeyUsingGroup), "Model": modelName},
|
|
FilterKind: kind, Channel: channel, NoAvailableChannel: true,
|
|
}
|
|
}
|
|
return channel, selectGroup, nil
|
|
}
|
|
|
|
// Origin-task pins report a fixed code so task polling can tell a retired
|
|
// channel from a malformed request.
|
|
func pinnedChannelUnavailable(pin dto.ChannelPin, statusCode int, messageID string) *ChannelSelectError {
|
|
if pin.Source == dto.PinSourceOriginTask {
|
|
return &ChannelSelectError{StatusCode: http.StatusBadRequest, Code: "origin_task_channel_disabled", Message: "origin_task_channel_disabled"}
|
|
}
|
|
return &ChannelSelectError{StatusCode: statusCode, MessageID: messageID}
|
|
}
|
|
|
|
// AppendUsedChannel records an attempted channel in the request's channel
|
|
// trail, which the retry log and the consume log's admin_info both read.
|
|
func AppendUsedChannel(c *gin.Context, channelID int) {
|
|
c.Set("use_channel", append(c.GetStringSlice("use_channel"), fmt.Sprintf("%d", channelID)))
|
|
}
|