Fix toctou race in UpdateActor

UpdateActor now takes an ActorRef, an ActorPrecondition, and a mutate
callback. The store reads the stored actor inside the WATCH
transaction, checks the precondition against that value, and hands it
to mutate, which edits it in place.

The storage layer retries up to 5 times when a concurrent write
invalidates it, re-running mutate against the newer state

I still need to migrate the other callers of store.UpdateActor to avoid calling GetActor outside of the closure function.

Author: laoj2

cmd/ateapi/internal/controlapi/crash.go

@@ -75,14 +75,16 @@ func crashActor(ctx context.Context, st store.Interface, actorRef resources.Acto
 	// the counter itself is emitted only after the transition commits.
 	crashAttrs := ateattr.ActorMetricAttributes(actor, sandboxClass, opName, reason)
 
-	actor.Status = ateapipb.Actor_STATUS_CRASHED
-
-	// InProgressSnapshotName and InProgressLocalSnapshotName are kept for
-	// debugging; failed workflow steps must never promote either of them to an
-	// ActorSnapshot or to LocalSnapshotInfo.
-	actor.WorkerAssignment = nil
-
-	_, err = st.UpdateActor(ctx, actor, actor.GetMetadata().GetVersion())
+	precondition := store.WithActorPrecondition(actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion())
+	_, err = st.UpdateActor(ctx, actorRef, precondition, func(dbActor *ateapipb.Actor) error {
+		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
+
+		// InProgressSnapshotName and InProgressLocalSnapshotName are kept for
+		// debugging; failed workflow steps must never promote either of them to an
+		// ActorSnapshot or to LocalSnapshotInfo.
+		dbActor.WorkerAssignment = nil
+		return nil
+	})
 	if err != nil {
 		errCollected = append(errCollected, fmt.Errorf("while marking actor crashed: %w", err))
 		return errors.Join(errCollected...)

cmd/ateapi/internal/controlapi/functional_test.go

@@ -26,6 +26,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache"
 	"github.com/agent-substrate/substrate/internal/ateinterceptors"
@@ -2447,7 +2448,7 @@ func TestUpdateActor_Preconditions(t *testing.T) {
 	// The uid from the deleted lifecycle must be rejected, even though the
 	// atespace/name it was observed under still resolves.
 	_, err := update(&ateapipb.ResourceMetadata{Uid: staleUID}, "other-lifecycle")
-	assertGrpcError(t, err, codes.Aborted, fmt.Sprintf("Actor %s/%s has uid %s, not %s", testAtespace, testActorID, uid, staleUID))
+	assertGrpcError(t, err, codes.Aborted, fmt.Sprintf("actor %s/%s not found with uid %s", testAtespace, testActorID, staleUID))
 
 	// An unguarded update is last-writer-wins, and moves the resource past the
 	// version observed above.
@@ -2494,7 +2495,7 @@ func TestUpdateActor_NotFound(t *testing.T) {
 		Actor:      &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "does-not-exist"}},
 		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"worker_selector"}},
 	})
