Refactored the worker eligibility check in workflow_resume.go, to also check the worker eligibility again if previous attempt failed after setting actor to RESUMING state.

If the worker becomes ineligible, set the actor to crashed state.

Author: zoez7

cmd/ateapi/internal/controlapi/functional_test.go

@@ -1731,7 +1731,8 @@ func TestUpdateActor_NotFound(t *testing.T) {
 // TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible verifies that
 // a worker claimed by a failed resume attempt is released back to the free
 // pool if, by the next resume attempt, the actor's worker_selector has
-// changed such that the worker's pool is no longer eligible.
+// changed such that the worker's pool is no longer eligible. The actor
+// itself is crashed rather than transparently migrated to another pool.
 // Workflow:
 //  1. Creates pool-a (tier=a) and pool-b (tier=b), and an actor narrowed to
 //     tier=a.
@@ -1740,8 +1741,9 @@ func TestUpdateActor_NotFound(t *testing.T) {
 //     worker is claimed, leaving worker-a's actor_id set and the actor
 //     stuck in RESUMING.
 //  3. Updates the actor's selector to tier=b, making pool-a ineligible.
-//  4. Resumes again; asserts it succeeds onto worker-b, and that worker-a
-//     has been released (actor_id cleared) rather than left dangling.
+//  4. Resumes again; asserts it fails and the actor is CRASHED, that
+//     worker-a has been released (actor_id cleared) rather than left
+//     dangling, and that worker-b was never claimed.
 func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T) {
 	ns := namespaceForTest("ns-resume-release-stale")
 	tc := setupTest(t, ns)
@@ -1780,19 +1782,16 @@ func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T)
 		t.Fatalf("UpdateActor failed: %v", err)
 	}
 
-	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}}); err != nil {
-		t.Fatalf("second ResumeActor failed: %v", err)
+	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}}); err == nil {
+		t.Fatalf("expected second ResumeActor to fail: the assigned worker's pool is no longer eligible")
 	}
 
 	getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{Actor: &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}})
 	if err != nil {
 		t.Fatalf("GetActor failed: %v", err)
 	}
-	if got := getResp.GetWorkerPoolName(); got != "pool-b" {
-		t.Errorf("expected actor to land on pool-b, got worker_pool_name=%q", got)
-	}
-	if got := getResp.GetStatus(); got != ateapipb.Actor_STATUS_RUNNING {
-		t.Errorf("expected actor status RUNNING, got %v", got)
+	if got := getResp.GetStatus(); got != ateapipb.Actor_STATUS_CRASHED {
+		t.Errorf("expected actor status CRASHED, got %v", got)
 	}
 
 	listResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
@@ -1813,16 +1812,12 @@ func TestResumeActor_ReleasesStaleWorkerWhenPoolBecomesIneligible(t *testing.T)
 				t.Errorf("expected worker-a (now-ineligible pool-a) to be released, got actor_id=%q", got)
 			}
 		case "pool-b":
-			if wass := w.Assignment; wass == nil {
-				t.Errorf("expected worker-b to be claimed by %q, got nil assignment", name)
-			} else {
-				if wact := wass.Actor; wact == nil {
-					t.Errorf("expected worker-b to be claimed by %q, got nil assignment.actor", name)
-				} else {
-					if got := wact.Name; got != name {
-						t.Errorf("expected worker-b to be claimed by %q, got actor_id=%q", name, got)
-					}
+			if wass := w.Assignment; wass != nil {
+				got := "<nil-actor>"
+				if wass.Actor != nil {
+					got = wass.Actor.Name
 				}
+				t.Errorf("expected worker-b to remain unclaimed (actor is crashed, not migrated), got actor_id=%q", got)
 			}
 		}
 	}

cmd/ateapi/internal/controlapi/workflow_pause.go

@@ -118,16 +118,16 @@ func (s *CallAteletPauseStep) CheckPrerequisite(ctx context.Context, input *Paus
 	if state.Actor.GetStatus() != ateapipb.Actor_STATUS_PAUSING {
 		return status.Errorf(codes.FailedPrecondition, "CallAteletPauseStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING)
 	}
-	return nil
-}
-func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	if state.Actor.GetAteomPodNamespace() == "" || state.Actor.GetAteomPodName() == "" {
 		if err := crashActor(ctx, s.store, state.Actor.GetMetadata().GetAtespace(), state.Actor.GetMetadata().GetName()); err != nil {
 			slog.Error("Failed to crash actor", slog.String("err", err.Error()))
 		}
-		return fmt.Errorf("actor is CRASHED because it was in PAUSING state but has no active worker")
+		return status.Errorf(codes.FailedPrecondition, "CallAteletPauseStep prerequisite not met for Actor: %s. AteomPodNamespace: %s, GetAteomPodName %s", input.ActorName, state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName())
 	}
+	return nil
+}
 
