atelet: stage local checkpoints by hard link instead of copying

cmd/atelet/main.go

@@ -577,6 +577,20 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe
 		return nil, err
 	}
 
+	// Drop earlier pause snapshots BEFORE checkpointing, not after. No earlier one
+	// can be restored again: the control plane tracks a single local snapshot, which
+	// this checkpoint either overwrites (pause) or clears (suspend). Ordering it here
+	// matters because restore staged this actor's memory image by linking to that
+	// snapshot, and MergeDeltaIntoBase can only take its cheap in-place path while
+	// the staged image is the sole name for those bytes. Releasing the other name
+	// first leaves the image untouched (the staged link still holds the inode) and
+	// keeps the merge off the copying path.
+	//
+	// The cost is that a checkpoint failing from here on leaves no earlier snapshot
+	// to fall back to. That is survivable: the guest stays paused when
+	// CheckpointWorkload fails, so the checkpoint can simply be retried.
+	pruneLocalCheckpoints(ctx, actorUID)
+
 	// Tell ateom to take the checkpoint and delete containers. ateom reports the
 	// exact files it wrote so we ship precisely that set (gVisor's image files,
 	// cloud-hypervisor's snapshot set, ...) rather than a hardcoded list.
@@ -611,15 +625,6 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe
 	sandboxRec.ActorTemplateName = req.GetActorTemplateName()
 	sandboxRec.Scope = ateattr.SnapshotScopeValue(req.GetScope())
 
