Initial vibeproxy implementation

Self-contained HTTP/HTTPS/TCP/UDP reverse proxy manager written in Go, with SQLite storage, a Bootstrap 5 admin UI (Go templates, no JS framework), and a REST API under /api/v1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QTXFrNP1imCW7jtfL5SZrw
2026-08-17 10:28:59 -05:00
commit ba1f7c4c03
105 changed files with 9142 additions and 0 deletions
@@ -0,0 +1,122 @@
+// 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
+	}
+}
@@ -0,0 +1,91 @@
+package loadbalancer
+
+import (
+	"testing"
+
+	"github.com/owen/vibeproxy/internal/models"
+)
+
+func TestRoundRobin(t *testing.T) {
+	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckNone})
+	b2 := NewBackend(models.Backend{Host: "b", Enabled: true, HealthCheckType: models.HealthCheckNone})
+	pool := NewPool("round_robin", []*Backend{b1, b2})
+
+	seen := map[string]int{}
+	for i := 0; i < 10; i++ {
+		b, err := pool.Pick()
+		if err != nil {
+			t.Fatal(err)
+		}
+		seen[b.Model.Host]++
+	}
+	if seen["a"] != 5 || seen["b"] != 5 {
+		t.Fatalf("expected even split, got %v", seen)
+	}
+}
+
+func TestLeastConn(t *testing.T) {
+	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckNone})
+	b2 := NewBackend(models.Backend{Host: "b", Enabled: true, HealthCheckType: models.HealthCheckNone})
+	pool := NewPool("least_conn", []*Backend{b1, b2})
+
+	b1.Acquire()
+	b1.Acquire()
+	picked, err := pool.Pick()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if picked.Model.Host != "b" {
+		t.Fatalf("expected b (fewer active conns), got %s", picked.Model.Host)
+	}
+}
+
+func TestWeighted(t *testing.T) {
+	b1 := NewBackend(models.Backend{Host: "a", Weight: 3, Enabled: true, HealthCheckType: models.HealthCheckNone})
+	b2 := NewBackend(models.Backend{Host: "b", Weight: 1, Enabled: true, HealthCheckType: models.HealthCheckNone})
+	pool := NewPool("weighted", []*Backend{b1, b2})
+
+	counts := map[string]int{}
+	for i := 0; i < 8; i++ {
+		b, err := pool.Pick()
+		if err != nil {
+			t.Fatal(err)
+		}
+		counts[b.Model.Host]++
+	}
+	if counts["a"] != 6 || counts["b"] != 2 {
+		t.Fatalf("expected 3:1 weighted split (6:2 of 8), got %v", counts)
+	}
+}
+
+func TestNoHealthyBackend(t *testing.T) {
+	b1 := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckTCP, FailureThreshold: 1})
+	pool := NewPool("round_robin", []*Backend{b1})
+	if _, err := pool.Pick(); err != ErrNoHealthyBackend {
+		t.Fatalf("expected ErrNoHealthyBackend, got %v", err)
+	}
+}
+
+func TestRecordCheckThresholds(t *testing.T) {
+	b := NewBackend(models.Backend{Host: "a", Enabled: true, HealthCheckType: models.HealthCheckTCP, FailureThreshold: 2, RecoveryThreshold: 2})
+	if b.Healthy() {
+		t.Fatal("expected backend to start unhealthy pending checks")
+	}
+	b.RecordCheck(true)
+	if b.Healthy() {
+		t.Fatal("should still be unhealthy after 1 of 2 required successes")
+	}
+	changed, healthy := b.RecordCheck(true)
+	if !changed || !healthy {
+		t.Fatal("expected backend to become healthy after 2 successes")
+	}
+
+	b.RecordCheck(false)
+	if !b.Healthy() {
+		t.Fatal("should still be healthy after 1 of 2 required failures")
+	}
+	changed, healthy = b.RecordCheck(false)
+	if !changed || healthy {
+		t.Fatal("expected backend to become unhealthy after 2 failures")
+	}
+}