rustlavel_discovery/
resolve.rs1use std::sync::atomic::{AtomicUsize, Ordering};
7
8use rustlavel_core::Result;
9
10use crate::instance::Instance;
11
12pub trait Resolve: Send + Sync {
14 fn instances<'a>(
20 &'a self,
21 service: &'a str,
22 ) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<Instance>>> + Send + 'a>>;
23}
24
25pub struct Static {
31 instances: Vec<Instance>,
32}
33
34impl Static {
35 pub fn new(instances: Vec<Instance>) -> Static {
36 Static { instances }
37 }
38}
39
40impl Resolve for Static {
41 fn instances<'a>(
42 &'a self,
43 service: &'a str,
44 ) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<Instance>>> + Send + 'a>> {
45 let service = service.to_ascii_uppercase();
46 Box::pin(async move {
47 Ok(self.instances.iter().filter(|held| held.service == service).cloned().collect())
48 })
49 }
50}
51
52pub struct Dns {
63 suffix: String,
64 port: u16,
65}
66
67impl Dns {
68 pub fn new(suffix: impl Into<String>, port: u16) -> Dns {
69 Dns { suffix: suffix.into(), port }
70 }
71
72 fn host_for(&self, service: &str) -> String {
73 format!("{}{}", service.to_ascii_lowercase(), self.suffix)
74 }
75}
76
77impl Resolve for Dns {
78 fn instances<'a>(
79 &'a self,
80 service: &'a str,
81 ) -> std::pin::Pin<Box<dyn Future<Output = Result<Vec<Instance>>> + Send + 'a>> {
82 let host = self.host_for(service);
83 let port = self.port;
84 let service = service.to_string();
85
86 Box::pin(async move {
87 Ok(vec![Instance::new(service, host, port)])
92 })
93 }
94}
95
96pub trait Balance: Send + Sync {
98 fn pick<'a>(&self, instances: &'a [Instance]) -> Option<&'a Instance>;
100}
101
102#[derive(Default)]
110pub struct RoundRobin {
111 next: AtomicUsize,
112}
113
114impl RoundRobin {
115 pub fn new() -> RoundRobin {
116 RoundRobin::default()
117 }
118}
119
120impl Balance for RoundRobin {
121 fn pick<'a>(&self, instances: &'a [Instance]) -> Option<&'a Instance> {
122 if instances.is_empty() {
123 return None;
124 }
125 let at = self.next.fetch_add(1, Ordering::Relaxed);
129 instances.get(at % instances.len())
130 }
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136
137 fn instances(count: usize) -> Vec<Instance> {
138 (0..count).map(|n| Instance::new("orders", format!("10.0.0.{n}"), 8080)).collect()
139 }
140
141 #[tokio::test]
142 async fn a_static_resolver_answers_only_for_the_service_asked_for() {
143 let resolver = Static::new(vec![
144 Instance::new("orders", "10.0.0.1", 8080),
145 Instance::new("billing", "10.0.0.2", 8080),
146 ]);
147
148 assert_eq!(resolver.instances("orders").await.unwrap().len(), 1);
149 assert_eq!(resolver.instances("ORDERS").await.unwrap().len(), 1);
151 assert_eq!(resolver.instances("nothing").await.unwrap().len(), 0);
152 }
153
154 #[tokio::test]
157 async fn dns_resolves_to_the_platform_name_rather_than_to_pods() {
158 let resolver = Dns::new(".default.svc.cluster.local", 8080);
159 let found = resolver.instances("orders").await.unwrap();
160
161 assert_eq!(found.len(), 1);
162 assert_eq!(found[0].url(), "http://orders.default.svc.cluster.local:8080");
163 }
164
165 #[test]
166 fn round_robin_gives_each_instance_a_turn() {
167 let pool = instances(3);
168 let balance = RoundRobin::new();
169
170 let picked: Vec<&str> =
171 (0..6).filter_map(|_| balance.pick(&pool)).map(|i| i.host.as_str()).collect();
172 assert_eq!(picked, ["10.0.0.0", "10.0.0.1", "10.0.0.2", "10.0.0.0", "10.0.0.1", "10.0.0.2"]);
173 }
174
175 #[test]
179 fn picking_from_nothing_is_nothing() {
180 assert!(RoundRobin::new().pick(&[]).is_none());
181 }
182
183 #[test]
187 fn the_counter_wraps_rather_than_saturating() {
188 let pool = instances(3);
189 let balance = RoundRobin::new();
190 balance.next.store(usize::MAX, Ordering::Relaxed);
191
192 assert!(balance.pick(&pool).is_some());
193 assert!(balance.pick(&pool).is_some());
194 }
195}