/
least_load.go
90 lines (77 loc) · 1.92 KB
/
least_load.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
// Copyright 2022 <mzh.scnu@qq.com>. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package balancer
import (
"sync"
fibHeap "github.com/starwander/GoFibonacciHeap"
)
func init() {
factories[LeastLoadBalancer] = NewLeastLoad
}
// Tag .
func (h *host) Tag() interface{} { return h.name }
// Key .
func (h *host) Key() float64 { return float64(h.load) }
// LeastLoad will choose a host based on the least load host
type LeastLoad struct {
sync.RWMutex
heap *fibHeap.FibHeap
}
// NewLeastLoad create new LeastLoad balancer
func NewLeastLoad(hosts []string) Balancer {
ll := &LeastLoad{heap: fibHeap.NewFibHeap()}
for _, h := range hosts {
ll.Add(h)
}
return ll
}
// Add new host to the balancer
func (l *LeastLoad) Add(hostName string) {
l.Lock()
defer l.Unlock()
if ok := l.heap.GetValue(hostName); ok != nil {
return
}
_ = l.heap.InsertValue(&host{hostName, 0})
}
// Remove new host from the balancer
func (l *LeastLoad) Remove(hostName string) {
l.Lock()
defer l.Unlock()
if ok := l.heap.GetValue(hostName); ok == nil {
return
}
_ = l.heap.Delete(hostName)
}
// Balance selects a suitable host according
func (l *LeastLoad) Balance(_ string) (string, error) {
l.RLock()
defer l.RUnlock()
if l.heap.Num() == 0 {
return "", NoHostError
}
return l.heap.MinimumValue().Tag().(string), nil
}
// Inc refers to the number of connections to the server `+1`
func (l *LeastLoad) Inc(hostName string) {
l.Lock()
defer l.Unlock()
if ok := l.heap.GetValue(hostName); ok == nil {
return
}
h := l.heap.GetValue(hostName)
h.(*host).load++
_ = l.heap.IncreaseKeyValue(h)
}
// Done refers to the number of connections to the server `-1`
func (l *LeastLoad) Done(hostName string) {
l.Lock()
defer l.Unlock()
if ok := l.heap.GetValue(hostName); ok == nil {
return
}
h := l.heap.GetValue(hostName)
h.(*host).load--
_ = l.heap.DecreaseKeyValue(h)
}