Fix TOCTOU race in UpdateActorSnapshotTag (#854)

This is a very similar fix to #829

We're also removing the precondition as the storage layer function
arguments and passing it as a wrapper around the closure functions.

https://github.com/agent-substrate/substrate/issues/763

- [x] Tests pass
- [x] Appropriate changes to documentation are included in the PR

cmd/ateapi/internal/controlapi/actor_snapshot.go

@@ -127,37 +127,24 @@ func (s *Service) UpdateActorSnapshotTag(ctx context.Context, req *ateapipb.Upda
 	}
 	in := req.GetTag()
 	atespace, name := in.GetMetadata().GetAtespace(), in.GetMetadata().GetName()
-	_, current, err := s.persistence.GetActorSnapshotByTag(ctx, atespace, name)
-	if errors.Is(err, store.ErrNotFound) {
-		return nil, status.Errorf(codes.NotFound, "ActorSnapshot tag %s/%s not found", atespace, name)
-	}
-	if err != nil {
-		return nil, fmt.Errorf("while getting actor snapshot tag: %w", err)
-	}
 
-	// UID and version preconditions.
-	if uid := in.GetMetadata().GetUid(); uid != "" && uid != current.GetMetadata().GetUid() {
-		return nil, status.Errorf(codes.Aborted, "ActorSnapshot tag %s/%s has uid %s, not %s", atespace, name, current.GetMetadata().GetUid(), uid)
-	}
-
-	expectedVersion := current.GetMetadata().GetVersion()
-	if version := in.GetMetadata().GetVersion(); version != 0 {
-		expectedVersion = version
-	}
-
-	fieldmask.Apply(current, in, req.GetUpdateMask())
-
-	updatedTag, err := s.persistence.UpdateActorSnapshotTag(ctx, atespace, name, current.GetScope(), expectedVersion)
-	if errors.Is(err, store.ErrNotFound) {
-		return nil, status.Errorf(codes.NotFound, "ActorSnapshot tag %s/%s not found", atespace, name)
-	}
-	if errors.Is(err, store.ErrVersionConflict) {
-		return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
-	}
+	storedTag, err := s.persistence.UpdateActorSnapshotTag(ctx, atespace, name, store.WithPrecondition(in, func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		fieldmask.Apply(toUpdate, in, req.GetUpdateMask())
+		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, "ActorSnapshot tag %s/%s not found with uid %s", atespace, name, in.GetMetadata().GetUid())
+		}
+		if errors.Is(err, store.ErrNotFound) {
+			return nil, status.Errorf(codes.NotFound, "ActorSnapshot tag %s/%s not found", atespace, name)
+		}
 		return nil, fmt.Errorf("while updating actor snapshot tag: %w", err)
 	}
-	return updatedTag, nil
+	return storedTag, nil
 }
 
 func validateUpdateActorSnapshotTagRequest(req *ateapipb.UpdateActorSnapshotTagRequest) field.ErrorList {

cmd/ateapi/internal/controlapi/actor_snapshot_test.go

@@ -19,6 +19,8 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"k8s.io/apimachinery/pkg/util/validation/field"
@@ -300,3 +302,145 @@ func serviceWithActorSnapshotTag(t *testing.T, tag *ateapipb.ActorSnapshotTag) (
 	}
 	return &Service{persistence: persistence}, created
 }
+
+// TestUpdateActorSnapshotTag_DeleteRecreateRace checks that an update is not
+// applied if a tag was deleted and re-created during the update operation.
+func TestUpdateActorSnapshotTag_DeleteRecreateRace(t *testing.T) {
+	ctx := context.Background()
+	persistence, cleanup := storetest.SetupTestStore(t)
+	t.Cleanup(cleanup)
+
+	for _, name := range []string{"snapshot-1", "snapshot-2"} {
+		if _, err := persistence.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{
+			Metadata:    &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: name},
+			SnapshotUri: "gs://bucket/root/snapshots/" + testAtespace + "/" + name,
+		}); err != nil {
+			t.Fatalf("Failed to CreateActorSnapshot(%s): %v", name, err)
+		}
+	}
+
+	const tagName = "before-upgrade"
+	// Tag A: what the client reads, and what its uid precondition names.
+	// Freshly created, so it sits at version 1.
+	originalTag, err := persistence.TagActorSnapshot(ctx, testAtespace, "snapshot-1", &ateapipb.ActorSnapshotTag{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: tagName},
+		Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE,
+	})
+	if err != nil {
+		t.Fatalf("Failed to TagActorSnapshot(snapshot-1): %v", err)
+	}
+
+	// A concurrent client deletes A and re-tags the same atespace/name as a
+	// brand new tag B, pointed at another snapshot.
+	var recreatedTag *ateapipb.ActorSnapshotTag
+	racing := &conflictInjectingStore{
+		Interface: persistence,
+		inject: func() {
+			if _, err := persistence.DeleteActorSnapshotTag(ctx, testAtespace, tagName); err != nil {
+				t.Fatalf("Racing writer: DeleteActorSnapshotTag: %v", err)
+			}
+			recreatedTag, err = persistence.TagActorSnapshot(ctx, testAtespace, "snapshot-2", &ateapipb.ActorSnapshotTag{
+				Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: tagName},
+				Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE,
+			})
+			if err != nil {
+				t.Fatalf("Racing writer: re-tag TagActorSnapshot: %v", err)
+			}
+		},
+	}
+	svc := &Service{persistence: racing}
+
+	// The client asserts "only update the tag with uid A". Its version guard is
+	// satisfied by B as well, because re-tagging resets the version to 1: the
+	// uid is the only thing that can tell the two lifecycles apart.
+	_, err = svc.UpdateActorSnapshotTag(ctx, &ateapipb.UpdateActorSnapshotTagRequest{
+		Tag: &ateapipb.ActorSnapshotTag{
+			Metadata: &ateapipb.ResourceMetadata{
+				Atespace: testAtespace,
+				Name:     tagName,
+				Uid:      originalTag.GetMetadata().GetUid(),
+				Version:  originalTag.GetMetadata().GetVersion(),
+			},
+			Scope: ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED,
+		},
+		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"scope"}},
+	})
+	if code := status.Code(err); code != codes.Aborted {
+		t.Errorf("UpdateActorSnapshotTag error = %v (code %v), want code Aborted: the tag holding uid %s was deleted mid-update",
+			err, code, originalTag.GetMetadata().GetUid())
+	}
+
+	_, storedTag, err := persistence.GetActorSnapshotByTag(ctx, testAtespace, tagName)
+	if err != nil {
+		t.Fatalf("GetActorSnapshotByTag: %v", err)
+	}
+	// The stored record must still be tag B as its creator left it. Any of A's
+	// state showing up here is the clobber.
+	if diff := cmp.Diff(recreatedTag, storedTag, protocmp.Transform()); diff != "" {
+		t.Errorf("Update meant for the deleted tag was applied to the recreated one (-recreated +stored):\n%s", diff)
+	}
+}
+
+// TestUpdateActorSnapshotTag_ConcurrentUnguardedUpdate checks that an update
+// pinning nothing is the server's conflict to resolve: a write landing in the
+// handler's read-modify-write window is absorbed, not reported as Aborted.
+func TestUpdateActorSnapshotTag_ConcurrentUnguardedUpdate(t *testing.T) {
+	ctx := context.Background()
+	persistence, cleanup := storetest.SetupTestStore(t)
+	t.Cleanup(cleanup)
+
+	if _, err := persistence.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{
+		Metadata:    &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "snapshot-1"},
+		SnapshotUri: "gs://bucket/root/snapshots/" + testAtespace + "/snapshot-1",
+	}); err != nil {
+		t.Fatalf("Failed to CreateActorSnapshot: %v", err)
+	}
+
+	const tagName = "before-upgrade"
+	originalTag, err := persistence.TagActorSnapshot(ctx, testAtespace, "snapshot-1", &ateapipb.ActorSnapshotTag{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: tagName},
+		Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE,
+	})
+	if err != nil {
+		t.Fatalf("Failed to TagActorSnapshot(snapshot-1): %v", err)
+	}
+
+	// A concurrent client moves the tag past the version the caller could have
+	// observed, in the window the handler used to leave open between its own
+	// read and the store's WATCH.
+	racing := &conflictInjectingStore{
+		Interface: persistence,
+		inject: func() {
+			if _, err := persistence.UpdateActorSnapshotTag(ctx, testAtespace, tagName, func(toUpdate *ateapipb.ActorSnapshotTag) error {
+				toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE
+				return nil
+			}); err != nil {
+				t.Fatalf("Racing writer: UpdateActorSnapshotTag: %v", err)
+			}
+		},
+	}
+	svc := &Service{persistence: racing}
+
+	if _, err := svc.UpdateActorSnapshotTag(ctx, &ateapipb.UpdateActorSnapshotTagRequest{
+		Tag: &ateapipb.ActorSnapshotTag{
+			Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: tagName},
+			Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED,
+		},
+		UpdateMask: &fieldmaskpb.FieldMask{Paths: []string{"scope"}},
+	}); err != nil {
+		t.Fatalf("UpdateActorSnapshotTag error = %v, want success: no precondition was set, so the conflict is the server's to resolve", err)
+	}
+
+	_, storedTag, err := persistence.GetActorSnapshotByTag(ctx, testAtespace, tagName)
+	if err != nil {
+		t.Fatalf("Failed to GetActorSnapshotByTag(%s/%s): %v", testAtespace, tagName, err)
+	}
+	if got, want := storedTag.GetScope(), ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED; got != want {
+		t.Errorf("Stored scope = %v, want %v", got, want)
+	}
+	// Seed, concurrent writer, then this update: the write landed on top of the
+	// concurrent one rather than on the state the handler read first.
+	if got, want := storedTag.GetMetadata().GetVersion(), originalTag.GetMetadata().GetVersion()+2; got != want {
+		t.Errorf("stored version = %d, want %d", got, want)
+	}
+}

