vibeproxy/internal/loadbalancer/loadbalancer.go

// Package loadbalancer picks a healthy backend for a proxy host using
// round-robin, weighted, or least-connections strategies.
package loadbalancer

import (
	"errors"
	"sync/atomic"

	"github.com/owen/vibeproxy/internal/models"
)

var ErrNoHealthyBackend = errors.New("no healthy backend available")

// Backend wraps a models.Backend with runtime-only state shared between the
// health checker and the load balancer.
type Backend struct {
	Model       models.Backend
	healthy     atomic.Bool
	activeConns atomic.Int64
	failCount   atomic.Int32
	okCount     atomic.Int32
}

func NewBackend(m models.Backend) *Backend {
	b := &Backend{Model: m}
	b.healthy.Store(m.HealthCheckType == models.HealthCheckNone || m.HealthCheckType == "")
	return b
}

func (b *Backend) Healthy() bool      { return b.healthy.Load() && b.Model.Enabled }
func (b *Backend) ActiveConns() int64 { return b.activeConns.Load() }
func (b *Backend) Acquire()           { b.activeConns.Add(1) }
func (b *Backend) Release()           { b.activeConns.Add(-1) }
func (b *Backend) Addr() string       { return b.Model.Host }

// RecordCheck updates the health state given a single check's outcome,
// applying the configured failure/recovery thresholds.
func (b *Backend) RecordCheck(ok bool) (changed bool, nowHealthy bool) {
	wasHealthy := b.healthy.Load()
	if ok {
		b.failCount.Store(0)
		n := b.okCount.Add(1)
		threshold := int32(b.Model.RecoveryThreshold)
		if threshold < 1 {
			threshold = 1
		}
		if !wasHealthy && n >= threshold {
			b.healthy.Store(true)
		}
	} else {
		b.okCount.Store(0)
		n := b.failCount.Add(1)
		threshold := int32(b.Model.FailureThreshold)
		if threshold < 1 {
			threshold = 1
		}
		if wasHealthy && n >= threshold {
			b.healthy.Store(false)
		}
	}
	return wasHealthy != b.healthy.Load(), b.healthy.Load()
}

// Pool is the set of backends for one proxy host plus a pick strategy.
type Pool struct {
	Method   string // round_robin | least_conn | weighted
	Backends []*Backend
	counter  atomic.Uint64
}

func NewPool(method string, backends []*Backend) *Pool {
	return &Pool{Method: method, Backends: backends}
}

// Pick selects the next backend to use. Callers must call Release() on the
// returned backend (via a deferred call) once the request/connection ends.
func (p *Pool) Pick() (*Backend, error) {
	healthy := make([]*Backend, 0, len(p.Backends))
	for _, b := range p.Backends {
		if b.Healthy() {
			healthy = append(healthy, b)
		}
	}
	if len(healthy) == 0 {
		return nil, ErrNoHealthyBackend
	}

	switch p.Method {
	case "least_conn":
		best := healthy[0]
		for _, b := range healthy[1:] {
			if b.ActiveConns() < best.ActiveConns() {
				best = b
			}
		}
		return best, nil
	case "weighted":
		total := 0
		for _, b := range healthy {
			w := b.Model.Weight
			if w < 1 {
				w = 1
			}
			total += w
		}
		n := int(p.counter.Add(1)-1) % total
		for _, b := range healthy {
			w := b.Model.Weight
			if w < 1 {
				w = 1
			}
			if n < w {
				return b, nil
			}
			n -= w
		}
		return healthy[0], nil
	default: // round_robin
		idx := int(p.counter.Add(1)-1) % len(healthy)
		return healthy[idx], nil
	}
}