LFG

internal/proxy/loadbalancer.go

@@ -0,0 +1,217 @@
+package proxy
+
+import (
+	"math/rand"
+	"sync"
+	"time"
+
+	"github.com/iwasforcedtobehere/goRZ/internal/config"
+)
+
+// RoundRobinLoadBalancer implements round-robin load balancing
+type RoundRobinLoadBalancer struct {
+	targets []*config.TargetConfig
+	current int
+	mu      sync.Mutex
+}
+
+// NewRoundRobinLoadBalancer creates a new round-robin load balancer
+func NewRoundRobinLoadBalancer(targets []config.TargetConfig) *RoundRobinLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	return &RoundRobinLoadBalancer{
+		targets: t,
+		current: 0,
+	}
+}
+
+// NextTarget returns the next target using round-robin algorithm
+func (lb *RoundRobinLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets
+	healthyTargets := make([]*config.TargetConfig, 0)
+	for _, target := range lb.targets {
+		if target.Healthy {
+			healthyTargets = append(healthyTargets, target)
+		}
+	}
+
+	if len(healthyTargets) == 0 {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Get next target
+	target := healthyTargets[lb.current%len(healthyTargets)]
+	lb.current = (lb.current + 1) % len(healthyTargets)
+
+	return target, nil
+}
+
+// UpdateTargets updates the targets list
+func (lb *RoundRobinLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	lb.targets = t
+	lb.current = 0
+}
+
+// RandomLoadBalancer implements random load balancing
+type RandomLoadBalancer struct {
+	targets []*config.TargetConfig
+	rand    *rand.Rand
+	mu      sync.Mutex
+}
+
+// NewRandomLoadBalancer creates a new random load balancer
+func NewRandomLoadBalancer(targets []config.TargetConfig) *RandomLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	return &RandomLoadBalancer{
+		targets: t,
+		rand:    rand.New(rand.NewSource(time.Now().UnixNano())),
+	}
+}
+
+// NextTarget returns a random target
+func (lb *RandomLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets
+	healthyTargets := make([]*config.TargetConfig, 0)
+	for _, target := range lb.targets {
+		if target.Healthy {
+			healthyTargets = append(healthyTargets, target)
+		}
+	}
+
+	if len(healthyTargets) == 0 {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Get random target
+	index := lb.rand.Intn(len(healthyTargets))
+	return healthyTargets[index], nil
+}
+
+// UpdateTargets updates the targets list
+func (lb *RandomLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+	lb.targets = t
+}
+
+// LeastConnectionsLoadBalancer implements least connections load balancing
+type LeastConnectionsLoadBalancer struct {
+	targets      []*config.TargetConfig
+	connections map[string]int
+	mu          sync.Mutex
+}
+
+// NewLeastConnectionsLoadBalancer creates a new least connections load balancer
+func NewLeastConnectionsLoadBalancer(targets []config.TargetConfig) *LeastConnectionsLoadBalancer {
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+
+	connections := make(map[string]int)
+	for _, target := range t {
+		connections[target.Name] = 0
+	}
+
+	return &LeastConnectionsLoadBalancer{
+		targets:      t,
+		connections: connections,
+	}
+}
+
+// NextTarget returns the target with the least connections
+func (lb *LeastConnectionsLoadBalancer) NextTarget() (*config.TargetConfig, error) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	// Filter healthy targets and find the one with least connections
+	var selectedTarget *config.TargetConfig
+	minConnections := -1
+
+	for _, target := range lb.targets {
+		if target.Healthy {
+			connections := lb.connections[target.Name]
+			if minConnections == -1 || connections < minConnections {
+				minConnections = connections
+				selectedTarget = target
+			}
+		}
+	}
+
+	if selectedTarget == nil {
+		return nil, ErrNoHealthyTargets
+	}
+
+	// Increment connection count
+	lb.connections[selectedTarget.Name]++
+
+	return selectedTarget, nil
+}
+
+// ReleaseConnection decrements the connection count for a target
+func (lb *LeastConnectionsLoadBalancer) ReleaseConnection(targetName string) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	if count, exists := lb.connections[targetName]; exists && count > 0 {
+		lb.connections[targetName] = count - 1
+	}
+}
+
+// UpdateTargets updates the targets list
+func (lb *LeastConnectionsLoadBalancer) UpdateTargets(targets []config.TargetConfig) {
+	lb.mu.Lock()
+	defer lb.mu.Unlock()
+
+	t := make([]*config.TargetConfig, len(targets))
+	for i := range targets {
+		t[i] = &targets[i]
+	}
+
+	// Update connections map
+	connections := make(map[string]int)
+	for _, target := range t {
+		// Preserve existing connection count if target exists
+		if count, exists := lb.connections[target.Name]; exists {
+			connections[target.Name] = count
+		} else {
+			connections[target.Name] = 0
+		}
+	}
+
+	lb.targets = t
+	lb.connections = connections
+}
+
+// ErrNoHealthyTargets is returned when no healthy targets are available
+var ErrNoHealthyTargets = errorString("no healthy targets available")
+
+// errorString is a simple string-based error type
+type errorString string
+
+func (e errorString) Error() string {
+	return string(e)
+}