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