openteam/server/internal/channel/channel.go

// Package channel 渠道仓储：候选选择、加权负载均衡、并发控制、密钥加解密。
// M5 起支持多候选（负载均衡 + 故障转移）与每渠道并发信号量。
package channel

import (
	"context"
	"errors"
	"log"
	mrand "math/rand/v2"
	"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")

const (
	// maxCooldown 运行期冷却上限：上游 Retry-After 再长也不让渠道缺席超过这个时间。
	maxCooldown = 10 * time.Minute
	// quotaCooldown 余额/配额耗尽的一次性冷却时长：这类失败不会自己恢复，
	// 冷却到上限也没意义，半小时后渠道大概率已被管理员处理或配额已重置。
	quotaCooldown = 30 * time.Minute
)

// FailKind 上游失败原因分类。冷却策略随原因不同：
// 余额/配额类第一次失败就长冷却（重试无意义）；其余按指数退避。
type FailKind uint8

const (
	FailNet    FailKind = iota // 网络错误、超时、连接失败
	FailServer                 // 5xx / 529 服务端错误
	FailRate                   // 429 限流
	FailQuota                  // 余额/配额耗尽（402、带配额特征的 429/403）
)

func (k FailKind) String() string {
	switch k {
	case FailNet:
		return "network"
	case FailServer:
		return "server_error"
	case FailRate:
		return "rate_limited"
	case FailQuota:
		return "quota_exhausted"
	}
	return "unknown"
}

type Service struct {
	db  *gorm.DB
	enc *crypto.Encryptor

	mu   sync.Mutex
	sems map[uint64]*semEntry
	runt map[uint64]*runtimeHealth

	// 运行期冷却参数；NewService 给默认值，测试可覆盖。
	failThreshold int
	cooldownFor   time.Duration

	// 目录缓存：模型→候选渠道的静态部分（渠道行、绑定、映射）缓存 catalogTTL，
	// 运行期状态（冷却、健康）每次读取时重新判定。管理端写操作与健康状态变化
	// 通过 InvalidateCatalog 立即失效。
	catMu      sync.Mutex
	cat        map[string]catEntry
	catalogTTL time.Duration
}

// catEntry 一次目录查询的结果（未按运行期状态过滤）。
type catEntry struct {
	rows []catRow
	at   time.Time
	// bound：该模型有渠道绑定（只用绑定渠道）；否则为全局渠道池
	bound bool
}

type catRow struct {
	ch            store.Channel
	upstreamModel string
	weight        int
}

// runtimeHealth 运行期健康：连续可重试失败（网络错误/超时/429/5xx）累计到阈值进入冷却。
// trips 记录连续进冷却的次数，退避时长按它翻倍；一次成功即清零。
// 与 HealthMonitor 的探测结果相互独立——探测只看 /models 是否 200，
// 覆盖不了"models 正常但对话接口一直报错"的渠道。
type runtimeHealth struct {
	fails int
	until time.Time
	kind  FailKind
	trips int
}

func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service {
	return &Service{
		db: db, enc: enc,
		sems:          map[uint64]*semEntry{},
		runt:          map[uint64]*runtimeHealth{},
		failThreshold: 3,
		cooldownFor:   30 * time.Second,
		cat:           map[string]catEntry{},
		catalogTTL:    10 * time.Second,
	}
}

// Candidate 一个候选渠道 + 该模型的映射关系。
type Candidate struct {
	Channel       *store.Channel
	UpstreamModel string // 全局模型在此渠道的映射名（无绑定则为空，用客户端模型名）
	Weight        int    // 绑定级权重（默认 1；均衡权重 = 绑定级 × 渠道级）
}

// Candidates 返回可用渠道候选：健康 + 启用 + 不在运行期冷却，按优先级、权重降序、id 升序排列。
// model 非空且该模型有渠道绑定时，只返回绑定渠道（携带 upstream_model 映射）；
// 绑定渠道全部不可用则返回空（上层报 no_channel），不会退化到不认识该模型的全局渠道。
// 仅当模型完全没有任何绑定时才退回全局渠道。
//
// 目录查询（模型/绑定/渠道）每请求原本要打 3 条 SQL，这里缓存 catalogTTL；
// 冷却与健康状态每次读取时重新判定，所以坏渠道照样立刻摘除。
func (s *Service) Candidates(model string) []Candidate {
	rows, bound := s.catalog(model)
	out := make([]Candidate, 0, len(rows))
	for i := range rows {
		ch := rows[i].ch
		if !ch.Enabled || ch.HealthStatus != store.ChannelHealthHealthy {
			continue // 缓存行里的健康/启用状态可能落后，这里按读到的值再过滤
		}
		if s.InCooldown(ch.ID) {
			continue
		}
		cand := Candidate{Channel: &ch}
		if bound {
			cand.UpstreamModel = rows[i].upstreamModel
			cand.Weight = rows[i].weight
		}
		out = append(out, cand)
	}
	return out
}

