-
Notifications
You must be signed in to change notification settings - Fork 1
/
random_w.go
64 lines (56 loc) · 1.18 KB
/
random_w.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
// SPDX-FileCopyrightText: 2018-2024 caixw
//
// SPDX-License-Identifier: MIT
package selector
import (
"math/rand/v2"
"sync"
)
type weightedRandom struct {
peers []WeightedPeer
weights []int
sumOfWeight int
mux sync.RWMutex
}
func newWeightedRandom(cap int) Updateable {
return &weightedRandom{
weights: make([]int, 0, cap),
peers: make([]WeightedPeer, 0, cap),
}
}
func (s *weightedRandom) Next() (string, error) {
s.mux.RLock()
defer s.mux.RUnlock()
l := len(s.peers)
switch l {
case 0:
return "", ErrNoPeer()
case 1:
return s.peers[0].Addr(), nil
default:
weight := rand.IntN(s.sumOfWeight) + 1 // 排除 0
for i, p := range s.peers {
weight -= s.weights[i]
if weight <= 0 {
return p.Addr(), nil
}
}
return s.peers[len(s.peers)-1].Addr(), nil
}
}
func (s *weightedRandom) Update(peers ...Peer) {
s.mux.Lock()
defer s.mux.Unlock()
s.peers = s.peers[:0]
s.weights = s.weights[:0]
s.sumOfWeight = 0
for _, p := range peers {
wp, ok := p.(WeightedPeer)
if !ok {
panic("p 必须实现 WeightedPeer 接口")
}
s.sumOfWeight += wp.Weight()
s.peers = append(s.peers, wp)
s.weights = append(s.weights, wp.Weight())
}
}