atelet: checkpoint idempotency

[Troy Chiu (troychiu)]

Towards #372

Checkpointing is a destructive operation that terminates the running sandbox. If a Checkpoint RPC fails or is interrupted after the sandbox is snapshotted or uploaded, subsequent retries would attempt to checkpoint a dead guest, resulting in false data-loss failures.

This PR makes checkpointing idempotent across Atelet and Ateom through two levels of completion detection:

1. Destination Fast-Forward (Atelet): If the destination (GCS or local disk) already holds the committed snapshot manifest with matching scope, Atelet fast-forwards and completes remaining node teardown.
2. Local Marker Replay (Ateom): Ateom writes an atomic completion marker after snapshotting. If a retried request arrives after the guest is down, Ateom replays the recorded snapshot files rather than failing against the destroyed sandbox.
3. Failure Classification: Distinguishes transient errors from unrecoverable guest crashes so the control plane can safely decide whether to retry.

AI has assisted with this PR and I have verified all the changes.

• Tests pass
• Appropriate changes to documentation are included in the PR

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[Troy Chiu (troychiu)]

Move checkpoint related functions to a separated file. Here is the diff from the original file for easier review

--- cmd/ateom-gvisor/main.go (original)
+++ cmd/ateom-gvisor/checkpoint.go (new)
@@ -7,11 +7,34 @@
 	s.setActiveRPC(rpcCheckpointWorkload, cancel)
 	defer s.clearActiveRPC()
 
+	attribution := ateomstats.ActorAttributionFromRequest(req)
+
+	// Replay a previously completed checkpoint for this actor if available.
+	if rec, ok, err := checkpointmarker.Read(req.GetActorUid(), req.GetScope().String()); err != nil {
+		return nil, err
+	} else if ok {
+		slog.InfoContext(ctx, "Checkpoint already completed for this actor; replaying its result",
+			"actor", attribution.Ref,
+			"actorUID", req.GetActorUid(),
+			"snapshotFiles", rec.SnapshotFiles)
+		// Finish any pending workload termination, unless the ateom now holds a different actor.
+		if held := s.activeActor.Load(); held != nil && held.UID != req.GetActorUid() {
+			slog.WarnContext(ctx, "Not running the post-checkpoint teardown: this ateom now holds a different actor",
+				slog.String("id", req.GetActorUid()), slog.String("active_actor_uid", held.UID))
+			return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil
+		}
+		if err := s.terminateWorkload(ctx, attribution.Ref, req.GetActorUid(), req.GetRunscPath(), req.GetSpec().GetContainers()); err != nil {
+			slog.WarnContext(ctx, "Failed to terminate workload while replaying checkpoint",
+				slog.String("actorUID", req.GetActorUid()), slog.Any("err", err))
+		}
+		s.activeSession = nil
+		return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil
+	}
+
 	if err := s.deactivateActorNetworking(ctx); err != nil {
 		return nil, err
 	}
 
-	attribution := ateomstats.ActorAttributionFromRequest(req)
 	s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointing", attribution)
 
 	// Contract with atelet:
@@ -26,6 +49,11 @@
 	}
 
 	checkpointPath := ateompath.CheckpointStateDir(req.GetActorUid())
+	// Start from a clean directory so retried attempts do not mix with stale
+	// or partially-written snapshot files from previous runs.
+	if err := os.RemoveAll(checkpointPath); err != nil {
+		return nil, fmt.Errorf("while clearing checkpoint directory: %w", err)
+	}
 	if err := os.MkdirAll(checkpointPath, 0o700); err != nil {
 		return nil, fmt.Errorf("while creating checkpoint directory: %w", err)
 	}
@@ -44,12 +72,12 @@
 			return nil, fmt.Errorf("no durable-dir volumes found for DATA snapshot")
 		}
 		if err := rcmd.cmdFsCheckpoint(ctx, "pause", checkpointPath, ddv); err != nil {
-			return nil, fmt.Errorf("while fscheckpointing durable-dir %q: %w", ddv[0], err)
+			return nil, classifyCheckpointFailure(ctx, rcmd, fmt.Errorf("while fscheckpointing durable-dir %q: %w", ddv[0], err))
 		}
 	case ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL:
 		// Checkpoint pause container (root of the sandbox)
 		if err := rcmd.cmdCheckpoint(ctx, "pause", checkpointPath); err != nil {
-			return nil, fmt.Errorf("while checkpointing pause: %w", err)
+			return nil, classifyCheckpointFailure(ctx, rcmd, fmt.Errorf("while checkpointing pause: %w", err))
 		}
 	default:
 		return nil, fmt.Errorf("unsupported snapshot scope: %v", req.GetScope())
