Add Weighted Round Robin Algorithm to Load Balancer (#591)

Author: sinadarbouy

cmd/run.go

@@ -851,6 +851,7 @@ var runCmd = &cobra.Command{
 				proxies[configGroupName][configBlockName] = network.NewProxy(
 					runCtx,
 					network.Proxy{
+						Name:                 configBlockName,
 						AvailableConnections: pools[configGroupName][configBlockName],
 						PluginRegistry:       pluginRegistry,
 						HealthCheckPeriod:    cfg.HealthCheckPeriod,
@@ -918,6 +919,7 @@ var runCmd = &cobra.Command{
 					KeyFile:                  cfg.KeyFile,
 					HandshakeTimeout:         cfg.HandshakeTimeout,
 					LoadbalancerStrategyName: cfg.LoadBalancer.Strategy,
+					LoadbalancerRules:        cfg.LoadBalancer.LoadBalancingRules,
 				},
 			)
 

config/config.go

@@ -6,6 +6,7 @@ import (
 	goerrors "errors"
 	"fmt"
 	"log"
+	"math"
 	"os"
 	"reflect"
 	"sort"
@@ -681,16 +682,26 @@ func (c *Config) ValidateGlobalConfig(ctx context.Context) *gerr.GatewayDError {
 	}
 
 	for configGroup := range globalConfig.Servers {
-		if globalConfig.Servers[configGroup] == nil {
+		serverConfig := globalConfig.Servers[configGroup]
+
+		if serverConfig == nil {
 			err := fmt.Errorf("\"servers.%s\" is nil or empty", configGroup)
 			span.RecordError(err)
 			errors = append(errors, gerr.ErrValidationFailed.Wrap(err))
+			continue
 		}
 		if configGroup != strings.ToLower(configGroup) {
 			err := fmt.Errorf(`"servers.%s" is not lowercase`, configGroup)
 			span.RecordError(err)
 			errors = append(errors, gerr.ErrValidationFailed.Wrap(err))
 		}
+
+		// Validate Load Balancing Rules
+		validatelBRulesErrors := ValidateLoadBalancingRules(serverConfig, configGroup, clientConfigGroups)
+		for _, err := range validatelBRulesErrors {
+			span.RecordError(err)
+			errors = append(errors, gerr.ErrValidationFailed.Wrap(err))
+		}
 	}
 
 	if len(globalConfig.Servers) > 1 {
@@ -797,3 +808,113 @@ func generateTagMapping(structs []interface{}, tagMapping map[string]string) {
 		}
 	}
 }
+
+// ValidateLoadBalancingRules validates the load balancing rules in the server configuration.
+func ValidateLoadBalancingRules(
+	serverConfig *Server,
+	configGroup string,
+	clientConfigGroups map[string]map[string]bool,
+) []error {
+	var errors []error
+
+	// Return early if there are no load balancing rules
+	if serverConfig.LoadBalancer.LoadBalancingRules == nil {
+		return errors
+	}
+
+	// Validate each load balancing rule
+	for _, rule := range serverConfig.LoadBalancer.LoadBalancingRules {
+		// Validate the condition of the rule
+		if err := validateRuleCondition(rule.Condition, configGroup); err != nil {
+			errors = append(errors, err)
+		}
+
+		// Validate the distribution of the rule
+		if err := validateDistribution(
+			rule.Distribution,
+			configGroup,
+			rule.Condition,
+			clientConfigGroups,
+		); err != nil {
+			errors = append(errors, err)
+		}
+	}
+
+	return errors
+}
+
+// validateRuleCondition checks if the rule condition is empty for LoadBalancingRules.
+func validateRuleCondition(condition string, configGroup string) error {
+	if condition == "" {
+		err := fmt.Errorf(`"servers.%s.loadBalancer.loadBalancingRules.condition" is nil or empty`, configGroup)
+		return err
+	}
+	return nil
+}
+
+// validateDistribution checks if the distribution list is valid for LoadBalancingRules.
+func validateDistribution(
+	distributionList []Distribution,
+	configGroup string,
+	condition string,
+	clientConfigGroups map[string]map[string]bool,
+) error {
+	// Check if the distribution list is empty
+	if len(distributionList) == 0 {
+		return fmt.Errorf(
+			`"servers.%s.loadBalancer.loadBalancingRules.distribution" is empty`,
+			configGroup,
+		)
+	}
+
+	var totalWeight int
+	for _, distribution := range distributionList {
+		// Validate each distribution entry
+		if err := validateDistributionEntry(distribution, configGroup, condition, clientConfigGroups); err != nil {
+			return err
+		}
+
+		// Check if adding the weight would exceed the maximum integer value
+		if totalWeight > math.MaxInt-distribution.Weight {
+			return fmt.Errorf(
+				`"servers.%s.loadBalancer.loadBalancingRules.%s" total weight exceeds maximum int value`,
+				configGroup,
+				condition,
+			)
+		}
+
+		totalWeight += distribution.Weight
+	}
+
+	return nil
+}
+
+// validateDistributionEntry validates a single distribution entry for LoadBalancingRules.
+func validateDistributionEntry(
+	distribution Distribution,
+	configGroup string,
+	condition string,
+	clientConfigGroups map[string]map[string]bool,
+) error {
+	// Check if the distribution.ProxyName is referenced in the proxy configuration
+	if !clientConfigGroups[configGroup][distribution.ProxyName] {
+		return fmt.Errorf(
+			`"servers.%s.loadBalancer.loadBalancingRules.%s.%s" not referenced in proxy configuration`,
+			configGroup,
+			condition,
+			distribution.ProxyName,
+		)
+	}
+
+	// Ensure that the distribution weight is positive
+	if distribution.Weight <= 0 {
+		return fmt.Errorf(
+			`"servers.%s.loadBalancer.loadBalancingRules.%s.%s.weight" must be positive`,
+			configGroup,
+			condition,
+			distribution.ProxyName,
+		)
+	}
+
+	return nil
+}

config/constants.go

@@ -90,11 +90,12 @@ const (
 	DefaultHealthCheckPeriod = 60 * time.Second // This must match PostgreSQL authentication timeout.
 
 	// Server constants.
-	DefaultListenNetwork        = "tcp"
-	DefaultListenAddress        = "0.0.0.0:15432"
-	DefaultTickInterval         = 5 * time.Second
-	DefaultHandshakeTimeout     = 5 * time.Second
-	DefaultLoadBalancerStrategy = "ROUND_ROBIN"
+	DefaultListenNetwork         = "tcp"
+	DefaultListenAddress         = "0.0.0.0:15432"
+	DefaultTickInterval          = 5 * time.Second
+	DefaultHandshakeTimeout      = 5 * time.Second
+	DefaultLoadBalancerStrategy  = "ROUND_ROBIN"
+	DefaultLoadBalancerCondition = "DEFAULT"
 
 	// Utility constants.
 	DefaultSeed = 1000
@@ -129,6 +130,7 @@ const (
 
 // Load balancing strategies.
 const (
-	RoundRobinStrategy = "ROUND_ROBIN"
-	RANDOMStrategy     = "RANDOM"
+	RoundRobinStrategy         = "ROUND_ROBIN"
+	RANDOMStrategy             = "RANDOM"
+	WeightedRoundRobinStrategy = "WEIGHTED_ROUND_ROBIN"
 )

config/types.go

@@ -96,8 +96,19 @@ type Proxy struct {
 	HealthCheckPeriod time.Duration `json:"healthCheckPeriod" jsonschema:"oneof_type=string;integer" yaml:"healthCheckPeriod"`
 }
 
+type Distribution struct {
+	ProxyName string `json:"proxyName"`
+	Weight    int    `json:"weight"`
+}
+
+type LoadBalancingRule struct {
+	Condition    string         `json:"condition"`
+	Distribution []Distribution `json:"distribution"`
+}
+
 type LoadBalancer struct {
-	Strategy string `json:"strategy"`
+	Strategy           string              `json:"strategy"`
+	LoadBalancingRules []LoadBalancingRule `json:"loadBalancingRules"`
 }
 
 type Server struct {

errors/errors.go

@@ -55,6 +55,7 @@ const (
 	ErrCodePublishAsyncAction
 	ErrCodeLoadBalancerStrategyNotFound
 	ErrCodeNoProxiesAvailable
+	ErrCodeNoLoadBalancerRules
 )
 
 var (
@@ -204,6 +205,10 @@ var (
 		ErrCodeNoProxiesAvailable, "No proxies available to select.", nil,
 	}
 
+	ErrNoLoadBalancerRules = &GatewayDError{
+		ErrCodeNoLoadBalancerRules, "No load balancer rules provided.", nil,
+	}
+
 	// Unwrapped errors.
 	ErrLoggerRequired = errors.New("terminate action requires a logger parameter")
 )

gatewayd.yaml

@@ -83,7 +83,15 @@ servers:
     address: 0.0.0.0:15432
     loadBalancer:
       # Load balancer strategies can be found in config/constants.go
-      strategy: RANDOM # ROUND_ROBIN, RANDOM
+      strategy: ROUND_ROBIN # ROUND_ROBIN, RANDOM, WEIGHTED_ROUND_ROBIN
+      # Optional configuration for strategies that support rules (e.g., WEIGHTED_ROUND_ROBIN)
+      # loadBalancingRules:
+      #   - condition: "DEFAULT" # Currently, only the "DEFAULT" condition is supported
+      #     distribution:
+      #       - proxyName: "writes"
+      #         weight: 70
+      #       - proxyName: "reads"
+      #         weight: 30
     enableTicker: False
     tickInterval: 5s # duration
     enableTLS: False

network/loadbalancer.go

@@ -9,13 +9,35 @@ type LoadBalancerStrategy interface {
 	NextProxy() (IProxy, *gerr.GatewayDError)
 }
 
+// NewLoadBalancerStrategy returns a LoadBalancerStrategy based on the server's load balancer strategy name.
+// If the server's load balancer strategy is weighted round-robin,
+// it selects a load balancer rule before returning the strategy.
+// Returns an error if the strategy is not found or if there are no load balancer rules when required.
 func NewLoadBalancerStrategy(server *Server) (LoadBalancerStrategy, *gerr.GatewayDError) {
 	switch server.LoadbalancerStrategyName {
 	case config.RoundRobinStrategy:
 		return NewRoundRobin(server), nil
 	case config.RANDOMStrategy:
 		return NewRandom(server), nil
+	case config.WeightedRoundRobinStrategy:
+		if server.LoadbalancerRules == nil {
+			return nil, gerr.ErrNoLoadBalancerRules
+		}
+		loadbalancerRule := selectLoadBalancerRule(server.LoadbalancerRules)
+		return NewWeightedRoundRobin(server, loadbalancerRule), nil
 	default:
 		return nil, gerr.ErrLoadBalancerStrategyNotFound
 	}
 }
+
+// selectLoadBalancerRule selects and returns the first load balancer rule that matches the default condition.
+// If no rule matches, it returns the first rule in the list as a fallback.
+func selectLoadBalancerRule(rules []config.LoadBalancingRule) config.LoadBalancingRule {
+	for _, rule := range rules {
+		if rule.Condition == config.DefaultLoadBalancerCondition {
+			return rule
+		}
+	}
+	// Return the first rule as a fallback
+	return rules[0]
+}

network/proxy.go

@@ -36,9 +36,11 @@ type IProxy interface {
 	Shutdown()
 	AvailableConnectionsString() []string
 	BusyConnectionsString() []string
+	GetName() string
 }
 
 type Proxy struct {
+	Name                 string
 	AvailableConnections pool.IPool
 	busyConnections      pool.IPool
 	Logger               zerolog.Logger
@@ -136,6 +138,10 @@ func NewProxy(
 	return &proxy
 }
 
+func (pr *Proxy) GetName() string {
+	return pr.Name
+}
+
 // Connect maps a server connection from the available connection pool to a incoming connection.
 // It returns an error if the pool is exhausted.
 func (pr *Proxy) Connect(conn *ConnWrapper) *gerr.GatewayDError {

network/server.go

@@ -77,6 +77,7 @@ type Server struct {
 	// loadbalancer
 	loadbalancerStrategy     LoadBalancerStrategy
 	LoadbalancerStrategyName string
+	LoadbalancerRules        []config.LoadBalancingRule
 	connectionToProxyMap     map[*ConnWrapper]IProxy
 }
 
@@ -696,6 +697,7 @@ func NewServer(
 		stopServer:               make(chan struct{}),
 		connectionToProxyMap:     make(map[*ConnWrapper]IProxy),
 		LoadbalancerStrategyName: srv.LoadbalancerStrategyName,
+		LoadbalancerRules:        srv.LoadbalancerRules,
 	}
 
 	// Try to resolve the address and log an error if it can't be resolved.