feat: add a sticky session latency-based load balancer (#37)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Session affinity across RPC/REST/gRPC using a sticky latency policy
for more consistent routing.
  * Support for client-provided session keys via the X-PROXY-KEY header.
* Short sticky latency window (short-term stickiness) to stabilize
backend selection.

* **Bug Fixes**
* Fixed REST routing that could report “no servers available” after a
successful selection.

* **Tests**
* Expanded coverage for session affinity, TTL/cleanup, concurrency,
cache semantics, and server updates.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: cloud-j-luna

cmd/main.go

@@ -165,9 +165,11 @@ func runProxy(cfg config.Config) {
 		},
 	}
 	seeder := seed.New(seederCfg, log, rpcListener, restListener, grpcListener)
-	rpcProxyHandler := proxy.NewRPCProxy(rpcListener, cfg.Health, log, proxy.NewLatencyBased(log))
-	restProxyHandler := proxy.NewRestProxy(restListener, cfg.Health, log, proxy.NewLatencyBased(log))
-	grpcProxyHandler := proxy.NewGRPCProxy(grpcListener, log, proxy.NewLatencyBased(log))
+
+	blockTime := 6 * time.Second
+	rpcProxyHandler := proxy.NewRPCProxy(rpcListener, cfg.Health, log, proxy.NewStickyLatencyBased(log, blockTime))
+	restProxyHandler := proxy.NewRestProxy(restListener, cfg.Health, log, proxy.NewStickyLatencyBased(log, blockTime))
+	grpcProxyHandler := proxy.NewGRPCProxy(grpcListener, log, proxy.NewStickyLatencyBased(log, blockTime))
 
 	ctx, proxyCtxCancel := context.WithCancel(context.Background())
 	defer proxyCtxCancel()
