Merge pull request #2 from Abdulrahman-Tayara/feature/weighted-round-robin

Feature/weighted round robin

Author: Abdulrahman-Tayara

README.md

@@ -40,12 +40,14 @@ Configuring Go-LB is a breeze, it supports both JSON and YAML configuration file
     {
       "name": "server1",
       "url": "http://localhost:8080",
-      "health_url": "/health"
+      "health_url": "/health",
+      "weight": 1,
     },
     {
       "name": "server2",
       "url": "http://localhost:8082",
-      "health_url": "/health-check"
+      "health_url": "/health-check",
+      "weight": 2
     }
   ],
   "tls_enabled": true, // default false
@@ -70,9 +72,11 @@ servers:
   - name: "server1"
     url: "http://localhost:8080"
     health_url: "/health"
+    weight: 1
   - name: "server2"
     url: "http://localhost:8082"
     health_url: "/health-check"
+    weight: 2
 tls_enabled: true # default false
 tls_cert_file: "/path/on/container/cert.pem"
 tls_key_file: "/path/on/container/key.pem"
@@ -93,12 +97,14 @@ tls_key_file: "/path/on/container/key.pem"
     {
       "name": "server1",
       "url": "http://localhost:8080",
-      "health_url": "/health"
+      "health_url": "/health",
+      "weight": 1
     },
     {
       "name": "server2",
       "url": "http://localhost:8082",
-      "health_url": "/health-check"
+      "health_url": "/health-check",
+      "weight": 2
     }
   ],
   "tls_enabled": true,
@@ -137,8 +143,10 @@ rate_limit_interval_seconds: 10
 servers:
   - url: "http://localhost:8080"
     health_url: "/health"
+    weight: 1
   - url: "http://localhost:8082"
     health_url: "/health-check"
+    weight: 2
 tls_enabled: true
 tls_cert_file: "/path/on/container/cert.pem"
 tls_key_file: "/path/on/container/key.pem"

main.go

