polars_io/cloud/
dns.rs

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

polars_io/cloud/
dns.rs
