From 001d71087da3b5d67e7c24153b78e3a16636decc Mon Sep 17 00:00:00 2001 From: abdulrahman Date: Sat, 20 Jan 2024 18:05:33 +0400 Subject: [PATCH] feature: implement sticky session strategy --- README.md | 9 ++++ configs/configs.go | 5 ++- configs/configs_test.go | 13 +++++- configs/configs_test.json | 4 ++ configs/configs_test.yml | 3 ++ main.go | 7 +-- strategy/factory.go | 2 +- strategy/sticky_session.go | 78 +++++++++++++++++++++++++++++++++ strategy/sticky_session_test.go | 15 +++++++ strategy/strategy.go | 5 +++ 10 files changed, 134 insertions(+), 7 deletions(-) create mode 100644 strategy/sticky_session.go create mode 100644 strategy/sticky_session_test.go diff --git a/README.md b/README.md index 0f69e0c..095dc07 100644 --- a/README.md +++ b/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 diff --git a/configs/configs.go b/configs/configs.go index 9c6f45d..df989de 100644 --- a/configs/configs.go +++ b/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"` diff --git a/configs/configs_test.go b/configs/configs_test.go index 954c12e..73c7912 100644 --- a/configs/configs_test.go +++ b/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, diff --git a/configs/configs_test.json b/configs/configs_test.json index df4d22d..4db2685 100644 --- a/configs/configs_test.json +++ b/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, diff --git a/configs/configs_test.yml b/configs/configs_test.yml index 7d11ebb..627ef36 100644 --- a/configs/configs_test.yml +++ b/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 diff --git a/main.go b/main.go index 6237c73..358c97c 100644 --- a/main.go +++ b/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), diff --git a/strategy/factory.go b/strategy/factory.go index a1a9630..12ab678 100644 --- a/strategy/factory.go +++ b/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() diff --git a/strategy/sticky_session.go b/strategy/sticky_session.go new file mode 100644 index 0000000..bc654fa --- /dev/null +++ b/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) +} diff --git a/strategy/sticky_session_test.go b/strategy/sticky_session_test.go new file mode 100644 index 0000000..7fcd6cd --- /dev/null +++ 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) +} diff --git a/strategy/strategy.go b/strategy/strategy.go index b20b8a0..69dad44 100644 --- a/strategy/strategy.go +++ b/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