-
Notifications
You must be signed in to change notification settings - Fork 0
/
random.go
116 lines (95 loc) · 2.05 KB
/
random.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
package pool
import (
"errors"
"github.com/dp0h/srp/app/config"
"github.com/rs/zerolog/log"
"math/rand"
"sync"
"time"
)
// RandomWeightedPool defines pool of services
type RandomWeightedPool struct {
refresh time.Duration
timeout time.Duration
services []*service
lock sync.RWMutex
}
type service struct {
host string
healthCheck string
weight int
alive bool
}
// NewRandomWeightedPool new random weighted pool of services
func NewRandomWeightedPool(config *config.ConfFile, refresh time.Duration, timeout time.Duration) *RandomWeightedPool {
rand.Seed(int64(time.Now().Nanosecond()))
res := RandomWeightedPool{
services: configToServices(config),
refresh: refresh,
timeout: timeout}
go res.healthCheck()
return &res
}
func configToServices(config *config.ConfFile) []*service {
var res []*service
for _, v := range config.Services {
svc := service{
host: v.Host,
healthCheck: v.HealthCheck,
weight: v.Weight,
alive: true,
}
res = append(res, &svc)
}
if len(res) == 0 {
log.Fatal().Msg("no services found")
}
return res
}
// Next returns next url
func (p *RandomWeightedPool) Next() (string, error) {
alive, err := p.getAlive()
if err != nil {
return "", err
}
if len(alive) == 1 {
return alive[0].host, nil
}
totalWeight := 0
for _, item := range alive {
totalWeight += item.weight
}
r := rand.Intn(totalWeight) + 1
for _, item := range alive {
r -= item.weight
if r <= 0 {
return item.host, nil
}
}
return alive[0].host, nil
}
type hostAndWeight struct {
host string
weight int
}
func (p *RandomWeightedPool) getAlive() ([]hostAndWeight, error) {
p.lock.RLock()
defer p.lock.RUnlock()
if len(p.services) == 0 {
return nil, errors.New("empty services list")
}
var res []hostAndWeight
for _, svc := range p.services {
if svc.alive && svc.weight > 0 {
hw := hostAndWeight{
host: svc.host,
weight: svc.weight,
}
res = append(res, hw)
}
}
if len(res) == 0 {
return nil, errors.New("no alive services")
}
return res, nil
}