From c538b68ba67fff61879dd10b6f1ada02939a242d Mon Sep 17 00:00:00 2001 From: Luiz Oliveira Date: Wed, 12 Aug 2026 10:00:22 -0400 Subject: [PATCH] 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 --- .../internal/controlapi/actor_snapshot.go | 41 +- .../controlapi/actor_snapshot_test.go | 144 +++++++ cmd/ateapi/internal/controlapi/crash.go | 11 +- .../internal/controlapi/functional_test.go | 6 +- cmd/ateapi/internal/controlapi/syncer.go | 15 +- .../internal/controlapi/update_actor.go | 13 +- .../internal/controlapi/update_actor_test.go | 8 +- .../internal/controlapi/workflow_delete.go | 13 +- .../internal/controlapi/workflow_pause.go | 36 +- .../internal/controlapi/workflow_resume.go | 37 +- .../controlapi/workflow_resume_test.go | 23 +- .../internal/controlapi/workflow_suspend.go | 36 +- .../internal/store/ateredis/ateredis.go | 124 ++++-- .../internal/store/ateredis/ateredis_test.go | 376 +++++++++++++++--- cmd/ateapi/internal/store/store.go | 62 ++- cmd/ateapi/internal/store/store_test.go | 119 ++++-- 16 files changed, 780 insertions(+), 284 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/actor_snapshot.go b/cmd/ateapi/internal/controlapi/actor_snapshot.go index 9c25caf79..ada2a2bdb 100644 --- a/cmd/ateapi/internal/controlapi/actor_snapshot.go +++ b/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 { diff --git a/cmd/ateapi/internal/controlapi/actor_snapshot_test.go b/cmd/ateapi/internal/controlapi/actor_snapshot_test.go index e3dc6a9c6..fe1c60896 100644 --- a/cmd/ateapi/internal/controlapi/actor_snapshot_test.go +++ b/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) + } +} diff --git a/cmd/ateapi/internal/controlapi/crash.go b/cmd/ateapi/internal/controlapi/crash.go index c40129efd..e12593c4b 100644 --- a/cmd/ateapi/internal/controlapi/crash.go +++ b/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...) diff --git a/cmd/ateapi/internal/controlapi/functional_test.go b/cmd/ateapi/internal/controlapi/functional_test.go index 84a87ad67..fd126c3ae 100644 --- a/cmd/ateapi/internal/controlapi/functional_test.go +++ b/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) diff --git a/cmd/ateapi/internal/controlapi/syncer.go b/cmd/ateapi/internal/controlapi/syncer.go index 3d0276dc8..77b642f23 100644 --- a/cmd/ateapi/internal/controlapi/syncer.go +++ b/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) diff --git a/cmd/ateapi/internal/controlapi/update_actor.go b/cmd/ateapi/internal/controlapi/update_actor.go index 887ebf135..60dfe8438 100644 --- a/cmd/ateapi/internal/controlapi/update_actor.go +++ b/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 { diff --git a/cmd/ateapi/internal/controlapi/update_actor_test.go b/cmd/ateapi/internal/controlapi/update_actor_test.go index 99f119484..e2696b1fe 100644 --- a/cmd/ateapi/internal/controlapi/update_actor_test.go +++ b/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) diff --git a/cmd/ateapi/internal/controlapi/workflow_delete.go b/cmd/ateapi/internal/controlapi/workflow_delete.go index e93516c41..867f69948 100644 --- a/cmd/ateapi/internal/controlapi/workflow_delete.go +++ b/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 diff --git a/cmd/ateapi/internal/controlapi/workflow_pause.go b/cmd/ateapi/internal/controlapi/workflow_pause.go index 138d596c2..b088233b5 100644 --- a/cmd/ateapi/internal/controlapi/workflow_pause.go +++ b/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 diff --git a/cmd/ateapi/internal/controlapi/workflow_resume.go b/cmd/ateapi/internal/controlapi/workflow_resume.go index 6483984a0..6666168d0 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume.go +++ b/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 } diff --git a/cmd/ateapi/internal/controlapi/workflow_resume_test.go b/cmd/ateapi/internal/controlapi/workflow_resume_test.go index 77e97044c..b9b1274b4 100644 --- a/cmd/ateapi/internal/controlapi/workflow_resume_test.go +++ b/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) } diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend.go b/cmd/ateapi/internal/controlapi/workflow_suspend.go index 14258aded..942b00473 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend.go +++ b/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 } diff --git a/cmd/ateapi/internal/store/ateredis/ateredis.go b/cmd/ateapi/internal/store/ateredis/ateredis.go index 6ab62b769..efa2db7ee 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis.go +++ b/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) { diff --git a/cmd/ateapi/internal/store/ateredis/ateredis_test.go b/cmd/ateapi/internal/store/ateredis/ateredis_test.go index 3d2b03dcf..93a997fad 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis_test.go +++ b/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) } } diff --git a/cmd/ateapi/internal/store/store.go b/cmd/ateapi/internal/store/store.go index 9c020d183..4b05d3caa 100644 --- a/cmd/ateapi/internal/store/store.go +++ b/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 diff --git a/cmd/ateapi/internal/store/store_test.go b/cmd/ateapi/internal/store/store_test.go index caf1b8ee2..38609b127 100644 --- a/cmd/ateapi/internal/store/store_test.go +++ b/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) } }) }