ateom: reuse the checkpoint image left by a cancelled attempt

A retried CheckpointWorkload ran `runsc checkpoint` against a sandbox the
first attempt had already checkpointed and stopped, got exit 128 and
crashed the actor. The first attempt had been cancelled by its caller
during the teardown, which ran under the RPC context.

Reuse the image when it exists and the pause container no longer runs,
and run the teardown after a checkpoint on a detached, bounded context.
cmdState now returns the container's status.

Author: dims

cmd/ateom-gvisor/cleanup_test.go

@@ -41,12 +41,12 @@ func newFakeRunsc(containers ...string) *fakeRunsc {
 	return f
 }
 
-func (f *fakeRunsc) cmdState(_ context.Context, name string) error {
+func (f *fakeRunsc) cmdState(_ context.Context, name string) (string, error) {
 	f.calls = append(f.calls, "state "+name)
 	if !f.containers[name] {
-		return errors.New("exit status 128")
+		return "", errors.New("exit status 128")
 	}
-	return nil
+	return "running", nil
 }
 
 func (f *fakeRunsc) cmdDelete(_ context.Context, name string) error {

cmd/ateom-gvisor/main.go

@@ -92,6 +92,9 @@ const minUsageSampleInterval = 50 * time.Second
 // kubelet SIGKILL of ateom itself.
 const workloadGracePeriod = 30 * time.Minute
 
+// checkpointTeardownTimeout bounds the teardown after a checkpoint, which runs detached from the caller.
+const checkpointTeardownTimeout = 60 * time.Second
+
 func main() {
 	pflag.Parse()
 	if *showVersion {
@@ -444,7 +447,7 @@ var containerKillTimeout = 5 * time.Second
 type containerRuntime interface {
 	cmdKill(ctx context.Context, containerName, signal string) error
 	cmdWait(ctx context.Context, containerName string) error
-	cmdState(ctx context.Context, containerName string) error
+	cmdState(ctx context.Context, containerName string) (string, error)
 	cmdDelete(ctx context.Context, containerName string) error
 	cmdList(ctx context.Context) ([]string, error)
 }
@@ -766,12 +769,28 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 			return nil, fmt.Errorf("while checkpointing durable-dir volumes: %w", err)
 		}
 	case ateompb.SnapshotFidelity_SNAPSHOT_FIDELITY_MEMORY:
-		// Checkpoint pause container (root of the sandbox), splitting durable-dir
-		// volumes into <checkpointPath>/fs when present.
-		err := rcmd.cmdCheckpoint(ctx, ocispec.PauseContainer, checkpointPath, durablePaths)
-		timing.checkpoint = lap(&tLast)
-		if err != nil {
-			return nil, fmt.Errorf("while checkpointing pause: %w", err)
+		// A retry after a cancelled teardown finds the image already written and
+		// the sandbox gone; checkpointing it again fails and crashes the actor.
+		pauseStatus := func(ctx context.Context) (string, error) { return rcmd.cmdState(ctx, ocispec.PauseContainer) }
+		if previousCheckpointReusable(ctx, pauseStatus, checkpointPath) {
+			slog.InfoContext(ctx, "Reusing the checkpoint image of a previous attempt",
+				slog.String("actor", attribution.Ref.String()),
+				slog.String("actorUID", attribution.UID))
+			timing.checkpoint = lap(&tLast)
+		} else {
+			// runsc opens its image files with O_EXCL; drop a partial earlier attempt.
+			if stale, err := listSnapshotFiles(checkpointPath); err == nil {
+				for _, f := range stale {
+					_ = os.Remove(filepath.Join(checkpointPath, f))
+				}
+			}
+			// Checkpoint pause container (root of the sandbox), splitting durable-dir
+			// volumes into <checkpointPath>/fs when present.
+			err := rcmd.cmdCheckpoint(ctx, ocispec.PauseContainer, checkpointPath, durablePaths)
+			timing.checkpoint = lap(&tLast)
+			if err != nil {
+				return nil, fmt.Errorf("while checkpointing pause: %w", err)
+			}
 		}
 	default:
 		return nil, fmt.Errorf("unsupported snapshot fidelity: %v", fidelity)
@@ -780,14 +799,17 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 	// Cleanup the containers after checkpointing. This also unhosts the actor,
 	// before the snapshot listing below can fail.
 	// 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.GetActorDirs(), req.GetSpec().GetContainers()); err != nil {
+	// Detached from the caller: a cancellation here leaves a half torn down sandbox.
+	teardownCtx, cancelTeardown := context.WithTimeout(context.WithoutCancel(ctx), checkpointTeardownTimeout)
+	defer cancelTeardown()
+	if err := s.terminateWorkload(teardownCtx, attribution.Ref, attribution.UID, req.GetRunscPath(), req.GetActorDirs(), req.GetSpec().GetContainers()); err != nil {
 		slog.WarnContext(ctx, "failed to terminate workload after checkpoint",
 			slog.String("actor", attribution.Ref.String()),
 			slog.String("actorUID", attribution.UID),
 			slog.Any("err", err))
 	}
 	timing.teardown = lap(&tLast)
-	s.recordFinalIfEnded(ctx, hosted)
+	s.recordFinalIfEnded(teardownCtx, hosted)
 
 	// Report exactly the files runsc wrote so atelet ships precisely this set
 	// (checkpoint.img plus any pages images and fs/ durable-dir files), rather
@@ -803,6 +825,21 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 	return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles, DataSnapshotFiles: durableFiles}, nil
 }
 
+// previousCheckpointReusable reports whether checkpointPath holds the image of
+// an earlier attempt whose sandbox no longer runs. status is the pause
+// container's OCI status, or an error when runsc has no record of it.
+func previousCheckpointReusable(ctx context.Context, status func(context.Context) (string, error), checkpointPath string) bool {
+	info, err := os.Stat(filepath.Join(checkpointPath, "checkpoint.img"))
+	if err != nil || info.Size() == 0 {
+		return false
+	}
+	st, err := status(ctx)
+	if err != nil {
+		return true
+	}
+	return st != "running"
+}
+
 // listSnapshotFiles returns the (relative) names of regular files directly under
 // dir and under dir/fs, which atelet ships to object storage as the snapshot.
 func listSnapshotFiles(dir string) ([]string, error) {
@@ -852,7 +889,7 @@ func cleanupContainers(ctx context.Context, rcmd containerRuntime, containers []
 	// Without this, `runsc delete` occasionally throws an error.
 	present := make([]string, 0, len(names))
 	for _, name := range names {
-		if err := rcmd.cmdState(ctx, name); err != nil {
+		if _, err := rcmd.cmdState(ctx, name); err != nil {
 			err = fmt.Errorf("while checking state of %q container: %w", name, err)
 			gone, listErr := isContainerAlreadyGone(ctx, rcmd, name)
 			if listErr != nil {

cmd/ateom-gvisor/main_test.go

@@ -18,6 +18,9 @@ package main
 
 import (
 	"context"
+	"errors"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"google.golang.org/grpc/codes"
@@ -89,3 +92,35 @@ func TestRPCsRejectUntrustedRunscPath(t *testing.T) {
 		}
 	}
 }
+
+func TestPreviousCheckpointReusable(t *testing.T) {
+	running := func(context.Context) (string, error) { return "running", nil }
+	stopped := func(context.Context) (string, error) { return "stopped", nil }
+	gone := func(context.Context) (string, error) { return "", errors.New("container does not exist") }
+	withImage := func(size int) string {
+		dir := t.TempDir()
+		if err := os.WriteFile(filepath.Join(dir, "checkpoint.img"), make([]byte, size), 0o600); err != nil {
+			t.Fatal(err)
+		}
+		return dir
+	}
+	tests := []struct {
+		name   string
+		dir    string
+		status func(context.Context) (string, error)
+		want   bool
+	}{
+		{"no image", t.TempDir(), gone, false},
+		{"empty image", withImage(0), gone, false},
+		{"image, container gone", withImage(1), gone, true},
+		{"image, container stopped", withImage(1), stopped, true},
+		{"image, container running", withImage(1), running, false},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := previousCheckpointReusable(context.Background(), tt.status, tt.dir); got != tt.want {
+				t.Errorf("previousCheckpointReusable() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

cmd/ateom-gvisor/runsc.go

@@ -19,6 +19,7 @@ package main
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -302,7 +303,9 @@ func (r *runsc) cmdDelete(ctx context.Context, containerName string) error {
 	return nil
 }
 
-func (r *runsc) cmdState(ctx context.Context, containerName string) error {
+// cmdState returns the container's OCI status ("running", "stopped", ...),
+// or an error when runsc has no record of it.
+func (r *runsc) cmdState(ctx context.Context, containerName string) (string, error) {
 	cmd := exec.CommandContext(
 		ctx,
 		r.path,
@@ -312,12 +315,19 @@ func (r *runsc) cmdState(ctx context.Context, containerName string) error {
 		"state",
 		containerName,
 	)
-	cmd.Stdout = os.Stdout
+	var out bytes.Buffer
+	cmd.Stdout = &out
 	cmd.Stderr = os.Stderr
 	if err := reaper.RunCommand(cmd); err != nil {
-		return fmt.Errorf("while running `runsc state`: %w", err)
+		return "", fmt.Errorf("while running `runsc state`: %w", err)
 	}
-	return nil
+	var state struct {
+		Status string `json:"status"`
+	}
+	if err := json.Unmarshal(out.Bytes(), &state); err != nil {
+		return "", fmt.Errorf("while parsing `runsc state` output: %w", err)
+	}
+	return state.Status, nil
 }
 
 // cmdList returns the container IDs runsc has a record of.