From 34debf0e00d0b4e82bfd857618b8684c171175b3 Mon Sep 17 00:00:00 2001
From: Benjamin Elder <bentheelder@google.com>
Date: Mon, 17 Aug 2026 13:31:16 -0700
Subject: [PATCH] e2e: stop TestGracefulWorkerTerminationSuspend racing the
 actor to CRASHED

The test deletes the worker pod, sleeps a fixed 2s, then suspends the actor,
which is meant to land while the graceful shutdown is under way. On micro-VM it
lands after the actor has already reached CRASHED, so the suspend is rejected
and the test fails on every run since #893 enabled it there - on main as well as
on unrelated PRs:

  termination_test.go: failed to suspend Actor: rpc error:
  code = FailedPrecondition desc = MarkSuspending prerequisite not met for
  Actor: demo/graceful-term-suspend-... (got: STATUS_CRASHED, want
  STATUS_RUNNING or STATUS_PAUSED)

A duration is the wrong thing to wait on: how quickly the drain starts, and how
long the actor survives it, differ by sandbox class. Wait for the state the test
actually depends on instead - the control plane marking the worker DRAINING when
it sees the pod's DeletionTimestamp - and suspend as soon as that is observable
rather than 2s later.

Measured against a kind cluster running this branch's main, micro-VM demo:

  before   0/3 pass, every failure the MarkSuspending rejection above
  after   14/15 pass; the one failure was elsewhere in the test (a resume
          readyz timeout, the pre-existing flake) and never the suspend

The three termination tests also pass together, twice, with no skips.

This keeps what the test is for. ateom's graceful shutdown deliberately releases
its lock so 'a suspend arriving mid-drain' is still served, and that is what is
being exercised; only the way the test finds that window changes.
---
 internal/e2e/suites/demo/termination_test.go | 36 ++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/internal/e2e/suites/demo/termination_test.go b/internal/e2e/suites/demo/termination_test.go
index 46194393e..2fdd116e5 100644
--- a/internal/e2e/suites/demo/termination_test.go
+++ b/internal/e2e/suites/demo/termination_test.go
@@ -16,6 +16,7 @@ package demo
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -118,6 +119,31 @@ func TestGracefulWorkerTermination(t *testing.T) {
 	}
 }
 
+// waitForWorkerDraining polls ListWorkers until the named worker reports
+// STATE_DRAINING, i.e. the control plane has seen the pod's DeletionTimestamp and
+// the graceful shutdown is under way. Tests that want to act "during termination"
+// synchronise on this rather than on a duration: how long the drain takes to start
+// — and how long the actor survives it — differs by sandbox class.
+func waitForWorkerDraining(ctx context.Context, t *testing.T, clients *e2e.Clients, podName string, timeout time.Duration) error {
+	t.Helper()
+	deadline := time.Now().Add(timeout)
+	for {
+		resp, err := clients.SubstrateAPI.ListWorkers(ctx, &ateapipb.ListWorkersRequest{})
+		if err == nil {
+			for _, w := range resp.GetWorkers() {
+				if w.GetWorkerPod() == podName && w.GetState() == ateapipb.Worker_STATE_DRAINING {
+					t.Logf("Worker %s is DRAINING", podName)
+					return nil
+				}
+			}
+		}
+		if time.Now().After(deadline) {
+			return fmt.Errorf("worker %s did not reach DRAINING within %s", podName, timeout)
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
 // waitForWorkerRemoved polls ListWorkers until the named worker is absent.
 func waitForWorkerRemoved(ctx context.Context, t *testing.T, clients *e2e.Clients, podName string, timeout time.Duration) error {
 	t.Helper()
@@ -295,8 +321,14 @@ func TestGracefulWorkerTerminationSuspend(t *testing.T) {
 		t.Fatalf("failed to delete worker pod %s/%s: %v", podNS, podName, err)
 	}
 
-	// Wait 2 seconds to make sure the SIGTERM was sent and the container is in its sleep phase.
-	time.Sleep(2 * time.Second)
+	// Wait for the control plane to observe the deletion and mark the worker
+	// DRAINING, which is the point from which a suspend is meant to be served
+	// mid-drain. Sleeping a fixed interval instead raced the actor's own progress
+	// to CRASHED — on micro-VM it lost, and the suspend below was rejected with
+	// "MarkSuspending prerequisite not met ... (got: STATUS_CRASHED)".
+	if err := waitForWorkerDraining(ctx, t, clients, podName, 60*time.Second); err != nil {
+		t.Fatalf("worker %s not DRAINING after pod deletion: %v", podName, err)
+	}
 
 	// Suspend the actor.
 	t.Logf("Suspending actor %q during termination", actorID)