Improve Firecracker fork concurrency

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{
@@ -683,6 +688,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

@@ -1261,6 +1261,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

@@ -193,6 +193,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:

cmd/api/config/config.go

@@ -198,6 +198,7 @@ type HypervisorConfig struct {
 	FirecrackerBinaryPath            string                 `koanf:"firecracker_binary_path"`
 	FirecrackerSnapshotMemoryBackend string                 `koanf:"firecracker_snapshot_memory_backend"`
 	FirecrackerUFFDCacheMaxBytes     string                 `koanf:"firecracker_uffd_cache_max_bytes"`
+	FirecrackerMaxConcurrentRestores int                    `koanf:"firecracker_max_concurrent_restores"`
 	Memory                           HypervisorMemoryConfig `koanf:"memory"`
 }
 
@@ -411,6 +412,7 @@ func defaultConfig() *Config {
 			FirecrackerBinaryPath:            "",
 			FirecrackerSnapshotMemoryBackend: "file",
 			FirecrackerUFFDCacheMaxBytes:     "4294967296",
+			FirecrackerMaxConcurrentRestores: 32,
 			Memory: HypervisorMemoryConfig{
 				Enabled:            false,
 				KernelPageInitMode: "hardened",
@@ -573,6 +575,9 @@ func (c *Config) Validate() error {
 	if c.Build.Timeout <= 0 {
 		return fmt.Errorf("build.timeout must be positive, got %d", c.Build.Timeout)
 	}
+	if c.Hypervisor.FirecrackerMaxConcurrentRestores < 0 {
+		return fmt.Errorf("hypervisor.firecracker_max_concurrent_restores must be >= 0, got %d", c.Hypervisor.FirecrackerMaxConcurrentRestores)
+	}
 	if c.Instances.LifecycleEventBufferSize <= 0 {
 		return fmt.Errorf("instances.lifecycle_event_buffer_size must be positive, got %d", c.Instances.LifecycleEventBufferSize)
 	}

cmd/api/config/config_test.go

@@ -61,6 +61,9 @@ func TestDefaultConfigIncludesMetricsSettings(t *testing.T) {
 	if cfg.Hypervisor.FirecrackerUFFDCacheMaxBytes != "4294967296" {
 		t.Fatalf("expected default firecracker uffd cache size to be 4294967296, got %q", cfg.Hypervisor.FirecrackerUFFDCacheMaxBytes)
 	}
+	if cfg.Hypervisor.FirecrackerMaxConcurrentRestores != 32 {
+		t.Fatalf("expected default firecracker max concurrent restores to be 32, got %d", cfg.Hypervisor.FirecrackerMaxConcurrentRestores)
+	}
 }
 
 func TestValidateFirecrackerSnapshotMemoryBackend(t *testing.T) {
@@ -86,6 +89,14 @@ func TestValidateFirecrackerSnapshotMemoryBackend(t *testing.T) {
 	}
 }
 
+func TestValidateFirecrackerMaxConcurrentRestores(t *testing.T) {
+	cfg := defaultConfig()
+	cfg.Hypervisor.FirecrackerMaxConcurrentRestores = -1
+	if err := cfg.Validate(); err == nil {
+		t.Fatalf("expected invalid firecracker max concurrent restores validation error")
+	}
+}
+
 func TestLoadEnvOverridesMetricsAndOtelInterval(t *testing.T) {
 	t.Setenv("METRICS__LISTEN_ADDRESS", "0.0.0.0")
 	t.Setenv("METRICS__PORT", "9999")

config.example.yaml

@@ -26,6 +26,8 @@ data_dir: /var/lib/hypeman
 #   default: cloud-hypervisor
 #   # Optional: use a custom Firecracker binary path instead of the embedded one.
 #   # firecracker_binary_path: /usr/local/bin/firecracker
+#   # Caps concurrent Firecracker snapshot restore/resume work to avoid host overload.
+#   # firecracker_max_concurrent_restores: 32
 #   memory:
 #     enabled: false
 #     # performance: init_on_alloc=0 init_on_free=0 (better density)

lib/hypervisor/firecracker/firecracker.go

@@ -56,15 +56,16 @@ func (f *Firecracker) Capabilities() hypervisor.Capabilities {
 
 func capabilities() hypervisor.Capabilities {
 	return hypervisor.Capabilities{
-		SupportsSnapshot:            true,
-		SupportsHotplugMemory:       false,
-		SupportsBalloonControl:      true,
-		SupportsPause:               true,
-		SupportsVsock:               true,
-		SupportsGPUPassthrough:      false,
-		SupportsDiskIOLimit:         true,
-		SupportsGracefulVMMShutdown: false,
-		SupportsSnapshotBaseReuse:   true,
+		SupportsSnapshot:              true,
+		SupportsHotplugMemory:         false,
+		SupportsBalloonControl:        true,
+		SupportsPause:                 true,
+		SupportsVsock:                 true,
+		SupportsGPUPassthrough:        false,
+		SupportsDiskIOLimit:           true,
+		SupportsGracefulVMMShutdown:   false,
+		SupportsSnapshotBaseReuse:     true,
+		SupportsConcurrentForkPrepare: true,
 	}
 }
 

lib/hypervisor/firecracker/fork.go

@@ -2,8 +2,6 @@ package firecracker
 
 import (
 	"context"
-	"fmt"
-	"os"
 	"path/filepath"
 
 	"github.com/kernel/hypeman/lib/hypervisor"
@@ -33,6 +31,7 @@ func (s *Starter) PrepareFork(ctx context.Context, req hypervisor.ForkPrepareReq
 	}
 
 	changed := false
+	result := hypervisor.ForkPrepareResult{}
 	if req.Network != nil && req.Network.TAPDevice != "" {
 		if len(meta.NetworkOverrides) == 0 {
 			meta.NetworkOverrides = []networkOverride{{
@@ -50,27 +49,12 @@ 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.
-		} else {
-			retainAlias := false
-			if _, err := os.Stat(req.SourceDataDir); err != nil {
-				if os.IsNotExist(err) {
-					retainAlias = true
-				} else {
-					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
-			}
+		updated, err := prepareSnapshotDataDirForFork(req.SnapshotConfigPath, meta, req.SourceDataDir, req.TargetDataDir)
+		if err != nil {
+			return hypervisor.ForkPrepareResult{}, err
 		}
+		changed = changed || updated
+		result.RequiresSnapshotSourceAlias = needsSnapshotSourceDirAlias(meta, req.TargetDataDir)
 	}
 
 	if changed {
@@ -79,5 +63,5 @@ func (s *Starter) PrepareFork(ctx context.Context, req hypervisor.ForkPrepareReq
 		}
 	}
 
-	return hypervisor.ForkPrepareResult{}, nil
+	return result, nil
 }

lib/hypervisor/firecracker/fork_test.go

@@ -2,6 +2,7 @@ package firecracker
 
 import (
 	"context"
+	"encoding/binary"
 	"os"
 	"path/filepath"
 	"testing"
@@ -65,6 +66,74 @@ 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")))
+
+	result, 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)
+	assert.False(t, result.RequiresSnapshotSourceAlias)
+
+	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")))
+
+	result, err := starter.PrepareFork(context.Background(), hypervisor.ForkPrepareRequest{
+		SnapshotConfigPath: filepath.Join(snapshotDir, "config.json"),
+		SourceDataDir:      sourceDir,
+		TargetDataDir:      targetDir,
+	})
+	require.NoError(t, err)
+	assert.True(t, result.RequiresSnapshotSourceAlias)
+
+	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()
@@ -81,6 +150,30 @@ func TestPrepareFork_ReturnsSourceStatErrors(t *testing.T) {
 	assert.Contains(t, err.Error(), "stat snapshot source data dir")
 }
 
+func TestPrepareFork_RetainedAliasSkipsSourceStat(t *testing.T) {
+	starter := NewStarter()
+	tmp := t.TempDir()
+	upstreamDir := filepath.Join(tmp, "upstream")
+	targetDir := filepath.Join(tmp, "target")
+	require.NoError(t, os.MkdirAll(targetDir, 0755))
+	require.NoError(t, saveRestoreMetadataState(targetDir, &restoreMetadata{
+		SnapshotSourceDataDir:            upstreamDir,
+		RetainSnapshotSourceDataDirAlias: true,
+	}))
+
+	_, err := starter.PrepareFork(context.Background(), hypervisor.ForkPrepareRequest{
+		SnapshotConfigPath: filepath.Join(targetDir, "snapshots", "snapshot-latest", "config.json"),
+		SourceDataDir:      filepath.Join(tmp, "source") + "\x00",
+		TargetDataDir:      targetDir,
+	})
+	require.NoError(t, err)
+
+	meta, err := loadRestoreMetadata(targetDir)
+	require.NoError(t, err)
+	assert.Equal(t, upstreamDir, meta.SnapshotSourceDataDir)
+	assert.True(t, meta.RetainSnapshotSourceDataDirAlias)
+}
+
 func TestPrepareFork_PreservesRetainedUpstreamAlias(t *testing.T) {
 	starter := NewStarter()
 	tmp := t.TempDir()
@@ -137,3 +230,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))
+}