1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133
use rand::{thread_rng, seq::SliceRandom}; use Backend; use std::{ rc::Rc, cell::RefCell }; use std::fmt::Debug; pub trait LoadBalancingAlgorithm: Debug { fn next_available_backend(&mut self, backends: &Vec<Rc<RefCell<Backend>>>) -> Option<Rc<RefCell<Backend>>>; } #[derive(Debug)] pub struct RoundRobinAlgorithm { pub next_backend: u32, } impl LoadBalancingAlgorithm for RoundRobinAlgorithm { fn next_available_backend(&mut self , backends: &Vec<Rc<RefCell<Backend>>>) -> Option<Rc<RefCell<Backend>>> { let res = backends.get(self.next_backend as usize % backends.len()) .map(|backend| (*backend).clone()); self.next_backend = (self.next_backend + 1) % backends.len() as u32; res } } impl RoundRobinAlgorithm { fn new() -> Self { Self { next_backend: 0, } } } #[derive(Debug)] pub struct RandomAlgorithm; impl LoadBalancingAlgorithm for RandomAlgorithm { fn next_available_backend(&mut self, backends: &Vec<Rc<RefCell<Backend>>>) -> Option<Rc<RefCell<Backend>>> { let mut rng = thread_rng(); (*backends).choose(&mut rng) .map(|backend| (*backend).clone()) } } #[derive(Debug)] pub struct LeastConnectionsAlgorithm; impl LoadBalancingAlgorithm for LeastConnectionsAlgorithm { fn next_available_backend(&mut self, backends: &Vec<Rc<RefCell<Backend>>>) -> Option<Rc<RefCell<Backend>>> { backends .iter() .min_by_key(|backend| backend.borrow().active_connections) .map(|backend| (*backend).clone()) } } #[cfg(test)] mod test { use super::*; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use BackendStatus; use retry::{RetryPolicyWrapper, ExponentialBackoffPolicy}; fn create_backend(id: String, connections: Option<usize>) -> Backend { Backend { sticky_id: None, backend_id: id, address: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080), status: BackendStatus::Normal, retry_policy: RetryPolicyWrapper::ExponentialBackoff(ExponentialBackoffPolicy::new(1)), active_connections: connections.unwrap_or(0), failures: 0, load_balancing_parameters: None, backup: false, } } #[test] fn it_should_find_the_backend_with_least_connections() { let backend_with_least_connection = Rc::new(RefCell::new(create_backend("yolo".to_string(), Some(1)))); let backends = vec![ Rc::new(RefCell::new(create_backend("nolo".to_string(), Some(10)))), Rc::new(RefCell::new(create_backend("philo".to_string(), Some(20)))), backend_with_least_connection.clone(), ]; let mut least_connection_algorithm = LeastConnectionsAlgorithm{}; let backend_res = least_connection_algorithm.next_available_backend(&backends).unwrap(); let backend = backend_res.borrow(); assert!(*backend == *backend_with_least_connection.borrow()); } #[test] fn it_shouldnt_find_backend_with_least_connections_when_list_is_empty() { let backends = vec![]; let mut least_connection_algorithm = LeastConnectionsAlgorithm{}; let backend = least_connection_algorithm.next_available_backend(&backends); assert!(backend.is_none()); } #[test] fn it_should_find_backend_with_roundrobin_when_some_backends_were_removed() { let mut backends = vec![ Rc::new(RefCell::new(create_backend("toto".to_string(), None))), Rc::new(RefCell::new(create_backend("voto".to_string(), None))), Rc::new(RefCell::new(create_backend("yoto".to_string(), None))) ]; let mut roundrobin = RoundRobinAlgorithm { next_backend: 1 }; let backend = roundrobin.next_available_backend(&backends); assert_eq!(backend.as_ref(), backends.get(1)); backends.remove(1); let backend2 = roundrobin.next_available_backend(&backends); assert_eq!(backend2.as_ref(), backends.get(0)); } }