tests: bind panic in runtime

runtime/Cargo.toml

@@ -44,7 +44,7 @@ rocket = "0.5.0-rc.2"
 
 [dependencies.shuttle-common]
 workspace = true
-features = ["service"]
+features = ["backend", "service"]
 
 [dependencies.shuttle-proto]
 workspace = true

runtime/src/legacy/mod.rs

@@ -38,7 +38,7 @@ use tonic::{
     Request, Response, Status,
 };
 use tower::ServiceBuilder;
-use tracing::{error, instrument, trace};
+use tracing::{error, info, trace};
 use uuid::Uuid;
 
 use crate::{provisioner_factory::ProvisionerFactory, Logger};
@@ -221,13 +221,56 @@ where
         let (kill_tx, kill_rx) = tokio::sync::oneshot::channel();
         *self.kill_tx.lock().unwrap() = Some(kill_tx);
 
+        let stopped_tx = self.stopped_tx.clone();
+
+        let handle = tokio::runtime::Handle::current();
+
         // start service as a background task with a kill receiver
-        tokio::spawn(run_until_stopped(
-            service,
-            service_address,
-            self.stopped_tx.clone(),
-            kill_rx,
-        ));
+        tokio::spawn(async move {
+            let mut background = handle.spawn(service.bind(service_address));
+
+            tokio::select! {
+                res = &mut background => {
+                    match res {
+                        Ok(_) => {
+                            info!("service stopped all on its own");
+                            stopped_tx.send((StopReason::End, String::new())).unwrap();
+                        },
+                        Err(error) => {
+                            if error.is_panic() {
+                                let panic = error.into_panic();
+                                let msg = panic.downcast_ref::<&str>()
+                                    .map(|x| x.to_string())
+                                    .unwrap_or_else(|| "<no panic message>".to_string());
+
+                                error!(error = msg, "service panicked");
+
+                                stopped_tx
+                                    .send((StopReason::Crash, msg))
+                                    .unwrap();
+                            } else {
+                                error!(%error, "service crashed");
+                                stopped_tx
+                                    .send((StopReason::Crash, error.to_string()))
+                                    .unwrap();
+                            }
+                        },
+                    }
+                },
+                message = kill_rx => {
+                    match message {
+                        Ok(_) => {
+                            stopped_tx.send((StopReason::Request, String::new())).unwrap();
+                        }
+                        Err(_) => trace!("the sender dropped")
+                    };
+
+                    info!("will now abort the service");
+                    background.abort();
+                    background.await.unwrap().expect("to stop service");
+                }
+            }
+        });
 
         let message = StartResponse { success: true };
 
@@ -297,35 +340,3 @@ where
         }
     }
 }
-
-/// Run the service until a stop signal is received
-#[instrument(skip(service, stopped_tx, kill_rx))]
-async fn run_until_stopped(
-    // service: LoadedService,
-    service: impl Service,
-    addr: SocketAddr,
-    stopped_tx: tokio::sync::broadcast::Sender<(StopReason, String)>,
-    kill_rx: tokio::sync::oneshot::Receiver<String>,
-) {
-    trace!("starting deployment on {}", &addr);
-    tokio::select! {
-        res = service.bind(addr) => {
-            match res {
-                Ok(_) => {
-                    stopped_tx.send((StopReason::End, String::new())).unwrap();
-                }
-                Err(error) => {
-                    stopped_tx.send((StopReason::Crash, error.to_string())).unwrap();
-                }
-            }
-        },
-        message = kill_rx => {
-            match message {
-                Ok(_) => {
-                    stopped_tx.send((StopReason::Request, String::new())).unwrap();
-                }
-                Err(_) => trace!("the sender dropped")
-            };
-        }
-    }
-}

runtime/src/next/mod.rs

@@ -426,7 +426,7 @@ pub mod tests {
             .unwrap();
     }
 
-    #[tokio::test]
+    #[tokio::test(flavor = "multi_thread")]
     async fn axum() {
         compile_module();
 

runtime/tests/integration/loader.rs

@@ -1,15 +1,9 @@
-use std::time::Duration;
-
-use shuttle_proto::runtime::{LoadRequest, StartRequest};
+use shuttle_proto::runtime::{LoadRequest, StartRequest, StopReason, SubscribeStopRequest};
 use uuid::Uuid;
 
 use crate::helpers::{spawn_runtime, TestRuntime};
 
-/// This test does panic, but the panic happens in a spawned task inside the project runtime,
-/// so we get this output: `thread 'tokio-runtime-worker' panicked at 'panic in bind', src/main.rs:6:9`,
-/// but `should_panic(expected = "panic in bind")` doesn't catch it.
 #[tokio::test]
-#[should_panic(expected = "panic in bind")]
 async fn bind_panic() {
     let project_path = format!("{}/tests/resources/bind-panic", env!("CARGO_MANIFEST_DIR"));
 
@@ -29,20 +23,24 @@ async fn bind_panic() {
 
     let _ = runtime_client.load(load_request).await.unwrap();
 
+    let mut stream = runtime_client
+        .subscribe_stop(tonic::Request::new(SubscribeStopRequest {}))
+        .await
+        .unwrap()
+        .into_inner();
+
     let start_request = StartRequest {
         deployment_id: Uuid::default().as_bytes().to_vec(),
         ip: runtime_address.to_string(),
     };
 
-    // I also tried this without spawning, but it gave the same result. Panic but it isn't caught.
-    tokio::spawn(async move {
-        runtime_client
-            .start(tonic::Request::new(start_request))
-            .await
-            .unwrap();
-        // Give it a second to panic.
-        tokio::time::sleep(Duration::from_secs(1)).await;
-    })
-    .await
-    .unwrap();
+    runtime_client
+        .start(tonic::Request::new(start_request))
+        .await
+        .unwrap();
+
+    let reason = stream.message().await.unwrap().unwrap();
+
+    assert_eq!(reason.reason, StopReason::Crash as i32);
+    assert_eq!(reason.message, "panic in bind");
 }