@@ -68,6 +96,23 @@
 	// reporting its usage is then the honest answer.
 	s.activeActor.Store(nil)
 
+	// Report exactly the files runsc wrote so atelet ships precisely this set
+	// (checkpoint.img plus any pages images), rather than a hardcoded list.
+	snapshotFiles, err := listSnapshotFiles(checkpointPath)
+	if err != nil {
+		return nil, fmt.Errorf("while listing checkpoint files: %w", err)
+	}
+
+	// Record checkpoint completion before answering. If writing the marker fails,
+	// log and continue since the snapshot files are already complete on disk.
+	if err := checkpointmarker.Write(req.GetActorUid(), req.GetScope().String(), snapshotFiles); err != nil {
+		slog.ErrorContext(ctx, "Failed to record the checkpoint completion marker; answering anyway, but a lost response can no longer be replayed",
+			"actor", attribution.Ref,
+			"actorUID", req.GetActorUid(),
+			"snapshotFiles", snapshotFiles,
+			"err", err)
+	}
+
 	// Cleanup the containers after checkpointing.
 	// This is best-effort cleanup for actor containers that may have been left behind after checkpointing.
 	if err := s.terminateWorkload(ctx, attribution.Ref, attribution.UID, req.GetRunscPath(), req.GetSpec().GetContainers()); err != nil {
@@ -77,21 +122,54 @@
 			slog.Any("err", err))
 	}
 
-	// Report exactly the files runsc wrote so atelet ships precisely this set
-	// (checkpoint.img plus any pages images), rather than a hardcoded list.
-	snapshotFiles, err := listSnapshotFiles(checkpointPath)
-	if err != nil {
-		return nil, fmt.Errorf("while listing checkpoint files: %w", err)
-	}
-
 	s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointed", attribution)
 	s.activeSession = nil
 
 	return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles}, nil
 }
 
-// listSnapshotFiles returns the (relative) names of regular files directly under
-// dir, which atelet ships to object storage as the snapshot.
+// stateProbeTimeout bounds probing runsc state during failure classification.
+const stateProbeTimeout = 15 * time.Second
+
+// classifyCheckpointFailure inspects runsc container state after a failure to
+// distinguish retriable transient errors from unrecoverable errors (where the
+// sandbox was destroyed and cannot be retried).
+func classifyCheckpointFailure(ctx context.Context, rcmd *runsc, err error) error {
+	// Probe with a separate timeout so an expired caller context is not
+	// mistaken for a missing sandbox.
+	probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), stateProbeTimeout)
+	defer cancel()
+
+	out, stateErr := rcmd.cmdStateOutput(probeCtx, "pause")
+	if stateErr == nil {
+		return err
+	}
+	if !sandboxNotFound(out) {
+		slog.WarnContext(ctx, "Checkpoint failed and the sandbox state could not be determined; leaving the failure retriable",
+			"actorUID", rcmd.actorUID, "stateErr", stateErr, "runscOutput", string(out), "err", err)
+		return err
+	}
+	slog.WarnContext(ctx, "Checkpoint failed and the sandbox is gone; the actor's state is unrecoverable",
+		"actorUID", rcmd.actorUID, "stateErr", stateErr, "err", err)
+	return ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonInvalidCheckpointResult, ateerrors.ActorCrashedMetadata(),
+		fmt.Errorf("%w: checkpoint failed and no sandbox remains to retry against: %w", ateerrors.ReasonInvalidCheckpointResult, err))
+}
+
+// sandboxNotFound checks if runsc output explicitly indicates the container does not exist.
+func sandboxNotFound(runscOutput []byte) bool {
+	for line := range strings.Lines(string(runscOutput)) {
+		msg, ok := strings.CutPrefix(strings.TrimSpace(line), "error:")
+		if !ok {
+			continue
+		}
+		if strings.Contains(strings.ToLower(msg), "does not exist") {
+			return true
+		}
+	}
+	return false
+}
+
+// listSnapshotFiles returns the (relative) names of regular files directly under dir.
 func listSnapshotFiles(dir string) ([]string, error) {
 	entries, err := os.ReadDir(dir)
 	if err != nil {
@@ -99,11 +177,12 @@
 	}
 	var files []string
 	for _, e := range entries {
-		if e.Type().IsRegular() {
+		// ateom's own completion marker shares the directory but is
+		// bookkeeping, not snapshot content, so it never joins the set.
+		if e.Type().IsRegular() && e.Name() != ateompath.CheckpointDoneFileName {
 			files = append(files, e.Name())
 		}
 	}
 	sort.Strings(files)
 	return files, nil
 }