package pool import ( "sync" "time" ) type accountStatus struct { configDir string activeRequests int lastUsed int64 rateLimitUntil int64 totalRequests int totalSuccess int totalErrors int totalRateLimits int totalLatencyMs int64 } type AccountStat struct { ConfigDir string ActiveRequests int TotalRequests int TotalSuccess int TotalErrors int TotalRateLimits int TotalLatencyMs int64 IsRateLimited bool RateLimitUntil int64 } type AccountPool struct { mu sync.Mutex accounts []*accountStatus } func NewAccountPool(configDirs []string) *AccountPool { accounts := make([]*accountStatus, 0, len(configDirs)) for _, dir := range configDirs { accounts = append(accounts, &accountStatus{configDir: dir}) } return &AccountPool{accounts: accounts} } func (p *AccountPool) GetNextConfigDir() string { p.mu.Lock() defer p.mu.Unlock() if len(p.accounts) == 0 { return "" } now := time.Now().UnixMilli() available := make([]*accountStatus, 0, len(p.accounts)) for _, a := range p.accounts { if a.rateLimitUntil < now { available = append(available, a) } } target := available if len(target) == 0 { target = make([]*accountStatus, len(p.accounts)) copy(target, p.accounts) // sort by earliest recovery for i := 1; i < len(target); i++ { for j := i; j > 0 && target[j].rateLimitUntil < target[j-1].rateLimitUntil; j-- { target[j], target[j-1] = target[j-1], target[j] } } } // pick least busy then least recently used best := target[0] for _, a := range target[1:] { if a.activeRequests < best.activeRequests { best = a } else if a.activeRequests == best.activeRequests && a.lastUsed < best.lastUsed { best = a } } best.lastUsed = now return best.configDir } func (p *AccountPool) find(configDir string) *accountStatus { for _, a := range p.accounts { if a.configDir == configDir { return a } } return nil } func (p *AccountPool) ReportRequestStart(configDir string) { if configDir == "" { return } p.mu.Lock() defer p.mu.Unlock() if a := p.find(configDir); a != nil { a.activeRequests++ a.totalRequests++ } } func (p *AccountPool) ReportRequestEnd(configDir string) { if configDir == "" { return } p.mu.Lock() defer p.mu.Unlock() if a := p.find(configDir); a != nil && a.activeRequests > 0 { a.activeRequests-- } } func (p *AccountPool) ReportRequestSuccess(configDir string, latencyMs int64) { if configDir == "" { return } p.mu.Lock() defer p.mu.Unlock() if a := p.find(configDir); a != nil { a.totalSuccess++ a.totalLatencyMs += latencyMs } } func (p *AccountPool) ReportRequestError(configDir string, latencyMs int64) { if configDir == "" { return } p.mu.Lock() defer p.mu.Unlock() if a := p.find(configDir); a != nil { a.totalErrors++ a.totalLatencyMs += latencyMs } } func (p *AccountPool) ReportRateLimit(configDir string, penaltyMs int64) { if configDir == "" { return } if penaltyMs <= 0 { penaltyMs = 60000 } p.mu.Lock() defer p.mu.Unlock() if a := p.find(configDir); a != nil { a.rateLimitUntil = time.Now().UnixMilli() + penaltyMs a.totalRateLimits++ } } func (p *AccountPool) GetStats() []AccountStat { p.mu.Lock() defer p.mu.Unlock() now := time.Now().UnixMilli() stats := make([]AccountStat, len(p.accounts)) for i, a := range p.accounts { stats[i] = AccountStat{ ConfigDir: a.configDir, ActiveRequests: a.activeRequests, TotalRequests: a.totalRequests, TotalSuccess: a.totalSuccess, TotalErrors: a.totalErrors, TotalRateLimits: a.totalRateLimits, TotalLatencyMs: a.totalLatencyMs, IsRateLimited: a.rateLimitUntil > now, RateLimitUntil: a.rateLimitUntil, } } return stats } func (p *AccountPool) Count() int { return len(p.accounts) } // ─── PoolHandle interface ────────────────────────────────────────────────── // PoolHandle 讓 handler 可以注入獨立的 pool 實例,避免多 port 模式共用全域 pool。 type PoolHandle interface { GetNextConfigDir() string ReportRequestStart(configDir string) ReportRequestEnd(configDir string) ReportRequestSuccess(configDir string, latencyMs int64) ReportRequestError(configDir string, latencyMs int64) ReportRateLimit(configDir string, penaltyMs int64) GetStats() []AccountStat } // GlobalPoolHandle 包裝全域函式以實作 PoolHandle 介面(單 port 模式使用) type GlobalPoolHandle struct{} func (GlobalPoolHandle) GetNextConfigDir() string { return GetNextAccountConfigDir() } func (GlobalPoolHandle) ReportRequestStart(d string) { ReportRequestStart(d) } func (GlobalPoolHandle) ReportRequestEnd(d string) { ReportRequestEnd(d) } func (GlobalPoolHandle) ReportRequestSuccess(d string, l int64) { ReportRequestSuccess(d, l) } func (GlobalPoolHandle) ReportRequestError(d string, l int64) { ReportRequestError(d, l) } func (GlobalPoolHandle) ReportRateLimit(d string, p int64) { ReportRateLimit(d, p) } func (GlobalPoolHandle) GetStats() []AccountStat { return GetAccountStats() } // ─── Global pool ─────────────────────────────────────────────────────────── var ( globalPool *AccountPool globalMu sync.Mutex ) func InitAccountPool(configDirs []string) { globalMu.Lock() defer globalMu.Unlock() globalPool = NewAccountPool(configDirs) } func GetNextAccountConfigDir() string { globalMu.Lock() p := globalPool globalMu.Unlock() if p == nil { return "" } return p.GetNextConfigDir() } func ReportRequestStart(configDir string) { globalMu.Lock() p := globalPool globalMu.Unlock() if p != nil { p.ReportRequestStart(configDir) } } func ReportRequestEnd(configDir string) { globalMu.Lock() p := globalPool globalMu.Unlock() if p != nil { p.ReportRequestEnd(configDir) } } func ReportRequestSuccess(configDir string, latencyMs int64) { globalMu.Lock() p := globalPool globalMu.Unlock() if p != nil { p.ReportRequestSuccess(configDir, latencyMs) } } func ReportRequestError(configDir string, latencyMs int64) { globalMu.Lock() p := globalPool globalMu.Unlock() if p != nil { p.ReportRequestError(configDir, latencyMs) } } func ReportRateLimit(configDir string, penaltyMs int64) { globalMu.Lock() p := globalPool globalMu.Unlock() if p != nil { p.ReportRateLimit(configDir, penaltyMs) } } func GetAccountStats() []AccountStat { globalMu.Lock() p := globalPool globalMu.Unlock() if p == nil { return nil } return p.GetStats() }