-	// No earlier pause snapshot can ever be restored again, so remove them
-	// all: the actor's current state was just captured by CheckpointWorkload,
-	// and the control plane tracks only a single local snapshot, which this
-	// checkpoint either overwrites (pause) or clears (suspend).
-	pruneLocalCheckpoints(ctx, actorUID)
-
-	// Pruning stays outside the persist window: it collects superseded
-	// snapshots on both paths, so timing it as part of an external upload would
-	// mix local disk deletion into the object-storage measurement.
 	tPersist := time.Now()
 	switch req.GetType() {
 	case ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL:
@@ -1156,6 +1161,26 @@ func (s *AteomHerder) copyLocalCheckpoint(ctx context.Context, snapshotName stri
 		}
 		src := filepath.Join(srcDir, snapshotName, fileName)
 		dst := filepath.Join(dstDir, fileName)
+		// Link rather than copy. The local checkpoint lives under the same actor dir
+		// as the restore staging area, so this stages the memory image in constant
+		// time instead of re-writing its whole working set. Nothing rewrites the
+		// shared inode: CH demand-pages from the staged image read-only,
+		// rewriteSnapshotSocketPaths renames its rewritten config.json into place
+		// rather than truncating, and MergeDeltaIntoBase refuses its in-place overlay
+		// once the image carries a second link.
+		//
+		// EXDEV alone falls back to copying, so an unexpected link failure surfaces
+		// instead of silently reverting to the full copy this exists to remove. It
+		// also keeps copyFile off a dst that is already a link to src, where its
+		// O_TRUNC would empty both and report a successful copy of the old size.
+		switch err := linkFile(src, dst); {
+		case err == nil:
+			continue
+		case !errors.Is(err, unix.EXDEV):
+			return fmt.Errorf("failed to link %s to %s: %w", src, dst, err)
+		}
+		slog.WarnContext(ctx, "local checkpoint and restore dir are on different filesystems; copying instead of linking",
+			slog.String("src", src), slog.String("dst", dst))
 		if _, err := copyFile(src, dst); err != nil {
 			return fmt.Errorf("failed to copy %s to %s: %w", src, dst, err)
 		}
@@ -1166,6 +1191,10 @@ func (s *AteomHerder) copyLocalCheckpoint(ctx context.Context, snapshotName stri
 
 var createDestFile = func(name string) (io.WriteCloser, error) { return os.Create(name) }
 
+// linkFile is os.Link, indirected so a test can force the cross-filesystem
+// fallback in copyLocalCheckpoint without mounting a second filesystem.
+var linkFile = os.Link
+
 // sparseDest is the part of *os.File a hole-preserving copy needs. Destinations that
 // do not implement it are copied densely instead.
 type sparseDest interface {

cmd/atelet/main_test.go

@@ -44,6 +44,7 @@ import (
 	"github.com/google/go-cmp/cmp"
 	"github.com/klauspost/compress/zstd"
 	"github.com/spf13/pflag"
+	"golang.org/x/sys/unix"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
@@ -197,6 +198,94 @@ func TestCopyFile(t *testing.T) {
 	}
 }
 
+// TestCopyLocalCheckpointLinks covers staging a local checkpoint into the restore
+// dir: the files must land as extra links to the cached snapshot rather than
+// copies, so a resume does not rewrite the image's working set. Sharing the inode
+// is what MergeDeltaIntoBase's Nlink check keys off to refuse its in-place overlay.
+func TestCopyLocalCheckpointLinks(t *testing.T) {
+	const snapshot = "snap-1"
+	want := []byte("checkpoint pages")
+
+	newDirs := func(t *testing.T) (srcDir, dstDir string) {
+		t.Helper()
+		root := t.TempDir()
+		srcDir, dstDir = filepath.Join(root, "local-checkpoint"), filepath.Join(root, "restore-state")
+		if err := os.MkdirAll(filepath.Join(srcDir, snapshot), 0o700); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.MkdirAll(dstDir, 0o700); err != nil {
+			t.Fatal(err)
+		}
+		if err := os.WriteFile(filepath.Join(srcDir, snapshot, "memory-ranges"), want, 0o600); err != nil {
+			t.Fatal(err)
+		}
+		return srcDir, dstDir
+	}
+	inode := func(t *testing.T, path string) uint64 {
+		t.Helper()
+		fi, err := os.Stat(path)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return fi.Sys().(*syscall.Stat_t).Ino
+	}
+
+	t.Run("links when it can", func(t *testing.T) {
+		srcDir, dstDir := newDirs(t)
+		s := &AteomHerder{}
+		if err := s.copyLocalCheckpoint(context.Background(), snapshot, srcDir, dstDir, []string{"memory-ranges"}); err != nil {
+			t.Fatalf("copyLocalCheckpoint: %v", err)
+		}
+		src := filepath.Join(srcDir, snapshot, "memory-ranges")
+		dst := filepath.Join(dstDir, "memory-ranges")
+		if got, err := os.ReadFile(dst); err != nil || !bytes.Equal(got, want) {
+			t.Fatalf("dst content = %q (err %v), want %q", got, err, want)
+		}
+		if inode(t, src) != inode(t, dst) {
+			t.Error("staged file is a copy; expected a link to the cached snapshot")
+		}
+	})
+
+	t.Run("falls back to copying across filesystems", func(t *testing.T) {
+		srcDir, dstDir := newDirs(t)
+		// EXDEV stands in for the mount boundary a unit test cannot produce.
+		orig := linkFile
+		linkFile = func(string, string) error { return unix.EXDEV }
+		t.Cleanup(func() { linkFile = orig })
+
+		s := &AteomHerder{}
+		if err := s.copyLocalCheckpoint(context.Background(), snapshot, srcDir, dstDir, []string{"memory-ranges"}); err != nil {
+			t.Fatalf("copyLocalCheckpoint: %v", err)
+		}
+		dst := filepath.Join(dstDir, "memory-ranges")
+		if got, err := os.ReadFile(dst); err != nil || !bytes.Equal(got, want) {
+			t.Fatalf("dst content = %q (err %v), want %q", got, err, want)
+		}
+		if inode(t, filepath.Join(srcDir, snapshot, "memory-ranges")) == inode(t, dst) {
+			t.Error("expected a copy on the fallback path, got a link")
+		}
+	})
+
+	// Only EXDEV may fall back. Copying on any other link failure would silently
+	// undo this optimization, and would hand copyFile a dst that may already be a
+	// link to src, where its O_TRUNC empties both before the copy reads a byte.
+	t.Run("other link failures are fatal", func(t *testing.T) {
+		srcDir, dstDir := newDirs(t)
+		dst := filepath.Join(dstDir, "memory-ranges")
+		if err := os.Link(filepath.Join(srcDir, snapshot, "memory-ranges"), dst); err != nil {
+			t.Fatal(err)
+		}
+		s := &AteomHerder{}
+		// dst already exists, so os.Link fails with EEXIST.
+		if err := s.copyLocalCheckpoint(context.Background(), snapshot, srcDir, dstDir, []string{"memory-ranges"}); err == nil {
+			t.Fatal("copyLocalCheckpoint accepted a non-EXDEV link failure, want an error")
+		}
+		if got, err := os.ReadFile(dst); err != nil || !bytes.Equal(got, want) {
+			t.Fatalf("dst content = %q (err %v), want the untouched %q", got, err, want)
+		}
+	})
+}
+
 type failingCloseFile struct{ *os.File }
 
 func (f failingCloseFile) Close() error {

cmd/ateom-microvm/internal/ch/merge.go

@@ -23,6 +23,7 @@ import (
 	"io"
 	"os"
 	"os/exec"
+	"syscall"
 
 	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/reaper"
 	"golang.org/x/sys/unix"
@@ -112,7 +113,8 @@ func MergeSparseOverlay(ctx context.Context, baseFile, deltaFile, outFile string
 // baseFile and deltaFile are siblings under the actor dir (restore-state/ and
 // checkpoint-state/), so the renames are same-filesystem (metadata-only). If they
 // straddle a mount boundary (EXDEV) it falls back to the copying MergeSparseOverlay
-// (baseFile is untouched until the first rename succeeds).
+// (baseFile is untouched until the first rename succeeds), as it does when baseFile
+// carries a second link and so cannot be overlaid in place.
 func MergeDeltaIntoBase(ctx context.Context, baseFile, deltaFile string) error {
 	bi, err := os.Stat(baseFile)
 	if err != nil {
@@ -128,6 +130,15 @@ func MergeDeltaIntoBase(ctx context.Context, baseFile, deltaFile string) error {
 		return fmt.Errorf("MergeDeltaIntoBase: size mismatch base=%d delta=%d", bi.Size(), di.Size())
 	}
 
+	// The fast path below MUTATES baseFile's inode in place, which is only safe while
+	// baseFile is the sole name for it. atelet stages restore-state by linking from
+	// this actor's local pause snapshot rather than copying it, so a second link means
+	// the overlay would rewrite that cached snapshot too, corrupting the actor's only
+	// local restore point. Copy in that case, which is what MergeSparseOverlay does.
+	if st, ok := bi.Sys().(*syscall.Stat_t); ok && st.Nlink > 1 {
+		return MergeSparseOverlay(ctx, baseFile, deltaFile, deltaFile)
+	}
+
 	// Move baseFile (with its already-on-disk working set) next to deltaFile. If this
 	// fails with EXDEV the two are on different filesystems and baseFile is still
 	// intact, so fall back to the copying merge.

cmd/ateom-microvm/internal/ch/merge_test.go

@@ -183,6 +183,54 @@ func TestMergeSparseOverlayNewOutFile(t *testing.T) {
 	}
 }
 
+// TestMergeDeltaIntoBaseHardlinkedBase asserts the fast path is refused when base
+// has a second link. The fast path overlays base's inode in place, so a shared
+// inode (atelet stages restore-state by linking from the actor's local pause
+// snapshot) would have that snapshot silently rewritten behind the merge.
+func TestMergeDeltaIntoBaseHardlinkedBase(t *testing.T) {
+	const size = 8 << 20 // 8 MiB logical
+	baseRegions := []region{{off: 0, data: fill(1, 4096)}}
+	deltaRegions := []region{{off: 4 << 20, data: fill(42, 12345)}}
+	want := make([]byte, size)
+	copy(want[baseRegions[0].off:], baseRegions[0].data)
+	copy(want[deltaRegions[0].off:], deltaRegions[0].data)
+	// What the shared inode must still hold afterwards: base, with no delta in it.
+	baseOnly := make([]byte, size)
+	copy(baseOnly[baseRegions[0].off:], baseRegions[0].data)
+
+	dir := t.TempDir()
+	base := filepath.Join(dir, "base")
+	delta := filepath.Join(dir, "delta")
+	cached := filepath.Join(dir, "cached-snapshot")
+	writeSparse(t, base, size, baseRegions)
+	writeSparse(t, delta, size, deltaRegions)
+	if err := os.Link(base, cached); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := MergeDeltaIntoBase(context.Background(), base, delta); err != nil {
+		t.Fatalf("MergeDeltaIntoBase: %v", err)
+	}
+	got, err := os.ReadFile(delta)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(got, want) {
+		t.Fatalf("merged result != expected (len got=%d want=%d)", len(got), len(want))
+	}
+	shared, err := os.ReadFile(cached)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(shared, baseOnly) {
+		t.Error("the hardlinked base was mutated; the cached snapshot behind it is now corrupt")
+	}
+	// The copying merge leaves base alone, unlike the fast path which consumes it.
+	if _, err := os.Stat(base); err != nil {
+		t.Errorf("base should survive a copying merge: %v", err)
+	}
+}
+
 // TestMergeDeltaIntoBaseSizeMismatch verifies a base/delta size mismatch is
 // refused (misaligned overlay would corrupt the image) rather than silently
 // producing garbage.

cmd/ateom-microvm/restore.go

@@ -450,8 +450,23 @@ func rewriteSnapshotSocketPaths(snapshotDir, id string) error {
 	if err != nil {
 		return err
 	}
-	if err := os.WriteFile(cfgPath, out, 0o600); err != nil {
+	// Temp file + rename, not os.WriteFile: atelet stages a local checkpoint by
+	// hard-linking it into this dir, so config.json can share an inode with the
+	// actor's cached pause snapshot. O_TRUNC would write straight through that link
+	// and rewrite the snapshot, and a crash mid-write would leave the actor's only
+	// local restore point holding a truncated config. Renaming replaces the name
+	// here and leaves the linked original whole.
+	tmp, err := os.CreateTemp(snapshotDir, ".config.json.tmp-*")
+	if err != nil {
 		return err
 	}
-	return nil
+	defer os.Remove(tmp.Name()) // no-op once the rename succeeds
+	if _, err := tmp.Write(out); err != nil {
+		tmp.Close()
+		return err
+	}
+	if err := tmp.Close(); err != nil {
+		return err
+	}
+	return os.Rename(tmp.Name(), cfgPath)
 }

cmd/ateom-microvm/restore_test.go

@@ -17,6 +17,7 @@
 package main
 
 import (
+	"bytes"
 	"encoding/json"
 	"os"
 	"path/filepath"
@@ -143,4 +144,46 @@ func TestRewriteSnapshotSocketPaths(t *testing.T) {
 			t.Errorf("serial file = %q, want %q", cfg.Serial.File, want)
 		}
 	})
+
+	// atelet stages a local checkpoint into the restore dir by hard-linking it, so
+	// the config.json rewritten here can share an inode with the actor's cached
+	// pause snapshot. Writing in place would rewrite that snapshot's config too.
+	t.Run("a hardlinked config is not written through", func(t *testing.T) {
+		dir := writeSnapshotConfig(t, []map[string]any{
+			{"tag": kata.FsTag, "socket": "/run/vc/vm/golden/virtiofsd.sock"},
+		})
+		cfgPath := filepath.Join(dir, "config.json")
+		before, err := os.ReadFile(cfgPath)
+		if err != nil {
+			t.Fatalf("reading config.json: %v", err)
+		}
+		cached := filepath.Join(t.TempDir(), "config.json")
+		if err := os.Link(cfgPath, cached); err != nil {
+			t.Fatalf("linking config.json: %v", err)
+		}
+
+		if err := rewriteSnapshotSocketPaths(dir, id); err != nil {
+			t.Fatalf("rewriteSnapshotSocketPaths: %v", err)
+		}
+		if got, want := readFsSockets(t, dir)[kata.FsTag], kata.VirtiofsdSocketPath(id); got != want {
+			t.Errorf("%s socket = %q, want %q", kata.FsTag, got, want)
+		}
+		got, err := os.ReadFile(cached)
+		if err != nil {
+			t.Fatalf("reading the linked config.json: %v", err)
+		}
+		if !bytes.Equal(got, before) {
+			t.Errorf("the linked config.json was rewritten: got %q, want the original %q", got, before)
+		}
+		// Nothing may be left behind for atelet to ship as part of the snapshot.
+		entries, err := os.ReadDir(dir)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for _, e := range entries {
+			if e.Name() != "config.json" {
+				t.Errorf("stray file left in the restore dir: %q", e.Name())
+			}
+		}
+	})
 }