From 3e7d9f741778a550fe0efec57291871217a40e26 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sat, 19 Sep 2026 22:39:04 -0700 Subject: [PATCH 1/4] fix(guest-agent): take the quote off the executor on every path get_quote in tdx-attest takes a process-global std::sync::Mutex and then blocks -- on configfs/TSM, or on a vsock connection to the host's QGS with no connect, read or write timeout. Until now exactly one caller treated it as blocking work: Attest on the v1 surface. GetQuote, the v0 Attest, Tappd.TdxQuote and RawQuote, Worker.GetAttestationForAppKey, issue_cert and get_info all ran it straight on the async executor. That parks a runtime worker for the whole quote, and the runtime is narrower than it looks: #[rocket::main] builds it before the configuration is read, so the workers = 8 in dstack.toml never reaches it and a CVM gets one worker per vCPU. Measured on TDX hardware with a 2-vCPU guest, one quote costs 1.02 s, and 24 concurrent GetQuote calls over the bind-mounted /var/run/dstack.sock stalled every other connection the agent served: 20 of 43 were dropped, and Version went from 0.4 ms to seconds. The stall reaches the agent's own systemd watchdog, whose heartbeat is a Worker.Version request to the external listener, so systemd hit WatchdogSec=30s, SIGABRTed the agent and restarted it. Any container in the CVM can do this -- the socket is bind-mounted into application containers by design -- and the same load over the external TCP listener does the same thing. Put the hop where the blocking is, not at each caller: quote_response and attest_cvm take it themselves, and issue_cert and get_info wrap their platform calls, so a new caller cannot forget. identity() stays on the executor, as its comment says: the retry throttle caps it at one quote per interval for the whole process. With the same load on the same guest, all 43 calls are answered, the Version probe stays at 2 ms p95 against a 1.0 s bound, and four consecutive 24-way loads leave the agent in the process it started as. --- dstack/guest-agent/src/rpc_service.rs | 126 ++++++++++++++++++++--- dstack/guest-agent/src/rpc_service_v1.rs | 16 ++- 2 files changed, 118 insertions(+), 24 deletions(-) diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 7d0159cd0..1a17dd48f 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -138,6 +138,31 @@ pub(crate) struct AppIdentity { pub(crate) key_provider_info: String, } +/// Run one blocking platform attestation off the async executor. +/// +/// `get_quote` in `dstack/tdx-attest/src/linux.rs` takes a process-global +/// `std::sync::Mutex` and then blocks -- on configfs/TSM, or on a vsock +/// connection to the host's QGS opened with no connect, read or write timeout. +/// Called straight from an `async fn` that parks a runtime worker thread for +/// the whole quote, and the runtime `#[rocket::main]` builds has one worker per +/// vCPU rather than the `workers` this crate's `dstack.toml` asks for. +/// +/// Measured on TDX hardware, one quote costs about a second, so concurrent +/// callers on a 2-vCPU CVM stall every other connection the agent serves. +/// `/var/run/dstack.sock` is bind-mounted into application containers by +/// design, which makes that reachable by any container in the CVM; and the +/// stall reaches the agent's own systemd watchdog, whose heartbeat is a +/// `Worker.Version` request to the external listener, so the service is +/// SIGABRTed and restarted at `WatchdogSec=30s`. +async fn attest_off_executor( + platform: Arc, + work: impl FnOnce(&dyn PlatformBackend) -> Result + Send + 'static, +) -> Result { + tokio::task::spawn_blocking(move || work(platform.as_ref())) + .await + .context("the attestation task panicked")? +} + impl AppStateInner { fn info_attestation(&self) -> Result { self.platform.attestation_for_info() @@ -150,10 +175,13 @@ impl AppStateInner { wire: AttestationWire, ) -> Result> { let pubkey = key.public_key_der(); + let attested = pubkey.clone(); let attestation = wire.apply( - self.platform - .certificate_attestation(&pubkey) - .context("Failed to get certificate attestation")?, + attest_off_executor(self.platform.clone(), move |platform| { + platform.certificate_attestation(&attested) + }) + .await + .context("Failed to get certificate attestation")?, ); let csr = CertSigningRequestV2 { confirm: "please sign cert:".to_string(), @@ -299,21 +327,25 @@ impl AppState { self.inner.health.as_deref() } - pub(crate) fn quote_response(&self, report_data: [u8; 64]) -> Result { - self.inner - .platform - .quote_response(report_data, &self.inner.vm_config) + pub(crate) async fn quote_response(&self, report_data: [u8; 64]) -> Result { + let vm_config = self.inner.vm_config.clone(); + attest_off_executor(self.inner.platform.clone(), move |platform| { + platform.quote_response(report_data, &vm_config) + }) + .await } /// Encoded attestation bytes for `Attest`, in the wire form the calling /// surface commits to. - pub(crate) fn attest_cvm( + pub(crate) async fn attest_cvm( &self, report_data: [u8; 64], wire: AttestationWire, ) -> Result> { - wire.apply(self.inner.platform.attest_cvm(report_data)?) - .to_bytes() + attest_off_executor(self.inner.platform.clone(), move |platform| { + wire.apply(platform.attest_cvm(report_data)?).to_bytes() + }) + .await } /// The application's root secp256k1 key, the root of every derived key and @@ -535,7 +567,9 @@ impl InternalRpcHandler { pub async fn get_info(state: &AppState, external: bool) -> Result { let hide_tcb_info = external && !state.config().app_compose.public_tcbinfo; - let versioned_attestation = state.inner.info_attestation()?; + let platform = state.inner.platform.clone(); + let versioned_attestation = + attest_off_executor(platform, |platform| platform.attestation_for_info()).await?; let attestation = versioned_attestation.into_v1(); let app_info = attestation .decode_app_info(false) @@ -678,7 +712,7 @@ impl DstackGuestRpc for InternalRpcHandler { async fn get_quote(self, request: RawQuoteArgs) -> Result { let report_data = pad64(&request.report_data).context("Report data is too long")?; - self.state.quote_response(report_data) + self.state.quote_response(report_data).await } /// Always fails. See the RPC's doc comment in agent_rpc.proto: the method @@ -805,7 +839,8 @@ impl DstackGuestRpc for InternalRpcHandler { Ok(AttestResponse { attestation: self .state - .attest_cvm(report_data, AttestationWire::Legacy)?, + .attest_cvm(report_data, AttestationWire::Legacy) + .await?, }) } @@ -912,7 +947,7 @@ impl TappdRpc for InternalRpcHandlerV0 { }; let report_data = content_type.to_report_data_with_hash(&request.report_data, &request.hash_algorithm)?; - let response = self.state.quote_response(report_data)?; + let response = self.state.quote_response(report_data).await?; Ok(TdxQuoteResponse { quote: response.quote, event_log: response.event_log, @@ -992,7 +1027,7 @@ impl WorkerRpc for ExternalRpcHandler { request: GetAttestationForAppKeyRequest, ) -> Result { let report_data = self.app_key_report_data(&request.algorithm).await?; - self.state.quote_response(report_data) + self.state.quote_response(report_data).await } } @@ -1093,6 +1128,67 @@ pub(crate) mod tests { use std::io::Write; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + /// A platform whose quote blocks the calling thread, as the real one does. + struct BlockingPlatform; + + impl PlatformBackend for BlockingPlatform { + fn attestation_for_info(&self) -> Result { + unimplemented!("only the quote path blocks in this fixture") + } + + fn certificate_attestation(&self, _pubkey: &[u8]) -> Result { + unimplemented!("only the quote path blocks in this fixture") + } + + fn quote_response( + &self, + _report_data: [u8; 64], + _vm_config: &str, + ) -> Result { + std::thread::sleep(Duration::from_millis(50)); + Ok(GetQuoteResponse::default()) + } + + fn attest_cvm(&self, _report_data: [u8; 64]) -> Result { + unimplemented!("only the quote path blocks in this fixture") + } + } + + /// The whole point of the hop: a quote must not park the runtime. + /// + /// A single-threaded runtime is the sharpest form of the production + /// arrangement, where the runtime has one worker per vCPU. If the platform + /// call ran on the executor the ticker below could never be polled while + /// the quote was in flight, which is how 24 concurrent callers used to + /// stall every other connection -- including the `Worker.Version` request + /// the systemd watchdog uses as its heartbeat, so the agent was SIGABRTed + /// and restarted at `WatchdogSec=30s`. + #[tokio::test(flavor = "current_thread")] + async fn a_blocking_quote_does_not_park_the_runtime() { + let platform: Arc = Arc::new(BlockingPlatform); + let ticks = Arc::new(AtomicUsize::new(0)); + let counter = ticks.clone(); + let quote = + attest_off_executor(platform, |platform| platform.quote_response([0u8; 64], "")); + tokio::pin!(quote); + tokio::select! { + result = &mut quote => { + result.expect("the quote must be answered"); + } + _ = async move { + loop { + counter.fetch_add(1, Ordering::Relaxed); + tokio::task::yield_now().await; + } + } => unreachable!("the ticker runs until the quote completes"), + } + assert!( + ticks.load(Ordering::Relaxed) > 1, + "the runtime was parked for the whole quote: the executor ran {} times", + ticks.load(Ordering::Relaxed) + ); + } + fn extract_pubkey_from_report_data(report_data: &[u8], prefix: &str) -> Result> { let end = report_data .iter() diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 91a5141de..57cb7a045 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -206,15 +206,12 @@ impl DstackGuestRpc for V1RpcHandler { async fn attest(self, request: AttestRequest) -> Result { let report_data = pad64(&request.report_data).context("report data is too long")?; - // Generating a quote takes a global mutex and then blocks in an ioctl. - // On the async executor that parks a worker thread for the duration and - // stalls every other connection this agent is serving. - let state = self.state.clone(); - let attestation = tokio::task::spawn_blocking(move || { - state.attest_cvm(report_data, AttestationWire::MsgpackV1) - }) - .await - .context("the attestation task panicked")??; + // `attest_cvm` takes the blocking hop itself, for every caller. See + // `attest_off_executor` in `rpc_service.rs`. + let attestation = self + .state + .attest_cvm(report_data, AttestationWire::MsgpackV1) + .await?; Ok(AttestResponse { attestation, boottime_gpu_evidence: boottime_gpu_evidence( @@ -637,6 +634,7 @@ mod tests { let (state, _guard) = state().await; let legacy = state .attest_cvm([0u8; 64], AttestationWire::Legacy) + .await .unwrap(); assert_eq!( legacy.first(), From a193b5eb70035391c4b7f14885eadf4fa82353a5 Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Sun, 20 Sep 2026 01:08:24 -0700 Subject: [PATCH 2/4] fix(guest-agent): coalesce identity retries off the executor --- dstack/guest-agent/src/rpc_service.rs | 98 +++++++++++++++++++++--- dstack/guest-agent/src/rpc_service_v1.rs | 12 +-- 2 files changed, 92 insertions(+), 18 deletions(-) diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 1a17dd48f..63a653ab2 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -113,6 +113,7 @@ struct AppStateInner { /// it is written only on the degraded path, and read only when the cache /// is empty. identity_last_failure: Mutex>, + identity_refresh: Arc>, /// `sys_vendor` and `product_name`, read once. Neither changes while the /// VM is running. cloud_vendor: String, @@ -299,6 +300,7 @@ impl AppState { app_root_signing_key, identity: RwLock::new(None), identity_last_failure: Mutex::new(None), + identity_refresh: Arc::default(), cloud_vendor: read_dmi_file("sys_vendor"), cloud_product: read_dmi_file("product_name"), }), @@ -308,7 +310,7 @@ impl AppState { // down, and `identity()` retries on demand. A failure here counts as // the first attempt and arms the retry throttle, so the boot attempt // and a request-driven one are on the same budget. - if let Err(err) = me.identity() { + if let Err(err) = me.identity().await { error!("failed to decode app identity at startup: {err:?}"); } me.maybe_request_demo_cert(); @@ -386,7 +388,29 @@ impl AppState { /// include anonymous ones. Within the window the caller is told the /// identity is unavailable and when the next attempt is, and the platform /// is not touched at all. - pub(crate) fn identity(&self) -> Result> { + pub(crate) async fn identity(&self) -> Result> { + if let Some(identity) = self + .inner + .identity + .read() + .or_panic("lock should never fail") + .as_ref() + { + return Ok(identity.clone()); + } + let refresh = self.inner.identity_refresh.clone().lock_owned().await; + let state = self.clone(); + tokio::task::spawn_blocking(move || { + // The task owns exclusion through cache publication, even if its + // requester disconnects while the platform is still quoting. + let _refresh = refresh; + state.identity_blocking() + }) + .await + .context("identity refresh task failed")? + } + + fn identity_blocking(&self) -> Result> { if let Some(identity) = self .inner .identity @@ -396,9 +420,6 @@ impl AppState { { return Ok(identity.clone()); } - // Two callers arriving together on a cold cache may both attempt once. - // That is the same race the cache has always had, and one extra quote - // is not worth holding a lock across the decode for. if let Some(failed_at) = *self .inner .identity_last_failure @@ -413,10 +434,6 @@ impl AppState { ); } } - // Blocking, and deliberately left on the executor: the throttle above - // caps this at one quote per interval for the whole process, which is - // far short of what would justify a `spawn_blocking` hop and making - // every caller of `identity()` async to reach it. let identity = match decode_identity(&self.inner) { Ok(identity) => Arc::new(identity), Err(err) => { @@ -1231,6 +1248,12 @@ pub(crate) mod tests { struct InfoAttestationProbe { calls: AtomicUsize, failing: AtomicBool, + gate: Mutex< + Option<( + tokio::sync::oneshot::Sender<()>, + std::sync::mpsc::Receiver<()>, + )>, + >, } impl InfoAttestationProbe { @@ -1238,6 +1261,7 @@ pub(crate) mod tests { Self { calls: AtomicUsize::new(0), failing: AtomicBool::new(true), + gate: Mutex::new(None), } } @@ -1424,6 +1448,10 @@ pNs85uhOZE8z2jr8Pg== impl PlatformBackend for TestSimulatorPlatform { fn attestation_for_info(&self) -> Result { self.probe.calls.fetch_add(1, Ordering::Relaxed); + if let Some((entered, release)) = self.probe.gate.lock().unwrap().take() { + entered.send(()).unwrap(); + release.recv_timeout(Duration::from_secs(5)).unwrap(); + } if self.probe.failing.load(Ordering::Relaxed) { anyhow::bail!("the platform cannot attest right now"); } @@ -1508,6 +1536,7 @@ pNs85uhOZE8z2jr8Pg== app_root_signing_key: SigningKey::from_slice(&DUMMY_K256_KEY).ok(), identity: RwLock::new(None), identity_last_failure: Mutex::new(None), + identity_refresh: Arc::default(), // Read the same way production does, so a test comparing v1 `Info` // against v0 `get_info` compares like with like. cloud_vendor: read_dmi_file("sys_vendor"), @@ -2108,6 +2137,46 @@ pNs85uhOZE8z2jr8Pg== assert!(err.contains("removed in dstack 0.6.0"), "{err}"); } + #[tokio::test(flavor = "current_thread")] + async fn identity_retry_is_off_executor_singleflight_and_cancellation_safe() { + let probe = Arc::new(InfoAttestationProbe::failing()); + let (state, _guard) = setup_test_state_with_probe(probe.clone()).await; + state.identity().await.unwrap_err(); + *state.inner.identity_last_failure.lock().unwrap() = + Some(Instant::now() - IDENTITY_RETRY_INTERVAL); + probe.set_failing(false); + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + *probe.gate.lock().unwrap() = Some((entered_tx, release_rx)); + let first_state = state.clone(); + let first = tokio::spawn(async move { first_state.identity().await }); + tokio::time::timeout(Duration::from_secs(2), entered_rx) + .await + .unwrap() + .unwrap(); + first.abort(); + assert!(first.await.unwrap_err().is_cancelled()); + let mut followers = Vec::new(); + for _ in 0..16 { + let state = state.clone(); + followers.push(tokio::spawn(async move { state.identity().await })); + } + tokio::time::sleep(Duration::from_millis(20)).await; + assert_eq!(probe.calls(), 2, "only one retry may enter the platform"); + assert!(followers.iter().all(|task| !task.is_finished())); + release_tx.send(()).unwrap(); + let identity = state.identity().await.unwrap(); + for task in followers { + let other = tokio::time::timeout(Duration::from_secs(2), task) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(Arc::ptr_eq(&identity, &other)); + } + assert_eq!(probe.calls(), 2); + } + /// A failed decode must be as cheap to repeat as a cached success is. /// `identity()` is reached from the anonymous `/prpc/v1/Info`, and each /// attempt is a hardware quote plus an RTMR replay under the global quote @@ -2120,12 +2189,13 @@ pNs85uhOZE8z2jr8Pg== let err = state .identity() + .await .expect_err("the platform refuses to attest"); assert!(err.to_string().contains("cannot attest"), "{err}"); assert_eq!(probe.calls(), 1); for _ in 0..8 { - let err = state.identity().expect_err("still throttled"); + let err = state.identity().await.expect_err("still throttled"); assert!( err.to_string().contains("the app identity is unavailable"), "{err}" @@ -2147,6 +2217,7 @@ pNs85uhOZE8z2jr8Pg== let (state, _guard) = setup_test_state_with_probe(probe.clone()).await; state .identity() + .await .expect_err("the platform refuses to attest"); // Age the recorded failure rather than sleeping out the interval. @@ -2161,10 +2232,13 @@ pNs85uhOZE8z2jr8Pg== ); probe.set_failing(false); - let identity = state.identity().expect("the platform recovered"); + let identity = state.identity().await.expect("the platform recovered"); assert_eq!(probe.calls(), 2); - let again = state.identity().expect("a decoded identity is cached"); + let again = state + .identity() + .await + .expect("a decoded identity is cached"); assert!(Arc::ptr_eq(&identity, &again)); assert_eq!( probe.calls(), diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 57cb7a045..69213538e 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -116,8 +116,8 @@ fn boottime_gpu_evidence(include: bool, path: &Path) -> Vec { /// `key_provider_info`; v1 blanks `key_provider_info` too, because it names the /// component that holds the app's keys and an external caller has no need for /// it. The frozen behaviour is unchanged on its own surface. -fn info_response(state: &AppState, hide_documents: bool) -> Result { - let identity = state.identity()?; +async fn info_response(state: &AppState, hide_documents: bool) -> Result { + let identity = state.identity().await?; let document = |value: &dyn Fn() -> String| { if hide_documents { String::new() @@ -243,7 +243,7 @@ impl DstackGuestRpc for V1RpcHandler { /// itself, so there is nobody to hide from. The external surface applies /// `public_tcbinfo`; see [`ExternalV1RpcHandler::info`]. async fn info(self, _request: InfoRequest) -> Result { - info_response(&self.state, false) + info_response(&self.state, false).await } async fn version(self, _request: VersionRequest) -> Result { @@ -279,7 +279,7 @@ impl ExternalV1RpcHandler { impl WorkerRpc for ExternalV1RpcHandler { async fn info(self, _request: InfoRequest) -> Result { let hide = !self.state.config().app_compose.public_tcbinfo; - info_response(&self.state, hide) + info_response(&self.state, hide).await } async fn version(self, _request: VersionRequest) -> Result { @@ -500,7 +500,7 @@ mod tests { async fn info_decodes_identity_at_most_once() { let (state, _guard) = state().await; - let first = state.identity().unwrap(); + let first = state.identity().await.unwrap(); for _ in 0..8 { V1RpcHandler::new(state.clone()) .info(InfoRequest {}) @@ -511,7 +511,7 @@ mod tests { .await .unwrap(); } - let last = state.identity().unwrap(); + let last = state.identity().await.unwrap(); assert!( std::sync::Arc::ptr_eq(&first, &last), From 3682d1121e3b3a93122bca72eea5157c9cd0686a Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Wed, 23 Sep 2026 23:55:02 -0700 Subject: [PATCH 3/4] refactor(guest-agent): queue platform calls behind one async lock Route every platform call through `AppStateInner::attest`, which takes an owned `tokio::sync::Mutex` guard and moves it into the `spawn_blocking` task. `dstack-attest` already serialises quotes under `QUOTE_LOCK`, so this costs no throughput, but a flood of callers now waits in the async queue instead of each parking one of up to 512 blocking-pool threads on that lock. The same lock gives `identity()` single-flight retries and cancellation safety for free, so the dedicated `identity_refresh` mutex goes. The throttle check moves ahead of the queue, so a throttled `Info` is rejected without waiting behind in-flight quotes. `decode_identity` takes the platform directly and `info_attestation` is dropped. The quote regression test reuses the probe gate instead of a separate `BlockingPlatform` fixture. --- dstack/guest-agent/src/rpc_service.rs | 267 ++++++++++------------- dstack/guest-agent/src/rpc_service_v1.rs | 2 - 2 files changed, 113 insertions(+), 156 deletions(-) diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index 63a653ab2..ecb048c8d 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -61,8 +61,8 @@ const IDENTITY_RETRY_INTERVAL: Duration = Duration::from_secs(30); /// /// Costs a quote and an event-log replay, which is why its result is cached for /// the life of the process. -fn decode_identity(inner: &AppStateInner) -> Result { - let attestation = inner.info_attestation()?.into_v1(); +fn decode_identity(platform: &dyn PlatformBackend) -> Result { + let attestation = platform.attestation_for_info()?.into_v1(); let app_info = attestation .decode_app_info(false) .context("failed to decode app info")?; @@ -113,7 +113,8 @@ struct AppStateInner { /// it is written only on the degraded path, and read only when the cache /// is empty. identity_last_failure: Mutex>, - identity_refresh: Arc>, + /// Queues platform calls; see [`AppStateInner::attest`]. + attest_lock: Arc>, /// `sys_vendor` and `product_name`, read once. Neither changes while the /// VM is running. cloud_vendor: String, @@ -139,34 +140,27 @@ pub(crate) struct AppIdentity { pub(crate) key_provider_info: String, } -/// Run one blocking platform attestation off the async executor. -/// -/// `get_quote` in `dstack/tdx-attest/src/linux.rs` takes a process-global -/// `std::sync::Mutex` and then blocks -- on configfs/TSM, or on a vsock -/// connection to the host's QGS opened with no connect, read or write timeout. -/// Called straight from an `async fn` that parks a runtime worker thread for -/// the whole quote, and the runtime `#[rocket::main]` builds has one worker per -/// vCPU rather than the `workers` this crate's `dstack.toml` asks for. -/// -/// Measured on TDX hardware, one quote costs about a second, so concurrent -/// callers on a 2-vCPU CVM stall every other connection the agent serves. -/// `/var/run/dstack.sock` is bind-mounted into application containers by -/// design, which makes that reachable by any container in the CVM; and the -/// stall reaches the agent's own systemd watchdog, whose heartbeat is a -/// `Worker.Version` request to the external listener, so the service is -/// SIGABRTed and restarted at `WatchdogSec=30s`. -async fn attest_off_executor( - platform: Arc, - work: impl FnOnce(&dyn PlatformBackend) -> Result + Send + 'static, -) -> Result { - tokio::task::spawn_blocking(move || work(platform.as_ref())) +impl AppStateInner { + /// Run a platform call on the blocking pool, one at a time. + /// + /// Platform calls block on the hardware quote (about a second on TDX), so + /// running one on the executor stalls every connection the agent serves, + /// including the systemd watchdog heartbeat. Callers queue here rather than + /// on `dstack-attest`'s quote lock, so a flood of requests holds no threads. + /// The guard moves into the task, so a cancelled caller cannot let the + /// next call in while its own is still running. + async fn attest( + &self, + work: impl FnOnce(&dyn PlatformBackend) -> Result + Send + 'static, + ) -> Result { + let guard = self.attest_lock.clone().lock_owned().await; + let platform = self.platform.clone(); + tokio::task::spawn_blocking(move || { + let _guard = guard; + work(platform.as_ref()) + }) .await .context("the attestation task panicked")? -} - -impl AppStateInner { - fn info_attestation(&self) -> Result { - self.platform.attestation_for_info() } async fn issue_cert( @@ -178,11 +172,9 @@ impl AppStateInner { let pubkey = key.public_key_der(); let attested = pubkey.clone(); let attestation = wire.apply( - attest_off_executor(self.platform.clone(), move |platform| { - platform.certificate_attestation(&attested) - }) - .await - .context("Failed to get certificate attestation")?, + self.attest(move |platform| platform.certificate_attestation(&attested)) + .await + .context("Failed to get certificate attestation")?, ); let csr = CertSigningRequestV2 { confirm: "please sign cert:".to_string(), @@ -300,7 +292,7 @@ impl AppState { app_root_signing_key, identity: RwLock::new(None), identity_last_failure: Mutex::new(None), - identity_refresh: Arc::default(), + attest_lock: Arc::default(), cloud_vendor: read_dmi_file("sys_vendor"), cloud_product: read_dmi_file("product_name"), }), @@ -331,10 +323,9 @@ impl AppState { pub(crate) async fn quote_response(&self, report_data: [u8; 64]) -> Result { let vm_config = self.inner.vm_config.clone(); - attest_off_executor(self.inner.platform.clone(), move |platform| { - platform.quote_response(report_data, &vm_config) - }) - .await + self.inner + .attest(move |platform| platform.quote_response(report_data, &vm_config)) + .await } /// Encoded attestation bytes for `Attest`, in the wire form the calling @@ -344,10 +335,9 @@ impl AppState { report_data: [u8; 64], wire: AttestationWire, ) -> Result> { - attest_off_executor(self.inner.platform.clone(), move |platform| { - wire.apply(platform.attest_cvm(report_data)?).to_bytes() - }) - .await + self.inner + .attest(move |platform| wire.apply(platform.attest_cvm(report_data)?).to_bytes()) + .await } /// The application's root secp256k1 key, the root of every derived key and @@ -389,28 +379,40 @@ impl AppState { /// identity is unavailable and when the next attempt is, and the platform /// is not touched at all. pub(crate) async fn identity(&self) -> Result> { - if let Some(identity) = self - .inner - .identity - .read() - .or_panic("lock should never fail") - .as_ref() - { - return Ok(identity.clone()); + if let Some(identity) = self.cached_identity()? { + return Ok(identity); } - let refresh = self.inner.identity_refresh.clone().lock_owned().await; let state = self.clone(); - tokio::task::spawn_blocking(move || { - // The task owns exclusion through cache publication, even if its - // requester disconnects while the platform is still quoting. - let _refresh = refresh; - state.identity_blocking() - }) - .await - .context("identity refresh task failed")? + self.inner + .attest(move |platform| { + // Callers queued behind a decode pick up its outcome here. + if let Some(identity) = state.cached_identity()? { + return Ok(identity); + } + let identity = match decode_identity(platform) { + Ok(identity) => Arc::new(identity), + Err(err) => { + *state + .inner + .identity_last_failure + .lock() + .or_panic("lock should never fail") = Some(Instant::now()); + return Err(err); + } + }; + *state + .inner + .identity + .write() + .or_panic("lock should never fail") = Some(identity.clone()); + Ok(identity) + }) + .await } - fn identity_blocking(&self) -> Result> { + /// The cached identity, or `None` if a decode may be attempted now. Fails + /// while the last failed decode is still throttled. + fn cached_identity(&self) -> Result>> { if let Some(identity) = self .inner .identity @@ -418,7 +420,7 @@ impl AppState { .or_panic("lock should never fail") .as_ref() { - return Ok(identity.clone()); + return Ok(Some(identity.clone())); } if let Some(failed_at) = *self .inner @@ -434,23 +436,7 @@ impl AppState { ); } } - let identity = match decode_identity(&self.inner) { - Ok(identity) => Arc::new(identity), - Err(err) => { - *self - .inner - .identity_last_failure - .lock() - .or_panic("lock should never fail") = Some(Instant::now()); - return Err(err); - } - }; - *self - .inner - .identity - .write() - .or_panic("lock should never fail") = Some(identity.clone()); - Ok(identity) + Ok(None) } /// The VM's hardware configuration, as the VMM produced it. @@ -584,9 +570,10 @@ impl InternalRpcHandler { pub async fn get_info(state: &AppState, external: bool) -> Result { let hide_tcb_info = external && !state.config().app_compose.public_tcbinfo; - let platform = state.inner.platform.clone(); - let versioned_attestation = - attest_off_executor(platform, |platform| platform.attestation_for_info()).await?; + let versioned_attestation = state + .inner + .attest(|platform| platform.attestation_for_info()) + .await?; let attestation = versioned_attestation.into_v1(); let app_info = attestation .decode_app_info(false) @@ -1145,67 +1132,6 @@ pub(crate) mod tests { use std::io::Write; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; - /// A platform whose quote blocks the calling thread, as the real one does. - struct BlockingPlatform; - - impl PlatformBackend for BlockingPlatform { - fn attestation_for_info(&self) -> Result { - unimplemented!("only the quote path blocks in this fixture") - } - - fn certificate_attestation(&self, _pubkey: &[u8]) -> Result { - unimplemented!("only the quote path blocks in this fixture") - } - - fn quote_response( - &self, - _report_data: [u8; 64], - _vm_config: &str, - ) -> Result { - std::thread::sleep(Duration::from_millis(50)); - Ok(GetQuoteResponse::default()) - } - - fn attest_cvm(&self, _report_data: [u8; 64]) -> Result { - unimplemented!("only the quote path blocks in this fixture") - } - } - - /// The whole point of the hop: a quote must not park the runtime. - /// - /// A single-threaded runtime is the sharpest form of the production - /// arrangement, where the runtime has one worker per vCPU. If the platform - /// call ran on the executor the ticker below could never be polled while - /// the quote was in flight, which is how 24 concurrent callers used to - /// stall every other connection -- including the `Worker.Version` request - /// the systemd watchdog uses as its heartbeat, so the agent was SIGABRTed - /// and restarted at `WatchdogSec=30s`. - #[tokio::test(flavor = "current_thread")] - async fn a_blocking_quote_does_not_park_the_runtime() { - let platform: Arc = Arc::new(BlockingPlatform); - let ticks = Arc::new(AtomicUsize::new(0)); - let counter = ticks.clone(); - let quote = - attest_off_executor(platform, |platform| platform.quote_response([0u8; 64], "")); - tokio::pin!(quote); - tokio::select! { - result = &mut quote => { - result.expect("the quote must be answered"); - } - _ = async move { - loop { - counter.fetch_add(1, Ordering::Relaxed); - tokio::task::yield_now().await; - } - } => unreachable!("the ticker runs until the quote completes"), - } - assert!( - ticks.load(Ordering::Relaxed) > 1, - "the runtime was parked for the whole quote: the executor ran {} times", - ticks.load(Ordering::Relaxed) - ); - } - fn extract_pubkey_from_report_data(report_data: &[u8], prefix: &str) -> Result> { let end = report_data .iter() @@ -1243,7 +1169,8 @@ pub(crate) mod tests { /// How many times the fixture platform was asked to attest for `Info`, and /// whether it should refuse. Counting the calls is the only way to see the /// identity throttle work: what it changes is how often the platform is - /// touched, not what any single call returns. + /// touched, not what any single call returns. The gate holds the next + /// `Info` or quote call inside the platform, as a slow quote does. #[derive(Default)] struct InfoAttestationProbe { calls: AtomicUsize, @@ -1259,9 +1186,8 @@ pub(crate) mod tests { impl InfoAttestationProbe { fn failing() -> Self { Self { - calls: AtomicUsize::new(0), failing: AtomicBool::new(true), - gate: Mutex::new(None), + ..Self::default() } } @@ -1272,6 +1198,27 @@ pub(crate) mod tests { fn set_failing(&self, failing: bool) { self.failing.store(failing, Ordering::Relaxed); } + + /// Arm the gate: the receiver fires once a platform call is blocked + /// in it, and sending on the sender lets that call finish. + fn arm_gate( + &self, + ) -> ( + tokio::sync::oneshot::Receiver<()>, + std::sync::mpsc::Sender<()>, + ) { + let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); + let (release_tx, release_rx) = std::sync::mpsc::channel(); + *self.gate.lock().unwrap() = Some((entered_tx, release_rx)); + (entered_rx, release_tx) + } + + fn pass_gate(&self) { + if let Some((entered, release)) = self.gate.lock().unwrap().take() { + entered.send(()).unwrap(); + release.recv_timeout(Duration::from_secs(5)).unwrap(); + } + } } /// The fixture state, built against a platform the probe watches. @@ -1448,10 +1395,7 @@ pNs85uhOZE8z2jr8Pg== impl PlatformBackend for TestSimulatorPlatform { fn attestation_for_info(&self) -> Result { self.probe.calls.fetch_add(1, Ordering::Relaxed); - if let Some((entered, release)) = self.probe.gate.lock().unwrap().take() { - entered.send(()).unwrap(); - release.recv_timeout(Duration::from_secs(5)).unwrap(); - } + self.probe.pass_gate(); if self.probe.failing.load(Ordering::Relaxed) { anyhow::bail!("the platform cannot attest right now"); } @@ -1474,6 +1418,7 @@ pNs85uhOZE8z2jr8Pg== report_data: [u8; 64], vm_config: &str, ) -> Result { + self.probe.pass_gate(); let attestation = patch_report_data(&self.attestation, report_data); let Some(quote) = attestation.platform.tdx_quote().map(ToOwned::to_owned) else { return Err(anyhow::anyhow!( @@ -1536,7 +1481,7 @@ pNs85uhOZE8z2jr8Pg== app_root_signing_key: SigningKey::from_slice(&DUMMY_K256_KEY).ok(), identity: RwLock::new(None), identity_last_failure: Mutex::new(None), - identity_refresh: Arc::default(), + attest_lock: Arc::default(), // Read the same way production does, so a test comparing v1 `Info` // against v0 `get_info` compares like with like. cloud_vendor: read_dmi_file("sys_vendor"), @@ -2137,6 +2082,22 @@ pNs85uhOZE8z2jr8Pg== assert!(err.contains("removed in dstack 0.6.0"), "{err}"); } + /// A quote must not park the runtime. On a single-threaded runtime the + /// test only sees the quote enter the platform if it left the executor. + #[tokio::test(flavor = "current_thread")] + async fn a_blocking_quote_does_not_park_the_runtime() { + let probe = Arc::new(InfoAttestationProbe::default()); + let (state, _guard) = setup_test_state_with_probe(probe.clone()).await; + let (entered_rx, release_tx) = probe.arm_gate(); + let quote = tokio::spawn(async move { state.quote_response([0u8; 64]).await }); + tokio::time::timeout(Duration::from_secs(2), entered_rx) + .await + .unwrap() + .unwrap(); + release_tx.send(()).unwrap(); + quote.await.unwrap().unwrap(); + } + #[tokio::test(flavor = "current_thread")] async fn identity_retry_is_off_executor_singleflight_and_cancellation_safe() { let probe = Arc::new(InfoAttestationProbe::failing()); @@ -2145,9 +2106,7 @@ pNs85uhOZE8z2jr8Pg== *state.inner.identity_last_failure.lock().unwrap() = Some(Instant::now() - IDENTITY_RETRY_INTERVAL); probe.set_failing(false); - let (entered_tx, entered_rx) = tokio::sync::oneshot::channel(); - let (release_tx, release_rx) = std::sync::mpsc::channel(); - *probe.gate.lock().unwrap() = Some((entered_tx, release_rx)); + let (entered_rx, release_tx) = probe.arm_gate(); let first_state = state.clone(); let first = tokio::spawn(async move { first_state.identity().await }); tokio::time::timeout(Duration::from_secs(2), entered_rx) diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 69213538e..4213d8110 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -206,8 +206,6 @@ impl DstackGuestRpc for V1RpcHandler { async fn attest(self, request: AttestRequest) -> Result { let report_data = pad64(&request.report_data).context("report data is too long")?; - // `attest_cvm` takes the blocking hop itself, for every caller. See - // `attest_off_executor` in `rpc_service.rs`. let attestation = self .state .attest_cvm(report_data, AttestationWire::MsgpackV1) From a201e3f56a95ac4bf3febbe4585f51a9554803bf Mon Sep 17 00:00:00 2001 From: Kevin Wang Date: Thu, 24 Sep 2026 00:40:41 -0700 Subject: [PATCH 4/4] refactor(guest-agent): cache identity in a tokio OnceCell OnceCell::get_or_try_init already gives single-flight initialisation with retry on failure, so the hand-rolled cache check inside the attest closure and the cached_identity() split are no longer needed. The identity is written once and never changes, so it is returned by reference instead of through an Arc. Drop the single-flight test and the Info half of the probe gate; the behaviour it pinned is now OnceCell's. --- dstack/guest-agent/src/rpc_service.rs | 131 +++++------------------ dstack/guest-agent/src/rpc_service_v1.rs | 2 +- 2 files changed, 28 insertions(+), 105 deletions(-) diff --git a/dstack/guest-agent/src/rpc_service.rs b/dstack/guest-agent/src/rpc_service.rs index ecb048c8d..cd0c2d700 100644 --- a/dstack/guest-agent/src/rpc_service.rs +++ b/dstack/guest-agent/src/rpc_service.rs @@ -107,11 +107,9 @@ struct AppStateInner { /// when each parsed the key itself. app_root_signing_key: Option, /// Identity as decoded from the boot attestation. See [`AppIdentity`]. - identity: RwLock>>, + identity: tokio::sync::OnceCell, /// When the last identity decode failed, so the retry can be throttled. - /// See [`IDENTITY_RETRY_INTERVAL`]. Separate from the cache above because - /// it is written only on the degraded path, and read only when the cache - /// is empty. + /// See [`IDENTITY_RETRY_INTERVAL`]. identity_last_failure: Mutex>, /// Queues platform calls; see [`AppStateInner::attest`]. attest_lock: Arc>, @@ -290,8 +288,8 @@ impl AppState { health, gpu_attestor, app_root_signing_key, - identity: RwLock::new(None), - identity_last_failure: Mutex::new(None), + identity: Default::default(), + identity_last_failure: Default::default(), attest_lock: Arc::default(), cloud_vendor: read_dmi_file("sys_vendor"), cloud_product: read_dmi_file("product_name"), @@ -378,67 +376,31 @@ impl AppState { /// include anonymous ones. Within the window the caller is told the /// identity is unavailable and when the next attempt is, and the platform /// is not touched at all. - pub(crate) async fn identity(&self) -> Result> { - if let Some(identity) = self.cached_identity()? { - return Ok(identity); - } - let state = self.clone(); + pub(crate) async fn identity(&self) -> Result<&AppIdentity> { self.inner - .attest(move |platform| { - // Callers queued behind a decode pick up its outcome here. - if let Some(identity) = state.cached_identity()? { - return Ok(identity); - } - let identity = match decode_identity(platform) { - Ok(identity) => Arc::new(identity), - Err(err) => { - *state - .inner - .identity_last_failure - .lock() - .or_panic("lock should never fail") = Some(Instant::now()); - return Err(err); + .identity + .get_or_try_init(|| async { + let last_failure = &self.inner.identity_last_failure; + if let Some(failed_at) = *last_failure.lock().or_panic("lock should never fail") { + let elapsed = failed_at.elapsed(); + if elapsed < IDENTITY_RETRY_INTERVAL { + let retry_in = (IDENTITY_RETRY_INTERVAL - elapsed).as_secs() + 1; + anyhow::bail!( + "the app identity is unavailable: decoding it failed and the next attempt is at most {retry_in}s away" + ); } - }; - *state - .inner - .identity - .write() - .or_panic("lock should never fail") = Some(identity.clone()); - Ok(identity) + } + self.inner + .attest(decode_identity) + .await + .inspect_err(|_| { + *last_failure.lock().or_panic("lock should never fail") = + Some(Instant::now()); + }) }) .await } - /// The cached identity, or `None` if a decode may be attempted now. Fails - /// while the last failed decode is still throttled. - fn cached_identity(&self) -> Result>> { - if let Some(identity) = self - .inner - .identity - .read() - .or_panic("lock should never fail") - .as_ref() - { - return Ok(Some(identity.clone())); - } - if let Some(failed_at) = *self - .inner - .identity_last_failure - .lock() - .or_panic("lock should never fail") - { - let elapsed = failed_at.elapsed(); - if elapsed < IDENTITY_RETRY_INTERVAL { - let retry_in = (IDENTITY_RETRY_INTERVAL - elapsed).as_secs() + 1; - anyhow::bail!( - "the app identity is unavailable: decoding it failed and the next attempt is at most {retry_in}s away" - ); - } - } - Ok(None) - } - /// The VM's hardware configuration, as the VMM produced it. pub(crate) fn vm_config(&self) -> &str { &self.inner.vm_config @@ -1170,7 +1132,7 @@ pub(crate) mod tests { /// whether it should refuse. Counting the calls is the only way to see the /// identity throttle work: what it changes is how often the platform is /// touched, not what any single call returns. The gate holds the next - /// `Info` or quote call inside the platform, as a slow quote does. + /// quote inside the platform, as a slow quote does. #[derive(Default)] struct InfoAttestationProbe { calls: AtomicUsize, @@ -1395,7 +1357,6 @@ pNs85uhOZE8z2jr8Pg== impl PlatformBackend for TestSimulatorPlatform { fn attestation_for_info(&self) -> Result { self.probe.calls.fetch_add(1, Ordering::Relaxed); - self.probe.pass_gate(); if self.probe.failing.load(Ordering::Relaxed) { anyhow::bail!("the platform cannot attest right now"); } @@ -1479,8 +1440,8 @@ pNs85uhOZE8z2jr8Pg== "/nonexistent/nvattest", ), app_root_signing_key: SigningKey::from_slice(&DUMMY_K256_KEY).ok(), - identity: RwLock::new(None), - identity_last_failure: Mutex::new(None), + identity: Default::default(), + identity_last_failure: Default::default(), attest_lock: Arc::default(), // Read the same way production does, so a test comparing v1 `Info` // against v0 `get_info` compares like with like. @@ -2098,44 +2059,6 @@ pNs85uhOZE8z2jr8Pg== quote.await.unwrap().unwrap(); } - #[tokio::test(flavor = "current_thread")] - async fn identity_retry_is_off_executor_singleflight_and_cancellation_safe() { - let probe = Arc::new(InfoAttestationProbe::failing()); - let (state, _guard) = setup_test_state_with_probe(probe.clone()).await; - state.identity().await.unwrap_err(); - *state.inner.identity_last_failure.lock().unwrap() = - Some(Instant::now() - IDENTITY_RETRY_INTERVAL); - probe.set_failing(false); - let (entered_rx, release_tx) = probe.arm_gate(); - let first_state = state.clone(); - let first = tokio::spawn(async move { first_state.identity().await }); - tokio::time::timeout(Duration::from_secs(2), entered_rx) - .await - .unwrap() - .unwrap(); - first.abort(); - assert!(first.await.unwrap_err().is_cancelled()); - let mut followers = Vec::new(); - for _ in 0..16 { - let state = state.clone(); - followers.push(tokio::spawn(async move { state.identity().await })); - } - tokio::time::sleep(Duration::from_millis(20)).await; - assert_eq!(probe.calls(), 2, "only one retry may enter the platform"); - assert!(followers.iter().all(|task| !task.is_finished())); - release_tx.send(()).unwrap(); - let identity = state.identity().await.unwrap(); - for task in followers { - let other = tokio::time::timeout(Duration::from_secs(2), task) - .await - .unwrap() - .unwrap() - .unwrap(); - assert!(Arc::ptr_eq(&identity, &other)); - } - assert_eq!(probe.calls(), 2); - } - /// A failed decode must be as cheap to repeat as a cached success is. /// `identity()` is reached from the anonymous `/prpc/v1/Info`, and each /// attempt is a hardware quote plus an RTMR replay under the global quote @@ -2198,7 +2121,7 @@ pNs85uhOZE8z2jr8Pg== .identity() .await .expect("a decoded identity is cached"); - assert!(Arc::ptr_eq(&identity, &again)); + assert!(std::ptr::eq(identity, again)); assert_eq!( probe.calls(), 2, diff --git a/dstack/guest-agent/src/rpc_service_v1.rs b/dstack/guest-agent/src/rpc_service_v1.rs index 4213d8110..354e94618 100644 --- a/dstack/guest-agent/src/rpc_service_v1.rs +++ b/dstack/guest-agent/src/rpc_service_v1.rs @@ -512,7 +512,7 @@ mod tests { let last = state.identity().await.unwrap(); assert!( - std::sync::Arc::ptr_eq(&first, &last), + std::ptr::eq(first, last), "identity was decoded again; every Info call is generating a quote" ); }