cmd/ateapi/internal/controlapi/crash.go

@@ -75,18 +75,15 @@ 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)
 
-	_, err = st.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
+	_, err = st.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.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
+		toUpdate.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

@@ -2534,7 +2534,7 @@ func TestUpdateActorSnapshotTag_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}, ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED)
-	assertGrpcError(t, err, codes.Aborted, fmt.Sprintf("ActorSnapshot tag %s/%s has uid %s, not %s", testAtespace, tagName, uid, staleUID))
+	assertGrpcError(t, err, codes.Aborted, fmt.Sprintf("ActorSnapshot tag %s/%s not found with uid %s", testAtespace, tagName, staleUID))
 
 	// An unguarded update is last-writer-wins, and moves the tag past the
 	// version observed above.
@@ -3219,8 +3219,8 @@ func TestDeleteActor_Crashed(t *testing.T) {
 	}
 
 	actorRef := resources.ActorRef{Atespace: testAtespace, Name: "id1"}
-	if _, err := tc.persistence.UpdateActor(context.Background(), actorRef, func(dbActor *ateapipb.Actor) error {
-		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
+	if _, err := tc.persistence.UpdateActor(context.Background(), actorRef, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_CRASHED
 		return nil
 	}); err != nil {
 		t.Fatalf("UpdateActor failed: %v", err)

cmd/ateapi/internal/controlapi/syncer.go

@@ -383,18 +383,15 @@ 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)
 
-	_, err = s.persistence.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_CRASHED
-		dbActor.WorkerAssignment = nil
+	_, err = s.persistence.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_CRASHED
+		toUpdate.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 = ""
+		toUpdate.InProgressSnapshotName = ""
+		toUpdate.InProgressLocalSnapshotName = ""
 		return nil
-	})
+	}))
 
 	if err == nil && !wasAlreadyCrashed {
 		recordActorCrash(ctx, crashAttrs)

cmd/ateapi/internal/controlapi/update_actor.go

@@ -43,13 +43,10 @@ func (s *Service) UpdateActor(ctx context.Context, req *ateapipb.UpdateActorRequ
 	actorRef := resources.ActorRefFromActor(in)
 	setSpanActorRefAttributes(ctx, actorRef)
 
-	updated, err := s.persistence.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, in.GetMetadata().GetUid(), in.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		fieldmask.Apply(dbActor, in, req.GetUpdateMask())
+	storedActor, err := s.persistence.UpdateActor(ctx, actorRef, store.WithPrecondition(in, func(toUpdate *ateapipb.Actor) error {
+		fieldmask.Apply(toUpdate, in, req.GetUpdateMask())
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
@@ -63,8 +60,8 @@ func (s *Service) UpdateActor(ctx context.Context, req *ateapipb.UpdateActorRequ
 		return nil, fmt.Errorf("while updating actor: %w", err)
 	}
 
-	setSpanActorAttributes(ctx, updated)
-	return updated, nil
+	setSpanActorAttributes(ctx, storedActor)
+	return storedActor, nil
 }
 
 func validateUpdateActorRequest(req *ateapipb.UpdateActorRequest) field.ErrorList {

cmd/ateapi/internal/controlapi/update_actor_test.go

@@ -300,8 +300,8 @@ func TestUpdateActor_DeleteRecreateRace(t *testing.T) {
 	racing := &conflictInjectingStore{
 		Interface: persistence,
 		inject: func() {
-			if _, err := persistence.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-				dbActor.Status = ateapipb.Actor_STATUS_DELETING
+			if _, err := persistence.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
+				toUpdate.Status = ateapipb.Actor_STATUS_DELETING
 				return nil
 			}); err != nil {
 				t.Fatalf("racing writer: mark deleting: %v", err)
@@ -383,8 +383,8 @@ func TestUpdateActor_ConcurrentDisjointUpdates(t *testing.T) {
 	racing := &conflictInjectingStore{
 		Interface: persistence,
 		inject: func() {
-			if _, err := persistence.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-				dbActor.Status = ateapipb.Actor_STATUS_SUSPENDING
+			if _, err := persistence.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
+				toUpdate.Status = ateapipb.Actor_STATUS_SUSPENDING
 				return nil
 			}); err != nil {
 				t.Fatalf("racing writer: mark suspending: %v", err)

cmd/ateapi/internal/controlapi/workflow_delete.go

@@ -78,23 +78,20 @@ func (w *ActorWorkflow) ensureMarkedDeleting(ctx context.Context, actorRef resou
 		return nil, status.Errorf(codes.FailedPrecondition, "Actor %s is not in a deletable status (status: %v)", actorRef, actor.GetStatus())
 	}
 
-	updated, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_DELETING
-		for _, vol := range dbActor.GetActorVolumes() {
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_DELETING
+		for _, vol := range toUpdate.GetActorVolumes() {
 			vol.Status = ateapipb.ExternalVolume_STATUS_DELETING
 		}
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, fmt.Errorf("while setting actor status to DELETING: %w", err)
 	}
-	return updated, nil
+	return storedActor, nil
 }
 
 // ensureVolumesDeleted removes the actor's external volumes. Volume deletion

cmd/ateapi/internal/controlapi/workflow_pause.go

@@ -121,21 +121,18 @@ func (w *ActorWorkflow) ensureMarkedPausing(ctx context.Context, actorRef resour
 	}
 
 	snapshotName := resources.NewSnapshotName()
-	updated, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_PAUSING
-		dbActor.InProgressLocalSnapshotName = snapshotName
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_PAUSING
+		toUpdate.InProgressLocalSnapshotName = snapshotName
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, err
 	}
-	return updated, nil
+	return storedActor, nil
 }
 
 // ensureAteletPaused checkpoints the workload locally on the worker node
@@ -266,27 +263,24 @@ func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef reso
 		latestActor.Status = newStatus
 		crashAttrs := ateattr.ActorMetricAttributes(latestActor, sandboxClass, ateattr.OperationPause, ateattr.ReasonCorruptedAssignment)
 
-		updatedActor, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-			if err := store.CheckActorPrecondition(dbActor, latestActor.GetMetadata().GetUid(), latestActor.GetMetadata().GetVersion()); err != nil {
-				return err
-			}
-			dbActor.Status = newStatus
+		storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(latestActor, func(toUpdate *ateapipb.Actor) error {
+			toUpdate.Status = newStatus
 			// TODO(dberkov) - what if InProgressLocalSnapshotName is empty? That shouldn't be possible.
-			if dbActor.GetInProgressLocalSnapshotName() != "" {
+			if toUpdate.GetInProgressLocalSnapshotName() != "" {
 				localInfo := &ateapipb.LocalSnapshotInfo{
-					SnapshotName: dbActor.GetInProgressLocalSnapshotName(),
+					SnapshotName: toUpdate.GetInProgressLocalSnapshotName(),
 					ContentScope: contentScope,
 				}
 				if newStatus != ateapipb.Actor_STATUS_CRASHED {
 					localInfo.NodeVmsWithLocalSnapshots = []string{nodeName}
 				}
-				dbActor.LocalSnapshotInfo = localInfo
-				dbActor.InProgressLocalSnapshotName = ""
+				toUpdate.LocalSnapshotInfo = localInfo
+				toUpdate.InProgressLocalSnapshotName = ""
 			}
-			dbActor.WorkerAssignment = nil
+			toUpdate.WorkerAssignment = nil
 			return nil
-		})
-		if err == nil && updatedActor.GetStatus() == ateapipb.Actor_STATUS_CRASHED && !wasAlreadyCrashed {
+		}))
+		if err == nil && storedActor.GetStatus() == ateapipb.Actor_STATUS_CRASHED && !wasAlreadyCrashed {
 			recordActorCrash(ctx, crashAttrs)
 		}
 		if err != nil {
@@ -295,7 +289,7 @@ func (w *ActorWorkflow) ensurePausedFinalized(ctx context.Context, actorRef reso
 			}
 			return nil, err
 		}
-		latestActor = updatedActor
+		latestActor = storedActor
 	}
 
 	return latestActor, nil

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -254,13 +254,10 @@ func (w *ActorWorkflow) ensureVolumesCreated(ctx context.Context, actorRef resou
 	volumes, createErr := createActorVolumes(ctx, w.pluginRegistry, w.storageClassLister, actor.GetMetadata().GetUid(), actorTemplate, actor.GetActorVolumes())
 	// createActorVolumes reports the state it got to even when it fails, so both
 	// paths persist the same field.
-	persistVolumes := func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.ActorVolumes = volumes
+	persistVolumes := store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.ActorVolumes = volumes
 		return nil
-	}
+	})
 	if createErr != nil {
 		// Even if volume creation failed, we still want to persist any updated volume state.
 		if _, updateErr := w.store.UpdateActor(ctx, actorRef, persistVolumes); updateErr != nil {
@@ -268,14 +265,14 @@ func (w *ActorWorkflow) ensureVolumesCreated(ctx context.Context, actorRef resou
 		}
 		return nil, createErr
 	}
-	updated, updateErr := w.store.UpdateActor(ctx, actorRef, persistVolumes)
+	storedActor, updateErr := w.store.UpdateActor(ctx, actorRef, persistVolumes)
 	if updateErr != nil {
 		if errors.Is(updateErr, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, fmt.Errorf("while updating actor after volume creation: %w", updateErr)
 	}
-	return updated, nil
+	return storedActor, nil
 }
 
 // ensureWorkerAssigned leaves the actor RESUMING with a validated, live,
@@ -511,14 +508,11 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour
 	}
 
 	newAssignment := workerAssignmentFrom(assignedWorker)
-	updatedActor, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_RESUMING
-		dbActor.WorkerAssignment = newAssignment
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_RESUMING
+		toUpdate.WorkerAssignment = newAssignment
 		return nil
-	})
+	}))
 	if err != nil {
 		if !errors.Is(err, store.ErrVersionConflict) {
 			return nil, nil, err
@@ -539,7 +533,7 @@ func (w *ActorWorkflow) assignWorkerAttempt(ctx context.Context, actorRef resour
 	}
 	pool = assignedWorker.GetWorkerPool()
 	outcome = ateattr.SchedulerOutcomeAssigned
-	return updatedActor, assignedWorker, nil
+	return storedActor, assignedWorker, nil
 }
 
 func workerAssignmentFrom(w *ateapipb.Worker) *ateapipb.WorkerAssignment {
@@ -735,18 +729,15 @@ func (w *ActorWorkflow) finalizeRunning(ctx context.Context, actorRef resources.
 		return nil, err
 	}
 
-	updatedActor, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, latestActor.GetMetadata().GetUid(), latestActor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(latestActor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, err
 	}
-	return updatedActor, nil
+	return storedActor, nil
 }

cmd/ateapi/internal/controlapi/workflow_resume_test.go

@@ -295,8 +295,8 @@ func TestAssignWorkerAttempt_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.
+// immediately before the first update, simulating a concurrent writer racing
+// the step's read-modify-write window.
 type conflictInjectingStore struct {
 	store.Interface
 	once   sync.Once
@@ -308,6 +308,11 @@ func (c *conflictInjectingStore) UpdateActor(ctx context.Context, actorRef resou
 	return c.Interface.UpdateActor(ctx, actorRef, mutate)
 }
 
+func (c *conflictInjectingStore) UpdateActorSnapshotTag(ctx context.Context, atespace, name string, mutate func(*ateapipb.ActorSnapshotTag) error) (*ateapipb.ActorSnapshotTag, error) {
+	c.once.Do(c.inject)
+	return c.Interface.UpdateActorSnapshotTag(ctx, atespace, name, mutate)
+}
+
 // 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) {
@@ -384,13 +389,15 @@ func TestAssignWorkerAttempt_ConflictRefreshesActor(t *testing.T) {
 					t.Errorf("inject GetActor: %v", err)
 					return
 				}
-				injected, err = persistence.UpdateActor(ctx, resources.ActorRef{Atespace: "team-a", Name: "id1"}, func(dbActor *ateapipb.Actor) error {
-					if err := store.CheckActorPrecondition(dbActor, store.AnyUID, fresh.GetMetadata().GetVersion()); err != nil {
-						return err
-					}
-					tc.mutate(dbActor)
+				// Pins the version alone: the observed actor carries the version
+				// just read and no uid.
+				pinVersion := &ateapipb.Actor{
+					Metadata: &ateapipb.ResourceMetadata{Uid: store.AnyUID, Version: fresh.GetMetadata().GetVersion()},
+				}
+				injected, err = persistence.UpdateActor(ctx, resources.ActorRef{Atespace: "team-a", Name: "id1"}, store.WithPrecondition(pinVersion, func(toUpdate *ateapipb.Actor) error {
+					tc.mutate(toUpdate)
 					return nil
-				})
+				}))
 				if err != nil {
 					t.Errorf("inject UpdateActor: %v", err)
 				}

cmd/ateapi/internal/controlapi/workflow_suspend.go

@@ -137,22 +137,19 @@ func (w *ActorWorkflow) ensureMarkedSuspending(ctx context.Context, actorRef res
 	if _, err := inProgressSnapshotURI(actorTemplate, actorRef.Atespace, name); err != nil {
 		return nil, err
 	}
-	updated, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, actor.GetMetadata().GetUid(), actor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_SUSPENDING
-		dbActor.InProgressSnapshotSourceActorVersion = dbActor.GetMetadata().GetVersion()
-		dbActor.InProgressSnapshotName = name
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(actor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_SUSPENDING
+		toUpdate.InProgressSnapshotSourceActorVersion = toUpdate.GetMetadata().GetVersion()
+		toUpdate.InProgressSnapshotName = name
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, err
 	}
-	return updated, nil
+	return storedActor, nil
 }
 
 // commitSnapshotScope returns the scope a commit (suspend) snapshot is taken
@@ -405,25 +402,22 @@ func (w *ActorWorkflow) ensureSuspendedFinalized(ctx context.Context, actorRef r
 			return nil, err
 		}
 	}
-	updatedActor, err := w.store.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, latestActor.GetMetadata().GetUid(), latestActor.GetMetadata().GetVersion()); err != nil {
-			return err
-		}
-		dbActor.Status = ateapipb.Actor_STATUS_SUSPENDED
+	storedActor, err := w.store.UpdateActor(ctx, actorRef, store.WithPrecondition(latestActor, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_SUSPENDED
 		if snapshotName != "" {
-			dbActor.LatestSnapshot = &ateapipb.ObjectRef{Atespace: actorRef.Atespace, Name: snapshotName}
-			dbActor.InProgressSnapshotName = ""
-			dbActor.InProgressSnapshotSourceActorVersion = 0
+			toUpdate.LatestSnapshot = &ateapipb.ObjectRef{Atespace: actorRef.Atespace, Name: snapshotName}
+			toUpdate.InProgressSnapshotName = ""
+			toUpdate.InProgressSnapshotSourceActorVersion = 0
 		}
-		dbActor.WorkerAssignment = nil
-		dbActor.LocalSnapshotInfo = nil
+		toUpdate.WorkerAssignment = nil
+		toUpdate.LocalSnapshotInfo = nil
 		return nil
-	})
+	}))
 	if err != nil {
 		if errors.Is(err, store.ErrVersionConflict) {
 			return nil, status.Error(codes.Aborted, "concurrent update conflict, please retry")
 		}
 		return nil, err
 	}
