Use UFFD only for initial snapshot fork restores (#270)

* Use UFFD only for initial snapshot fork restore

* Defer UFFD fork memory clone until standby

* Document one-shot UFFD fork lifecycle

* Add UFFD fork standby lifecycle test

* Clean up UFFD snapshot ownership

* Reproduce chained UFFD fork backing path bug

* Preserve deferred UFFD backing path across chained forks

* Fix UFFD one-shot standby lifecycle

* Guard Firecracker snapshot source alias during UFFD forks

* Bump UFFD pager version for one-shot restore

* Install isolated UFFD pager units in CI

* Release UFFD source lock before fork restore

* Remove UFFD restore source lock

* Clear deferred Firecracker memory on stop

* Assert UFFD lifecycle guest state persists

* Use regexp for UFFD systemd instance sanitizer

* Consolidate UFFD restore state cleanup

Author: sjmiller609

.github/workflows/test.yml

@@ -95,6 +95,29 @@ jobs:
       - name: Build
         run: make build
 
+      - name: Install UFFD pager systemd template
+        run: |
+          sudo mkdir -p /run/hypeman/uffd
+          sudo tee /etc/systemd/system/hypeman-uffd@.service > /dev/null <<'EOF'
+          [Unit]
+          Description=Hypeman UFFD Pager (%i)
+          Documentation=https://github.com/kernel/hypeman
+          After=network.target
+
+          [Service]
+          Type=simple
+          Environment="HYPEMAN_UFFD_BINARY=/opt/hypeman/bin/hypeman-uffd-pager"
+          Environment="HYPEMAN_UFFD_DATA_DIR=/var/lib/hypeman"
+          Environment="HYPEMAN_UFFD_VERSION_KEY=%i"
+          Environment="HYPEMAN_UFFD_CACHE_MAX_BYTES=4294967296"
+          EnvironmentFile=-/run/hypeman/uffd/%i.env
+          ExecStart=/bin/sh -c 'exec "$HYPEMAN_UFFD_BINARY" --data-dir "$HYPEMAN_UFFD_DATA_DIR" --version-key "$HYPEMAN_UFFD_VERSION_KEY" --cache-max-bytes "$HYPEMAN_UFFD_CACHE_MAX_BYTES"'
+          Restart=on-failure
+          RestartSec=5
+          KillMode=process
+          EOF
+          sudo systemctl daemon-reload
+
       - name: Prewarm test cache
         env:
           HYPEMAN_TEST_REGISTRY: 127.0.0.1:5001
@@ -131,13 +154,25 @@ jobs:
           TLS_ALLOWED_DOMAINS: '*.hypeman-development.com'
           HYPEMAN_TEST_PREWARM_STRICT: "1"
           HYPEMAN_TEST_REGISTRY: 127.0.0.1:5001
+          HYPEMAN_UFFD_PAGER_BINARY: ${{ runner.temp }}/hypeman-uffd-pager-${{ github.run_id }}-${{ github.run_attempt }}
+          HYPEMAN_UFFD_SYSTEMD_INSTANCE_PREFIX: ci-${{ github.run_id }}-${{ github.run_attempt }}
         run: |
+          cp "$PWD/bin/hypeman-uffd-pager" "$HYPEMAN_UFFD_PAGER_BINARY"
+          chmod +x "$HYPEMAN_UFFD_PAGER_BINARY"
           export HYPEMAN_TEST_PREWARM_DIR="$HOME/.cache/hypeman-ci/linux-amd64"
           make test TEST_TIMEOUT=20m
 
       - name: Cleanup
         if: always()
-        run: sudo rm -rf "/tmp/hm-net-${{ github.run_id }}-${{ github.run_attempt }}"
+        run: |
+          units="$(systemctl list-units --all --full --plain 'hypeman-uffd@ci-${{ github.run_id }}-${{ github.run_attempt }}-*.service' --no-legend | awk '{print $1}' || true)"
+          if [ -n "$units" ]; then
+            echo "$units" | xargs -r sudo systemctl stop || true
+            echo "$units" | xargs -r sudo systemctl reset-failed || true
+          fi
+          sudo rm -f /run/hypeman/uffd/ci-${{ github.run_id }}-${{ github.run_attempt }}-*.env
+          sudo rm -rf "/tmp/hm-net-${{ github.run_id }}-${{ github.run_attempt }}"
+          rm -f "${{ runner.temp }}/hypeman-uffd-pager-${{ github.run_id }}-${{ github.run_attempt }}"
 
   test-darwin:
     runs-on: [self-hosted, macos, arm64]

Makefile

@@ -4,6 +4,7 @@ SHELL := /bin/bash
 # Directory where local binaries will be installed
 BIN_DIR ?= $(CURDIR)/bin
 GO_TEST_TIMEOUT ?= 300s
+UFFD_PAGER_BINARY ?= $(BIN_DIR)/hypeman-uffd-pager
 
 $(BIN_DIR):
 	mkdir -p $(BIN_DIR)
@@ -256,6 +257,9 @@ build-linux: ensure-ch-binaries ensure-firecracker-binaries ensure-caddy-binarie
 	go build -tags containers_image_openpgp -o $(BIN_DIR)/hypeman ./cmd/api
 	go build -o $(BIN_DIR)/hypeman-uffd-pager ./cmd/uffd-pager
 
+$(BIN_DIR)/hypeman-uffd-pager: | $(BIN_DIR)
+	go build -o $@ ./cmd/uffd-pager
+
 # Build all binaries
 build-all: build
 
@@ -289,19 +293,23 @@ else
 endif
 
 # Linux tests (as root for network capabilities)
-test-linux: ensure-ch-binaries ensure-firecracker-binaries ensure-caddy-binaries build-embedded
+test-linux: ensure-ch-binaries ensure-firecracker-binaries ensure-caddy-binaries build-embedded $(BIN_DIR)/hypeman-uffd-pager
 	@VERBOSE_FLAG=""; \
 	TEST_PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:$$PATH"; \
 	if [ -n "$(VERBOSE)" ]; then VERBOSE_FLAG="-v"; fi; \
 	if [ -n "$(TEST)" ]; then \
 		echo "Running specific test: $(TEST)"; \
 		sudo env "PATH=$$TEST_PATH" "DOCKER_CONFIG=$${DOCKER_CONFIG:-$$HOME/.docker}" "CI=$${CI:-}" \
+			"HYPEMAN_UFFD_PAGER_BINARY=$${HYPEMAN_UFFD_PAGER_BINARY:-$(UFFD_PAGER_BINARY)}" \
+			"HYPEMAN_UFFD_SYSTEMD_INSTANCE_PREFIX=$${HYPEMAN_UFFD_SYSTEMD_INSTANCE_PREFIX:-}" \
 			"HYPEMAN_TEST_PREWARM_DIR=$${HYPEMAN_TEST_PREWARM_DIR:-}" \
 			"HYPEMAN_TEST_PREWARM_STRICT=$${HYPEMAN_TEST_PREWARM_STRICT:-}" \
 			"HYPEMAN_TEST_REGISTRY=$${HYPEMAN_TEST_REGISTRY:-}" \
 			go test -tags containers_image_openpgp -run=$(TEST) $$VERBOSE_FLAG -timeout=$(TEST_TIMEOUT) ./...; \
 	else \
 		sudo env "PATH=$$TEST_PATH" "DOCKER_CONFIG=$${DOCKER_CONFIG:-$$HOME/.docker}" "CI=$${CI:-}" \
+			"HYPEMAN_UFFD_PAGER_BINARY=$${HYPEMAN_UFFD_PAGER_BINARY:-$(UFFD_PAGER_BINARY)}" \
+			"HYPEMAN_UFFD_SYSTEMD_INSTANCE_PREFIX=$${HYPEMAN_UFFD_SYSTEMD_INSTANCE_PREFIX:-}" \
 			"HYPEMAN_TEST_PREWARM_DIR=$${HYPEMAN_TEST_PREWARM_DIR:-}" \
 			"HYPEMAN_TEST_PREWARM_STRICT=$${HYPEMAN_TEST_PREWARM_STRICT:-}" \
 			"HYPEMAN_TEST_REGISTRY=$${HYPEMAN_TEST_REGISTRY:-}" \

lib/forkvm/README.md

@@ -59,6 +59,20 @@ instead of reusing the source identity.
 - Network override fields are supplied at snapshot load to bind the fork to its
   own TAP device.
 - Vsock CID remains stable for snapshot-based flows.
+- When the Firecracker snapshot memory backend is configured as UFFD, UFFD is
+  used as a one-shot acceleration for the first restore of a newly forked
+  standby snapshot. The fork initially reuses the source snapshot memory as the
+  pager backing file instead of cloning the large memory file during fanout.
+- That deferred memory clone is paid when the fork later enters standby. Before
+  Firecracker writes the fork's diff snapshot, Hypeman materializes the fork's
+  own `snapshot-latest/memory` file from the original backing memory. After that
+  point the fork has a normal on-disk snapshot base, independent from the source.
+- Subsequent direct restores of that same fork use Firecracker's normal
+  file-backed memory backend. If that standby fork is itself forked again, the
+  new child gets its own one-shot UFFD restore.
+- This keeps UFFD on the high-fanout path where shared snapshot cache is most
+  useful, while preserving the normal Firecracker diff-snapshot lifecycle for
+  per-instance standby/resume cycles.
 
 ## VZ (Virtualization.framework)
 

lib/forkvm/copy.go

@@ -33,12 +33,22 @@ type copyState struct {
 	reflinkDead bool
 }
 
+// CopyOptions controls which guest-directory files are copied.
+type CopyOptions struct {
+	SkipRelativePaths map[string]struct{}
+}
+
 // CopyGuestDirectory recursively copies a guest directory to a new destination.
 // Regular files are cloned via reflink (FICLONE) when the underlying filesystem
 // supports it; otherwise we fall back to a sparse extent copy
 // (SEEK_DATA/SEEK_HOLE). Runtime sockets and logs are skipped because they are
 // host-runtime artifacts.
 func CopyGuestDirectory(srcDir, dstDir string) error {
+	return CopyGuestDirectoryWithOptions(srcDir, dstDir, CopyOptions{})
+}
+
+// CopyGuestDirectoryWithOptions is CopyGuestDirectory with optional path skips.
+func CopyGuestDirectoryWithOptions(srcDir, dstDir string, opts CopyOptions) error {
 	srcInfo, err := os.Stat(srcDir)
 	if err != nil {
 		return fmt.Errorf("stat source directory: %w", err)
@@ -68,9 +78,15 @@ func CopyGuestDirectory(srcDir, dstDir string) error {
 		if relPath == "." {
 			return nil
 		}
+		if _, ok := opts.SkipRelativePaths[filepath.Clean(relPath)]; ok {
+			return nil
+		}
 		if d.IsDir() && shouldSkipDirectory(relPath) {
 			return filepath.SkipDir
 		}
+		if !d.IsDir() && shouldSkipRuntimeSocket(relPath) {
+			return nil
+		}
 		if !d.IsDir() && shouldSkipRegularFile(relPath) {
 			return nil
 		}
@@ -115,6 +131,26 @@ func CopyGuestDirectory(srcDir, dstDir string) error {
 	})
 }
 
+// CopyRegularFile copies one regular file using the same reflink-first behavior
+// as CopyGuestDirectory.
+func CopyRegularFile(srcPath, dstPath string) error {
+	info, err := os.Stat(srcPath)
+	if err != nil {
+		return fmt.Errorf("stat source file: %w", err)
+	}
+	if !info.Mode().IsRegular() {
+		return fmt.Errorf("source path is not a regular file: %s", srcPath)
+	}
+	if err := os.MkdirAll(filepath.Dir(dstPath), 0755); err != nil {
+		return fmt.Errorf("create destination parent: %w", err)
+	}
+	state := &copyState{}
+	if reflinkDisabled.Load() {
+		state.reflinkDead = true
+	}
+	return copyRegularFile(state, srcPath, dstPath, info.Mode().Perm())
+}
+
 // copyRegularFile clones path to dstPath, preferring FICLONE reflink and
 // falling back to sparse extent copy. The state object lets us short-circuit
 // future reflink attempts once we observe an "unsupported" signal from the
@@ -142,3 +178,7 @@ func shouldSkipDirectory(relPath string) bool {
 func shouldSkipRegularFile(relPath string) bool {
 	return strings.HasSuffix(relPath, ".lz4.tmp") || strings.HasSuffix(relPath, ".zst.tmp")
 }
+
+func shouldSkipRuntimeSocket(relPath string) bool {
+	return strings.HasSuffix(filepath.Base(relPath), ".sock")
+}

lib/forkvm/copy_sparse_unix_test.go

@@ -81,6 +81,7 @@ func TestCopyGuestDirectory_SkipsSocketRuntimeArtifacts(t *testing.T) {
 	dst := filepath.Join(base, "dst")
 	require.NoError(t, os.MkdirAll(src, 0755))
 	require.NoError(t, os.WriteFile(filepath.Join(src, "metadata.json"), []byte(`{"id":"abc"}`), 0644))
+	require.NoError(t, os.WriteFile(filepath.Join(src, "vsock.sock"), []byte("stale regular artifact"), 0644))
 
 	socketPath := filepath.Join(src, fmt.Sprintf("vz-%d.sock", time.Now().UnixNano()))
 	listener, err := net.Listen("unix", socketPath)
@@ -90,6 +91,7 @@ func TestCopyGuestDirectory_SkipsSocketRuntimeArtifacts(t *testing.T) {
 	require.NoError(t, CopyGuestDirectory(src, dst))
 
 	assert.NoFileExists(t, filepath.Join(dst, filepath.Base(socketPath)))
+	assert.NoFileExists(t, filepath.Join(dst, "vsock.sock"))
 	assert.FileExists(t, filepath.Join(dst, "metadata.json"))
 }
 

lib/forkvm/copy_test.go

@@ -44,6 +44,43 @@ func TestCopyGuestDirectory(t *testing.T) {
 	assert.Equal(t, "metadata.json", linkTarget)
 }
 
+func TestCopyGuestDirectoryWithOptionsSkipsRelativePaths(t *testing.T) {
+	src := filepath.Join(t.TempDir(), "src")
+	dst := filepath.Join(t.TempDir(), "dst")
+
+	require.NoError(t, os.MkdirAll(filepath.Join(src, "snapshots", "snapshot-latest"), 0755))
+	require.NoError(t, os.WriteFile(filepath.Join(src, "snapshots", "snapshot-latest", "memory"), []byte("memory"), 0644))
+	require.NoError(t, os.WriteFile(filepath.Join(src, "snapshots", "snapshot-latest", "state"), []byte("state"), 0644))
+	require.NoError(t, os.WriteFile(filepath.Join(src, "overlay.raw"), []byte("overlay"), 0644))
+
+	require.NoError(t, CopyGuestDirectoryWithOptions(src, dst, CopyOptions{
+		SkipRelativePaths: map[string]struct{}{
+			filepath.Join("snapshots", "snapshot-latest", "memory"): {},
+		},
+	}))
+
+	assert.NoFileExists(t, filepath.Join(dst, "snapshots", "snapshot-latest", "memory"))
+	assert.FileExists(t, filepath.Join(dst, "snapshots", "snapshot-latest", "state"))
+	assert.FileExists(t, filepath.Join(dst, "overlay.raw"))
+}
+
+func TestCopyRegularFile(t *testing.T) {
+	src := filepath.Join(t.TempDir(), "src", "memory")
+	dst := filepath.Join(t.TempDir(), "dst", "snapshots", "snapshot-latest", "memory")
+
+	require.NoError(t, os.MkdirAll(filepath.Dir(src), 0755))
+	require.NoError(t, os.WriteFile(src, []byte("memory"), 0640))
+
+	require.NoError(t, CopyRegularFile(src, dst))
+
+	got, err := os.ReadFile(dst)
+	require.NoError(t, err)
+	assert.Equal(t, []byte("memory"), got)
+	info, err := os.Stat(dst)
+	require.NoError(t, err)
+	assert.Equal(t, os.FileMode(0640), info.Mode().Perm())
+}
+
 func TestCopyGuestDirectory_DoesNotSkipTmpSuffixedDirectories(t *testing.T) {
 	src := filepath.Join(t.TempDir(), "src")
 	dst := filepath.Join(t.TempDir(), "dst")

lib/guestmemory/controller_test.go

@@ -47,9 +47,11 @@ func (s *stubHypervisor) Shutdown(ctx context.Context) error { return nil }
 func (s *stubHypervisor) GetVMInfo(ctx context.Context) (*hypervisor.VMInfo, error) {
 	return &hypervisor.VMInfo{State: hypervisor.StateRunning}, nil
 }
-func (s *stubHypervisor) Pause(ctx context.Context) error                     { return nil }
-func (s *stubHypervisor) Resume(ctx context.Context) error                    { return nil }
-func (s *stubHypervisor) Snapshot(ctx context.Context, destPath string) error { return nil }
+func (s *stubHypervisor) Pause(ctx context.Context) error  { return nil }
+func (s *stubHypervisor) Resume(ctx context.Context) error { return nil }
+func (s *stubHypervisor) Snapshot(ctx context.Context, destPath string, _ hypervisor.SnapshotOptions) error {
+	return nil
+}
 func (s *stubHypervisor) ResizeMemory(ctx context.Context, bytes int64) error { return nil }
 func (s *stubHypervisor) ResizeMemoryAndWait(ctx context.Context, bytes int64, timeout time.Duration) error {
 	return nil

lib/hypervisor/README.md

@@ -33,7 +33,7 @@ Before using optional features, callers check capabilities:
 
 ```go
 if hv.Capabilities().SupportsSnapshot {
-    hv.Snapshot(ctx, path)
+    hv.Snapshot(ctx, path, hypervisor.SnapshotOptions{})
 }
 ```
 

lib/hypervisor/cloudhypervisor/cloudhypervisor.go

@@ -160,7 +160,7 @@ func (c *CloudHypervisor) Resume(ctx context.Context) error {
 }
 
 // Snapshot creates a VM snapshot.
-func (c *CloudHypervisor) Snapshot(ctx context.Context, destPath string) error {
+func (c *CloudHypervisor) Snapshot(ctx context.Context, destPath string, _ hypervisor.SnapshotOptions) error {
 	snapshotURL := "file://" + destPath
 	snapshotConfig := vmm.VmSnapshotConfig{DestinationUrl: &snapshotURL}
 	resp, err := c.client.PutVmSnapshotWithResponse(ctx, snapshotConfig)

lib/hypervisor/firecracker/config_test.go

@@ -101,6 +101,37 @@ func TestSnapshotLoadParamsSupportsUFFDBackend(t *testing.T) {
 	assert.Equal(t, "/tmp/pager.sock", load.MemBackend.BackendPath)
 }
 
+func TestMaterializeDeferredSnapshotMemory(t *testing.T) {
+	t.Parallel()
+
+	sourcePath := filepath.Join(t.TempDir(), "source-memory")
+	snapshotDir := filepath.Join(t.TempDir(), "snapshot-latest")
+	require.NoError(t, os.WriteFile(sourcePath, []byte("memory"), 0644))
+
+	require.NoError(t, materializeDeferredSnapshotMemory(snapshotDir, sourcePath))
+
+	got, err := os.ReadFile(filepath.Join(snapshotDir, "memory"))
+	require.NoError(t, err)
+	assert.Equal(t, []byte("memory"), got)
+}
+
+func TestMaterializeDeferredSnapshotMemoryUsesRetainedSnapshotAlternate(t *testing.T) {
+	t.Parallel()
+
+	root := t.TempDir()
+	sourcePath := filepath.Join(root, "snapshots", "snapshot-base", "memory")
+	alternatePath := filepath.Join(root, "snapshots", "snapshot-latest", "memory")
+	destPath := filepath.Join(t.TempDir(), "snapshot-latest")
+	require.NoError(t, os.MkdirAll(filepath.Dir(alternatePath), 0755))
+	require.NoError(t, os.WriteFile(alternatePath, []byte("memory"), 0644))
+
+	require.NoError(t, materializeDeferredSnapshotMemory(destPath, sourcePath))
+
+	got, err := os.ReadFile(filepath.Join(destPath, "memory"))
+	require.NoError(t, err)
+	assert.Equal(t, []byte("memory"), got)
+}
+
 func TestToBalloonConfig(t *testing.T) {
 	cfg := hypervisor.VMConfig{
 		GuestMemory: hypervisor.GuestMemoryConfig{