polars_io/cloud/
dns.rs

1use std::net::{SocketAddr, ToSocketAddrs};
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::{Arc, LazyLock};
4use std::time::{Duration, Instant};
5
6use futures::stream::{FuturesUnordered, StreamExt};
7use hashbrown::HashMap;
8use polars_core::runtime::ASYNC;
9use reqwest::dns::{Addrs, Name, Resolve, Resolving};
10use tokio::sync::RwLock;
11
12type DynErr = Box<dyn std::error::Error + Send + Sync>;
13
14const DEFAULT_DNS_CACHE_TTL_SECS: u64 = 5;
15
16/// Process-wide DNS cache shared by all `CachingResolver` instances, so resolved
17/// addrs survive client teardown/rebuild (e.g. object-store cache eviction).
18/// In addition, addrs will be shared across object stores that share the same
19/// domain name, e.g., in the case of different buckets in the same region with
20/// path-style hosts in AWS.
21///
22/// The hostname-only key assumes all clients resolve names identically (true
23/// today); per-client config that changes *answers* (e.g. custom nameservers)
24/// would require keying or splitting this cache.
25///
26/// Entries are never evicted. This is ok for object-store endpoints (low cardinality).
27/// Revisit with a size cap if keys become externally driven (e.g. per-bucket
28/// virtual-hosted hosts at scale).
29static DNS_CACHE: LazyLock<RwLock<HashMap<String, CachedAddrs>>> = LazyLock::new(Default::default);
30
31/// Hard-coded DNS TTL cache, as the operating system does not return it with the
32/// calls used. Defaults to AWS TTL.
33pub(crate) fn get_dns_cache_ttl() -> Duration {
34    Duration::from_secs(
35        std::env::var("POLARS_DNS_CACHE_TTL_SECS")
36            .ok()
37            .and_then(|s| s.parse::<u64>().ok())
38            .unwrap_or(DEFAULT_DNS_CACHE_TTL_SECS),
39    )
40}
41
42const DEFAULT_DNS_MAX_STALE_SECS: u64 = 300;
43
44/// Upper limit for serving stale DNS while refresh is happening in the background.
45pub(crate) fn get_dns_max_stale() -> Option<Duration> {
46    let max_stale = Duration::from_secs(
47        std::env::var("POLARS_DNS_MAX_STALE_SECS")
48            .ok()
49            .and_then(|s| s.parse::<u64>().ok())
50            .unwrap_or(DEFAULT_DNS_MAX_STALE_SECS),
51    );
52    if max_stale.is_zero() {
53        None
54    } else {
55        Some(max_stale)
56    }
57}
58
59const DEFAULT_DNS_LOOKUP_ATTEMPTS: u64 = 3;
60
61/// Total DNS lookup attempts (timeout-bounded retries + one final unbounded).
62pub(crate) fn get_dns_lookup_attempts() -> u64 {
63    std::env::var("POLARS_DNS_LOOKUP_ATTEMPTS")
64        .ok()
65        .and_then(|s| s.trim().parse::<u64>().ok())
66        .unwrap_or(DEFAULT_DNS_LOOKUP_ATTEMPTS)
67        .max(1)
68}
69
70const DEFAULT_DNS_ATTEMPT_TIMEOUT_MS: u64 = 500;
71
72/// DNS lookup attempt timeout before hedging kicks in.
73pub(crate) fn get_dns_attempt_timeout() -> Duration {
74    let timeout_ms = std::env::var("POLARS_DNS_ATTEMPT_TIMEOUT_MS")
75        .ok()
76        .and_then(|s| s.trim().parse::<u64>().ok())
77        .unwrap_or(DEFAULT_DNS_ATTEMPT_TIMEOUT_MS);
78
79    Duration::from_millis(timeout_ms)
80}
81
82#[derive(Debug)]
83struct CachedAddrs {
84    addrs: Arc<Vec<SocketAddr>>,
85    fetched_at: Instant,
86    /// True while a background refresh for this host is in flight (single-flight gate).
87    refreshing: Arc<AtomicBool>,
88}
89
90#[derive(Debug, Clone)]
91pub struct DnsResolverConfig {
92    pub ttl: Duration,
93    pub max_stale: Option<Duration>,
94    pub lookup_attempts: u64,
95    pub attempt_timeout: Duration,
96}
97
98impl DnsResolverConfig {
99    pub fn from_env() -> Self {
100        Self {
101            ttl: get_dns_cache_ttl(),
102            max_stale: get_dns_max_stale(),
103            lookup_attempts: get_dns_lookup_attempts(),
104            attempt_timeout: get_dns_attempt_timeout(),
105        }
106    }
107}
108
109/// Shuffle resolver with basic DNS cache. TTL is fixed and set by the calling site.
110/// The resolver serve policy:
111/// - case fresh: serve from cache;
112/// - case expired and within max_stale (> ttl): serve stale + single-flight background refresh;
113/// - case beyond max_stale (or max_stale = None): blocking resolve
114#[derive(Clone, Debug)]
115pub struct CachingResolver {
116    cache: &'static RwLock<HashMap<String, CachedAddrs>>,
117    config: DnsResolverConfig,
118}
119
120impl CachingResolver {
121    pub fn new(config: DnsResolverConfig) -> Self {
122        if polars_config::config().verbose() {
123            let max_stale = config
124                .max_stale
125                .map_or("disabled".to_string(), |m| format!("{}s", m.as_secs()));
126            eprintln!(
127                "[dns_cache] ttl: {}s, max_stale: {}, lookup_attempts: {}, attempt_timeout: {}ms",
128                config.ttl.as_secs(),
129                max_stale,
130                config.lookup_attempts,
131                config.attempt_timeout.as_millis()
132            );
133        }
134
135        Self {
136            cache: &DNS_CACHE,
137            config,
138        }
139    }
140}
141
142impl Resolve for CachingResolver {
143    fn resolve(&self, name: Name) -> Resolving {
144        let cache = self.cache;
145        let DnsResolverConfig {
146            ttl,
147            max_stale,
148            lookup_attempts,
149            attempt_timeout,
150        } = self.config;
151
152        let key = name.as_str().to_string();
153
154        Box::pin(async move {
155            {
156                let read_guard = cache.read().await;
157
158                if let Some(entry) = read_guard.get(&key) {
159                    let age = entry.fetched_at.elapsed();
160                    if age < ttl {
161                        return Ok(shuffle_addrs(&entry.addrs));
162                    }
163
164                    if let Some(max_stale) = max_stale
165                        && age < max_stale
166                    {
167                        // Expired: serve stale immediately, refresh in the background.
168                        // CAS ensures a burst of stale hits spawns exactly one refresh.
169                        if entry
170                            .refreshing
171                            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
172                            .is_ok()
173                        {
174                            let refreshing = entry.refreshing.clone();
175                            ASYNC.spawn(async move {
176                                let result =
177                                    lookup_hedged(&key, lookup_attempts, attempt_timeout).await;
178
179                                // Swap in the new set only on success; on failure keep
180                                // serving the old addrs (next stale hit re-triggers).
181                                if let Ok(addrs) = result {
182                                    let mut write_guard = cache.write().await;
183                                    write_guard.insert(
184                                        key,
185                                        CachedAddrs {
186                                            addrs: Arc::new(addrs),
187                                            fetched_at: Instant::now(),
188                                            refreshing: refreshing.clone(),
189                                        },
190                                    );
191                                }
192                                refreshing.store(false, Ordering::Release);
193                            });
194                        }
195
196                        return Ok(shuffle_addrs(&entry.addrs));
197                    }
198                }
199            }
200
201            // Cache miss or expired
202            let mut write_guard = cache.write().await;
203
204            // Re-check in case the cache has been populated in the meanwhile
205            if let Some(entry) = write_guard.get(&key) {
206                let age = entry.fetched_at.elapsed();
207                if max_stale.is_some_and(|m| age < m) || age < ttl {
208                    return Ok(shuffle_addrs(&entry.addrs));
209                }
210            }
211
212            let addrs = Arc::new(lookup_hedged(&key, lookup_attempts, attempt_timeout).await?);
213
214            write_guard.insert(
215                key,
216                CachedAddrs {
217                    addrs: addrs.clone(),
218                    fetched_at: Instant::now(),
219                    refreshing: Arc::new(AtomicBool::new(false)),
220                },
221            );
222            drop(write_guard);
223
224            Ok(shuffle_addrs(&addrs))
225        })
226    }
227}
228
229/// DNS lookup with hedged attempts once the timeout has been exceeded.
230async fn lookup_hedged(
231    key: &str,
232    lookup_attempts: u64,
233    attempt_timeout: Duration,
234) -> Result<Vec<SocketAddr>, DynErr> {
235    let spawn_lookup = |key: String| {
236        ASYNC.spawn_blocking(move || {
237            (key.as_str(), 0u16)
238                .to_socket_addrs()
239                .map(|it| it.collect::<Vec<_>>())
240        })
241    };
242
243    let mut in_flight = FuturesUnordered::new();
244    in_flight.push(spawn_lookup(key.to_string()));
245    let mut launched = 1;
246
247    let t0 = Instant::now();
248
249    loop {
250        let can_hedge = launched < lookup_attempts;
251
252        tokio::select! {
253            biased;
254
255            completed = in_flight.next() => {
256
257                let completed: Option<Result<Vec<SocketAddr>, DynErr>> = completed.map(|joined| {
258                    joined
259                        .map_err(DynErr::from)
260                        .and_then(|res| res.map_err(DynErr::from))
261                });
262
263                match completed {
264                    // First success wins.
265                    Some(Ok(addrs)) => {
266                        let elapsed = t0.elapsed();
267
268                        if let Some(threshold) = polars_config::config().dns_log_threshold()
269                            && elapsed.gt(&threshold)
270                        {
271                            let display_key = if polars_config::config().verbose_sensitive() {
272                                key
273                            } else {
274                                "<name suppressed>"
275                            };
276                            eprintln!(
277                                "[dns_cache] dns lookup for {} launched {} attempt(s), took {:.1} ms, exceeded threshold of {} ms",
278                                display_key,
279                                launched,
280                                elapsed.as_secs_f64() * 1000.0,
281                                threshold.as_secs_f64() * 1000.0,
282                            )
283                        };
284
285                        return Ok(addrs)},
286                    Some(Err(err)) => {
287                        if in_flight.is_empty() {
288                            if can_hedge {
289                                in_flight.push(spawn_lookup(key.to_string()));
290                                launched += 1;
291                            } else {
292                                return Err(err);
293                            }
294                        }
295                    },
296                    None => unreachable!("in_flight drained while still looping"),
297                }
298            }
299
300            _ = tokio::time::sleep(attempt_timeout), if can_hedge => {
301                in_flight.push(spawn_lookup(key.to_string()));
302                launched += 1;
303            }
304        }
305    }
306}
307
308fn shuffle_addrs(addrs: &Arc<Vec<SocketAddr>>) -> Addrs {
309    let mut indices: Vec<usize> = (0..addrs.len()).collect();
310    fastrand::shuffle(&mut indices);
311    let addrs = addrs.clone();
312    Box::new(indices.into_iter().map(move |i| addrs[i]))
313}
polars_io/cloud/dns.rs

polars_io/cloud/
dns.rs