-	return updatedActor, nil
+	return storedActor, nil
 }

cmd/ateapi/internal/store/ateredis/ateredis.go

@@ -529,50 +529,96 @@ func (s *Persistence) TagActorSnapshot(ctx context.Context, atespace, name strin
 	return dbTag, nil
 }
 
-func (s *Persistence) UpdateActorSnapshotTag(ctx context.Context, atespace, name string, scope ateapipb.ActorSnapshotTagScope, expectedVersion int64) (*ateapipb.ActorSnapshotTag, error) {
+func validateUpdateActorSnapshotTagMutation(storedTag, mutatedTag *ateapipb.ActorSnapshotTag) error {
+	if stored, mutated := storedTag.GetMetadata().GetAtespace(), mutatedTag.GetMetadata().GetAtespace(); stored != mutated {
+		return fmt.Errorf("metadata.atespace is immutable: mutation changed it from %q to %q", stored, mutated)
+	}
+	if stored, mutated := storedTag.GetMetadata().GetName(), mutatedTag.GetMetadata().GetName(); stored != mutated {
+		return fmt.Errorf("metadata.name is immutable: mutation changed it from %q to %q", stored, mutated)
+	}
+	if stored, mutated := storedTag.GetSnapshot().GetAtespace(), mutatedTag.GetSnapshot().GetAtespace(); stored != mutated {
+		return fmt.Errorf("snapshot.atespace is immutable: mutation changed it from %q to %q", stored, mutated)
+	}
+	if stored, mutated := storedTag.GetSnapshot().GetName(), mutatedTag.GetSnapshot().GetName(); stored != mutated {
+		return fmt.Errorf("snapshot.name is immutable: mutation changed it from %q to %q", stored, mutated)
+	}
+	return nil
+}
+
+// updateActorSnapshotTagMaxAttempts bounds how many times UpdateActorSnapshotTag
+// re-runs its read-modify-write after a concurrent writer invalidates the
+// transaction.
+const updateActorSnapshotTagMaxAttempts = 5
+
+func (s *Persistence) UpdateActorSnapshotTag(ctx context.Context, atespace, name string, mutate func(*ateapipb.ActorSnapshotTag) error) (*ateapipb.ActorSnapshotTag, error) {
 	tagKey := actorSnapshotTagDBKey(atespace, name)
-	var updated *ateapipb.ActorSnapshotTag
-	err := s.rdb.Watch(ctx, func(tx *redis.Tx) error {
-		b, err := tx.Get(ctx, tagKey).Bytes()
-		if err != nil {
-			if errors.Is(err, redis.Nil) {
-				return store.ErrNotFound
+	for range updateActorSnapshotTagMaxAttempts {
+		var dbTag *ateapipb.ActorSnapshotTag
+		var abortErr error
+
+		err := s.rdb.Watch(ctx, func(tx *redis.Tx) error {
+			currentVal, err := tx.Get(ctx, tagKey).Bytes()
+			if err != nil {
+				if errors.Is(err, redis.Nil) {
+					return store.ErrNotFound
+				}
+				return fmt.Errorf("while getting actor snapshot tag %s/%s: %w", atespace, name, err)
 			}
-			return err
-		}
-		tag := &ateapipb.ActorSnapshotTag{}
-		if err := protojson.Unmarshal(b, tag); err != nil {
-			return fmt.Errorf("while unmarshaling actor snapshot tag %s/%s: %w", atespace, name, err)
-		}
-		if tag.GetMetadata().GetVersion() != expectedVersion {
-			return store.ErrVersionConflict
-		}
-		if tag.GetScope() == scope {
-			updated = tag
+
+			currentTag := &ateapipb.ActorSnapshotTag{}
+			if err := protojson.Unmarshal(currentVal, currentTag); err != nil {
+				return fmt.Errorf("while unmarshaling actor snapshot tag %s/%s: %w", atespace, name, err)
+			}
+
+			// Snapshot the stored state before handing the tag to mutate.
+			// mutate is free to edit anything it is given.
+			tagBeforeMutation := proto.Clone(currentTag).(*ateapipb.ActorSnapshotTag)
+			if err := mutate(currentTag); err != nil {
+				abortErr = err
+				return err
+			}
+			if err := validateUpdateActorSnapshotTagMutation(tagBeforeMutation, currentTag); err != nil {
+				abortErr = err
+				return err
+			}
+			// The stored metadata is authoritative; derive the next metadata
+			// from it, discarding whatever mutate made of it.
+			currentTag.Metadata = newUpdateMetadata(tagBeforeMutation.GetMetadata())
+
+			newVal, err := protojson.Marshal(currentTag)
+			if err != nil {
+				return fmt.Errorf("while marshaling actor snapshot tag: %w", err)
+			}
+
+			if _, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
+				pipe.Set(ctx, tagKey, newVal, 0)
+				return nil
+			}); err != nil {
+				return err
+			}
+			dbTag = currentTag
 			return nil
+		}, tagKey)
+
+		switch {
+		case err == nil:
+			return dbTag, nil
+		case abortErr != nil:
+			return nil, abortErr
+		case errors.Is(err, store.ErrNotFound):
+			return nil, store.ErrNotFound
+		case errors.Is(err, redis.TxFailedErr):
+			// A concurrent write landed before we could commit.
+			// Retry.
+			continue
+		default:
+			return nil, fmt.Errorf("while executing update actor snapshot tag transaction: %w", err)
 		}
-		tag.Scope = scope
-		tag.Metadata = newUpdateMetadata(tag.GetMetadata())
-		b, err = protojson.Marshal(tag)
-		if err != nil {
-			return fmt.Errorf("while marshaling actor snapshot tag: %w", err)
-		}
-		if _, err := tx.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
-			pipe.Set(ctx, tagKey, b, 0)
-			return nil
-		}); err != nil {
-			return err
-		}
-		updated = tag
-		return nil
-	}, tagKey)
-	if errors.Is(err, redis.TxFailedErr) {
-		return nil, store.ErrVersionConflict
 	}
-	if err != nil {
-		return nil, fmt.Errorf("while updating actor snapshot tag: %w", err)
-	}
-	return updated, nil
+
+	// Only the TxFailedErr branch continues the loop, so getting here means every
+	// attempt lost the race.
+	return nil, store.ErrVersionConflict
 }
 
 func (s *Persistence) DeleteActorSnapshotTag(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshotTag, error) {

cmd/ateapi/internal/store/ateredis/ateredis_test.go

@@ -162,8 +162,8 @@ func TestUpdateActor_Success(t *testing.T) {
 	}
 
 	actorRef := resources.ActorRefFromActor(actor)
-	updated, err := s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+	updated, err := s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
 	})
 	if err != nil {
@@ -207,9 +207,9 @@ func TestUpdateActor_MutateErrorAreNotRetried(t *testing.T) {
 
 	actorRef := resources.ActorRefFromActor(actor)
 	callsToMutateFn := 0
-	_, err = s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
+	_, err = s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
 		callsToMutateFn++
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return fmt.Errorf("actor %s: %w", actorRef, mutationError)
 	})
 	// The error must arrive intact
