Merge pull request #869 from Dstack-TEE/codex/fix-gateway-exit-response-order

fix(gateway): support graceful and forced exit

Author: kvinwang

dstack/gateway/rpc/proto/gateway_rpc.proto

@@ -394,13 +394,18 @@ message GetNodeStatusesResponse {
   repeated NodeStatusEntry statuses = 1;
 }
 
+message ExitRequest {
+  // Exit immediately without waiting for in-flight requests to drain.
+  bool force = 1;
+}
+
 service Admin {
   // Get the status of the gateway.
   rpc Status(google.protobuf.Empty) returns (StatusResponse) {}
   // Find Proxied HostInfo by instance ID
   rpc GetInfo(GetInfoRequest) returns (GetInfoResponse) {}
   // Exit the Gateway process.
-  rpc Exit(google.protobuf.Empty) returns (google.protobuf.Empty) {}
+  rpc Exit(ExitRequest) returns (google.protobuf.Empty) {}
   // Renew the proxy TLS certificate if certbot is enabled
   rpc RenewCert(google.protobuf.Empty) returns (RenewCertResponse) {}
   // Reload the proxy TLS certificate from files

dstack/gateway/src/admin_service.rs

@@ -10,7 +10,7 @@ use dstack_gateway_rpc::{
     admin_server::{AdminRpc, AdminServer},
     CertAttestationInfo, CertbotConfigResponse, ClearInstancePortPolicyRequest,
     CreateDnsCredentialRequest, DeleteDnsCredentialRequest, DeleteZtDomainRequest,
-    DnsCredentialInfo, ForceReleaseCertLockRequest, GetDefaultDnsCredentialResponse,
+    DnsCredentialInfo, ExitRequest, ForceReleaseCertLockRequest, GetDefaultDnsCredentialResponse,
     GetDnsCredentialRequest, GetInfoRequest, GetInfoResponse, GetInstanceHandshakesRequest,
     GetInstanceHandshakesResponse, GetInstancePortPolicyRequest, GetInstancePortPolicyResponse,
     GetMetaResponse, GetNodeStatusesResponse, GetZtDomainRequest, GlobalConnectionsStats,
@@ -82,8 +82,8 @@ impl AdminRpcHandler {
 }
 
 impl AdminRpc for AdminRpcHandler {
-    async fn exit(self) -> Result<()> {
-        self.state.lock().exit();
+    async fn exit(self, request: ExitRequest) -> Result<()> {
+        self.state.lock().exit(request.force)
     }
 
     async fn renew_cert(self) -> Result<RenewCertResponse> {

dstack/gateway/src/main.rs

@@ -347,15 +347,19 @@ async fn main() -> Result<()> {
             } else {
                 tracing::info!("admin server authentication enabled");
             }
-            rocket::custom(admin_figment)
+            let admin_rocket = rocket::custom(admin_figment)
                 .attach(auth_fairing)
                 .mount("/", admin_auth::routes())
                 .mount("/", web_routes::routes())
                 .mount("/", prpc!(Proxy, AdminRpcHandler, trim: "Admin."))
                 .mount("/prpc", prpc!(Proxy, AdminRpcHandler, trim: "Admin."))
-                .manage(admin_state)
-                .launch()
-                .await
+                .manage(admin_state.clone())
+                .ignite()
+                .await?;
+            admin_state
+                .lock()
+                .set_admin_shutdown(admin_rocket.shutdown());
+            admin_rocket.launch().await
         } else {
             std::future::pending().await
         }

dstack/gateway/src/main_service.rs

@@ -112,6 +112,7 @@ pub(crate) struct ProxyState {
     /// Reference to KvStore for syncing changes
     kv_store: Arc<KvStore>,
     handshake_cache: Arc<LatestHandshakesCache>,
+    admin_shutdown: Option<rocket::Shutdown>,
 }
 
 /// Options for creating a Proxy instance
@@ -244,6 +245,7 @@ impl ProxyInner {
             state,
             kv_store: kv_store.clone(),
             handshake_cache: handshake_cache.clone(),
+            admin_shutdown: None,
         });
         let auth_client = AuthClient::new(config.auth.clone());
         // Bootstrap WaveKV first if sync is enabled, so certbot can load certs from peers
@@ -1324,8 +1326,21 @@ impl ProxyState {
         Ok(())
     }
 
-    pub(crate) fn exit(&mut self) -> ! {
-        std::process::exit(0);
+    pub(crate) fn set_admin_shutdown(&mut self, shutdown: rocket::Shutdown) {
+        self.admin_shutdown = Some(shutdown);
+    }
+
+    pub(crate) fn exit(&self, force: bool) -> Result<()> {
+        if force {
+            std::process::exit(0);
+        }
+
+        let shutdown = self
+            .admin_shutdown
+            .as_ref()
+            .context("admin server shutdown handle is not initialized")?;
+        shutdown.notify();
+        Ok(())
     }
 
     pub(crate) fn refresh_state(&mut self) -> Result<()> {