157 lines
4.6 KiB
Go
157 lines
4.6 KiB
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{}{}}
|
|
}
|
|
|
|
// Candidate 一个候选渠道 + 该模型的映射关系。
|
|
type Candidate struct {
|
|
Channel *store.Channel
|
|
UpstreamModel string // 全局模型在此渠道的映射名（无绑定则为空，用客户端模型名）
|
|
}
|
|
|
|
// Candidates 返回可用渠道候选：健康 + 启用，按优先级、权重降序、id 升序排列。
|
|
// model 非空时优先取绑定该模型的渠道（携带 upstream_model 映射）；无绑定则退回全局。
|
|
func (s *Service) Candidates(model string) []Candidate {
|
|
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)
|
|
if cands := s.loadBound(b); len(cands) > 0 {
|
|
return cands
|
|
}
|
|
}
|
|
}
|
|
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 {
|
|
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 {
|
|
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)
|
|
}
|