李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
280 lines
9.3 KiB
Go
280 lines
9.3 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
)
|
|
|
|
// cnQuotaProber 抽象额度探测（*CNProviderQuotaService 实现，测试可替换）。
|
|
type cnQuotaProber interface {
|
|
QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error)
|
|
}
|
|
|
|
// cnQuotaProbeConcurrency 周期任务并发探测额度账号的并发度。
|
|
const cnQuotaProbeConcurrency = 4
|
|
|
|
// CNProviderBalanceCheckService 周期性探测国产供应商账号：
|
|
// - payg（按量付费）：余额低于阈值则临时停调，恢复则清除（仅清除本服务写入的停调）；
|
|
// - coding plan：调用 CNProviderQuotaService 探测 5h/weekly 滚动窗口并落 extra 快照，
|
|
// 调度阈值评估（cnProviderThresholdCandidates）据此自动停调/恢复。
|
|
//
|
|
// 克隆自 AccountExpiryService 的 Start/Stop/runOnce + ticker 骨架。
|
|
// 余额探测仅覆盖有公开余额端点的 kimi / deepseek；智谱无余额端点，仅靠响应式 429/402。
|
|
// 额度探测覆盖 kimi / zhipu 的 coding plan 账号（deepseek 无 coding 套餐）。
|
|
type CNProviderBalanceCheckService struct {
|
|
accountRepo AccountRepository
|
|
balanceService *CNProviderBalanceService
|
|
quotaService cnQuotaProber
|
|
cfg *config.Config
|
|
interval time.Duration
|
|
stopCh chan struct{}
|
|
stopOnce sync.Once
|
|
wg sync.WaitGroup
|
|
}
|
|
|
|
// NewCNProviderBalanceCheckService 构造周期余额/额度检测服务。
|
|
// interval <= 0 时 Start() 直接返回（不启动），便于通过配置关闭。
|
|
func NewCNProviderBalanceCheckService(
|
|
accountRepo AccountRepository,
|
|
balanceService *CNProviderBalanceService,
|
|
quotaService *CNProviderQuotaService,
|
|
cfg *config.Config,
|
|
interval time.Duration,
|
|
) *CNProviderBalanceCheckService {
|
|
return &CNProviderBalanceCheckService{
|
|
accountRepo: accountRepo,
|
|
balanceService: balanceService,
|
|
quotaService: quotaService,
|
|
cfg: cfg,
|
|
interval: interval,
|
|
stopCh: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
func (s *CNProviderBalanceCheckService) Start() {
|
|
if s == nil || s.accountRepo == nil || s.balanceService == nil || s.cfg == nil {
|
|
return
|
|
}
|
|
if !s.cfg.Gateway.CNProviders.BalanceCheckEnabled {
|
|
return
|
|
}
|
|
if s.interval <= 0 {
|
|
return
|
|
}
|
|
log.Printf("[CNBalance] started (interval=%s threshold=%.2f)", s.interval, s.cfg.Gateway.CNProviders.BalanceThreshold)
|
|
s.wg.Add(1)
|
|
go func() {
|
|
defer s.wg.Done()
|
|
ticker := time.NewTicker(s.interval)
|
|
defer ticker.Stop()
|
|
|
|
// 启动后先等待一个周期再首次探测，避免与进程启动峰重叠。
|
|
for {
|
|
select {
|
|
case <-ticker.C:
|
|
s.runOnce()
|
|
case <-s.stopCh:
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
}
|
|
|
|
func (s *CNProviderBalanceCheckService) Stop() {
|
|
if s == nil {
|
|
return
|
|
}
|
|
s.stopOnce.Do(func() {
|
|
close(s.stopCh)
|
|
})
|
|
s.wg.Wait()
|
|
}
|
|
|
|
func (s *CNProviderBalanceCheckService) runOnce() {
|
|
// 收集 coding 探测目标（kimi/deepseek + 智谱）与 payg 检查队列。
|
|
// coding 探测统一在收集完成后按 4 并发执行：单账号探测 15-20s，串行 ×
|
|
// 多账号会耗尽整体预算（120s 上限），排在后面的账号快照会饥饿，
|
|
// 连锁影响阈值停调的新鲜度判定。
|
|
type quotaTarget struct {
|
|
id int64
|
|
platform string
|
|
}
|
|
var quotaTargets []quotaTarget
|
|
var paygTargets []*Account
|
|
collect := func(platform string, accounts []Account) {
|
|
for i := range accounts {
|
|
account := &accounts[i]
|
|
if !account.IsActive() {
|
|
continue
|
|
}
|
|
// coding 账号：探测滚动窗口并落快照（不要求 Schedulable——已被
|
|
// 阈值停调的账号也需要新鲜快照决定是否续停）。
|
|
if account.IsCodingPlan() {
|
|
quotaTargets = append(quotaTargets, quotaTarget{id: account.ID, platform: account.Platform})
|
|
continue
|
|
}
|
|
// payg 余额探测仅 kimi/deepseek（智谱无公开余额端点，payg 账号
|
|
// 依赖响应式 402/429 处理）。
|
|
if platform != PlatformZhipu && account.Schedulable {
|
|
paygTargets = append(paygTargets, account)
|
|
}
|
|
}
|
|
}
|
|
for _, platform := range s.platforms() {
|
|
accounts, err := s.accountRepo.ListByPlatform(context.Background(), platform)
|
|
if err != nil {
|
|
log.Printf("[CNBalance] list %s accounts failed: %v", platform, err)
|
|
continue
|
|
}
|
|
collect(platform, accounts)
|
|
}
|
|
// 智谱无余额端点，仅进额度探测。
|
|
if s.quotaService != nil {
|
|
accounts, err := s.accountRepo.ListByPlatform(context.Background(), PlatformZhipu)
|
|
if err != nil {
|
|
log.Printf("[CNBalance] list %s accounts failed: %v", PlatformZhipu, err)
|
|
} else {
|
|
collect(PlatformZhipu, accounts)
|
|
}
|
|
}
|
|
|
|
// 预算按工作量放大：4 并发 × 15s/批 + payg 每账号 5s，下限 30s 上限 300s。
|
|
batches := (len(quotaTargets) + cnQuotaProbeConcurrency - 1) / cnQuotaProbeConcurrency
|
|
timeout := 30*time.Second + time.Duration(batches)*15*time.Second + time.Duration(len(paygTargets))*5*time.Second
|
|
if timeout > 300*time.Second {
|
|
timeout = 300 * time.Second
|
|
}
|
|
ctx, cancel := context.WithTimeout(context.Background(), timeout)
|
|
defer cancel()
|
|
|
|
threshold := s.cfg.Gateway.CNProviders.BalanceThreshold
|
|
paused, cleared := 0, 0
|
|
for _, account := range paygTargets {
|
|
switch s.checkOne(ctx, account, threshold) {
|
|
case cnBalancePaused:
|
|
paused++
|
|
case cnBalanceCleared:
|
|
cleared++
|
|
}
|
|
}
|
|
|
|
if len(quotaTargets) > 0 && s.quotaService != nil {
|
|
sem := make(chan struct{}, cnQuotaProbeConcurrency)
|
|
var wg sync.WaitGroup
|
|
for _, target := range quotaTargets {
|
|
wg.Add(1)
|
|
go func(t quotaTarget) {
|
|
defer wg.Done()
|
|
sem <- struct{}{}
|
|
defer func() { <-sem }()
|
|
s.probeQuota(ctx, t.id, t.platform)
|
|
}(target)
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
if paused > 0 || cleared > 0 {
|
|
log.Printf("[CNBalance] paused=%d cleared=%d (threshold=%.2f)", paused, cleared, threshold)
|
|
}
|
|
}
|
|
|
|
// probeQuota 探测单个 coding plan 账号的滚动窗口用量并落 extra 快照。
|
|
// 不在此处做停调/恢复决策：调度阈值评估读取快照统一判定（含暂停账号的续停）。
|
|
func (s *CNProviderBalanceCheckService) probeQuota(ctx context.Context, accountID int64, platform string) {
|
|
if s.quotaService == nil {
|
|
return
|
|
}
|
|
result, err := s.quotaService.QueryUsage(ctx, accountID)
|
|
if err != nil {
|
|
log.Printf("[CNBalance] quota probe account %d (%s) failed: %v", accountID, platform, err)
|
|
return
|
|
}
|
|
if result != nil && !result.Success && result.Error != "" {
|
|
log.Printf("[CNBalance] quota probe account %d (%s) error: %s", accountID, platform, result.Error)
|
|
}
|
|
}
|
|
|
|
type cnBalanceCheckOutcome int
|
|
|
|
const (
|
|
cnBalanceNoChange cnBalanceCheckOutcome = iota
|
|
cnBalancePaused
|
|
cnBalanceCleared
|
|
)
|
|
|
|
// checkOne 探测单账号余额并决定停调/恢复。探测失败时不动现状（避免瞬时网络抖动
|
|
// 误解除或误停调）。
|
|
func (s *CNProviderBalanceCheckService) checkOne(ctx context.Context, account *Account, threshold float64) cnBalanceCheckOutcome {
|
|
result, err := s.balanceService.QueryBalance(ctx, account.ID)
|
|
if err != nil || result == nil || !result.Success {
|
|
return cnBalanceNoChange
|
|
}
|
|
|
|
// 双币种（deepseek CNY+USD）任一币种余额达标即可继续调度；仅当全部低于
|
|
// 阈值（或不可用）才停调。
|
|
low := !result.Available || allCNBalancesBelowThreshold(result, threshold)
|
|
if low {
|
|
// 已被（任何来源）停调时不覆盖其 reason。
|
|
if !account.IsSchedulable() {
|
|
return cnBalanceNoChange
|
|
}
|
|
reason := cnBalanceLowReason(fmt.Sprintf("余额 %.4g %s 低于阈值 %.2f", result.Balance, result.Currency, threshold))
|
|
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, time.Now().Add(s.cooldown()), reason); err != nil {
|
|
log.Printf("[CNBalance] pause account %d failed: %v", account.ID, err)
|
|
return cnBalanceNoChange
|
|
}
|
|
log.Printf("[CNBalance] paused account %d (%s): balance=%.4g %s", account.ID, account.Platform, result.Balance, result.Currency)
|
|
return cnBalancePaused
|
|
}
|
|
|
|
// 余额健康：仅清除「本服务写入」的临时停调（reason 前缀匹配），不触碰其他子系统。
|
|
if account.TempUnschedulableUntil != nil && strings.HasPrefix(account.TempUnschedulableReason, cnBalanceLowReasonPrefix) {
|
|
if err := s.accountRepo.ClearTempUnschedulable(ctx, account.ID); err != nil {
|
|
log.Printf("[CNBalance] clear account %d failed: %v", account.ID, err)
|
|
return cnBalanceNoChange
|
|
}
|
|
log.Printf("[CNBalance] reactivated account %d (%s): balance=%.4g %s", account.ID, account.Platform, result.Balance, result.Currency)
|
|
return cnBalanceCleared
|
|
}
|
|
return cnBalanceNoChange
|
|
}
|
|
|
|
func (s *CNProviderBalanceCheckService) platforms() []string {
|
|
return []string{PlatformKimi, PlatformDeepseek}
|
|
}
|
|
|
|
// allCNBalancesBelowThreshold 判断全部币种余额是否均低于阈值。
|
|
// 无明细时退回主币种判定（与旧行为一致）。
|
|
func allCNBalancesBelowThreshold(result *CNProviderBalanceResult, threshold float64) bool {
|
|
if len(result.Balances) == 0 {
|
|
return result.Balance < threshold
|
|
}
|
|
for _, entry := range result.Balances {
|
|
if entry.Balance >= threshold {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
// cooldown 返回临时停调持续时长（= 2× 检测周期），与响应式 402/429 路径一致。
|
|
func (s *CNProviderBalanceCheckService) cooldown() time.Duration {
|
|
minutes := 10
|
|
if s.cfg != nil {
|
|
if cfgMin := s.cfg.Gateway.CNProviders.BalanceCheckIntervalMinutes; cfgMin > 0 {
|
|
minutes = cfgMin
|
|
}
|
|
}
|
|
cooldown := time.Duration(minutes) * time.Minute * 2
|
|
if cooldown < time.Minute {
|
|
cooldown = 10 * time.Minute
|
|
}
|
|
return cooldown
|
|
}
|