From ef13f3a1e0d4458618fd6c6c5735f06bb44c1bcf Mon Sep 17 00:00:00 2001
From: Martin Raszyk <martin.raszyk@dfinity.org>
Date: Wed, 20 Nov 2024 16:54:38 +0100
Subject: [PATCH 1/2] chore: stop HTTP gateway gracefully

---
 src/dfx/src/actors/pocketic_proxy.rs | 60 ++++++++++++++++++++--------
 1 file changed, 44 insertions(+), 16 deletions(-)

diff --git a/src/dfx/src/actors/pocketic_proxy.rs b/src/dfx/src/actors/pocketic_proxy.rs
index b41a04454f..81dec6fb41 100644
--- a/src/dfx/src/actors/pocketic_proxy.rs
+++ b/src/dfx/src/actors/pocketic_proxy.rs
@@ -284,23 +284,26 @@ fn pocketic_proxy_start_thread(
                         }
                     }
                 };
-            if let Err(e) = initialize_gateway(
+            let instance = match initialize_gateway(
                 format!("http://localhost:{port}").parse().unwrap(),
                 replica_url.clone(),
                 domains.clone(),
                 address,
                 logger.clone(),
             ) {
-                error!(logger, "Failed to initialize HTTP gateway: {e:#}");
-                let _ = child.kill();
-                let _ = child.wait();
-                if receiver.try_recv().is_ok() {
-                    debug!(logger, "Got signal to stop.");
-                    break;
-                } else {
-                    continue;
+                Err(e) => {
+                    error!(logger, "Failed to initialize HTTP gateway: {e:#}");
+                    let _ = child.kill();
+                    let _ = child.wait();
+                    if receiver.try_recv().is_ok() {
+                        debug!(logger, "Got signal to stop.");
+                        break;
+                    } else {
+                        continue;
+                    }
                 }
-            }
+                Ok(i) => i,
+            };
             info!(logger, "Replica API running on {address}");
 
             // Send PocketIcProxyReadySignal to PocketIcProxy.
@@ -314,6 +317,9 @@ fn pocketic_proxy_start_thread(
                         logger,
                         "Got signal to stop. Killing pocket-ic gateway process..."
                     );
+                    if let Err(e) = shutdown_pocketic_proxy(port, instance, logger.clone()) {
+                        error!(logger, "Error shutting down PocketIC gracefully: {e}");
+                    }
                     let _ = child.kill();
                     let _ = child.wait();
                     break;
@@ -349,7 +355,7 @@ async fn initialize_gateway(
     domains: Option<Vec<String>>,
     addr: SocketAddr,
     logger: Logger,
-) -> DfxResult {
+) -> DfxResult<usize> {
     use pocket_ic::common::rest::{
         CreateHttpGatewayResponse, HttpGatewayBackend, HttpGatewayConfig,
     };
@@ -369,11 +375,12 @@ async fn initialize_gateway(
         .await?
         .error_for_status()?;
     let resp = resp.json::<CreateHttpGatewayResponse>().await?;
-    if let CreateHttpGatewayResponse::Error { message } = resp {
-        bail!("Gateway init error: {message}")
-    }
+    let instance = match resp {
+        CreateHttpGatewayResponse::Created(info) => info.instance_id,
+        CreateHttpGatewayResponse::Error { message } => bail!("Gateway init error: {message}"),
+    };
     info!(logger, "Initialized HTTP gateway.");
-    Ok(())
+    Ok(instance)
 }
 
 #[cfg(not(unix))]
@@ -383,6 +390,27 @@ fn initialize_gateway(
     _: Option<Vec<String>>,
     _: SocketAddr,
     _: Logger,
-) -> DfxResult {
+) -> DfxResult<usize> {
     bail!("PocketIC gateway not supported on this platform")
 }
+
+#[cfg(unix)]
+#[tokio::main(flavor = "current_thread")]
+async fn shutdown_pocketic_proxy(port: u16, instance: usize, logger: Logger) -> DfxResult {
+    use reqwest::Client;
+    let shutdown_client = Client::new();
+    debug!(logger, "Sending shutdown request to HTTP gateway");
+    shutdown_client
+        .post(format!(
+            "http://localhost:{port}/http_gateway/{instance}/stop"
+        ))
+        .send()
+        .await?
+        .error_for_status()?;
+    Ok(())
+}
+
+#[cfg(not(unix))]
+fn shutdown_pocketic(_: u16, _: usize, _: Logger) -> DfxResult {
+    bail!("PocketIC not supported on this platform")
+}

From 7e652d81736f59ea7b153cebde32f9a04c2e0e8f Mon Sep 17 00:00:00 2001
From: Martin Raszyk <martin.raszyk@dfinity.org>
Date: Wed, 20 Nov 2024 18:25:01 +0100
Subject: [PATCH 2/2] windows

---
 src/dfx/src/actors/pocketic_proxy.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dfx/src/actors/pocketic_proxy.rs b/src/dfx/src/actors/pocketic_proxy.rs
index 81dec6fb41..8d5cc16526 100644
--- a/src/dfx/src/actors/pocketic_proxy.rs
+++ b/src/dfx/src/actors/pocketic_proxy.rs
@@ -411,6 +411,6 @@ async fn shutdown_pocketic_proxy(port: u16, instance: usize, logger: Logger) ->
 }
 
 #[cfg(not(unix))]
-fn shutdown_pocketic(_: u16, _: usize, _: Logger) -> DfxResult {
+fn shutdown_pocketic_proxy(_: u16, _: usize, _: Logger) -> DfxResult {
     bail!("PocketIC not supported on this platform")
 }