+func (s *CallAteletPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error {
 	ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName())
 	if err != nil {
 		if errors.Is(err, ErrWorkerPodNotFound) {

cmd/ateapi/internal/controlapi/workflow_pause_test.go

@@ -24,11 +24,10 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-// TestPauseActorWorkflow exercises the pause workflow end-to-end against
-// seeded actor statuses, covering both the rejected and the idempotent-success
-// paths. The atelet dialer is nil, so any step that unexpectedly reaches it
-// panics.
-func TestPauseActorWorkflow(t *testing.T) {
+// TestPauseActorWorkflow_RejectedAndIdempotentPaths covers the two
+// short-circuit paths of the pause workflow: rejection by MarkPausingStep's
+// CheckPrerequisite and the IsComplete idempotent fast-forward.
+func TestPauseActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) {
 	tests := []struct {
 		name       string
 		seedStatus ateapipb.Actor_Status
@@ -136,9 +135,38 @@ func TestPauseSteps_CheckPrerequisite(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			ctx := context.Background()
 			for _, st := range allActorStatuses {
-				err := tc.step.CheckPrerequisite(ctx, &PauseInput{ActorName: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st}})
+				// Worker pod fields are populated so CallAteletPauseStep's
+				// missing-worker crash branch is not taken; this test only
+				// verifies status gating.
+				err := tc.step.CheckPrerequisite(ctx, &PauseInput{ActorName: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st, AteomPodNamespace: "ns", AteomPodName: "worker-1"}})
 				assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st])
 			}
 		})
 	}
 }
+
+// TestPauseActor_CrashesWhenPausingActorMissingWorkerPod verifies that a
+// PAUSING actor with no worker pod recorded is moved to CRASHED by
+// CallAteletPauseStep's prerequisite check and the pause fails with
+// FailedPrecondition.
+func TestPauseActor_CrashesWhenPausingActorMissingWorkerPod(t *testing.T) {
+	ctx := context.Background()
+	st, cleanup := storetest.SetupTestStore(t)
+	defer cleanup()
+	w := newTestActorWorkflow(t, st, "ns", "tmpl1")
+
+	seedWorkflowActor(t, ctx, st, "team-a", "id1", "ns", "tmpl1", ateapipb.Actor_STATUS_PAUSING)
+
+	_, err := w.PauseActor(ctx, "team-a", "id1")
+	if got := status.Code(err); got != codes.FailedPrecondition {
+		t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err)
+	}
+
+	got, err := st.GetActor(ctx, "team-a", "id1")
+	if err != nil {
+		t.Fatalf("GetActor failed: %v", err)
+	}
+	if got.GetStatus() != ateapipb.Actor_STATUS_CRASHED {
+		t.Errorf("stored status = %v, want %v", got.GetStatus(), ateapipb.Actor_STATUS_CRASHED)
+	}
+}

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -48,6 +48,7 @@ type ResumeInput struct {
 // ResumeState holds the mutable state loaded and modified during execution.
 type ResumeState struct {
 	Actor         *ateapipb.Actor
+	Worker        *ateapipb.Worker
 	ActorTemplate *atev1alpha1.ActorTemplate
 }
 
@@ -80,6 +81,36 @@ func (s *LoadActorForResumeStep) Execute(ctx context.Context, input *ResumeInput
 	}
 	state.ActorTemplate = actorTemplate
 
+	// If the Actor is in Resuming state, it means a previous attempt crashed after AssignWorkerStep.
+	// We don't need to repeat the AssignWorkerStep, load the Worker now.
+	if actor.Status == ateapipb.Actor_STATUS_RESUMING {
+		allPopulated := actor.AteomPodUid != "" && actor.WorkerPoolName != "" && actor.AteomPodName != ""
+		if !allPopulated {
+			slog.ErrorContext(ctx, "expected all of AteomPodUid, WorkerPoolName and AteomPodName to be populated, found",
+				slog.String("AteomPodUid", actor.AteomPodUid),
+				slog.String("WorkerPoolName", actor.WorkerPoolName),
+				slog.String("AteomPodName", actor.AteomPodName))
+
+			// Crash the actor if its worker assignment is corrupted. We should never be in this state.
+			if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil {
+				return cerr
+			}
+			return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName)
+		}
+
+		wk, err := s.store.GetWorker(ctx, actor.AteomPodNamespace, actor.WorkerPoolName, actor.AteomPodName)
+		if err != nil {
+			// Crash the actor if it was assigned to a deleted pod.
+			if errors.Is(err, ErrWorkerPodNotFound) {
+				if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil {
+					return cerr
+				}
+				return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName)
+			}
+			return fmt.Errorf("failed to get already assigned worker for actor %w", err)
+		}
+		state.Worker = wk
+	}
 	return nil
 }
 
