Merge pull request #1 from Abdulrahman-Tayara/feature/sticky-session

feature: implement sticky session strategy

Author: Abdulrahman-Tayara

README.md

@@ -27,6 +27,11 @@ Configuring Go-LB is a breeze, it supports both JSON and YAML configuration file
 {
   "port": "load balancer port",
   "strategy": "round_robin | random | least_connections", // default round_robin
+  "strategy_configs": {
+    // Sticky session configs
+    "sticky_session_cookie_name": "example",
+    "sticky_session_ttl_seconds": 100
+  },
   "health_check_interval_seconds": 2,
   "rate_limiter_enabled": true,
   "rate_limit_tokens": 10, // default 10
@@ -53,6 +58,10 @@ Configuring Go-LB is a breeze, it supports both JSON and YAML configuration file
 ```yaml
 port: "load balancer port"
 strategy: "round_robin | random | least_connections" # default round_robin
+strategy_configs:
+  # Sticky session configs
+  sticky_session_cookie_name: "example"
+  sticky_session_ttl_seconds: 100
 health_check_interval_seconds: 2
 rate_limiter_enabled: True
 rate_limit_tokens: 10 # default 10

configs/configs.go

@@ -1,8 +1,10 @@
 package configs
 
 import (
-	"github.com/spf13/viper"
 	"tayara/go-lb/models"
+	"tayara/go-lb/strategy"
+
+	"github.com/spf13/viper"
 )
 
 const (
@@ -12,6 +14,7 @@ const (
 type Configs struct {
 	Port                       string           `mapstructure:"port" json:"port" yaml:"port"`
 	LoadBalancerStrategy       string           `mapstructure:"load_balancer_strategy" json:"load_balancer_strategy" yaml:"load_balancer_strategy"`
+	StrategyConfigs            strategy.Configs `mapstructure:"strategy_configs" json:"strategy_configs" yaml:"strategy_configs"`
 	Servers                    []*models.Server `mapstructure:"servers" json:"servers" yaml:"servers"`
 	Routing                    models.Routing   `mapstructure:"routing" json:"routing" yaml:"routing"`
 	HealthCheckIntervalSeconds int              `mapstructure:"health_check_interval_seconds" json:"health_check_interval_seconds" yaml:"health_check_interval_seconds"`

configs/configs_test.go

@@ -3,6 +3,7 @@ package configs
 import (
 	"reflect"
 	"tayara/go-lb/models"
+	"tayara/go-lb/strategy"
 	"testing"
 )
 
@@ -31,7 +32,11 @@ func TestLoadConfigs(t *testing.T) {
 						HealthUrl: "http://localhost:8080/health",
 					},
 				},
-				LoadBalancerStrategy:       "round_robin",
+				LoadBalancerStrategy: "round_robin",
+				StrategyConfigs: strategy.Configs{
+					StickySessionCookieName: "example",
+					StickySessionTTLSeconds: 100,
+				},
 				HealthCheckIntervalSeconds: 5,
 				RateLimiterEnabled:         true,
 				RateLimitIntervalSeconds:   10,
@@ -75,7 +80,11 @@ func TestLoadConfigs(t *testing.T) {
 						},
 					},
 				},
-				LoadBalancerStrategy:       "round_robin",
+				LoadBalancerStrategy: "round_robin",
+				StrategyConfigs: strategy.Configs{
+					StickySessionCookieName: "example",
+					StickySessionTTLSeconds: 100,
+				},
 				HealthCheckIntervalSeconds: 3,
 				RateLimiterEnabled:         true,
 				RateLimitIntervalSeconds:   10,

configs/configs_test.json

@@ -1,6 +1,10 @@
 {
   "port": "9090",
   "load_balancer_strategy": "round_robin",
+  "strategy_configs": {
+    "sticky_session_cookie_name": "example",
+    "sticky_session_ttl_seconds": 100
+  },
   "health_check_interval_seconds": 5,
   "rate_limiter_enabled": true,
   "rate_limit_tokens": 10,

configs/configs_test.yml

@@ -1,5 +1,8 @@
 port: 8900
 load_balancer_strategy: "round_robin"
+strategy_configs:
+  sticky_session_cookie_name: "example"
+  sticky_session_ttl_seconds: 100
 health_check_interval_seconds: 3
 rate_limiter_enabled: True
 rate_limit_tokens: 10

main.go

@@ -3,8 +3,6 @@ package main
 import (
 	"flag"
 	"fmt"
-	"github.com/pkg/errors"
-	"golang.org/x/exp/slog"
 	"net/http"
 	"slices"
 	"tayara/go-lb/configs"
@@ -13,6 +11,9 @@ import (
 	"tayara/go-lb/ratelimiter/buckettokens"
 	"tayara/go-lb/strategy"
 	"time"
+
+	"github.com/pkg/errors"
+	"golang.org/x/exp/slog"
 )
 
 var (
@@ -38,7 +39,7 @@ func main() {
 
 	slog.Info("configs were loaded", "configs", *cfg)
 
-	selectedStrategy := strategy.GetLoadBalancerStrategy(cfg.LoadBalancerStrategy)
+	selectedStrategy := strategy.GetLoadBalancerStrategy(cfg.LoadBalancerStrategy, cfg.StrategyConfigs)
 
 	loadBalancer := lb.NewLoadBalancer(
 		slices.Clone(cfg.Servers),

strategy/factory.go

@@ -6,7 +6,7 @@ const (
 	LeastConnectionsLoadBalancerStrategy = "least_connections"
 )
 
-func GetLoadBalancerStrategy(strategy string) ILoadBalancerStrategy {
+func GetLoadBalancerStrategy(strategy string, cfg Configs) ILoadBalancerStrategy {
 	switch strategy {
 	case RoundRobinLoadBalancerStrategy:
 		return NewRoundRobinStrategy()

strategy/sticky_session.go

@@ -0,0 +1,78 @@
+package strategy
+
+import (
+	"math/rand"
+	"net/http"
+	"tayara/go-lb/models"
+	"time"
+)
+
+var (
+	defaultSessionName = "lbsession"
+	defaultTTLSeconds  = 300
+)
+
+type StickySessionStrategy struct {
+	cfg Configs
+
+	servers []*models.Server
+}
+
+func NewStickySessionStrategy(cfg Configs) ILoadBalancerStrategy {
+	if cfg.StickySessionCookieName == "" {
+		cfg.StickySessionCookieName = defaultSessionName
+	}
+	if cfg.StickySessionTTLSeconds <= 0 {
+		cfg.StickySessionTTLSeconds = defaultTTLSeconds
+	}
+
+	return &StickySessionStrategy{
+		cfg: cfg,
+	}
+}
+
+func (s *StickySessionStrategy) Next(request *http.Request) *models.Server {
+	cookie, err := request.Cookie(s.cfg.StickySessionCookieName)
+	if err != nil || cookie.Value == "" {
+		cookieValue := generateSessionID()
+		cookie.Value = cookieValue
+		request.AddCookie(&http.Cookie{
+			Name:     s.cfg.StickySessionCookieName,
+			Value:    cookieValue,
+			Expires:  time.Now().Add(time.Second * time.Duration(s.cfg.StickySessionTTLSeconds)),
+			HttpOnly: true,
+		})
+	}
+
+	return s.getServer(cookie.Value)
+}
+
+func (s *StickySessionStrategy) getServer(sessionId string) *models.Server {
+	hash := hashSessionToInt(sessionId)
+	index := hash % len(s.servers)
+	return s.servers[index]
+}
+
+func hashSessionToInt(sessionId string) int {
+	hash := 0
+	for _, char := range sessionId {
+		hash += int(char)
+	}
+	return hash
+}
+
+func (*StickySessionStrategy) RequestServed(server *models.Server, request *http.Request) {
+}
+
+func (*StickySessionStrategy) UpdateServers(servers []*models.Server) {
+}
+
+var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
+
+func generateSessionID() string {
+	b := make([]rune, 20)
+	for i := range b {
+		b[i] = letterRunes[rand.Intn(len(letterRunes))]
+	}
+	return string(b)
+}

strategy/sticky_session_test.go

@@ -0,0 +1,15 @@
+package strategy
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasSessionToInt(t *testing.T) {
+	session := "ABCDEFG"
+
+	num := hashSessionToInt(session)
+
+	assert.Equal(t, 476, num)
+}

strategy/strategy.go

@@ -5,6 +5,11 @@ import (
 	"tayara/go-lb/models"
 )
 
+type Configs struct {
+	StickySessionCookieName string `mapstructure:"sticky_session_cookie_name" json:"sticky_session_cookie_name" yaml:"sticky_session_cookie_name"`
+	StickySessionTTLSeconds int    `mapstructure:"sticky_session_ttl_seconds" json:"sticky_session_ttl_seconds" yaml:"sticky_session_ttl_seconds"`
+}
+
 type ILoadBalancerStrategy interface {
 	Next(request *http.Request) *models.Server