From 552b08374267c83f46315d4465cc364f9413e38a Mon Sep 17 00:00:00 2001 From: Jonathan <64296013+Joncallim@users.noreply.github.com> Date: Thu, 3 Sep 2026 01:37:31 +0800 Subject: [PATCH] fix: retain telemetry claim ownership through timeout --- crates/dockermap-daemon/src/cache_refresh.rs | 230 +++++++++++++++++-- 1 file changed, 215 insertions(+), 15 deletions(-) diff --git a/crates/dockermap-daemon/src/cache_refresh.rs b/crates/dockermap-daemon/src/cache_refresh.rs index 703737ab..2e0e98d4 100644 --- a/crates/dockermap-daemon/src/cache_refresh.rs +++ b/crates/dockermap-daemon/src/cache_refresh.rs @@ -103,7 +103,16 @@ struct ResourceTelemetryCache { collection_state: ObservedResourceTelemetryCollectionState, samples: BTreeMap, rotation_cursor: usize, - in_flight: usize, + next_flight_id: u64, + in_flight: BTreeMap, +} + +/// Private cancellation-recovery record. This never reaches a response: it +/// lets a subsequent refresh reclaim a cache slot if the entire refresh task +/// was aborted while an attempt was pending. +#[derive(Clone)] +struct ResourceTelemetryFlight { + deadline: Duration, } #[derive(Clone)] @@ -124,6 +133,7 @@ struct TelemetryCounters { #[derive(Clone)] struct ResourceTelemetryClaim { + flight_id: u64, raw_container_id: String, public_container_id: String, source_generation: u64, @@ -137,7 +147,8 @@ impl Default for ResourceTelemetryCache { collection_state: ObservedResourceTelemetryCollectionState::Unavailable, samples: BTreeMap::new(), rotation_cursor: 0, - in_flight: 0, + next_flight_id: 0, + in_flight: BTreeMap::new(), } } } @@ -146,6 +157,24 @@ impl ResourceTelemetryCache { fn source_reset() -> Self { Self::default() } + + /// A timeout normally completes in the refresh-owned future. This is the + /// secondary recovery path for cancellation of that entire future: stale + /// claims cannot consume the fixed two-slot budget forever, and a late + /// completion cannot remove a newer claim because every flight is opaque + /// and unique. + fn reclaim_expired_flights(&mut self, now: Duration) { + self.in_flight.retain(|_, flight| flight.deadline > now); + if self.in_flight.is_empty() + && self.collection_state == ObservedResourceTelemetryCollectionState::Collecting + { + self.collection_state = if self.samples.is_empty() { + ObservedResourceTelemetryCollectionState::Stale + } else { + ObservedResourceTelemetryCollectionState::Fresh + }; + } + } } /// Daemon-lifetime inventory deltas only. This contains no Docker event stream @@ -737,9 +766,16 @@ pub(crate) async fn refresh_cache(state: &AppState) { } } -/// Launch at most two finite, snapshot-selected stats requests per refresh. +/// Complete at most two finite, snapshot-selected stats requests per refresh. /// The source generation and both revisions are captured before any I/O; a /// completion may publish only if all three still attest the same Docker view. +/// +/// These futures deliberately remain owned by this refresh call rather than +/// being detached. A detached task can be cancelled during shutdown or become +/// non-cooperative below the HTTP boundary after its cache claim increments; +/// in either case there would be no owner left to release the fixed in-flight +/// slot. Waiting for the bounded attempts keeps the claim and its completion +/// in one lifetime while the two requests still run concurrently. async fn refresh_resource_telemetry( state: &AppState, snapshot: DockerSnapshot, @@ -758,10 +794,9 @@ async fn refresh_resource_telemetry( return; } let claims = claim_resource_telemetry(&state.cache, &snapshot, source_generation).await; - for claim in claims { - let state = state.clone(); + let completions = futures_util::future::join_all(claims.into_iter().map(|claim| { let collector = collector.clone(); - tokio::spawn(async move { + async move { let observed_at_ms = wall_clock_millis(); let result = timeout( RESOURCE_TELEMETRY_TIMEOUT, @@ -771,8 +806,12 @@ async fn refresh_resource_telemetry( .ok() .and_then(Result::ok) .and_then(|stats| observed_at_ms.map(|at| (at, stats))); - apply_resource_telemetry(&state.cache, claim, result).await; - }); + (claim, result) + } + })) + .await; + for (claim, result) in completions { + apply_resource_telemetry(&state.cache, claim, result).await; } } @@ -788,7 +827,10 @@ async fn claim_resource_telemetry( let model_revision = cache.snapshot.model_revision.clone(); let observation_revision = cache.docker_observation_token(); let telemetry = &mut cache.resource_telemetry; - let remaining = MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS.saturating_sub(telemetry.in_flight); + let now = monotonic_now(); + telemetry.reclaim_expired_flights(now); + let remaining = + MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS.saturating_sub(telemetry.in_flight.len()); if remaining == 0 { return Vec::new(); } @@ -814,7 +856,21 @@ async fn claim_resource_telemetry( .map(|offset| { let (public_container_id, raw_container_id) = containers[(telemetry.rotation_cursor + offset) % containers.len()].clone(); + telemetry.next_flight_id = telemetry + .next_flight_id + .checked_add(1) + .expect("resource telemetry flight identifier overflow"); + let flight_id = telemetry.next_flight_id; + telemetry.in_flight.insert( + flight_id, + ResourceTelemetryFlight { + deadline: now + .checked_add(RESOURCE_TELEMETRY_TIMEOUT) + .expect("resource telemetry flight deadline overflow"), + }, + ); ResourceTelemetryClaim { + flight_id, raw_container_id, public_container_id, source_generation, @@ -824,7 +880,6 @@ async fn claim_resource_telemetry( }) .collect::>(); telemetry.rotation_cursor = (telemetry.rotation_cursor + count) % containers.len(); - telemetry.in_flight += claims.len(); telemetry.collection_state = ObservedResourceTelemetryCollectionState::Collecting; claims } @@ -849,12 +904,17 @@ async fn apply_resource_telemetry( let revision_changed = cache.snapshot.model_revision != claim.model_revision || cache.docker_observation_token() != claim.observation_revision; let telemetry = &mut cache.resource_telemetry; - telemetry.in_flight = telemetry.in_flight.saturating_sub(1); + // An expired claim was reclaimed after a cancelled refresh. Do not let + // its delayed response publish into a newer flight or consume that + // flight's capacity. + if telemetry.in_flight.remove(&claim.flight_id).is_none() { + return; + } // A newer Docker publication supersedes this sample. Release the fixed // concurrency claim, but never merge measurements from a different model // or observation revision into the current public envelope. if revision_changed || !still_current { - telemetry.collection_state = if telemetry.in_flight == 0 { + telemetry.collection_state = if telemetry.in_flight.is_empty() { ObservedResourceTelemetryCollectionState::Stale } else { ObservedResourceTelemetryCollectionState::Collecting @@ -872,7 +932,7 @@ async fn apply_resource_telemetry( telemetry.samples.insert(claim.public_container_id, record); } } - telemetry.collection_state = if telemetry.in_flight == 0 { + telemetry.collection_state = if telemetry.in_flight.is_empty() { if telemetry.samples.is_empty() { ObservedResourceTelemetryCollectionState::Stale } else { @@ -3969,6 +4029,147 @@ mod scheduler_tests { } } + #[tokio::test] + async fn resource_telemetry_timeout_is_completed_before_the_refresh_releases_its_claim() { + let directory = tempfile::tempdir().unwrap(); + let socket = directory.path().join("telemetry-timeout.sock"); + let listener = UnixListener::bind(&socket).unwrap(); + let server = tokio::spawn(async move { + let (mut connection, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut bytes = [0; 1_024]; + loop { + let read = connection.read(&mut bytes).await.unwrap(); + assert!(read > 0, "client closed before the fixed request head"); + request.extend_from_slice(&bytes[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + assert!( + request.starts_with(b"GET /containers/"), + "only the snapshot-selected stats route is attempted" + ); + // Headers prove the fixed request reached the gateway-like peer, + // while the intentionally absent response body keeps Bollard's + // stream pending until the daemon's own 750 ms bound cancels it. + connection + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n", + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + }); + + let snapshot = mock_snapshot(); + let state = AppState { + cache: Arc::new(RwLock::new(docker_cache(snapshot.clone()))), + docker: Arc::new(RwLock::new(Some(DockerCollector::with_client( + Docker::connect_with_unix(socket.to_str().unwrap(), 5, API_DEFAULT_VERSION) + .unwrap(), + None, + )))), + provider_slot_in_flight: Arc::new(ProviderSlotFlights::default()), + }; + + let started = tokio::time::Instant::now(); + refresh_resource_telemetry(&state, snapshot, 0).await; + assert!( + started.elapsed() >= RESOURCE_TELEMETRY_TIMEOUT, + "the refresh owns the request until its timeout completes" + ); + assert!( + started.elapsed() < Duration::from_secs(2), + "a non-cooperative response body cannot outlive the finite bound" + ); + let cache = state.cache.read().await; + assert!(cache.resource_telemetry.in_flight.is_empty()); + assert!(cache.resource_telemetry.samples.is_empty()); + assert_eq!( + cache.resource_telemetry.collection_state, + ObservedResourceTelemetryCollectionState::Stale + ); + drop(cache); + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + } + + #[tokio::test] + async fn aborted_telemetry_refresh_reclaims_only_its_expired_claim_before_retrying() { + let directory = tempfile::tempdir().unwrap(); + let socket = directory.path().join("telemetry-abort.sock"); + let listener = UnixListener::bind(&socket).unwrap(); + let (claimed_tx, mut claimed_rx) = mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + let (mut connection, _) = listener.accept().await.unwrap(); + let mut request = Vec::new(); + let mut bytes = [0; 1_024]; + loop { + let read = connection.read(&mut bytes).await.unwrap(); + assert!(read > 0, "client closed before the fixed request head"); + request.extend_from_slice(&bytes[..read]); + if request.windows(4).any(|window| window == b"\r\n\r\n") { + break; + } + } + claimed_tx.send(()).unwrap(); + connection + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ntransfer-encoding: chunked\r\n\r\n", + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_secs(5)).await; + }); + + let snapshot = mock_snapshot(); + let state = AppState { + cache: Arc::new(RwLock::new(docker_cache(snapshot.clone()))), + docker: Arc::new(RwLock::new(Some(DockerCollector::with_client( + Docker::connect_with_unix(socket.to_str().unwrap(), 5, API_DEFAULT_VERSION) + .unwrap(), + None, + )))), + provider_slot_in_flight: Arc::new(ProviderSlotFlights::default()), + }; + let refresh_state = state.clone(); + let refresh_snapshot = snapshot.clone(); + let refresh = tokio::spawn(async move { + refresh_resource_telemetry(&refresh_state, refresh_snapshot, 0).await; + }); + tokio::time::timeout(Duration::from_secs(2), claimed_rx.recv()) + .await + .expect("the fixed stats request is claimed before cancellation") + .expect("test gateway reports the claim"); + refresh.abort(); + assert!(refresh.await.unwrap_err().is_cancelled()); + assert_eq!( + state.cache.read().await.resource_telemetry.in_flight.len(), + MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS + ); + + tokio::time::sleep(RESOURCE_TELEMETRY_TIMEOUT + Duration::from_millis(25)).await; + let retry = claim_resource_telemetry(&state.cache, &snapshot, 0).await; + assert_eq!( + retry.len(), + MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS, + "the abandoned claim no longer consumes capacity" + ); + let cache = state.cache.read().await; + assert_eq!( + cache.resource_telemetry.in_flight.len(), + MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS + ); + assert_eq!( + cache.resource_telemetry.next_flight_id, + (MAX_CONCURRENT_RESOURCE_TELEMETRY_REQUESTS * 2) as u64 + ); + drop(cache); + server.abort(); + assert!(server.await.unwrap_err().is_cancelled()); + } + #[test] fn telemetry_sanitizer_retains_only_numeric_metrics_and_derives_rates() { let first = sanitize_resource_telemetry( @@ -4077,10 +4278,9 @@ mod scheduler_tests { telemetry .samples .insert(record.public.container_id.clone(), record); - telemetry.in_flight = 1; let reset = ResourceTelemetryCache::source_reset(); assert!(reset.samples.is_empty()); - assert_eq!(reset.in_flight, 0); + assert!(reset.in_flight.is_empty()); assert_eq!( reset.collection_state, ObservedResourceTelemetryCollectionState::Unavailable