@@ -116,19 +147,16 @@ type AssignWorkerStep struct {
 func (s *AssignWorkerStep) Name() string { return "AssignWorker" }
 
 func (s *AssignWorkerStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) {
-	// Only RUNNING is past this step. RESUMING intentionally re-runs because
-	// a retry must be able to release a stale worker whose pool became
-	// ineligible and pick a fresh one.
-	return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil
+	// RESUMING means a previous attempt already assigned a worker (loaded by
+	// LoadActorForResumeStep); RUNNING is past this step entirely.
+	return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RESUMING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil
 }
 func (s *AssignWorkerStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error {
-	// The resume edge exists from SUSPENDED and PAUSED.
-	// RESUMING is allowed for retrying this step.
 	switch state.Actor.GetStatus() {
-	case ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED, ateapipb.Actor_STATUS_RESUMING:
+	case ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED:
 		return nil
 	default:
-		return status.Errorf(codes.FailedPrecondition, "AssignWorkerStep prerequisite not met for Actor: %s (got: %v, want %s, %s or %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED, ateapipb.Actor_STATUS_RESUMING)
+		return status.Errorf(codes.FailedPrecondition, "AssignWorkerStep prerequisite not met for Actor: %s (got: %v, want %s or %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED)
 	}
 }
 
@@ -138,13 +166,9 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat
 		return fmt.Errorf("while listing workers: %w", err)
 	}
 
-	var assignedWorker *ateapipb.Worker
-
 	// Check if we already have a worker assigned from a previous failed attempt.
-	// If that worker's pool is no longer eligible (e.g. the actor's
-	// worker_selector was updated after the failed attempt), release it back
-	// to the free pool instead of leaving it claimed forever — nothing else
-	// reclaims a healthy worker whose actor moved on to a different pool.
+	// This can happen if ateapi crashed after updating worker with actor assignment,
+	// but has not yet updated the actor.
 	for _, worker := range workers {
 		if worker.Assignment == nil {
 			continue
@@ -157,20 +181,28 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat
 			return fmt.Errorf("while checking worker eligibility: %w", err)
 		}
 		if eligible {
-			assignedWorker = worker
+			state.Worker = worker
 			break
 		}
 		// Workers() returns pointers directly from the cache so we need to clone before
 		// mutating so that the cache is not corrupted if UpdateWorker fails.
-		release := proto.Clone(worker).(*ateapipb.Worker)
-		release.Assignment = nil
-		if err := s.store.UpdateWorker(ctx, release, release.Version); err != nil {
-			return fmt.Errorf("while releasing stale worker assignment: %w", err)
-		}
+		releaseWorker := proto.Clone(worker).(*ateapipb.Worker)
+		releaseWorker.Assignment = nil
+		// The claimed worker is no longer eligible (e.g. the actor's
+		// worker_selector changed after the failed attempt); release it back
+		// to the free pool — nothing else reclaims a healthy worker whose
+		// actor moved on to a different pool. Best effort in the background.
+		go func(release *ateapipb.Worker) {
+			bgCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+			defer cancel()
+			if err := s.store.UpdateWorker(bgCtx, release, release.Version); err != nil {
+				slog.ErrorContext(bgCtx, "Failed to release stale worker assignment",
+					slog.String("worker", release.GetWorkerNamespace()+"/"+release.GetWorkerPod()),
+					slog.Any("err", err))
+			}
+		}(releaseWorker)
 	}
-
-	// If not, find a free one using randomized shuffling
-	if assignedWorker == nil {
+	if state.Worker == nil {
 		pickedWorker, err := s.findFreeWorker(workers, state.ActorTemplate.Spec.SandboxClass, state.ActorTemplate.Spec.WorkerSelector, state.Actor.GetWorkerSelector(), state.Actor.GetLatestSnapshotInfo().GetLocal().GetNodeVmsWithLocalSnapshots())
 		if err != nil {
 			return err
@@ -179,14 +211,14 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat
 			return status.Errorf(codes.FailedPrecondition, "no free workers available")
 		}
 
-		assignedWorker = pickedWorker
+		state.Worker = pickedWorker
 		slog.InfoContext(ctx, "Picked worker", slog.Any("worker", pickedWorker.String()))
 	}
 
 	// Workers() returns pointers directly from the cache so we need to clone before
 	// mutating so that the cache is not corrupted if UpdateWorker fails.
-	assignedWorker = proto.Clone(assignedWorker).(*ateapipb.Worker)
-	assignedWorker.Assignment = &ateapipb.Assignment{
+	state.Worker = proto.Clone(state.Worker).(*ateapipb.Worker)
+	state.Worker.Assignment = &ateapipb.Assignment{
 		ActorTemplate: &ateapipb.KubeNamespacedObjectRef{
 			Namespace: state.Actor.GetActorTemplateNamespace(),
 			Name:      state.Actor.GetActorTemplateName(),
@@ -197,16 +229,16 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat
 		},
 	}
 
-	if err := s.store.UpdateWorker(ctx, assignedWorker, assignedWorker.Version); err != nil {
+	if err := s.store.UpdateWorker(ctx, state.Worker, state.Worker.Version); err != nil {
 		return err
 	}
 
 	state.Actor.Status = ateapipb.Actor_STATUS_RESUMING
-	state.Actor.AteomPodNamespace = assignedWorker.GetWorkerNamespace()
-	state.Actor.AteomPodName = assignedWorker.GetWorkerPod()
-	state.Actor.AteomPodIp = assignedWorker.GetIp()
-	state.Actor.AteomPodUid = assignedWorker.GetWorkerPodUid()
-	state.Actor.WorkerPoolName = assignedWorker.GetWorkerPool()
+	state.Actor.AteomPodNamespace = state.Worker.GetWorkerNamespace()
+	state.Actor.AteomPodName = state.Worker.GetWorkerPod()
+	state.Actor.AteomPodIp = state.Worker.GetIp()
+	state.Actor.AteomPodUid = state.Worker.GetWorkerPodUid()
+	state.Actor.WorkerPoolName = state.Worker.GetWorkerPool()
 
 	updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion())
 	if err != nil {
@@ -275,6 +307,40 @@ func (s *CallAteletRestoreStep) CheckPrerequisite(ctx context.Context, input *Re
 	if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING {
 		return status.Errorf(codes.FailedPrecondition, "CallAteletRestoreStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING)
 	}
+	if state.Worker == nil {
+		return status.Errorf(codes.FailedPrecondition, "Assigned worker is nil")
+	}
+	// Verify if the worker is still assigned to the same Actor.
+	assigned := state.Worker.GetAssignment().GetActor()
+	if assigned.GetAtespace() != input.Atespace || assigned.GetName() != input.ActorName {
+		slog.ErrorContext(ctx, "crashing actor because its assigned worker no longer belongs to it",
+			slog.String("worker", state.Worker.GetWorkerPod()),
+			slog.Any("assignment", state.Worker.GetAssignment()))
+		if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil {
+			return fmt.Errorf("while crashing actor: %w", cerr)
+		}
+		return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName)
+	}
+	eligible, err := isWorkerEligibleForActor(state.Worker, state.ActorTemplate.Spec.SandboxClass, state.ActorTemplate.Spec.WorkerSelector, state.Actor.GetWorkerSelector())
+	if err != nil {
+		return fmt.Errorf("while calling isWorkerEligbleForActor :%w", err)
+	}
+	if !eligible {
+		slog.ErrorContext(ctx, "crashing actor because previously assigned worker is not eligible anymore")
+		release := proto.Clone(state.Worker).(*ateapipb.Worker)
+		release.Assignment = nil
+		// If that worker's pool is no longer eligible (e.g. the actor's
+		// worker_selector was updated after the failed attempt), release it back
+		// to the free pool instead of leaving it claimed forever — nothing else
+		// reclaims a healthy worker whose actor moved on to a different pool.
+		if err := s.store.UpdateWorker(ctx, release, release.Version); err != nil {
+			return fmt.Errorf("while releasing stale worker assignment: %w", err)
+		}
+		if cerr := crashActor(ctx, s.store, input.Atespace, input.ActorName); cerr != nil {
+			return fmt.Errorf("while crashing actor: %w", cerr)
+		}
+		return status.Errorf(codes.Aborted, "actor %s crashed", input.ActorName)
+	}
 	return nil
 }
 func (s *CallAteletRestoreStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error {