atelet: Reject stale snapshot generations

cmd/atelet/internal/snapshottracker/actor.go

@@ -191,6 +191,10 @@ type actorSnapshots struct {
 	// the actor in progress on this node. While it is nonzero, atelet does not
 	// detach the actor, and the entry stays in memory.
 	holds int
+	// inFlight holds the snapshots of Checkpoints and PromoteSnapshots of the
+	// actor currently in progress on this node. newest includes them, so an
+	// older generation cannot start while a newer one is in flight.
+	inFlight []*SnapshotRef
 	// promotes counts the PromoteSnapshots of the actor in progress on this
 	// node. While it is nonzero, atelet does not start a background upload of
 	// the actor.
@@ -207,6 +211,28 @@ type actorSnapshots struct {
 	unsaved bool
 }
 
+// localNeeded reports whether CommitLocal should record snap as LatestLocal:
+// LatestLocal does not already hold snap or a newer snapshot, and no other
+// committed snapshot has a higher generation or the same generation under
+// another UUID.
+func (a *actorSnapshots) localNeeded(snap SnapshotRef) bool {
+	if a.LatestLocal != nil && a.LatestLocal.generation() >= snap.generation() {
+		return false
+	}
+	for _, r := range []*SnapshotRef{a.Uploading, a.LatestDurable} {
+		if r == nil {
+			continue
+		}
+		if r.generation() > snap.generation() {
+			return false
+		}
+		if r.generation() == snap.generation() && !sameUUID(r, &snap) {
+			return false
+		}
+	}
+	return true
+}
+
 // uploadNeeded reports whether ref still needs uploading: neither it nor a
 // newer snapshot is durable.
 func (a *actorSnapshots) uploadNeeded(ref SnapshotRef) bool {
@@ -331,6 +357,52 @@ func (a *actorSnapshots) runningDetach() *SnapshotRef {
 	return &a.op.target
 }
 
+// opTarget returns the snapshot an operation in progress acts on, or nil.
+func (a *actorSnapshots) opTarget() *SnapshotRef {
+	if a.op == nil {
+		return nil
+	}
+	return &a.op.target
+}
+
+// newest returns the actor's snapshot with the highest generation, or nil
+// when atelet holds no snapshot for the actor. It includes snapshots whose
+// Checkpoint or PromoteSnapshot is in flight.
+func (a *actorSnapshots) newest() *SnapshotRef {
+	var newest *SnapshotRef
+	consider := func(r *SnapshotRef) {
+		if r != nil && (newest == nil || r.generation() > newest.generation()) {
+			newest = r
+		}
+	}
+	consider(a.LatestLocal)
+	consider(a.Uploading)
+	consider(a.LatestDurable)
+	consider(a.opTarget())
+	for _, r := range a.inFlight {
+		consider(r)
+	}
+	return newest
+}
+
+// checkGeneration returns a FailedPrecondition error if snap has a stale
+// generation for the actor: below the actor's newest snapshot, or equal to
+// it under another snapshot_uuid.
+func (a *actorSnapshots) checkGeneration(snap *ateletpb.Snapshot) error {
+	newest := a.newest()
+	if newest == nil {
+		return nil
+	}
+	switch {
+	case snap.GetGeneration() < newest.generation():
+		return apierror.FailedPrecondition("snapshot %q has generation %d, below the actor's newest snapshot %q at generation %d", snap.GetSnapshotUuid(), snap.GetGeneration(), newest.uuid(), newest.generation())
+	case snap.GetGeneration() == newest.generation() && snap.GetSnapshotUuid() != newest.uuid():
+		return apierror.FailedPrecondition("snapshot %q has generation %d, which the actor's newest snapshot %q already holds", snap.GetSnapshotUuid(), snap.GetGeneration(), newest.uuid())
+	default:
+		return nil
+	}
+}
+
 // uploadDue returns when atelet next runs a background upload of the actor:
 // uploadDelay after UndurableSince, or at retryAt if that is later. An upload
 // or promote in progress can delay it. ok is false when there is nothing

cmd/atelet/internal/snapshottracker/background.go

@@ -395,12 +395,18 @@ func (t *Tracker) waitIdleLocked(ctx context.Context, actorUID string, check fun
 // waitForDetachLocked waits, as long as ctx allows, until no detach of the
 // actor is in progress, and returns the actor's entry, which may be nil. It
 // releases the actor's lock while it waits, and holds it again when it
-// returns. Checkpoint, Run, and Restore call this first so they never race
-// with a detach pruning the actor's local directory, without waiting behind
-// an upload.
-func (t *Tracker) waitForDetachLocked(ctx context.Context, actorUID string) (*actorSnapshots, error) {
+// returns. check, if not nil, vets the entry on every pass: an error from it
+// is returned at once, without waiting further. Checkpoint, Run, and Restore
+// call this first so they never race with a detach pruning the actor's local
+// directory, without waiting behind an upload.
+func (t *Tracker) waitForDetachLocked(ctx context.Context, actorUID string, check func(*actorSnapshots) error) (*actorSnapshots, error) {
 	for {
 		a := t.findLocked(actorUID)
+		if a != nil && check != nil {
+			if err := check(a); err != nil {
+				return nil, err
+			}
+		}
 		if a == nil || a.runningDetach() == nil {
 			return a, nil
 		}
@@ -412,11 +418,19 @@ func (t *Tracker) waitForDetachLocked(ctx context.Context, actorUID string) (*ac
 
 // waitForDetachOrSnapshotUploadLocked waits, as long as ctx allows, until no
 // detach of the actor and no background upload of snap is in progress, and
-// returns the actor's entry, which may be nil. PromoteSnapshot calls this
-// first so it never races with a detach or uploads the same snapshot twice.
+// returns the actor's entry, which may be nil. It checks snap's generation on
+// every pass, so a stale generation fails at once without waiting, and one
+// superseded while waiting fails as soon as the wait ends. PromoteSnapshot
+// calls this first so it never races with a detach or uploads the same
+// snapshot twice.
 func (t *Tracker) waitForDetachOrSnapshotUploadLocked(ctx context.Context, actorUID string, snap SnapshotRef) (*actorSnapshots, error) {
 	for {
 		a := t.findLocked(actorUID)
+		if a != nil {
+			if err := a.checkGeneration(snap.Snapshot); err != nil {
+				return nil, err
+			}
+		}
 		if a == nil || a.op == nil {
 			return a, nil
 		}

cmd/atelet/internal/snapshottracker/background_test.go

@@ -335,9 +335,9 @@ func TestTrackerBackoff(t *testing.T) {
 		name string
 		call func()
 	}{
-		{"CommitLocal", func() { tr.CommitLocal(trackedActorUID, testRef("s0", 1, clk.Now())) }},
-		{"CommitDurable", func() { tr.CommitDurable(trackedActorUID, testRef("s1", 1, clk.Now())) }},
-		{"CommitPromoted", func() { tr.CommitPromoted(trackedActorUID, testRef("s2", 2, clk.Now())) }},
+		{"CommitLocal", func() { tr.CommitLocal(trackedActorUID, testRef("s1", 1, clk.Now())) }},
+		{"CommitDurable", func() { tr.CommitDurable(trackedActorUID, testRef("s2", 2, clk.Now())) }},
+		{"CommitPromoted", func() { tr.CommitPromoted(trackedActorUID, testRef("s3", 3, clk.Now())) }},
 		{"Forget", func() {
 			if err := tr.Forget(trackedActorUID); err != nil {
 				t.Fatalf("Forget: %v", err)

cmd/atelet/internal/snapshottracker/detach_test.go

@@ -170,7 +170,6 @@ func TestTrackerDetach(t *testing.T) {
 		store := newFakeSnapshotStore()
 		tr := newTestTracker(t, store)
 		local := testRef("s1", 1, tr.clock.Now())
-		local.Snapshot.Object = nil
 		tr.CommitLocal(trackedActorUID, local)
 
 		promoted := testRef("s1", 1, time.Time{})
@@ -258,9 +257,10 @@ func TestTrackerDetach(t *testing.T) {
 
 	// A refused detach keeps the local snapshot and uploads nothing.
 	for _, tc := range []struct {
-		name     string
-		setup    func(t *testing.T, tr *Tracker)
-		wantCode codes.Code
+		name        string
+		modifyLocal func(*SnapshotRef)
+		setup       func(t *testing.T, tr *Tracker)
+		wantCode    codes.Code
 	}{
 		{
 			name: "refused after a start",
@@ -287,7 +287,7 @@ func TestTrackerDetach(t *testing.T) {
 		{
 			name: "refused while a checkpoint is in progress",
 			setup: func(t *testing.T, tr *Tracker) {
-				end, err := tr.BeginCheckpoint(t.Context(), trackedActorUID)
+				end, err := tr.BeginCheckpoint(t.Context(), trackedActorUID, testRef("s2", 2, time.Time{}).Snapshot)
 				if err != nil {
 					t.Fatalf("beginCheckpoint: %v", err)
 				}
@@ -305,29 +305,33 @@ func TestTrackerDetach(t *testing.T) {
 		},
 		{
 			name: "refused when durable URI mismatches local",
+			modifyLocal: func(local *SnapshotRef) {
+				local.Snapshot.Object.SnapshotUri += "-other"
+			},
 			setup: func(t *testing.T, tr *Tracker) {
 				tr.CommitDurable(trackedActorUID, testRef("s1", 1, tr.clock.Now()))
-				local := testRef("s1", 1, tr.clock.Now())
-				local.Snapshot.Object.SnapshotUri += "-other"
-				tr.CommitLocal(trackedActorUID, local)
 			},
 			wantCode: codes.FailedPrecondition,
 		},
 		{
 			name: "refused without an object storage location",
-			setup: func(t *testing.T, tr *Tracker) {
-				local := testRef("s1", 1, tr.clock.Now())
+			modifyLocal: func(local *SnapshotRef) {
 				local.Snapshot.Object = nil
-				tr.CommitLocal(trackedActorUID, local)
 			},
 			wantCode: codes.FailedPrecondition,
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			store := newFakeSnapshotStore()
 			tr := newTestTracker(t, store)
-			tr.CommitLocal(trackedActorUID, testRef("s1", 1, tr.clock.Now()))
-			tc.setup(t, tr)
+			local := testRef("s1", 1, tr.clock.Now())
+			if tc.modifyLocal != nil {
+				tc.modifyLocal(&local)
+			}
+			tr.CommitLocal(trackedActorUID, local)
+			if tc.setup != nil {
+				tc.setup(t, tr)
+			}
 			if _, err := tr.Detach(t.Context(), trackedActorUID); apierror.Code(err) != tc.wantCode {
 				t.Errorf("detach = %v, want %v", err, tc.wantCode)
 			}
@@ -409,7 +413,7 @@ func TestTrackerDetachTimer(t *testing.T) {
 
 	tr.CommitLocal(trackedActorUID, testRef("s2", 2, tr.clock.Now()))
 	due := detachDueOf(tr, trackedActorUID)
-	endCheckpoint, err := tr.BeginCheckpoint(t.Context(), trackedActorUID)
+	endCheckpoint, err := tr.BeginCheckpoint(t.Context(), trackedActorUID, testRef("s3", 3, time.Time{}).Snapshot)
 	if err != nil {
 		t.Fatalf("beginCheckpoint: %v", err)
 	}
@@ -485,8 +489,10 @@ func TestTrackerWaitsForDetach(t *testing.T) {
 	waitStarted(t, store)
 
 	for name, begin := range map[string]func(context.Context, string) (func(), error){
-		"beginStart":      tr.BeginStart,
-		"beginCheckpoint": tr.BeginCheckpoint,
+		"beginStart": tr.BeginStart,
+		"beginCheckpoint": func(ctx context.Context, uid string) (func(), error) {
+			return tr.BeginCheckpoint(ctx, uid, testRef("s2", 2, time.Time{}).Snapshot)
+		},
 	} {
 		result := make(chan error, 1)
 		results[name] = result
@@ -522,6 +528,75 @@ func TestTrackerWaitsForDetach(t *testing.T) {
 	}
 }
 
+// BeginCheckpoint and BeginPromote fail at once during a detach when the
+// requested snapshot is already stale, and fail once the detach ends when a
+// newer snapshot was committed while they waited.
+func TestTrackerDetachWaitChecksGeneration(t *testing.T) {
+	store := newFakeSnapshotStore()
+	store.setBlocking(true)
+	tr := newTestTracker(t, store)
+	tr.CommitLocal(trackedActorUID, testRef("s1", 2, tr.clock.Now()))
+	detached := startDetach(t, tr)
+	waitStarted(t, store)
+
+	stale := testRef("s0", 1, time.Time{})
+	returnsPromptly(t, "BeginCheckpoint of stale generation during detach", func() {
+		if _, err := tr.BeginCheckpoint(t.Context(), trackedActorUID, stale.Snapshot); apierror.Code(err) != codes.FailedPrecondition {
+			t.Errorf("BeginCheckpoint = %v, want FailedPrecondition", err)
+		}
+	})
+	returnsPromptly(t, "BeginPromote of stale generation during detach", func() {
+		if _, _, err := tr.BeginPromote(t.Context(), trackedActorUID, stale); apierror.Code(err) != codes.FailedPrecondition {
+			t.Errorf("BeginPromote = %v, want FailedPrecondition", err)
+		}
+	})
+
+	superseded := testRef("s2", 3, time.Time{})
+	results := make(map[string]chan error)
+	for name, begin := range map[string]func() (func(), error){
+		"beginCheckpoint": func() (func(), error) {
+			return tr.BeginCheckpoint(t.Context(), trackedActorUID, superseded.Snapshot)
+		},
+		"beginPromote": func() (func(), error) {
+			_, end, err := tr.BeginPromote(t.Context(), trackedActorUID, superseded)
+			return end, err
+		},
+	} {
+		result := make(chan error, 1)
+		results[name] = result
+		go func() {
+			end, err := begin()
+			if err == nil {
+				end()
+			}
+			result <- err
+		}()
+	}
+	for name, result := range results {
+		select {
+		case err := <-result:
+			t.Fatalf("%s returned (%v) during a detach", name, err)
+		case <-time.After(50 * time.Millisecond):
+		}
+	}
+
+	tr.CommitDurable(trackedActorUID, testRef("s3", 4, tr.clock.Now()))
+	store.release <- nil
+	if err := <-detached; err != nil {
+		t.Fatalf("detach: %v", err)
+	}
+	for name, result := range results {
+		select {
+		case err := <-result:
+			if apierror.Code(err) != codes.FailedPrecondition {
+				t.Errorf("%s = %v, want FailedPrecondition", name, err)
+			}
+		case <-time.After(10 * time.Second):
+			t.Fatalf("%s did not return after the detach", name)
+		}
+	}
+}
+
 // A failed detach keeps the local snapshot, and is retried no sooner than the
 // retry delay after it, nor than it was due. The upload is held back the retry
 // delay too.
@@ -618,11 +693,12 @@ func TestTrackerDetachWithoutStore(t *testing.T) {
 func TestTrackerHoldsFirstCheckpoint(t *testing.T) {
 	store := newFakeSnapshotStore()
 	tr := newTestTracker(t, store)
-	end, err := tr.BeginCheckpoint(t.Context(), trackedActorUID)
+	s1 := testRef("s1", 1, tr.clock.Now())
+	end, err := tr.BeginCheckpoint(t.Context(), trackedActorUID, s1.Snapshot)
 	if err != nil {
 		t.Fatalf("beginCheckpoint: %v", err)
 	}
-	tr.CommitLocal(trackedActorUID, testRef("s1", 1, tr.clock.Now()))
+	tr.CommitLocal(trackedActorUID, s1)
 	if _, err := tr.Detach(t.Context(), trackedActorUID); apierror.Code(err) != codes.Aborted {
 		t.Errorf("detach during the checkpoint = %v, want Aborted", err)
 	}