├── .github
    └── logo.png
├── .gitignore
├── LICENSE
├── Makefile
├── README.md
├── balance.go
├── balance_test.go
├── benchmark_test.go
├── example
    └── main.go
└── go.mod


/.github/logo.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/mr-karan/balance/d32c6ade6cf1476c2fe5e19c6a1021a27027c87d/.github/logo.png


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
 1 | # Binaries for programs and plugins
 2 | *.exe
 3 | *.exe~
 4 | *.dll
 5 | *.so
 6 | *.dylib
 7 | 
 8 | # Test binary, built with `go test -c`
 9 | *.test
10 | 
11 | # Output of the go coverage tool, specifically when used with LiteIDE
12 | *.out
13 | 
14 | # Dependency directories (remove the comment below to include it)
15 | # vendor/
16 | 


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 Karan Sharma
 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 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
1 | .PHONY: test
2 | test:
3 | 	go test -v -failfast -race -coverpkg=./... -covermode=atomic
4 | 
5 | .PHONY: bench
6 | bench:
7 | 	go test -v -failfast -bench=. -benchmem -run=^$$
8 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | <p align="center">
 2 | <img src="./.github/logo.png" alt="logo" width="15%" />
 3 | </p>
 4 | 
 5 | # balance
 6 | 
 7 | A minimal Golang library for implementing weighted round-robin load balancing for a given set of items.
 8 | 
 9 | ## Installation