@@ -240,13 +240,13 @@ func TestUpdateActor_DiscardsServerOwnedFieldsEdits(t *testing.T) {
 	}
 
 	actorRef := resources.ActorRefFromActor(actor)
-	updated, err := s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
+	updated, err := s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
 		// Metadata is server-owned: a closure must not be able to change it.
-		dbActor.Metadata.Uid = "forged-uid"
-		dbActor.Metadata.Version = 99
-		dbActor.Metadata.CreateTime = nil
-		dbActor.Metadata.UpdateTime = nil
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+		toUpdate.Metadata.Uid = "forged-uid"
+		toUpdate.Metadata.Version = 99
+		toUpdate.Metadata.CreateTime = nil
+		toUpdate.Metadata.UpdateTime = nil
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
 	})
 	if err != nil {
@@ -274,27 +274,27 @@ func TestUpdateActor_DiscardsServerOwnedFieldsEdits(t *testing.T) {
 func TestUpdateActor_RejectsImmutableFieldChange(t *testing.T) {
 	tests := []struct {
 		name      string
-		mutate    func(dbActor *ateapipb.Actor)
+		mutate    func(toUpdate *ateapipb.Actor)
 		wantField string
 	}{
 		{
 			name:      "atespace",
-			mutate:    func(dbActor *ateapipb.Actor) { dbActor.Metadata.Atespace = "other-atespace" },
+			mutate:    func(toUpdate *ateapipb.Actor) { toUpdate.Metadata.Atespace = "other-atespace" },
 			wantField: "metadata.atespace",
 		},
 		{
 			name:      "name",
-			mutate:    func(dbActor *ateapipb.Actor) { dbActor.Metadata.Name = "other-name" },
+			mutate:    func(toUpdate *ateapipb.Actor) { toUpdate.Metadata.Name = "other-name" },
 			wantField: "metadata.name",
 		},
 		{
 			name:      "actor template namespace",
-			mutate:    func(dbActor *ateapipb.Actor) { dbActor.ActorTemplateNamespace = "other-ns" },
+			mutate:    func(toUpdate *ateapipb.Actor) { toUpdate.ActorTemplateNamespace = "other-ns" },
 			wantField: "actor_template_namespace",
 		},
 		{
 			name:      "actor template name",
-			mutate:    func(dbActor *ateapipb.Actor) { dbActor.ActorTemplateName = "other-template" },
+			mutate:    func(toUpdate *ateapipb.Actor) { toUpdate.ActorTemplateName = "other-template" },
 			wantField: "actor_template_name",
 		},
 	}
@@ -308,11 +308,11 @@ func TestUpdateActor_RejectsImmutableFieldChange(t *testing.T) {
 			}
 
 			actorRef := resources.ActorRefFromActor(actor)
-			_, err = s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
+			_, err = s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
 				// Paired with a legitimate edit, so the rejection cannot be
 				// mistaken for a no-op mutation.
-				dbActor.Status = ateapipb.Actor_STATUS_RUNNING
-				tt.mutate(dbActor)
+				toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
+				tt.mutate(toUpdate)
 				return nil
 			})
 			// The message must name the offending field: the closure is buggy,
@@ -383,9 +383,9 @@ func TestUpdateActor_RetriesOnConcurrentWrite(t *testing.T) {
 	}}
 	racing := &Persistence{rdb: interceptor, lockTTL: defaultLockTTL}
 
