feature: adding content based routing option to the load balancer

Author: abdulrahman

README.md

@@ -33,10 +33,12 @@ Configuring Go-LB is a breeze, it supports both JSON and YAML configuration file
   "rate_limit_interval_seconds": 1 // default 2
   "servers": [
     {
+      "name": "server1",
       "url": "http://localhost:8080",
       "health_url": "/health"
     },
     {
+      "name": "server2",
       "url": "http://localhost:8082",
       "health_url": "/health-check"
     }
@@ -55,16 +57,98 @@ health_check_interval_seconds: 2
 rate_limiter_enabled: True
 rate_limit_tokens: 10 # default 10
 rate_limit_interval_seconds: 10 # default 2
+servers:
+  - name: "server1"
+    url: "http://localhost:8080"
+    health_url: "/health"
+  - name: "server2"
+    url: "http://localhost:8082"
+    health_url: "/health-check"
+tls_enabled: true # default false
+tls_cert_file: "/path/on/container/cert.pem"
+tls_key_file: "/path/on/container/key.pem"
+```
+
+## Content Based Routing (CBR) Configuration
+
+### JSON
+```json
+{
+  "port": "load balancer port",
+  "strategy": "round_robin | random | least_connections",
+  "health_check_interval_seconds": 2,
+  "rate_limiter_enabled": true,
+  "rate_limit_tokens": 10,
+  "rate_limit_interval_seconds": 1,
+  "servers": [
+    {
+      "name": "server1",
+      "url": "http://localhost:8080",
+      "health_url": "/health"
+    },
+    {
+      "name": "server2",
+      "url": "http://localhost:8082",
+      "health_url": "/health-check"
+    }
+  ],
+  "tls_enabled": true,
+  "tls_cert_file": "/path/to/cert.pem",
+  "tls_key_file": "/path/to/key.pem",
+  "routing": {
+    "default_server": "server1",
+    "rules": [
+      {
+        "conditions": [
+          {
+            "path_prefix": "/api/v1 (Optional)",
+            "method": "GET | post | Put (Optional)",
+            "headers": { // Optional
+              "MyHeader": "my-value"
+            }
+          }
+        ],
+        "action": {
+          "route_to": "server2"
+        }
+      }
+    ]
+  }
+}
+```
+
+### Yaml
+```yaml
+port: "load balancer port"
+strategy: "round_robin | random | least_connections"
+health_check_interval_seconds: 2
+rate_limiter_enabled: True
+rate_limit_tokens: 10
+rate_limit_interval_seconds: 10
 servers:
   - url: "http://localhost:8080"
     health_url: "/health"
   - url: "http://localhost:8082"
     health_url: "/health-check"
-tls_enabled: true # default false
+tls_enabled: true
 tls_cert_file: "/path/on/container/cert.pem"
 tls_key_file: "/path/on/container/key.pem"
+routing:
+  default_server: "server1"
+  rules:
+    - conditions:
+        - path_prefix: "/api/v1"
+          method: "GET"
+          headers:
+            MyHeader: "my-value"
+      action:
+        route_to: "server2"
 ```
 
+
+
+
+
 ## Getting Started
 
 ### Docker

configs/configs.go

@@ -13,6 +13,7 @@ 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"`
 	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"`
 	RateLimiterEnabled         bool             `mapstructure:"rate_limiter_enabled" json:"rate_limiter_enabled" yaml:"rate_limiter_enabled"`
 	RateLimitTokens            int              `mapstructure:"rate_limit_tokens" json:"rate_limit_tokens" yaml:"rate_limit_tokens"`

configs/configs_test.go

@@ -56,6 +56,25 @@ func TestLoadConfigs(t *testing.T) {
 						HealthUrl: "http://localhost:8081/health",
 					},
 				},
+				Routing: models.Routing{
+					DefaultServer: "server2",
+					Rules: models.RouteRules{
+						{
+							Conditions: []models.RouteCondition{
+								{
+									PathPrefix: "/api/v1",
+									Method:     "GET",
+									Headers: map[string]string{
+										"useragent": "Mobile",
+									},
+								},
+							},
+							Action: models.RouteAction{
+								RouteTo: "server1",
+							},
+						},
+					},
+				},
 				LoadBalancerStrategy:       "round_robin",
 				HealthCheckIntervalSeconds: 3,
 				RateLimiterEnabled:         true,

configs/configs_test.yml

@@ -10,4 +10,14 @@ servers:
   health_url: http://localhost:8080/health
 - name: server2
   url: http://localhost:8081
-  health_url: http://localhost:8081/health
+  health_url: http://localhost:8081/health
+routing:
+  rules:
+    - conditions:
+        - path_prefix: "/api/v1"
+          method: "GET"
+          headers:
+            useragent: "Mobile"
+      action:
+        route_to: server1
+  default_server: server2

lb/load_balancer.go