// catalog 取该模型的目录行（缓存命中返回副本，避免调用方改写缓存内容）。
// bound 表示这些行来自「该模型的渠道绑定」而不是全局渠道池。
func (s *Service) catalog(model string) (rows []catRow, bound bool) {
	if s.catalogTTL > 0 {
		s.catMu.Lock()
		if e, ok := s.cat[model]; ok && time.Since(e.at) < s.catalogTTL {
			s.catMu.Unlock()
			return e.rows, e.bound
		}
		s.catMu.Unlock()
	}
	rows, bound = s.loadCatalog(model)
	if s.catalogTTL > 0 {
		s.catMu.Lock()
		if s.cat == nil {
			s.cat = map[string]catEntry{}
		}
		s.cat[model] = catEntry{rows: rows, at: time.Now(), bound: bound}
		s.catMu.Unlock()
	}
	return rows, bound
}

// loadCatalog 查询模型/绑定/渠道，组装目录行（不判运行期状态）。
func (s *Service) loadCatalog(model string) (rows []catRow, bound bool) {
	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.boundRows(b), true
		}
	}
	var chs []store.Channel
	s.db.Where("enabled = ? AND health_status = ?", true, store.ChannelHealthHealthy).
		Order("priority ASC, weight DESC, id ASC").Find(&chs)
	rows = make([]catRow, 0, len(chs))
	for i := range chs {
		rows = append(rows, catRow{ch: chs[i]})
	}
	return rows, false
}

// boundRows 按绑定组装目录行（保持 priority ASC 的故障转移顺序）。
func (s *Service) boundRows(bindings []store.ChannelModelBinding) []catRow {
	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)
	rows := make([]catRow, 0, len(chs))
	for i := range chs {
		b := byChannel[chs[i].ID]
		rows = append(rows, catRow{ch: chs[i], upstreamModel: b.UpstreamModel, weight: b.Weight})
	}
	return rows
}

// InvalidateCatalog 清空目录缓存：管理端改了渠道/绑定/模型，或渠道健康状态变化时调用。
func (s *Service) InvalidateCatalog() {
	s.catMu.Lock()
	s.cat = map[string]catEntry{}
	s.catMu.Unlock()
}

// 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 按权重加权随机选一个候选渠道（负载均衡）。
// 只在最高优先级层（priority 最小，候选已按 priority ASC 排序）内随机——
// 低优先级候选不参与首个请求的分配，仅作为故障转移备选。
// 权重 = 绑定级 Weight × 渠道级 Weight（见 pickWeight）。
func (s *Service) Pick(cands []Candidate) *store.Channel {
	if len(cands) == 0 {
		return nil
	}
	top := cands[0].Channel.Priority
	pool := cands
	for i, c := range cands {
		if c.Channel.Priority != top {
			pool = cands[:i]
			break
		}
	}
	total := 0
	for _, c := range pool {
		total += pickWeight(c)
	}
	// 负载均衡不需要密码学随机：math/rand/v2 的 IntN 无锁、零分配
	// （原来是 crypto/rand + big.Int，每请求一次系统调用加两次分配）。
	n := mrand.IntN(total)
	acc := 0
	for _, c := range pool {
		acc += pickWeight(c)
		if n < acc {
			return c.Channel
		}
	}
	return pool[len(pool)-1].Channel
}

// pickWeight 候选的均衡权重 = 绑定级权重 × 渠道级权重（各 ≤0 按 1）。
// 绑定级默认 1，未显式配置时等价于纯渠道权重；绑定上显式加权可对该模型单独放大。
func pickWeight(c Candidate) int {
	bw := c.Weight
	if bw <= 0 {
		bw = 1
	}
	cw := c.Channel.Weight
	if cw <= 0 {
		cw = 1
	}
	return bw * cw
}

// semEntry 渠道并发信号量。容量变更（管理后台改 MaxConcurrency）时在旧槽排空后
// 就地换新，改配置不再需要重启。release 闭包持有旧 chan：换新后旧持有者释放进旧
// chan 是无害的（新 chan 重新从 0 计数）。
type semEntry struct {
	sem chan struct{}
	cap int
}

// TryAcquire 尝试获取渠道并发槽；渠道满载返回 false（调用方可溢出到其他渠道）。
// MaxConcurrency<=0 视为不限制。
func (s *Service) TryAcquire(ch *store.Channel) (func(), bool) {
	if ch.MaxConcurrency <= 0 {
		return func() {}, true
	}
	s.mu.Lock()
	e, ok := s.sems[ch.ID]
	// 容量变了且旧槽已排空（len 即当前持有数）→ 换新；还有在途请求就先沿用旧槽
	if !ok {
		e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
		s.sems[ch.ID] = e
	} else if e.cap != ch.MaxConcurrency && len(e.sem) == 0 {
		e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
		s.sems[ch.ID] = e
	}
	sem := e.sem
	s.mu.Unlock()
	select {
	case sem <- struct{}{}:
		return func() { <-sem }, true
	default:
		return nil, false
	}
}

