diff --git a/internal/cmd/cmd.go b/internal/cmd/cmd.go index 4f69805..ecfba9e 100644 --- a/internal/cmd/cmd.go +++ b/internal/cmd/cmd.go @@ -35,7 +35,6 @@ import ( "github.com/qianfree/team-api/internal/handler/relay" setupHandler "github.com/qianfree/team-api/internal/handler/setup" "github.com/qianfree/team-api/internal/plugin" - "github.com/qianfree/team-api/relay/scheduler" "github.com/qianfree/team-api/web" ) @@ -178,10 +177,6 @@ var ( billing.CleanSettledPreDeductTracks(ctx) return nil }) - cs.Register("affinity_cache_cleanup", "*/5 * * * *", func(ctx context.Context) error { - scheduler.GetGlobalAffinity().CleanExpired() - return nil - }) cs.Register("update_check", "0 */6 * * *", func(ctx context.Context) error { if common.Config().GetBool(ctx, "update_auto_check_enabled") { return update.BackgroundCheck(ctx) diff --git a/internal/handler/relay/task.go b/internal/handler/relay/task.go index dacd2c3..82077f1 100644 --- a/internal/handler/relay/task.go +++ b/internal/handler/relay/task.go @@ -90,6 +90,7 @@ func HandleTaskSubmit(r *ghttp.Request) { r.Context(), body, r.URL.Path, r.Header, rc, taskDataProvider, taskBillingProvider, channelMeta, ) + commitTaskAffinity(r.Context(), rc, modelName, channelMeta, capture.StatusCode()) // 提交成功后切换为任务 ID 跟踪（任务生命周期远超 HTTP 请求） if rc.TaskID != "" { @@ -176,8 +177,9 @@ func selectTaskChannel(r *ghttp.Request, body []byte) (*relay_common.ChannelMeta ctx := r.Context() tenantID := middleware.GetTenantID(ctx) userID := middleware.GetUserID(ctx) + apiKeyID := middleware.GetApiKeyID(ctx) - selection, err := relayDataProvider.GetChannelForModel(ctx, tenantID, userID, req.Model, nil) + selection, err := relayDataProvider.GetChannelForModel(ctx, tenantID, userID, apiKeyID, req.Model, nil) if err != nil { return nil, err } @@ -245,6 +247,7 @@ func HandleMjSubmit(r *ghttp.Request) { r.Context(), body, r.URL.Path, r.Header, rc, taskDataProvider, taskBillingProvider, channelMeta, ) + commitTaskAffinity(r.Context(), rc, modelName, channelMeta, capture.StatusCode()) if rc.TaskID != "" { monitor.SwitchToTaskID(rc.RequestID, rc.TaskID) @@ -327,6 +330,7 @@ func HandleAliImageSubmit(r *ghttp.Request) { writeSyncImageErrorWithCode(rc.Writer, 400, "image_async_disabled", "async image generation is disabled for this model; call POST /v1/images/generations instead") } + commitTaskAffinity(r.Context(), rc, modelName, channelMeta, capture.StatusCode()) if rc.TaskID != "" { monitor.SwitchToTaskID(rc.RequestID, rc.TaskID) @@ -338,6 +342,12 @@ func HandleAliImageSubmit(r *ghttp.Request) { go recordTaskSubmitAudit(r, rc, capture, body) } +func commitTaskAffinity(ctx context.Context, rc *relay_handler.TaskRelayContext, modelName string, channelMeta *relay_common.ChannelMeta, statusCode int) { + if statusCode >= 200 && statusCode < 300 { + relayDataProvider.SetChannelAffinity(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName, channelMeta.ChannelID) + } +} + // registerAsyncTask 注册异步任务到实时监控 func registerAsyncTask(requestID string, tenantID, userID, projectID int64, modelName string, channelMeta *relay_common.ChannelMeta, path string) { monitor.RegisterRequest(&monitor.TrackedRequest{ diff --git a/internal/logic/admin/channel.go b/internal/logic/admin/channel.go index 193c5bb..db53a2d 100644 --- a/internal/logic/admin/channel.go +++ b/internal/logic/admin/channel.go @@ -358,6 +358,7 @@ func (s *sAdmin) UpdateChannel(ctx context.Context, req *v1.ChannelUpdateReq) (* return nil, err } } + relay.InvalidateChannelAffinities(ctx, req.ID) return nil, nil } @@ -382,6 +383,7 @@ func (s *sAdmin) DeleteChannel(ctx context.Context, req *v1.ChannelDeleteReq) (* if err != nil { return nil, err } + relay.InvalidateChannelAffinities(ctx, req.ID) return nil, nil } @@ -492,6 +494,7 @@ func (s *sAdmin) DeleteChannelKey(ctx context.Context, req *v1.ChannelKeyDeleteR if err != nil { return nil, err } + relay.InvalidateChannelAffinities(ctx, req.ChannelID) return nil, nil } @@ -516,6 +519,7 @@ func (s *sAdmin) SetChannelAbilities(ctx context.Context, req *v1.ChannelAbility if err != nil { return nil, err } + relay.InvalidateChannelAffinities(ctx, req.ChannelID) return nil, nil } diff --git a/internal/logic/common/settings_registry.go b/internal/logic/common/settings_registry.go index 5709c39..0c44c95 100644 --- a/internal/logic/common/settings_registry.go +++ b/internal/logic/common/settings_registry.go @@ -215,6 +215,14 @@ var Registry = []SettingDef{ Label: "渠道自动禁用", Description: "连续失败达到阈值时自动禁用渠道"}, {Key: "channel_auto_disable_threshold", Type: SettingTypeInt, Default: "5", Category: "channel", Label: "自动禁用失败阈值", Validation: "min:2,max:50"}, + {Key: "channel_scheduler_v2_enabled", Type: SettingTypeBool, Default: "true", Category: "channel", + Label: "新版渠道调度", Description: "启用优先级、健康度、权重和渠道亲和统一调度；关闭后回退旧版顺序选择"}, + {Key: "channel_affinity_enabled", Type: SettingTypeBool, Default: "true", Category: "channel", + Label: "渠道亲和", Description: "同一调用主体和模型优先复用上一次成功渠道，提高上游缓存命中率"}, + {Key: "channel_affinity_ttl_seconds", Type: SettingTypeInt, Default: "1800", Category: "channel", + Label: "渠道亲和有效期", Validation: "min:60,max:86400", Description: "成功请求刷新亲和有效期，单位秒"}, + {Key: "channel_capacity_enabled", Type: SettingTypeBool, Default: "true", Category: "channel", + Label: "渠道并发容量", Description: "使用 Redis 分布式租约限制渠道最大并发，并在饱和时自动溢出到其他渠道"}, {Key: "health_snapshot_retention_days", Type: SettingTypeInt, Default: "7", Category: "channel", Label: "健康快照保留天数", Validation: "min:1,max:90"}, {Key: "channel_proxy_url", Type: SettingTypeString, Default: "", Category: "channel", diff --git a/internal/logic/relay/affinity.go b/internal/logic/relay/affinity.go index f2af4b2..9c4e069 100644 --- a/internal/logic/relay/affinity.go +++ b/internal/logic/relay/affinity.go @@ -2,54 +2,89 @@ package relay import ( "context" - "github.com/qianfree/team-api/internal/dao" + "crypto/sha256" + "encoding/hex" + "fmt" + "strconv" "time" "github.com/gogf/gf/v2/frame/g" + lcommon "github.com/qianfree/team-api/internal/logic/common" ) -// UpsertAffinity 更新亲和性记录（成功后调用） -func UpsertAffinity(ctx context.Context, tenantID, userID int64, modelName string, channelID int64) { - now := time.Now() - expiresAt := now.Add(1800 * time.Second) - g.DB().Exec(ctx, - `INSERT INTO chn_channel_affinities (tenant_id, user_id, model_name, channel_id, hit_count, expires_at, updated_at) - VALUES (?, ?, ?, ?, 1, ?, ?) - ON CONFLICT (tenant_id, user_id, model_name) - DO UPDATE SET channel_id = ?, hit_count = chn_channel_affinities.hit_count + 1, expires_at = ?, updated_at = ?`, - tenantID, userID, modelName, channelID, expiresAt, now, channelID, expiresAt, now) +const affinityRedisPrefix = "relay:affinity:v2" + +func affinityIdentity(tenantID, userID, apiKeyID int64, modelName string) (redisKey, seed string) { + seed = fmt.Sprintf("%d:%d:%d:%s", tenantID, userID, apiKeyID, modelName) + digest := sha256.Sum256([]byte(seed)) + return affinityRedisPrefix + ":key:" + hex.EncodeToString(digest[:]), seed } -// GetAffinity 获取亲和性渠道 -func GetAffinity(ctx context.Context, tenantID, userID int64, modelName string) (int64, bool) { - var result *struct { - ChannelID int64 `json:"channel_id"` +func affinityTTL(ctx context.Context) time.Duration { + seconds := lcommon.Config().GetInt(ctx, "channel_affinity_ttl_seconds") + if seconds <= 0 { + seconds = 1800 } - err := dao.ChnChannelAffinities.Ctx(ctx). - Where("tenant_id", tenantID). - Where("user_id", userID). - Where("model_name", modelName). - Where("expires_at > ?", time.Now()). - Fields("channel_id"). - Scan(&result) - if err != nil || result == nil { + return time.Duration(seconds) * time.Second +} + +func affinityEnabled(ctx context.Context) bool { + return lcommon.Config().GetBool(ctx, "channel_affinity_enabled") +} + +func getAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string) (int64, bool) { + if !affinityEnabled(ctx) { return 0, false } - return result.ChannelID, true + key, _ := affinityIdentity(tenantID, userID, apiKeyID, modelName) + value, err := g.Redis().Do(ctx, "GET", key) + if err != nil || value.IsNil() { + return 0, false + } + channelID, err := strconv.ParseInt(value.String(), 10, 64) + return channelID, err == nil && channelID > 0 } -// DeleteAffinity 删除亲和性记录 -func DeleteAffinity(ctx context.Context, tenantID, userID int64, modelName string) { - dao.ChnChannelAffinities.Ctx(ctx). - Where("tenant_id", tenantID). - Where("user_id", userID). - Where("model_name", modelName). - Delete() +func setAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, channelID int64) { + if !affinityEnabled(ctx) || channelID <= 0 { + return + } + key, _ := affinityIdentity(tenantID, userID, apiKeyID, modelName) + ttlSeconds := int64(affinityTTL(ctx).Seconds()) + reversePrefix := affinityRedisPrefix + ":channel:" + const script = ` +local old = redis.call('GET', KEYS[1]) +if old and old ~= ARGV[1] then + redis.call('SREM', ARGV[3] .. old, KEYS[1]) +end +redis.call('SETEX', KEYS[1], ARGV[2], ARGV[1]) +redis.call('SADD', ARGV[3] .. ARGV[1], KEYS[1]) +redis.call('EXPIRE', ARGV[3] .. ARGV[1], ARGV[2] + 60) +return 1` + if _, err := g.Redis().Do(ctx, "EVAL", script, 1, key, channelID, ttlSeconds, reversePrefix); err != nil { + g.Log().Warningf(ctx, "[Affinity] set failed: %v", err) + } } -// DeleteAffinityByChannel 删除某渠道的所有亲和性记录 -func DeleteAffinityByChannel(ctx context.Context, channelID int64) { - dao.ChnChannelAffinities.Ctx(ctx). - Where("channel_id", channelID). - Delete() +func deleteAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string) { + key, _ := affinityIdentity(tenantID, userID, apiKeyID, modelName) + reversePrefix := affinityRedisPrefix + ":channel:" + const script = ` +local old = redis.call('GET', KEYS[1]) +if old then redis.call('SREM', ARGV[1] .. old, KEYS[1]) end +return redis.call('DEL', KEYS[1])` + _, _ = g.Redis().Do(ctx, "EVAL", script, 1, key, reversePrefix) +} + +// InvalidateChannelAffinities removes all distributed affinity records bound to a channel. +func InvalidateChannelAffinities(ctx context.Context, channelID int64) { + reverseKey := affinityRedisPrefix + ":channel:" + strconv.FormatInt(channelID, 10) + const script = ` +local keys = redis.call('SMEMBERS', KEYS[1]) +for _, key in ipairs(keys) do redis.call('DEL', key) end +redis.call('DEL', KEYS[1]) +return #keys` + if _, err := g.Redis().Do(ctx, "EVAL", script, 1, reverseKey); err != nil { + g.Log().Warningf(ctx, "[Affinity] invalidate channel %d failed: %v", channelID, err) + } } diff --git a/internal/logic/relay/affinity_test.go b/internal/logic/relay/affinity_test.go new file mode 100644 index 0000000..ca0e28e --- /dev/null +++ b/internal/logic/relay/affinity_test.go @@ -0,0 +1,32 @@ +package relay + +import "testing" + +func TestAffinityIdentity(t *testing.T) { + key, seed := affinityIdentity(1, 2, 3, "gpt-4o") + sameKey, sameSeed := affinityIdentity(1, 2, 3, "gpt-4o") + if key != sameKey || seed != sameSeed { + t.Fatal("affinity identity is not deterministic") + } + + tests := []struct { + name string + tenantID int64 + userID int64 + apiKeyID int64 + modelName string + }{ + {name: "tenant", tenantID: 9, userID: 2, apiKeyID: 3, modelName: "gpt-4o"}, + {name: "user", tenantID: 1, userID: 9, apiKeyID: 3, modelName: "gpt-4o"}, + {name: "api key", tenantID: 1, userID: 2, apiKeyID: 9, modelName: "gpt-4o"}, + {name: "model", tenantID: 1, userID: 2, apiKeyID: 3, modelName: "claude-3"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + otherKey, _ := affinityIdentity(tt.tenantID, tt.userID, tt.apiKeyID, tt.modelName) + if otherKey == key { + t.Fatal("distinct affinity identity produced the same key") + } + }) + } +} diff --git a/internal/logic/relay/capacity.go b/internal/logic/relay/capacity.go new file mode 100644 index 0000000..f886d95 --- /dev/null +++ b/internal/logic/relay/capacity.go @@ -0,0 +1,52 @@ +package relay + +import ( + "context" + "strconv" + "time" + + "github.com/gogf/gf/v2/frame/g" + lcommon "github.com/qianfree/team-api/internal/logic/common" +) + +const channelLeaseDuration = 90 * time.Second + +func channelLeaseKey(channelID int64) string { + return "relay:channel:leases:v1:" + strconv.FormatInt(channelID, 10) +} + +func (p *DataProviderImpl) AcquireChannelSlot(ctx context.Context, channelID int64, maxConcurrency int, requestID string) bool { + if maxConcurrency <= 0 || !lcommon.Config().GetBool(ctx, "channel_capacity_enabled") { + return true + } + now := time.Now().UnixMilli() + expiresAt := now + channelLeaseDuration.Milliseconds() + const script = ` +redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) +if redis.call('ZCARD', KEYS[1]) >= tonumber(ARGV[3]) then return 0 end +redis.call('ZADD', KEYS[1], ARGV[2], ARGV[4]) +redis.call('PEXPIRE', KEYS[1], ARGV[5]) +return 1` + result, err := g.Redis().Do(ctx, "EVAL", script, 1, channelLeaseKey(channelID), now, expiresAt, maxConcurrency, requestID, channelLeaseDuration.Milliseconds()+60000) + if err != nil { + g.Log().Warningf(ctx, "[ChannelCapacity] acquire failed open: channel=%d err=%v", channelID, err) + return true + } + return result.Int() == 1 +} + +func (p *DataProviderImpl) RefreshChannelSlot(ctx context.Context, channelID int64, requestID string) { + expiresAt := time.Now().Add(channelLeaseDuration).UnixMilli() + const script = ` +if redis.call('ZSCORE', KEYS[1], ARGV[1]) then + redis.call('ZADD', KEYS[1], ARGV[2], ARGV[1]) + redis.call('PEXPIRE', KEYS[1], ARGV[3]) + return 1 +end +return 0` + _, _ = g.Redis().Do(ctx, "EVAL", script, 1, channelLeaseKey(channelID), requestID, expiresAt, channelLeaseDuration.Milliseconds()+60000) +} + +func (p *DataProviderImpl) ReleaseChannelSlot(ctx context.Context, channelID int64, requestID string) { + _, _ = g.Redis().Do(ctx, "ZREM", channelLeaseKey(channelID), requestID) +} diff --git a/internal/logic/relay/health.go b/internal/logic/relay/health.go index 083fb1c..135dcc4 100644 --- a/internal/logic/relay/health.go +++ b/internal/logic/relay/health.go @@ -10,7 +10,6 @@ import ( "github.com/qianfree/team-api/internal/dao" "github.com/qianfree/team-api/internal/logic/common" do "github.com/qianfree/team-api/internal/model/do" - "github.com/qianfree/team-api/relay/scheduler" ) // UpdateHealthScoreDirect 同步更新健康度（供测试使用） @@ -197,7 +196,7 @@ func checkAutoDisable(ctx context.Context, channelID int64, consecutiveFailures g.Log().Warningf(ctx, "[AutoDisable] channel %d auto-disabled after %d consecutive failures", channelID, consecutiveFailures) // 清除该渠道的所有亲和性记录，避免后续请求继续路由到已禁用渠道 - scheduler.GetGlobalAffinity().DeleteByChannel(channelID) + InvalidateChannelAffinities(ctx, channelID) // 查询渠道名称用于通知 var chName *string diff --git a/internal/logic/relay/provider.go b/internal/logic/relay/provider.go index a01c7d6..0babb6a 100644 --- a/internal/logic/relay/provider.go +++ b/internal/logic/relay/provider.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "sort" "strconv" "strings" "time" @@ -58,7 +59,7 @@ func (p *DataProviderImpl) ValidateApiKey(ctx context.Context, rawKey string) (* // GetChannelForModel 实现 DataProvider.GetChannelForModel // 使用调度引擎：亲和性优先 → 租户模型权限 → 渠道范围过滤 → 优先级+权重选择 → 健康度过滤 -func (p *DataProviderImpl) GetChannelForModel(ctx context.Context, tenantID, userID int64, modelName string, excludeChannelIDs []int64) (*common.ChannelSelection, error) { +func (p *DataProviderImpl) GetChannelForModel(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, excludeChannelIDs []int64) (*common.ChannelSelection, error) { // 检查租户模型权限和渠道范围 enabled, channelScope, err := p.CheckTenantModelAccess(ctx, tenantID, modelName) if err != nil { @@ -68,90 +69,19 @@ func (p *DataProviderImpl) GetChannelForModel(ctx context.Context, tenantID, use return nil, common.ErrTenantModelNotEnabled } - // 亲和性查找：优先使用上次成功的渠道，保持会话连续性 - if preferredChannelID, ok := scheduler.GetGlobalAffinity().Get(tenantID, userID, modelName); ok { - if selection, err := p.tryAffinityChannel(ctx, tenantID, modelName, preferredChannelID, channelScope, excludeChannelIDs); err == nil && selection != nil { - return selection, nil - } - // 亲和性渠道不可用，清除记录，继续正常调度 - scheduler.GetGlobalAffinity().Delete(tenantID, userID, modelName) - } - - return p.selectChannelFromDB(ctx, tenantID, modelName, channelScope, excludeChannelIDs) + return p.selectChannelFromDB(ctx, tenantID, userID, apiKeyID, modelName, channelScope, excludeChannelIDs) } -// tryAffinityChannel 尝试使用亲和性渠道 -func (p *DataProviderImpl) tryAffinityChannel(ctx context.Context, tenantID int64, modelName string, preferredChannelID int64, channelScope []int64, excludeChannelIDs []int64) (*common.ChannelSelection, error) { - // 检查是否在排除列表中 - for _, id := range excludeChannelIDs { - if id == preferredChannelID { - return nil, common.ErrChannelUnavailable - } - } +func (p *DataProviderImpl) SetChannelAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, channelID int64) { + setAffinity(ctx, tenantID, userID, apiKeyID, modelName, channelID) +} - // 检查渠道范围 - if len(channelScope) > 0 { - found := false - for _, id := range channelScope { - if id == preferredChannelID { - found = true - break - } - } - if !found { - return nil, common.ErrChannelUnavailable - } - } - - type channelRow struct { - ChannelID int64 `json:"channel_id"` - ChannelName string `json:"channel_name"` - ChannelType int `json:"channel_type"` - BaseURL string `json:"base_url"` - UpstreamModel string `json:"upstream_model"` - MaxConcurrency int `json:"max_concurrency"` - Settings string `json:"settings"` - } - - var ch *channelRow - err := dao.ChnChannels.Ctx(ctx).As("c"). - LeftJoin("chn_abilities a ON a.channel_id = c.id"). - Where("c.id", preferredChannelID). - Where("c.status", "active"). - Where("a.model_name", modelName). - Where("a.enabled", true). - Fields("c.id as channel_id, c.name as channel_name, c.type as channel_type, c.base_url, a.upstream_model, c.max_concurrency, c.settings"). - Scan(&ch) - if err != nil || ch == nil { - return nil, common.ErrChannelUnavailable - } - - key, err := getChannelKey(ctx, ch.ChannelID) - if err != nil || key == "" { - return nil, common.ErrChannelUnavailable - } - - settings := ParseChannelSettings(ch.Settings) - upstreamModel := ch.UpstreamModel - if upstreamModel == "" { - upstreamModel = modelName - } - - return &common.ChannelSelection{ - ChannelID: ch.ChannelID, - ChannelType: ch.ChannelType, - ChannelName: ch.ChannelName, - BaseURL: ch.BaseURL, - ApiKey: key, - UpstreamModelName: upstreamModel, - IsModelMapped: ch.UpstreamModel != "" && ch.UpstreamModel != modelName, - MaxConcurrency: ch.MaxConcurrency, - Settings: settings, - }, nil +func (p *DataProviderImpl) DeleteChannelAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string) { + deleteAffinity(ctx, tenantID, userID, apiKeyID, modelName) } // selectChannelFromDB 正常渠道调度（优先级+权重+健康度） -func (p *DataProviderImpl) selectChannelFromDB(ctx context.Context, tenantID int64, modelName string, channelScope []int64, excludeChannelIDs []int64) (*common.ChannelSelection, error) { +func (p *DataProviderImpl) selectChannelFromDB(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, channelScope []int64, excludeChannelIDs []int64) (*common.ChannelSelection, error) { type channelRow struct { ChannelID int64 `json:"channel_id"` ChannelName string `json:"channel_name"` @@ -172,9 +102,7 @@ func (p *DataProviderImpl) selectChannelFromDB(ctx context.Context, tenantID int Where("a.model_name", modelName). Where("a.enabled", true). Where("c.status", "active"). - Fields("c.id as channel_id, c.name as channel_name, c.type as channel_type, c.base_url, a.upstream_model, c.priority, c.weight, c.max_concurrency, c.settings, h.health_score, h.consecutive_failures"). - OrderDesc("c.priority"). - OrderDesc("c.weight") + Fields("c.id as channel_id, c.name as channel_name, c.type as channel_type, c.base_url, a.upstream_model, c.priority, c.weight, c.max_concurrency, c.settings, h.health_score, h.consecutive_failures") // 渠道范围过滤 if len(channelScope) > 0 { @@ -195,17 +123,22 @@ func (p *DataProviderImpl) selectChannelFromDB(ctx context.Context, tenantID int return nil, common.ErrChannelUnavailable } - // 尝试按 last_used_at 选择最早使用的 Key - for _, ch := range channels { - key, err := getChannelKey(ctx, ch.ChannelID) - if err == nil && key != "" { - settings := ParseChannelSettings(ch.Settings) - + if !lcommon.Config().GetBool(ctx, "channel_scheduler_v2_enabled") { + sort.SliceStable(channels, func(i, j int) bool { + if channels[i].Priority == channels[j].Priority { + return channels[i].Weight > channels[j].Weight + } + return channels[i].Priority > channels[j].Priority + }) + for _, ch := range channels { + key, keyErr := getChannelKey(ctx, ch.ChannelID) + if keyErr != nil || key == "" { + continue + } upstreamModel := ch.UpstreamModel if upstreamModel == "" { upstreamModel = modelName } - return &common.ChannelSelection{ ChannelID: ch.ChannelID, ChannelType: ch.ChannelType, @@ -215,14 +148,118 @@ func (p *DataProviderImpl) selectChannelFromDB(ctx context.Context, tenantID int UpstreamModelName: upstreamModel, IsModelMapped: ch.UpstreamModel != "" && ch.UpstreamModel != modelName, MaxConcurrency: ch.MaxConcurrency, - Settings: settings, + Settings: ParseChannelSettings(ch.Settings), + SelectionReason: "legacy", + Priority: ch.Priority, + Weight: ch.Weight, + HealthScore: 100, }, nil } + return nil, common.ErrChannelUnavailable + } + + candidates := make([]scheduler.ChannelCandidate, 0, len(channels)) + for _, ch := range channels { + healthScore := 100.0 + if ch.HealthScore != nil { + healthScore = *ch.HealthScore + } + upstreamModel := ch.UpstreamModel + if upstreamModel == "" { + upstreamModel = modelName + } + candidates = append(candidates, scheduler.ChannelCandidate{ + ChannelID: ch.ChannelID, + ChannelType: ch.ChannelType, + ChannelName: ch.ChannelName, + BaseURL: ch.BaseURL, + Priority: ch.Priority, + Weight: ch.Weight, + HealthScore: healthScore, + UpstreamModel: upstreamModel, + IsModelMapped: ch.UpstreamModel != "" && ch.UpstreamModel != modelName, + Settings: ParseChannelSettings(ch.Settings), + MaxConcurrency: ch.MaxConcurrency, + }) + } + + preferredChannelID, _ := getAffinity(ctx, tenantID, userID, apiKeyID, modelName) + _, seed := affinityIdentity(tenantID, userID, apiKeyID, modelName) + bucketSeconds := int64(affinityTTL(ctx).Seconds()) + seed += ":" + strconv.FormatInt(time.Now().Unix()/bucketSeconds, 10) + + for len(candidates) > 0 { + var result *scheduler.SchedulerResult + if affinityEnabled(ctx) { + result = scheduler.SelectStable(candidates, preferredChannelID, seed) + } else { + result = scheduler.Select(candidates) + } + if result == nil { + break + } + affinityHit := preferredChannelID > 0 && result.ChannelID == preferredChannelID + preserveAffinity := false + if preferredChannelID > 0 && !affinityHit { + if containsChannelID(excludeChannelIDs, preferredChannelID) { + preserveAffinity = true + } else { + deleteAffinity(ctx, tenantID, userID, apiKeyID, modelName) + preferredChannelID = 0 + } + } + + key, keyErr := getChannelKey(ctx, result.ChannelID) + if keyErr == nil && key != "" { + reason := "weighted" + if affinityHit { + reason = "affinity" + } + return &common.ChannelSelection{ + ChannelID: result.ChannelID, + ChannelType: result.ChannelType, + ChannelName: result.ChannelName, + BaseURL: result.BaseURL, + ApiKey: key, + UpstreamModelName: result.UpstreamModelName, + IsModelMapped: result.IsModelMapped, + MaxConcurrency: result.MaxConcurrency, + Settings: result.Settings, + SelectionReason: reason, + PreserveAffinity: preserveAffinity, + Priority: result.Priority, + Weight: result.Weight, + HealthScore: result.HealthScore, + }, nil + } + if affinityHit { + deleteAffinity(ctx, tenantID, userID, apiKeyID, modelName) + preferredChannelID = 0 + } + candidates = removeChannelCandidate(candidates, result.ChannelID) } return nil, common.ErrChannelUnavailable } +func containsChannelID(channelIDs []int64, channelID int64) bool { + for _, id := range channelIDs { + if id == channelID { + return true + } + } + return false +} + +func removeChannelCandidate(candidates []scheduler.ChannelCandidate, channelID int64) []scheduler.ChannelCandidate { + for i := range candidates { + if candidates[i].ChannelID == channelID { + return append(candidates[:i], candidates[i+1:]...) + } + } + return candidates +} + // CheckTenantModelAccess 实现 DataProvider.CheckTenantModelAccess // 检查租户是否有权使用指定模型，返回是否启用和渠道范围 func (p *DataProviderImpl) CheckTenantModelAccess(ctx context.Context, tenantID int64, modelName string) (bool, []int64, error) { diff --git a/internal/logic/task/sync_image_worker.go b/internal/logic/task/sync_image_worker.go index 5bcba84..ba49da0 100644 --- a/internal/logic/task/sync_image_worker.go +++ b/internal/logic/task/sync_image_worker.go @@ -431,7 +431,7 @@ func processSyncImageJob(job *SyncImageJob) { var exclude []int64 lastErr := "no available channel" for attempt := 0; attempt < syncImageMaxAttempts; attempt++ { - sel, err := syncImageRelayProv.GetChannelForModel(ctx, job.TenantID, job.UserID, job.Model, exclude) + sel, err := syncImageRelayProv.GetChannelForModel(ctx, job.TenantID, job.UserID, job.ApiKeyID, job.Model, exclude) if err != nil { lastErr = fmt.Sprintf("select channel: %v", err) break // 无更多候选 @@ -443,9 +443,18 @@ func processSyncImageJob(job *SyncImageJob) { syncImageJitter() continue } + if !syncImageRelayProv.AcquireChannelSlot(ctx, sel.ChannelID, sel.MaxConcurrency, job.RequestID) { + decInflight(sel.ChannelID) + exclude = append(exclude, sel.ChannelID) + syncImageJitter() + continue + } ok, memW, perr := runImagePipelineWithRelease(ctx, job, sel) if ok { + if !sel.PreserveAffinity { + syncImageRelayProv.SetChannelAffinity(ctx, job.TenantID, job.UserID, job.ApiKeyID, job.Model, sel.ChannelID) + } settleSyncImageSuccess(ctx, job, sel, memW) return } @@ -453,6 +462,7 @@ func processSyncImageJob(job *SyncImageJob) { // 失败：降健康度 + 排除该渠道 + 抖动退避重选 syncImageRelayProv.UpdateChannelHealth(ctx, sel.ChannelID, false, 0) syncImageRelayProv.IncrementConsecutiveFailure(ctx, sel.ChannelID) + syncImageRelayProv.DeleteChannelAffinity(ctx, job.TenantID, job.UserID, job.ApiKeyID, job.Model) exclude = append(exclude, sel.ChannelID) lastErr = perr syncImageJitter() @@ -464,10 +474,29 @@ func processSyncImageJob(job *SyncImageJob) { // runImagePipelineWithRelease 保证无论成功/失败/panic 都释放 per-channel 槽（defer）。 func runImagePipelineWithRelease(ctx context.Context, job *SyncImageJob, sel *common.ChannelSelection) (ok bool, memW *memResponseWriter, failReason string) { + stopHeartbeat := make(chan struct{}) + go refreshSyncImageChannelLease(sel.ChannelID, job.RequestID, stopHeartbeat) + defer close(stopHeartbeat) defer decInflight(sel.ChannelID) + defer syncImageRelayProv.ReleaseChannelSlot(context.Background(), sel.ChannelID, job.RequestID) return runImagePipeline(ctx, job, sel) } +func refreshSyncImageChannelLease(channelID int64, requestID string, stop <-chan struct{}) { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + syncImageRelayProv.RefreshChannelSlot(ctx, channelID, requestID) + cancel() + case <-stop: + return + } + } +} + // runImagePipeline 复刻 RelayHandler 的最小管线，把上游响应捕获到内存 writer。 func runImagePipeline(ctx context.Context, job *SyncImageJob, sel *common.ChannelSelection) (bool, *memResponseWriter, string) { info := &common.RelayInfo{ diff --git a/relay/common/forwarding_trace.go b/relay/common/forwarding_trace.go index 22ec013..7680fed 100644 --- a/relay/common/forwarding_trace.go +++ b/relay/common/forwarding_trace.go @@ -13,16 +13,20 @@ type ForwardingTrace struct { // ForwardingHop 单次转发跳转记录 type ForwardingHop struct { - Attempt int `json:"attempt"` - ChannelID int64 `json:"channel_id"` - ChannelName string `json:"channel_name"` - ChannelType int `json:"channel_type"` - Provider string `json:"provider"` - BaseURL string `json:"base_url"` - UpstreamURL string `json:"upstream_url"` - UpstreamModel string `json:"upstream_model"` - ModelMapped bool `json:"model_mapped"` - Success bool `json:"success"` - Error string `json:"error,omitempty"` - LatencyMs float64 `json:"latency_ms"` + Attempt int `json:"attempt"` + ChannelID int64 `json:"channel_id"` + ChannelName string `json:"channel_name"` + ChannelType int `json:"channel_type"` + Provider string `json:"provider"` + BaseURL string `json:"base_url"` + UpstreamURL string `json:"upstream_url"` + UpstreamModel string `json:"upstream_model"` + ModelMapped bool `json:"model_mapped"` + SelectionReason string `json:"selection_reason,omitempty"` + Priority int `json:"priority"` + Weight int `json:"weight"` + HealthScore float64 `json:"health_score"` + Success bool `json:"success"` + Error string `json:"error,omitempty"` + LatencyMs float64 `json:"latency_ms"` } diff --git a/relay/common/provider.go b/relay/common/provider.go index f7e2075..42c0f94 100644 --- a/relay/common/provider.go +++ b/relay/common/provider.go @@ -37,9 +37,24 @@ type DataProvider interface { ValidateApiKey(ctx context.Context, rawKey string) (*ApiKeyInfo, error) // GetChannelForModel 为指定模型选择最佳渠道。 - // tenantID 和 userID 用于亲和性计算。 + // tenantID、userID 和 apiKeyID 用于亲和性计算。 // excludeChannelIDs 是本次请求已尝试过且失败的渠道，用于重试时排除。 - GetChannelForModel(ctx context.Context, tenantID, userID int64, modelName string, excludeChannelIDs []int64) (*ChannelSelection, error) + GetChannelForModel(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, excludeChannelIDs []int64) (*ChannelSelection, error) + + // SetChannelAffinity 在请求成功后提交并刷新渠道亲和。 + SetChannelAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string, channelID int64) + + // DeleteChannelAffinity 在亲和渠道硬失败后清除绑定。 + DeleteChannelAffinity(ctx context.Context, tenantID, userID, apiKeyID int64, modelName string) + + // AcquireChannelSlot 原子占用渠道并发槽；Redis 异常时降级放行。 + AcquireChannelSlot(ctx context.Context, channelID int64, maxConcurrency int, requestID string) bool + + // RefreshChannelSlot 刷新长请求的渠道容量租约。 + RefreshChannelSlot(ctx context.Context, channelID int64, requestID string) + + // ReleaseChannelSlot 释放渠道并发槽。 + ReleaseChannelSlot(ctx context.Context, channelID int64, requestID string) // GetModelMapping 获取模型映射信息。 // 返回标准模型名和分类（chat/embedding/image 等）。 @@ -120,8 +135,13 @@ type ChannelSelection struct { ApiKey string // 解密后的上游 API Key UpstreamModelName string IsModelMapped bool - MaxConcurrency int // 该渠道最大并发（0/负值表示不限），供 sync_image worker 做 per-channel 容量控制 + MaxConcurrency int // 该渠道最大并发（0/负值表示不限），供各转发入口做容量控制 Settings ChannelSettings + SelectionReason string // affinity / weighted / legacy + PreserveAffinity bool // 临时容量溢出时不覆盖原亲和 + Priority int + Weight int + HealthScore float64 } // UsageRecord API 调用用量记录 diff --git a/relay/handler/channel_lease.go b/relay/handler/channel_lease.go new file mode 100644 index 0000000..2315627 --- /dev/null +++ b/relay/handler/channel_lease.go @@ -0,0 +1,61 @@ +package handler + +import ( + "context" + "sync" + "time" + + "github.com/qianfree/team-api/relay/common" +) + +type channelLease struct { + provider common.DataProvider + channelID int64 + requestID string + stop chan struct{} + once sync.Once +} + +func acquireChannelLease(ctx context.Context, provider common.DataProvider, selection *common.ChannelSelection, requestID string) (*channelLease, bool) { + if selection.MaxConcurrency <= 0 { + return nil, true + } + if !provider.AcquireChannelSlot(ctx, selection.ChannelID, selection.MaxConcurrency, requestID) { + return nil, false + } + lease := &channelLease{ + provider: provider, + channelID: selection.ChannelID, + requestID: requestID, + stop: make(chan struct{}), + } + go lease.refreshLoop() + return lease, true +} + +func (l *channelLease) refreshLoop() { + ticker := time.NewTicker(30 * time.Second) + defer ticker.Stop() + for { + select { + case <-ticker.C: + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + l.provider.RefreshChannelSlot(ctx, l.channelID, l.requestID) + cancel() + case <-l.stop: + return + } + } +} + +func (l *channelLease) Release() { + if l == nil { + return + } + l.once.Do(func() { + close(l.stop) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + l.provider.ReleaseChannelSlot(ctx, l.channelID, l.requestID) + }) +} diff --git a/relay/handler/realtime_handler.go b/relay/handler/realtime_handler.go index 1ca86d6..9d29a07 100644 --- a/relay/handler/realtime_handler.go +++ b/relay/handler/realtime_handler.go @@ -152,10 +152,24 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, } // 4. 渠道选择 - selection, err := provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, modelName, nil) - if err != nil { - return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, err) + var ( + selection *common.ChannelSelection + lease *channelLease + excluded []int64 + ) + for { + selection, err = provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName, excluded) + if err != nil { + return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, err) + } + var acquired bool + lease, acquired = acquireChannelLease(ctx, provider, selection, rc.RequestID) + if acquired { + break + } + excluded = append(excluded, selection.ChannelID) } + defer lease.Release() // 5. 构造 RelayInfo info := &common.RelayInfo{ @@ -227,6 +241,7 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, // 8. 建立 Realtime 代理 proxy := openai.NewRealtimeProxy(info) if err := proxy.DialUpstream(); err != nil { + provider.DeleteChannelAffinity(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName) if billing != nil && preDeductAmount > 0 { _ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount) } @@ -244,6 +259,9 @@ func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, // 9. 启动双向代理 usage, proxyErr := proxy.Proxy(ctx) + if proxyErr == nil && !selection.PreserveAffinity { + provider.SetChannelAffinity(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName, selection.ChannelID) + } // 10. 结算费用 if billing != nil && preDeductAmount > 0 { diff --git a/relay/handler/relay_handler.go b/relay/handler/relay_handler.go index b2b3223..f63827f 100644 --- a/relay/handler/relay_handler.go +++ b/relay/handler/relay_handler.go @@ -21,7 +21,6 @@ import ( "github.com/qianfree/team-api/relay/constant" "github.com/qianfree/team-api/relay/helper" "github.com/qianfree/team-api/relay/override" - "github.com/qianfree/team-api/relay/scheduler" ) // RelayContext relay 请求上下文（从 GoFrame handler 传入） @@ -257,7 +256,9 @@ func settleSuccessfulRequest( // 17. 更新健康度 + 记录用量 + 更新亲和性 provider.UpdateChannelHealth(postCtx, selection.ChannelID, true, info.LatencyMs()) provider.ResetConsecutiveFailure(postCtx, selection.ChannelID) - scheduler.GetGlobalAffinity().Set(rc.TenantID, rc.UserID, v.modelName, selection.ChannelID) + if !selection.PreserveAffinity { + provider.SetChannelAffinity(postCtx, rc.TenantID, rc.UserID, rc.ApiKeyID, v.lookupModel, selection.ChannelID) + } firstTokenMs := 0 if !info.FirstResponseTime.IsZero() { @@ -414,7 +415,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He } for attempt := 0; attempt <= maxRetries; attempt++ { - selection, err := provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, v.lookupModel, excludeChannelIDs) + selection, err := provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, v.lookupModel, excludeChannelIDs) if err != nil { result := handleChannelUnavailable(ctx, billing, provider, rc, v, preDeductAmount, channelErrors, err) return result.usage, result.billingResult, result.err @@ -466,6 +467,13 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He return nil, v.billingResult, err } + lease, acquired := acquireChannelLease(ctx, provider, selection, rc.RequestID) + if !acquired { + excludeChannelIDs = append(excludeChannelIDs, selection.ChannelID) + attempt-- // 容量溢出不消耗上游重试次数 + continue + } + upstreamCtx := context.WithoutCancel(ctx) settleCtx, settleCancel := context.WithTimeout(context.Background(), 30*time.Second) defer settleCancel() @@ -473,6 +481,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He // 发送请求到上游 resp, err := adaptor.DoRequest(upstreamCtx, info, convertedBody) if err != nil { + lease.Release() // 高成本非幂等生成（图片/视频）：DoRequest 阶段「可能已送达上游」的模糊网络错误 // （EOF/RST/超时）不重试，避免上游重复生成 + 重复计费；连接被拒/DNS 失败等确定 // 未送达的错误与状态码类错误仍照常重试。chat 等端点行为不变。 @@ -488,7 +497,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He provider.UpdateChannelHealth(settleCtx, selection.ChannelID, false, info.LatencyMs()) provider.IncrementConsecutiveFailure(settleCtx, selection.ChannelID) excludeChannelIDs = append(excludeChannelIDs, selection.ChannelID) - scheduler.GetGlobalAffinity().Delete(rc.TenantID, rc.UserID, v.modelName) + provider.DeleteChannelAffinity(settleCtx, rc.TenantID, rc.UserID, rc.ApiKeyID, v.lookupModel) if retryable && attempt < maxRetries { recordChannelError(rc, selection, v.modelName, attempt, false, err, info.LatencyMs()) @@ -512,6 +521,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He // 处理上游响应 usage, err := adaptor.DoResponse(ctx, resp, info, rc.Writer) + lease.Release() if err != nil { err = helper.RemapStatusCode(err, info.ChannelMeta.Settings.StatusCodeMapping) @@ -552,7 +562,7 @@ func RelayHandler(ctx context.Context, body []byte, path string, headers http.He failReason := fmt.Sprintf("attempt=%d channel=%d(%s) model=%s doResponse_error=[%v] latency=%.0fms", attempt, selection.ChannelID, selection.ChannelName, v.modelName, err, info.LatencyMs()) channelErrors = append(channelErrors, failReason) - scheduler.GetGlobalAffinity().Delete(rc.TenantID, rc.UserID, v.modelName) + provider.DeleteChannelAffinity(settleCtx, rc.TenantID, rc.UserID, rc.ApiKeyID, v.lookupModel) // 响应已由 adaptor 直接写入客户端（原生错误透传）时，即使 StatusCode 可重试也必须 // 终止重试：否则重试成功会把成功体追加到已写的错误体后，造成响应污染（错误体+成功体拼接）。 @@ -729,15 +739,19 @@ func buildTraceHop(attempt int, selection *common.ChannelSelection, adaptor comm upstreamURL = u } return common.ForwardingHop{ - Attempt: attempt, - ChannelID: selection.ChannelID, - ChannelName: selection.ChannelName, - ChannelType: selection.ChannelType, - Provider: constant.ProviderType(selection.ChannelType).String(), - BaseURL: selection.BaseURL, - UpstreamURL: upstreamURL, - UpstreamModel: selection.UpstreamModelName, - ModelMapped: selection.IsModelMapped, + Attempt: attempt, + ChannelID: selection.ChannelID, + ChannelName: selection.ChannelName, + ChannelType: selection.ChannelType, + Provider: constant.ProviderType(selection.ChannelType).String(), + BaseURL: selection.BaseURL, + UpstreamURL: upstreamURL, + UpstreamModel: selection.UpstreamModelName, + ModelMapped: selection.IsModelMapped, + SelectionReason: selection.SelectionReason, + Priority: selection.Priority, + Weight: selection.Weight, + HealthScore: selection.HealthScore, } } diff --git a/relay/scheduler/affinity.go b/relay/scheduler/affinity.go deleted file mode 100644 index 5099991..0000000 --- a/relay/scheduler/affinity.go +++ /dev/null @@ -1,121 +0,0 @@ -package scheduler - -import ( - "sync" - "time" -) - -const ( - // AffinityTTL 亲和性缓存默认 TTL（秒） - AffinityTTL = 1800 -) - -// affinityKey 亲和性缓存键 -type affinityKey struct { - TenantID int64 - UserID int64 - ModelName string -} - -// AffinityStore 亲和性内存缓存 -// 用户+模型 → 渠道ID 映射，避免频繁切换渠道导致上下文丢失 -type AffinityStore struct { - mu sync.RWMutex - entries map[affinityKey]*affinityEntry -} - -type affinityEntry struct { - ChannelID int64 - HitCount int - ExpiresAt time.Time -} - -// NewAffinityStore 创建亲和性缓存 -func NewAffinityStore() *AffinityStore { - return &AffinityStore{ - entries: make(map[affinityKey]*affinityEntry), - } -} - -// Get 获取亲和性渠道 -func (s *AffinityStore) Get(tenantID, userID int64, modelName string) (int64, bool) { - s.mu.RLock() - defer s.mu.RUnlock() - - key := affinityKey{TenantID: tenantID, UserID: userID, ModelName: modelName} - entry, ok := s.entries[key] - if !ok || time.Now().After(entry.ExpiresAt) { - return 0, false - } - return entry.ChannelID, true -} - -// Set 设置亲和性渠道 -func (s *AffinityStore) Set(tenantID, userID int64, modelName string, channelID int64) { - s.mu.Lock() - defer s.mu.Unlock() - - key := affinityKey{TenantID: tenantID, UserID: userID, ModelName: modelName} - if existing, ok := s.entries[key]; ok { - existing.ChannelID = channelID - existing.HitCount++ - existing.ExpiresAt = time.Now().Add(AffinityTTL * time.Second) - } else { - s.entries[key] = &affinityEntry{ - ChannelID: channelID, - HitCount: 1, - ExpiresAt: time.Now().Add(AffinityTTL * time.Second), - } - } -} - -// Delete 删除亲和性记录 -func (s *AffinityStore) Delete(tenantID, userID int64, modelName string) { - s.mu.Lock() - defer s.mu.Unlock() - - delete(s.entries, affinityKey{TenantID: tenantID, UserID: userID, ModelName: modelName}) -} - -// DeleteByChannel 删除某渠道的所有亲和性记录 -func (s *AffinityStore) DeleteByChannel(channelID int64) { - s.mu.Lock() - defer s.mu.Unlock() - - for key, entry := range s.entries { - if entry.ChannelID == channelID { - delete(s.entries, key) - } - } -} - -// CleanExpired 清理过期记录 -func (s *AffinityStore) CleanExpired() int { - s.mu.Lock() - defer s.mu.Unlock() - - now := time.Now() - count := 0 - for key, entry := range s.entries { - if now.After(entry.ExpiresAt) { - delete(s.entries, key) - count++ - } - } - return count -} - -// Size 返回缓存条目数 -func (s *AffinityStore) Size() int { - s.mu.RLock() - defer s.mu.RUnlock() - return len(s.entries) -} - -// 全局亲和性缓存实例 -var globalAffinity = NewAffinityStore() - -// GetGlobalAffinity 获取全局亲和性缓存 -func GetGlobalAffinity() *AffinityStore { - return globalAffinity -} diff --git a/relay/scheduler/affinity_test.go b/relay/scheduler/affinity_test.go deleted file mode 100644 index 54fea73..0000000 --- a/relay/scheduler/affinity_test.go +++ /dev/null @@ -1,144 +0,0 @@ -package scheduler - -import ( - "testing" - "time" -) - -func TestAffinityStore_SetGet(t *testing.T) { - store := NewAffinityStore() - - store.Set(1, 1, "gpt-4o", 100) - ch, ok := store.Get(1, 1, "gpt-4o") - if !ok { - t.Fatal("expected affinity to be found") - } - if ch != 100 { - t.Errorf("expected channel 100, got %d", ch) - } -} - -func TestAffinityStore_NotFound(t *testing.T) { - store := NewAffinityStore() - - _, ok := store.Get(1, 999, "nonexistent") - if ok { - t.Error("expected affinity not to be found") - } -} - -func TestAffinityStore_Expiry(t *testing.T) { - store := NewAffinityStore() - - // Manually set an expired entry - store.mu.Lock() - store.entries[affinityKey{TenantID: 1, UserID: 1, ModelName: "test"}] = &affinityEntry{ - ChannelID: 100, - ExpiresAt: time.Now().Add(-1 * time.Second), - } - store.mu.Unlock() - - _, ok := store.Get(1, 1, "test") - if ok { - t.Error("expected expired affinity not to be found") - } -} - -func TestAffinityStore_Delete(t *testing.T) { - store := NewAffinityStore() - store.Set(1, 1, "gpt-4o", 100) - store.Delete(1, 1, "gpt-4o") - - _, ok := store.Get(1, 1, "gpt-4o") - if ok { - t.Error("expected affinity to be deleted") - } -} - -func TestAffinityStore_DeleteByChannel(t *testing.T) { - store := NewAffinityStore() - store.Set(1, 1, "gpt-4o", 100) - store.Set(1, 2, "gpt-4o", 100) // same model, different user - store.Set(1, 1, "claude-3", 200) - - store.DeleteByChannel(100) - - _, ok1 := store.Get(1, 1, "gpt-4o") - _, ok2 := store.Get(1, 2, "gpt-4o") - _, ok3 := store.Get(1, 1, "claude-3") - - if ok1 || ok2 { - t.Error("expected channel 100 affinities to be deleted") - } - if !ok3 { - t.Error("expected channel 200 affinity to remain") - } -} - -func TestAffinityStore_HitCount(t *testing.T) { - store := NewAffinityStore() - - store.Set(1, 1, "gpt-4o", 100) - store.Set(1, 1, "gpt-4o", 100) // same key, should increment - - store.mu.RLock() - entry := store.entries[affinityKey{TenantID: 1, UserID: 1, ModelName: "gpt-4o"}] - store.mu.RUnlock() - - if entry.HitCount != 2 { - t.Errorf("expected hit count 2, got %d", entry.HitCount) - } -} - -func TestAffinityStore_TenantIsolation(t *testing.T) { - store := NewAffinityStore() - - // Tenant 1, User 1 → channel 100 - store.Set(1, 1, "gpt-4o", 100) - // Tenant 2, User 1 → channel 200 (same userID, different tenant) - store.Set(2, 1, "gpt-4o", 200) - - ch1, ok1 := store.Get(1, 1, "gpt-4o") - ch2, ok2 := store.Get(2, 1, "gpt-4o") - - if !ok1 || ch1 != 100 { - t.Errorf("tenant 1 expected channel 100, got %d, ok=%v", ch1, ok1) - } - if !ok2 || ch2 != 200 { - t.Errorf("tenant 2 expected channel 200, got %d, ok=%v", ch2, ok2) - } - - // Deleting tenant 1's affinity should not affect tenant 2 - store.Delete(1, 1, "gpt-4o") - _, ok1 = store.Get(1, 1, "gpt-4o") - _, ok2 = store.Get(2, 1, "gpt-4o") - if ok1 { - t.Error("expected tenant 1 affinity to be deleted") - } - if !ok2 { - t.Error("expected tenant 2 affinity to remain") - } -} - -func TestAffinityStore_CleanExpired(t *testing.T) { - store := NewAffinityStore() - - store.Set(1, 1, "active", 100) - store.Set(1, 2, "active", 200) - - // Manually add an expired entry - store.mu.Lock() - store.entries[affinityKey{TenantID: 1, UserID: 3, ModelName: "expired"}] = &affinityEntry{ - ChannelID: 300, - ExpiresAt: time.Now().Add(-1 * time.Second), - } - store.mu.Unlock() - - cleaned := store.CleanExpired() - if cleaned != 1 { - t.Errorf("expected 1 expired entry cleaned, got %d", cleaned) - } - if store.Size() != 2 { - t.Errorf("expected 2 remaining entries, got %d", store.Size()) - } -} diff --git a/relay/scheduler/scheduler.go b/relay/scheduler/scheduler.go index 44aabaf..9200989 100644 --- a/relay/scheduler/scheduler.go +++ b/relay/scheduler/scheduler.go @@ -1,24 +1,29 @@ package scheduler import ( + "crypto/sha256" + "encoding/binary" + "math" "math/rand/v2" "sort" + "strconv" "github.com/qianfree/team-api/relay/common" ) // ChannelCandidate 调度候选渠道 type ChannelCandidate struct { - ChannelID int64 - ChannelType int - ChannelName string - BaseURL string - Priority int - Weight int - HealthScore float64 - UpstreamModel string - IsModelMapped bool - Settings common.ChannelSettings + ChannelID int64 + ChannelType int + ChannelName string + BaseURL string + Priority int + Weight int + HealthScore float64 + UpstreamModel string + IsModelMapped bool + Settings common.ChannelSettings + MaxConcurrency int } // SchedulerResult 调度结果 @@ -33,6 +38,7 @@ type SchedulerResult struct { UpstreamModelName string IsModelMapped bool Settings common.ChannelSettings + MaxConcurrency int } // Select 从候选渠道中选择最佳渠道 @@ -41,28 +47,10 @@ type SchedulerResult struct { // 2. 按优先级分组，取最高优先级组 // 3. 组内按权重加权随机选择 func Select(candidates []ChannelCandidate) *SchedulerResult { - if len(candidates) == 0 { + highestPriority := highestPriorityCandidates(candidates) + if len(highestPriority) == 0 { return nil } - - // 排除健康度过低的渠道 - healthy := make([]ChannelCandidate, 0, len(candidates)) - for _, c := range candidates { - if c.HealthScore >= 20 { - healthy = append(healthy, c) - } - } - - // 如果全部不健康，降级使用所有渠道 - if len(healthy) == 0 { - healthy = candidates - } - - // 按优先级分组 - groups := groupByPriority(healthy) - - // 取最高优先级组 - highestPriority := groups[len(groups)-1] if len(highestPriority) == 1 { return candidateToResult(highestPriority[0]) } @@ -72,6 +60,51 @@ func Select(candidates []ChannelCandidate) *SchedulerResult { return candidateToResult(selected) } +// SelectStable 在最高可用优先级组内优先使用亲和渠道，否则按权重做稳定选择。 +// 相同 seed 和候选集在不同进程中会得到相同结果，适合并发冷请求和多实例部署。 +func SelectStable(candidates []ChannelCandidate, preferredChannelID int64, seed string) *SchedulerResult { + highestPriority := highestPriorityCandidates(candidates) + if len(highestPriority) == 0 { + return nil + } + + if preferredChannelID > 0 { + for _, candidate := range highestPriority { + if candidate.ChannelID == preferredChannelID { + return candidateToResult(candidate) + } + } + } + + selected := weightedRendezvousSelect(highestPriority, seed) + return candidateToResult(selected) +} + +func highestPriorityCandidates(candidates []ChannelCandidate) []ChannelCandidate { + selectable := make([]ChannelCandidate, 0, len(candidates)) + for _, candidate := range candidates { + if candidate.Weight > 0 { + selectable = append(selectable, candidate) + } + } + if len(selectable) == 0 { + return nil + } + + healthy := make([]ChannelCandidate, 0, len(selectable)) + for _, candidate := range selectable { + if candidate.HealthScore >= 20 { + healthy = append(healthy, candidate) + } + } + if len(healthy) == 0 { + healthy = selectable + } + + groups := groupByPriority(healthy) + return groups[len(groups)-1] +} + // groupByPriority 按优先级分组，返回按优先级升序排列的组 func groupByPriority(candidates []ChannelCandidate) [][]ChannelCandidate { if len(candidates) == 0 { @@ -128,10 +161,31 @@ func weightedRandomSelect(candidates []ChannelCandidate) ChannelCandidate { return candidates[len(candidates)-1] } +func weightedRendezvousSelect(candidates []ChannelCandidate, seed string) ChannelCandidate { + selected := candidates[0] + bestScore := math.Inf(1) + for _, candidate := range candidates { + h := sha256.New() + _, _ = h.Write([]byte(seed)) + _, _ = h.Write([]byte{0}) + _, _ = h.Write([]byte(strconv.FormatInt(candidate.ChannelID, 10))) + sum := h.Sum(nil) + raw := binary.BigEndian.Uint64(sum[:8]) + // 使用 53 位有效精度构造 (0,1) 区间，避免 log(0)。 + u := (float64(raw>>11) + 1) / (float64(uint64(1)<<53) + 1) + score := -math.Log(u) / float64(effectiveWeight(candidate)) + if score < bestScore || (score == bestScore && candidate.ChannelID < selected.ChannelID) { + bestScore = score + selected = candidate + } + } + return selected +} + func effectiveWeight(c ChannelCandidate) int { w := c.Weight if w <= 0 { - w = 1 + return 0 } // 健康度降权：50-79 权重减半，20-49 权重降到 1/4 @@ -160,5 +214,6 @@ func candidateToResult(c ChannelCandidate) *SchedulerResult { UpstreamModelName: c.UpstreamModel, IsModelMapped: c.IsModelMapped, Settings: c.Settings, + MaxConcurrency: c.MaxConcurrency, } } diff --git a/relay/scheduler/scheduler_test.go b/relay/scheduler/scheduler_test.go index 203ff9b..16e6f58 100644 --- a/relay/scheduler/scheduler_test.go +++ b/relay/scheduler/scheduler_test.go @@ -1,6 +1,7 @@ package scheduler import ( + "fmt" "testing" ) @@ -116,9 +117,9 @@ func TestSelect_HealthDegradation(t *testing.T) { expect: 25, }, { - name: "zero weight defaults to selectable", + name: "zero weight is disabled", input: ChannelCandidate{Weight: 0, HealthScore: 80}, - expect: 1, + expect: 0, }, { name: "degraded low weight remains selectable", @@ -136,6 +137,78 @@ func TestSelect_HealthDegradation(t *testing.T) { } } +func TestSelect_ZeroWeightExcluded(t *testing.T) { + candidates := []ChannelCandidate{ + {ChannelID: 1, Priority: 10, Weight: 0, HealthScore: 100}, + {ChannelID: 2, Priority: 10, Weight: 10, HealthScore: 100}, + } + for i := 0; i < 20; i++ { + if result := Select(candidates); result == nil || result.ChannelID != 2 { + t.Fatalf("zero-weight channel selected: %+v", result) + } + } +} + +func TestSelectStable_AffinityAndDeterminism(t *testing.T) { + candidates := []ChannelCandidate{ + {ChannelID: 1, Priority: 10, Weight: 7, HealthScore: 100}, + {ChannelID: 2, Priority: 10, Weight: 3, HealthScore: 100}, + } + + first := SelectStable(candidates, 0, "tenant:user:model:bucket") + for i := 0; i < 100; i++ { + result := SelectStable(candidates, 0, "tenant:user:model:bucket") + if result.ChannelID != first.ChannelID { + t.Fatalf("stable selection changed: first=%d current=%d", first.ChannelID, result.ChannelID) + } + } + if result := SelectStable(candidates, 2, "ignored-for-affinity"); result.ChannelID != 2 { + t.Fatalf("affinity channel not selected: got %d", result.ChannelID) + } +} + +func TestSelectStable_AffinityMustRemainEligible(t *testing.T) { + candidates := []ChannelCandidate{ + {ChannelID: 1, Priority: 1, Weight: 100, HealthScore: 100}, + {ChannelID: 2, Priority: 10, Weight: 100, HealthScore: 100}, + } + if result := SelectStable(candidates, 1, "subject"); result.ChannelID != 2 { + t.Fatalf("lower-priority affinity bypassed priority: got %d", result.ChannelID) + } + + candidates[0].Priority = 10 + candidates[0].HealthScore = 10 + if result := SelectStable(candidates, 1, "subject"); result.ChannelID != 2 { + t.Fatalf("unhealthy affinity bypassed health filter: got %d", result.ChannelID) + } +} + +func TestSelect_AllZeroWeightReturnsNil(t *testing.T) { + candidates := []ChannelCandidate{ + {ChannelID: 1, Priority: 10, Weight: 0, HealthScore: 100}, + {ChannelID: 2, Priority: 10, Weight: 0, HealthScore: 100}, + } + if result := SelectStable(candidates, 1, "subject"); result != nil { + t.Fatalf("expected no selectable channel, got %+v", result) + } +} + +func TestSelectStable_WeightedDistribution(t *testing.T) { + candidates := []ChannelCandidate{ + {ChannelID: 1, Priority: 10, Weight: 7, HealthScore: 100}, + {ChannelID: 2, Priority: 10, Weight: 3, HealthScore: 100}, + } + counts := map[int64]int{} + for i := 0; i < 10000; i++ { + result := SelectStable(candidates, 0, fmt.Sprintf("subject-%d", i)) + counts[result.ChannelID]++ + } + ratio := float64(counts[1]) / 10000 + if ratio < 0.67 || ratio > 0.73 { + t.Fatalf("7:3 distribution out of tolerance: counts=%v", counts) + } +} + func TestGroupByPriority(t *testing.T) { candidates := []ChannelCandidate{ {Priority: 1, ChannelID: 1}, diff --git a/web/admin/src/views/ai/ChannelDetailPage.vue b/web/admin/src/views/ai/ChannelDetailPage.vue index a676314..63909de 100644 --- a/web/admin/src/views/ai/ChannelDetailPage.vue +++ b/web/admin/src/views/ai/ChannelDetailPage.vue @@ -631,10 +631,16 @@ function formatHeaders(headers: Record): string { - + + + + - + + + + @@ -726,4 +732,5 @@ function formatHeaders(headers: Record): string { .test-debug-label { font-size: 12px; color: var(--color-text-3); min-width: 60px; flex-shrink: 0; padding-top: 2px; } .test-debug-value { font-size: 12px; color: var(--color-text-2); word-break: break-all; background: var(--color-fill-1); padding: 2px 6px; border-radius: 3px; } .test-debug-code { font-size: 12px; background: var(--color-fill-1); padding: 6px 8px; border-radius: 4px; max-height: 150px; overflow: auto; color: var(--color-text-3); white-space: pre-wrap; word-break: break-all; margin: 0; flex: 1; } +.field-help { color: var(--color-text-3); font-size: 12px; } diff --git a/web/admin/src/views/ai/ChannelsPage.vue b/web/admin/src/views/ai/ChannelsPage.vue index 1ac161b..5ea5991 100644 --- a/web/admin/src/views/ai/ChannelsPage.vue +++ b/web/admin/src/views/ai/ChannelsPage.vue @@ -327,23 +327,38 @@ const { exporting, exportFile } = useExport({
- 了解优先级和权重如何影响渠道调度 → + 了解优先级、权重、亲和与容量如何共同调度 →
- -
-

请求到达时，调度器按以下规则从启用渠道中选择一个处理：

-
    -
  1. 排除健康度 < 20 的渠道（如果全部不健康则降级使用全部）
    2. -
  2. 按优先级分组，只取最高优先级组的渠道（数字越大越优先）
    3. -
  3. 组内按权重加权随机选择（权重越高被选中的概率越大）
    4. -
  4. 同时根据健康度对权重降权：健康度 50-79 权重减半，20-49 权重降至 1/4
    5. -
-
-
💡 配置示例
-

两个同供应商渠道分别设 优先级=100（主力）和 优先级=50（备用），调度器始终优先使用主力渠道，仅在主力不可用时切换到备用。

-

同一优先级内两个渠道分别设 权重=70 和 权重=30，则流量按 7:3 比例分发。

+ +
+

新版渠道调度开启时，请求按以下规则选择渠道：

+ +
+
1候选渠道
+

先筛选已启用、支持当前模型、符合租户范围且存在可用 Key 的渠道。健康度低于 20 的渠道通常会被排除；如果全部低于 20，则保留作为降级候选。

+
+ +
+
2优先级与权重
+

数字越大优先级越高，仅最高可用优先级组参与分配。同组通过加权稳定选路建立新亲和，权重=0 的渠道不参与调度；健康度 50-79 时有效权重减半，20-49 时降至 1/4。

+
+ +
+
3渠道亲和
+

首次请求或亲和失效后，调度器按权重稳定选择渠道。请求成功后，同一租户、用户、API Key 和模型会尽量复用该渠道，默认有效期为 30 分钟，后续成功会刷新有效期，以提高上游缓存命中率。

+
+ +
+
4故障与容量
+

亲和渠道发生硬失败时会清除绑定并重新选择。渠道达到并发上限时，仅当前请求临时溢出到其他可用渠道，不覆盖原亲和；后续请求仍会优先尝试原渠道。

+
+ +
+
配置示例
+

优先级=100 的主力渠道可用时，不会选择 优先级=50 的备用渠道；主力组不可用后才会降级。

+

同一优先级内设置 权重=70 和 权重=30，表示大量新亲和主体约按 7:3 建立绑定。由于后续请求保持亲和，用户请求量不同会使总请求比例偏离 7:3。

 @@ -392,10 +407,16 @@ const { exporting, exportFile } = useExport({ - + + + + - + + + + @@ -513,4 +534,61 @@ const { exporting, exportFile } = useExport({ .scheduling-hint:hover { color: rgb(var(--arcoblue-6)); } +.scheduling-guide { + color: var(--color-text-1); + font-size: 14px; + line-height: 1.7; +} +.guide-intro { + margin: 0 0 14px; + color: var(--color-text-2); +} +.guide-section { + margin-bottom: 14px; +} +.guide-title { + display: flex; + align-items: center; + gap: 8px; + margin-bottom: 4px; + font-weight: 600; +} +.guide-index { + display: inline-flex; + width: 22px; + height: 22px; + align-items: center; + justify-content: center; + flex: 0 0 22px; + border-radius: 50%; + color: rgb(var(--arcoblue-6)); + background: var(--color-primary-light-1); + font-size: 12px; +} +.guide-section p, +.guide-example p { + margin: 0; +} +.guide-example { + padding: 14px 16px; + border-radius: 6px; + background: var(--color-fill-2); +} +.guide-example-title { + margin-bottom: 6px; + font-weight: 600; +} +.guide-example p + p { + margin-top: 8px; +} +.scheduling-guide code { + padding: 1px 4px; + border-radius: 3px; + color: rgb(var(--arcoblue-6)); + background: var(--color-primary-light-1); +} +.field-help { + color: var(--color-text-3); + font-size: 12px; +}