├── .github └── workflows │ └── go.yml ├── .gitignore ├── LICENSE ├── README.md ├── aperture ├── aperture.go ├── aperture_test.go ├── ring.go └── ring_test.go ├── go.mod ├── go.sum ├── internal └── internal.go ├── loadbalance.go ├── p2c ├── least_loaded.go ├── least_loaded_test.go ├── pewma.go └── pewma_test.go ├── roundrobin ├── smooth_weighted.go └── smooth_weighted_test.go └── set ├── set.go └── set_test.go /.github/workflows/go.yml: -------------------------------------------------------------------------------- 1 | name: Go 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | pull_request: 7 | branches: [ master ] 8 | 9 | jobs: 10 | 11 | build: 12 | name: Build 13 | runs-on: ubuntu-latest 14 | steps: 15 | 16 | - name: Set up Go 1.x 17 | uses: actions/setup-go@v2 18 | with: 19 | go-version: ^1.13 20 | id: go 21 | 22 | - name: Check out code into the Go module directory 23 | uses: actions/checkout@v2 24 | 25 | - name: Get dependencies 26 | run: | 27 | go get -v -t -d ./... 28 | if [ -f Gopkg.toml ]; then 29 | curl https://raw.githubusercontent.com/golang/dep/master/install.sh | sh 30 | dep ensure 31 | fi 32 | 33 | - name: Build 34 | run: go build -v . 35 | 36 | - name: Test 37 | run: go test -v ./... -coverprofile coverage.out -covermode=atomic 38 | 39 | - name: Codecov 40 | uses: codecov/codecov-action@v1.0.12 41 | with: 42 | token: ${{ secrets.CODECOV_TOKEN }} 43 | file: ./coverage.out -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/sysulq/go-loadbalance/d51f523d82513cc02253bb1c3059b9d298c9a94c/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Sophos 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | go-loadbalance 2 | === 3 | 4 | [![codecov](https://codecov.io/gh/sysulq/go-loadbalance/branch/master/graph/badge.svg?token=NMTC2ENZQA)](https://codecov.io/gh/hnlq715/go-loadbalance) 5 | -------------------------------------------------------------------------------- /aperture/aperture.go: -------------------------------------------------------------------------------- 1 | package aperture 2 | 3 | import ( 4 | "math" 5 | 6 | "github.com/hnlq715/go-loadbalance" 7 | "github.com/hnlq715/go-loadbalance/p2c" 8 | "github.com/hnlq715/go-loadbalance/roundrobin" 9 | "google.golang.org/grpc/balancer" 10 | ) 11 | 12 | // Aperture support map local peers to remote peers 13 | // to divide remote peers into subsets 14 | // to reduce the connections and separate services into small sets 15 | type aperture struct { 16 | localID string 17 | localPeers []string 18 | localPeersMap map[string]int 19 | remotePeers []interface{} 20 | logicalAperture int 21 | 22 | picker loadbalance.Picker 23 | apertureIdxes []int 24 | } 25 | 26 | const ( 27 | // defaultLogicalAperture means the max logic aperture size 28 | // to control the stability for aperture load balance algorithm 29 | defaultLogicalAperture int = 12 30 | ) 31 | 32 | // NewLeastLoadedApeture returns an Apeture interface with least loaded p2c 33 | func NewLeastLoadedApeture() loadbalance.Aperture { 34 | return &aperture{ 35 | logicalAperture: defaultLogicalAperture, 36 | localPeers: make([]string, 0), 37 | localPeersMap: make(map[string]int), 38 | remotePeers: make([]interface{}, 0), 39 | picker: p2c.NewLeastLoaded(), 40 | } 41 | } 42 | 43 | // NewPeakEwmaAperture returns an Apeture interface with pewma p2c 44 | func NewPeakEwmaAperture() loadbalance.Aperture { 45 | return &aperture{ 46 | logicalAperture: defaultLogicalAperture, 47 | localPeers: make([]string, 0), 48 | localPeersMap: make(map[string]int), 49 | remotePeers: make([]interface{}, 0), 50 | picker: p2c.NewPeakEwma(), 51 | } 52 | } 53 | 54 | // NewSmoothRoundrobin returns an Apeture interface with smooth roundrobin 55 | func NewSmoothRoundrobin() loadbalance.Aperture { 56 | return &aperture{ 57 | logicalAperture: defaultLogicalAperture, 58 | localPeers: make([]string, 0), 59 | localPeersMap: make(map[string]int), 60 | remotePeers: make([]interface{}, 0), 61 | picker: roundrobin.NewSmoothRoundrobin(), 62 | } 63 | } 64 | 65 | // SetLogicalAperture sets the logical aperture size 66 | func (a *aperture) SetLogicalAperture(width int) { 67 | if width > 0 { 68 | a.logicalAperture = width 69 | a.rebuild() 70 | } 71 | } 72 | 73 | // SetLocalPeerID sets the local peer id 74 | func (a *aperture) SetLocalPeerID(id string) { 75 | a.localID = id 76 | a.rebuild() 77 | } 78 | 79 | // SetLocalPeers sets the local peers 80 | func (a *aperture) SetLocalPeers(localPeers []string) { 81 | a.localPeers = localPeers 82 | for idx, local := range localPeers { 83 | a.localPeersMap[local] = idx 84 | } 85 | 86 | a.rebuild() 87 | } 88 | 89 | // SetRemotePeers sets the remote peers 90 | func (a *aperture) SetRemotePeers(remotePeers []interface{}) { 91 | a.remotePeers = remotePeers 92 | a.rebuild() 93 | } 94 | 95 | // Next returns the next selected item 96 | func (a *aperture) Next() (interface{}, func(balancer.DoneInfo)) { 97 | return a.picker.Next() 98 | } 99 | 100 | // List returns the remote peers for the local peer id 101 | // NOTE: current for test/debug only 102 | func (a *aperture) List() []int { 103 | return a.apertureIdxes 104 | } 105 | 106 | // rebuild just rebuilds the aperture when any arguments changed 107 | func (a *aperture) rebuild() { 108 | if len(a.localPeers) == 0 { 109 | return 110 | } 111 | 112 | if len(a.remotePeers) == 0 { 113 | return 114 | } 115 | 116 | idx, ok := a.localPeersMap[a.localID] 117 | if !ok { 118 | return 119 | } 120 | 121 | localWidth := floatOne / float64(len(a.localPeers)) 122 | remoteWidth := floatOne / float64(len(a.remotePeers)) 123 | 124 | if a.logicalAperture > len(a.remotePeers) { 125 | a.logicalAperture = len(a.remotePeers) 126 | } 127 | 128 | apertureWidth := dApertureWidth(localWidth, remoteWidth, a.logicalAperture) 129 | offset := float64(idx) * apertureWidth 130 | 131 | ring := newRing(len(a.remotePeers)) 132 | a.apertureIdxes = ring.Slice(offset, apertureWidth) 133 | 134 | a.picker.Reset() 135 | for _, apertureIdx := range a.apertureIdxes { 136 | weight := ring.Weight(apertureIdx, offset, apertureWidth) 137 | a.picker.Add(a.remotePeers[apertureIdx], weight) 138 | } 139 | } 140 | 141 | // dApertureWidth calculates the actual aperture size base on logic aperture size 142 | func dApertureWidth(localWidth, remoteWidth float64, logicalAperture int) float64 { 143 | unitWidth := localWidth 144 | unitAperture := float64(logicalAperture) * remoteWidth 145 | n := math.Ceil(unitAperture / unitWidth) 146 | width := n * unitWidth 147 | 148 | return math.Min(floatOne, width) 149 | } 150 | -------------------------------------------------------------------------------- /aperture/aperture_test.go: -------------------------------------------------------------------------------- 1 | package aperture 2 | 3 | import ( 4 | "sync" 5 | "testing" 6 | "time" 7 | 8 | "github.com/stretchr/testify/assert" 9 | "google.golang.org/grpc/balancer" 10 | ) 11 | 12 | func TestAperture(t *testing.T) { 13 | t.Run("0 item", func(t *testing.T) { 14 | ll := NewPeakEwmaAperture() 15 | item, done := ll.Next() 16 | done(balancer.DoneInfo{}) 17 | assert.Nil(t, item) 18 | }) 19 | 20 | t.Run("1 client 1 server", func(t *testing.T) { 21 | ll := NewSmoothRoundrobin() 22 | ll.SetLocalPeers(nil) 23 | ll.SetLocalPeers([]string{"1"}) 24 | ll.SetRemotePeers([]interface{}{"8"}) 25 | ll.SetLocalPeerID("1") 26 | 27 | item, done := ll.Next() 28 | done(balancer.DoneInfo{}) 29 | assert.Equal(t, "8", item) 30 | }) 31 | 32 | t.Run("1 client 1 server", func(t *testing.T) { 33 | ll := NewLeastLoadedApeture() 34 | ll.SetLocalPeers(nil) 35 | ll.SetLocalPeers([]string{"1"}) 36 | ll.SetRemotePeers([]interface{}{"8"}) 37 | ll.SetLocalPeerID("1") 38 | 39 | item, done := ll.Next() 40 | done(balancer.DoneInfo{}) 41 | assert.Equal(t, "8", item) 42 | }) 43 | 44 | t.Run("3 client 3 server", func(t *testing.T) { 45 | ll := NewLeastLoadedApeture() 46 | ll.SetLocalPeers([]string{"1", "2", "3"}) 47 | ll.SetRemotePeers([]interface{}{"8", "9", "10"}) 48 | ll.SetLocalPeerID("1") 49 | ll.SetLogicalAperture(1) 50 | 51 | item, done := ll.Next() 52 | done(balancer.DoneInfo{}) 53 | assert.Equal(t, "8", item) 54 | 55 | ll.SetLocalPeerID("2") 56 | 57 | item, done = ll.Next() 58 | done(balancer.DoneInfo{}) 59 | assert.Equal(t, "9", item) 60 | 61 | ll.SetLocalPeerID("3") 62 | 63 | item, done = ll.Next() 64 | done(balancer.DoneInfo{}) 65 | assert.Equal(t, "10", item) 66 | }) 67 | 68 | t.Run("count", func(t *testing.T) { 69 | ll := NewLeastLoadedApeture() 70 | ll.SetLocalPeers([]string{"1", "2", "3"}) 71 | ll.SetRemotePeers([]interface{}{"8", "9", "10", "11", "12"}) 72 | ll.SetLocalPeerID("1") 73 | ll.SetLogicalAperture(2) 74 | 75 | countMap := make(map[interface{}]int) 76 | 77 | totalCount := 5000 78 | wg := sync.WaitGroup{} 79 | wg.Add(totalCount) 80 | 81 | mu := sync.Mutex{} 82 | for i := 0; i < totalCount; i++ { 83 | go func() { 84 | defer wg.Done() 85 | item, done := ll.Next() 86 | time.Sleep(1 * time.Second) 87 | done(balancer.DoneInfo{}) 88 | 89 | mu.Lock() 90 | countMap[item]++ 91 | mu.Unlock() 92 | }() 93 | } 94 | 95 | wg.Wait() 96 | 97 | total := 0 98 | for _, count := range countMap { 99 | total += count 100 | } 101 | assert.Less(t, totalCount*3/10-10, countMap["8"]) 102 | assert.Less(t, totalCount*3/10-10, countMap["9"]) 103 | assert.Less(t, totalCount*3/10-10, countMap["10"]) 104 | assert.Less(t, totalCount*1/10-10, countMap["11"]) 105 | 106 | assert.Equal(t, totalCount, total) 107 | }) 108 | } 109 | 110 | func TestDynamic(t *testing.T) { 111 | t.Run("1client-3client", func(t *testing.T) { 112 | ll := NewLeastLoadedApeture() 113 | ll.SetLocalPeers([]string{"1"}) 114 | ll.SetRemotePeers([]interface{}{"8", "9", "10"}) 115 | ll.SetLocalPeerID("1") 116 | ll.SetLogicalAperture(2) 117 | 118 | assert.Equal(t, []int{0, 1, 2}, ll.(*aperture).List()) 119 | 120 | ll.SetLocalPeers([]string{"1", "2", "3"}) 121 | assert.Equal(t, []int{0, 1}, ll.(*aperture).List()) 122 | 123 | }) 124 | 125 | t.Run("3server-4server", func(t *testing.T) { 126 | ll := NewLeastLoadedApeture() 127 | ll.SetLocalPeers([]string{"1", "2", "3"}) 128 | ll.SetRemotePeers([]interface{}{"8", "9", "10"}) 129 | ll.SetLocalPeerID("1") 130 | ll.SetLogicalAperture(2) 131 | 132 | assert.Equal(t, []int{0, 1}, ll.(*aperture).List()) 133 | 134 | ll.SetRemotePeers([]interface{}{"1", "2", "3", "4"}) 135 | assert.Equal(t, []int{0, 1, 2}, ll.(*aperture).List()) 136 | 137 | }) 138 | } 139 | -------------------------------------------------------------------------------- /aperture/ring.go: -------------------------------------------------------------------------------- 1 | package aperture 2 | 3 | import ( 4 | "math" 5 | ) 6 | 7 | // ring maps the indices [0, `size`) uniformly around a coordinate space [0.0, 1.0). 8 | type ring struct { 9 | size int 10 | unitWidth float64 11 | } 12 | 13 | const ( 14 | floatOne float64 = 1.0 15 | intOne int = 1 16 | ) 17 | 18 | func newRing(size int) *ring { 19 | return &ring{ 20 | size: size, 21 | unitWidth: floatOne / float64(size), 22 | } 23 | } 24 | 25 | // Range returns the total number of indices that [offset, offset + width) intersects with. 26 | func (r *ring) Range(offset, width float64) int { 27 | begin := r.Index(offset) 28 | end := r.Index(math.Mod(offset+width, 1.0)) 29 | 30 | if width < floatOne { 31 | if begin == end && width > r.unitWidth { 32 | return r.size 33 | } else if begin == end { 34 | return intOne 35 | } 36 | 37 | beginWeight := r.Weight(begin, offset, width) 38 | endWeight := r.Weight(end, offset, width) 39 | 40 | adjustedBegin := begin 41 | if beginWeight <= 0 { 42 | adjustedBegin++ 43 | } 44 | 45 | adjustedEnd := end 46 | if endWeight > 0 { 47 | adjustedEnd++ 48 | } 49 | 50 | diff := adjustedEnd - adjustedBegin 51 | if diff <= 0 { 52 | return diff + r.size 53 | } 54 | 55 | return diff 56 | } 57 | 58 | return r.size 59 | } 60 | 61 | // Slice returns the indices where [offset, offset + width) intersects. 62 | func (r *ring) Slice(offset, width float64) []int { 63 | seq := make([]int, 0) 64 | i := r.Index(offset) 65 | rr := r.Range(offset, width) 66 | 67 | for rr > 0 { 68 | idx := i % r.size 69 | seq = append(seq, idx) 70 | i++ 71 | rr-- 72 | } 73 | 74 | return seq 75 | } 76 | 77 | // Index returns the (zero-based) index between [0, `size`) which the 78 | // position `offset` maps to. 79 | func (r *ring) Index(offset float64) int { 80 | return int(math.Floor(offset*float64(r.size))) % r.size 81 | } 82 | 83 | // Weight returns the ratio of the intersection between `index` and [offset, offset + width). 84 | func (r *ring) Weight(index int, offset, width float64) float64 { 85 | ab := float64(index) * r.unitWidth 86 | if ab+1 < offset+width { 87 | ab++ 88 | } 89 | 90 | ae := ab + r.unitWidth 91 | 92 | return intersect(ab, ae, offset, offset+width) / r.unitWidth 93 | } 94 | 95 | // intersect returns the length of the intersection between the two ranges. 96 | func intersect(b0, e0, b1, e1 float64) float64 { 97 | len := math.Min(e0, e1) - math.Max(b0, b1) 98 | return math.Max(0, len) 99 | } 100 | -------------------------------------------------------------------------------- /aperture/ring_test.go: -------------------------------------------------------------------------------- 1 | package aperture 2 | 3 | import ( 4 | "math" 5 | "testing" 6 | 7 | "gotest.tools/assert" 8 | ) 9 | 10 | func TestRing(t *testing.T) { 11 | r := newRing(3) 12 | 13 | len := float64(3) 14 | width := 1.0 / len 15 | 16 | offset := math.Mod(0*width, 1.0) 17 | assert.DeepEqual(t, []int{0}, r.Slice(offset, width)) 18 | 19 | offset = math.Mod(1*width, 1.0) 20 | assert.DeepEqual(t, []int{1}, r.Slice(offset, width)) 21 | 22 | offset = math.Mod(2*width, 1.0) 23 | assert.DeepEqual(t, []int{2}, r.Slice(offset, width)) 24 | } 25 | 26 | func TestRing1(t *testing.T) { 27 | r := newRing(5) 28 | 29 | len := float64(3) 30 | width := 1.0 / len 31 | 32 | offset := math.Mod(0*width, 1.0) 33 | assert.DeepEqual(t, []int{0, 1}, r.Slice(offset, width)) 34 | 35 | offset = math.Mod(1*width, 1.0) 36 | assert.DeepEqual(t, []int{1, 2, 3}, r.Slice(offset, width)) 37 | 38 | offset = math.Mod(2*width, 1.0) 39 | assert.DeepEqual(t, []int{3, 4}, r.Slice(offset, width)) 40 | 41 | } 42 | 43 | func TestRing2(t *testing.T) { 44 | r := newRing(5) 45 | 46 | len := float64(3) 47 | width := 1.0 / len 48 | 49 | offset := float64(0) * width 50 | assert.Equal(t, float64(10), math.Round(r.Weight(0, offset, width)*10)) 51 | assert.Equal(t, float64(7), math.Round(r.Weight(1, offset, width)*10)) 52 | assert.Equal(t, float64(0), math.Round(r.Weight(2, offset, width)*10)) 53 | assert.Equal(t, float64(0), math.Round(r.Weight(3, offset, width)*10)) 54 | assert.Equal(t, float64(0), math.Round(r.Weight(4, offset, width)*10)) 55 | 56 | offset = float64(1) * width 57 | assert.Equal(t, float64(0), math.Round(r.Weight(0, offset, width)*10)) 58 | assert.Equal(t, float64(3), math.Round(r.Weight(1, offset, width)*10)) 59 | assert.Equal(t, float64(10), math.Round(r.Weight(2, offset, width)*10)) 60 | assert.Equal(t, float64(3), math.Round(r.Weight(3, offset, width)*10)) 61 | assert.Equal(t, float64(0), math.Round(r.Weight(5, offset, width)*10)) 62 | 63 | offset = float64(2) * width 64 | assert.Equal(t, float64(0), math.Round(r.Weight(0, offset, width)*10)) 65 | assert.Equal(t, float64(0), math.Round(r.Weight(1, offset, width)*10)) 66 | assert.Equal(t, float64(0), math.Round(r.Weight(2, offset, width)*10)) 67 | assert.Equal(t, float64(7), math.Round(r.Weight(3, offset, width)*10)) 68 | assert.Equal(t, float64(10), math.Round(r.Weight(4, offset, width)*10)) 69 | } 70 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/hnlq715/go-loadbalance 2 | 3 | go 1.14 4 | 5 | require ( 6 | github.com/davecgh/go-spew v1.1.1 // indirect 7 | github.com/golang/protobuf v1.4.2 // indirect 8 | github.com/kr/pretty v0.2.0 // indirect 9 | github.com/pkg/errors v0.9.1 // indirect 10 | github.com/stretchr/testify v1.6.1 11 | google.golang.org/grpc v1.31.0 12 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect 13 | gotest.tools v2.2.0+incompatible 14 | ) 15 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= 2 | github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= 3 | github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= 4 | github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= 5 | github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= 6 | github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 7 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 8 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 9 | github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= 10 | github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= 11 | github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= 12 | github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= 13 | github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= 14 | github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 15 | github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= 16 | github.com/golang/protobuf v1.3.3 h1:gyjaxf+svBWX08ZjK86iN9geUJF0H6gp2IRKX6Nf6/I= 17 | github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw= 18 | github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= 19 | github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= 20 | github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= 21 | github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= 22 | github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= 23 | github.com/golang/protobuf v1.4.2 h1:+Z5KGCizgyZCbGh1KZqA0fcLLkwbsjIzS4aV2v7wJX0= 24 | github.com/golang/protobuf v1.4.2/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= 25 | github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= 26 | github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= 27 | github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= 28 | github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4= 29 | github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= 30 | github.com/kr/pretty v0.2.0 h1:s5hAObm+yFO5uHYt5dYjxi2rXrsnmRpJx4OYvIWUaQs= 31 | github.com/kr/pretty v0.2.0/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI= 32 | github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ= 33 | github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE= 34 | github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= 35 | github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= 36 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 37 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 38 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 39 | github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= 40 | github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= 41 | github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= 42 | github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= 43 | golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= 44 | golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= 45 | golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= 46 | golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= 47 | golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= 48 | golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 49 | golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 50 | golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= 51 | golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= 52 | golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= 53 | golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 54 | golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 55 | golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= 56 | golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 57 | golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= 58 | golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= 59 | golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= 60 | golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= 61 | golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= 62 | golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= 63 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4= 64 | golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= 65 | google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= 66 | google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= 67 | google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= 68 | google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= 69 | google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= 70 | google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= 71 | google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY= 72 | google.golang.org/grpc v1.31.0 h1:T7P4R73V3SSDPhH7WW7ATbfViLtmamH0DKrP3f9AuDI= 73 | google.golang.org/grpc v1.31.0/go.mod h1:N36X2cJ7JwdamYAgDz+s+rVMFjt3numwzf/HckM8pak= 74 | google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= 75 | google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= 76 | google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= 77 | google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= 78 | google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= 79 | google.golang.org/protobuf v1.23.0 h1:4MY060fB1DLGMB/7MBTLnwQUY6+F09GEiz6SsrNqyzM= 80 | google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= 81 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 82 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= 83 | gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 84 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= 85 | gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 86 | gotest.tools v2.2.0+incompatible h1:VsBPFP1AI068pPrMxtb/S8Zkgf9xEmTLJjfM+P5UIEo= 87 | gotest.tools v2.2.0+incompatible/go.mod h1:DsYFclhRJ6vuDpmuTbkuFWG+y2sxOXAzmJt81HFBacw= 88 | honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= 89 | honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= 90 | -------------------------------------------------------------------------------- /internal/internal.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import "google.golang.org/grpc/balancer" 4 | 5 | var ( 6 | // EmptyDoneFunc is a empty done function 7 | EmptyDoneFunc = func(balancer.DoneInfo) {} 8 | ) 9 | -------------------------------------------------------------------------------- /loadbalance.go: -------------------------------------------------------------------------------- 1 | package loadbalance 2 | 3 | import ( 4 | "google.golang.org/grpc/balancer" 5 | ) 6 | 7 | // Aperture support map local peers to remote peers 8 | // to divide remote peers into subsets 9 | // to separate services into small sets and reduce the total connections 10 | type Aperture interface { 11 | // Next returns next selected item. 12 | Next() (interface{}, func(balancer.DoneInfo)) 13 | // Set logical aperture 14 | SetLogicalAperture(int) 15 | // Set local peer id 16 | SetLocalPeerID(string) 17 | // Set local peers. 18 | SetLocalPeers([]string) 19 | // Set remote peers. 20 | SetRemotePeers([]interface{}) 21 | } 22 | 23 | // SetInfo contains region, zone and set 24 | type SetInfo struct { 25 | // Name, app name defined as set 26 | Name string 27 | // Region, like `bj(beijing)` or `sh(shanghai)` 28 | Region string 29 | // UnitName, unit name defined as subsets 30 | UnitName string 31 | } 32 | 33 | // Set supports divide remote peers into subsets 34 | // based on region, zone and set info 35 | type Set interface { 36 | // Next returns next selected item. 37 | Next() (interface{}, func(balancer.DoneInfo)) 38 | // Add a weighted item with set info. 39 | Add(interface{}, float64, SetInfo) 40 | // Reset this picker 41 | Reset() 42 | } 43 | 44 | // Picker supports multiple algorithms for load balance, 45 | // uses the ideas behind the "power of 2 choices" 46 | // to select two nodes from the underlying vector. 47 | type Picker interface { 48 | // Next returns next selected item. 49 | Next() (interface{}, func(balancer.DoneInfo)) 50 | // Add a weighted item. 51 | Add(interface{}, float64) 52 | // Reset this picker 53 | Reset() 54 | } 55 | -------------------------------------------------------------------------------- /p2c/least_loaded.go: -------------------------------------------------------------------------------- 1 | package p2c 2 | 3 | import ( 4 | "math/rand" 5 | "sync" 6 | "sync/atomic" 7 | "time" 8 | 9 | "github.com/hnlq715/go-loadbalance" 10 | "github.com/hnlq715/go-loadbalance/internal" 11 | "google.golang.org/grpc/balancer" 12 | ) 13 | 14 | type leastLoadedNode struct { 15 | item interface{} 16 | inflight int64 17 | weight float64 18 | } 19 | 20 | type leastLoaded struct { 21 | items []*leastLoadedNode 22 | mu sync.Mutex 23 | rand *rand.Rand 24 | } 25 | 26 | func NewLeastLoaded() loadbalance.Picker { 27 | return &leastLoaded{ 28 | items: make([]*leastLoadedNode, 0), 29 | rand: rand.New(rand.NewSource(time.Now().Unix())), 30 | } 31 | } 32 | 33 | func (p *leastLoaded) Add(item interface{}, weight float64) { 34 | p.items = append(p.items, &leastLoadedNode{item: item, weight: weight}) 35 | } 36 | 37 | func (p *leastLoaded) Reset() { 38 | p.items = p.items[:0] 39 | } 40 | 41 | func (p *leastLoaded) Next() (interface{}, func(balancer.DoneInfo)) { 42 | var sc, backsc *leastLoadedNode 43 | 44 | switch len(p.items) { 45 | case 0: 46 | return nil, internal.EmptyDoneFunc 47 | case 1: 48 | sc = p.items[0] 49 | default: 50 | // rand needs lock 51 | p.mu.Lock() 52 | a := p.rand.Intn(len(p.items)) 53 | b := p.rand.Intn(len(p.items) - 1) 54 | p.mu.Unlock() 55 | 56 | if b >= a { 57 | b++ 58 | } 59 | 60 | sc, backsc = p.items[a], p.items[b] 61 | 62 | // choose the least loaded item based on inflight and weight 63 | scInflight := atomic.LoadInt64(&sc.inflight) 64 | backscInflight := atomic.LoadInt64(&backsc.inflight) 65 | 66 | if float64(scInflight)*backsc.weight > float64(backscInflight)*sc.weight { 67 | sc, backsc = backsc, sc 68 | } 69 | } 70 | 71 | atomic.AddInt64(&sc.inflight, 1) 72 | 73 | return sc.item, func(balancer.DoneInfo) { 74 | atomic.AddInt64(&sc.inflight, -1) 75 | } 76 | } 77 | -------------------------------------------------------------------------------- /p2c/least_loaded_test.go: -------------------------------------------------------------------------------- 1 | package p2c_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/hnlq715/go-loadbalance/p2c" 7 | "github.com/stretchr/testify/assert" 8 | "google.golang.org/grpc/balancer" 9 | ) 10 | 11 | func TestLeastLoaded(t *testing.T) { 12 | t.Run("0 item", func(t *testing.T) { 13 | ll := p2c.NewLeastLoaded() 14 | ll.Reset() 15 | item, done := ll.Next() 16 | done(balancer.DoneInfo{}) 17 | assert.Nil(t, item) 18 | }) 19 | 20 | t.Run("1 item", func(t *testing.T) { 21 | ll := p2c.NewLeastLoaded() 22 | ll.Add(1, 1) 23 | item, done := ll.Next() 24 | done(balancer.DoneInfo{}) 25 | assert.Equal(t, 1, item) 26 | }) 27 | 28 | t.Run("3 items", func(t *testing.T) { 29 | ll := p2c.NewLeastLoaded() 30 | ll.Add(1, 1) 31 | ll.Add(2, 1) 32 | ll.Add(3, 1) 33 | 34 | countMap := make(map[interface{}]int) 35 | 36 | totalCount := 10000 37 | for i := 0; i < totalCount; i++ { 38 | item, done := ll.Next() 39 | done(balancer.DoneInfo{}) 40 | 41 | countMap[item]++ 42 | } 43 | 44 | total := 0 45 | for _, count := range countMap { 46 | total += count 47 | assert.Less(t, totalCount/3-200, count) 48 | } 49 | 50 | assert.Equal(t, totalCount, total) 51 | }) 52 | } 53 | 54 | func TestLeastLoadedAbnormal(t *testing.T) { 55 | t.Run("fixed inflight", func(t *testing.T) { 56 | ll := p2c.NewLeastLoaded() 57 | ll.Add(1, 1) 58 | ll.Add(2, 1) 59 | ll.Add(3, 1) 60 | 61 | item, _ := ll.Next() 62 | 63 | for i := 0; i < 1000; i++ { 64 | next, done := ll.Next() 65 | done(balancer.DoneInfo{}) 66 | 67 | assert.NotEqual(t, item, next) 68 | } 69 | }) 70 | } 71 | -------------------------------------------------------------------------------- /p2c/pewma.go: -------------------------------------------------------------------------------- 1 | package p2c 2 | 3 | import ( 4 | "math" 5 | "math/rand" 6 | "sync" 7 | "sync/atomic" 8 | "time" 9 | 10 | "github.com/hnlq715/go-loadbalance" 11 | "github.com/hnlq715/go-loadbalance/internal" 12 | "google.golang.org/grpc/balancer" 13 | ) 14 | 15 | type peakEwma struct { 16 | stamp int64 17 | value int64 18 | tau time.Duration 19 | } 20 | 21 | const ( 22 | defaultTau = 10000 * time.Millisecond 23 | ) 24 | 25 | func newPEWMA() *peakEwma { 26 | return &peakEwma{ 27 | tau: defaultTau, 28 | } 29 | } 30 | 31 | // Observe 计算peak指数加权移动平均值 32 | func (p *peakEwma) Observe(rtt int64) { 33 | now := time.Now().UnixNano() 34 | 35 | stamp := atomic.SwapInt64(&p.stamp, now) 36 | td := now - stamp 37 | 38 | if td < 0 { 39 | td = 0 40 | } 41 | 42 | w := math.Exp(float64(-td) / float64(p.tau)) 43 | latency := atomic.LoadInt64(&p.value) 44 | 45 | if rtt > latency { 46 | atomic.StoreInt64(&p.value, rtt) 47 | } else { 48 | atomic.StoreInt64(&p.value, int64(float64(latency)*w+float64(rtt)*(1.0-w))) 49 | } 50 | } 51 | 52 | func (p *peakEwma) Value() int64 { 53 | return atomic.LoadInt64(&p.value) 54 | } 55 | 56 | type peakEwmaNode struct { 57 | item interface{} 58 | latency *peakEwma 59 | weight float64 60 | } 61 | 62 | type pewma struct { 63 | items []*peakEwmaNode 64 | mu sync.Mutex 65 | rand *rand.Rand 66 | } 67 | 68 | func NewPeakEwma() loadbalance.Picker { 69 | return &pewma{ 70 | items: make([]*peakEwmaNode, 0), 71 | rand: rand.New(rand.NewSource(time.Now().Unix())), 72 | } 73 | } 74 | 75 | func (p *pewma) Add(item interface{}, weight float64) { 76 | p.items = append(p.items, &peakEwmaNode{item: item, latency: newPEWMA(), weight: weight}) 77 | } 78 | 79 | func (p *pewma) Reset() { 80 | *p = pewma{ 81 | items: make([]*peakEwmaNode, 0), 82 | rand: rand.New(rand.NewSource(time.Now().Unix())), 83 | } 84 | } 85 | 86 | func (p *pewma) Next() (interface{}, func(balancer.DoneInfo)) { 87 | var sc, backsc *peakEwmaNode 88 | begin := time.Now().UnixNano() 89 | 90 | switch len(p.items) { 91 | case 0: 92 | return nil, internal.EmptyDoneFunc 93 | case 1: 94 | sc = p.items[0] 95 | default: 96 | // rand needs lock 97 | p.mu.Lock() 98 | a := p.rand.Intn(len(p.items)) 99 | b := p.rand.Intn(len(p.items) - 1) 100 | p.mu.Unlock() 101 | 102 | if b >= a { 103 | b++ 104 | } 105 | 106 | sc, backsc = p.items[a], p.items[b] 107 | 108 | // choose the least loaded item based on inflight and weight 109 | if float64(sc.latency.Value())*backsc.weight > float64(backsc.latency.Value())*sc.weight { 110 | sc, backsc = backsc, sc 111 | } 112 | } 113 | 114 | return sc.item, func(balancer.DoneInfo) { 115 | end := time.Now().UnixNano() 116 | sc.latency.Observe(end - begin) 117 | } 118 | } 119 | -------------------------------------------------------------------------------- /p2c/pewma_test.go: -------------------------------------------------------------------------------- 1 | package p2c 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "testing" 7 | "time" 8 | 9 | "github.com/stretchr/testify/assert" 10 | "google.golang.org/grpc/balancer" 11 | ) 12 | 13 | func TestPEWMA(t *testing.T) { 14 | p := newPEWMA() 15 | // p.tau = 600 * time.Millisecond 16 | p.Observe(int64(1 * time.Second)) 17 | assert.Equal(t, p.Value(), int64(1*time.Second)) 18 | p.Observe(int64(1 * time.Second)) 19 | assert.Equal(t, p.Value(), int64(1*time.Second)) 20 | p.Observe(int64(1 * time.Second)) 21 | assert.Equal(t, p.Value(), int64(1*time.Second)) 22 | time.Sleep(1 * time.Second) 23 | p.Observe(int64(1 * time.Second)) 24 | assert.Equal(t, p.Value(), int64(1*time.Second)) 25 | p.Observe(int64(2 * time.Second)) 26 | assert.Equal(t, p.Value(), int64(2*time.Second)) 27 | for i := 0; i <= 1000; i++ { 28 | time.Sleep(1 * time.Microsecond) 29 | p.Observe(int64(1 * time.Second)) 30 | } 31 | assert.True(t, p.Value() > int64(1800*time.Millisecond) && p.Value() < int64(2000*time.Millisecond), fmt.Sprintf("%d", p.Value())) 32 | } 33 | 34 | func BenchmarkPeakEwma(b *testing.B) { 35 | b.ResetTimer() 36 | p := newPEWMA() 37 | for i := 0; i < b.N; i++ { 38 | p.Observe(int64(time.Duration(rand.Intn(10)) * time.Second)) 39 | } 40 | // b.Error(p.Value()) 41 | } 42 | 43 | func TestPeakEwma(t *testing.T) { 44 | t.Run("0 item", func(t *testing.T) { 45 | ll := NewPeakEwma() 46 | item, done := ll.Next() 47 | done(balancer.DoneInfo{}) 48 | assert.Nil(t, item) 49 | }) 50 | 51 | t.Run("1 item", func(t *testing.T) { 52 | ll := NewPeakEwma() 53 | ll.Reset() 54 | ll.Add(1, 1) 55 | item, done := ll.Next() 56 | done(balancer.DoneInfo{}) 57 | assert.Equal(t, 1, item) 58 | }) 59 | 60 | t.Run("3 items", func(t *testing.T) { 61 | ll := NewPeakEwma() 62 | ll.Add(1, 1) 63 | ll.Add(2, 1) 64 | ll.Add(3, 1) 65 | 66 | countMap := make(map[interface{}]int) 67 | 68 | totalCount := 1000 69 | for i := 0; i < totalCount; i++ { 70 | item, done := ll.Next() 71 | time.Sleep(time.Millisecond) 72 | done(balancer.DoneInfo{}) 73 | 74 | countMap[item]++ 75 | } 76 | 77 | total := 0 78 | for _, count := range countMap { 79 | total += count 80 | assert.Less(t, totalCount/3-2000, count) 81 | } 82 | 83 | assert.Equal(t, totalCount, total) 84 | }) 85 | } 86 | -------------------------------------------------------------------------------- /roundrobin/smooth_weighted.go: -------------------------------------------------------------------------------- 1 | package roundrobin 2 | 3 | import ( 4 | "math" 5 | 6 | "github.com/hnlq715/go-loadbalance" 7 | "github.com/hnlq715/go-loadbalance/internal" 8 | "google.golang.org/grpc/balancer" 9 | ) 10 | 11 | // smoothRoundrobinNode is a wrapped weighted item. 12 | type smoothRoundrobinNode struct { 13 | Item interface{} 14 | Weight int64 15 | CurrentWeight int64 16 | EffectiveWeight int64 17 | } 18 | 19 | type smoothRoundrobin struct { 20 | items []*smoothRoundrobinNode 21 | n int64 22 | } 23 | 24 | // NewSmoothRoundrobin (Smooth Weighted) contains weighted items and provides methods to select a weighted item. 25 | // It is used for the smooth weighted round-robin balancing algorithm. 26 | // This algorithm is implemented in Nginx: 27 | // https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35. 28 | // 29 | // Algorithm is as follows: on each peer selection we increase current_weight 30 | // of each eligible peer by its weight, select peer with greatest current_weight 31 | // and reduce its current_weight by total number of weight points distributed 32 | // among peers. 33 | // In case of { 5, 1, 1 } weights this gives the following sequence of 34 | // current_weight's: (a, a, b, a, c, a, a) 35 | func NewSmoothRoundrobin() loadbalance.Picker { 36 | return &smoothRoundrobin{} 37 | } 38 | 39 | // Add a weighted server. 40 | func (w *smoothRoundrobin) Add(item interface{}, weight float64) { 41 | wt := int64(math.Floor(weight)) 42 | weighted := &smoothRoundrobinNode{Item: item, Weight: wt, EffectiveWeight: wt} 43 | w.items = append(w.items, weighted) 44 | w.n++ 45 | } 46 | 47 | func (w *smoothRoundrobin) Reset() { 48 | w.items = w.items[:0] 49 | w.n = 0 50 | } 51 | 52 | // Next returns next selected server. 53 | func (w *smoothRoundrobin) Next() (interface{}, func(balancer.DoneInfo)) { 54 | if w.n == 0 { 55 | return nil, internal.EmptyDoneFunc 56 | } 57 | 58 | if w.n == 1 { 59 | return w.items[0].Item, internal.EmptyDoneFunc 60 | } 61 | 62 | return nextSmoothWeighted(w.items).Item, internal.EmptyDoneFunc 63 | } 64 | 65 | // nextSmoothWeighted selects the best node through the smooth weighted roundrobin . 66 | func nextSmoothWeighted(items []*smoothRoundrobinNode) (best *smoothRoundrobinNode) { 67 | total := int64(0) 68 | 69 | for i := 0; i < len(items); i++ { 70 | w := items[i] 71 | 72 | w.CurrentWeight += w.EffectiveWeight 73 | total += w.EffectiveWeight 74 | 75 | if w.EffectiveWeight < w.Weight { 76 | w.EffectiveWeight++ 77 | } 78 | 79 | if best == nil || w.CurrentWeight > best.CurrentWeight { 80 | best = w 81 | } 82 | } 83 | 84 | best.CurrentWeight -= total 85 | 86 | return best 87 | } 88 | -------------------------------------------------------------------------------- /roundrobin/smooth_weighted_test.go: -------------------------------------------------------------------------------- 1 | package roundrobin 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | "google.golang.org/grpc/balancer" 8 | ) 9 | 10 | func TestSW_Next(t *testing.T) { 11 | w := NewSmoothRoundrobin() 12 | 13 | s, done := w.Next() 14 | assert.Nil(t, s) 15 | done(balancer.DoneInfo{}) 16 | 17 | w.Add("server1", 5) 18 | s, _ = w.Next() 19 | assert.Equal(t, "server1", s.(string)) 20 | 21 | w.Reset() 22 | s, _ = w.Next() 23 | assert.Nil(t, s) 24 | 25 | w.Add("server1", 5) 26 | s, _ = w.Next() 27 | assert.Equal(t, "server1", s.(string)) 28 | 29 | w.Add("server2", 2) 30 | w.Add("server3", 3) 31 | 32 | results := make(map[string]int) 33 | 34 | for i := 0; i < 1000; i++ { 35 | s, _ := w.Next() 36 | results[s.(string)]++ 37 | } 38 | 39 | if results["server1"] != 500 || results["server2"] != 200 || results["server3"] != 300 { 40 | t.Error("the algorithm is wrong") 41 | } 42 | 43 | w.(*smoothRoundrobin).items[0].EffectiveWeight = w.(*smoothRoundrobin).items[0].CurrentWeight - 1 44 | s, _ = w.Next() 45 | assert.Equal(t, "server3", s.(string)) 46 | } 47 | -------------------------------------------------------------------------------- /set/set.go: -------------------------------------------------------------------------------- 1 | package set 2 | 3 | import ( 4 | "github.com/hnlq715/go-loadbalance" 5 | "github.com/hnlq715/go-loadbalance/roundrobin" 6 | "google.golang.org/grpc/balancer" 7 | ) 8 | 9 | type Set struct { 10 | info loadbalance.SetInfo 11 | picker loadbalance.Picker 12 | } 13 | 14 | func New(info loadbalance.SetInfo) loadbalance.Set { 15 | return &Set{ 16 | info: info, 17 | picker: roundrobin.NewSmoothRoundrobin(), 18 | } 19 | } 20 | 21 | func (s *Set) Next() (interface{}, func(balancer.DoneInfo)) { 22 | return s.picker.Next() 23 | } 24 | 25 | func (s *Set) Add(item interface{}, weigth float64, info loadbalance.SetInfo) { 26 | if info.Name != s.info.Name { 27 | return 28 | } 29 | 30 | if info.Region != s.info.Region { 31 | return 32 | } 33 | 34 | if s.info.UnitName != "*" { 35 | if info.UnitName != s.info.UnitName { 36 | return 37 | } 38 | } 39 | 40 | s.picker.Add(item, weigth) 41 | } 42 | 43 | func (s *Set) Reset() { 44 | s.picker.Reset() 45 | } 46 | -------------------------------------------------------------------------------- /set/set_test.go: -------------------------------------------------------------------------------- 1 | package set_test 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/hnlq715/go-loadbalance" 7 | "github.com/hnlq715/go-loadbalance/set" 8 | "gotest.tools/assert" 9 | ) 10 | 11 | func TestSet(t *testing.T) { 12 | t.Run("same set", func(t *testing.T) { 13 | s := set.New(loadbalance.SetInfo{ 14 | Name: "app", 15 | Region: "bj", 16 | UnitName: "01", 17 | }) 18 | 19 | s.Add(1, 1, loadbalance.SetInfo{ 20 | Name: "app", 21 | Region: "bj", 22 | UnitName: "01", 23 | }) 24 | 25 | item, _ := s.Next() 26 | assert.Equal(t, 1, item) 27 | 28 | s.Reset() 29 | 30 | item, _ = s.Next() 31 | assert.Equal(t, nil, item) 32 | }) 33 | 34 | t.Run("different region", func(t *testing.T) { 35 | s := set.New(loadbalance.SetInfo{ 36 | Name: "app", 37 | Region: "bj", 38 | UnitName: "01", 39 | }) 40 | 41 | s.Add(1, 1, loadbalance.SetInfo{ 42 | Name: "app", 43 | Region: "sh", 44 | UnitName: "02", 45 | }) 46 | 47 | item, _ := s.Next() 48 | assert.Equal(t, nil, item) 49 | }) 50 | 51 | t.Run("different name", func(t *testing.T) { 52 | s := set.New(loadbalance.SetInfo{ 53 | Name: "app01", 54 | Region: "bj", 55 | UnitName: "01", 56 | }) 57 | 58 | s.Add(1, 1, loadbalance.SetInfo{ 59 | Name: "app02", 60 | Region: "bj", 61 | UnitName: "01", 62 | }) 63 | 64 | item, _ := s.Next() 65 | assert.Equal(t, nil, item) 66 | }) 67 | 68 | t.Run("different set", func(t *testing.T) { 69 | s := set.New(loadbalance.SetInfo{ 70 | Name: "app", 71 | Region: "bj", 72 | UnitName: "01", 73 | }) 74 | 75 | s.Add(1, 1, loadbalance.SetInfo{ 76 | Name: "app", 77 | Region: "bj", 78 | UnitName: "02", 79 | }) 80 | 81 | item, _ := s.Next() 82 | assert.Equal(t, nil, item) 83 | }) 84 | 85 | t.Run("* set", func(t *testing.T) { 86 | s := set.New(loadbalance.SetInfo{ 87 | Name: "app", 88 | Region: "bj", 89 | UnitName: "*", 90 | }) 91 | 92 | s.Add(1, 1, loadbalance.SetInfo{ 93 | Name: "app", 94 | Region: "bj", 95 | UnitName: "*", 96 | }) 97 | 98 | s.Add(2, 1, loadbalance.SetInfo{ 99 | Name: "app", 100 | Region: "bj", 101 | UnitName: "01", 102 | }) 103 | 104 | item, _ := s.Next() 105 | assert.Equal(t, 1, item) 106 | 107 | item, _ = s.Next() 108 | assert.Equal(t, 2, item) 109 | }) 110 | } 111 | --------------------------------------------------------------------------------