loadwise_core/strategy/
mod.rs

1//! Load balancing strategies and the [`Strategy`] trait.
2//!
3//! Built-in strategies:
4//!
5//! | Strategy                  | Requires        | Algorithm                              |
6//! |---------------------------|-----------------|----------------------------------------|
7//! | [`RoundRobin`]            | —               | Sequential cycling                     |
8//! | [`WeightedRoundRobin`]    | `Weighted+Node` | Nginx smooth weighted round-robin      |
9//! | [`Random`]                | —               | Uniform random                         |
10//! | [`WeightedRandom`]        | `Weighted`      | Proportional random                    |
11//! | [`LeastLoad`]             | `LoadMetric`    | Pick lowest `load_score()`             |
12//! | [`PowerOfTwoChoices`]     | `LoadMetric`    | Random 2, pick lower load              |
13//! | [`ConsistentHash`]        | `Node`          | Virtual-node ring (cached)             |
14//! | [`MostAvailable`]         | `RateMetric`    | Pick highest remaining, epsilon tie-break |
15//!
16//! Combinators: [`WithFallback`], [`FallbackChain`].
17
18mod chain;
19mod consistent_hash;
20mod least_load;
21mod most_available;
22mod p2c;
23mod random;
24mod round_robin;
25mod weighted_round_robin;
26
27pub use chain::{FallbackChain, WithFallback};
28pub use consistent_hash::ConsistentHash;
29pub use least_load::LeastLoad;
30pub use most_available::MostAvailable;
31pub use p2c::PowerOfTwoChoices;
32pub use random::{Random, WeightedRandom};
33pub use round_robin::RoundRobin;
34pub use weighted_round_robin::WeightedRoundRobin;
35
36use std::collections::HashMap;
37use bon::Builder;
38
39/// Context provided to strategies during node selection.
40///
41/// # Examples
42///
43/// ```
44/// # extern crate loadwise_core as loadwise;
45/// use loadwise::SelectionContext;
46///
47/// // Empty context (most strategies work fine without it).
48/// let ctx = SelectionContext::default();
49///
50/// // With a hash key for consistent hashing.
51/// let ctx = SelectionContext::builder()
52///     .hash_key(42)
53///     .build();
54///
55/// // With estimated request cost and tags.
56/// let ctx = SelectionContext::builder()
57///     .estimated_cost(150.0)
58///     .tags([("model".into(), "gpt-4".into())].into())
59///     .build();
60///
61/// assert_eq!(ctx.tag("model"), Some("gpt-4"));
62///
63/// // Exclude specific indices (e.g., retry after failure).
64/// let ctx = SelectionContext::builder()
65///     .exclude(vec![0, 2])
66///     .build();
67/// assert!(ctx.is_excluded(0));
68/// assert!(!ctx.is_excluded(1));
69/// ```
70#[derive(Debug, Clone, Default, Builder)]
71pub struct SelectionContext {
72    /// Estimated cost/weight of this request (e.g., token count).
73    pub estimated_cost: Option<f64>,
74    /// Hash key for consistent hashing strategies.
75    pub hash_key: Option<u64>,
76    /// Arbitrary tags for routing decisions. `None` by default to avoid allocation.
77    pub tags: Option<HashMap<String, String>>,
78    /// Indices to skip during selection (for retry scenarios).
79    /// `None` by default to avoid allocation.
80    pub exclude: Option<Vec<usize>>,
81}
82
83impl SelectionContext {
84    /// Look up a tag value by key.
85    pub fn tag(&self, key: &str) -> Option<&str> {
86        self.tags.as_ref()?.get(key).map(|s| s.as_str())
87    }
88
89    /// Returns `true` if the given index should be skipped.
90    ///
91    /// Uses linear search — O(n) in the size of the exclude list.
92    /// Fine for typical retry scenarios (1–3 entries); for larger lists
93    /// consider pre-filtering candidates instead.
94    pub fn is_excluded(&self, idx: usize) -> bool {
95        self.exclude.as_ref().is_some_and(|e| e.contains(&idx))
96    }
97}
98
99/// A load balancing strategy that selects a node from a list of candidates.
100///
101/// Returns the **index** of the selected node within the `candidates` slice.
102/// Implementations must use interior mutability for any state (e.g., `AtomicUsize`).
103///
104/// # Examples
105///
106/// ```
107/// # extern crate loadwise_core as loadwise;
108/// use loadwise::{Strategy, SelectionContext};
109/// use loadwise::strategy::RoundRobin;
110///
111/// let rr = RoundRobin::new();
112/// let nodes = ["a", "b", "c"];
113/// let ctx = SelectionContext::default();
114///
115/// assert_eq!(rr.select(&nodes, &ctx), Some(0));
116/// assert_eq!(rr.select(&nodes, &ctx), Some(1));
117/// assert_eq!(rr.select(&nodes, &ctx), Some(2));
118/// assert_eq!(rr.select(&nodes, &ctx), Some(0)); // wraps around
119/// ```
120pub trait Strategy<N>: Send + Sync {
121    /// Pick the next node. Returns `None` when `candidates` is empty or all
122    /// candidates are excluded via [`SelectionContext::exclude`].
123    fn select(&self, candidates: &[N], ctx: &SelectionContext) -> Option<usize>;
124}
125
126// ── internal helpers ──
127
128use std::hash::{DefaultHasher, Hash, Hasher};
129use crate::Node;
130
131/// Order-sensitive fingerprint of a candidate set based on node IDs.
132///
133/// Used by strategies that cache internal state (e.g., hash ring, WRR weights)
134/// to detect when the candidate list has changed.
135///
136/// NOTE: this is O(n) per call. For very large candidate lists, consider
137/// letting the caller notify the strategy of changes instead of probing
138/// every time. Good enough for typical pool sizes (10–100 nodes).
139pub(crate) fn node_set_fingerprint<N: Node>(candidates: &[N]) -> u64 {
140    let mut hasher = DefaultHasher::new();
141    candidates.len().hash(&mut hasher);
142    for node in candidates {
143        node.id().hash(&mut hasher);
144    }
145    hasher.finish()
146}
loadwise_core/strategy/mod.rs

loadwise_core/strategy/
mod.rs
