diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 32122ea51..b1814fb2e 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume.go @@ -222,7 +222,23 @@ func (s *AssignWorkerStep) Execute(ctx context.Context, input *ResumeInput, stat updatedActor, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion()) if err != nil { - return err + if !errors.Is(err, store.ErrPersistenceRetry) { + return err + } + // refresh the version of actor to avoid always failure in rest retries. + fresh, gerr := s.store.GetActor(ctx, input.Atespace, input.ActorName) + if gerr != nil { + slog.WarnContext(ctx, "Failed to refresh actor after assignment conflict", slog.Any("err", gerr)) + return err + } + switch fresh.GetStatus() { + case ateapipb.Actor_STATUS_SUSPENDED, ateapipb.Actor_STATUS_PAUSED: + slog.InfoContext(ctx, "Retrying assignment due to actor version conflict", slog.String("actor", input.Atespace+"/"+input.ActorName)) + state.Actor = fresh + return err + default: + return status.Errorf(codes.Aborted, "actor %s is %s and can no longer be resumed", input.ActorName, fresh.GetStatus()) + } } state.Actor = updatedActor state.Worker = assignedWorker diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 05a37bbb3..94c12c3ba 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_resume_test.go @@ -16,10 +16,13 @@ package controlapi import ( "context" + "errors" + "sync" "testing" "time" "github.com/agent-substrate/substrate/cmd/ateapi/internal/scheduling" + "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest" "github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache" atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" @@ -263,6 +266,137 @@ func TestAssignWorkerStep_RetryAfterConflictPicksFreshWorker(t *testing.T) { } } +// conflictInjectingStore wraps a store and runs inject exactly once, +// immediately before the first UpdateActor, simulating a concurrent writer +// racing the step's read-modify-write window. +type conflictInjectingStore struct { + store.Interface + once sync.Once + inject func() +} + +func (c *conflictInjectingStore) UpdateActor(ctx context.Context, actor *ateapipb.Actor, expectedVersion int64) (*ateapipb.Actor, error) { + c.once.Do(c.inject) + return c.Interface.UpdateActor(ctx, actor, expectedVersion) +} + +// seedAssignFixture stores one free gvisor worker and a SUSPENDED actor and +// returns the actor plus a started worker cache. +func seedAssignFixture(t *testing.T, ctx context.Context, persistence store.Interface) (*ateapipb.Actor, *workercache.Cache) { + t.Helper() + if err := persistence.CreateWorker(ctx, &ateapipb.Worker{ + WorkerNamespace: "worker-ns", + WorkerPool: "pool", + WorkerPod: "pod-1", + SandboxClass: "gvisor", + }); err != nil { + t.Fatalf("CreateWorker: %v", 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) + t.Cleanup(cancel) + wc := workercache.New(persistence, time.Minute) + if err := wc.Start(cacheCtx); err != nil { + t.Fatalf("workercache.Start: %v", err) + } + return actor, wc +} + +// TestAssignWorkerStep_ConflictRefreshesActor verifies the actor write's +// conflict handling within a single Execute: a concurrent spec write leaves +// ErrPersistenceRetry with state.Actor refreshed. +func TestAssignWorkerStep_ConflictRefreshesActor(t *testing.T) { + t.Parallel() + tests := []struct { + name string + // mutate is the racing concurrent write applied to the fresh actor. + mutate func(fresh *ateapipb.Actor) + // wantRetry means Execute surfaces ErrPersistenceRetry with + // state.Actor refreshed to the injected write; otherwise Aborted. + wantRetry bool + // wantStoredStatus is the persisted status after Execute. + wantStoredStatus ateapipb.Actor_Status + }{ + { + name: "another writer refreshes state.Actor - can recover", + mutate: func(fresh *ateapipb.Actor) { + fresh.WorkerSelector = &ateapipb.Selector{MatchLabels: map[string]string{"team": "blue"}} + }, + wantRetry: true, + wantStoredStatus: ateapipb.Actor_STATUS_SUSPENDED, + }, + { + name: "another writer crash the Actor", + mutate: func(fresh *ateapipb.Actor) { + fresh.Status = ateapipb.Actor_STATUS_CRASHED + }, + wantRetry: false, + wantStoredStatus: ateapipb.Actor_STATUS_CRASHED, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := context.Background() + persistence := newTestPersistence(t) + actor, wc := seedAssignFixture(t, ctx, persistence) + + var injected *ateapipb.Actor + st := &conflictInjectingStore{Interface: persistence, inject: func() { + fresh, err := persistence.GetActor(ctx, "team-a", "id1") + if err != nil { + t.Errorf("inject GetActor: %v", err) + return + } + tc.mutate(fresh) + injected, err = persistence.UpdateActor(ctx, fresh, fresh.GetMetadata().GetVersion()) + if err != nil { + t.Errorf("inject UpdateActor: %v", err) + } + }} + + step := &AssignWorkerStep{store: st, workerCache: wc, scheduler: scheduling.New(wc)} + state := &ResumeState{ + Actor: actor, + ActorTemplate: &atev1alpha1.ActorTemplate{ + Spec: atev1alpha1.ActorTemplateSpec{SandboxClass: atev1alpha1.SandboxClassGvisor}, + }, + } + err := step.Execute(ctx, &ResumeInput{ActorName: "id1", Atespace: "team-a"}, state) + + if tc.wantRetry { + if !errors.Is(err, store.ErrPersistenceRetry) { + t.Fatalf("Execute: %v, want ErrPersistenceRetry", err) + } + if got := state.Actor.GetMetadata().GetVersion(); got != injected.GetMetadata().GetVersion() { + t.Errorf("state.Actor version = %d, want %d (refreshed for the retry)", got, injected.GetMetadata().GetVersion()) + } + if !proto.Equal(state.Actor.GetWorkerSelector(), injected.GetWorkerSelector()) { + t.Errorf("state.Actor WorkerSelector = %v, want %v (concurrent write must survive)", state.Actor.GetWorkerSelector(), injected.GetWorkerSelector()) + } + } else { + if got := status.Code(err); got != codes.Aborted { + t.Fatalf("status.Code(err) = %v, want %v (err: %v)", got, codes.Aborted, err) + } + } + + stored, err := persistence.GetActor(ctx, "team-a", "id1") + if err != nil { + t.Fatalf("GetActor: %v", err) + } + if stored.GetStatus() != tc.wantStoredStatus { + t.Errorf("stored status = %v, want %v", stored.GetStatus(), tc.wantStoredStatus) + } + }) + } +} + // TestResumeActorWorkflow_RejectedAndIdempotentPaths covers the two // short-circuit paths of the resume workflow: rejection by AssignWorkerStep's // CheckPrerequisite and the IsComplete idempotent fast-forward.