forked from luraproject/lura
/
loadbalancing.go
106 lines (90 loc) · 2.51 KB
/
loadbalancing.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
// SPDX-License-Identifier: Apache-2.0
package sd
import (
"errors"
"runtime"
"sync/atomic"
"github.com/valyala/fastrand"
)
// Balancer applies a balancing stategy in order to select the backend host to be used
type Balancer interface {
Host() (string, error)
}
// ErrNoHosts is the error the balancer must return when there are 0 hosts ready
var ErrNoHosts = errors.New("no hosts available")
// NewBalancer returns the best perfomant balancer depending on the number of available processors.
// If GOMAXPROCS = 1, it returns a round robin LB due there is no contention over the atomic counter.
// If GOMAXPROCS > 1, it returns a pseudo random LB optimized for scaling over the number of CPUs.
func NewBalancer(subscriber Subscriber) Balancer {
if p := runtime.GOMAXPROCS(-1); p == 1 {
return NewRoundRobinLB(subscriber)
}
return NewRandomLB(subscriber)
}
// NewRoundRobinLB returns a new balancer using a round robin strategy and starting at a random
// position in the set of hosts.
func NewRoundRobinLB(subscriber Subscriber) Balancer {
s, ok := subscriber.(FixedSubscriber)
start := uint64(0)
if ok {
if l := len(s); l == 1 {
return nopBalancer(s[0])
} else if l > 1 {
start = uint64(fastrand.Uint32n(uint32(l)))
}
}
return &roundRobinLB{
balancer: balancer{subscriber: subscriber},
counter: start,
}
}
type roundRobinLB struct {
balancer
counter uint64
}
// Host implements the balancer interface
func (r *roundRobinLB) Host() (string, error) {
hosts, err := r.hosts()
if err != nil {
return "", err
}
offset := (atomic.AddUint64(&r.counter, 1) - 1) % uint64(len(hosts))
return hosts[offset], nil
}
// NewRandomLB returns a new balancer using a fastrand pseudorandom generator
func NewRandomLB(subscriber Subscriber) Balancer {
if s, ok := subscriber.(FixedSubscriber); ok && len(s) == 1 {
return nopBalancer(s[0])
}
return &randomLB{
balancer: balancer{subscriber: subscriber},
rand: fastrand.Uint32n,
}
}
type randomLB struct {
balancer
rand func(uint32) uint32
}
// Host implements the balancer interface
func (r *randomLB) Host() (string, error) {
hosts, err := r.hosts()
if err != nil {
return "", err
}
return hosts[int(r.rand(uint32(len(hosts))))], nil
}
type balancer struct {
subscriber Subscriber
}
func (b *balancer) hosts() ([]string, error) {
hs, err := b.subscriber.Hosts()
if err != nil {
return hs, err
}
if len(hs) <= 0 {
return hs, ErrNoHosts
}
return hs, nil
}
type nopBalancer string
func (b nopBalancer) Host() (string, error) { return string(b), nil }