/
strategy_weighted_roundrobin.go
91 lines (84 loc) · 2.49 KB
/
strategy_weighted_roundrobin.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
package cslb
import (
"math/big"
"sync/atomic"
"unsafe"
)
type weightedRoundRobinStrategy struct {
*roundRobinStrategy
weightFunc func(Node) int
}
// NewWeightedRoundRobinStrategy creates a new strategy with weighted round-robin method.
// This strategy will store a slice of node order in memory as cache, the length of the slice will be the sum of the weight
// of all nodes divided by the greatest common divisor of all weights.
// For example, a nodes configuration like this:
// | Node | Weight |
// | ---- | ------ |
// | A | 10 |
// | B | 10 |
// | C | 15 |
// | D | 20 |
// The greatest common divisor of [10, 10, 15, 20] is 5, so equivalent weights are [2, 2, 3, 4], which means the length
// of cache slice will be 2 + 2 + 3 + 4 = 11. Actually, the cache slice will be [D, C, A, B, D, C, D, A, B, C, D].
//
// Be careful when nodes have large weights and co-prime with each other, cache size might be very large.
// For example, a nodes configuration like this:
// | Node | Weight |
// | ---- | ------ |
// | A | 100000 |
// | B | 100001 |
// | C | 100002 |
// will cause a cache slice with length 300003, which will consume 2.3MB memory.
func NewWeightedRoundRobinStrategy(weightFunc func(Node) int) *weightedRoundRobinStrategy {
return &weightedRoundRobinStrategy{
roundRobinStrategy: NewRoundRobinStrategy(),
weightFunc: weightFunc,
}
}
func (s *weightedRoundRobinStrategy) SetNodes(nodes []Node) {
internalNodes := &roundRobinInternalNodes{
nodes: nodes,
order: s.generateWeightedOrder(nodes),
}
atomic.StorePointer(&s.internalNodes, unsafe.Pointer(internalNodes))
}
func (s *weightedRoundRobinStrategy) generateWeightedOrder(nodes []Node) []int {
if len(nodes) < 1 {
return []int{}
}
weights := make([]int, len(nodes))
for i, node := range nodes {
weights[i] = s.weightFunc(node)
if weights[i] < 0 {
weights[i] = 0
}
}
gcd := big.NewInt(int64(weights[0]))
for i := 1; i < len(weights); i++ {
gcd = gcd.GCD(nil, nil, gcd, big.NewInt(int64(weights[i])))
}
sum := 0
gcdi := int(gcd.Int64())
for i := 0; i < len(weights); i++ {
weights[i] /= gcdi
sum += weights[i]
}
result := make([]int, 0, sum)
curr := make([]int, len(weights))
for i := 0; i < sum; i++ {
for j := 0; j < len(weights); j++ {
curr[j] += weights[j]
}
max := 0
maxIndex := 0
for j := 0; j < len(weights); j++ {
if curr[j] > max {
max = curr[j]
maxIndex = j
}
}
result = append(result, maxIndex)
curr[maxIndex] -= sum
}
return result
}