atelet: preserve holes when copying a local checkpoint

cmd/atelet/copyrange_linux.go

@@ -0,0 +1,65 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"errors"
+
+	"golang.org/x/sys/unix"
+)
+
+// maxKernelCopy caps a single copy_file_range request. The syscall takes an int
+// length, and a bounded request keeps one call from monopolising the thread.
+const maxKernelCopy = 1 << 30
+
+// kernelCopyRange copies up to length bytes at off from srcFd to dstFd without the
+// data crossing into userspace, and reports how much it copied (short copies are
+// normal, so callers must loop).
+//
+// It reports errKernelCopyUnsupported when the kernel or filesystem cannot do the
+// copy — most commonly EXDEV, when source and destination are on different
+// filesystems — so the caller can fall back to a userspace copy.
+func kernelCopyRange(srcFd, dstFd int, off, length int64) (int64, error) {
+	if length > maxKernelCopy {
+		length = maxKernelCopy
+	}
+	for {
+		// The offsets are only advanced by bytes the call actually copied, so a
+		// retry after EINTR resumes from the right place.
+		roff, woff := off, off
+		n, err := unix.CopyFileRange(srcFd, &roff, dstFd, &woff, int(length), 0)
+		switch {
+		case err == nil:
+			if n == 0 {
+				// No error and no progress: treat as unsupported rather than spin.
+				return 0, errKernelCopyUnsupported
+			}
+			return int64(n), nil
+		case errors.Is(err, unix.EINTR):
+			// A copy this large is interruptible, and the Go runtime signals
+			// goroutines for preemption, so this is expected rather than fatal.
+			continue
+		case errors.Is(err, unix.ENOSYS), // pre-4.5 kernel, or blocked by seccomp
+			errors.Is(err, unix.EXDEV),      // different filesystems
+			errors.Is(err, unix.EOPNOTSUPP), // filesystem does not implement it
+			errors.Is(err, unix.EPERM),      // e.g. append-only destination
+			errors.Is(err, unix.EINVAL),     // ranges or flags this kernel rejects
+			errors.Is(err, unix.EBADF):      // not both regular files
+			return 0, errKernelCopyUnsupported
+		default:
+			return 0, err
+		}
+	}
+}

cmd/atelet/copyrange_linux_test.go

