trillium_proxy/upstream/
connection_counting.rs

1//! Upstream selectors
2use super::{IntoUpstreamSelector, UpstreamSelector};
3use std::{
4    cmp::Ordering::*,
5    fmt::Debug,
6    ops::{Deref, DerefMut},
7    sync::Arc,
8};
9use trillium::Conn;
10use url::Url;
11
12#[derive(Debug)]
13/// an upstream selector that attempts to send requests to the upstream with the fewest open
14/// connections.
15///
16/// if there are several with the same lowest number of connections, a random upstream is chosen
17/// from them.
18pub struct ConnectionCounting<T>(Vec<(T, Arc<()>)>);
19impl<T> ConnectionCounting<T>
20where
21    T: UpstreamSelector,
22{
23    /// Constructs a new connection counting upstream.
24    pub fn new<I, U>(urls: I) -> Self
25    where
26        I: IntoIterator<Item = U>,
27        U: IntoUpstreamSelector<UpstreamSelector = T>,
28    {
29        Self(
30            urls.into_iter()
31                .map(|u| (u.into_upstream(), Arc::new(())))
32                .collect(),
33        )
34    }
35}
36
37#[allow(dead_code)]
38struct ConnectionCount(Arc<()>);
39
40impl<T> UpstreamSelector for ConnectionCounting<T>
41where
42    T: UpstreamSelector,
43{
44    fn determine_upstream(&self, conn: &mut Conn) -> Option<Url> {
45        let mut current_lowest = usize::MAX;
46        let mut current_selection = vec![];
47        for (u, c) in &self.0 {
48            let current = Arc::strong_count(c);
49            match current.cmp(&current_lowest) {
50                Less => {
51                    current_lowest = current;
52                    current_selection = vec![(u, c)];
53                }
54
55                Equal => {
56                    current_selection.push((u, c));
57                }
58
59                Greater => {}
60            }
61        }
62
63        fastrand::choice(current_selection).and_then(|(u, cc)| {
64            conn.insert_state(ConnectionCount(Arc::clone(cc)));
65            u.determine_upstream(conn)
66        })
67    }
68}
69
70impl<T> Deref for ConnectionCounting<T>
71where
72    T: UpstreamSelector,
73{
74    type Target = [(T, Arc<()>)];
75
76    fn deref(&self) -> &Self::Target {
77        &self.0
78    }
79}
80impl<T> DerefMut for ConnectionCounting<T>
81where
82    T: UpstreamSelector,
83{
84    fn deref_mut(&mut self) -> &mut Self::Target {
85        &mut self.0
86    }
87}
88impl<U, T> Extend<U> for ConnectionCounting<T>
89where
90    T: UpstreamSelector,
91    U: IntoUpstreamSelector<UpstreamSelector = T>,
92{
93    fn extend<I: IntoIterator<Item = U>>(&mut self, iter: I) {
94        self.0
95            .extend(iter.into_iter().map(|i| (i.into_upstream(), Arc::new(()))));
96    }
97}
98
99impl<U, V> FromIterator<U> for ConnectionCounting<V>
100where
101    U: IntoUpstreamSelector<UpstreamSelector = V>,
102    V: UpstreamSelector,
103{
104    fn from_iter<T>(urls: T) -> Self
105    where
106        T: IntoIterator<Item = U>,
107    {
108        Self::new(urls)
109    }
110}
trillium_proxy/upstream/connection_counting.rs

trillium_proxy/upstream/
connection_counting.rs
