forked from caddyserver/caddy
/
policy.go
133 lines (116 loc) · 3.29 KB
/
policy.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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
package proxy
import (
"hash/fnv"
"math"
"math/rand"
"net"
"net/http"
"sync"
)
// HostPool is a collection of UpstreamHosts.
type HostPool []*UpstreamHost
// Policy decides how a host will be selected from a pool.
type Policy interface {
Select(pool HostPool, r *http.Request) *UpstreamHost
}
func init() {
RegisterPolicy("random", func() Policy { return &Random{} })
RegisterPolicy("least_conn", func() Policy { return &LeastConn{} })
RegisterPolicy("round_robin", func() Policy { return &RoundRobin{} })
RegisterPolicy("ip_hash", func() Policy { return &IPHash{} })
}
// Random is a policy that selects up hosts from a pool at random.
type Random struct{}
// Select selects an up host at random from the specified pool.
func (r *Random) Select(pool HostPool, request *http.Request) *UpstreamHost {
// Because the number of available hosts isn't known
// up front, the host is selected via reservoir sampling
// https://en.wikipedia.org/wiki/Reservoir_sampling
var randHost *UpstreamHost
count := 0
for _, host := range pool {
if !host.Available() {
continue
}
// (n % 1 == 0) holds for all n, therefore randHost
// will always get assigned a value if there is
// at least 1 available host
count++
if (rand.Int() % count) == 0 {
randHost = host
}
}
return randHost
}
// LeastConn is a policy that selects the host with the least connections.
type LeastConn struct{}
// Select selects the up host with the least number of connections in the
// pool. If more than one host has the same least number of connections,
// one of the hosts is chosen at random.
func (r *LeastConn) Select(pool HostPool, request *http.Request) *UpstreamHost {
var bestHost *UpstreamHost
count := 0
leastConn := int64(math.MaxInt64)
for _, host := range pool {
if !host.Available() {
continue
}
if host.Conns < leastConn {
leastConn = host.Conns
count = 0
}
// Among hosts with same least connections, perform a reservoir
// sample: https://en.wikipedia.org/wiki/Reservoir_sampling
if host.Conns == leastConn {
count++
if (rand.Int() % count) == 0 {
bestHost = host
}
}
}
return bestHost
}
// RoundRobin is a policy that selects hosts based on round robin ordering.
type RoundRobin struct {
robin uint32
mutex sync.Mutex
}
// Select selects an up host from the pool using a round robin ordering scheme.
func (r *RoundRobin) Select(pool HostPool, request *http.Request) *UpstreamHost {
poolLen := uint32(len(pool))
r.mutex.Lock()
defer r.mutex.Unlock()
// Return next available host
for i := uint32(0); i < poolLen; i++ {
r.robin++
host := pool[r.robin%poolLen]
if host.Available() {
return host
}
}
return nil
}
// IPHash is a policy that selects hosts based on hashing the request ip
type IPHash struct{}
func hash(s string) uint32 {
h := fnv.New32a()
h.Write([]byte(s))
return h.Sum32()
}
// Select selects an up host from the pool using a round robin ordering scheme.
func (r *IPHash) Select(pool HostPool, request *http.Request) *UpstreamHost {
poolLen := uint32(len(pool))
clientIP, _, err := net.SplitHostPort(request.RemoteAddr)
if err != nil {
clientIP = request.RemoteAddr
}
index := hash(clientIP) % poolLen
for i := uint32(0); i < poolLen; i++ {
index += i
host := pool[index%poolLen]
if host.Available() {
return host
}
}
return nil
}