-	updated, err := racing.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
+	updated, err := racing.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
 		attempts++
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
 	})
 	if err != nil {
@@ -406,7 +406,7 @@ func TestUpdateActor_RetriesOnConcurrentWrite(t *testing.T) {
 
 func TestUpdateActor_NotFound(t *testing.T) {
 	_, s, ctx := setupTest(t)
-	_, err := s.UpdateActor(ctx, resources.ActorRef{Atespace: testAtespace, Name: "non-existent"}, func(dbActor *ateapipb.Actor) error {
+	_, err := s.UpdateActor(ctx, resources.ActorRef{Atespace: testAtespace, Name: "non-existent"}, func(toUpdate *ateapipb.Actor) error {
 		t.Error("mutate must not run for a missing actor")
 		return nil
 	})
@@ -423,8 +423,8 @@ func TestUpdateActor_RejectsStaleUID(t *testing.T) {
 		t.Fatalf("CreateActor failed: %v", err)
 	}
 	actorRef := resources.ActorRefFromActor(original)
-	if _, err := s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		dbActor.Status = ateapipb.Actor_STATUS_DELETING
+	if _, err := s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_DELETING
 		return nil
 	}); err != nil {
 		t.Fatalf("marking actor deleting failed: %v", err)
@@ -440,14 +440,16 @@ func TestUpdateActor_RejectsStaleUID(t *testing.T) {
 		t.Fatalf("recreated actor reused uid %s, want a fresh one", recreated.GetMetadata().GetUid())
 	}
 
-	_, err = s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, original.GetMetadata().GetUid(), store.AnyVersion); err != nil {
-			return err
-		}
+	// Pins the incarnation alone: the observed actor carries the original uid and
+	// no version.
+	pinUID := &ateapipb.Actor{
+		Metadata: &ateapipb.ResourceMetadata{Uid: original.GetMetadata().GetUid(), Version: store.AnyVersion},
+	}
+	_, err = s.UpdateActor(ctx, actorRef, store.WithPrecondition(pinUID, func(toUpdate *ateapipb.Actor) error {
 		t.Error("mutate ran past its precondition once the pinned incarnation was gone")
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
-	})
+	}))
 	if !errors.Is(err, store.ErrUIDConflict) {
 		t.Errorf("UpdateActor error = %v, want one matching store.ErrUIDConflict", err)
 	}
@@ -474,23 +476,20 @@ func TestUpdateActor_RejectsStaleVersion(t *testing.T) {
 		t.Fatalf("CreateActor failed: %v", err)
 	}
 	actorRef := resources.ActorRefFromActor(created)
-	staleVersion := created.GetMetadata().GetVersion()
 
-	if _, err := s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		dbActor.Status = ateapipb.Actor_STATUS_RUNNING
+	if _, err := s.UpdateActor(ctx, actorRef, func(toUpdate *ateapipb.Actor) error {
+		toUpdate.Status = ateapipb.Actor_STATUS_RUNNING
 		return nil
 	}); err != nil {
 		t.Fatalf("UpdateActor failed: %v", err)
 	}
 
-	_, err = s.UpdateActor(ctx, actorRef, func(dbActor *ateapipb.Actor) error {
-		if err := store.CheckActorPrecondition(dbActor, created.GetMetadata().GetUid(), staleVersion); err != nil {
-			return err
-		}
+	// The write above moved the version, so created is now a stale observation.
+	_, err = s.UpdateActor(ctx, actorRef, store.WithPrecondition(created, func(toUpdate *ateapipb.Actor) error {
 		t.Error("mutate ran past its precondition once the pinned version had moved")
-		dbActor.Status = ateapipb.Actor_STATUS_SUSPENDED
+		toUpdate.Status = ateapipb.Actor_STATUS_SUSPENDED
 		return nil
-	})
+	}))
 	if !errors.Is(err, store.ErrVersionConflict) {
 		t.Errorf("UpdateActor error = %v, want one matching store.ErrVersionConflict", err)
 	}
@@ -869,7 +868,10 @@ func TestActorSnapshotLifecycle(t *testing.T) {
 	if _, err := s.TagActorSnapshot(ctx, testAtespace, "snapshot-1", differentScope); !errors.Is(err, store.ErrAlreadyExists) {
 		t.Fatalf("re-tag with different scope error = %v, want ErrAlreadyExists", err)
 	}
-	tagged, err = s.UpdateActorSnapshotTag(ctx, testAtespace, "before-upgrade", ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED, tagged.GetMetadata().GetVersion())
+	tagged, err = s.UpdateActorSnapshotTag(ctx, testAtespace, "before-upgrade", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return nil
+	})
 	if err != nil || tagged.GetScope() != ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED {
 		t.Fatalf("UpdateActorSnapshotTag = (%v, %v), want published", tagged, err)
 	}
@@ -893,34 +895,292 @@ func TestActorSnapshotLifecycle(t *testing.T) {
 	}
 }
 
