opencode-cursor-agent/internal/pool/pool.go

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