1use crate::config::{Config, Strategy};
12use crate::error::{Error, ProviderError};
13use crate::provider::BoxedProvider;
14use crate::types::ProviderResult;
15use std::collections::HashMap;
16use std::future::Future;
17use std::net::IpAddr;
18use std::pin::Pin;
19use std::task::{Context, Poll};
20use std::time::Instant;
21use tokio::time::timeout;
22
23type BoxFut<'a> = Pin<Box<dyn Future<Output = Result<ProviderResult, ProviderError>> + Send + 'a>>;
25
26pub struct Resolver {
31 config: Config,
32}
33
34impl Resolver {
35 pub fn new(config: Config) -> Self {
37 Self { config }
38 }
39
40 #[inline]
42 fn matching_providers(&self) -> impl Iterator<Item = &BoxedProvider> {
43 self.config
44 .providers
45 .iter()
46 .filter(|p| p.supports_version(self.config.version))
47 }
48
49 fn make_provider_future<'a>(&'a self, provider: &'a BoxedProvider) -> BoxFut<'a> {
52 let provider_name = provider.name().to_string();
53 let protocol = provider.protocol();
54 let start = Instant::now();
55 let fut = provider.get_ip(self.config.version);
56 let timeout_duration = self.config.timeout;
57 let version = self.config.version;
58
59 Box::pin(async move {
60 match timeout(timeout_duration, fut).await {
61 Ok(Ok(ip)) => {
62 if !version.matches(ip) {
63 return Err(ProviderError::message(
64 provider_name,
65 "provider returned unexpected IP version",
66 ));
67 }
68 let latency = start.elapsed();
69 Ok(ProviderResult {
70 ip,
71 provider: provider_name,
72 protocol,
73 latency,
74 })
75 }
76 Ok(Err(e)) => Err(e),
77 Err(_) => Err(ProviderError::message(provider_name, "timeout")),
78 }
79 })
80 }
81
82 pub async fn resolve(&self) -> Result<ProviderResult, Error> {
90 if self.matching_providers().next().is_none() {
91 return Err(Error::NoProvidersForVersion);
92 }
93
94 match self.config.strategy {
95 Strategy::First => self.resolve_first().await,
96 Strategy::Race => self.resolve_race().await,
97 Strategy::Consensus { min_agree } => self.resolve_consensus(min_agree).await,
98 }
99 }
100
101 async fn resolve_first(&self) -> Result<ProviderResult, Error> {
103 let mut errors = Vec::new();
104
105 for provider in self.matching_providers() {
106 match self.make_provider_future(provider).await {
107 Ok(result) => return Ok(result),
108 Err(e) => errors.push(e),
109 }
110 }
111
112 Err(Error::AllProvidersFailed(errors))
113 }
114
115 async fn resolve_race(&self) -> Result<ProviderResult, Error> {
117 let mut futures: Vec<BoxFut<'_>> = self
118 .matching_providers()
119 .map(|p| self.make_provider_future(p))
120 .collect();
121
122 if futures.is_empty() {
125 return Err(Error::NoProvidersForVersion);
126 }
127
128 let mut errors = Vec::new();
129
130 while !futures.is_empty() {
131 let (result, _index, remaining) = select_first(futures).await;
132 futures = remaining;
133
134 match result {
135 Ok(provider_result) => return Ok(provider_result),
136 Err(e) => errors.push(e),
137 }
138 }
139
140 Err(Error::AllProvidersFailed(errors))
141 }
142
143 async fn resolve_consensus(&self, min_agree: usize) -> Result<ProviderResult, Error> {
145 let futures: Vec<BoxFut<'_>> = self
146 .matching_providers()
147 .map(|p| self.make_provider_future(p))
148 .collect();
149
150 if futures.is_empty() {
151 return Err(Error::NoProvidersForVersion);
152 }
153
154 let all_results = join_all_vec(futures).await;
155
156 let mut ip_results: HashMap<IpAddr, Vec<ProviderResult>> = HashMap::new();
157 let mut errors = Vec::new();
158
159 for result in all_results {
160 match result {
161 Ok(pr) => ip_results.entry(pr.ip).or_default().push(pr),
162 Err(e) => errors.push(e),
163 }
164 }
165
166 let mut best: Option<(IpAddr, usize)> = None;
167 for (ip, providers) in &ip_results {
168 if providers.len() >= min_agree {
169 match &best {
170 None => best = Some((*ip, providers.len())),
171 Some((_, current_len)) if providers.len() > *current_len => {
172 best = Some((*ip, providers.len()))
173 }
174 _ => {}
175 }
176 }
177 }
178
179 match best {
180 Some((ip, _)) => {
181 if let Some(providers) = ip_results.remove(&ip) {
182 if let Some(fastest) = providers.into_iter().min_by_key(|p| p.latency) {
183 return Ok(fastest);
184 }
185 }
186 Err(Error::ConsensusNotReached {
187 required: min_agree,
188 got: 0,
189 errors,
190 })
191 }
192 None => {
193 let max_agreement = ip_results.values().map(|v| v.len()).max().unwrap_or(0);
194 Err(Error::ConsensusNotReached {
195 required: min_agree,
196 got: max_agreement,
197 errors,
198 })
199 }
200 }
201 }
202}
203
204async fn select_first<F: Future + Unpin>(mut futures: Vec<F>) -> (F::Output, usize, Vec<F>) {
219 std::future::poll_fn(|cx: &mut Context<'_>| {
220 for (i, fut) in futures.iter_mut().enumerate() {
221 if let Poll::Ready(output) = Pin::new(fut).poll(cx) {
222 futures.swap_remove(i);
223 return Poll::Ready((output, i, std::mem::take(&mut futures)));
224 }
225 }
226 Poll::Pending
227 })
228 .await
229}
230
231async fn join_all_vec<T, F: Future<Output = T> + Unpin>(mut futures: Vec<F>) -> Vec<T> {
240 let total = futures.len();
241 let mut results: Vec<Option<T>> = (0..total).map(|_| None).collect();
242 let mut done = 0;
243
244 std::future::poll_fn(|cx: &mut Context<'_>| {
245 for (i, fut) in futures.iter_mut().enumerate() {
246 if results[i].is_some() {
247 continue;
248 }
249 if let Poll::Ready(output) = Pin::new(fut).poll(cx) {
250 results[i] = Some(output);
251 done += 1;
252 }
253 }
254 if done == total {
255 Poll::Ready(())
256 } else {
257 Poll::Pending
258 }
259 })
260 .await;
261
262 results
263 .into_iter()
264 .map(|r| r.expect("bug: future completed but result slot is empty"))
265 .collect()
266}