ollama-proxy/internal/proxy/balancer.go

// Package proxy implements the Ollama Cloud reverse proxy: account selection
// (round-robin with per-account cooldown), the HTTP reverse-proxy handler,
// and the pre-stream failover logic.
package proxy

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/Atte149/ollama-proxy/internal/config"
)

// Balancer selects the next upstream account using round-robin and skips
// accounts currently in cooldown. It is safe for concurrent use.
type Balancer struct {
	accounts []config.Account
	rr       atomic.Uint64 // round-robin counter

	mu          sync.RWMutex
	cooldown    map[string]time.Time // account ID -> until
	last429     map[string]time.Time // account ID -> last 429 time (for display)
	cooldownDur time.Duration
}

// NewBalancer builds a Balancer from a list of accounts. The cooldown window
// is applied uniformly to all accounts. At least one account is required;
// otherwise Next returns ErrNoAccounts on every call.
func NewBalancer(accounts []config.Account, cooldown time.Duration) *Balancer {
	// copy to avoid external mutation; normalise empty provider to ollama-cloud.
	accts := make([]config.Account, len(accounts))
	for i, a := range accounts {
		if a.Provider == "" {
			a.Provider = config.ProviderOllamaCloud
		}
		accts[i] = a
	}
	return &Balancer{
		accounts:    accts,
		cooldown:    make(map[string]time.Time),
		last429:     make(map[string]time.Time),
		cooldownDur: cooldown,
	}
}

// Len returns the number of accounts the balancer knows about.
func (b *Balancer) Len() int { return len(b.accounts) }

// Account returns the i-th account (mostly for tests / display).
func (b *Balancer) Account(i int) config.Account { return b.accounts[i] }

// Accounts returns a shallow copy of the account list.
func (b *Balancer) Accounts() []config.Account {
	out := make([]config.Account, len(b.accounts))
	copy(out, b.accounts)
	return out
}

// ErrNoAccounts is returned when the balancer has no accounts configured.
type ErrNoAccounts struct{}

func (ErrNoAccounts) Error() string { return "no accounts configured" }

// ErrAllCooldown is returned when every account is currently in cooldown.
type ErrAllCooldown struct {
	// Until is the earliest time at which any account becomes available again.
	Until time.Time
}

func (e ErrAllCooldown) Error() string {
	return "all accounts are in cooldown until " + e.Until.Format(time.RFC3339)
}

// Next picks the next available account, skipping accounts in cooldown. It
// rotates starting from the round-robin counter so consecutive calls land on
// different accounts when possible. Returns ErrNoAccounts or ErrAllCooldown
// when nothing is available.
//
// Next considers all accounts regardless of provider. Use NextFor to restrict
// to a set of eligible providers (e.g. for model-based routing).
func (b *Balancer) Next() (config.Account, error) {
	return b.NextFor(nil)
}

// NextFor picks the next available account whose provider is in the eligible
// list. When eligible is nil or empty, all providers are eligible (matches
// the original Next behaviour). Accounts in cooldown are skipped.
func (b *Balancer) NextFor(eligible []config.ProviderType) (config.Account, error) {
	if len(b.accounts) == 0 {
		return config.Account{}, ErrNoAccounts{}
	}
	allowed := make(map[config.ProviderType]bool, len(eligible))
	for _, p := range eligible {
		allowed[p] = true
	}
	allAllowed := len(eligible) == 0

	now := time.Now()
	var earliest time.Time
	for i := 0; i < len(b.accounts); i++ {
		idx := int(b.rr.Add(1)) % len(b.accounts)
		acct := b.accounts[idx]
		// idx is computed from the *new* counter value; we Add-then-mod so each
		// call advances even if this loop iteration rejects the candidate.
		if !allAllowed && !allowed[acct.Provider] {
			continue
		}
		if !b.isCooldown(acct.ID, now) {
			return acct, nil
		}
		if earliest.IsZero() || b.cooldownUntil(acct.ID).Before(earliest) {
			earliest = b.cooldownUntil(acct.ID)
		}
	}
	return config.Account{}, ErrAllCooldown{Until: earliest}
}

// isCooldown reports whether the account is currently rate-limited.
func (b *Balancer) isCooldown(id string, now time.Time) bool {
	b.mu.RLock()
	defer b.mu.RUnlock()
	until, ok := b.cooldown[id]
	if !ok {
		return false
	}
	return now.Before(until)
}

// cooldownUntil returns the cooldown expiry for an account (zero if none).
func (b *Balancer) cooldownUntil(id string) time.Time {
	b.mu.RLock()
	defer b.mu.RUnlock()
	return b.cooldown[id]
}

// MarkCooldown puts the given account into cooldown for the configured window
// starting from now. Safe to call concurrently; idempotent (extends the window).
func (b *Balancer) MarkCooldown(id string) {
	b.MarkCooldownFor(id, b.cooldownDur)
}

// MarkCooldownFor puts the account into cooldown for an explicit duration.
func (b *Balancer) MarkCooldownFor(id string, d time.Duration) {
	b.mu.Lock()
	defer b.mu.Unlock()
	until := time.Now().Add(d)
	b.cooldown[id] = until
	b.last429[id] = time.Now()
}

// ClearCooldown removes the cooldown for an account (e.g. on a successful
// request after earlier failures).
func (b *Balancer) ClearCooldown(id string) {
	b.mu.Lock()
	defer b.mu.Unlock()
	delete(b.cooldown, id)
}

// Status returns a snapshot of per-account cooldown state, suitable for
// display in `accounts list`. The map key is the account ID.
func (b *Balancer) Status() map[string]AccountStatus {
	b.mu.RLock()
	defer b.mu.RUnlock()
	out := make(map[string]AccountStatus, len(b.accounts))
	now := time.Now()
	for _, a := range b.accounts {
		until := b.cooldown[a.ID]
		st := AccountStatus{
			InCooldown: !until.IsZero() && now.Before(until),
			Until:      until,
			Last429:    b.last429[a.ID],
		}
		out[a.ID] = st
	}
	return out
}

// AccountStatus is the per-account cooldown snapshot returned by Status.
type AccountStatus struct {
	InCooldown bool
	Until      time.Time
	Last429    time.Time
}