169 lines
3.5 KiB
Go
169 lines
3.5 KiB
Go
package routing
|
|
|
|
import (
|
|
"errors"
|
|
"math/rand"
|
|
"sync"
|
|
)
|
|
|
|
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 *synchronizedRandomSource
|
|
}
|
|
|
|
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 {
|
|
mu sync.Mutex
|
|
next int
|
|
}
|
|
|
|
func NewRoundRobin() Selector {
|
|
return &roundRobinSelector{}
|
|
}
|
|
|
|
func (s *roundRobinSelector) Select(candidates []Candidate) (Candidate, error) {
|
|
if len(candidates) == 0 {
|
|
return Candidate{}, ErrNoCandidate
|
|
}
|
|
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
start := s.next % len(candidates)
|
|
for offset := range len(candidates) {
|
|
selected := (start + offset) % len(candidates)
|
|
if !candidates[selected].Eligible {
|
|
continue
|
|
}
|
|
s.next = (selected + 1) % len(candidates)
|
|
return candidates[selected], nil
|
|
}
|
|
return Candidate{}, ErrNoCandidate
|
|
}
|
|
|
|
type weightedSelector struct {
|
|
source *synchronizedRandomSource
|
|
}
|
|
|
|
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) *synchronizedRandomSource {
|
|
var source RandomSource = globalRandomSource{}
|
|
if len(sources) > 0 && sources[0] != nil {
|
|
source = sources[0]
|
|
}
|
|
return &synchronizedRandomSource{source: source}
|
|
}
|