trillium_proxy/upstream/
round_robin.rs

1//! Upstream selectors
2use super::{IntoUpstreamSelector, UpstreamSelector};
3use std::{
4    fmt::Debug,
5    ops::{Deref, DerefMut},
6    sync::atomic::{AtomicU8, Ordering},
7};
8use trillium::Conn;
9use url::Url;
10
11#[derive(Debug)]
12/// an upstream selector that dispatches to upstreams in the order that they're specified, wrapping
13pub struct RoundRobin<T>(Vec<T>, AtomicU8);
14impl<T> UpstreamSelector for RoundRobin<T>
15where
16    T: UpstreamSelector,
17{
18    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
19        let wrapping_count = self.1.fetch_add(1, Ordering::SeqCst);
20        let index = (wrapping_count as usize) % self.0.len();
21        self.0.get(index).and_then(|u| u.determine_upstream(conn))
22    }
23}
24
25impl<T> RoundRobin<T>
26where
27    T: UpstreamSelector,
28{
29    /// construct a new round robin
30    pub fn new<I, U>(urls: I) -> Self
31    where
32        I: IntoIterator<Item = U>,
33        U: IntoUpstreamSelector<UpstreamSelector = T>,
34    {
35        Self(
36            urls.into_iter().map(|u| u.into_upstream()).collect(),
37            0.into(),
38        )
39    }
40}
41
42impl<T> Deref for RoundRobin<T>
43where
44    T: UpstreamSelector,
45{
46    type Target = [T];
47
48    fn deref(&self) -> &Self::Target {
49        &self.0
50    }
51}
52impl<T> DerefMut for RoundRobin<T>
53where
54    T: UpstreamSelector,
55{
56    fn deref_mut(&mut self) -> &mut Self::Target {
57        &mut self.0
58    }
59}
60impl<U, T> Extend<U> for RoundRobin<T>
61where
62    T: UpstreamSelector,
63    U: IntoUpstreamSelector<UpstreamSelector = T>,
64{
65    fn extend<I: IntoIterator<Item = U>>(&mut self, iter: I) {
66        self.0.extend(iter.into_iter().map(|i| i.into_upstream()));
67    }
68}
69
70impl<U, V> FromIterator<U> for RoundRobin<V>
71where
72    U: IntoUpstreamSelector<UpstreamSelector = V>,
73    V: UpstreamSelector,
74{
75    fn from_iter<T>(urls: T) -> Self
76    where
77        T: IntoIterator<Item = U>,
78    {
79        Self::new(urls)
80    }
81}
trillium_proxy/upstream/round_robin.rs

trillium_proxy/upstream/
round_robin.rs
