From d7c2a202fca6832a1949fe96a31e420e32b221aa Mon Sep 17 00:00:00 2001
From: Ben Duncan <5178243+benduncan@users.noreply.github.com>
Date: Mon, 14 Sep 2026 15:38:22 -0500
Subject: [PATCH] fix: two flaky tests failing the Race Detection job (#1030)

* fix: stop the process-group test misreading a reaped child as alive

The test polled until the child was gone, then checked once more. A task
being reaped shows state X for an instant, which processRunning counted as
running, so that second look failed about 1 run in 50 under -race.

* fix: wait for the recovery retry loop before removing a test's data dir

restoreInstances starts the retry loop untracked, and it rewrites the local
state file when it ends. A write landing during t.TempDir cleanup failed
TestRestoreInstances_* with RemoveAll: directory not empty.
---
 spinifex/daemon/daemon.go      |  8 ++++++--
 spinifex/daemon/state_test.go  |  1 +
 spinifex/utils/command_test.go | 27 ++++++++++++++++-----------
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go
index 8d9c5d733..cc00d1b54 100644
--- a/spinifex/daemon/daemon.go
+++ b/spinifex/daemon/daemon.go
@@ -174,6 +174,10 @@ type Daemon struct {
 	cancel                context.CancelFunc
 	shutdownWg            sync.WaitGroup
 
+	// recoveryRetryWg tracks the post-restore retry loop, which writes the
+	// local state file when it finishes. Tests wait on it before cleanup.
+	recoveryRetryWg sync.WaitGroup
+
 	// systemDispatchWg tracks in-flight system.LaunchInstance / TerminateInstance
 	// handlers. Each runs in its own goroutine so a slow VM boot never blocks
 	// the NATS subscription. Used by tests to await dispatch completion.
@@ -2451,12 +2455,12 @@ func (d *Daemon) restoreInstances() error {
 		slog.Error("Failed to persist local state after restore", "error", err)
 	}
 
-	go func() {
+	d.recoveryRetryWg.Go(func() {
 		d.vmMgr.RunRecoveryRetryLoop(d.shuttingDown.Load)
 		if err := d.WriteState(); err != nil {
 			slog.Error("Failed to persist local state after recovery retry", "error", err)
 		}
-	}()
+	})
 
 	return nil
 }
diff --git a/spinifex/daemon/state_test.go b/spinifex/daemon/state_test.go
index 266bc2469..f4d86c623 100644
--- a/spinifex/daemon/state_test.go
+++ b/spinifex/daemon/state_test.go
@@ -65,6 +65,7 @@ func createDaemonWithJetStream(t *testing.T) *Daemon {
 	// Registered after t.TempDir() so it runs first (LIFO): goroutines
 	// spawned by MarkFailed/MarkRecoveryFailed finish before RemoveAll runs.
 	t.Cleanup(daemon.vmMgr.WaitForBackgroundWork)
+	t.Cleanup(daemon.recoveryRetryWg.Wait)
 
 	return daemon
 }
diff --git a/spinifex/utils/command_test.go b/spinifex/utils/command_test.go
index 219de90cb..e9430ace7 100644
--- a/spinifex/utils/command_test.go
+++ b/spinifex/utils/command_test.go
@@ -27,11 +27,7 @@ func TestRunCommandWithTimeoutKillsProcessGroup(t *testing.T) {
 	}
 
 	childPID := childPIDFromOutput(t, output)
-	deadline := time.Now().Add(time.Second)
-	for processRunning(childPID) && time.Now().Before(deadline) {
-		time.Sleep(10 * time.Millisecond)
-	}
-	if processRunning(childPID) {
+	if !processExitsWithin(childPID, time.Second) {
 		t.Errorf("child process %d survived command timeout", childPID)
 	}
 }
@@ -47,11 +43,7 @@ func TestRunCommandWithTimeoutKillsDetachedOutputHolder(t *testing.T) {
 	}
 
 	childPID := childPIDFromOutput(t, output)
-	deadline := time.Now().Add(time.Second)
-	for processRunning(childPID) && time.Now().Before(deadline) {
-		time.Sleep(10 * time.Millisecond)
-	}
-	if processRunning(childPID) {
+	if !processExitsWithin(childPID, time.Second) {
 		t.Errorf("detached child process %d survived command cleanup", childPID)
 	}
 }
@@ -85,11 +77,24 @@ func childPIDFromOutput(t *testing.T, output string) int {
 	return 0
 }
 
+// processExitsWithin polls until pid is gone. One not-running observation is
+// final: a re-check can catch a reaped task's brief "X" state and misreport it.
+func processExitsWithin(pid int, timeout time.Duration) bool {
+	deadline := time.Now().Add(timeout)
+	for processRunning(pid) {
+		if time.Now().After(deadline) {
+			return false
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	return true
+}
+
 func processRunning(pid int) bool {
 	data, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
 	if err == nil {
 		fields := strings.Fields(string(data))
-		return len(fields) < 3 || fields[2] != "Z"
+		return len(fields) < 3 || (fields[2] != "Z" && fields[2] != "X")
 	}
 	if os.IsNotExist(err) {
 		return false