// Package channel 渠道仓储:候选选择、加权负载均衡、并发控制、密钥加解密。 // M5 起支持多候选(负载均衡 + 故障转移)与每渠道并发信号量。 package channel import ( "crypto/rand" "errors" "log" "math/big" "sync" "time" "github.com/openteam/server/internal/pkg/crypto" "github.com/openteam/server/internal/store" "gorm.io/gorm" ) var ErrNoChannel = errors.New("no available channel") // maxCooldown 运行期冷却上限:上游 Retry-After 再长也不让渠道缺席超过这个时间。 const maxCooldown = 10 * time.Minute type Service struct { db *gorm.DB enc *crypto.Encryptor mu sync.Mutex sems map[uint64]chan struct{} runt map[uint64]*runtimeHealth // 运行期冷却参数;NewService 给默认值,测试可覆盖。 failThreshold int cooldownFor time.Duration } // runtimeHealth 运行期健康:连续可重试失败(网络错误/超时/429/5xx)累计到阈值进入冷却。 // 与 HealthMonitor 的探测结果相互独立——探测只看 /models 是否 200, // 覆盖不了"models 正常但对话接口一直报错"的渠道。 type runtimeHealth struct { fails int until time.Time } func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service { return &Service{ db: db, enc: enc, sems: map[uint64]chan struct{}{}, runt: map[uint64]*runtimeHealth{}, failThreshold: 3, cooldownFor: 30 * time.Second, } } // Candidate 一个候选渠道 + 该模型的映射关系。 type Candidate struct { Channel *store.Channel UpstreamModel string // 全局模型在此渠道的映射名(无绑定则为空,用客户端模型名) } // Candidates 返回可用渠道候选:健康 + 启用 + 不在运行期冷却,按优先级、权重降序、id 升序排列。 // model 非空且该模型有渠道绑定时,只返回绑定渠道(携带 upstream_model 映射); // 绑定渠道全部不可用则返回空(上层报 no_channel),不会退化到不认识该模型的全局渠道。 // 仅当模型完全没有任何绑定时才退回全局渠道。 func (s *Service) Candidates(model string) []Candidate { if model != "" { var modelIDs []uint64 var b []store.ChannelModelBinding s.db.Model(&store.Model{}).Where("name = ? AND enabled = ?", model, true).Pluck("id", &modelIDs) if len(modelIDs) > 0 { s.db.Where("model_id IN ?", modelIDs).Find(&b) } if len(b) > 0 { return s.loadBound(b) } } var chs []store.Channel s.db.Where("enabled = ? AND health_status = ?", true, store.ChannelHealthHealthy). Order("priority ASC, weight DESC, id ASC").Find(&chs) out := make([]Candidate, 0, len(chs)) for i := range chs { if s.InCooldown(chs[i].ID) { continue } out = append(out, Candidate{Channel: &chs[i]}) } return out } // loadBound 按绑定顺序加载渠道候选,过滤健康/启用/运行期冷却,携带 upstream_model 映射。 func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate { if len(bindings) == 0 { return nil } // channel_id -> 绑定(取该渠道对该模型的映射) byChannel := map[uint64]store.ChannelModelBinding{} ids := make([]uint64, 0, len(bindings)) for _, b := range bindings { if _, ok := byChannel[b.ChannelID]; !ok { ids = append(ids, b.ChannelID) } byChannel[b.ChannelID] = b } var chs []store.Channel s.db.Where("id IN ? AND enabled = ? AND health_status = ?", ids, true, store.ChannelHealthHealthy). Order("priority ASC, weight DESC, id ASC").Find(&chs) byID := map[uint64]*store.Channel{} for i := range chs { byID[chs[i].ID] = &chs[i] } out := make([]Candidate, 0, len(ids)) for _, id := range ids { if ch, ok := byID[id]; ok && !s.InCooldown(id) { out = append(out, Candidate{Channel: ch, UpstreamModel: byChannel[id].UpstreamModel}) } } return out } // AvailableModelIDs 返回对外可见的模型 ID:启用的模型且至少绑定到一个启用且健康的渠道。 // 与 Candidates 的过滤口径一致(enabled + health_status=healthy),避免暴露绑定到已停用渠道的模型。 func (s *Service) AvailableModelIDs() []uint64 { var ids []uint64 s.db.Model(&store.ChannelModelBinding{}). Joins("JOIN channels ON channels.id = channel_model_bindings.channel_id"). Joins("JOIN models ON models.id = channel_model_bindings.model_id"). Where("channels.enabled = ? AND channels.health_status = ?", true, store.ChannelHealthHealthy). Where("models.enabled = ?", true). Distinct("channel_model_bindings.model_id"). Pluck("channel_model_bindings.model_id", &ids) return ids } // Pick 按权重加权随机选一个候选渠道(负载均衡)。 func (s *Service) Pick(cands []Candidate) *store.Channel { if len(cands) == 0 { return nil } total := 0 for _, c := range cands { w := c.Channel.Weight if w <= 0 { w = 1 } total += w } n, _ := rand.Int(rand.Reader, big.NewInt(int64(total))) acc := 0 for _, c := range cands { w := c.Channel.Weight if w <= 0 { w = 1 } acc += w if int(n.Int64()) < acc { return c.Channel } } return cands[len(cands)-1].Channel } // TryAcquire 尝试获取渠道并发槽;渠道满载返回 false(调用方可溢出到其他渠道)。 // MaxConcurrency<=0 视为不限制。 func (s *Service) TryAcquire(ch *store.Channel) (func(), bool) { if ch.MaxConcurrency <= 0 { return func() {}, true } s.mu.Lock() sem, ok := s.sems[ch.ID] if !ok { sem = make(chan struct{}, ch.MaxConcurrency) s.sems[ch.ID] = sem } s.mu.Unlock() select { case sem <- struct{}{}: return func() { <-sem }, true default: return nil, false } } // UpstreamKey 解密渠道上游密钥。 func (s *Service) UpstreamKey(ch *store.Channel) (string, error) { return s.enc.Decrypt(ch.APIKeyEnc) } // --------------------------------------------------------------------------- // 运行期健康(进程内,与 /v1/models 探测独立) // ReportFailure 记一次可重试的上游失败;连续达 failThreshold 次则该渠道进入冷却, // 期间不出现在候选集里。冷却到期后需再攒满阈值才会被剔除。 // cooldown>0 时用它覆盖默认时长(上游 429 带的 Retry-After),并夹到 maxCooldown // ——不能让一个写错的重试头把渠道钉死几小时。 func (s *Service) ReportFailure(ch *store.Channel, cooldown time.Duration) { if cooldown > s.cooldownFor { if cooldown > maxCooldown { cooldown = maxCooldown } } else { cooldown = s.cooldownFor } s.mu.Lock() st := s.runt[ch.ID] if st == nil { st = &runtimeHealth{} s.runt[ch.ID] = st } st.fails++ tripped := false if st.fails >= s.failThreshold { st.fails = 0 st.until = time.Now().Add(cooldown) tripped = true } s.mu.Unlock() if tripped { log.Printf("channel: %q entered runtime cooldown (%s)", ch.Name, cooldown) } } // ReportSuccess 上游正常返回响应头,清空该渠道的连续失败计数。 func (s *Service) ReportSuccess(ch *store.Channel) { s.mu.Lock() if st, ok := s.runt[ch.ID]; ok { st.fails = 0 } s.mu.Unlock() } // RuntimeState 渠道运行期健康快照:连续失败次数与冷却剩余时间(管理后台展示用)。 // 只读,不做过期清理——清理仍由 InCooldown 在选路时惰性完成。 func (s *Service) RuntimeState(id uint64) (fails int, remaining time.Duration) { s.mu.Lock() defer s.mu.Unlock() st, ok := s.runt[id] if !ok { return 0, 0 } if !st.until.IsZero() { if d := time.Until(st.until); d > 0 { return 0, d // 冷却中:计数在写入时已归零 } return 0, 0 // 已到期,等下次选路清理 } return st.fails, 0 } // ClearRuntime 清除运行期状态(管理后台"立即恢复")。 // 不改 DB 的 health_status——那条由探测/管理端点负责,两者语义不同不要混在一起。 func (s *Service) ClearRuntime(id uint64) { s.mu.Lock() delete(s.runt, id) s.mu.Unlock() } // InCooldown 渠道是否处于运行期冷却;过期则清掉状态、恢复可用。 // until 为零值表示只在累计失败、尚未触发冷却,此时不能清状态(否则计数永远攒不满阈值)。 func (s *Service) InCooldown(id uint64) bool { s.mu.Lock() defer s.mu.Unlock() st, ok := s.runt[id] if !ok { return false } if st.until.IsZero() { return false } if time.Now().Before(st.until) { return true } delete(s.runt, id) return false }