Wait for test server to shutdown (#11485)

Author: emilk

Cargo.toml

@@ -647,6 +647,7 @@ iter_on_empty_collections = "warn"
 iter_on_single_items = "warn"
 iter_without_into_iter = "warn"
 large_digit_groups = "warn"
+large_futures = "warn"
 large_include_file = "warn"
 large_stack_arrays = "warn"
 large_stack_frames = "warn"
@@ -738,6 +739,7 @@ unnecessary_semicolon = "warn"
 unnecessary_struct_initialization = "warn"
 unnecessary_wraps = "warn"
 unnested_or_patterns = "warn"
+unused_async = "warn"
 unused_peekable = "warn"
 unused_rounding = "warn"
 unused_self = "warn"

crates/store/re_server/src/server.rs

@@ -39,7 +39,7 @@ pub struct Server {
 /// `ServerHandle` is a tiny helper abstraction that enables us to
 /// deal with the gRPC server lifecycle more easily.
 pub struct ServerHandle {
-    shutdown: Sender<()>,
+    shutdown: Option<Sender<()>>,
     ready: mpsc::Receiver<SocketAddr>,
     failed: mpsc::Receiver<String>,
 }
@@ -73,9 +73,12 @@ impl ServerHandle {
         self.failed.recv().await;
     }
 
-    /// Signal to the gRPC server to shutdown.
-    pub async fn shutdown(self) {
-        let _ = self.shutdown.send(());
+    /// Signal to the gRPC server to shutdown, and then wait for it.
+    pub async fn shutdown_and_wait(mut self) {
+        if let Some(shutdown) = self.shutdown.take() {
+            shutdown.send(()).ok();
+            self.wait_for_shutdown().await;
+        }
     }
 }
 
@@ -151,7 +154,7 @@ impl Server {
         });
 
         ServerHandle {
-            shutdown: shutdown_tx,
+            shutdown: Some(shutdown_tx),
             ready: ready_rx,
             failed: failed_rx,
         }

tests/rust/re_integration_test/src/lib.rs

@@ -70,7 +70,7 @@ impl Drop for TestServer {
             .expect("Server handle not initialized");
         tokio::task::block_in_place(move || {
             tokio::runtime::Handle::current().block_on(async move {
-                server_handle.shutdown().await;
+                server_handle.shutdown_and_wait().await;
             });
         });
     }