openteam/server/internal/channel/channel.go

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

import (
	"crypto/rand"
	"errors"
	"math/big"
	"sync"

	"github.com/openteam/server/internal/pkg/crypto"
	"github.com/openteam/server/internal/store"
	"gorm.io/gorm"
)

var ErrNoChannel = errors.New("no available channel")

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

	mu   sync.Mutex
	sems map[uint64]chan struct{}
}

func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service {
	return &Service{db: db, enc: enc, sems: map[uint64]chan struct{}{}}
}

// Candidates 返回可用渠道候选：健康 + 启用，按优先级、权重降序、id 升序排列。
// model 非空时优先取绑定该模型的渠道；无绑定则退回全局。
func (s *Service) Candidates(model string) []*store.Channel {
	if model != "" {
		var b []store.ChannelModelBinding
		var modelIDs []uint64
		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)
			chs := s.loadBound(b)
			if len(chs) > 0 {
				return chs
			}
		}
	}
	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([]*store.Channel, 0, len(chs))
	for i := range chs {
		out = append(out, &chs[i])
	}
	return out
}

// loadBound 按绑定顺序加载渠道，过滤健康/启用。
func (s *Service) loadBound(bindings []store.ChannelModelBinding) []*store.Channel {
	if len(bindings) == 0 {
		return nil
	}
	ids := make([]uint64, 0, len(bindings))
	seen := map[uint64]bool{}
	for _, b := range bindings {
		if !seen[b.ChannelID] {
			seen[b.ChannelID] = true
			ids = append(ids, b.ChannelID)
		}
	}
	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([]*store.Channel, 0, len(ids))
	for _, id := range ids {
		if ch, ok := byID[id]; ok {
			out = append(out, ch)
		}
	}
	return out
}

// Pick 按权重加权随机选一个候选（负载均衡）。
func (s *Service) Pick(cands []*store.Channel) *store.Channel {
	if len(cands) == 0 {
		return nil
	}
	total := 0
	for _, c := range cands {
		w := c.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.Weight
		if w <= 0 {
			w = 1
		}
		acc += w
		if int(n.Int64()) < acc {
			return c
		}
	}
	return cands[len(cands)-1]
}

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