From 1795b424aae2babbd6b80d77617f9f7832042b16 Mon Sep 17 00:00:00 2001 From: zoezhao Date: Wed, 8 Jul 2026 17:26:15 -0700 Subject: [PATCH 1/3] Add a CheckPrerequisite step to each workflow --- cmd/ateapi/internal/controlapi/workflow.go | 15 ++ .../internal/controlapi/workflow_pause.go | 29 +++- .../controlapi/workflow_pause_test.go | 144 +++++++++++++++++ .../internal/controlapi/workflow_resume.go | 29 ++++ .../controlapi/workflow_resume_test.go | 132 ++++++++++++++++ .../internal/controlapi/workflow_suspend.go | 28 +++- .../controlapi/workflow_suspend_test.go | 145 ++++++++++++++++++ .../internal/controlapi/workflow_test.go | 106 +++++++++++++ .../controlapi/workflow_testutil_test.go | 100 ++++++++++++ 9 files changed, 718 insertions(+), 10 deletions(-) create mode 100644 cmd/ateapi/internal/controlapi/workflow_pause_test.go create mode 100644 cmd/ateapi/internal/controlapi/workflow_suspend_test.go create mode 100644 cmd/ateapi/internal/controlapi/workflow_test.go create mode 100644 cmd/ateapi/internal/controlapi/workflow_testutil_test.go diff --git a/cmd/ateapi/internal/controlapi/workflow.go b/cmd/ateapi/internal/controlapi/workflow.go index 5c193f992..0f3debfc0 100644 --- a/cmd/ateapi/internal/controlapi/workflow.go +++ b/cmd/ateapi/internal/controlapi/workflow.go @@ -44,6 +44,14 @@ type WorkflowStep[Params any, Context any] interface { // If it returns true, the engine skips Execute() and fast-forwards to the next step. IsComplete(ctx context.Context, params Params, wCtx Context) (bool, error) + // CheckPrerequisite validates that the current state permits executing this + // step (e.g. the actor's status allows this state-machine edge). The engine + // calls it only when IsComplete returned false, immediately before Execute, + // so completed steps of a retried workflow fast-forward without + // re-validation. Return a gRPC status error with + // codes.FailedPrecondition to abort the workflow if prereqs are not met. + CheckPrerequisite(ctx context.Context, params Params, wCtx Context) error + // Execute performs the step's business logic and persists any state changes. // If an error is returned, the workflow stops and relies on the client to retry. Execute(ctx context.Context, params Params, wCtx Context) error @@ -79,6 +87,13 @@ func RunWorkflow[Params any, Context any](ctx context.Context, params Params, wC continue } + if err := step.CheckPrerequisite(ctx, params, wCtx); err != nil { + span.RecordError(err) + span.SetStatus(codes.Error, err.Error()) + span.End() + return fmt.Errorf("prerequisite not met at step %s: %w", step.Name(), err) + } + err = runStep(ctx, params, wCtx, step) if err != nil { span.RecordError(err) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 14c28ca5f..93e6779bc 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -27,6 +27,8 @@ import ( atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "k8s.io/apimachinery/pkg/util/wait" ) @@ -52,6 +54,9 @@ func (s *LoadActorForPauseStep) IsComplete(ctx context.Context, input *PauseInpu // Always run to get the freshest state return false, nil } +func (s *LoadActorForPauseStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { + return nil +} func (s *LoadActorForPauseStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error { actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { @@ -79,11 +84,14 @@ func (s *MarkPausingStep) IsComplete(ctx context.Context, input *PauseInput, sta // Fast forward if we've already marked our intent or if we are further along. return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, nil } -func (s *MarkPausingStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error { +func (s *MarkPausingStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { + // The pause edge only exists from RUNNING; PAUSING/PAUSED are fast-forwarded by IsComplete. if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - return nil + return status.Errorf(codes.FailedPrecondition, "MarkPausingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) } - + return nil +} +func (s *MarkPausingStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error { state.Actor.Status = ateapipb.Actor_STATUS_PAUSING state.Actor.InProgressSnapshot = fmt.Sprintf("%s-%s-%s", state.Actor.GetMetadata().GetName(), time.Now().Format(time.RFC3339), rand.Text()) updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion()) @@ -106,6 +114,12 @@ func (s *CallAteletPauseStep) IsComplete(ctx context.Context, input *PauseInput, // If we are already PAUSED, we've already called Atelet return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, nil } +func (s *CallAteletPauseStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { + 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.ActorID, 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 { @@ -157,8 +171,13 @@ type FinalizePausedStep struct { func (s *FinalizePausedStep) Name() string { return "FinalizePaused" } func (s *FinalizePausedStep) IsComplete(ctx context.Context, input *PauseInput, state *PauseState) (bool, error) { - // The workflow is completely done ONLY if the status is PAUSED *and* we've successfully freed the worker. - return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED && state.Actor.GetAteomPodNamespace() == "", nil + return state.Actor.GetStatus() == ateapipb.Actor_STATUS_PAUSED, nil +} +func (s *FinalizePausedStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_PAUSING { + return status.Errorf(codes.FailedPrecondition, "FinalizePausedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING) + } + return nil } func (s *FinalizePausedStep) Execute(ctx context.Context, input *PauseInput, state *PauseState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause_test.go b/cmd/ateapi/internal/controlapi/workflow_pause_test.go new file mode 100644 index 000000000..ba089a777 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_pause_test.go @@ -0,0 +1,144 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "testing" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "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) { + tests := []struct { + name string + seedStatus ateapipb.Actor_Status + // wantErr true means PauseActor must fail with FailedPrecondition. + wantErr bool + // wantStatus is the stored status after the call. + wantStatus ateapipb.Actor_Status + }{ + { + // Pausing a SUSPENDED actor is rejected by MarkPausingStep's + // CheckPrerequisite and the actor's status is left untouched. + name: "not running rejected", + seedStatus: ateapipb.Actor_STATUS_SUSPENDED, + wantErr: true, + wantStatus: ateapipb.Actor_STATUS_SUSPENDED, + }, + { + // Pausing a PAUSED actor succeeds idempotently via IsComplete + // fast-forward without calling atelet. + name: "already paused succeeds", + seedStatus: ateapipb.Actor_STATUS_PAUSED, + wantStatus: ateapipb.Actor_STATUS_PAUSED, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(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", tc.seedStatus) + + actor, err := w.PauseActor(ctx, "team-a", "id1") + if tc.wantErr { + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err) + } + } else { + if err != nil { + t.Fatalf("PauseActor failed: %v", err) + } + if actor.GetStatus() != tc.wantStatus { + t.Errorf("returned status = %v, want %v", actor.GetStatus(), tc.wantStatus) + } + } + + got, err := st.GetActor(ctx, "team-a", "id1") + if err != nil { + t.Fatalf("GetActor failed: %v", err) + } + if got.GetStatus() != tc.wantStatus { + t.Errorf("stored status = %v, want %v", got.GetStatus(), tc.wantStatus) + } + }) + } +} + +// TestPauseSteps_CheckPrerequisite verifies each pause step's CheckPrerequisite +// against every actor status: nil for the step's allowed statuses, +// FailedPrecondition for all others. +func TestPauseSteps_CheckPrerequisite(t *testing.T) { + tests := []struct { + name string + step WorkflowStep[*PauseInput, *PauseState] + // allowed lists the statuses CheckPrerequisite accepts; nil means + // every status is accepted. + allowed map[ateapipb.Actor_Status]bool + }{ + { + // Loading has no prerequisite: it is allowed from every status. + name: "LoadActorForPauseStep", + step: &LoadActorForPauseStep{}, + allowed: nil, + }, + { + // Pausing is allowed only from RUNNING. + name: "MarkPausingStep", + step: &MarkPausingStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RUNNING: true, + }, + }, + { + // The checkpoint call is allowed only from PAUSING (PAUSED is + // fast-forwarded by IsComplete). + name: "CallAteletPauseStep", + step: &CallAteletPauseStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_PAUSING: true, + }, + }, + { + // Finalizing is allowed only from PAUSING: a persisted PAUSED + // actor always has its worker pod fields cleared and is + // fast-forwarded by IsComplete. + name: "FinalizePausedStep", + step: &FinalizePausedStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_PAUSING: true, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + for _, st := range allActorStatuses { + err := tc.step.CheckPrerequisite(ctx, &PauseInput{ActorID: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st}}) + assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) + } + }) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 7cd768ac4..79f4cdf05 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -61,6 +61,9 @@ func (s *LoadActorForResumeStep) IsComplete(ctx context.Context, input *ResumeIn // Always run this step to get the latest state from the DB return false, nil } +func (s *LoadActorForResumeStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + return nil +} func (s *LoadActorForResumeStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { @@ -113,8 +116,22 @@ 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 } +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: + return nil + default: + return status.Errorf(codes.FailedPrecondition, "AssignWorkerStep prerequisite not met for Actor: %s (got: %v, want %s, %s or %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED, ateapipb.Actor_STATUS_RESUMING) + } +} + func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { workers, err := s.workerCache.Workers() if err != nil { @@ -254,6 +271,12 @@ func (s *CallAteletRestoreStep) Name() string { return "CallAteletRestore" } func (s *CallAteletRestoreStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) { return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil } +func (s *CallAteletRestoreStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + 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.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) + } + return nil +} func (s *CallAteletRestoreStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName()) if err != nil { @@ -358,6 +381,12 @@ func (s *FinalizeRunningStep) Name() string { return "FinalizeRunning" } func (s *FinalizeRunningStep) IsComplete(ctx context.Context, input *ResumeInput, state *ResumeState) (bool, error) { return state.Actor.GetStatus() == ateapipb.Actor_STATUS_RUNNING, nil } +func (s *FinalizeRunningStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING { + return status.Errorf(codes.FailedPrecondition, "FinalizeRunningStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) + } + return nil +} func (s *FinalizeRunningStep) Execute(ctx context.Context, input *ResumeInput, state *ResumeState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index b87bd2305..2b380573e 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -15,10 +15,14 @@ package controlapi import ( + "context" "testing" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) @@ -162,3 +166,131 @@ func TestIsWorkerEligibleForActor(t *testing.T) { }) } } + +// TestResumeActorWorkflow exercises the resume workflow end-to-end against +// seeded actor statuses, covering both the rejected and the idempotent-success +// paths. The worker cache and atelet dialer are nil, so any step that +// unexpectedly reaches them panics. +func TestResumeActorWorkflow(t *testing.T) { + tests := []struct { + name string + seedStatus ateapipb.Actor_Status + // wantErr true means ResumeActor must fail with FailedPrecondition. + wantErr bool + // wantStatus is the stored status after the call. + wantStatus ateapipb.Actor_Status + }{ + { + // The resume edge only exists from SUSPENDED, PAUSED, and + // RESUMING; a CRASHED actor is rejected by AssignWorkerStep's + // CheckPrerequisite and its status is left untouched. + name: "crashed rejected", + seedStatus: ateapipb.Actor_STATUS_CRASHED, + wantErr: true, + wantStatus: ateapipb.Actor_STATUS_CRASHED, + }, + { + // Resuming a RUNNING actor succeeds idempotently: every step + // fast-forwards via IsComplete. + name: "already running succeeds", + seedStatus: ateapipb.Actor_STATUS_RUNNING, + wantStatus: ateapipb.Actor_STATUS_RUNNING, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(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", tc.seedStatus, func(a *ateapipb.Actor) { + a.AteomPodNamespace = "wns" + a.AteomPodName = "wpod" + a.AteomPodIp = "1.2.3.4" + a.AteomPodUid = "uid" + a.WorkerPoolName = "pool1" + }) + + actor, err := w.ResumeActor(ctx, "team-a", "id1", false) + if tc.wantErr { + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err) + } + } else { + if err != nil { + t.Fatalf("ResumeActor failed: %v", err) + } + if actor.GetStatus() != tc.wantStatus { + t.Errorf("returned status = %v, want %v", actor.GetStatus(), tc.wantStatus) + } + } + + got, err := st.GetActor(ctx, "team-a", "id1") + if err != nil { + t.Fatalf("GetActor failed: %v", err) + } + if got.GetStatus() != tc.wantStatus { + t.Errorf("stored status = %v, want %v", got.GetStatus(), tc.wantStatus) + } + }) + } +} + +// TestResumeSteps_CheckPrerequisite verifies each resume step's +// CheckPrerequisite against every actor status: nil for the step's allowed +// statuses, FailedPrecondition for all others. +func TestResumeSteps_CheckPrerequisite(t *testing.T) { + tests := []struct { + name string + step WorkflowStep[*ResumeInput, *ResumeState] + // allowed lists the statuses CheckPrerequisite accepts; nil means + // every status is accepted. + allowed map[ateapipb.Actor_Status]bool + }{ + { + // Loading has no prerequisite: it is allowed from every status. + name: "LoadActorForResumeStep", + step: &LoadActorForResumeStep{}, + allowed: nil, + }, + { + // Resuming is allowed from SUSPENDED, PAUSED, and RESUMING + // (retry of this step). + name: "AssignWorkerStep", + step: &AssignWorkerStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDED: true, + ateapipb.Actor_STATUS_PAUSED: true, + ateapipb.Actor_STATUS_RESUMING: true, + }, + }, + { + // The restore call is allowed only from RESUMING (RUNNING is + // fast-forwarded by IsComplete). + name: "CallAteletRestoreStep", + step: &CallAteletRestoreStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RESUMING: true, + }, + }, + { + // Finalizing transitions RESUMING -> RUNNING; RUNNING itself is + // fast-forwarded by IsComplete before the prerequisite is checked. + name: "FinalizeRunningStep", + step: &FinalizeRunningStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RESUMING: true, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + for _, st := range allActorStatuses { + err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorID: "id1"}, &ResumeState{Actor: &ateapipb.Actor{Status: st}}) + assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) + } + }) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index cacb8c307..79fc7c158 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -28,6 +28,8 @@ import ( atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" "k8s.io/apimachinery/pkg/util/wait" ) @@ -53,6 +55,9 @@ func (s *LoadActorForSuspendStep) IsComplete(ctx context.Context, input *Suspend // Always run to get the freshest state return false, nil } +func (s *LoadActorForSuspendStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + return nil +} func (s *LoadActorForSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { actor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) if err != nil { @@ -80,11 +85,13 @@ func (s *MarkSuspendingStep) IsComplete(ctx context.Context, input *SuspendInput // Fast forward if we've already marked our intent or if we are further along. return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDING || state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil } -func (s *MarkSuspendingStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { +func (s *MarkSuspendingStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - return nil + return status.Errorf(codes.FailedPrecondition, "MarkSuspendingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) } - + return nil +} +func (s *MarkSuspendingStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { state.Actor.Status = ateapipb.Actor_STATUS_SUSPENDING snapshotID := time.Now().Format(time.RFC3339) + "-" + rand.Text() state.Actor.InProgressSnapshot = strings.TrimSuffix(state.ActorTemplate.Spec.SnapshotsConfig.Location, "/") + "/" + input.ActorName + "/" + snapshotID @@ -108,6 +115,12 @@ func (s *CallAteletSuspendStep) IsComplete(ctx context.Context, input *SuspendIn // If we are already SUSPENDED, we've already called Atelet return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil } +func (s *CallAteletSuspendStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { + return status.Errorf(codes.FailedPrecondition, "CallAteletSuspendStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + } + return nil +} func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetAteomPodNamespace() == "" || state.Actor.GetAteomPodName() == "" { if err := crashActor(ctx, s.store, state.Actor.GetMetadata().GetAtespace(), state.Actor.GetMetadata().GetName()); err != nil { @@ -159,8 +172,13 @@ type FinalizeSuspendedStep struct { func (s *FinalizeSuspendedStep) Name() string { return "FinalizeSuspended" } func (s *FinalizeSuspendedStep) IsComplete(ctx context.Context, input *SuspendInput, state *SuspendState) (bool, error) { - // The workflow is completely done ONLY if the status is SUSPENDED *and* we've successfully freed the worker. - return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED && state.Actor.GetAteomPodNamespace() == "", nil + return state.Actor.GetStatus() == ateapipb.Actor_STATUS_SUSPENDED, nil +} +func (s *FinalizeSuspendedStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { + if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { + return status.Errorf(codes.FailedPrecondition, "FinalizeSuspendedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + } + return nil } func (s *FinalizeSuspendedStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { latestActor, err := s.store.GetActor(ctx, input.Atespace, input.ActorName) diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go new file mode 100644 index 000000000..b8edf96c1 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -0,0 +1,145 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "testing" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// TestSuspendActorWorkflow exercises the suspend 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 TestSuspendActorWorkflow(t *testing.T) { + tests := []struct { + name string + seedStatus ateapipb.Actor_Status + // wantErr true means SuspendActor must fail with FailedPrecondition. + wantErr bool + // wantStatus is the stored status after the call. + wantStatus ateapipb.Actor_Status + }{ + { + // The state machine's PAUSED->SUSPENDED commit edge is rejected + // (suspending needs a live worker to checkpoint from) and the + // actor's status is left untouched. + name: "paused rejected", + seedStatus: ateapipb.Actor_STATUS_PAUSED, + wantErr: true, + wantStatus: ateapipb.Actor_STATUS_PAUSED, + }, + { + // Suspending a SUSPENDED actor succeeds idempotently via + // IsComplete fast-forward without calling atelet. + name: "already suspended succeeds", + seedStatus: ateapipb.Actor_STATUS_SUSPENDED, + wantStatus: ateapipb.Actor_STATUS_SUSPENDED, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(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", tc.seedStatus) + + actor, err := w.SuspendActor(ctx, "team-a", "id1") + if tc.wantErr { + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err) + } + } else { + if err != nil { + t.Fatalf("SuspendActor failed: %v", err) + } + if actor.GetStatus() != tc.wantStatus { + t.Errorf("returned status = %v, want %v", actor.GetStatus(), tc.wantStatus) + } + } + + got, err := st.GetActor(ctx, "team-a", "id1") + if err != nil { + t.Fatalf("GetActor failed: %v", err) + } + if got.GetStatus() != tc.wantStatus { + t.Errorf("stored status = %v, want %v", got.GetStatus(), tc.wantStatus) + } + }) + } +} + +// TestSuspendSteps_CheckPrerequisite verifies each suspend step's +// CheckPrerequisite against every actor status: nil for the step's allowed +// statuses, FailedPrecondition for all others. +func TestSuspendSteps_CheckPrerequisite(t *testing.T) { + tests := []struct { + name string + step WorkflowStep[*SuspendInput, *SuspendState] + // allowed lists the statuses CheckPrerequisite accepts; nil means + // every status is accepted. + allowed map[ateapipb.Actor_Status]bool + }{ + { + // Loading has no prerequisite: it is allowed from every status. + name: "LoadActorForSuspendStep", + step: &LoadActorForSuspendStep{}, + allowed: nil, + }, + { + // Suspending is allowed only from RUNNING. + name: "MarkSuspendingStep", + step: &MarkSuspendingStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_RUNNING: true, + }, + }, + { + // The checkpoint call is allowed only from SUSPENDING (SUSPENDED + // is fast-forwarded by IsComplete). + name: "CallAteletSuspendStep", + step: &CallAteletSuspendStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDING: true, + }, + }, + { + // Finalizing is allowed only from SUSPENDING: a persisted + // SUSPENDED actor always has its worker pod fields cleared and is + // fast-forwarded by IsComplete. + name: "FinalizeSuspendedStep", + step: &FinalizeSuspendedStep{}, + allowed: map[ateapipb.Actor_Status]bool{ + ateapipb.Actor_STATUS_SUSPENDING: true, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + for _, st := range allActorStatuses { + err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorID: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st}}) + assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) + } + }) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_test.go b/cmd/ateapi/internal/controlapi/workflow_test.go new file mode 100644 index 000000000..f4a81b728 --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_test.go @@ -0,0 +1,106 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "testing" + + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "k8s.io/apimachinery/pkg/util/wait" +) + +// stubStep records the order of engine callbacks and returns canned results. +type stubStep struct { + name string + complete bool + prereqErr error + executeErr error + calls *[]string +} + +func (s *stubStep) Name() string { return s.name } +func (s *stubStep) IsComplete(ctx context.Context, params struct{}, wCtx struct{}) (bool, error) { + *s.calls = append(*s.calls, s.name+".IsComplete") + return s.complete, nil +} +func (s *stubStep) CheckPrerequisite(ctx context.Context, params struct{}, wCtx struct{}) error { + *s.calls = append(*s.calls, s.name+".CheckPrerequisite") + return s.prereqErr +} +func (s *stubStep) Execute(ctx context.Context, params struct{}, wCtx struct{}) error { + *s.calls = append(*s.calls, s.name+".Execute") + return s.executeErr +} +func (s *stubStep) RetryBackoff() *wait.Backoff { return nil } + +func TestRunWorkflow_CheckPrerequisiteOrdering(t *testing.T) { + ctx := context.Background() + + t.Run("prerequisite checked after IsComplete and before Execute", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", calls: &calls}, + } + if err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + want := []string{"s1.IsComplete", "s1.CheckPrerequisite", "s1.Execute"} + if len(calls) != len(want) { + t.Fatalf("calls = %v, want %v", calls, want) + } + for i := range want { + if calls[i] != want[i] { + t.Fatalf("calls = %v, want %v", calls, want) + } + } + }) + + t.Run("prerequisite skipped when step is complete", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", complete: true, prereqErr: status.Error(codes.FailedPrecondition, "must not be called"), calls: &calls}, + } + if err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps); err != nil { + t.Fatalf("RunWorkflow: %v", err) + } + for _, c := range calls { + if c == "s1.CheckPrerequisite" || c == "s1.Execute" { + t.Fatalf("unexpected call %q; calls = %v", c, calls) + } + } + }) + + t.Run("prerequisite error aborts before Execute and preserves status code", func(t *testing.T) { + var calls []string + steps := []WorkflowStep[struct{}, struct{}]{ + &stubStep{name: "s1", prereqErr: status.Error(codes.FailedPrecondition, "nope"), calls: &calls}, + &stubStep{name: "s2", calls: &calls}, + } + err := RunWorkflow(ctx, struct{}{}, struct{}{}, steps) + if err == nil { + t.Fatal("RunWorkflow: expected error, got nil") + } + if got := status.Code(err); got != codes.FailedPrecondition { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.FailedPrecondition, err) + } + for _, c := range calls { + if c == "s1.Execute" || c == "s2.IsComplete" { + t.Fatalf("unexpected call %q after failed prerequisite; calls = %v", c, calls) + } + } + }) +} diff --git a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go new file mode 100644 index 000000000..6522a865b --- /dev/null +++ b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go @@ -0,0 +1,100 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package controlapi + +import ( + "context" + "slices" + "testing" + + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" + atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" + listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1" + "github.com/agent-substrate/substrate/pkg/proto/ateapipb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/client-go/tools/cache" +) + +// newTestActorWorkflow builds an ActorWorkflow backed by the given store and a +// lister serving one minimal ActorTemplate. Dependencies the unit tests never +// reach (worker cache, atelet dialer, k8s clients) are nil, so a step that +// unexpectedly executes against them fails the test loudly. +func newTestActorWorkflow(t *testing.T, st store.Interface, tmplNamespace, tmplName string) *ActorWorkflow { + t.Helper() + indexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{}) + if err := indexer.Add(&atev1alpha1.ActorTemplate{ + ObjectMeta: metav1.ObjectMeta{Namespace: tmplNamespace, Name: tmplName}, + }); err != nil { + t.Fatalf("add template to indexer: %v", err) + } + return NewActorWorkflow(st, nil, nil, listersv1alpha1.NewActorTemplateLister(indexer), nil, nil, nil) +} + +// seedWorkflowActor stores an actor with the given status, bound to the given +// template (pass the same tmplNamespace/tmplName as newTestActorWorkflow). +// opts mutate the actor before it is stored. +func seedWorkflowActor(t *testing.T, ctx context.Context, st store.Interface, atespace, id, tmplNamespace, tmplName string, actorStatus ateapipb.Actor_Status, opts ...func(*ateapipb.Actor)) { + t.Helper() + actor := &ateapipb.Actor{ + ActorId: id, + Atespace: atespace, + Status: actorStatus, + ActorTemplateNamespace: tmplNamespace, + ActorTemplateName: tmplName, + } + for _, opt := range opts { + opt(actor) + } + if err := st.CreateActor(ctx, actor); err != nil { + t.Fatalf("seed actor: %v", err) + } +} + +// allActorStatuses enumerates every Actor_Status value, for exhaustive +// CheckPrerequisite table tests. It is derived from the generated enum map so +// statuses added to the proto are covered automatically. +var allActorStatuses = func() []ateapipb.Actor_Status { + nums := make([]int32, 0, len(ateapipb.Actor_Status_name)) + for n := range ateapipb.Actor_Status_name { + nums = append(nums, n) + } + slices.Sort(nums) + statuses := make([]ateapipb.Actor_Status, 0, len(nums)) + for _, n := range nums { + statuses = append(statuses, ateapipb.Actor_Status(n)) + } + return statuses +}() + +// assertPrerequisiteResult verifies a CheckPrerequisite outcome for one +// status: nil when allowed, FailedPrecondition otherwise. +func assertPrerequisiteResult(t *testing.T, st ateapipb.Actor_Status, err error, wantAllowed bool) { + t.Helper() + if wantAllowed { + if err != nil { + t.Errorf("status %v: CheckPrerequisite = %v, want nil", st, err) + } + return + } + if err == nil { + t.Errorf("status %v: CheckPrerequisite = nil, want FailedPrecondition", st) + return + } + if got := status.Code(err); got != codes.FailedPrecondition { + t.Errorf("status %v: status.Code = %v, want %v", st, got, codes.FailedPrecondition) + } +} From 4104a7d4d0633ce7fb8bf91da7e4cf00bddab9c7 Mon Sep 17 00:00:00 2001 From: zoezhao Date: Wed, 8 Jul 2026 17:38:13 -0700 Subject: [PATCH 2/3] Adopting the new Actor API changes. --- cmd/ateapi/internal/controlapi/workflow_pause.go | 6 +++--- cmd/ateapi/internal/controlapi/workflow_pause_test.go | 2 +- cmd/ateapi/internal/controlapi/workflow_resume.go | 6 +++--- cmd/ateapi/internal/controlapi/workflow_resume_test.go | 2 +- cmd/ateapi/internal/controlapi/workflow_suspend.go | 6 +++--- cmd/ateapi/internal/controlapi/workflow_suspend_test.go | 2 +- cmd/ateapi/internal/controlapi/workflow_testutil_test.go | 5 ++--- 7 files changed, 14 insertions(+), 15 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 93e6779bc..e8057d560 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause.go @@ -87,7 +87,7 @@ func (s *MarkPausingStep) IsComplete(ctx context.Context, input *PauseInput, sta func (s *MarkPausingStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { // The pause edge only exists from RUNNING; PAUSING/PAUSED are fast-forwarded by IsComplete. if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - return status.Errorf(codes.FailedPrecondition, "MarkPausingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) + return status.Errorf(codes.FailedPrecondition, "MarkPausingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) } return nil } @@ -116,7 +116,7 @@ func (s *CallAteletPauseStep) IsComplete(ctx context.Context, input *PauseInput, } func (s *CallAteletPauseStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { 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.ActorID, 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 } @@ -175,7 +175,7 @@ func (s *FinalizePausedStep) IsComplete(ctx context.Context, input *PauseInput, } func (s *FinalizePausedStep) CheckPrerequisite(ctx context.Context, input *PauseInput, state *PauseState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_PAUSING { - return status.Errorf(codes.FailedPrecondition, "FinalizePausedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING) + return status.Errorf(codes.FailedPrecondition, "FinalizePausedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_PAUSING) } return nil } diff --git a/cmd/ateapi/internal/controlapi/workflow_pause_test.go b/cmd/ateapi/internal/controlapi/workflow_pause_test.go index ba089a777..d270568ab 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_pause_test.go @@ -136,7 +136,7 @@ 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{ActorID: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st}}) + err := tc.step.CheckPrerequisite(ctx, &PauseInput{ActorName: "id1"}, &PauseState{Actor: &ateapipb.Actor{Status: st}}) assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) } }) diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 79f4cdf05..0a7f8115e 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -128,7 +128,7 @@ func (s *AssignWorkerStep) CheckPrerequisite(ctx context.Context, input *ResumeI case ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED, ateapipb.Actor_STATUS_RESUMING: return nil default: - return status.Errorf(codes.FailedPrecondition, "AssignWorkerStep prerequisite not met for Actor: %s (got: %v, want %s, %s or %s)", input.ActorID, 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, %s or %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED, ateapipb.Actor_STATUS_RESUMING) } } @@ -273,7 +273,7 @@ func (s *CallAteletRestoreStep) IsComplete(ctx context.Context, input *ResumeInp } func (s *CallAteletRestoreStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { 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.ActorID, 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) } return nil } @@ -383,7 +383,7 @@ func (s *FinalizeRunningStep) IsComplete(ctx context.Context, input *ResumeInput } func (s *FinalizeRunningStep) CheckPrerequisite(ctx context.Context, input *ResumeInput, state *ResumeState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING { - return status.Errorf(codes.FailedPrecondition, "FinalizeRunningStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) + return status.Errorf(codes.FailedPrecondition, "FinalizeRunningStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RESUMING) } return nil } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 2b380573e..f82859968 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -288,7 +288,7 @@ func TestResumeSteps_CheckPrerequisite(t *testing.T) { t.Run(tc.name, func(t *testing.T) { ctx := context.Background() for _, st := range allActorStatuses { - err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorID: "id1"}, &ResumeState{Actor: &ateapipb.Actor{Status: st}}) + err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorName: "id1"}, &ResumeState{Actor: &ateapipb.Actor{Status: st}}) assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) } }) diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index 79fc7c158..87214653f 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -87,7 +87,7 @@ func (s *MarkSuspendingStep) IsComplete(ctx context.Context, input *SuspendInput } func (s *MarkSuspendingStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING { - return status.Errorf(codes.FailedPrecondition, "MarkSuspendingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) + return status.Errorf(codes.FailedPrecondition, "MarkSuspendingStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_RUNNING) } return nil } @@ -117,7 +117,7 @@ func (s *CallAteletSuspendStep) IsComplete(ctx context.Context, input *SuspendIn } func (s *CallAteletSuspendStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { - return status.Errorf(codes.FailedPrecondition, "CallAteletSuspendStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + return status.Errorf(codes.FailedPrecondition, "CallAteletSuspendStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) } return nil } @@ -176,7 +176,7 @@ func (s *FinalizeSuspendedStep) IsComplete(ctx context.Context, input *SuspendIn } func (s *FinalizeSuspendedStep) CheckPrerequisite(ctx context.Context, input *SuspendInput, state *SuspendState) error { if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { - return status.Errorf(codes.FailedPrecondition, "FinalizeSuspendedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorID, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) + return status.Errorf(codes.FailedPrecondition, "FinalizeSuspendedStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) } return nil } diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index b8edf96c1..2ab879d62 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -137,7 +137,7 @@ func TestSuspendSteps_CheckPrerequisite(t *testing.T) { t.Run(tc.name, func(t *testing.T) { ctx := context.Background() for _, st := range allActorStatuses { - err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorID: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st}}) + err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorName: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st}}) assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) } }) diff --git a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go index 6522a865b..c6962694a 100644 --- a/cmd/ateapi/internal/controlapi/workflow_testutil_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_testutil_test.go @@ -50,8 +50,7 @@ func newTestActorWorkflow(t *testing.T, st store.Interface, tmplNamespace, tmplN func seedWorkflowActor(t *testing.T, ctx context.Context, st store.Interface, atespace, id, tmplNamespace, tmplName string, actorStatus ateapipb.Actor_Status, opts ...func(*ateapipb.Actor)) { t.Helper() actor := &ateapipb.Actor{ - ActorId: id, - Atespace: atespace, + Metadata: &ateapipb.ResourceMetadata{Name: id, Atespace: atespace}, Status: actorStatus, ActorTemplateNamespace: tmplNamespace, ActorTemplateName: tmplName, @@ -59,7 +58,7 @@ func seedWorkflowActor(t *testing.T, ctx context.Context, st store.Interface, at for _, opt := range opts { opt(actor) } - if err := st.CreateActor(ctx, actor); err != nil { + if _, err := st.CreateActor(ctx, actor); err != nil { t.Fatalf("seed actor: %v", err) } } From 6d976cf16f6d30589e2ce1e0674afa1c3fd9f3f5 Mon Sep 17 00:00:00 2001 From: zoezhao Date: Sat, 11 Jul 2026 17:28:11 -0700 Subject: [PATCH 3/3] 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. --- .../internal/controlapi/functional_test.go | 33 +-- .../internal/controlapi/workflow_pause.go | 8 +- .../controlapi/workflow_pause_test.go | 40 ++- .../internal/controlapi/workflow_resume.go | 130 ++++++--- .../controlapi/workflow_resume_test.go | 255 +++++++++++++++++- .../internal/controlapi/workflow_suspend.go | 7 +- .../controlapi/workflow_suspend_test.go | 39 ++- 7 files changed, 432 insertions(+), 80 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/functional_test.go b/cmd/ateapi/internal/controlapi/functional_test.go index 57c08a9c5..29e87acdf 100644 --- a/cmd/ateapi/internal/controlapi/functional_test.go +++ b/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 := "" + 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) } } } diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index aaa65bec9..8ecf797aa 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/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) { diff --git a/cmd/ateapi/internal/controlapi/workflow_pause_test.go b/cmd/ateapi/internal/controlapi/workflow_pause_test.go index d270568ab..176b8a52d 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause_test.go +++ b/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) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 95685a377..3f4bb30e9 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/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 { diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 0b797889d..32ba25171 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -25,6 +25,7 @@ import ( "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "google.golang.org/protobuf/proto" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" ) @@ -218,11 +219,89 @@ func TestAssignWorkerStep_SkipsWorkerAssignedInOtherAtespace(t *testing.T) { } } -// TestResumeActorWorkflow exercises the resume workflow end-to-end against -// seeded actor statuses, covering both the rejected and the idempotent-success -// paths. The worker cache and atelet dialer are nil, so any step that -// unexpectedly reaches them panics. -func TestResumeActorWorkflow(t *testing.T) { +// TestAssignWorkerStep_ReleasesIneligibleStaleWorkerInBackground verifies +// that a worker claimed by a previous failed attempt whose pool is no longer +// eligible is released back to the free pool asynchronously, without failing +// the resume, while a fresh eligible worker is assigned. +func TestAssignWorkerStep_ReleasesIneligibleStaleWorkerInBackground(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + // stale-pod is claimed by this actor from a failed attempt but its sandbox + // class no longer matches the template; free-pod is eligible and free. + stale := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool-a", + WorkerPod: "stale-pod", + SandboxClass: "microvm", + Assignment: &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "id1"}, + }, + } + free := &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool-b", + WorkerPod: "free-pod", + SandboxClass: "gvisor", + } + for _, w := range []*ateapipb.Worker{stale, free} { + if err := persistence.CreateWorker(ctx, w); err != nil { + t.Fatalf("CreateWorker(%s): %v", w.GetWorkerPod(), err) + } + } + + actor, err := persistence.CreateActor(ctx, &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "id1"}, + Status: ateapipb.Actor_STATUS_SUSPENDED, + }) + if err != nil { + t.Fatalf("CreateActor: %v", err) + } + + cacheCtx, cancel := context.WithCancel(ctx) + defer cancel() + wc := workercache.New(persistence, time.Minute) + if err := wc.Start(cacheCtx); err != nil { + t.Fatalf("workercache.Start: %v", err) + } + + step := &AssignWorkerStep{store: persistence, workerCache: wc} + state := &ResumeState{ + Actor: actor, + ActorTemplate: &atev1alpha1.ActorTemplate{ + Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}, + }, + } + if err := step.Execute(ctx, &ResumeInput{ActorName: "id1", Atespace: "team-a"}, state); err != nil { + t.Fatalf("Execute() error = %v, want nil (release must not fail the resume)", err) + } + + if got := state.Worker.GetWorkerPod(); got != "free-pod" { + t.Errorf("assigned worker = %q, want %q", got, "free-pod") + } + + // The stale worker is released in the background; poll until its + // assignment is cleared. + deadline := time.Now().Add(5 * time.Second) + for { + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool-a", "stale-pod") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if stored.GetAssignment() == nil { + break + } + if time.Now().After(deadline) { + t.Fatalf("stale worker still assigned after %v: %v", 5*time.Second, stored.GetAssignment()) + } + time.Sleep(10 * time.Millisecond) + } +} + +// TestResumeActorWorkflow_RejectedAndIdempotentPaths covers the two +// short-circuit paths of the resume workflow: rejection by AssignWorkerStep's +// CheckPrerequisite and the IsComplete idempotent fast-forward. +func TestResumeActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) { tests := []struct { name string seedStatus ateapipb.Actor_Status @@ -306,14 +385,13 @@ func TestResumeSteps_CheckPrerequisite(t *testing.T) { allowed: nil, }, { - // Resuming is allowed from SUSPENDED, PAUSED, and RESUMING - // (retry of this step). + // Resuming is allowed from SUSPENDED and PAUSED (RESUMING and + // RUNNING are fast-forwarded by IsComplete). name: "AssignWorkerStep", step: &AssignWorkerStep{}, allowed: map[ateapipb.Actor_Status]bool{ ateapipb.Actor_STATUS_SUSPENDED: true, ateapipb.Actor_STATUS_PAUSED: true, - ateapipb.Actor_STATUS_RESUMING: true, }, }, { @@ -339,9 +417,168 @@ func TestResumeSteps_CheckPrerequisite(t *testing.T) { t.Run(tc.name, func(t *testing.T) { ctx := context.Background() for _, st := range allActorStatuses { - err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorName: "id1"}, &ResumeState{Actor: &ateapipb.Actor{Status: st}}) + // An eligible Worker assigned to this actor is provided so + // CallAteletRestoreStep's worker checks pass; this test only + // verifies status gating. + state := &ResumeState{ + Actor: &ateapipb.Actor{Status: st}, + Worker: &ateapipb.Worker{ + SandboxClass: string(atev1alpha1.SandboxClassGvisor), + Assignment: &ateapipb.Assignment{Actor: &ateapipb.ObjectRef{Name: "id1"}}, + }, + ActorTemplate: &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}}, + } + err := tc.step.CheckPrerequisite(ctx, &ResumeInput{ActorName: "id1"}, state) assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) } }) } } + +// TestResumeActor_CrashesOnCorruptWorkerAssignment verifies that a RESUMING +// actor with only some of its worker assignment fields populated is moved to +// CRASHED by LoadActorForResumeStep and the resume fails with Aborted. +func TestResumeActor_CrashesOnCorruptWorkerAssignment(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_RESUMING, func(a *ateapipb.Actor) { + a.AteomPodName = "worker-1" // AteomPodUid and WorkerPoolName left empty + }) + + _, err := w.ResumeActor(ctx, "team-a", "id1", false) + if got := status.Code(err); got != codes.Aborted { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.Aborted, 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) + } +} + +// TestCallAteletRestoreStep_CheckPrerequisite_WorkerOwnership verifies that +// the restore prerequisite only proceeds on a worker whose assignment still +// names this actor: the recovery path loads the worker by pod name only, so +// the assignment may have been cleared and the worker re-claimed by another +// actor in the meantime. On a mismatch the actor is crashed and the worker — +// which is not ours — must not be written. +func TestCallAteletRestoreStep_CheckPrerequisite_WorkerOwnership(t *testing.T) { + ownAssignment := &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-a", Name: "shared"}, + } + otherAssignment := &ateapipb.Assignment{ + Actor: &ateapipb.ObjectRef{Atespace: "team-b", Name: "shared"}, + } + + tests := []struct { + name string + sandboxClass string + assignment *ateapipb.Assignment + // wantCode is codes.OK when CheckPrerequisite must return nil. + wantCode codes.Code + wantActorStatus ateapipb.Actor_Status + // wantAssignment is the assignment expected on the stored worker + // afterwards; wantWorkerWrite false additionally asserts the worker + // version did not move (no write at all). + wantAssignment *ateapipb.Assignment + wantWorkerWrite bool + }{ + { + name: "crashes actor and leaves worker untouched when assigned to another actor", + sandboxClass: "gvisor", + assignment: otherAssignment, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: otherAssignment, + }, + { + name: "crashes actor and leaves worker untouched when assignment is cleared", + sandboxClass: "gvisor", + assignment: nil, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: nil, + }, + { + name: "passes for own eligible worker", + sandboxClass: "gvisor", + assignment: ownAssignment, + wantCode: codes.OK, + wantActorStatus: ateapipb.Actor_STATUS_RESUMING, + wantAssignment: ownAssignment, + }, + { + name: "releases own ineligible worker and crashes actor", + sandboxClass: "microvm", + assignment: ownAssignment, + wantCode: codes.Aborted, + wantActorStatus: ateapipb.Actor_STATUS_CRASHED, + wantAssignment: nil, + wantWorkerWrite: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx := context.Background() + persistence := newTestPersistence(t) + + if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool", + WorkerPod: "pod-1", + SandboxClass: tt.sandboxClass, + Assignment: tt.assignment, + }); err != nil { + t.Fatalf("CreateWorker: %v", err) + } + // Re-fetch so state.Worker carries the stored version (needed by + // the release path's optimistic update). + seeded, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + + seedWorkflowActor(t, ctx, persistence, "team-a", "shared", "ns", "tmpl1", ateapipb.Actor_STATUS_RESUMING) + + step := &CallAteletRestoreStep{store: persistence} + state := &ResumeState{ + Actor: &ateapipb.Actor{ + Metadata: &ateapipb.ResourceMetadata{Atespace: "team-a", Name: "shared"}, + Status: ateapipb.Actor_STATUS_RESUMING, + }, + Worker: seeded, + ActorTemplate: &atev1alpha1.ActorTemplate{Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}}, + } + err = step.CheckPrerequisite(ctx, &ResumeInput{Atespace: "team-a", ActorName: "shared"}, state) + if got := status.Code(err); got != tt.wantCode { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, tt.wantCode, err) + } + + actor, err := persistence.GetActor(ctx, "team-a", "shared") + if err != nil { + t.Fatalf("GetActor: %v", err) + } + if actor.GetStatus() != tt.wantActorStatus { + t.Errorf("stored actor status = %v, want %v", actor.GetStatus(), tt.wantActorStatus) + } + + stored, err := persistence.GetWorker(ctx, "worker-ns", "pool", "pod-1") + if err != nil { + t.Fatalf("GetWorker: %v", err) + } + if !proto.Equal(stored.GetAssignment(), tt.wantAssignment) { + t.Errorf("stored worker assignment = %v, want %v", stored.GetAssignment(), tt.wantAssignment) + } + if !tt.wantWorkerWrite && stored.GetVersion() != seeded.GetVersion() { + t.Errorf("worker version moved %d -> %d, want no write", seeded.GetVersion(), stored.GetVersion()) + } + }) + } +} diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index 853739aa9..095676618 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend.go @@ -119,16 +119,15 @@ func (s *CallAteletSuspendStep) CheckPrerequisite(ctx context.Context, input *Su if state.Actor.GetStatus() != ateapipb.Actor_STATUS_SUSPENDING { return status.Errorf(codes.FailedPrecondition, "CallAteletSuspendStep prerequisite not met for Actor: %s (got: %v, want %s)", input.ActorName, state.Actor.GetStatus(), ateapipb.Actor_STATUS_SUSPENDING) } - return nil -} -func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) 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 SUSPENDING state but has no active worker") } - + return nil +} +func (s *CallAteletSuspendStep) Execute(ctx context.Context, input *SuspendInput, state *SuspendState) error { ateletConn, err := s.dialer.DialForWorker(state.Actor.GetAteomPodNamespace(), state.Actor.GetAteomPodName()) if err != nil { if errors.Is(err, ErrWorkerPodNotFound) { diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index 8a4d08f4b..a3995793b 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -28,11 +28,11 @@ import ( "google.golang.org/grpc/status" ) -// TestSuspendActorWorkflow exercises the suspend 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 TestSuspendActorWorkflow(t *testing.T) { +// TestSuspendActorWorkflow_RejectedAndIdempotentPaths covers the two +// short-circuit paths of the suspend workflow: rejection by +// MarkSuspendingStep's CheckPrerequisite and the IsComplete idempotent +// fast-forward. +func TestSuspendActorWorkflow_RejectedAndIdempotentPaths(t *testing.T) { tests := []struct { name string seedStatus ateapipb.Actor_Status @@ -141,13 +141,40 @@ func TestSuspendSteps_CheckPrerequisite(t *testing.T) { t.Run(tc.name, func(t *testing.T) { ctx := context.Background() for _, st := range allActorStatuses { - err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorName: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st}}) + // Worker pod fields are populated so CallAteletSuspendStep's + // missing-worker crash branch is not taken; this test only + // verifies status gating. + err := tc.step.CheckPrerequisite(ctx, &SuspendInput{ActorName: "id1"}, &SuspendState{Actor: &ateapipb.Actor{Status: st, AteomPodNamespace: "ns", AteomPodName: "worker-1"}}) assertPrerequisiteResult(t, st, err, tc.allowed == nil || tc.allowed[st]) } }) } } +// TestSuspendActor_CrashesWhenSuspendingActorMissingWorkerPod verifies that a +// SUSPENDING actor with no worker pod recorded is moved to CRASHED by +// CallAteletSuspendStep's prerequisite check and the suspend fails. +func TestSuspendActor_CrashesWhenSuspendingActorMissingWorkerPod(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_SUSPENDING) + + if _, err := w.SuspendActor(ctx, "team-a", "id1"); err == nil { + t.Fatal("SuspendActor succeeded, want error for SUSPENDING actor with no worker pod") + } + + 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) + } +} + // newTestPersistence returns a store backed by a throwaway miniredis. func newTestPersistence(t *testing.T) store.Interface { t.Helper()