refactor:重构渠道调度功能

This commit is contained in:
martin committed 2026-07-25 20:48:15 +08:00
1 parent 75bc4038db
commit 6665b1e113
21 files changed
+749 -483

No files matched your search

-5
View File
@@ -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)
+11 -1
View File
@@ -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{
+4
View File
@@ -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
}
@@ -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",
+71 -36
View File
@@ -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)
}
}
+32
View File
@@ -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")
}
})
}
}
+52
View File
@@ -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)
}
+1 -2
View File
@@ -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
+127 -90
View File
@@ -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) {
+30 -1
View File
@@ -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{
+16 -12
View File
@@ -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"`
}
+23 -3
View File
@@ -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 调用用量记录
+61
View File
@@ -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)
})
}
+21 -3
View File
@@ -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 {
+28 -14
View File
@@ -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,
}
}
-121
View File
@@ -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
}
-144
View File
@@ -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())
}
}
+86 -31
View File
@@ -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,
}
}
+75 -2
View File
@@ -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},
+9 -2
View File
@@ -631,10 +631,16 @@ function formatHeaders(headers: Record<string, string>): string {
</ARow>
<ARow :gutter="16">
<ACol :span="6">
<AFormItem label="优先级"><AInputNumber v-model="editForm.priority" :min="0" class="w-full" /></AFormItem>
<AFormItem label="优先级">
<AInputNumber v-model="editForm.priority" :min="0" class="w-full" />
<template #extra><span class="field-help">仅最高可用优先级组参与分配</span></template>
</AFormItem>
</ACol>
<ACol :span="6">
<AFormItem label="权重"><AInputNumber v-model="editForm.weight" :min="1" :max="100" class="w-full" /></AFormItem>
<AFormItem label="权重">
<AInputNumber v-model="editForm.weight" :min="0" :max="100" class="w-full" />
<template #extra><span class="field-help">新亲和绑定比例;0 表示不参与调度</span></template>
</AFormItem>
</ACol>
<ACol :span="6">
<AFormItem label="测试模型"><AInput v-model="editForm.test_model" /></AFormItem>
@@ -726,4 +732,5 @@ function formatHeaders(headers: Record<string, string>): 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; }
</style>
+94 -16
View File
@@ -327,23 +327,38 @@ const { exporting, exportFile } = useExport({
<!-- 渠道调度说明(点击弹窗) -->
<div class="scheduling-hint" @click="showSchedulingGuide = true">
<icon-info-circle /> 了解优先级和权重如何影响渠道调度 →
<icon-info-circle /> 了解优先级、权重、亲和与容量如何共同调度 →
</div>
<!-- 调度说明弹窗 -->
<AModal v-model:visible="showSchedulingGuide" title="渠道调度逻辑说明" :width="640" :footer="false">
<div style="font-size:14px;line-height:2;color:var(--color-text-1)">
<p>请求到达时,调度器按以下规则从启用渠道中<strong>选择一个</strong>处理:</p>
<ol style="margin:8px 0 16px 20px;padding:0">
<li style="margin-bottom:8px"><strong>排除</strong>健康度 &lt; 20 的渠道(如果全部不健康则降级使用全部)</li>
<li style="margin-bottom:8px">按<strong>优先级</strong>分组,<strong>只取最高优先级组</strong>的渠道(数字<strong>越大越优先</strong>)</li>
<li style="margin-bottom:8px">组内按<strong>权重加权随机选择</strong>(权重越高被选中的概率越大)</li>
<li style="margin-bottom:8px">同时根据<strong>健康度对权重降权</strong>:健康度 50-79 权重减半,20-49 权重降至 1/4</li>
</ol>
<div style="background:var(--color-fill-2);border-radius:8px;padding:16px;margin-top:8px">
<div style="font-weight:600;margin-bottom:8px">💡 配置示例</div>
<p style="margin:0 0 4px">两个同供应商渠道分别设 <code>优先级=100</code>(主力)和 <code>优先级=50</code>(备用),调度器始终优先使用主力渠道,仅在主力不可用时切换到备用。</p>
<p style="margin:8px 0 0">同一优先级内两个渠道分别设 <code>权重=70</code> 和 <code>权重=30</code>,则流量按 7:3 比例分发。</p>
<AModal v-model:visible="showSchedulingGuide" title="新版渠道调度说明" :width="680" :footer="false">
<div class="scheduling-guide">
<p class="guide-intro">新版渠道调度开启时,请求按以下规则选择渠道:</p>
<section class="guide-section">
<div class="guide-title"><span class="guide-index">1</span>候选渠道</div>
<p>先筛选已启用、支持当前模型、符合租户范围且存在可用 Key 的渠道。健康度低于 20 的渠道通常会被排除;如果全部低于 20,则保留作为降级候选。</p>
</section>
<section class="guide-section">
<div class="guide-title"><span class="guide-index">2</span>优先级与权重</div>
<p>数字越大优先级越高,仅最高可用优先级组参与分配。同组通过加权稳定选路建立新亲和,<code>权重=0</code> 的渠道不参与调度;健康度 50-79 时有效权重减半,20-49 时降至 1/4。</p>
</section>
<section class="guide-section">
<div class="guide-title"><span class="guide-index">3</span>渠道亲和</div>
<p>首次请求或亲和失效后,调度器按权重稳定选择渠道。请求成功后,同一租户、用户、API Key 和模型会尽量复用该渠道,默认有效期为 30 分钟,后续成功会刷新有效期,以提高上游缓存命中率。</p>
</section>
<section class="guide-section">
<div class="guide-title"><span class="guide-index">4</span>故障与容量</div>
<p>亲和渠道发生硬失败时会清除绑定并重新选择。渠道达到并发上限时,仅当前请求临时溢出到其他可用渠道,不覆盖原亲和;后续请求仍会优先尝试原渠道。</p>
</section>
<div class="guide-example">
<div class="guide-example-title">配置示例</div>
<p><code>优先级=100</code> 的主力渠道可用时,不会选择 <code>优先级=50</code> 的备用渠道;主力组不可用后才会降级。</p>
<p>同一优先级内设置 <code>权重=70</code> 和 <code>权重=30</code>,表示大量新亲和主体约按 7:3 建立绑定。由于后续请求保持亲和,用户请求量不同会使总请求比例偏离 7:3。</p>
</div>
</div>
</AModal>
@@ -392,10 +407,16 @@ const { exporting, exportFile } = useExport({
<!-- 调度 & 状态 -->
<ARow :gutter="16">
<ACol :span="8">
<AFormItem label="优先级"><AInputNumber v-model="form.priority" :min="0" placeholder="越高越优先" class="w-full" /></AFormItem>
<AFormItem label="优先级">
<AInputNumber v-model="form.priority" :min="0" placeholder="越高越优先" class="w-full" />
<template #extra><span class="field-help">仅最高可用优先级组参与分配</span></template>
</AFormItem>
</ACol>
<ACol :span="8">
<AFormItem label="权重"><AInputNumber v-model="form.weight" :min="1" :max="100" class="w-full" /></AFormItem>
<AFormItem label="权重">
<AInputNumber v-model="form.weight" :min="0" :max="100" class="w-full" />
<template #extra><span class="field-help">新亲和绑定比例;0 表示不参与调度</span></template>
</AFormItem>
</ACol>
<ACol :span="8">
<AFormItem label="状态">
@@ -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;
}
</style>