-func TestUpdateActorSnapshotTag_Conflict(t *testing.T) {
-	_, s, ctx := setupTest(t)
+// seedTaggedSnapshot stores a snapshot and an Atespace-scoped tag pointing at
+// it, and returns the stored tag.
+func seedTaggedSnapshot(t *testing.T, s *Persistence, ctx context.Context, snapshotName, tagName string) *ateapipb.ActorSnapshotTag {
+	t.Helper()
 	if _, err := s.CreateActorSnapshot(ctx, &ateapipb.ActorSnapshot{
-		Metadata:    &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "snapshot-1"},
-		SnapshotUri: "gs://bucket/root/snapshots/" + testAtespace + "/snapshot-1",
+		Metadata:    &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: snapshotName},
+		SnapshotUri: "gs://bucket/root/snapshots/" + testAtespace + "/" + snapshotName,
 	}); err != nil {
-		t.Fatal(err)
+		t.Fatalf("CreateActorSnapshot(%s) failed: %v", snapshotName, err)
 	}
-	if _, err := s.TagActorSnapshot(ctx, testAtespace, "snapshot-1", &ateapipb.ActorSnapshotTag{
-		Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "tag-1"},
-	}); err != nil {
-		t.Fatal(err)
-	}
-	_, tag1, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
+	tagged, err := s.TagActorSnapshot(ctx, testAtespace, snapshotName, &ateapipb.ActorSnapshotTag{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: tagName},
+		Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE,
+	})
 	if err != nil {
-		t.Fatal(err)
+		t.Fatalf("TagActorSnapshot(%s) failed: %v", tagName, err)
 	}
-	_, tag2, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
-	if err != nil {
-		t.Fatal(err)
+	return tagged
+}
+
+func TestUpdateActorSnapshotTag_MutateErrorAreNotRetried(t *testing.T) {
+	_, s, ctx := setupTest(t)
+	tagged := seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+
+	var mutationError = errors.New("mutation error")
+
+	callsToMutateFn := 0
+	_, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		callsToMutateFn++
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return fmt.Errorf("tag %s/%s: %w", testAtespace, "tag-1", mutationError)
+	})
+	// The error must arrive intact
+	if !errors.Is(err, mutationError) {
+		t.Errorf("UpdateActorSnapshotTag error = %v, want one wrapping mutationError", err)
 	}
-	if _, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED, tag1.GetMetadata().GetVersion()); err != nil {
-		t.Fatal(err)
+	// Mutation errors are non-retriable
+	if callsToMutateFn != 1 {
+		t.Errorf("mutate ran %d times, want exactly 1 (a rejected precondition must not be retried)", callsToMutateFn)
 	}
 
-	_, err = s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE, tag2.GetMetadata().GetVersion())
+	_, got, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
+	if err != nil {
+		t.Fatalf("GetActorSnapshotByTag failed: %v", err)
+	}
+	if diff := cmp.Diff(tagged, got, protocmp.Transform()); diff != "" {
+		t.Errorf("aborted mutation was persisted (-tagged +got):\n%s", diff)
+	}
+}
+
+func TestUpdateActorSnapshotTag_DiscardsServerOwnedFieldsEdits(t *testing.T) {
+	_, s, ctx := setupTest(t)
+	tagged := seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+
+	updated, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		// Metadata is server-owned: a closure must not be able to change it.
+		toUpdate.Metadata.Uid = "forged-uid"
+		toUpdate.Metadata.Version = 99
+		toUpdate.Metadata.CreateTime = nil
+		toUpdate.Metadata.UpdateTime = nil
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("UpdateActorSnapshotTag failed: %v", err)
+	}
+
+	if got := updated.GetMetadata().GetUid(); got != tagged.GetMetadata().GetUid() {
+		t.Errorf("uid = %q, want the server-assigned %q", got, tagged.GetMetadata().GetUid())
+	}
+	if got := updated.GetMetadata().GetVersion(); got != tagged.GetMetadata().GetVersion()+1 {
+		t.Errorf("version = %d, want %d (one past the stored version, not the forged value)", got, tagged.GetMetadata().GetVersion()+1)
+	}
+	if got := updated.GetMetadata().GetCreateTime(); got == nil || !got.AsTime().Equal(tagged.GetMetadata().GetCreateTime().AsTime()) {
+		t.Errorf("create_time = %v, want the creation value %v", got, tagged.GetMetadata().GetCreateTime())
+	}
+	if got, want := updated.GetScope(), ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED; got != want {
+		t.Errorf("scope = %v, want %v: discarding metadata edits must not discard the mutation", got, want)
+	}
+}
+
+// TestUpdateActorSnapshotTag_RejectsImmutableFieldChange covers the fields a
+// mutation may not touch. Unlike the server-owned metadata, which is silently
+// restored, these fail the call: a caller that renamed a tag or repointed it at
+// another snapshot asked for something the store cannot do, and must hear about
+// it.
+func TestUpdateActorSnapshotTag_RejectsImmutableFieldChange(t *testing.T) {
+	tests := []struct {
+		name      string
+		mutate    func(toUpdate *ateapipb.ActorSnapshotTag)
+		wantField string
+	}{
+		{
+			name:      "atespace",
+			mutate:    func(toUpdate *ateapipb.ActorSnapshotTag) { toUpdate.Metadata.Atespace = "other-atespace" },
+			wantField: "metadata.atespace",
+		},
+		{
+			name:      "name",
+			mutate:    func(toUpdate *ateapipb.ActorSnapshotTag) { toUpdate.Metadata.Name = "other-name" },
+			wantField: "metadata.name",
+		},
+		{
+			name:      "snapshot atespace",
+			mutate:    func(toUpdate *ateapipb.ActorSnapshotTag) { toUpdate.Snapshot.Atespace = "other-atespace" },
+			wantField: "snapshot.atespace",
+		},
+		{
+			name:      "snapshot name",
+			mutate:    func(toUpdate *ateapipb.ActorSnapshotTag) { toUpdate.Snapshot.Name = "other-snapshot" },
+			wantField: "snapshot.name",
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, s, ctx := setupTest(t)
+			tagged := seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+
+			_, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+				// Paired with a legitimate edit, so the rejection cannot be
+				// mistaken for a no-op mutation.
+				toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+				tt.mutate(toUpdate)
+				return nil
+			})
+			// The message must name the offending field: the closure is buggy,
+			// and whoever has to fix it only has this error to go on.
+			if want := tt.wantField + " is immutable"; err == nil || !strings.Contains(err.Error(), want) {
+				t.Errorf("UpdateActorSnapshotTag changing %s = %v, want an error containing %q", tt.name, err, want)
+			}
+
+			_, got, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
+			if err != nil {
+				t.Fatalf("GetActorSnapshotByTag failed: %v", err)
+			}
+			if diff := cmp.Diff(tagged, got, protocmp.Transform()); diff != "" {
+				t.Errorf("rejected mutation was persisted anyway (-tagged +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestUpdateActorSnapshotTag_RetriesOnConcurrentWrite(t *testing.T) {
+	mr, s, ctx := setupTest(t)
+	seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+	tagKey := actorSnapshotTagDBKey(testAtespace, "tag-1")
+
+	// A separate client, so its write lands outside the transaction's connection.
+	otherClient := redis.NewClusterClient(&redis.ClusterOptions{Addrs: []string{mr.Addr()}})
+	t.Cleanup(func() { otherClient.Close() })
+
+	attempts := 0
+	interceptor := &watchInterceptor{redisClient: s.rdb, before: func() {
+		// Only the first attempt races. We do this to make sure the second retry
+		// will succeed.
+		if attempts > 0 {
+			return
+		}
+		_, concurrent, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
+		if err != nil {
+			t.Errorf("GetActorSnapshotByTag for concurrent write failed: %v", err)
+			return
+		}
+		// Repointing the tag is not something a mutation may do, but a writer
+		// holding the key can: the retry must carry it forward, not revert it.
+		concurrent.Snapshot = &ateapipb.ObjectRef{Atespace: testAtespace, Name: "snapshot-2"}
+		val, err := protojson.Marshal(concurrent)
+		if err != nil {
+			t.Errorf("protojson.Marshal failed: %v", err)
+			return
+		}
+		if err := otherClient.Set(ctx, tagKey, val, 0).Err(); err != nil {
+			t.Errorf("concurrent Set failed: %v", err)
+		}
+	}}
+	racing := &Persistence{rdb: interceptor, lockTTL: defaultLockTTL}
+
+	updated, err := racing.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		attempts++
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return nil
+	})
+	if err != nil {
+		t.Fatalf("UpdateActorSnapshotTag failed: %v", err)
+	}
+	if attempts < 2 {
+		t.Errorf("mutate ran %d times, want at least 2: the first write is racey and must be rejected", attempts)
+	}
+	if got, want := updated.GetScope(), ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED; got != want {
+		t.Errorf("scope = %v, want %v", got, want)
+	}
+	// The concurrent tx repointed the tag at snapshot-2. That change should
+	// survive instead of being reverted by a mutation computed against the
+	// older state.
+	if got := updated.GetSnapshot().GetName(); got != "snapshot-2" {
+		t.Errorf("snapshot.name = %q, want %q: the retry clobbered the concurrent write", got, "snapshot-2")
+	}
+}
+
+func TestUpdateActorSnapshotTag_NotFound(t *testing.T) {
+	_, s, ctx := setupTest(t)
+	_, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "does-not-exist", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		t.Error("mutate must not run for a missing tag")
+		return nil
+	})
+	if !errors.Is(err, store.ErrNotFound) {
+		t.Errorf("expected store.ErrNotFound, got %v", err)
+	}
+}
+
+func TestUpdateActorSnapshotTag_RejectsStaleUID(t *testing.T) {
+	_, s, ctx := setupTest(t)
+
+	original := seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+	if _, err := s.DeleteActorSnapshotTag(ctx, testAtespace, "tag-1"); err != nil {
+		t.Fatalf("DeleteActorSnapshotTag failed: %v", err)
+	}
+	recreated, err := s.TagActorSnapshot(ctx, testAtespace, "snapshot-1", &ateapipb.ActorSnapshotTag{
+		Metadata: &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: "tag-1"},
+		Scope:    ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE,
+	})
+	if err != nil {
+		t.Fatalf("re-tag TagActorSnapshot failed: %v", err)
+	}
+	if recreated.GetMetadata().GetUid() == original.GetMetadata().GetUid() {
+		t.Fatalf("recreated tag reused uid %s, want a fresh one", recreated.GetMetadata().GetUid())
+	}
+	// The version reset to 1 along with the uid, so a version guard alone would
+	// have waved this write through. Only the uid distinguishes the lifecycles.
+	if got, want := recreated.GetMetadata().GetVersion(), original.GetMetadata().GetVersion(); got != want {
+		t.Fatalf("recreated version = %d, want %d: the version cannot tell the lifecycles apart", got, want)
+	}
+
+	// Pins the incarnation alone: the observed tag carries the original uid and
+	// no version.
+	pinUID := &ateapipb.ActorSnapshotTag{
+		Metadata: &ateapipb.ResourceMetadata{Uid: original.GetMetadata().GetUid(), Version: store.AnyVersion},
+	}
+	_, err = s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", store.WithPrecondition(pinUID, func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		t.Error("mutate ran past its precondition once the pinned incarnation was gone")
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return nil
+	}))
+	if !errors.Is(err, store.ErrUIDConflict) {
+		t.Errorf("UpdateActorSnapshotTag error = %v, want one matching store.ErrUIDConflict", err)
+	}
+
+	// The version guard was waived, so this is the incarnation failure alone.
+	if errors.Is(err, store.ErrVersionConflict) {
+		t.Errorf("UpdateActorSnapshotTag error = %v, want no store.ErrVersionConflict match: no version was pinned", err)
+	}
+
+	_, stored, err := s.GetActorSnapshotByTag(ctx, testAtespace, "tag-1")
+	if err != nil {
+		t.Fatalf("GetActorSnapshotByTag failed: %v", err)
+	}
+	if diff := cmp.Diff(recreated, stored, protocmp.Transform()); diff != "" {
+		t.Errorf("the rejected update still wrote (-recreated +stored):\n%s", diff)
+	}
+}
+
+func TestUpdateActorSnapshotTag_RejectsStaleVersion(t *testing.T) {
+	_, s, ctx := setupTest(t)
+
+	tagged := seedTaggedSnapshot(t, s, ctx, "snapshot-1", "tag-1")
+
+	if _, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_PUBLISHED
+		return nil
+	}); err != nil {
+		t.Fatalf("UpdateActorSnapshotTag failed: %v", err)
+	}
+
+	// The write above moved the version, so tagged is now a stale observation.
+	_, err := s.UpdateActorSnapshotTag(ctx, testAtespace, "tag-1", store.WithPrecondition(tagged, func(toUpdate *ateapipb.ActorSnapshotTag) error {
+		t.Error("mutate ran past its precondition once the pinned version had moved")
+		toUpdate.Scope = ateapipb.ActorSnapshotTagScope_ACTOR_SNAPSHOT_TAG_SCOPE_ATESPACE
+		return nil
+	}))
 	if !errors.Is(err, store.ErrVersionConflict) {
-		t.Fatalf("UpdateActorSnapshotTag error = %v, want ErrVersionConflict", err)
+		t.Errorf("UpdateActorSnapshotTag error = %v, want one matching store.ErrVersionConflict", err)
+	}
+	// The uid still matches, so this is not the incarnation failure: callers key
+	// their retry decision off the difference.
+	if errors.Is(err, store.ErrUIDConflict) {
+		t.Errorf("UpdateActorSnapshotTag error = %v, want no store.ErrUIDConflict match: the incarnation is unchanged", err)
 	}
 }
 