-	assertGrpcError(t, err, codes.NotFound, "Actor test-atespace/does-not-exist not found")
+	assertGrpcError(t, err, codes.NotFound, "actor test-atespace/does-not-exist not found")
 }
 
 func TestUpdateActorSnapshotTag_Success(t *testing.T) {
@@ -3287,7 +3288,7 @@ func TestDeleteActor_Crashed(t *testing.T) {
 
 	createTemplate(t, tc, ns)
 
-	_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{Actor: &ateapipb.Actor{
+	actor, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{Actor: &ateapipb.Actor{
 		Metadata:               &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "id1"},
 		ActorTemplateNamespace: ns,
 		ActorTemplateName:      "tmpl1",
@@ -3296,12 +3297,12 @@ func TestDeleteActor_Crashed(t *testing.T) {
 		t.Fatalf("CreateActor failed: %v", err)
 	}
 
-	actor, err := tc.persistence.GetActor(context.Background(), resources.ActorRef{Atespace: testAtespace, Name: "id1"})
-	if err != nil {
-		t.Fatalf("GetActor failed: %v", err)
-	}
-	actor.Status = ateapipb.Actor_STATUS_CRASHED
-	if _, err := tc.persistence.UpdateActor(context.Background(), actor, actor.GetMetadata().GetVersion()); err != nil {
+	actorRef := resources.ActorRef{Atespace: testAtespace, Name: "id1"}
+	precondition := store.WithActorPrecondition(actor.GetMetadata().Uid, actor.GetMetadata().GetVersion())
+	if _, err := tc.persistence.UpdateActor(context.Background(), actorRef, precondition, func(dbActor *ateapipb.Actor) error {
+		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
+		return nil
+	}); err != nil {
 		t.Fatalf("UpdateActor failed: %v", err)
 	}
 

cmd/ateapi/internal/controlapi/syncer.go

@@ -383,14 +383,16 @@ func (s *WorkerPoolSyncer) releaseActorOnDeadWorker(ctx context.Context, namespa
 	// Snapshot crash attributes before pod and pool pointers are cleared on actor.
 	crashAttrs := ateattr.ActorMetricAttributes(actor, worker.GetSandboxClass(), opName, ateattr.ReasonWorkerPodGone)
 
-	actor.Status = ateapipb.Actor_STATUS_CRASHED
-	actor.WorkerAssignment = nil
-	// Both in-progress checkpoints die with the worker: the durable one was
-	// never uploaded, the local one lived on the node that went away.
-	actor.InProgressSnapshotName = ""
-	actor.InProgressLocalSnapshotName = ""
-
-	_, err = s.persistence.UpdateActor(ctx, actor, actor.GetMetadata().GetVersion())
+	precondition := store.WithActorPrecondition(actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion())
+	_, err = s.persistence.UpdateActor(ctx, actorRef, precondition, func(dbActor *ateapipb.Actor) error {
+		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
+		dbActor.WorkerAssignment = nil
+		// Both in-progress checkpoints die with the worker: the durable one was
+		// never uploaded, the local one lived on the node that went away.
+		dbActor.InProgressSnapshotName = ""
+		dbActor.InProgressLocalSnapshotName = ""
+		return nil
+	})
 
 	if err == nil && !wasAlreadyCrashed {
 		recordActorCrash(ctx, crashAttrs)

cmd/ateapi/internal/controlapi/update_actor.go

@@ -41,31 +41,21 @@ func (s *Service) UpdateActor(ctx context.Context, req *ateapipb.UpdateActorRequ
 	actorRef := resources.ActorRefFromActor(in)
 	setSpanActorRefAttributes(ctx, actorRef)
 
-	actor, err := s.persistence.GetActor(ctx, actorRef)
-	if err != nil {
-		if errors.Is(err, store.ErrNotFound) {
-			return nil, status.Errorf(codes.NotFound, "Actor %s not found", actorRef)
-		}
-		return nil, fmt.Errorf("while getting actor: %w", err)
-	}
-
-	// UID and version preconditions
-	if uid := in.GetMetadata().GetUid(); uid != "" && uid != actor.GetMetadata().GetUid() {
-		return nil, status.Errorf(codes.Aborted, "Actor %s has uid %s, not %s", actorRef, actor.GetMetadata().GetUid(), uid)
-	}
-
-	expectedVersion := actor.GetMetadata().GetVersion()
-	if version := in.GetMetadata().GetVersion(); version != 0 {
-		expectedVersion = version
-	}
-
-	applyUpdateMask(actor, in, req.GetUpdateMask(), actorMutableFields)
-
-	updated, err := s.persistence.UpdateActor(ctx, actor, expectedVersion)
+	precondition := store.WithActorPrecondition(in.GetMetadata().GetUid(), in.GetMetadata().GetVersion())
+	updated, err := s.persistence.UpdateActor(ctx, actorRef, precondition, func(dbActor *ateapipb.Actor) error {
+		applyUpdateMask(dbActor, in, req.GetUpdateMask(), actorMutableFields)
+		return nil
+	})
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
+		if errors.Is(err, store.ErrUIDConflict) {
+			return nil, status.Errorf(codes.Aborted, "actor %s/%s not found with uid %s", in.GetMetadata().GetAtespace(), in.GetMetadata().GetName(), in.GetMetadata().GetUid())
+		}
+		if errors.Is(err, store.ErrNotFound) {
+			return nil, status.Errorf(codes.NotFound, "actor %s not found", actorRef)
+		}
 		return nil, fmt.Errorf("while updating actor: %w", err)
 	}
 

cmd/ateapi/internal/controlapi/update_actor_test.go

@@ -24,8 +24,10 @@ import (
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"k8s.io/apimachinery/pkg/util/validation/field"
 
+	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/storetest"
 	"github.com/agent-substrate/substrate/internal/ateattr"
+	"github.com/agent-substrate/substrate/internal/resources"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 )
 
@@ -218,6 +220,152 @@ func TestUpdateActor_FailedLookupStampsRefIdentityOnly(t *testing.T) {
 	}
 }
 
+// TestUpdateActor_DeleteRecreateRace checks that an update is not applied
+// if an actor was deleted and recreated during the update operation.
+func TestUpdateActor_DeleteRecreateRace(t *testing.T) {
+	ctx := context.Background()
+	persistence, cleanup := storetest.SetupTestStore(t)
+	t.Cleanup(cleanup)
+
+	actorRef := resources.ActorRef{Atespace: testAtespace, Name: testActorID}
+
+	// Actor A: what the client reads, and what its uid precondition names.
+	// Freshly created, so it sits at version 1.
+	original, err := persistence.CreateActor(ctx, &ateapipb.Actor{
+		Metadata:               &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: testActorID},
+		ActorTemplateNamespace: "ns1",
+		ActorTemplateName:      "tmpl1",
+		Status:                 ateapipb.Actor_STATUS_RUNNING,
+		WorkerAssignment:       &ateapipb.WorkerAssignment{WorkerPod: "pod-a"},
+	})
+	if err != nil {
+		t.Fatalf("seed CreateActor: %v", err)
+	}
+
+	// A concurrent client deletes A and recreates the same atespace/name as a
+	// brand new actor B, in the window the handler used to leave open between
+	// its own read and the store's WATCH.
+	var recreated *ateapipb.Actor
+	racing := &conflictInjectingStore{
+		Interface: persistence,
+		inject: func() {
+			if _, err := persistence.UpdateActor(ctx, actorRef, store.WithActorPrecondition(store.AnyUID, store.AnyVersion), func(dbActor *ateapipb.Actor) error {
+				dbActor.Status = ateapipb.Actor_STATUS_DELETING
+				return nil
+			}); err != nil {
+				t.Fatalf("racing writer: mark deleting: %v", err)
+			}
+			if _, err := persistence.DeleteActor(ctx, actorRef); err != nil {
+				t.Fatalf("racing writer: DeleteActor: %v", err)
+			}
+			recreated, err = persistence.CreateActor(ctx, &ateapipb.Actor{
+				Metadata:               &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: testActorID},
+				ActorTemplateNamespace: "ns1",
+				ActorTemplateName:      "tmpl1",
+				Status:                 ateapipb.Actor_STATUS_SUSPENDED,
+			})
+			if err != nil {
+				t.Fatalf("racing writer: recreate CreateActor: %v", err)
+			}
+		},
+	}
+	svc := &Service{persistence: racing}
+
+	// The client asserts "only update the actor with uid A".
+	_, err = svc.UpdateActor(ctx, &ateapipb.UpdateActorRequest{
+		Actor: &ateapipb.Actor{
+			Metadata: &ateapipb.ResourceMetadata{
+				Atespace: testAtespace,
+				Name:     testActorID,
+				Uid:      original.GetMetadata().GetUid(),
+			},
+			WorkerSelector: &ateapipb.Selector{MatchLabels: map[string]string{"tier": "paid"}},
+		},
+		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"worker_selector"}},
+	})
+	if code := status.Code(err); code != codes.Aborted {
+		t.Errorf("UpdateActor error = %v (code %v), want code Aborted: the actor holding uid %s was deleted mid-update",
+			err, code, original.GetMetadata().GetUid())
+	}
+
+	stored, err := persistence.GetActor(ctx, actorRef)
+	if err != nil {
+		t.Fatalf("GetActor: %v", err)
+	}
+	if got, want := stored.GetMetadata().GetUid(), recreated.GetMetadata().GetUid(); got != want {
+		t.Fatalf("stored uid = %s, want recreated actor's uid %s", got, want)
+	}
+	// The stored record must still be actor B as its creator left it. Any of A's
+	// state showing up here is the clobber.
+	if got := stored.GetStatus(); got != ateapipb.Actor_STATUS_SUSPENDED {
+		t.Errorf("stored status = %v, want %v: recreated actor was overwritten with the deleted actor's state",
+			got, ateapipb.Actor_STATUS_SUSPENDED)
+	}
+	if got := stored.GetWorkerAssignment(); got != nil {
+		t.Errorf("stored worker_assignment = %v, want nil: recreated actor inherited the deleted actor's worker", got)
+	}
+	if got := stored.GetWorkerSelector(); got != nil {
+		t.Errorf("stored worker_selector = %v, want nil: update meant for the deleted actor was applied", got)
+	}
+}
+
+// TestUpdateActor_ConcurrentDisjointUpdates checks that concurrent write
+// to a disjoint field is resolved by the store and both fields survive the update.
+func TestUpdateActor_ConcurrentDisjointUpdates(t *testing.T) {
+	ctx := context.Background()
+	persistence, cleanup := storetest.SetupTestStore(t)
+	t.Cleanup(cleanup)
+
+	actorRef := resources.ActorRef{Atespace: testAtespace, Name: testActorID}
+
+	if _, err := persistence.CreateActor(ctx, &ateapipb.Actor{
+		Metadata:               &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: testActorID},
+		ActorTemplateNamespace: "ns1",
+		ActorTemplateName:      "tmpl1",
+		Status:                 ateapipb.Actor_STATUS_RUNNING,
+	}); err != nil {
+		t.Fatalf("seed CreateActor: %v", err)
+	}
+
+	// A suspend workflow bumps status (a field that a later update operation will not touch)
+	// inside the handler's read-modify-write window.
+	racing := &conflictInjectingStore{
+		Interface: persistence,
+		inject: func() {
+			if _, err := persistence.UpdateActor(ctx, actorRef, store.WithActorPrecondition(store.AnyUID, store.AnyVersion), func(dbActor *ateapipb.Actor) error {
+				dbActor.Status = ateapipb.Actor_STATUS_SUSPENDING
+				return nil
+			}); err != nil {
+				t.Fatalf("racing writer: mark suspending: %v", err)
+			}
+		},
+	}
+	svc := &Service{persistence: racing}
+
+	// Update operation is changing the worker_selector field, not the actor's status (like the concurrent op)
+	if _, err := svc.UpdateActor(ctx, &ateapipb.UpdateActorRequest{
+		Actor: &ateapipb.Actor{
+			Metadata:       &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: testActorID},
+			WorkerSelector: &ateapipb.Selector{MatchLabels: map[string]string{"tier": "paid"}},
+		},
+		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"worker_selector"}},
+	}); err != nil {
+		t.Fatalf("UpdateActor error = %v, want success: no version precondition was set, so the conflict is the server's to resolve", err)
+	}
+
+	stored, err := persistence.GetActor(ctx, actorRef)
+	if err != nil {
+		t.Fatalf("GetActor: %v", err)
+	}
+	// Both worker selector and status updates survive
+	if got := stored.GetWorkerSelector().GetMatchLabels()["tier"]; got != "paid" {
+		t.Errorf("stored worker_selector[tier] = %q, want %q", got, "paid")
+	}
+	if got := stored.GetStatus(); got != ateapipb.Actor_STATUS_SUSPENDING {
+		t.Errorf("stored status = %v, want %v: the concurrent writer's field must survive", got, ateapipb.Actor_STATUS_SUSPENDING)
+	}
+}
+
 // updateActorReq builds a minimal valid UpdateActorRequest, then applies the
 // given mutations.
 func updateActorReq(mutate ...func(*ateapipb.UpdateActorRequest)) *ateapipb.UpdateActorRequest {

cmd/ateapi/internal/controlapi/workflow_delete.go

@@ -81,11 +81,14 @@ func (s *MarkDeletingStep) CheckPrerequisite(ctx context.Context, input *DeleteI
 	return nil
 }
 func (s *MarkDeletingStep) Execute(ctx context.Context, input *DeleteInput, state *DeleteState) error {
-	state.Actor.Status = ateapipb.Actor_STATUS_DELETING
-	for _, vol := range state.Actor.GetActorVolumes() {
-		vol.Status = ateapipb.ExternalVolume_STATUS_DELETING
-	}
-	updated, err := s.store.UpdateActor(ctx, state.Actor, state.Actor.GetMetadata().GetVersion())
+	precondition := store.WithActorPrecondition(state.Actor.GetMetadata().GetUid(), state.Actor.GetMetadata().GetVersion())
+	updated, err := s.store.UpdateActor(ctx, input.ActorRef, precondition, func(dbActor *ateapipb.Actor) error {
+		dbActor.Status = ateapipb.Actor_STATUS_DELETING
+		for _, vol := range dbActor.GetActorVolumes() {
+			vol.Status = ateapipb.ExternalVolume_STATUS_DELETING
+		}
+		return nil
+	})
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return status.Error(codes.Aborted, "concurrent update conflict, please retry")