Skip to main content

praxis_core/config/cluster/
load_balancer_strategy.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2024 Praxis Contributors
3
4//! Load-balancing strategy types for upstream clusters.
5
6use serde::{Deserialize, Serialize};
7
8// -----------------------------------------------------------------------------
9// LoadBalancerStrategy
10// -----------------------------------------------------------------------------
11
12/// Load-balancing algorithm used by a cluster.
13#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Serialize)]
14#[serde(untagged)]
15pub enum LoadBalancerStrategy {
16    /// Plain-string strategies: `"round_robin"` or `"least_connections"`.
17    Simple(SimpleStrategy),
18
19    /// Consistent-hash strategy with an optional hash-key header.
20    Parameterised(ParameterisedStrategy),
21}
22
23impl Default for LoadBalancerStrategy {
24    fn default() -> Self {
25        Self::Simple(SimpleStrategy::RoundRobin)
26    }
27}
28
29/// String-serialisable load-balancing strategies.
30#[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq, Serialize)]
31#[serde(rename_all = "snake_case")]
32pub enum SimpleStrategy {
33    /// Cycle through endpoints in order, respecting weights.
34    #[default]
35    RoundRobin,
36
37    /// Pick the endpoint with the fewest active in-flight requests.
38    LeastConnections,
39
40    /// Sample two random endpoints; pick the less loaded one.
41    #[serde(rename = "p2c")]
42    PowerOfTwoChoices,
43}
44
45/// Load-balancing strategies that carry parameters.
46#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Serialize)]
47pub enum ParameterisedStrategy {
48    /// Hash a request attribute to route requests to a stable endpoint.
49    #[serde(rename = "consistent_hash")]
50    ConsistentHash(ConsistentHashOpts),
51}
52
53/// Options for the `consistent_hash` load-balancing strategy.
54#[derive(Debug, Clone, Default, Deserialize, PartialEq, Eq, Serialize)]
55#[serde(deny_unknown_fields)]
56pub struct ConsistentHashOpts {
57    /// Name of the request header to use as the hash key.
58    ///
59    /// Falls back to the request URI path when the header is absent or when this field is `None`.
60    #[serde(default)]
61    pub header: Option<String>,
62}
63
64// -----------------------------------------------------------------------------
65// Tests
66// -----------------------------------------------------------------------------
67
68#[cfg(test)]
69#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
70#[allow(
71    clippy::unwrap_used,
72    clippy::expect_used,
73    clippy::indexing_slicing,
74    clippy::needless_raw_strings,
75    clippy::needless_raw_string_hashes,
76    reason = "tests use unwrap/expect/indexing/raw strings for brevity"
77)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn load_balancer_strategy_defaults_to_round_robin() {
83        assert_eq!(
84            LoadBalancerStrategy::default(),
85            LoadBalancerStrategy::Simple(SimpleStrategy::RoundRobin),
86            "default strategy should be round_robin"
87        );
88    }
89
90    #[test]
91    fn load_balancer_strategy_parses_round_robin() {
92        let yaml = "round_robin";
93        let strategy: LoadBalancerStrategy = serde_yaml::from_str(yaml).unwrap();
94        assert_eq!(
95            strategy,
96            LoadBalancerStrategy::Simple(SimpleStrategy::RoundRobin),
97            "should parse 'round_robin' string"
98        );
99    }
100
101    #[test]
102    fn load_balancer_strategy_parses_least_connections() {
103        let yaml = "least_connections";
104        let strategy: LoadBalancerStrategy = serde_yaml::from_str(yaml).unwrap();
105        assert_eq!(
106            strategy,
107            LoadBalancerStrategy::Simple(SimpleStrategy::LeastConnections),
108            "should parse 'least_connections' string"
109        );
110    }
111
112    #[test]
113    fn load_balancer_strategy_parses_consistent_hash() {
114        let yaml = r#"
115consistent_hash:
116  header: "X-User-Id"
117"#;
118        let strategy: LoadBalancerStrategy = serde_yaml::from_str(yaml).unwrap();
119        assert_eq!(
120            strategy,
121            LoadBalancerStrategy::Parameterised(ParameterisedStrategy::ConsistentHash(ConsistentHashOpts {
122                header: Some("X-User-Id".into()),
123            })),
124            "should parse consistent_hash with header"
125        );
126    }
127
128    #[test]
129    fn load_balancer_strategy_parses_p2c() {
130        let yaml = "p2c";
131        let strategy: LoadBalancerStrategy = serde_yaml::from_str(yaml).unwrap();
132        assert_eq!(
133            strategy,
134            LoadBalancerStrategy::Simple(SimpleStrategy::PowerOfTwoChoices),
135            "should parse 'p2c' string"
136        );
137    }
138
139    #[test]
140    fn consistent_hash_without_header() {
141        let yaml = "consistent_hash: {}";
142        let strategy: LoadBalancerStrategy = serde_yaml::from_str(yaml).unwrap();
143        assert_eq!(
144            strategy,
145            LoadBalancerStrategy::Parameterised(ParameterisedStrategy::ConsistentHash(ConsistentHashOpts {
146                header: None,
147            })),
148            "should parse consistent_hash with no header"
149        );
150    }
151}