@@ -264,7 +266,7 @@ func runProxy(cfg config.Config) {
 }
 
 func main() {
-	var v = viper.New()
+	v := viper.New()
 
 	if err := NewRootCmd(v).Execute(); err != nil {
 		log.Fatalf("failed to execute command: %v", err)

internal/proxy/balancer.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"log/slog"
 	"math/rand"
+	"net/http"
 	"sync"
 	"time"
 )
@@ -11,14 +12,17 @@ import (
 // used as toleration when comparing floats.
 const epsilon = 1e-9
 
+// ProxyKeyHeader is the HTTP header used for sticky session identification
+const ProxyKeyHeader = "X-PROXY-KEY"
+
 // LoadBalancer is an interface for load balancing algorithms. It provides
 // methods to update the list of available servers and to select the next
 // server to be used.
 type LoadBalancer interface {
 	// Update updates the list of available servers.
 	Update([]*Server)
-	// Next returns the next server to be used based on the load balancing algorithm.
-	Next() *Server
+	// NextServer returns the next server to be used based on the load balancing algorithm.
+	NextServer(*http.Request) *Server
 }
 
 // RoundRobin is a simple load balancer that distributes incoming requests
@@ -41,9 +45,9 @@ func NewRoundRobin(log *slog.Logger) *RoundRobin {
 	}
 }
 
-// Next returns the next server to be used based on the round-robin algorithm.
+// NextServer returns the next server to be used based on the round-robin algorithm.
 // If the selected server is unhealthy, it will recursively try the next server.
-func (rr *RoundRobin) Next() *Server {
+func (rr *RoundRobin) NextServer(r *http.Request) *Server {
 	rr.mu.Lock()
 	if len(rr.servers) == 0 {
 		return nil
@@ -56,7 +60,7 @@ func (rr *RoundRobin) Next() *Server {
 		return server
 	}
 	rr.log.Warn("server is unhealthy, trying next", "name", server.name)
-	return rr.Next()
+	return rr.NextServer(r)
 }
 
 // Update updates the list of available servers.
@@ -96,15 +100,15 @@ func NewLatencyBased(log *slog.Logger) *LatencyBased {
 	}
 }
 
-// Next returns the next server based on the weighted random selection,
+// NextServer returns the next server based on the weighted random selection,
 // where the weight is determined by the latency Rate of each server. The cumulative
 // approach is used to select a server, effectively creating a "range" for each
 // server in the interval [0, 1]. For example, if the rates are [0.5, 0.3, 0.2],
 // the ranges would be: Server 1: [0, 0.5), Server 2: [0.5, 0.8), Server 3: [0.8, 1).
 // The random number will fall into one of these ranges, effectively selecting
 // a server based on its latency rate. This approach works regardless of the order of
 // the servers, so there's no need to sort them based on latency or rate.
-func (rr *LatencyBased) Next() *Server {
+func (rr *LatencyBased) NextServer(_ *http.Request) *Server {
 	rr.mu.Lock()
 	defer rr.mu.Unlock()
 
@@ -159,3 +163,176 @@ func (rr *LatencyBased) Update(servers []*Server) {
 		rr.servers[i].Rate /= totalInverse
 	}
 }
+
+// StickyLatencyBased is a load balancer that combines session affinity with latency-based routing.
+// It embeds LatencyBased to reuse latency calculation and server management functionality,
+// while adding session stickiness using industry-standard headers and cookies.
+// Warning: This load balancer type is not effective if running alongside other replicas as
+// the state is not shared between replicas.
+type StickyLatencyBased struct {
+	// LatencyBased provides the core latency-based selection functionality
+	*LatencyBased
+	// sessionMap maps session identifiers to server references for sticky sessions.
+	sessionMap map[string]*Server
+	// sessionMu is a separate mutex for session-specific operations to avoid lock contention
+	sessionMu sync.RWMutex
+	// sessionTimeout defines how long sessions are kept in memory.
+	sessionTimeout time.Duration
+	// sessionCleanupTicker periodically cleans up expired sessions.
+	sessionCleanupTicker *time.Ticker
+	// sessionTimestamps tracks when sessions were last accessed.
+	sessionTimestamps map[string]time.Time
+}
+
+// NewStickyLatencyBased returns a new StickyLatencyBased load balancer instance.
+// It embeds a LatencyBased load balancer and adds session management functionality.
+func NewStickyLatencyBased(log *slog.Logger, sessionTimeout time.Duration) *StickyLatencyBased {
+	if sessionTimeout == 0 {
+		sessionTimeout = 30 * time.Minute // Default session timeout
+	}
+
+	slb := &StickyLatencyBased{
+		LatencyBased:      NewLatencyBased(log),
+		sessionMap:        make(map[string]*Server),
+		sessionTimestamps: make(map[string]time.Time),
+		sessionTimeout:    sessionTimeout,
+	}
+
+	slb.sessionCleanupTicker = time.NewTicker(5 * time.Minute)
+	go slb.cleanupExpiredSessions()
+
+	return slb
+}
+
+// NextServer returns the next server based on session affinity and latency.
+// It first checks for existing session identifiers in headers or cookies,
+// then falls back to the embedded LatencyBased selection for new sessions.
+func (slb *StickyLatencyBased) NextServer(req *http.Request) *Server {
+	if req == nil {
+		slb.log.Warn("provided request is nil")
+		return slb.LatencyBased.NextServer(req)
+	}
+
+	slb.LatencyBased.mu.Lock()
+	if len(slb.LatencyBased.servers) == 0 {
+		slb.LatencyBased.mu.Unlock()
+		return nil
+	}
+	slb.LatencyBased.mu.Unlock()
+
+	sessionID := slb.extractSessionID(req)
+
+	if sessionID != "" {
+		slb.sessionMu.RLock()
+		if server, exists := slb.sessionMap[sessionID]; exists {
+			lastAccessed := slb.sessionTimestamps[sessionID]
+			if time.Since(lastAccessed) > slb.sessionTimeout {
+				slb.sessionMu.RUnlock()
+
+				slb.sessionMu.Lock()
+				delete(slb.sessionMap, sessionID)
+				delete(slb.sessionTimestamps, sessionID)
+				slb.sessionMu.Unlock()
+
+				slb.log.Info("session timed out, removed",
+					"session_id", sessionID,
+					"last_accessed", lastAccessed)
+			} else {
+				slb.sessionMu.RUnlock()
+				if server != nil && server.Healthy() { // serve only healthy cached servers.
+					slb.sessionMu.Lock()
+					slb.sessionTimestamps[sessionID] = time.Now()
+					slb.sessionMu.Unlock()
+					return server
+				}
+				slb.sessionMu.Lock()
+				delete(slb.sessionMap, sessionID)
+				delete(slb.sessionTimestamps, sessionID)
+				slb.sessionMu.Unlock()
+			}
+		} else {
+			slb.sessionMu.RUnlock()
+		}
+	}
+
+	server := slb.LatencyBased.NextServer(req)
+
+	if server != nil && sessionID != "" {
+		slb.sessionMu.Lock()
+		slb.sessionMap[sessionID] = server
+		slb.sessionTimestamps[sessionID] = time.Now()
+		slb.sessionMu.Unlock()
+
+		slb.log.Debug("created new sticky session",
+			"session_id", sessionID,
+			"server", server.name)
+	}
+
+	return server
+}
+
+// extractSessionID extracts session identifier from HTTP request.
+// It only checks for the X-PROXY-KEY header. If not provided, returns empty string
+// which will cause the load balancer to use normal latency-based selection.
+func (slb *StickyLatencyBased) extractSessionID(req *http.Request) string {
+	return req.Header.Get(ProxyKeyHeader)
+}
+
+// Update updates the list of available servers using the embedded LatencyBased functionality
+// and cleans up session mappings for servers that no longer exist.
+func (slb *StickyLatencyBased) Update(servers []*Server) {
+	slb.LatencyBased.Update(servers)
+	slb.pruneInvalidSessions(servers)
+}
+
+// pruneInvalidSessions removes session mappings for servers that no longer exist
+// in the provided server list.
+func (slb *StickyLatencyBased) pruneInvalidSessions(servers []*Server) {
+	// Create a lookup map from the input servers for O(1) lookups
+	serverExists := make(map[string]bool, len(servers))
+	for _, s := range servers {
+		if s != nil {
+			serverExists[s.name] = true
+		}
+	}
+
+	slb.sessionMu.Lock()
+	defer slb.sessionMu.Unlock()
+
+	for sid, srv := range slb.sessionMap {
+		if srv == nil || !serverExists[srv.name] {
+			delete(slb.sessionMap, sid)
+			delete(slb.sessionTimestamps, sid)
+			if srv != nil {
+				slb.log.Debug(
+					"removed sticky session for deleted server",
+					"session_id", sid,
+					"server", srv.name,
+				)
+			}
+		}
+	}
+}
+
+// cleanupExpiredSessions runs in a background goroutine to clean up expired sessions
+func (slb *StickyLatencyBased) cleanupExpiredSessions() {
+	for range slb.sessionCleanupTicker.C {
+		slb.sessionMu.Lock()
+		now := time.Now()
+		for sessionID, timestamp := range slb.sessionTimestamps {
+			if now.Sub(timestamp) > slb.sessionTimeout {
+				delete(slb.sessionMap, sessionID)
+				delete(slb.sessionTimestamps, sessionID)
+				slb.log.Debug("cleaned up expired session", "session_id", sessionID)
+			}
+		}
+		slb.sessionMu.Unlock()
+	}
+}
+
+// Stop stops the cleanup ticker and releases resources
+func (slb *StickyLatencyBased) Stop() {
+	if slb.sessionCleanupTicker != nil {
+		slb.sessionCleanupTicker.Stop()
+	}
+}