-
Notifications
You must be signed in to change notification settings - Fork 0
/
round_robin.go
92 lines (75 loc) · 1.93 KB
/
round_robin.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
package loadbalance
import (
"sync"
"github.com/alphadose/haxmap"
)
// 轮询负载均衡
type RoundRobin[T Hashable, I Instance[T]] struct {
curIndex int
mutex sync.Mutex
instancesMap *haxmap.Map[T, I] //map[T]I
instances []I
}
func NewRoundRobin[T Hashable, I Instance[T]]() *RoundRobin[T, I] {
return &RoundRobin[T, I]{
instancesMap: haxmap.New[T, I](8),
instances: make([]I, 0, 8),
}
}
func (rr *RoundRobin[T, I]) Add(instances ...I) int {
rr.mutex.Lock()
defer rr.mutex.Unlock()
count := 0
for _, instance := range instances {
if _, ok := rr.instancesMap.Get(instance.InstanceID()); !ok {
rr.instancesMap.Set(instance.InstanceID(), instance)
rr.instances = append(rr.instances, instance)
count++
}
}
return count
}
func haxMapGetVal[K Hashable, V any](m *haxmap.Map[K, V], key K) (V, bool) {
return m.Get(key)
}
func haxMapForEach[K Hashable, V any](m *haxmap.Map[K, V], callback func(K, V) bool) {
m.ForEach(callback)
}
func (rr *RoundRobin[T, I]) Get(key T) (I, bool) {
return haxMapGetVal(rr.instancesMap, key)
}
func (rr *RoundRobin[T, I]) ForEach(callback func(T, I) bool) {
haxMapForEach(rr.instancesMap, callback)
}
func (rr *RoundRobin[T, I]) Select() (ins I) {
if rr.instancesMap.Len() == 0 {
return ins
}
rr.mutex.Lock()
defer rr.mutex.Unlock()
c := (rr.curIndex + 1) % len(rr.instances)
rr.curIndex = c
return rr.instances[c]
}
func (rr *RoundRobin[T, I]) Del(instances ...I) int {
rr.mutex.Lock()
defer rr.mutex.Unlock()
count := 0
for _, instance := range instances {
id := instance.InstanceID()
if _, ok := rr.instancesMap.Get(id); ok {
for i := 0; i < len(rr.instances); i++ {
if rr.instances[i].InstanceID() == id {
rr.instancesMap.Del(instance.InstanceID())
rr.instances = append(rr.instances[:i], rr.instances[i+1:]...)
break
}
}
count++
}
}
return count
}
func (rr *RoundRobin[T, I]) Size() int {
return int(rr.instancesMap.Len())
}