├── .gitignore ├── go.mod ├── go.sum ├── traffic.go ├── weighted_round_robin_test.go ├── smooth_weighted_round_robin_test.go ├── .travis.yml ├── CONTRIBUTING.md ├── LICENSE.md ├── weighted_round_robin.go ├── README.md └── smooth_weighted_rr.go /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | *.DS_Store 3 | -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/thedevsaddam/traffic 2 | 3 | go 1.19 4 | 5 | require github.com/davecgh/go-spew v1.1.1 // indirect 6 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 2 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 3 | -------------------------------------------------------------------------------- /traffic.go: -------------------------------------------------------------------------------- 1 | package traffic 2 | 3 | import ( 4 | "fmt" 5 | ) 6 | 7 | // Traffic represents a round robin algorithm contract 8 | type Traffic interface { 9 | Add(peer interface{}, weight int) error 10 | Next() interface{} 11 | Reset() 12 | fmt.Stringer 13 | } 14 | 15 | type Peer struct { 16 | Name interface{} 17 | Weight int 18 | CurrentWeight int 19 | EffectiveWeight int 20 | } 21 | -------------------------------------------------------------------------------- /weighted_round_robin_test.go: -------------------------------------------------------------------------------- 1 | package traffic 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestNewWeightedRoundRobin(t *testing.T) { 8 | rr := NewWeightedRoundRobin() 9 | rr.Add("a", 5) 10 | rr.Add("b", 2) 11 | rr.Add("c", 3) 12 | 13 | m := map[interface{}]int{} 14 | for i := 0; i < 100; i++ { 15 | m[rr.Next()]++ 16 | } 17 | if m["a"] != 50 && m["b"] != 20 && m["c"] != 30 { 18 | t.Fail() 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /smooth_weighted_round_robin_test.go: -------------------------------------------------------------------------------- 1 | package traffic 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestNewSmoothWeightedRoundRobin(t *testing.T) { 8 | rr := NewSmoothWeightedRoundRobin() 9 | rr.Add("a", 5) 10 | rr.Add("b", 2) 11 | rr.Add("c", 3) 12 | 13 | m := map[interface{}]int{} 14 | for i := 0; i < 100; i++ { 15 | m[rr.Next()]++ 16 | } 17 | if m["a"] != 50 && m["b"] != 20 && m["c"] != 30 { 18 | t.Fail() 19 | } 20 | } 21 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: go 2 | sudo: false 3 | 4 | matrix: 5 | include: 6 | - go: 1.6 7 | - go: 1.7 8 | - go: 1.8 9 | - go: 1.9 10 | - go: 1.10.x 11 | - go: 1.11.x 12 | - go: 1.12.x 13 | - go: 1.13.x 14 | allow_failures: 15 | - go: tip 16 | before_install: 17 | - go get github.com/mattn/goveralls 18 | script: 19 | - $GOPATH/bin/goveralls -service=travis-ci 20 | - go get -t -v ./... 21 | - diff -u <(echo -n) <(gofmt -d .) 22 | - go vet $(go list ./... | grep -v /vendor/) 23 | - go test -v -race ./... 24 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing 2 | 3 | ## Must follow the guide for issues 4 | - Use the search tool before opening a new issue. 5 | - Please provide source code and stack trace if you found a bug. 6 | - Please review the existing issues and provide feedback to them 7 | 8 | ## Pull Request Process 9 | - Open your pull request against `dev` branch 10 | - It should pass all tests in the available continuous integrations systems such as TravisCI. 11 | - You should add/modify tests to cover your proposed code changes. 12 | - If your pull request contains a new feature, please document it on the README. 13 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | # The MIT License (MIT) 2 | 3 | Copyright (c) 2023 Saddam H 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 13 | > all 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 21 | > THE SOFTWARE. 22 | -------------------------------------------------------------------------------- /weighted_round_robin.go: -------------------------------------------------------------------------------- 1 | package traffic 2 | 3 | import ( 4 | "fmt" 5 | "math/rand" 6 | "sync" 7 | "time" 8 | ) 9 | 10 | func NewWeightedRoundRobin() Traffic { 11 | return &WeightedRoundRobin{ 12 | mu: sync.Mutex{}, 13 | peers: make([]*Peer, 0), 14 | counter: 0, 15 | peersCount: 0, 16 | } 17 | } 18 | 19 | // WeightedRoundRobin ... 20 | type WeightedRoundRobin struct { 21 | mu sync.Mutex 22 | peers []*Peer 23 | peersCount int 24 | counter int64 25 | } 26 | 27 | func (w *WeightedRoundRobin) String() string { 28 | str := fmt.Sprintf("\nTotal: %d\nCounter: %d\n", w.peersCount, w.counter) 29 | for i, p := range w.peers { 30 | str += fmt.Sprintf("Peer[%d]: %s (%d)\n", i+1, p.Name, p.Weight) 31 | } 32 | return str 33 | } 34 | 35 | func (w *WeightedRoundRobin) Reset() { 36 | w.counter = 0 37 | w.peersCount = 0 38 | w.peers = make([]*Peer, 0) 39 | } 40 | 41 | func (w *WeightedRoundRobin) Add(peer interface{}, weight int) error { 42 | w.mu.Lock() 43 | defer w.mu.Unlock() 44 | 45 | for _, p := range w.peers { 46 | if p == peer { 47 | return fmt.Errorf("wrr: peer already exist") 48 | } 49 | } 50 | 51 | for i := 0; i < weight; i++ { 52 | w.peers = append(w.peers, &Peer{Name: peer, Weight: weight}) 53 | } 54 | 55 | w.peersCount += weight 56 | 57 | rand.Seed(time.Now().UnixNano()) 58 | rand.Shuffle(w.peersCount, func(i, j int) { 59 | w.peers[i], w.peers[j] = w.peers[j], w.peers[i] 60 | }) 61 | 62 | return nil 63 | } 64 | 65 | func (w *WeightedRoundRobin) Next() interface{} { 66 | w.mu.Lock() 67 | defer w.mu.Unlock() 68 | 69 | if w.peersCount == 0 { 70 | return nil 71 | } 72 | 73 | p := w.peers[w.counter%int64(w.peersCount)] 74 | w.counter++ 75 | 76 | if w.counter >= int64(w.peersCount) { 77 | w.counter = 0 78 | } 79 | 80 | return p.Name 81 | } 82 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | traffic 2 | ================== 3 | 4 | [![Build Status](https://travis-ci.org/thedevsaddam/traffic.svg?branch=master)](https://travis-ci.org/thedevsaddam/traffic) 5 | [![Project status](https://img.shields.io/badge/version-1.0-green.svg)](https://github.com/thedevsaddam/traffic/releases) 6 | [![Go Report Card](https://goreportcard.com/badge/github.com/thedevsaddam/traffic)](https://goreportcard.com/report/github.com/thedevsaddam/traffic) 7 | [![Coverage Status](https://coveralls.io/repos/github/thedevsaddam/traffic/badge.svg?branch=master)](https://coveralls.io/github/thedevsaddam/traffic?branch=master) 8 | [![GoDoc](https://godoc.org/github.com/thedevsaddam/traffic?status.svg)](https://pkg.go.dev/github.com/thedevsaddam/traffic) 9 | [![License](https://img.shields.io/dub/l/vibe-d.svg)](https://github.com/thedevsaddam/traffic/blob/main/LICENSE.md) 10 | 11 | 12 | Thread safe load-balancer package for Golang 13 | 14 | ### Installation 15 | 16 | Install the package using 17 | ```go 18 | $ go get github.com/thedevsaddam/traffic 19 | ``` 20 | 21 | ### Usage 22 | 23 | To use the package import it in your `*.go` code 24 | ```go 25 | import "github.com/thedevsaddam/traffic" 26 | ``` 27 | 28 | ### Example 29 | 30 | ```go 31 | 32 | package main 33 | 34 | import ( 35 | "fmt" 36 | 37 | "github.com/thedevsaddam/traffic" 38 | ) 39 | 40 | func main() { 41 | t := traffic.NewWeightedRoundRobin() 42 | t.Add("a", 5) 43 | t.Add("b", 2) 44 | t.Add("c", 3) 45 | 46 | for i := 0; i < 100; i++ { 47 | fmt.Println(t.Next()) 48 | } 49 | } 50 | 51 | ``` 52 | 53 | ### **Contribution** 54 | If you are interested to make the package better please send pull requests or create an issue so that others can fix. Read the [contribution guide here](CONTRIBUTING.md). 55 | 56 | ### **License** 57 | The **traffic** is an open-source software licensed under the [MIT License](LICENSE.md). 58 | -------------------------------------------------------------------------------- /smooth_weighted_rr.go: -------------------------------------------------------------------------------- 1 | package traffic 2 | 3 | import ( 4 | "fmt" 5 | "sync" 6 | ) 7 | 8 | /* 9 | SW (Smooth Weighted) is a struct that contains weighted items and provides methods to select a weighted item. 10 | It is used for the smooth weighted round-robin balancing algorithm. This algorithm is implemented in Nginx: 11 | https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35. 12 | http://hg.nginx.org/nginx/rev/c90801720a0c 13 | 14 | Algorithm is as follows: on each peer selection we increase current_weight 15 | of each eligible peer by its weight, select peer with greatest current_weight 16 | and reduce its current_weight by total number of weight points distributed 17 | among peers. 18 | 19 | In case of { 5, 1, 1 } weights this gives the following sequence of 20 | current_weight's: (a, a, b, a, c, a, a) 21 | */ 22 | 23 | func NewSmoothWeightedRoundRobin() Traffic { 24 | return &SmoothWeightedRoundRobin{ 25 | mu: sync.Mutex{}, 26 | peers: make([]*Peer, 0), 27 | peersCount: 0, 28 | } 29 | } 30 | 31 | type SmoothWeightedRoundRobin struct { 32 | mu sync.Mutex 33 | peers []*Peer 34 | peersCount int 35 | } 36 | 37 | func (w *SmoothWeightedRoundRobin) String() string { 38 | str := fmt.Sprintf("\nCounter: %d\n", w.peersCount) 39 | for i, p := range w.peers { 40 | str += fmt.Sprintf("Peer[%d]: %s (%d)\n", i+1, p.Name, p.Weight) 41 | } 42 | return str 43 | } 44 | 45 | // Add a weighted server. 46 | func (w *SmoothWeightedRoundRobin) Add(peer interface{}, weight int) error { 47 | w.mu.Lock() 48 | defer w.mu.Unlock() 49 | 50 | for _, p := range w.peers { 51 | if p == peer { 52 | return fmt.Errorf("wrr: peer already exist") 53 | } 54 | } 55 | 56 | w.peers = append(w.peers, &Peer{Name: peer, Weight: weight, EffectiveWeight: weight}) 57 | w.peersCount++ 58 | 59 | return nil 60 | } 61 | 62 | // Next returns selected server. 63 | // https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35 64 | func (w *SmoothWeightedRoundRobin) Next() interface{} { 65 | w.mu.Lock() 66 | defer w.mu.Unlock() 67 | 68 | if w.peersCount == 0 { 69 | return nil 70 | } 71 | 72 | if w.peersCount == 1 { 73 | return w.peers[0].Name 74 | } 75 | 76 | total := 0 77 | var best *Peer 78 | 79 | for i := 0; i < w.peersCount; i++ { 80 | p := w.peers[i] 81 | 82 | if p == nil { 83 | continue 84 | } 85 | 86 | p.CurrentWeight += p.EffectiveWeight 87 | total += p.EffectiveWeight 88 | if p.EffectiveWeight < p.Weight { 89 | p.EffectiveWeight++ 90 | } 91 | 92 | if best == nil || 93 | p.CurrentWeight > best.CurrentWeight { 94 | best = p 95 | } 96 | 97 | } 98 | 99 | if best == nil { 100 | return nil 101 | } 102 | 103 | best.CurrentWeight -= total 104 | 105 | return best.Name 106 | } 107 | 108 | func (w *SmoothWeightedRoundRobin) Reset() { 109 | for _, p := range w.peers { 110 | p.EffectiveWeight = p.Weight 111 | p.CurrentWeight = 0 112 | } 113 | } 114 | --------------------------------------------------------------------------------