cmd/ateapi/internal/store/store.go

@@ -58,19 +58,19 @@ type Interface interface {
 	// mutated. Returns ErrAlreadyExists if key is taken.
 	CreateActor(ctx context.Context, actor *ateapipb.Actor) (*ateapipb.Actor, error)
 
-	// UpdateActor performs a transactional read-modify-write and returns the updated
+	// UpdateActor performs a transactional read-modify-write and returns the stored
 	// actor with advanced metadata (version, update_time).
 	//
 	// mutate receives the stored actor and edits it in place. The mutated actor is
 	// written iff mutate returns nil. A mutate that must only land on the actor the
-	// caller observed guards itself with CheckActorPrecondition.
+	// caller observed wraps itself in WithPrecondition.
 	//
 	// mutate may run more than once, because the store retries when a concurrent
 	// write invalidates the transaction.
 	//
 	// Returns ErrNotFound if missing, ErrVersionConflict if the retry budget is
 	// exhausted, or the mutate's error verbatim otherwise.
-	UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(dbActor *ateapipb.Actor) error) (*ateapipb.Actor, error)
+	UpdateActor(ctx context.Context, actorRef resources.ActorRef, mutate func(toUpdate *ateapipb.Actor) error) (*ateapipb.Actor, error)
 
 	// Removes an actor and returns the deleted resource. Returns ErrNotFound if
 	// missing, or ErrFailedPrecondition if not suspended.
