-
Notifications
You must be signed in to change notification settings - Fork 4.6k
/
server_picker.go
89 lines (71 loc) · 1.41 KB
/
server_picker.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
package protocol
import (
"sync"
)
type ServerList struct {
sync.RWMutex
servers []*ServerSpec
}
func NewServerList() *ServerList {
return &ServerList{}
}
func (sl *ServerList) AddServer(server *ServerSpec) {
sl.Lock()
defer sl.Unlock()
sl.servers = append(sl.servers, server)
}
func (sl *ServerList) Size() uint32 {
sl.RLock()
defer sl.RUnlock()
return uint32(len(sl.servers))
}
func (sl *ServerList) GetServer(idx uint32) *ServerSpec {
sl.Lock()
defer sl.Unlock()
for {
if idx >= uint32(len(sl.servers)) {
return nil
}
server := sl.servers[idx]
if !server.IsValid() {
sl.removeServer(idx)
continue
}
return server
}
}
func (sl *ServerList) removeServer(idx uint32) {
n := len(sl.servers)
sl.servers[idx] = sl.servers[n-1]
sl.servers = sl.servers[:n-1]
}
type ServerPicker interface {
PickServer() *ServerSpec
}
type RoundRobinServerPicker struct {
sync.Mutex
serverlist *ServerList
nextIndex uint32
}
func NewRoundRobinServerPicker(serverlist *ServerList) *RoundRobinServerPicker {
return &RoundRobinServerPicker{
serverlist: serverlist,
nextIndex: 0,
}
}
func (p *RoundRobinServerPicker) PickServer() *ServerSpec {
p.Lock()
defer p.Unlock()
next := p.nextIndex
server := p.serverlist.GetServer(next)
if server == nil {
next = 0
server = p.serverlist.GetServer(0)
}
next++
if next >= p.serverlist.Size() {
next = 0
}
p.nextIndex = next
return server
}