@@ -10,6 +10,7 @@ import (
 	"tayara/go-lb/lb"
 	"tayara/go-lb/ratelimiter/buckettokens"
 	"tayara/go-lb/strategy"
+	"tayara/go-lb/utils"
 	"time"
 
 	"github.com/pkg/errors"
@@ -78,10 +79,10 @@ func runHTTPServer(cfg *configs.Configs, handler http.Handler) {
 	)
 
 	if cfg.TLSEnabled {
-		if !IsFileExist(cfg.TLSCertPath) {
+		if !utils.IsFileExist(cfg.TLSCertPath) {
 			panic(fmt.Errorf("TLS cert filepath doesn't exist %v", cfg.TLSCertPath))
 		}
-		if !IsFileExist(cfg.TLSKeyPath) {
+		if !utils.IsFileExist(cfg.TLSKeyPath) {
 			panic(fmt.Errorf("TLS key filepath doesn't exist %v", cfg.TLSKeyPath))
 		}
 

models/server.go

@@ -13,6 +13,7 @@ type Server struct {
 	Name      string `mapstructure:"name" json:"name" yaml:"name"`
 	Url       string `mapstructure:"url" json:"url" yaml:"url"`
 	HealthUrl string `mapstructure:"health_url" json:"health_url" yaml:"health_url"`
+	Weight    int    `mapstructure:"weight" json:"weight" yaml:"weight"`
 
 	host              string
 	absoluteHealthUrl string

strategy/weighted_round_robin.go

@@ -0,0 +1,96 @@
+package strategy
+
+import (
+	"net/http"
+	"slices"
+	"sync"
+	"tayara/go-lb/models"
+)
+
+type WeightedRoundRobinStrategy struct {
+	weightsBucket map[string]int
+
+	servers []*models.Server
+
+	sync.RWMutex
+}
+
+func NewWeightedRoundRobinStrategy() ILoadBalancerStrategy {
+	return &WeightedRoundRobinStrategy{
+		weightsBucket: make(map[string]int),
+	}
+}
+
+func (s *WeightedRoundRobinStrategy) Next(request *http.Request) *models.Server {
+	defer s.RUnlock()
+
+	s.RLock()
+
+	return s.selectServer(s.servers)
+}
+
+func (s *WeightedRoundRobinStrategy) selectServer(servers []*models.Server) *models.Server {
+	if len(servers) == 0 {
+		return nil
+	}
+
+	s.ensureWightsBucketValid()
+
+	selectedServer := s.findMaxServerWeight()
+
+	return selectedServer
+}
+
+func (s *WeightedRoundRobinStrategy) findMaxServerWeight() *models.Server {
+	var selectedServerName string
+	var maxWeight = -1000
+
+	for serverName, tokens := range s.weightsBucket {
+		if tokens > maxWeight {
+			maxWeight = tokens
+			selectedServerName = serverName
+		}
+	}
+
+	if index := slices.IndexFunc(s.servers, func(s *models.Server) bool {
+		return s.Name == selectedServerName
+	}); index >= 0 && index < len(s.servers) {
+		return s.servers[index]
+	} else {
+		return nil
+	}
+}
+
+func (s *WeightedRoundRobinStrategy) ensureWightsBucketValid() {
+	for _, tokens := range s.weightsBucket {
+		if tokens > 0 {
+			return
+		}
+	}
+
+	s.resetBucket()
+}
+
+func (s *WeightedRoundRobinStrategy) UpdateServers(servers []*models.Server) {
+	defer s.Unlock()
+
+	s.Lock()
+
+	s.servers = servers
+
+	s.resetBucket()
+}
+
+func (s *WeightedRoundRobinStrategy) resetBucket() {
+	clear(s.weightsBucket)
+
+	for _, server := range s.servers {
+		s.weightsBucket[server.Name] = server.Weight
+	}
+}
+
+func (s *WeightedRoundRobinStrategy) RequestServed(server *models.Server, request *http.Request) {
+	if tokens, ok := s.weightsBucket[server.Name]; ok {
+		s.weightsBucket[server.Name] = max(tokens-1, 0)
+	}
+}

strategy/weighted_round_robin_test.go

@@ -0,0 +1,53 @@
+package strategy
+
+import (
+	"tayara/go-lb/models"
+	"tayara/go-lb/utils"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWeightedRoundRobin(t *testing.T) {
+	strategy := NewWeightedRoundRobinStrategy()
+
+	servers := []*models.Server{
+		{
+			Name:   "Server1",
+			Url:    "http://localhost:8080",
+			Weight: 2,
+		},
+		{
+			Name:   "Server2",
+			Url:    "http://localhost:8081",
+			Weight: 3,
+		},
+		{
+			Name:   "Server3",
+			Url:    "http://localhost:8082",
+			Weight: 1,
+		},
+	}
+
+	strategy.UpdateServers(servers)
+
+	requestsCount := 6
+
+	actualCounts := map[string]int{}
+
+	for i := 0; i < requestsCount; i++ {
+		server := strategy.Next(nil)
+
+		actualCounts[server.Name] = utils.GetOrDefault(actualCounts, server.Name, 0) + 1
+
+		strategy.RequestServed(server, nil)
+	}
+
+	expectedCounts := map[string]int{
+		servers[0].Name: 2,
+		servers[1].Name: 3,
+		servers[2].Name: 1,
+	}
+
+	assert.Equal(t, expectedCounts, actualCounts)
+}

utils.go

@@ -0,0 +1,15 @@
+package utils
+
+import "os"
+
+func IsFileExist(filePath string) bool {
+	_, err := os.Stat(filePath)
+	return err == nil
+}
+
+func GetOrDefault[K comparable, V any](m map[K]V, key K, defaultValue V) V {
+	if v, ok := m[key]; ok {
+		return v
+	}
+	return defaultValue
+}