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