proxy-pool/internal/domain/routing/strategy.go

package routing

import (
	"errors"
	"math/rand"
	"sync"
	"sync/atomic"
)

var ErrNoCandidate = errors.New("no eligible routing candidate")

type Candidate struct {
	Name     string
	Weight   int
	Active   int64
	Eligible bool
}

type Selector interface {
	Select([]Candidate) (Candidate, error)
}

type RandomSource interface {
	Intn(int) int
}

type randomSelector struct {
	source RandomSource
}

type globalRandomSource struct{}

func (globalRandomSource) Intn(n int) int {
	return rand.Intn(n)
}

type synchronizedRandomSource struct {
	mu     sync.Mutex
	source RandomSource
}

func (s *synchronizedRandomSource) Intn(n int) int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.source.Intn(n)
}

func NewRandom(source ...RandomSource) Selector {
	return &randomSelector{source: newSynchronizedRandomSource(source)}
}

func (s *randomSelector) Select(candidates []Candidate) (Candidate, error) {
	eligibleCount := 0
	for _, candidate := range candidates {
		if candidate.Eligible {
			eligibleCount++
		}
	}
	if eligibleCount == 0 {
		return Candidate{}, ErrNoCandidate
	}

	selected := s.source.Intn(eligibleCount)
	for _, candidate := range candidates {
		if !candidate.Eligible {
			continue
		}
		if selected == 0 {
			return candidate, nil
		}
		selected--
	}
	return Candidate{}, ErrNoCandidate
}

type roundRobinSelector struct {
	next atomic.Uint64
}

func NewRoundRobin() Selector {
	return &roundRobinSelector{}
}

func (s *roundRobinSelector) Select(candidates []Candidate) (Candidate, error) {
	if len(candidates) == 0 {
		return Candidate{}, ErrNoCandidate
	}

	for {
		next := s.next.Load()
		start := int(next % uint64(len(candidates)))
		for offset := range len(candidates) {
			selected := (start + offset) % len(candidates)
			if !candidates[selected].Eligible {
				continue
			}
			if s.next.CompareAndSwap(next, uint64(selected+1)) {
				return candidates[selected], nil
			}
			break
		}
		if allCandidatesIneligible(candidates) {
			return Candidate{}, ErrNoCandidate
		}
	}
}

func allCandidatesIneligible(candidates []Candidate) bool {
	for _, candidate := range candidates {
		if candidate.Eligible {
			return false
		}
	}
	return true
}

type weightedSelector struct {
	source RandomSource
}

func NewWeighted(source ...RandomSource) Selector {
	return &weightedSelector{source: newSynchronizedRandomSource(source)}
}

func (s *weightedSelector) Select(candidates []Candidate) (Candidate, error) {
	totalWeight := 0
	for _, candidate := range candidates {
		if candidate.Eligible && candidate.Weight > 0 {
			if candidate.Weight > int(^uint(0)>>1)-totalWeight {
				return Candidate{}, ErrNoCandidate
			}
			totalWeight += candidate.Weight
		}
	}
	if totalWeight == 0 {
		return Candidate{}, ErrNoCandidate
	}

	target := s.source.Intn(totalWeight)
	for _, candidate := range candidates {
		if !candidate.Eligible || candidate.Weight <= 0 {
			continue
		}
		if target < candidate.Weight {
			return candidate, nil
		}
		target -= candidate.Weight
	}
	return Candidate{}, ErrNoCandidate
}

type leastConnectionsSelector struct{}

func NewLeastConnections() Selector {
	return leastConnectionsSelector{}
}

func (leastConnectionsSelector) Select(candidates []Candidate) (Candidate, error) {
	var selected Candidate
	found := false
	for _, candidate := range candidates {
		if !candidate.Eligible {
			continue
		}
		if !found || candidate.Active < selected.Active {
			selected = candidate
			found = true
		}
	}
	if !found {
		return Candidate{}, ErrNoCandidate
	}
	return selected, nil
}

func newSynchronizedRandomSource(sources []RandomSource) RandomSource {
	if len(sources) == 0 || sources[0] == nil {
		return globalRandomSource{}
	}
	return &synchronizedRandomSource{source: sources[0]}
}