M5: 渠道体系(负载均衡+并发控制+健康检查+故障转移)
- channel.Candidates 按模型绑定取候选 + Pick 加权随机负载均衡 - TryAcquire 每渠道并发信号量, 满载溢出到其他候选 - HealthMonitor 后台定时探测, 连续失败进 cooldown, 恢复放回(可配 interval/threshold) - doProxy 遍历候选故障转移: 网络错误/429/5xx/超时且未写出响应头时安全重试; 400 等业务错误透传, 流式写出首字节后放弃重试 - 单测覆盖候选过滤/加权/并发/健康状态机 - E2E: 杀上游自动切换、cooldown、恢复、并发溢出 10/10 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Sakurasan
Claude
@@ -1,10 +1,12 @@
|
||||
// Package channel 渠道仓储：选择、密钥加解密、模型解析。
|
||||
// M1 实现最小选择逻辑（优先级+权重取第一个健康启用的渠道）；
|
||||
// 负载均衡/健康检查/故障转移在 M5 完善。
|
||||
// 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"
|
||||
@@ -16,47 +18,120 @@ 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}
|
||||
return &Service{db: db, enc: enc, sems: map[uint64]chan struct{}{}}
|
||||
}
|
||||
|
||||
// Select 选择处理请求的渠道：启用 + 健康，按 priority 升序、weight 降序。
|
||||
func (s *Service) Select() (*store.Channel, error) {
|
||||
// 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
|
||||
if err := s.db.Where("enabled = ? AND health_status = ?", true, store.ChannelHealthHealthy).
|
||||
Order("priority ASC, weight DESC, id ASC").Find(&chs).Error; err != nil {
|
||||
return nil, err
|
||||
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])
|
||||
}
|
||||
if len(chs) == 0 {
|
||||
return nil, ErrNoChannel
|
||||
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
|
||||
}
|
||||
return &chs[0], nil
|
||||
}
|
||||
|
||||
// UpstreamKey 解密渠道上游密钥。
|
||||
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
|
||||
return s.enc.Decrypt(ch.APIKeyEnc)
|
||||
}
|
||||
|
||||
// ResolveModel 按全局模型名找到绑定渠道；M1 简化：返回绑定该模型的第一个健康渠道。
|
||||
func (s *Service) ResolveModel(modelName string) (*store.Channel, *store.ChannelModelBinding, error) {
|
||||
var m store.Model
|
||||
if err := s.db.Where("name = ? AND enabled = ?", modelName, true).First(&m).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var b store.ChannelModelBinding
|
||||
if err := s.db.Where("model_id = ?", m.ID).
|
||||
Order("weight DESC, id ASC").First(&b).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var ch store.Channel
|
||||
if err := s.db.First(&ch, b.ChannelID).Error; err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if !ch.Enabled || ch.HealthStatus != store.ChannelHealthHealthy {
|
||||
return nil, nil, ErrNoChannel
|
||||
}
|
||||
return &ch, &b, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
package channel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/pkg/crypto"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func newTestDB(t *testing.T) *gorm.DB {
|
||||
t.Helper()
|
||||
dsn := "/tmp/channel_test_" + time.Now().Format("150405.000000000") + ".db"
|
||||
db, err := store.Open("sqlite", dsn)
|
||||
if err != nil {
|
||||
t.Fatalf("open db: %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
sqlDB, _ := db.DB()
|
||||
sqlDB.Close()
|
||||
os.Remove(dsn)
|
||||
})
|
||||
return db
|
||||
}
|
||||
|
||||
func seedChannel(t *testing.T, db *gorm.DB, name string, weight, priority, maxConc int) uint64 {
|
||||
t.Helper()
|
||||
enc := crypto.NewEncryptor("test-master-key-1234567890")
|
||||
key, _ := enc.Encrypt("upstream-key")
|
||||
ch := store.Channel{
|
||||
Name: name, Provider: store.ChannelProviderOpenAI, BaseURL: "http://localhost:9000",
|
||||
APIKeyEnc: key, Weight: weight, Priority: priority, MaxConcurrency: maxConc,
|
||||
HealthStatus: store.ChannelHealthHealthy, Enabled: true,
|
||||
}
|
||||
if err := db.Create(&ch).Error; err != nil {
|
||||
t.Fatalf("create channel: %v", err)
|
||||
}
|
||||
return ch.ID
|
||||
}
|
||||
|
||||
func TestCandidatesFiltersUnhealthy(t *testing.T) {
|
||||
db := newTestDB(t)
|
||||
seedChannel(t, db, "a", 1, 0, 4)
|
||||
seedChannel(t, db, "b", 5, 1, 4)
|
||||
// 禁用渠道不入候选
|
||||
bad := store.Channel{Name: "bad", Provider: store.ChannelProviderOpenAI, BaseURL: "http://x", APIKeyEnc: "x", HealthStatus: store.ChannelHealthCooldown, Enabled: true}
|
||||
db.Create(&bad)
|
||||
|
||||
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
||||
cands := s.Candidates("")
|
||||
if len(cands) != 2 {
|
||||
t.Fatalf("candidates = %d, want 2", len(cands))
|
||||
}
|
||||
if cands[0].Name != "a" {
|
||||
t.Fatalf("first by priority should be a, got %s", cands[0].Name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPickWeighted(t *testing.T) {
|
||||
db := newTestDB(t)
|
||||
seedChannel(t, db, "only", 1, 0, 4)
|
||||
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
||||
cands := s.Candidates("")
|
||||
chosen := s.Pick(cands)
|
||||
if chosen == nil || chosen.Name != "only" {
|
||||
t.Fatalf("pick single: %+v", chosen)
|
||||
}
|
||||
// 空候选返回 nil
|
||||
if p := s.Pick(nil); p != nil {
|
||||
t.Fatal("expected nil for empty candidates")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryAcquire(t *testing.T) {
|
||||
db := newTestDB(t)
|
||||
id := seedChannel(t, db, "c", 1, 0, 1)
|
||||
var ch store.Channel
|
||||
db.First(&ch, id)
|
||||
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
|
||||
release, ok := s.TryAcquire(&ch)
|
||||
if !ok {
|
||||
t.Fatal("first acquire should succeed")
|
||||
}
|
||||
if _, ok := s.TryAcquire(&ch); ok {
|
||||
t.Fatal("second acquire should fail (capacity 1)")
|
||||
}
|
||||
release()
|
||||
if _, ok := s.TryAcquire(&ch); !ok {
|
||||
t.Fatal("acquire after release should succeed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthCooldownAndRecovery(t *testing.T) {
|
||||
db := newTestDB(t)
|
||||
id := seedChannel(t, db, "h", 1, 0, 4)
|
||||
var ch store.Channel
|
||||
db.First(&ch, id)
|
||||
|
||||
m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{
|
||||
Interval: time.Minute, FailThreshold: 2, Timeout: time.Second,
|
||||
})
|
||||
|
||||
// 两次失败 → cooldown
|
||||
m.record(&ch, false)
|
||||
if ch.HealthStatus != store.ChannelHealthHealthy {
|
||||
t.Fatal("one failure should stay healthy")
|
||||
}
|
||||
m.record(&ch, false)
|
||||
if ch.HealthStatus != store.ChannelHealthCooldown {
|
||||
t.Fatalf("two failures should cooldown, got %s", ch.HealthStatus)
|
||||
}
|
||||
// 恢复
|
||||
m.record(&ch, true)
|
||||
if ch.HealthStatus != store.ChannelHealthHealthy {
|
||||
t.Fatalf("success should recover, got %s", ch.HealthStatus)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthStartStops(t *testing.T) {
|
||||
db := newTestDB(t)
|
||||
seedChannel(t, db, "probe", 1, 0, 4)
|
||||
m := NewHealthMonitor(db, crypto.NewEncryptor("test-master-key-1234567890"), HealthConfig{
|
||||
Interval: time.Hour, FailThreshold: 2, Timeout: time.Second,
|
||||
})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
m.Start(ctx)
|
||||
cancel() // 不应 panic
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package channel
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/openteam/server/internal/pkg/crypto"
|
||||
"github.com/openteam/server/internal/store"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// HealthConfig 健康检查参数。
|
||||
type HealthConfig struct {
|
||||
Interval time.Duration // 探测周期
|
||||
FailThreshold int // 连续失败 N 次进入 cooldown
|
||||
Timeout time.Duration // 单次探测超时
|
||||
}
|
||||
|
||||
// HealthMonitor 定时探测渠道健康状态。
|
||||
type HealthMonitor struct {
|
||||
db *gorm.DB
|
||||
enc *crypto.Encryptor
|
||||
cfg HealthConfig
|
||||
hc *http.Client
|
||||
mu sync.Mutex
|
||||
fail map[uint64]int // 渠道连续失败次数
|
||||
}
|
||||
|
||||
func NewHealthMonitor(db *gorm.DB, enc *crypto.Encryptor, cfg HealthConfig) *HealthMonitor {
|
||||
if cfg.Interval <= 0 {
|
||||
cfg.Interval = 60 * time.Second
|
||||
}
|
||||
if cfg.FailThreshold <= 0 {
|
||||
cfg.FailThreshold = 2
|
||||
}
|
||||
if cfg.Timeout <= 0 {
|
||||
cfg.Timeout = 5 * time.Second
|
||||
}
|
||||
return &HealthMonitor{
|
||||
db: db, enc: enc, cfg: cfg,
|
||||
hc: &http.Client{Timeout: cfg.Timeout},
|
||||
fail: map[uint64]int{},
|
||||
}
|
||||
}
|
||||
|
||||
// Start 启动后台探测循环（ctx 取消即停止）。
|
||||
func (h *HealthMonitor) Start(ctx context.Context) {
|
||||
go func() {
|
||||
ticker := time.NewTicker(h.cfg.Interval)
|
||||
defer ticker.Stop()
|
||||
h.probeAll()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
h.probeAll()
|
||||
}
|
||||
}
|
||||
}()
|
||||
log.Printf("health: monitor started (interval=%s, threshold=%d)", h.cfg.Interval, h.cfg.FailThreshold)
|
||||
}
|
||||
|
||||
// probeAll 探测所有启用渠道并更新健康状态。
|
||||
func (h *HealthMonitor) probeAll() {
|
||||
var chs []store.Channel
|
||||
if err := h.db.Where("enabled = ?", true).Find(&chs).Error; err != nil {
|
||||
return
|
||||
}
|
||||
for i := range chs {
|
||||
h.probe(&chs[i])
|
||||
}
|
||||
}
|
||||
|
||||
// probe 单渠道探测：GET {base}/v1/models。
|
||||
func (h *HealthMonitor) probe(ch *store.Channel) {
|
||||
key, err := h.enc.Decrypt(ch.APIKeyEnc)
|
||||
if err != nil {
|
||||
h.record(ch, false)
|
||||
return
|
||||
}
|
||||
url := strings.TrimRight(ch.BaseURL, "/") + "/v1/models"
|
||||
req, _ := http.NewRequest(http.MethodGet, url, nil)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := h.hc.Do(req)
|
||||
ok := err == nil && resp.StatusCode >= 200 && resp.StatusCode < 300
|
||||
if resp != nil {
|
||||
resp.Body.Close()
|
||||
}
|
||||
h.record(ch, ok)
|
||||
}
|
||||
|
||||
// record 记录一次探测结果，按阈值切换健康状态。
|
||||
func (h *HealthMonitor) record(ch *store.Channel, ok bool) {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
|
||||
prev := ch.HealthStatus
|
||||
if ok {
|
||||
h.fail[ch.ID] = 0
|
||||
if prev != store.ChannelHealthHealthy {
|
||||
h.setStatus(ch, store.ChannelHealthHealthy)
|
||||
log.Printf("health: channel %q recovered -> healthy", ch.Name)
|
||||
}
|
||||
return
|
||||
}
|
||||
h.fail[ch.ID]++
|
||||
if h.fail[ch.ID] >= h.cfg.FailThreshold && prev != store.ChannelHealthCooldown {
|
||||
h.setStatus(ch, store.ChannelHealthCooldown)
|
||||
log.Printf("health: channel %q -> cooldown (%d consecutive failures)", ch.Name, h.fail[ch.ID])
|
||||
}
|
||||
}
|
||||
|
||||
func (h *HealthMonitor) setStatus(ch *store.Channel, status string) {
|
||||
h.db.Model(&store.Channel{}).Where("id = ?", ch.ID).Update("health_status", status)
|
||||
ch.HealthStatus = status
|
||||
}
|
||||
Reference in New Issue
Block a user