Parallelize Firecracker fork restore hot path

cmd/api/api/instances.go

@@ -620,6 +620,11 @@ func (s *ApiService) RestoreInstance(ctx context.Context, request oapi.RestoreIn
 				Code:    "invalid_state",
 				Message: err.Error(),
 			}, nil
+		case errors.Is(err, instances.ErrInsufficientResources):
+			return oapi.RestoreInstance409JSONResponse{
+				Code:    "insufficient_resources",
+				Message: err.Error(),
+			}, nil
 		default:
 			log.ErrorContext(ctx, "failed to restore instance", "error", err)
 			return oapi.RestoreInstance500JSONResponse{
@@ -684,6 +689,11 @@ func (s *ApiService) ForkInstance(ctx context.Context, request oapi.ForkInstance
 				Code:    "name_conflict",
 				Message: err.Error(),
 			}, nil
+		case errors.Is(err, instances.ErrInsufficientResources):
+			return oapi.ForkInstance409JSONResponse{
+				Code:    "insufficient_resources",
+				Message: err.Error(),
+			}, nil
 		case errors.Is(err, instances.ErrNotSupported):
 			return oapi.ForkInstance501JSONResponse{
 				Code:    "not_supported",

cmd/api/api/instances_test.go

@@ -1265,6 +1265,43 @@ func TestForkInstance_InvalidRequest(t *testing.T) {
 	assert.Equal(t, "invalid_request", badReq.Code)
 }
 
+func TestForkInstance_InsufficientResources(t *testing.T) {
+	t.Parallel()
+	svc := newTestService(t)
+
+	source := instances.Instance{
+		StoredMetadata: instances.StoredMetadata{
+			Id:             "src-instance",
+			Name:           "src-instance",
+			Image:          "docker.io/library/alpine:latest",
+			CreatedAt:      time.Now(),
+			HypervisorType: hypervisor.TypeFirecracker,
+		},
+		State: instances.StateStandby,
+	}
+
+	mockMgr := &captureForkManager{
+		Manager: svc.InstanceManager,
+		err:     fmt.Errorf("apply fork target state: %w: insufficient network bandwidth", instances.ErrInsufficientResources),
+	}
+	svc.InstanceManager = mockMgr
+
+	resp, err := svc.ForkInstance(
+		mw.WithResolvedInstance(ctx(), source.Id, source),
+		oapi.ForkInstanceRequestObject{
+			Id: source.Id,
+			Body: &oapi.ForkInstanceRequest{
+				Name: "forked-instance",
+			},
+		},
+	)
+	require.NoError(t, err)
+
+	conflict, ok := resp.(oapi.ForkInstance409JSONResponse)
+	require.True(t, ok, "expected 409 response")
+	assert.Equal(t, "insufficient_resources", conflict.Code)
+}
+
 func TestStandbyInstance_InvalidRequest(t *testing.T) {
 	t.Parallel()
 	svc := newTestService(t)

cmd/api/api/snapshots.go

@@ -196,6 +196,8 @@ func (s *ApiService) ForkSnapshot(ctx context.Context, request oapi.ForkSnapshot
 			return oapi.ForkSnapshot400JSONResponse{Code: "invalid_request", Message: err.Error()}, nil
 		case errors.Is(err, instances.ErrInvalidState), errors.Is(err, instances.ErrAlreadyExists), errors.Is(err, network.ErrNameExists):
 			return oapi.ForkSnapshot409JSONResponse{Code: "conflict", Message: err.Error()}, nil
+		case errors.Is(err, instances.ErrInsufficientResources):
+			return oapi.ForkSnapshot409JSONResponse{Code: "insufficient_resources", Message: err.Error()}, nil
 		case errors.Is(err, instances.ErrNotSupported):
 			return oapi.ForkSnapshot501JSONResponse{Code: "not_supported", Message: err.Error()}, nil
 		default:

lib/hypervisor/firecracker/fork.go

@@ -50,9 +50,20 @@ func (s *Starter) PrepareFork(ctx context.Context, req hypervisor.ForkPrepareReq
 		}
 	}
 	if req.SourceDataDir != "" && req.TargetDataDir != "" && req.SourceDataDir != req.TargetDataDir {
-		if meta.RetainSnapshotSourceDataDirAlias && meta.SnapshotSourceDataDir != "" {
-			// Keep the upstream source path for snapshot-derived forks. The retained
-			// Firecracker base can still reference that path after later diff snapshots.
+		statePath := snapshotStatePath(filepath.Dir(req.SnapshotConfigPath))
+		rewritten, err := rewriteSnapshotStatePathsForFork(statePath, snapshotStatePathRewritesForFork(meta, req.SourceDataDir, req.TargetDataDir))
+		if err != nil {
+			return hypervisor.ForkPrepareResult{}, err
+		}
+		if rewritten {
+			if meta.SnapshotSourceDataDir != "" {
+				meta.SnapshotSourceDataDir = ""
+				changed = true
+			}
+			if meta.RetainSnapshotSourceDataDirAlias {
+				meta.RetainSnapshotSourceDataDirAlias = false
+				changed = true
+			}
 		} else {
 			retainAlias := false
 			if _, err := os.Stat(req.SourceDataDir); err != nil {
@@ -62,13 +73,18 @@ func (s *Starter) PrepareFork(ctx context.Context, req hypervisor.ForkPrepareReq
 					return hypervisor.ForkPrepareResult{}, fmt.Errorf("stat snapshot source data dir %q: %w", req.SourceDataDir, err)
 				}
 			}
-			if meta.SnapshotSourceDataDir != req.SourceDataDir {
-				meta.SnapshotSourceDataDir = req.SourceDataDir
-				changed = true
-			}
-			if meta.RetainSnapshotSourceDataDirAlias != retainAlias {
-				meta.RetainSnapshotSourceDataDirAlias = retainAlias
-				changed = true
+			if meta.RetainSnapshotSourceDataDirAlias && meta.SnapshotSourceDataDir != "" {
+				// Keep the upstream source path for snapshot-derived forks. The retained
+				// Firecracker base can still reference that path after later diff snapshots.
+			} else {
+				if meta.SnapshotSourceDataDir != req.SourceDataDir {
+					meta.SnapshotSourceDataDir = req.SourceDataDir
+					changed = true
+				}
+				if meta.RetainSnapshotSourceDataDirAlias != retainAlias {
+					meta.RetainSnapshotSourceDataDirAlias = retainAlias
+					changed = true
+				}
 			}
 		}
 	}
@@ -81,3 +97,34 @@ func (s *Starter) PrepareFork(ctx context.Context, req hypervisor.ForkPrepareReq
 
 	return hypervisor.ForkPrepareResult{}, nil
 }
+
+func snapshotStatePathRewritesForFork(meta *restoreMetadata, sourceDataDir, targetDataDir string) []snapshotStatePathRewrite {
+	var rewrites []snapshotStatePathRewrite
+	add := func(source, target string) {
+		if source == "" || target == "" {
+			return
+		}
+		rewrites = append(rewrites, snapshotStatePathRewrite{Source: source, Target: target})
+	}
+
+	add(sourceDataDir, targetDataDir)
+	if meta != nil {
+		add(meta.SnapshotSourceDataDir, targetDataDir)
+	}
+	if resolvedTarget, err := filepath.EvalSymlinks(targetDataDir); err == nil {
+		add(sourceDataDir, resolvedTarget)
+		if meta != nil {
+			add(meta.SnapshotSourceDataDir, resolvedTarget)
+		}
+		if resolvedSource, err := filepath.EvalSymlinks(sourceDataDir); err == nil {
+			add(resolvedSource, resolvedTarget)
+		}
+		if meta != nil && meta.SnapshotSourceDataDir != "" {
+			if resolvedSource, err := filepath.EvalSymlinks(meta.SnapshotSourceDataDir); err == nil {
+				add(resolvedSource, resolvedTarget)
+			}
+		}
+	}
+
+	return rewrites
+}

lib/hypervisor/firecracker/fork_test.go

@@ -2,6 +2,7 @@ package firecracker
 
 import (
 	"context"
+	"encoding/binary"
 	"os"
 	"path/filepath"
 	"testing"
@@ -65,6 +66,72 @@ func TestPrepareFork_DoesNotRetainExistingSourceAlias(t *testing.T) {
 	assert.False(t, meta.RetainSnapshotSourceDataDirAlias)
 }
 
+func TestPrepareFork_RewritesSnapshotStatePathsWithoutAlias(t *testing.T) {
+	starter := NewStarter()
+	tmp := t.TempDir()
+	sourceDir := filepath.Join(tmp, "source-12345678901234567890")
+	targetDir := filepath.Join(tmp, "target-12345678901234567890")
+	snapshotDir := filepath.Join(targetDir, "snapshots", "snapshot-latest")
+	require.NoError(t, os.MkdirAll(sourceDir, 0755))
+	require.NoError(t, os.MkdirAll(snapshotDir, 0755))
+	require.NoError(t, saveRestoreMetadata(targetDir, []networkInterface{{IfaceID: "eth0", HostDevName: "tap-old"}}))
+	writeSnapshotStateForTest(t, snapshotStatePath(snapshotDir), []byte("drive="+filepath.Join(sourceDir, "overlay.raw")+" vsock="+filepath.Join(sourceDir, "vsock.sock")))
+
+	_, err := starter.PrepareFork(context.Background(), hypervisor.ForkPrepareRequest{
+		SnapshotConfigPath: filepath.Join(snapshotDir, "config.json"),
+		SourceDataDir:      sourceDir,
+		TargetDataDir:      targetDir,
+		Network: &hypervisor.ForkNetworkConfig{
+			TAPDevice: "tap-new",
+		},
+	})
+	require.NoError(t, err)
+
+	meta, err := loadRestoreMetadata(targetDir)
+	require.NoError(t, err)
+	require.Len(t, meta.NetworkOverrides, 1)
+	assert.Equal(t, "tap-new", meta.NetworkOverrides[0].HostDevName)
+	assert.Empty(t, meta.SnapshotSourceDataDir)
+	assert.False(t, meta.RetainSnapshotSourceDataDirAlias)
+
+	data, err := os.ReadFile(snapshotStatePath(snapshotDir))
+	require.NoError(t, err)
+	assert.Equal(t, uint64(0), firecrackerSnapshotCRC64(data))
+	assert.NotContains(t, string(data), sourceDir)
+	assert.Contains(t, string(data), filepath.Join(targetDir, "overlay.raw"))
+	assert.Contains(t, string(data), filepath.Join(targetDir, "vsock.sock"))
+}
+
+func TestPrepareFork_FallsBackToAliasWhenSnapshotStatePathLengthDiffers(t *testing.T) {
+	starter := NewStarter()
+	tmp := t.TempDir()
+	sourceDir := filepath.Join(tmp, "source-short")
+	targetDir := filepath.Join(tmp, "target-much-longer")
+	snapshotDir := filepath.Join(targetDir, "snapshots", "snapshot-latest")
+	require.NoError(t, os.MkdirAll(sourceDir, 0755))
+	require.NoError(t, os.MkdirAll(snapshotDir, 0755))
+	require.NoError(t, saveRestoreMetadata(targetDir, nil))
+	writeSnapshotStateForTest(t, snapshotStatePath(snapshotDir), []byte("drive="+filepath.Join(sourceDir, "overlay.raw")))
+
+	_, err := starter.PrepareFork(context.Background(), hypervisor.ForkPrepareRequest{
+		SnapshotConfigPath: filepath.Join(snapshotDir, "config.json"),
+		SourceDataDir:      sourceDir,
+		TargetDataDir:      targetDir,
+	})
+	require.NoError(t, err)
+
+	meta, err := loadRestoreMetadata(targetDir)
+	require.NoError(t, err)
+	assert.Equal(t, sourceDir, meta.SnapshotSourceDataDir)
+	assert.False(t, meta.RetainSnapshotSourceDataDirAlias)
+
+	data, err := os.ReadFile(snapshotStatePath(snapshotDir))
+	require.NoError(t, err)
+	assert.Equal(t, uint64(0), firecrackerSnapshotCRC64(data))
+	assert.Contains(t, string(data), sourceDir)
+	assert.NotContains(t, string(data), targetDir)
+}
+
 func TestPrepareFork_ReturnsSourceStatErrors(t *testing.T) {
 	starter := NewStarter()
 	tmp := t.TempDir()
@@ -137,3 +204,11 @@ func TestPrepareFork_NetworkRewritePreservesRetainedAlias(t *testing.T) {
 	assert.Equal(t, upstreamDir, meta.SnapshotSourceDataDir)
 	assert.True(t, meta.RetainSnapshotSourceDataDirAlias)
 }
+
+func writeSnapshotStateForTest(t *testing.T, path string, body []byte) {
+	t.Helper()
+	data := make([]byte, len(body)+8)
+	copy(data, body)
+	binary.LittleEndian.PutUint64(data[len(data)-8:], firecrackerSnapshotCRC64(body))
+	require.NoError(t, os.WriteFile(path, data, 0644))
+}

lib/hypervisor/firecracker/process.go

@@ -8,7 +8,6 @@ import (
 	"os/exec"
 	"path/filepath"
 	"strings"
-	"sync"
 	"syscall"
 	"time"
 
@@ -62,8 +61,6 @@ func WithUFFDClient(client UFFDClient) StarterOption {
 
 var _ hypervisor.VMStarter = (*Starter)(nil)
 
-var snapshotSourceAliasMu sync.Mutex
-
 func (s *Starter) SocketName() string {
 	return "fc.sock"
 }
@@ -158,13 +155,19 @@ func (s *Starter) RestoreVM(ctx context.Context, p *paths.Paths, version string,
 		createdUFFDSession = resp.SessionID
 		backend = snapshotMemBackend{BackendType: "Uffd", BackendPath: resp.UFFDSocketPath}
 	}
-	err = func() error {
-		snapshotSourceAliasMu.Lock()
-		defer snapshotSourceAliasMu.Unlock()
-		return withSnapshotSourceDirAlias(meta, filepath.Dir(socketPath), func() error {
-			return hv.loadSnapshot(ctx, snapshotPath, meta.NetworkOverrides, resumeOnLoad, backend)
-		})
-	}()
+	targetDataDir := filepath.Dir(socketPath)
+	loadSnapshot := func() error {
+		return hv.loadSnapshot(ctx, snapshotPath, meta.NetworkOverrides, resumeOnLoad, backend)
+	}
+	if needsSnapshotSourceDirAlias(meta, targetDataDir) {
+		err = func() error {
+			unlockAlias := hypervisor.LockSnapshotSourceAliasMutation()
+			defer unlockAlias()
+			return withSnapshotSourceDirAlias(meta, targetDataDir, loadSnapshot)
+		}()
+	} else {
+		err = loadSnapshot()
+	}
 	if err != nil {
 		if createdUFFDSession != "" {
 			_ = s.uffd.CloseSession(context.Background(), createdUFFDSession)
@@ -247,6 +250,15 @@ func withSnapshotSourceDirAlias(meta *restoreMetadata, targetDataDir string, run
 	return nil
 }
 
+func needsSnapshotSourceDirAlias(meta *restoreMetadata, targetDataDir string) bool {
+	if meta == nil || meta.SnapshotSourceDataDir == "" {
+		return false
+	}
+	sourceDataDir := filepath.Clean(meta.SnapshotSourceDataDir)
+	targetDataDir = filepath.Clean(targetDataDir)
+	return sourceDataDir != targetDataDir
+}
+
 func (s *Starter) startProcess(_ context.Context, p *paths.Paths, version string, socketPath string) (int, error) {
 	binaryPath, err := s.GetBinaryPath(p, version)
 	if err != nil {