pingora_load_balancing/selection/
mod.rs

1// Copyright 2026 Cloudflare, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Backend selection interfaces and algorithms
16
17pub mod algorithms;
18pub mod consistent;
19pub mod weighted;
20
21use super::Backend;
22use std::collections::{BTreeSet, HashSet};
23use std::sync::Arc;
24use weighted::Weighted;
25
26/// [BackendSelection] is the interface to implement backend selection mechanisms.
27pub trait BackendSelection: Sized {
28    /// The [BackendIter] returned from iter() below.
29    type Iter;
30
31    /// The configuration type constructing [BackendSelection]
32    type Config;
33
34    /// Create a [BackendSelection] from a set of backends and the given configuration. The
35    /// default implementation ignores the configuration and simply calls [Self::build]
36    fn build_with_config(backends: &BTreeSet<Backend>, _config: &Self::Config) -> Self {
37        Self::build(backends)
38    }
39
40    /// The function to create a [BackendSelection] implementation.
41    fn build(backends: &BTreeSet<Backend>) -> Self;
42    /// Select backends for a given key.
43    ///
44    /// An [BackendIter] should be returned. The first item in the iter is the first
45    /// choice backend. The user should continue to iterate over it if the first backend
46    /// cannot be used due to its health or other reasons.
47    fn iter(self: &Arc<Self>, key: &[u8]) -> Self::Iter
48    where
49        Self::Iter: BackendIter;
50}
51
52/// An iterator to find the suitable backend
53///
54/// Similar to [Iterator] but allow self referencing.
55pub trait BackendIter {
56    /// Return `Some(&Backend)` when there are more backends left to choose from.
57    fn next(&mut self) -> Option<&Backend>;
58}
59
60/// [SelectionAlgorithm] is the interface to implement selection algorithms.
61///
62/// All [std::hash::Hasher] + [Default] can be used directly as a selection algorithm.
63pub trait SelectionAlgorithm {
64    /// Create a new implementation
65    fn new() -> Self;
66    /// Return the next index of backend. The caller should perform modulo to get
67    /// the valid index of the backend.
68    fn next(&self, key: &[u8]) -> u64;
69}
70
71/// [FNV](https://en.wikipedia.org/wiki/Fowler%E2%80%93Noll%E2%80%93Vo_hash_function) hashing
72/// on weighted backends
73pub type FNVHash = Weighted<fnv::FnvHasher>;
74
75/// Alias of [`FNVHash`] for backwards compatibility until the next breaking change
76#[doc(hidden)]
77pub type FVNHash = Weighted<fnv::FnvHasher>;
78/// Random selection on weighted backends
79pub type Random = Weighted<algorithms::Random>;
80/// Round robin selection on weighted backends
81pub type RoundRobin = Weighted<algorithms::RoundRobin>;
82/// Consistent Ketama hashing on weighted backends
83pub type Consistent = consistent::KetamaHashing;
84
85// TODO: least conn
86
87/// An iterator which wraps another iterator and yields unique items. It optionally takes a max
88/// number of iterations if the wrapped iterator never returns.
89pub struct UniqueIterator<I>
90where
91    I: BackendIter,
92{
93    iter: I,
94    seen: HashSet<u64>,
95    max_iterations: usize,
96    steps: usize,
97}
98
99impl<I> UniqueIterator<I>
100where
101    I: BackendIter,
102{
103    /// Wrap a new iterator and specify the maximum number of times we want to iterate.
104    pub fn new(iter: I, max_iterations: usize) -> Self {
105        Self {
106            iter,
107            max_iterations,
108            seen: HashSet::new(),
109            steps: 0,
110        }
111    }
112
113    pub fn get_next(&mut self) -> Option<Backend> {
114        while let Some(item) = self.iter.next() {
115            if self.steps >= self.max_iterations {
116                return None;
117            }
118            self.steps += 1;
119
120            let hash_key = item.hash_key();
121            if !self.seen.contains(&hash_key) {
122                self.seen.insert(hash_key);
123                return Some(item.clone());
124            }
125        }
126
127        None
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    struct TestIter {
136        seq: Vec<Backend>,
137        idx: usize,
138    }
139    impl TestIter {
140        fn new(input: &[&Backend]) -> Self {
141            Self {
142                seq: input.iter().cloned().cloned().collect(),
143                idx: 0,
144            }
145        }
146    }
147    impl BackendIter for TestIter {
148        fn next(&mut self) -> Option<&Backend> {
149            let idx = self.idx;
150            self.idx += 1;
151            self.seq.get(idx)
152        }
153    }
154
155    #[test]
156    fn unique_iter_max_iterations_is_correct() {
157        let b1 = Backend::new("1.1.1.1:80").unwrap();
158        let b2 = Backend::new("1.0.0.1:80").unwrap();
159        let b3 = Backend::new("1.0.0.255:80").unwrap();
160        let items = [&b1, &b2, &b3];
161
162        let mut all = UniqueIterator::new(TestIter::new(&items), 3);
163        assert_eq!(all.get_next(), Some(b1.clone()));
164        assert_eq!(all.get_next(), Some(b2.clone()));
165        assert_eq!(all.get_next(), Some(b3.clone()));
166        assert_eq!(all.get_next(), None);
167
168        let mut stop = UniqueIterator::new(TestIter::new(&items), 1);
169        assert_eq!(stop.get_next(), Some(b1));
170        assert_eq!(stop.get_next(), None);
171    }
172
173    #[test]
174    fn unique_iter_duplicate_items_are_filtered() {
175        let b1 = Backend::new("1.1.1.1:80").unwrap();
176        let b2 = Backend::new("1.0.0.1:80").unwrap();
177        let b3 = Backend::new("1.0.0.255:80").unwrap();
178        let items = [&b1, &b1, &b2, &b2, &b2, &b3];
179
180        let mut uniq = UniqueIterator::new(TestIter::new(&items), 10);
181        assert_eq!(uniq.get_next(), Some(b1));
182        assert_eq!(uniq.get_next(), Some(b2));
183        assert_eq!(uniq.get_next(), Some(b3));
184    }
185}
pingora_load_balancing/selection/mod.rs

pingora_load_balancing/selection/
mod.rs
