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

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