gorz/internal/proxy/loadbalancer.go

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