pingora_load_balancing/selection/
mod.rs

1// Copyright 2025 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 {
28    /// The [BackendIter] returned from iter() below.
29    type Iter;
30    /// The function to create a [BackendSelection] implementation.
31    fn build(backends: &BTreeSet<Backend>) -> Self;
32    /// Select backends for a given key.
33    ///
34    /// An [BackendIter] should be returned. The first item in the iter is the first
35    /// choice backend. The user should continue to iterate over it if the first backend
36    /// cannot be used due to its health or other reasons.
37    fn iter(self: &Arc<Self>, key: &[u8]) -> Self::Iter
38    where
39        Self::Iter: BackendIter;
40}
41
42/// An iterator to find the suitable backend
43///
44/// Similar to [Iterator] but allow self referencing.
45pub trait BackendIter {
46    /// Return `Some(&Backend)` when there are more backends left to choose from.
47    fn next(&mut self) -> Option<&Backend>;
48}
49
50/// [SelectionAlgorithm] is the interface to implement selection algorithms.
51///
52/// All [std::hash::Hasher] + [Default] can be used directly as a selection algorithm.
53pub trait SelectionAlgorithm {
54    /// Create a new implementation
55    fn new() -> Self;
56    /// Return the next index of backend. The caller should perform modulo to get
57    /// the valid index of the backend.
58    fn next(&self, key: &[u8]) -> u64;
59}
60
61/// [FNV](https://en.wikipedia.org/wiki/Fowler%E2%80%93Noll%E2%80%93Vo_hash_function) hashing
62/// on weighted backends
63pub type FNVHash = Weighted<fnv::FnvHasher>;
64
65/// Alias of [`FNVHash`] for backwards compatibility until the next breaking change
66#[doc(hidden)]
67pub type FVNHash = Weighted<fnv::FnvHasher>;
68/// Random selection on weighted backends
69pub type Random = Weighted<algorithms::Random>;
70/// Round robin selection on weighted backends
71pub type RoundRobin = Weighted<algorithms::RoundRobin>;
72/// Consistent Ketama hashing on weighted backends
73pub type Consistent = consistent::KetamaHashing;
74
75// TODO: least conn
76
77/// An iterator which wraps another iterator and yields unique items. It optionally takes a max
78/// number of iterations if the wrapped iterator never returns.
79pub struct UniqueIterator<I>
80where
81    I: BackendIter,
82{
83    iter: I,
84    seen: HashSet<u64>,
85    max_iterations: usize,
86    steps: usize,
87}
88
89impl<I> UniqueIterator<I>
90where
91    I: BackendIter,
92{
93    /// Wrap a new iterator and specify the maximum number of times we want to iterate.
94    pub fn new(iter: I, max_iterations: usize) -> Self {
95        Self {
96            iter,
97            max_iterations,
98            seen: HashSet::new(),
99            steps: 0,
100        }
101    }
102
103    pub fn get_next(&mut self) -> Option<Backend> {
104        while let Some(item) = self.iter.next() {
105            if self.steps >= self.max_iterations {
106                return None;
107            }
108            self.steps += 1;
109
110            let hash_key = item.hash_key();
111            if !self.seen.contains(&hash_key) {
112                self.seen.insert(hash_key);
113                return Some(item.clone());
114            }
115        }
116
117        None
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    struct TestIter {
126        seq: Vec<Backend>,
127        idx: usize,
128    }
129    impl TestIter {
130        fn new(input: &[&Backend]) -> Self {
131            Self {
132                seq: input.iter().cloned().cloned().collect(),
133                idx: 0,
134            }
135        }
136    }
137    impl BackendIter for TestIter {
138        fn next(&mut self) -> Option<&Backend> {
139            let idx = self.idx;
140            self.idx += 1;
141            self.seq.get(idx)
142        }
143    }
144
145    #[test]
146    fn unique_iter_max_iterations_is_correct() {
147        let b1 = Backend::new("1.1.1.1:80").unwrap();
148        let b2 = Backend::new("1.0.0.1:80").unwrap();
149        let b3 = Backend::new("1.0.0.255:80").unwrap();
150        let items = [&b1, &b2, &b3];
151
152        let mut all = UniqueIterator::new(TestIter::new(&items), 3);
153        assert_eq!(all.get_next(), Some(b1.clone()));
154        assert_eq!(all.get_next(), Some(b2.clone()));
155        assert_eq!(all.get_next(), Some(b3.clone()));
156        assert_eq!(all.get_next(), None);
157
158        let mut stop = UniqueIterator::new(TestIter::new(&items), 1);
159        assert_eq!(stop.get_next(), Some(b1));
160        assert_eq!(stop.get_next(), None);
161    }
162
163    #[test]
164    fn unique_iter_duplicate_items_are_filtered() {
165        let b1 = Backend::new("1.1.1.1:80").unwrap();
166        let b2 = Backend::new("1.0.0.1:80").unwrap();
167        let b3 = Backend::new("1.0.0.255:80").unwrap();
168        let items = [&b1, &b1, &b2, &b2, &b2, &b3];
169
170        let mut uniq = UniqueIterator::new(TestIter::new(&items), 10);
171        assert_eq!(uniq.get_next(), Some(b1));
172        assert_eq!(uniq.get_next(), Some(b2));
173        assert_eq!(uniq.get_next(), Some(b3));
174    }
175}
pingora_load_balancing/selection/mod.rs

pingora_load_balancing/selection/
mod.rs