10 | 
11 | ```bash
12 | go get github.com/mr-karan/balance
13 | ```
14 | 
15 | ## Usage
16 | 
17 | ```go
18 | package main
19 | 
20 | import (
21 |     "fmt"
22 |     "github.com/mr-karan/balance"
23 |     
24 |     // Create a new load balancer.
25 |     b := balance.NewBalance()
26 | 
27 |     // Add items to the load balancer with their corresponding weights.
28 |     b.Add("a", 5)
29 |     b.Add("b", 3)
30 |     b.Add("c", 2)    
31 |     
32 |     // Get the next item from the load balancer.
33 |     fmt.Println(b.Get())
34 | 
35 |     // For 10 requests, the output sequence will be: [a b c a a b a c b a]
36 | )
37 | ```
38 | 
39 | ## Algorithm
40 | 
41 | The algorithm is based on the [Smooth Weighted Round Robin](https://github.com/phusion/nginx/commit/27e94984486058d73157038f7950a0a36ecc6e35) used by NGINX.
42 | 
43 | > Algorithm is as follows: on each peer selection we increase current_weight
44 | of each eligible peer by its weight, select peer with greatest current_weight
45 | and reduce its current_weight by total number of weight points distributed
46 | among peers.
47 | 
48 | ## Examples
49 | 
50 | ### Round Robin
51 | 
52 | For implementing an equal weighted round-robin load balancer for a set of servers, you can use the following config:
53 | 
54 | ```go
55 | b.Add("server1", 1)
56 | b.Add("server2", 1)
57 | b.Add("server3", 1)
58 | ```
59 | 
60 | Since the weights of all 3 servers are equal, the load balancer will distribute the load equally among all 3 servers.
61 | 
62 | ### Weighted Round Robin
63 | 
64 | For implementing a weighted round-robin load balancer for a set of servers, you can use the following config:
65 | 
66 | ```go
67 | b.Add("server1", 5)
68 | b.Add("server2", 3)
69 | b.Add("server3", 2)
70 | ```
71 | 
72 | The load balancer will distribute the load in the ratio of 5:3:2 among the 3 servers.
73 | 
74 | ## Benchmark
75 | 
76 | ```bash
77 | go test -v -failfast -bench=. -benchmem -run=^$
78 | goos: linux
79 | goarch: amd64
80 | pkg: github.com/mr-karan/balance
81 | cpu: 11th Gen Intel(R) Core(TM) i7-1165G7 @ 2.80GHz
82 | BenchmarkBalance
83 | BenchmarkBalance/items-10
84 | BenchmarkBalance/items-10-8             18249529                63.82 ns/op            0 B/op          0 allocs/op
85 | BenchmarkBalance/items-100
86 | BenchmarkBalance/items-100-8             9840943               119.5 ns/op             0 B/op          0 allocs/op
87 | BenchmarkBalance/items-1000
88 | BenchmarkBalance/items-1000-8            1608460               767.1 ns/op             0 B/op          0 allocs/op
89 | BenchmarkBalance/items-10000
90 | BenchmarkBalance/items-10000-8            123394              9621 ns/op               0 B/op          0 allocs/op
91 | BenchmarkBalance/items-100000
92 | BenchmarkBalance/items-100000-8            10000            102295 ns/op               0 B/op          0 allocs/op
93 | PASS
94 | ok      github.com/mr-karan/balance     7.927s
95 | ```
96 | 


--------------------------------------------------------------------------------
/balance.go:
--------------------------------------------------------------------------------
  1 | package balance
  2 | 
  3 | import (
  4 | 	"errors"
  5 | 	"sync"
  6 | )
  7 | 
  8 | var (
  9 | 	// ErrDuplicateID error is thrown when attempt to add an ID
 10 | 	// which is already added to the balancer.
 11 | 	ErrDuplicateID = errors.New("entry already added")
 12 | 	// ErrIDNotFound is thrown when removing a non-existent ID.
 13 | 	ErrIDNotFound = errors.New("id not found")
 14 | )
 15 | 
 16 | // Balance represents a smooth weighted round-robin load balancer.
 17 | type Balance struct {
 18 | 	sync.RWMutex
 19 | 
 20 | 	// items is the list of items to balance
 21 | 	items []*Item
 22 | 	// next is the index of the next item to use.
 23 | 	next *Item
 24 | }
 25 | 
 26 | // NewBalance creates a new load balancer.
 27 | func NewBalance() *Balance {
 28 | 	return &Balance{
 29 | 		items: make([]*Item, 0),
 30 | 	}
 31 | }
 32 | 
 33 | // Item represents the item in the list.
 34 | type Item struct {
 35 | 	// id is the id of the item.
 36 | 	id string
 37 | 	// weight is the weight of the item that is given by the user.
 38 | 	weight int
 39 | 	// current is the current weight of the item.
 40 | 	current int
 41 | }
 42 | 
 43 | func NewItem(id string, weight int) *Item {
 44 | 	return &Item{
 45 | 		id:      id,
 46 | 		weight:  weight,
 47 | 		current: 0,
 48 | 	}
 49 | }
 50 | 
 51 | func (b *Balance) Add(id string, weight int) error {
 52 | 	b.Lock()
 53 | 	defer b.Unlock()
 54 | 	for _, v := range b.items {
 55 | 		if v.id == id {
 56 | 			return ErrDuplicateID
 57 | 		}
 58 | 	}
 59 | 
 60 | 	b.items = append(b.items, NewItem(id, weight))
 61 | 
 62 | 	return nil
 63 | }
 64 | 
 65 | func (b *Balance) Get() string {
 66 | 	b.Lock()
 67 | 	defer b.Unlock()
 68 | 
 69 | 	if len(b.items) == 0 {
 70 | 		return ""
 71 | 	}
 72 | 
 73 | 	// Total weight of all items.
 74 | 	var total int
 75 | 
 76 | 	// Loop through the list of items and add the item's weight to the current weight.
 77 | 	// Also increment the total weight counter.
 78 | 	var max *Item
 79 | 	for _, item := range b.items {
 80 | 		item.current += item.weight
 81 | 		total += item.weight
 82 | 
 83 | 		// Select the item with max weight.
 84 | 		if max == nil || item.current > max.current {
 85 | 			max = item
 86 | 		}
 87 | 	}
 88 | 
 89 | 	// Select the item with the max weight.
 90 | 	b.next = max
 91 | 	// Reduce the current weight of the selected item by the total weight.
 92 | 	max.current -= total
 93 | 
 94 | 	return max.id
 95 | }
 96 | 
 97 | // Remove deletes an item by ID from the balancer.
 98 | func (b *Balance) Remove(id string) error {
 99 | 	b.Lock()
100 | 	defer b.Unlock()
101 | 
102 | 	for i, item := range b.items {
103 | 		if item.id == id {
104 | 			b.items = append(b.items[:i], b.items[i+1:]...)
105 | 			return nil
106 | 		}
107 | 	}
108 | 
109 | 	return ErrIDNotFound
110 | }
111 | 
112 | // ItemIDs returns a list of all item IDs in the balancer.
113 | func (b *Balance) ItemIDs() []string {
114 | 	b.RLock()
115 | 	defer b.RUnlock()
116 | 
117 | 	ids := make([]string, len(b.items))
118 | 	for i, item := range b.items {
119 | 		ids[i] = item.id
120 | 	}
121 | 	return ids
122 | }
123 | 


--------------------------------------------------------------------------------
/balance_test.go:
--------------------------------------------------------------------------------
  1 | package balance_test
  2 | 
  3 | import (
  4 | 	"errors"
  5 | 	"sync"
  6 | 	"sync/atomic"
  7 | 	"testing"
  8 | 
  9 | 	"github.com/mr-karan/balance"
 10 | )
 11 | 
 12 | func TestBalance(t *testing.T) {
 13 | 	// Test Init.
 14 | 	t.Run("init", func(t *testing.T) {
 15 | 		bl := balance.NewBalance()
 16 | 		if bl.Get() != "" {
 17 | 			t.Error("Expected empty string")
 18 | 		}
 19 | 	})
 20 | 
 21 | 	// Test round robin.
 22 | 	t.Run("round robin", func(t *testing.T) {
 23 | 		bl := balance.NewBalance()
 24 | 		bl.Add("a", 1)
 25 | 		bl.Add("b", 1)
 26 | 		bl.Add("c", 1)
 27 | 		result := make(map[string]int)
 28 | 		for i := 0; i < 999; i++ {
 29 | 			result[bl.Get()]++
 30 | 		}
 31 | 
 32 | 		if result["a"] != 333 || result["b"] != 333 || result["c"] != 333 {
 33 | 			t.Error("Wrong counts", result)
 34 | 		}
 35 | 	})
 36 | 
 37 | 	t.Run("adding duplicate entry", func(t *testing.T) {
 38 | 		bl := balance.NewBalance()
 39 | 		err := bl.Add("c", 1)
 40 | 		if !errors.Is(err, nil) {
 41 | 			t.Error("Wrong error received", err.Error())
 42 | 		}
 43 | 
 44 | 		err = bl.Add("c", 1)
 45 | 		if !errors.Is(err, balance.ErrDuplicateID) {
 46 | 			t.Error("Wrong error received", err.Error())
 47 | 		}
 48 | 	})
 49 | 
 50 | 	// Test weighted.
 51 | 	t.Run("weighted custom split", func(t *testing.T) {
 52 | 		bl := balance.NewBalance()
 53 | 		bl.Add("a", 2)
 54 | 		bl.Add("b", 1)
 55 | 		bl.Add("c", 1)
 56 | 		result := make(map[string]int)
 57 | 		for i := 0; i < 1000; i++ {
 58 | 			result[bl.Get()]++
 59 | 		}
 60 | 
 61 | 		if result["a"] != 500 || result["b"] != 250 || result["c"] != 250 {
 62 | 			t.Error("Wrong counts", result)
 63 | 		}
 64 | 	})
 65 | 
 66 | 	t.Run("weighted another custom split", func(t *testing.T) {
 67 | 		bl := balance.NewBalance()
 68 | 		bl.Add("a", 5)
 69 | 		bl.Add("b", 3)
 70 | 		bl.Add("c", 2)
 71 | 		result := make(map[string]int)
 72 | 		for i := 0; i < 1000; i++ {
 73 | 			result[bl.Get()]++
 74 | 		}
 75 | 
 76 | 		if result["a"] != 500 || result["b"] != 300 || result["c"] != 200 {
 77 | 			t.Error("Wrong counts", result)
 78 | 		}
 79 | 	})
 80 | 
 81 | 	// Test with one item as zero weight.
 82 | 	t.Run("weighted with zero", func(t *testing.T) {
 83 | 		bl := balance.NewBalance()
 84 | 		bl.Add("a", 0)
 85 | 		bl.Add("b", 1)
 86 | 		bl.Add("c", 1)
 87 | 		result := make(map[string]int)
 88 | 		for i := 0; i < 1000; i++ {
 89 | 			result[bl.Get()]++
 90 | 		}
 91 | 
 92 | 		if result["a"] != 0 || result["b"] != 500 || result["c"] != 500 {
 93 | 			t.Error("Wrong counts", result)
 94 | 		}
 95 | 	})
 96 | 
 97 | 	// Test remove item.
 98 | 	t.Run("remove item", func(t *testing.T) {
 99 | 		bl := balance.NewBalance()
100 | 		bl.Add("a", 1)
101 | 		bl.Add("b", 1)
102 | 		bl.Add("c", 1)
103 | 
104 | 		err := bl.Remove("b")
105 | 		if err != nil {
106 | 			t.Error("Expected no error, got", err)
107 | 		}
108 | 
109 | 		ids := bl.ItemIDs()
110 | 		expected := map[string]bool{"a": true, "c": true}
111 | 		for _, id := range ids {
112 | 			if !expected[id] {
113 | 				t.Error("Unexpected ID in list", id)
114 | 			}
115 | 		}
116 | 
117 | 		// Ensure removed item isn't returned by a Get.
118 | 		for i := 0; i < 100; i++ {
119 | 			if bl.Get() == "b" {
120 | 				t.Error("Removed item 'b' still returned by Get")
121 | 			}
122 | 		}
123 | 	})
124 | 
125 | 	// Test remove non-existent item.
126 | 	t.Run("remove non-existent item", func(t *testing.T) {
127 | 		bl := balance.NewBalance()
128 | 		bl.Add("a", 1)
129 | 		err := bl.Remove("x")
130 | 		if !errors.Is(err, balance.ErrIDNotFound) {
131 | 			t.Error("Expected ErrIDNotFound, got", err)
132 | 		}
133 | 	})
134 | 
135 | 	// Test list items ids.
136 | 	t.Run("list items", func(t *testing.T) {
137 | 		bl := balance.NewBalance()
138 | 		bl.Add("x", 3)
139 | 		bl.Add("y", 2)
140 | 
141 | 		ids := bl.ItemIDs()
142 | 		expected := map[string]bool{"x": true, "y": true}
143 | 		for _, id := range ids {
144 | 			if !expected[id] {
145 | 				t.Error("Unexpected ID in list", id)
146 | 			}
147 | 		}
148 | 
149 | 		if len(ids) != 2 {
150 | 			t.Error("Expected 2 items, got", len(ids))
151 | 		}
152 | 	})
153 | }
154 | 
155 | func TestBalance_Concurrent(t *testing.T) {
156 | 	t.Run("concurrent", func(t *testing.T) {
157 | 		var (
158 | 			a, b, c int64
159 | 		)
160 | 		bl := balance.NewBalance()
161 | 		bl.Add("a", 1)
162 | 		bl.Add("b", 1)
163 | 		bl.Add("c", 1)
164 | 
165 | 		var wg sync.WaitGroup
166 | 
167 | 		for i := 0; i < 999; i++ {
168 | 			wg.Add(1)
169 | 			go func() {
170 | 				defer wg.Done()
171 | 				switch bl.Get() {
172 | 				case "a":
173 | 					atomic.AddInt64(&a, 1)
174 | 				case "b":
175 | 					atomic.AddInt64(&b, 1)
176 | 				case "c":
177 | 					atomic.AddInt64(&c, 1)
178 | 				default:
179 | 					t.Error("Wrong item")
180 | 				}
181 | 			}()
182 | 		}
183 | 
184 | 		wg.Wait()
185 | 
186 | 		if a != 333 || b != 333 || c != 333 {
187 | 			t.Error("Wrong counts", a, b, c)
188 | 		}
189 | 	})
190 | }
191 | 


--------------------------------------------------------------------------------
/benchmark_test.go:
--------------------------------------------------------------------------------
 1 | package balance_test
 2 | 
 3 | import (
 4 | 	"math/rand"
 5 | 	"strconv"
 6 | 	"testing"
 7 | 	"time"
 8 | 
 9 | 	"github.com/mr-karan/balance"
10 | )
11 | 
12 | func BenchmarkBalance(b *testing.B) {
13 | 	b.ReportAllocs()
14 | 	rand.Seed(time.Now().UnixNano())
15 | 
16 | 	for n := 10; n <= 100000; n *= 10 {
17 | 		b.Run("items-"+strconv.Itoa(n), func(b *testing.B) {
18 | 			bl := balance.NewBalance()
19 | 			items := generateItems(n)
20 | 			for i, w := range items {
21 | 				bl.Add(i, w)
22 | 			}
23 | 
24 | 			b.ResetTimer()
25 | 			b.RunParallel(func(p *testing.PB) {
26 | 				for p.Next() {
27 | 					_ = bl.Get()
28 | 				}
29 | 			})
30 | 		})
31 | 	}
32 | 
33 | }
34 | 
35 | func generateItems(n int) map[string]int {
36 | 	items := make(map[string]int)
37 | 	for i := 0; i < n; i++ {
38 | 		items["server-"+strconv.Itoa(i)] = rand.Intn(100) + 50
39 | 	}
40 | 	return items
41 | }
42 | 


--------------------------------------------------------------------------------
/example/main.go:
--------------------------------------------------------------------------------
 1 | package main
 2 | 
 3 | import (
 4 | 	"fmt"
 5 | 
 6 | 	balance "github.com/mr-karan/balance"
 7 | )
 8 | 
 9 | func main() {
10 | 	// Create a new load balancer.
11 | 	b := balance.NewBalance()
12 | 
13 | 	// Add items to the load balancer.
14 | 	b.Add("a", 5)
15 | 	b.Add("b", 3)
16 | 	b.Add("c", 2)
17 | 
18 | 	for i := 0; i < 10; i++ {
19 | 		item := b.Get()
20 | 		fmt.Printf("%s ", item)
21 | 	}
22 | }
23 | 


--------------------------------------------------------------------------------
/go.mod:
--------------------------------------------------------------------------------
1 | module github.com/mr-karan/balance
2 | 
3 | go 1.19
4 | 


--------------------------------------------------------------------------------