diff --git a/dstack/kms/src/main.rs b/dstack/kms/src/main.rs index 478c5c52a..74921a698 100644 --- a/dstack/kms/src/main.rs +++ b/dstack/kms/src/main.rs @@ -63,13 +63,18 @@ async fn run_onboard_service(kms_config: KmsConfig, figment: Figment) -> Result< // Remove section tls - let _ = rocket::custom(figment) + let rocket = rocket::custom(figment) .mount("/", rocket::routes![index, finish]) .mount( "/prpc", ra_rpc::prpc_routes!(OnboardState, OnboardHandler, trim: "Onboard."), ) - .manage(state) + .manage(state.clone()) + .ignite() + .await + .map_err(|err| anyhow!(err.to_string()))?; + state.set_shutdown(rocket.shutdown())?; + let _ = rocket .launch() .await .map_err(|err| anyhow!(err.to_string()))?; diff --git a/dstack/kms/src/onboard_service.rs b/dstack/kms/src/onboard_service.rs index 334a6c6db..58651833a 100644 --- a/dstack/kms/src/onboard_service.rs +++ b/dstack/kms/src/onboard_service.rs @@ -2,7 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use anyhow::{bail, Context, Result}; use dstack_kms_rpc::{ @@ -45,6 +45,7 @@ pub struct OnboardState { config: KmsConfig, attestation_verifier: Arc, bootstrap_lock: Arc>, + shutdown: Arc>, } impl OnboardState { @@ -57,8 +58,17 @@ impl OnboardState { config, attestation_verifier, bootstrap_lock: Arc::new(AsyncMutex::new(())), + shutdown: Arc::new(OnceLock::new()), }) } + + /// Hand the Rocket shutdown handle to the service so `finish` can stop the + /// server after its response has been sent. + pub fn set_shutdown(&self, shutdown: rocket::Shutdown) -> Result<()> { + self.shutdown + .set(shutdown) + .map_err(|_| anyhow::anyhow!("onboard shutdown handle is already set")) + } } pub struct OnboardHandler { @@ -199,7 +209,15 @@ impl OnboardRpc for OnboardHandler { } async fn finish(self) -> anyhow::Result<()> { - std::process::exit(0); + let shutdown = self + .state + .shutdown + .get() + .context("onboard shutdown handle is unavailable")?; + // Graceful shutdown lets Rocket finish sending this response before the + // server stops, so the client learns that onboarding succeeded. + shutdown.clone().notify(); + Ok(()) } }