@@ -96,8 +96,20 @@ type Interface interface {
 	// Adds an immutable Atespace-owned tag to an ActorSnapshot.
 	TagActorSnapshot(ctx context.Context, atespace, name string, tag *ateapipb.ActorSnapshotTag) (*ateapipb.ActorSnapshotTag, error)
 
-	// Updates a tag's reuse scope.
-	UpdateActorSnapshotTag(ctx context.Context, atespace, name string, scope ateapipb.ActorSnapshotTagScope, expectedVersion int64) (*ateapipb.ActorSnapshotTag, error)
+	// UpdateActorSnapshotTag performs a transactional read-modify-write on the tag
+	// addressed by atespace and name, and returns the stored ActorSnapshotTag with
+	// advanced metadata (version, update_time).
+	//
+	// mutate receives the stored tag and edits it in place. The mutated tag is
+	// written iff mutate returns nil. A mutate that must only land on the tag the
+	// caller observed wraps itself in WithPrecondition.
+	//
+	// mutate may run more than once, because the store retries when a concurrent
+	// write invalidates the transaction.
+	//
+	// Returns ErrNotFound if missing, ErrVersionConflict if the retry budget is
+	// exhausted, or the mutate's error verbatim otherwise.
+	UpdateActorSnapshotTag(ctx context.Context, atespace, name string, mutate func(toUpdate *ateapipb.ActorSnapshotTag) error) (*ateapipb.ActorSnapshotTag, error)
 
 	// Deletes and returns a tag.
 	DeleteActorSnapshotTag(ctx context.Context, atespace, name string) (*ateapipb.ActorSnapshotTag, error)
@@ -153,23 +165,20 @@ type Interface interface {
 }
 
 const (
-	// AnyUID accepts whichever actor holds the atespace and name at write time.
+	// AnyUID accepts whichever object holds the atespace and name at write time.
 	AnyUID = ""
 	// AnyVersion accepts whatever revision the store is at.
 	AnyVersion int64 = 0
 )
 
-// CheckActorPrecondition reports whether dbActor is still the actor the caller
+// checkPrecondition reports whether md still describes the object the caller
 // observed, pinned on the uid and version it read, each waivable with AnyUID or
-// AnyVersion. Version guards against concurrent writes, uid against actor
-// atespace/name re-use across actor lifecycles.
+// AnyVersion. Version guards against concurrent writes, uid against
+// atespace/name re-use across object lifecycles.
 //
-// Call it at the top of an UpdateActor mutation so the write is conditional on
-// the stored actor the transaction actually read, not on one read earlier
-// outside of it. Returns ErrUIDConflict or ErrVersionConflict, which UpdateActor
-// surfaces verbatim.
-func CheckActorPrecondition(dbActor *ateapipb.Actor, uid string, version int64) error {
-	md := dbActor.GetMetadata()
+// Returns ErrUIDConflict or ErrVersionConflict, which the update surfaces
+// verbatim.
+func checkPrecondition(md *ateapipb.ResourceMetadata, uid string, version int64) error {
 	if uid != AnyUID && uid != md.GetUid() {
 		return ErrUIDConflict
 	}
@@ -179,6 +188,29 @@ func CheckActorPrecondition(dbActor *ateapipb.Actor, uid string, version int64)
 	return nil
 }
 
+// hasResourceMetadata is an object the store addresses by atespace and name,
+// and whose identity a caller can pin with WithPrecondition.
+type hasResourceMetadata interface {
+	GetMetadata() *ateapipb.ResourceMetadata
+}
+
+// WithPrecondition returns a mutation that runs mutate only if the stored
+// object is still the one the caller observed, pinned on observed's uid and
+// version.
+//
+// An observed object carrying no uid or version pins nothing, so an unguarded
+// client update stays unguarded. To pin one of the two and waive the other,
+// pass an observed object carrying only the field to pin.
+func WithPrecondition[T hasResourceMetadata](observed T, mutate func(stored T) error) func(stored T) error {
+	uid, version := observed.GetMetadata().GetUid(), observed.GetMetadata().GetVersion()
+	return func(stored T) error {
+		if err := checkPrecondition(stored.GetMetadata(), uid, version); err != nil {
+			return err
+		}
+		return mutate(stored)
+	}
+}
+
 // WorkerEventType indicates the type of change to a Worker.
 type WorkerEventType int
 

cmd/ateapi/internal/store/store_test.go

@@ -21,14 +21,15 @@ import (
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 )
 
-func TestCheckActorPrecondition(t *testing.T) {
+func TestWithPrecondition(t *testing.T) {
 	const (
 		storedUID = "stored-uid"
 		staleUID  = "stale-uid"
 		storedVer = int64(7)
 		staleVer  = int64(6)
 	)
-	dbActor := &ateapipb.Actor{
+	// The actor the transaction reads, which the mutation is checked against.
+	stored := &ateapipb.Actor{
 		Metadata: &ateapipb.ResourceMetadata{
 			Atespace: "test-atespace",
 			Name:     "actor-1",
@@ -36,58 +37,110 @@ func TestCheckActorPrecondition(t *testing.T) {
 			Version:  storedVer,
 		},
 	}
+	errMutate := errors.New("mutate failed")
 
 	tests := []struct {
-		name    string
-		uid     string
-		version int64
-		wantErr error
+		name        string
+		observed    *ateapipb.Actor
+		mutateErr   error
+		wantErr     error
+		wantMutated bool
 	}{
 		{
-			name:    "both waived",
-			uid:     AnyUID,
-			version: AnyVersion,
-			wantErr: nil,
+			name:        "the observed object is still the stored one",
+			observed:    stored,
+			wantErr:     nil,
+			wantMutated: true,
 		},
 		{
-			name:    "both guarded and both match",
-			uid:     storedUID,
-			version: storedVer,
-			wantErr: nil,
+			name:        "the mutate error is surfaced verbatim",
+			observed:    stored,
+			mutateErr:   errMutate,
+			wantErr:     errMutate,
+			wantMutated: true,
 		},
 		{
-			name:    "uid guarded, version waived, tolerates the moved version",
-			uid:     storedUID,
-			version: AnyVersion,
-			wantErr: nil,
+			name:        "an unguarded observed object pins nothing",
+			observed:    &ateapipb.Actor{Metadata: &ateapipb.ResourceMetadata{Atespace: "test-atespace", Name: "actor-1"}},
+			wantErr:     nil,
+			wantMutated: true,
 		},
 		{
-			name:    "version guarded, uid waived, still catches the moved version",
-			uid:     AnyUID,
-			version: staleVer,
-			wantErr: ErrVersionConflict,
+			name: "uid guarded, version waived, tolerates the moved version",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: storedUID, Version: AnyVersion},
+			},
+			wantErr:     nil,
+			wantMutated: true,
 		},
 		{
-			name:    "uid guarded, version waived, still catches the new incarnation",
-			uid:     staleUID,
-			version: AnyVersion,
-			wantErr: ErrUIDConflict,
+			name: "version guarded, uid waived, tolerates the new incarnation",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: AnyUID, Version: storedVer},
+			},
+			wantErr:     nil,
+			wantMutated: true,
+		},
+		{
+			name: "the name now addresses a different incarnation",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: staleUID, Version: storedVer},
+			},
+			wantErr:     ErrUIDConflict,
+			wantMutated: false,
+		},
+		{
+			name: "the version moved under the caller",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: storedUID, Version: staleVer},
+			},
+			wantErr:     ErrVersionConflict,
+			wantMutated: false,
+		},
+		{
+			name: "uid guarded, version waived, still catches the new incarnation",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: staleUID, Version: AnyVersion},
+			},
+			wantErr:     ErrUIDConflict,
+			wantMutated: false,
+		},
+		{
+			name: "version guarded, uid waived, still catches the moved version",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: AnyUID, Version: staleVer},
+			},
+			wantErr:     ErrVersionConflict,
+			wantMutated: false,
 		},
 		{
 			// The uid is reported first: a new incarnation makes the version
 			// meaningless, and it is the failure a retry can never resolve.
-			name:    "both stale reports the uid conflict",
-			uid:     staleUID,
-			version: staleVer,
-			wantErr: ErrUIDConflict,
+			name: "both stale reports the uid conflict",
+			observed: &ateapipb.Actor{
+				Metadata: &ateapipb.ResourceMetadata{Uid: staleUID, Version: staleVer},
+			},
+			wantErr:     ErrUIDConflict,
+			wantMutated: false,
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			err := CheckActorPrecondition(dbActor, tt.uid, tt.version)
-			if !errors.Is(err, tt.wantErr) {
-				t.Errorf("CheckActorPrecondition(dbActor, %q, %d) = %v, want one matching %v", tt.uid, tt.version, err, tt.wantErr)
+			mutated := false
+			mutate := WithPrecondition(tt.observed, func(toUpdate *ateapipb.Actor) error {
+				mutated = true
+				if toUpdate != stored {
+					t.Errorf("mutate got actor %v, want the stored one", toUpdate)
+				}
+				return tt.mutateErr
+			})
+
+			if err := mutate(stored); !errors.Is(err, tt.wantErr) {
+				t.Errorf("mutate(stored) = %v, want one matching %v", err, tt.wantErr)
+			}
+			if mutated != tt.wantMutated {
+				t.Errorf("mutate ran = %t, want %t", mutated, tt.wantMutated)
 			}
 		})
 	}