@@ -0,0 +1,99 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build linux
+
+package main
+
+import (
+	"bytes"
+	"errors"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+// TestKernelCopyRange exercises the copy_file_range path directly. copySparse falls
+// back to userspace when the kernel refuses, so without this the syscall could be
+// broken and every other test would still pass.
+func TestKernelCopyRange(t *testing.T) {
+	const (
+		size = 1 << 20
+		off  = 64 << 10
+	)
+	dir := t.TempDir()
+
+	payload := bytes.Repeat([]byte{0x5A}, 128<<10)
+	srcPath := filepath.Join(dir, "src")
+	src, err := os.Create(srcPath)
+	if err != nil {
+		t.Fatalf("creating src: %v", err)
+	}
+	defer src.Close()
+	if err := src.Truncate(size); err != nil {
+		t.Fatalf("sizing src: %v", err)
+	}
+	if _, err := src.WriteAt(payload, off); err != nil {
+		t.Fatalf("writing src: %v", err)
+	}
+	if err := src.Sync(); err != nil {
+		t.Fatalf("syncing src: %v", err)
+	}
+
+	dstPath := filepath.Join(dir, "dst")
+	dst, err := os.Create(dstPath)
+	if err != nil {
+		t.Fatalf("creating dst: %v", err)
+	}
+	defer dst.Close()
+	if err := dst.Truncate(size); err != nil {
+		t.Fatalf("sizing dst: %v", err)
+	}
+
+	// Short copies are legal, so loop like copySparse does.
+	remaining := int64(len(payload))
+	pos := int64(off)
+	for remaining > 0 {
+		n, err := kernelCopyRange(int(src.Fd()), int(dst.Fd()), pos, remaining)
+		if errors.Is(err, errKernelCopyUnsupported) {
+			t.Skipf("copy_file_range unavailable on this filesystem after %d of %d bytes",
+				int64(len(payload))-remaining, int64(len(payload)))
+		}
+		if err != nil {
+			t.Fatalf("kernelCopyRange at %d: %v", pos, err)
+		}
+		if n <= 0 {
+			t.Fatalf("kernelCopyRange reported %d bytes copied", n)
+		}
+		pos += n
+		remaining -= n
+	}
+
+	got := make([]byte, len(payload))
+	if _, err := dst.ReadAt(got, off); err != nil {
+		t.Fatalf("reading dst: %v", err)
+	}
+	if !bytes.Equal(got, payload) {
+		t.Error("copied range does not match the source")
+	}
+
+	// Everything outside the copied range must still be untouched.
+	head := make([]byte, off)
+	if _, err := dst.ReadAt(head, 0); err != nil {
+		t.Fatalf("reading dst head: %v", err)
+	}
+	if !bytes.Equal(head, make([]byte, off)) {
+		t.Error("copy wrote outside the requested range")
+	}
+}

cmd/atelet/copyrange_other.go

@@ -0,0 +1,24 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build !linux
+
+package main
+
+// kernelCopyRange has no implementation off Linux (atelet runs on Linux; this keeps
+// the package building for local development on other platforms), so callers copy
+// through userspace.
+func kernelCopyRange(_, _ int, _, _ int64) (int64, error) {
+	return 0, errKernelCopyUnsupported
+}

cmd/atelet/main.go

@@ -54,6 +54,7 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/metric"
 	"golang.org/x/sync/errgroup"
+	"golang.org/x/sys/unix"
 	"google.golang.org/api/option"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -765,6 +766,29 @@ func (s *AteomHerder) copyLocalCheckpoint(ctx context.Context, snapshotPrefix st
 
 var createDestFile = func(name string) (io.WriteCloser, error) { return os.Create(name) }
 
+// 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 {
+	Truncate(size int64) error
+	WriteAt(b []byte, off int64) (int, error)
+}
+
+// errSparseUnsupported means the source filesystem cannot report holes, so the caller
+// should fall back to a dense copy.
+var errSparseUnsupported = errors.New("filesystem cannot report holes")
+
+// errKernelCopyUnsupported means this platform, kernel or filesystem cannot copy a
+// range in the kernel, so the caller should copy through userspace instead.
+var errKernelCopyUnsupported = errors.New("kernel range copy unsupported")
+
+// copyFile copies src to dst, preserving holes where it can, and returns the number of
+// logical bytes copied.
+//
+// Preserving holes matters because the biggest thing copied here is a guest memory
+// image, which is mostly unallocated: a plain io.Copy reads holes as zeroes and writes
+// them as data, inflating a snapshot to its full logical size. That costs disk on every
+// local checkpoint restore, and it destroys the sparseness that later stages rely on to
+// tell which parts of guest RAM actually hold anything.
 func copyFile(src, dst string) (int64, error) {
 	sourceFileStat, err := os.Stat(src)
 	if err != nil {
@@ -785,10 +809,113 @@ func copyFile(src, dst string) (int64, error) {
 	if err != nil {
 		return 0, err
 	}
+
+	if sd, ok := destination.(sparseDest); ok {
+		switch err := copySparse(source, sd, sourceFileStat.Size()); {
+		case err == nil:
+			return sourceFileStat.Size(), destination.Close()
+		case !errors.Is(err, errSparseUnsupported):
+			return 0, errors.Join(err, destination.Close())
+		}
+		// Unsupported: nothing has been written yet, but probing moved the read
+		// offset, so rewind before the dense copy below.
+		if _, err := source.Seek(0, io.SeekStart); err != nil {
+			return 0, errors.Join(err, destination.Close())
+		}
+	}
+
 	nBytes, err := io.Copy(destination, source)
 	return nBytes, errors.Join(err, destination.Close())
 }
 
+// copySparse writes only src's populated extents to dst, located with SEEK_DATA and
+// SEEK_HOLE, leaving the rest of dst unallocated. It reports errSparseUnsupported
+// before writing anything if the filesystem cannot report holes.
+//
+// Extents are copied in the kernel where possible. The dense io.Copy this replaces got
+// that for free (os.File's ReadFrom uses copy_file_range), so without it a fully
+// populated file — a guest that really did touch all its RAM — would copy slower than
+// before.
+func copySparse(src *os.File, dst sparseDest, size int64) error {
+	fd := int(src.Fd())
+
+	// Probe first so an unsupported filesystem falls back with dst untouched. ENXIO
+	// means the seek ran but found no data at all, i.e. the file is one big hole.
+	if _, err := unix.Seek(fd, 0, unix.SEEK_DATA); err != nil {
+		if errors.Is(err, unix.ENXIO) {
+			return dst.Truncate(size)
+		}
+		return errSparseUnsupported
+	}
+	if err := dst.Truncate(size); err != nil {
+		return err
+	}
+
+	// A destination that exposes its descriptor can be written by the kernel; anything
+	// else (the test seam substitutes plain writers) goes through userspace.
+	dstFd := -1
+	if f, ok := dst.(interface{ Fd() uintptr }); ok {
+		dstFd = int(f.Fd())
+	}
+	var buf []byte
+
+	for off := int64(0); off < size; {
+		dataOff, err := unix.Seek(fd, off, unix.SEEK_DATA)
+		if err != nil {
+			if errors.Is(err, unix.ENXIO) {
+				break // no data past off; the tail is a hole
+			}
+			return fmt.Errorf("seeking to data at %d: %w", off, err)
+		}
+		if dataOff >= size {
+			break // data starts past the size we were asked to copy
+		}
+		holeOff, err := unix.Seek(fd, dataOff, unix.SEEK_HOLE)
+		if err != nil {
+			return fmt.Errorf("seeking to hole at %d: %w", dataOff, err)
+		}
+		// Refuse to spin: every iteration must move off forward, which a
+		// filesystem reporting a hole at or before where we started would not.
+		if holeOff <= off {
+			return fmt.Errorf("seeking to hole at %d returned non-advancing offset %d", dataOff, holeOff)
+		}
+		if holeOff > size {
+			holeOff = size
+		}
+		for pos := dataOff; pos < holeOff; {
+			if dstFd >= 0 {
+				copied, err := kernelCopyRange(fd, dstFd, pos, holeOff-pos)
+				if err == nil {
+					pos += copied
+					continue
+				}
+				if !errors.Is(err, errKernelCopyUnsupported) {
+					return fmt.Errorf("copying %d bytes at %d: %w", holeOff-pos, pos, err)
+				}
+				// Give up on the kernel path for the rest of this file, but redo
+				// this chunk below: nothing was copied.
+				dstFd = -1
+			}
+			if buf == nil {
+				buf = make([]byte, 4<<20)
+			}
+			n := int64(len(buf))
+			if rem := holeOff - pos; rem < n {
+				n = rem
+			}
+			if _, err := src.ReadAt(buf[:n], pos); err != nil {
+				return fmt.Errorf("reading %d bytes at %d: %w", n, pos, err)
+			}
+			if _, err := dst.WriteAt(buf[:n], pos); err != nil {
+				return fmt.Errorf("writing %d bytes at %d: %w", n, pos, err)
+			}
+			pos += n
+		}
+		off = holeOff
+	}
+	return nil
+}
+
 // goldenOnlyFiles returns the golden snapshot files not shadowed by the
 // actor's own snapshot: on a DATA_ON_GOLDEN restore the actor's files (the
 // durable-dir data) win name collisions, and the golden snapshot supplies