From 8f4d4f43ddc3c041ad417841ac43414ae2444832 Mon Sep 17 00:00:00 2001 From: cfal Date: Sun, 25 Oct 2026 00:20:20 +0000 Subject: [PATCH] Harden native DNS coalescing wakeups and cancellation tests --- src/resolver.rs | 117 +++++++++--------- src/resolver/native_tests.rs | 225 +++++++++++++++++++++++++---------- 2 files changed, 227 insertions(+), 115 deletions(-) diff --git a/src/resolver.rs b/src/resolver.rs index 13a43dde..8a816694 100644 --- a/src/resolver.rs +++ b/src/resolver.rs @@ -471,13 +471,15 @@ pub async fn resolve_location( } /// Native resolver with application-level caching. -/// The default for servers without DNS configuration. OS lookups expose no DNS +/// The default for servers without configured DNS servers. OS lookups expose no DNS /// TTL; successful results are reused for at most 60 seconds by default, measured /// from lookup completion. Errors and empty results are never cached. /// /// At most 1,024 callers participate in coalesced misses (owners plus subscribers). /// Excess callers perform independent lookups, preserving existing failure /// semantics without adding an unbounded admission queue or detached tasks. +/// Cancelling an owner makes surviving subscribers start a new lookup, which +/// may exceed their remaining deadlines. Already-running OS calls cannot be stopped. pub struct CachingNativeResolver { state: Arc>, participants: Arc, @@ -493,14 +495,14 @@ struct NativeResolverState { } struct CachedResolveResult { - timestamp: tokio::time::Instant, + completed_at: tokio::time::Instant, addrs: Arc<[SocketAddr]>, } impl NativeResolverState { fn cached(&mut self, address: &Address, max_age: Duration) -> Option> { if let Some(cached) = self.cache.get(address) - && cached.timestamp.elapsed() < max_age + && cached.completed_at.elapsed() < max_age { return Some(cached.addrs.clone()); } @@ -515,6 +517,11 @@ struct NativeLookupOwner { registered: bool, } +enum NativeLookupRole { + Owner(NativeLookupOwner), + Subscriber(broadcast::Receiver), +} + impl NativeLookupOwner { fn publish(mut self, result: &NativeLookupResult) { let mut state = self.state.lock(); @@ -522,16 +529,18 @@ impl NativeLookupOwner { state.cache.put( self.address.clone(), CachedResolveResult { - timestamp: tokio::time::Instant::now(), + completed_at: tokio::time::Instant::now(), addrs: addrs.clone(), }, ); } - // Cache insertion, publication, and removal are atomic with subscription. + // Cache insertion and flight removal are atomic with subscription. if self.registered { let sender = state.inflight.remove(&self.address).unwrap(); - let _ = sender.send(result.clone()); + // A panicking waker must not let Drop remove a replacement flight. self.registered = false; + drop(state); + let _ = sender.send(result.clone()); } } } @@ -539,8 +548,9 @@ impl NativeLookupOwner { impl Drop for NativeLookupOwner { fn drop(&mut self) { if self.registered { - // Closing the channel wakes subscribers to elect a replacement owner. - self.state.lock().inflight.remove(&self.address); + // Closing outside the lock wakes subscribers to elect a replacement owner. + let sender = self.state.lock().inflight.remove(&self.address); + drop(sender); } } } @@ -603,16 +613,16 @@ impl Default for CachingNativeResolver { impl Resolver for CachingNativeResolver { fn resolve_location(&self, location: &NetLocation) -> ResolveFuture { let max_age = Duration::from_secs(self.result_timeout_secs); - let port = location.port(); + let caller_port = location.port(); if let Some(addrs) = self.state.lock().cached(location.address(), max_age) { - let addrs = native_addrs_with_port(&addrs, port); + let addrs = native_addrs_with_port(&addrs, caller_port); return Box::pin(async move { Ok(addrs) }); } // Tokio's (host, u16) lookup delegates to std, which passes a null service // to getaddrinfo and sets ports afterward. Only the exact host is a key; // socket candidates retain order, duplicates, IPv6 scope IDs and flowinfo. - let location = NetLocation::new(location.address().clone(), 0); + let query_location = NetLocation::new(location.address().clone(), 0); let state = self.state.clone(); let participants = self.participants.clone(); let lookup = self.lookup.clone(); @@ -620,63 +630,64 @@ impl Resolver for CachingNativeResolver { Box::pin(async move { let permit = participants.try_acquire().ok(); loop { - let (owner, subscriber) = { + let role = { let mut locked = state.lock(); // Rechecks after future creation or cancellation of an owner. - if let Some(addrs) = locked.cached(location.address(), max_age) { - return Ok(native_addrs_with_port(&addrs, port)); + if let Some(addrs) = locked.cached(query_location.address(), max_age) { + return Ok(native_addrs_with_port(&addrs, caller_port)); } - if let Some(sender) = locked.inflight.get(location.address()) + if let Some(sender) = locked.inflight.get(query_location.address()) && permit.is_some() { - (None, Some(sender.subscribe())) + NativeLookupRole::Subscriber(sender.subscribe()) } else { let registered = permit.is_some(); if registered { let (sender, _) = broadcast::channel(1); - locked.inflight.insert(location.address().clone(), sender); + locked + .inflight + .insert(query_location.address().clone(), sender); } - ( - Some(NativeLookupOwner { - state: state.clone(), - address: location.address().clone(), - registered, - }), - None, - ) + NativeLookupRole::Owner(NativeLookupOwner { + state: state.clone(), + address: query_location.address().clone(), + registered, + }) } }; - if let Some(mut subscriber) = subscriber { - match subscriber.recv().await { - Ok(result) => { - return result - .map(|addrs| native_addrs_with_port(&addrs, port)) - .map_err(native_lookup_error); - } + let result = match role { + NativeLookupRole::Subscriber(mut subscriber) => match subscriber.recv().await { + Ok(result) => result, Err(_) => continue, + }, + NativeLookupRole::Owner(owner) => { + let result = lookup + .resolve_location(&query_location) + .await + .and_then(|mut addrs| { + addrs.retain(|addr| !addr.ip().is_unspecified()); + if addrs.is_empty() { + Err(std::io::Error::other(format!( + "DNS lookup returned no addresses for {}", + query_location.address() + ))) + } else { + Ok(Arc::from(addrs)) + } + }) + .map_err(Arc::new); + owner.publish(&result); + result } - } - - let result = lookup - .resolve_location(&location) - .await - .and_then(|mut addrs| { - addrs.retain(|addr| !addr.ip().is_unspecified()); - if addrs.is_empty() { - Err(std::io::Error::other(format!( - "DNS lookup returned no addresses for {}", - location.address() - ))) - } else { - Ok(Arc::from(addrs)) - } - }) - .map_err(Arc::new); - owner.unwrap().publish(&result); - return result - .map(|addrs| native_addrs_with_port(&addrs, port)) - .map_err(native_lookup_error); + }; + let addrs = result.map_err(native_lookup_error)?; + let addrs = native_addrs_with_port(&addrs, caller_port); + debug!( + "CachingNativeResolver resolved {}:{caller_port} -> {addrs:?}", + query_location.address() + ); + return Ok(addrs); } }) } diff --git a/src/resolver/native_tests.rs b/src/resolver/native_tests.rs index 18641736..8e4fb685 100644 --- a/src/resolver/native_tests.rs +++ b/src/resolver/native_tests.rs @@ -88,6 +88,10 @@ fn assert_idle(resolver: &CachingNativeResolver, lookup: &FakeLookup) { assert_eq!(lookup.active.load(Ordering::Relaxed), 0); } +fn ready_result(future: ResolveFuture) -> LookupReply { + future.now_or_never().expect("lookup unexpectedly pending") +} + async fn complete_lookup(resolver: &CachingNativeResolver, lookup: &FakeLookup, port: u16) { let mut future = resolver.resolve_location(&location(port)); assert!(futures::poll!(&mut future).is_pending()); @@ -119,19 +123,13 @@ async fn cold_misses_coalesce_across_ports_and_preserve_complete_candidates() { let owner = futures.remove(0); assert_eq!(owner.await.unwrap(), candidates()); assert!(resolver.state.lock().inflight.is_empty()); - assert_eq!( - resolver - .state - .lock() - .cache - .peek(location(0).address()) - .unwrap() - .addrs - .as_ref(), - candidates() - ); + { + let state = resolver.state.lock(); + let cached = state.cache.peek(location(0).address()).unwrap(); + assert_eq!(cached.addrs.as_ref(), candidates()); + } // A new caller sees the cache even before subscribers consume the publication. - let mut hit = resolver.resolve_location(&location(8080)).await.unwrap(); + let mut hit = ready_result(resolver.resolve_location(&location(8080))).unwrap(); assert_eq!(hit, native_addrs_with_port(&candidates(), 8080)); hit.clear(); for (future, &port) in futures.into_iter().zip(&ports[1..]) { @@ -152,7 +150,7 @@ async fn unpolled_miss_rechecks_cache_and_allocates_no_inflight_state() { assert_idle(&resolver, &lookup); complete_lookup(&resolver, &lookup, 443).await; assert_eq!( - unpolled.await.unwrap(), + ready_result(unpolled).unwrap(), native_addrs_with_port(&candidates(), 80) ); assert_eq!(lookup.calls(), 1); @@ -188,6 +186,67 @@ async fn publisher_cancellation_reelects_exactly_one_subscriber() { assert_idle(&resolver, &lookup); } +#[tokio::test] +async fn publication_and_cancellation_wake_subscribers_outside_resolver_lock() { + struct WakeProbe { + state: Arc>, + wakes: AtomicUsize, + locked_wakes: AtomicUsize, + } + + impl futures::task::ArcWake for WakeProbe { + fn wake_by_ref(probe: &Arc) { + probe.wakes.fetch_add(1, Ordering::Relaxed); + if probe.state.try_lock().is_none() { + probe.locked_wakes.fetch_add(1, Ordering::Relaxed); + } + } + } + + for cancel_owner in [false, true] { + let (resolver, lookup) = fixture(60); + let mut owner = resolver.resolve_location(&location(80)); + assert!(futures::poll!(&mut owner).is_pending()); + let mut waiter = resolver.resolve_location(&location(443)); + let probe = Arc::new(WakeProbe { + state: resolver.state.clone(), + wakes: AtomicUsize::new(0), + locked_wakes: AtomicUsize::new(0), + }); + let waker = futures::task::waker_ref(&probe); + assert!( + waiter + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + + if cancel_owner { + drop(owner); + lookup.discard_cancelled(); + } else { + lookup.reply(Ok(candidates())); + ready_result(owner).unwrap(); + } + assert!(probe.wakes.load(Ordering::Relaxed) > 0); + assert_eq!( + probe.locked_wakes.load(Ordering::Relaxed), + 0, + "cancel_owner={cancel_owner}" + ); + + if cancel_owner { + assert!(futures::poll!(&mut waiter).is_pending()); + lookup.reply(Ok(candidates())); + } + assert_eq!( + ready_result(waiter).unwrap(), + native_addrs_with_port(&candidates(), 443) + ); + assert_idle(&resolver, &lookup); + } +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn concurrent_subscribers_wake_after_publisher_cancellation() { #[derive(Debug, Default)] @@ -212,9 +271,10 @@ async fn concurrent_subscribers_wake_after_publisher_cancellation() { let mut owner = resolver.resolve_location(&location(80)); assert!(futures::poll!(&mut owner).is_pending()); - let barrier = Arc::new(tokio::sync::Barrier::new(129)); + let ports = 0..128; + let barrier = Arc::new(tokio::sync::Barrier::new(ports.len() + 1)); let mut tasks = Vec::new(); - for port in 0..128 { + for port in ports { let resolver = resolver.clone(); let barrier = barrier.clone(); tasks.push(tokio::spawn(async move { @@ -227,16 +287,16 @@ async fn concurrent_subscribers_wake_after_publisher_cancellation() { ); })); } - barrier.wait().await; - assert_eq!(lookup.0.load(Ordering::Relaxed), 1); - drop(owner); tokio::time::timeout(Duration::from_secs(5), async { + barrier.wait().await; + assert_eq!(lookup.0.load(Ordering::Relaxed), 1); + drop(owner); for task in tasks { task.await.unwrap(); } }) .await - .expect("subscribers remained stuck after cancellation"); + .expect("subscribers failed to rendezvous or finish after cancellation"); assert_eq!(lookup.0.load(Ordering::Relaxed), 2); assert!(resolver.state.lock().inflight.is_empty()); assert_eq!( @@ -379,6 +439,7 @@ async fn failures_share_errno_kind_and_message_without_negative_caching() { let mut waiter = resolver.resolve_location(&location(443)); assert!(futures::poll!(&mut owner).is_pending()); assert!(futures::poll!(&mut waiter).is_pending()); + assert_eq!(lookup.calls(), 1); lookup.reply(failure); let first_error = owner.await.unwrap_err(); // A new flight can start before an old subscriber consumes its error. @@ -395,7 +456,7 @@ async fn failures_share_errno_kind_and_message_without_negative_caching() { } #[tokio::test(start_paused = true)] -async fn result_age_starts_at_completion_hits_do_not_extend_it_and_zero_disables_reuse() { +async fn result_age_starts_at_completion_and_hits_do_not_extend_it() { assert_eq!(CachingNativeResolver::DEFAULT_RESULT_TIMEOUT_SECS, 60); let (resolver, lookup) = fixture(60); let mut first = resolver.resolve_location(&location(80)); @@ -405,7 +466,7 @@ async fn result_age_starts_at_completion_hits_do_not_extend_it_and_zero_disables first.await.unwrap(); tokio::time::advance(Duration::from_secs(59)).await; assert_eq!( - resolver.resolve_location(&location(443)).await.unwrap(), + ready_result(resolver.resolve_location(&location(443))).unwrap(), native_addrs_with_port(&candidates(), 443) ); assert_eq!(lookup.calls(), 1); @@ -420,12 +481,16 @@ async fn result_age_starts_at_completion_hits_do_not_extend_it_and_zero_disables ); assert_eq!(lookup.calls(), 2); assert_idle(&resolver, &lookup); +} +#[tokio::test(start_paused = true)] +async fn zero_result_age_disables_reuse_but_coalesces_misses() { let (resolver, lookup) = fixture(0); let mut first = resolver.resolve_location(&location(80)); let mut second = resolver.resolve_location(&location(443)); assert!(futures::poll!(&mut first).is_pending()); assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(lookup.calls(), 1); lookup.reply(Ok(candidates())); first.await.unwrap(); second.await.unwrap(); @@ -443,43 +508,39 @@ async fn cache_evicts_least_recent_host_and_keeps_exact_host_identity() { .lock() .cache .resize(NonZeroUsize::new(2).unwrap()); - for host in [ - "example.test", - "EXAMPLE.test", - "example.test", - "example.test.", + for (host, expect_lookup) in [ + ("example.test", true), + ("EXAMPLE.test", true), + ("example.test", false), + ("example.test.", true), ] { let target = NetLocation::new(Address::Hostname(host.into()), 443); let mut future = resolver.resolve_location(&target); - if host == "example.test" && lookup.calls() == 2 { - future.await.unwrap(); - continue; + if expect_lookup { + assert!(futures::poll!(&mut future).is_pending()); + lookup.reply(Ok(candidates())); } - assert!(futures::poll!(&mut future).is_pending()); - lookup.reply(Ok(candidates())); - future.await.unwrap(); + ready_result(future).unwrap(); + } + { + let state = resolver.state.lock(); + assert_eq!(state.cache.len(), 2); + assert!( + state + .cache + .contains(&Address::Hostname("example.test".into())) + ); + assert!( + !state + .cache + .contains(&Address::Hostname("EXAMPLE.test".into())) + ); + assert!( + state + .cache + .contains(&Address::Hostname("example.test.".into())) + ); } - let state = resolver.state.lock(); - assert_eq!(state.cache.len(), 2); - assert!( - state - .cache - .peek(&Address::Hostname("example.test".into())) - .is_some() - ); - assert!( - state - .cache - .peek(&Address::Hostname("EXAMPLE.test".into())) - .is_none() - ); - assert!( - state - .cache - .peek(&Address::Hostname("example.test.".into())) - .is_some() - ); - drop(state); assert_eq!(lookup.calls(), 3); let evicted = NetLocation::new(Address::Hostname("EXAMPLE.test".into()), 443); let mut retry = resolver.resolve_location(&evicted); @@ -498,18 +559,13 @@ fn default_cache_capacity_is_bounded() { state.cache.put( Address::Hostname(format!("{index}.test")), CachedResolveResult { - timestamp: tokio::time::Instant::now(), + completed_at: tokio::time::Instant::now(), addrs: Arc::from(candidates()), }, ); } assert_eq!(state.cache.len(), CACHING_NATIVE_RESOLVER_CACHE_CAPACITY); - assert!( - state - .cache - .peek(&Address::Hostname("0.test".into())) - .is_none() - ); + assert!(!state.cache.contains(&Address::Hostname("0.test".into()))); } #[tokio::test] @@ -553,7 +609,7 @@ async fn participants_and_inflight_are_bounded_without_new_overload_errors() { // Overflow completion must not remove another caller's flight. assert_eq!(resolver.state.lock().inflight.len(), expected); assert_eq!( - resolver.resolve_location(&location(80)).await.unwrap(), + ready_result(resolver.resolve_location(&location(80))).unwrap(), native_addrs_with_port(&candidates(), 80) ); assert_eq!(resolver.participants.available_permits(), 0); @@ -568,6 +624,51 @@ async fn participants_and_inflight_are_bounded_without_new_overload_errors() { } } +#[tokio::test] +async fn overflow_cancellation_preserves_flight_and_publication_retains_waiter_permits() { + let (resolver, lookup) = fixture(60); + let mut callers: Vec<_> = (0..CACHING_NATIVE_RESOLVER_MAX_PARTICIPANTS) + .map(|port| resolver.resolve_location(&location(port as u16))) + .collect(); + for caller in &mut callers { + assert!(futures::poll!(caller).is_pending()); + } + assert_eq!(resolver.participants.available_permits(), 0); + let mut overflow = resolver.resolve_location(&location(8443)); + assert!(futures::poll!(&mut overflow).is_pending()); + assert_eq!(lookup.calls(), 2); + drop(overflow); + assert!(lookup.replies.lock().pop_back().unwrap().is_closed()); + assert_eq!(lookup.active.load(Ordering::Relaxed), 1); + assert_eq!(resolver.state.lock().inflight.len(), 1); + assert_eq!(resolver.participants.available_permits(), 0); + + lookup.reply(Ok(candidates())); + let owner = callers.remove(0); + assert_eq!(ready_result(owner).unwrap(), candidates()); + assert!(resolver.state.lock().inflight.is_empty()); + assert_eq!(resolver.participants.available_permits(), 1); + assert_eq!(lookup.active.load(Ordering::Relaxed), 0); + assert_eq!( + ready_result(resolver.resolve_location(&location(443))).unwrap(), + native_addrs_with_port(&candidates(), 443) + ); + assert_eq!(resolver.participants.available_permits(), 1); + + let waiter = callers.remove(0); + assert_eq!( + ready_result(waiter).unwrap(), + native_addrs_with_port(&candidates(), 1) + ); + assert_eq!(resolver.participants.available_permits(), 2); + drop(callers); + assert_idle(&resolver, &lookup); + assert_eq!(lookup.calls(), 2); + resolver.state.lock().cache.clear(); + complete_lookup(&resolver, &lookup, 443).await; + assert_idle(&resolver, &lookup); +} + #[tokio::test] async fn numeric_literals_preserve_native_filtering_and_ports() { let resolver = CachingNativeResolver::new(); @@ -596,7 +697,7 @@ async fn numeric_literals_preserve_native_filtering_and_ports() { } #[tokio::test] -#[ignore = "manual native-cache microbenchmark; run under the shared heavy.lock"] +#[ignore = "manual native-cache microbenchmark; run alone in release mode"] async fn benchmark_native_cache_hits_and_cold_miss_suppression() { let (resolver, lookup) = fixture(60); const CONCURRENCY: usize = 512;