From c433be132f119e51b6f2e16d5480951a4aece490 Mon Sep 17 00:00:00 2001
From: Rain <rain@sunshowers.io>
Date: Sat, 14 Dec 2024 16:58:17 -0800
Subject: [PATCH] [nextest-runner] wait for dispatcher to acknowledge retries

Gives the dispatcher a chance to control communication.
---
 nextest-runner/src/runner/dispatcher.rs      | 26 +++++++++++++++++---
 nextest-runner/src/runner/executor.rs        | 13 ++++++++++
 nextest-runner/src/runner/internal_events.rs |  2 ++
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/nextest-runner/src/runner/dispatcher.rs b/nextest-runner/src/runner/dispatcher.rs
index 7d76d8583e3..ebe2a76243d 100644
--- a/nextest-runner/src/runner/dispatcher.rs
+++ b/nextest-runner/src/runner/dispatcher.rs
@@ -471,10 +471,28 @@ where
             InternalEvent::Executor(ExecutorEvent::RetryStarted {
                 test_instance,
                 retry_data,
-            }) => self.callback_none_response(TestEventKind::TestRetryStarted {
-                test_instance,
-                retry_data,
-            }),
+                tx,
+            }) => {
+                if self.cancel_state.is_some() {
+                    // The run has been cancelled: don't send a message over the tx and don't start
+                    // any new units.
+                    return HandleEventResponse::None;
+                }
+
+                match tx.send(()) {
+                    Ok(_) => {}
+                    Err(_) => {
+                        // The test task died?
+                        debug!(test = ?test_instance.id(), "test task died, ignoring");
+                        return HandleEventResponse::None;
+                    }
+                }
+
+                self.callback_none_response(TestEventKind::TestRetryStarted {
+                    test_instance,
+                    retry_data,
+                })
+            }
             InternalEvent::Executor(ExecutorEvent::Finished {
                 test_instance,
                 success_output,
diff --git a/nextest-runner/src/runner/executor.rs b/nextest-runner/src/runner/executor.rs
index 1d87ee7dac3..f9d051df1d6 100644
--- a/nextest-runner/src/runner/executor.rs
+++ b/nextest-runner/src/runner/executor.rs
@@ -230,10 +230,23 @@ impl<'a> ExecutorContext<'a> {
             // is empty.
 
             if retry_data.attempt > 1 {
+                // Ensure that the dispatcher believes the run is still ongoing. If the run is
+                // cancelled, the dispatcher will let us know.
+                let (tx, rx) = oneshot::channel();
                 _ = resp_tx.send(ExecutorEvent::RetryStarted {
                     test_instance,
                     retry_data,
+                    tx,
                 });
+
+                match rx.await {
+                    Ok(()) => {}
+                    Err(_) => {
+                        // The receiver was dropped -- the dispatcher has signaled that this unit
+                        // should exit.
+                        return;
+                    }
+                }
             }
 
             // Some of this information is only useful for event reporting, but
diff --git a/nextest-runner/src/runner/internal_events.rs b/nextest-runner/src/runner/internal_events.rs
index e92366a18dc..3d47783a015 100644
--- a/nextest-runner/src/runner/internal_events.rs
+++ b/nextest-runner/src/runner/internal_events.rs
@@ -89,6 +89,8 @@ pub(super) enum ExecutorEvent<'a> {
     RetryStarted {
         test_instance: TestInstance<'a>,
         retry_data: RetryData,
+        // This is used to indicate that the dispatcher still wants to run the test.
+        tx: oneshot::Sender<()>,
     },
     Finished {
         test_instance: TestInstance<'a>,