@@ -26,8 +26,10 @@ type ILoadBalancer interface {
 // -----------------------
 
 type loadBalancer struct {
-	servers  []*models.Server
-	strategy lbs.ILoadBalancerStrategy
+	servers    []*models.Server
+	strategy   lbs.ILoadBalancerStrategy
+	cbr        *models.Routing
+	cbrEnabled bool
 
 	sync.RWMutex
 }
@@ -105,13 +107,47 @@ func (l *loadBalancer) Next(request *http.Request) *models.Server {
 
 	l.RLock()
 
+	if l.cbrEnabled {
+		return l.getServerFromCbr(request)
+	}
+
 	return l.strategy.Next(request)
 }
 
-func NewLoadBalancer(servers []*models.Server, strategy lbs.ILoadBalancerStrategy) ILoadBalancer {
+func (l *loadBalancer) getServerFromCbr(req *http.Request) *models.Server {
+	serverName := l.cbr.Rules.GetRouteTo(&models.RequestProps{
+		Method:  req.Method,
+		Headers: req.Header,
+		Path:    req.URL.Path,
+	})
+
+	if serverName == "" {
+		serverName = l.cbr.DefaultServer
+	}
+
+	return l.serverByName(serverName)
+}
+
+func (l *loadBalancer) serverByName(name string) *models.Server {
+	for _, server := range l.servers {
+		if server.Name == name {
+			return server
+		}
+	}
+
+	return nil
+}
+
+func NewLoadBalancer(
+	servers []*models.Server,
+	cbr *models.Routing,
+	strategy lbs.ILoadBalancerStrategy,
+) ILoadBalancer {
 	strategy.UpdateServers(servers)
 	return &loadBalancer{
-		servers:  servers,
-		strategy: strategy,
+		servers:    servers,
+		strategy:   strategy,
+		cbr:        cbr,
+		cbrEnabled: cbr != nil && len(cbr.Rules) > 0,
 	}
 }

lb/load_balancer_test.go

@@ -18,7 +18,7 @@ func TestLoadBalancer_ServerDown(t *testing.T) {
 		},
 	}
 
-	lb := NewLoadBalancer(slices.Clone(servers), strategy.NewRoundRobinStrategy()).(*loadBalancer)
+	lb := NewLoadBalancer(slices.Clone(servers), nil, strategy.NewRoundRobinStrategy()).(*loadBalancer)
 
 	assert.Equal(t, len(servers), len(lb.servers))
 
@@ -36,7 +36,7 @@ func TestLoadBalancer_ServerUp(t *testing.T) {
 		},
 	}
 
-	lb := NewLoadBalancer(slices.Clone(servers), strategy.NewRoundRobinStrategy()).(*loadBalancer)
+	lb := NewLoadBalancer(slices.Clone(servers), nil, strategy.NewRoundRobinStrategy()).(*loadBalancer)
 
 	assert.Equal(t, len(servers), len(lb.servers))
 

main.go

@@ -42,6 +42,7 @@ func main() {
 
 	loadBalancer := lb.NewLoadBalancer(
 		slices.Clone(cfg.Servers),
+		&cfg.Routing,
 		selectedStrategy,
 	)
 

models/routing.go

@@ -0,0 +1,68 @@
+package models
+
+import (
+	"slices"
+	"strings"
+)
+
+type Routing struct {
+	Rules         RouteRules `mapstructure:"rules" json:"rules" yaml:"rules"`
+	DefaultServer string     `mapstructure:"default_server" json:"default_server" yaml:"default_server"`
+}
+
+type RequestProps struct {
+	Path    string
+	Headers map[string][]string
+	Method  string
+}
+
+type RouteRule struct {
+	Conditions []RouteCondition `mapstructure:"conditions" json:"conditions" yaml:"conditions"`
+	Action     RouteAction      `mapstructure:"action" json:"action" yaml:"action"`
+}
+
+type RouteAction struct {
+	RouteTo string `mapstructure:"route_to" json:"route_to" yaml:"route_to"`
+}
+
+type RouteCondition struct {
+	PathPrefix string            `mapstructure:"path_prefix" json:"path_prefix" yaml:"path_prefix"`
+	Headers    map[string]string `mapstructure:"headers" json:"headers" yaml:"headers"`
+	Method     string            `mapstructure:"method" json:"method" yaml:"method"`
+}
+
+type RouteRules []RouteRule
+
+func (rc *RouteCondition) DoesMatch(req *RequestProps) bool {
+	if rc.PathPrefix != "" && !strings.HasPrefix(req.Path, rc.PathPrefix) {
+		return false
+	}
+	if rc.Method != "" && strings.ToLower(rc.Method) != strings.ToLower(req.Method) {
+		return false
+	}
+	for k, v := range rc.Headers {
+		if !slices.Contains(req.Headers[k], v) {
+			return false
+		}
+	}
+	return true
+}
+
+func (rr *RouteRule) DoesMatch(req *RequestProps) bool {
+	for _, condition := range rr.Conditions {
+		if !condition.DoesMatch(req) {
+			return false
+		}
+	}
+	return true
+}
+
+func (rr *RouteRules) GetRouteTo(req *RequestProps) string {
+	for _, rule := range *rr {
+		if rule.DoesMatch(req) {
+			return rule.Action.RouteTo
+		}
+	}
+
+	return ""
+}