// Acquire 获取渠道并发槽；满载时最多等 wait（0=立刻返回 false，等价 TryAcquire）。
// ctx 结束（客户端断开/请求取消）即放弃等待。
//
// 满载时直接跳过会让超出并发上限的请求立刻失败：实测 16 并发、渠道上限 4 时
// 有 18–25% 的请求直接 502。LLM 请求本来就慢，短暂排队远好于把错误返回给 agent
// （一次 502 可能毁掉一整轮工具链）。
func (s *Service) Acquire(ctx context.Context, ch *store.Channel, wait time.Duration) (func(), bool) {
	if ch.MaxConcurrency <= 0 {
		return func() {}, true
	}
	s.mu.Lock()
	e, ok := s.sems[ch.ID]
	if !ok {
		e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
		s.sems[ch.ID] = e
	} else if e.cap != ch.MaxConcurrency && len(e.sem) == 0 {
		e = &semEntry{sem: make(chan struct{}, ch.MaxConcurrency), cap: ch.MaxConcurrency}
		s.sems[ch.ID] = e
	}
	sem := e.sem
	s.mu.Unlock()
	select {
	case sem <- struct{}{}:
		return func() { <-sem }, true
	default:
	}
	if wait <= 0 {
		return nil, false
	}
	t := time.NewTimer(wait)
	defer t.Stop()
	select {
	case sem <- struct{}{}:
		return func() { <-sem }, true
	case <-t.C:
		return nil, false
	case <-ctx.Done():
		return nil, false
	}
}

// UpstreamKey 解密渠道上游密钥。
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
	return s.enc.Decrypt(ch.APIKeyEnc)
}

// ---------------------------------------------------------------------------
// 运行期健康（进程内，与 /v1/models 探测独立）

// ReportFailure 记一次可重试的上游失败，按原因分类冷却：
//   - FailQuota（余额/配额耗尽）：渠道在充值或配额重置前不会恢复，第一次失败就冷却
//     quotaCooldown，不等阈值——429 不一定是限流（智谱 429 "余额不足" 也在此列）；
//   - 其余原因：连续 failThreshold 次进入冷却，时长为基准冷却按连续进冷却次数的
//     指数退避（30s→60s→120s…），坏渠道不会以固定节拍被反复打。
//
// retryAfter>0（上游 429 带的 Retry-After）且长于退避值时覆盖退避时长，仍夹到
// maxCooldown——不能让一个写错的重试头把渠道钉死几小时；比退避值短的 Retry-After
// 不缩短冷却，否则 1s 就把坏渠道放回轮转。
func (s *Service) ReportFailure(ch *store.Channel, kind FailKind, retryAfter time.Duration) {
	s.mu.Lock()
	st := s.runt[ch.ID]
	if st == nil {
		st = &runtimeHealth{}
		s.runt[ch.ID] = st
	}
	st.kind = kind
	if kind == FailQuota {
		st.fails = 0
		st.until = time.Now().Add(quotaCooldown)
		s.mu.Unlock()
		log.Printf("channel: %q cooled down for %s: %s", ch.Name, quotaCooldown, kind)
		return
	}
	d := s.cooldownFor << st.trips // 指数退避
	if d > maxCooldown || d <= 0 {
		d = maxCooldown
	}
	if retryAfter > d {
		d = retryAfter
		if d > maxCooldown {
			d = maxCooldown
		}
	}
	st.fails++
	if st.fails < s.failThreshold {
		s.mu.Unlock()
		return
	}
	st.fails = 0
	st.trips++
	st.until = time.Now().Add(d)
	s.mu.Unlock()
	log.Printf("channel: %q entered runtime cooldown (%s, %s)", ch.Name, d, kind)
}

// ReportSuccess 上游正常返回响应，清空该渠道的连续失败与冷却退避计数。
func (s *Service) ReportSuccess(ch *store.Channel) {
	s.mu.Lock()
	if st, ok := s.runt[ch.ID]; ok {
		st.fails = 0
		st.trips = 0
	}
	s.mu.Unlock()
}

// RuntimeState 渠道运行期健康快照：连续失败次数、冷却剩余时间与最近失败原因（管理后台展示用）。
// 只读，不做过期清理——清理仍由 InCooldown 在选路时惰性完成。
func (s *Service) RuntimeState(id uint64) (fails int, remaining time.Duration, reason string) {
	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, st.kind.String() // 冷却中：计数在写入时已归零
		}
		return 0, 0, st.kind.String() // 已到期，等下次选路清理
	}
	return st.fails, 0, st.kind.String()
}

// 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
}