fix: address P1/P2 review findings from PR #16

P1 GC: Implement RegisterGC for FC backend — protects blob IDs
referenced by FC VMs from garbage collection, mirroring CH's GC module.

P1 Clone paths: Save cocoon.json metadata (StorageConfigs + BootConfig)
in snapshot tar. Create temporary symlinks from source drive paths to
clone paths before snapshot/load so FC finds drives at expected locations.
Symlinks are cleaned up after load + reconfigure.

P2 Rebuild: Replace fragile rebuildFromSnapshot (searched live VM records)
with self-contained metadata from cocoon.json. Clones no longer depend
on the source VM or any sibling VM existing in the DB.

P2 Console relay: Add 3s timeout on second goroutine wait after client
disconnect to prevent blocking the accept loop when PTY read is stuck.

Author: CMGS

hypervisor/firecracker/clone.go

@@ -65,36 +65,34 @@ func (fc *Firecracker) cloneSetup(ctx context.Context, vmID string, vmCfg *types
 func (fc *Firecracker) cloneAfterExtract(ctx context.Context, vmID string, vmCfg *types.VMConfig, networkConfigs []*types.NetworkConfig, runDir, logDir string, now time.Time) (*types.VM, error) {
 	logger := log.WithFunc("firecracker.Clone")
 
-	// Rebuild storage/boot configs from the snapshot's cow.raw and existing record context.
-	// FC stores StorageConfigs on the VMRecord (not in a config.json like CH).
-	// We need to find the COW file in the snapshot and update paths.
+	// Read snapshot metadata (cocoon.json) to reconstruct storage/boot config.
+	// This makes the clone self-contained — no dependency on live VM records.
+	meta, err := loadSnapshotMeta(runDir)
+	if err != nil {
+		return nil, fmt.Errorf("load snapshot metadata: %w", err)
+	}
+
 	cowPath := fc.conf.COWRawPath(vmID)
 	snapshotCOW := filepath.Join(runDir, cowFileName)
-
-	// Move the extracted COW to its canonical location.
-	if err := os.Rename(snapshotCOW, cowPath); err != nil {
-		return nil, fmt.Errorf("move COW to canonical path: %w", err)
+	if renameErr := os.Rename(snapshotCOW, cowPath); renameErr != nil {
+		return nil, fmt.Errorf("move COW to canonical path: %w", renameErr)
 	}
 
-	// Rebuild storage configs: read-only layers from snapshot config (via blob IDs),
-	// plus the new COW disk.
-	storageConfigs, bootCfg, blobIDs, err := fc.rebuildFromSnapshot(ctx, vmID, vmCfg, cowPath)
-	if err != nil {
-		return nil, fmt.Errorf("rebuild from snapshot: %w", err)
-	}
+	// Rebuild storage configs: reuse layer paths from metadata, update COW path.
+	storageConfigs := rebuildCloneStorage(meta, cowPath)
+	bootCfg := meta.BootConfig
+	blobIDs := hypervisor.ExtractBlobIDs(storageConfigs, bootCfg)
 
 	if verifyErr := verifyBaseFiles(storageConfigs, bootCfg); verifyErr != nil {
 		return nil, fmt.Errorf("verify base files: %w", verifyErr)
 	}
 
-	// Expand COW if vmCfg requests larger storage.
 	if vmCfg.Storage > 0 {
 		if expandErr := expandRawImage(cowPath, vmCfg.Storage); expandErr != nil {
 			return nil, fmt.Errorf("resize COW: %w", expandErr)
 		}
 	}
 
-	// Update bootCfg.Cmdline for the new clone (new VM name, IP, DNS).
 	if bootCfg != nil {
 		dns, dnsErr := fc.conf.DNSServers()
 		if dnsErr != nil {
@@ -103,9 +101,13 @@ func (fc *Firecracker) cloneAfterExtract(ctx context.Context, vmID string, vmCfg
 		bootCfg.Cmdline = buildCmdline(storageConfigs, networkConfigs, vmCfg.Name, dns)
 	}
 
-	// Launch FC process, load snapshot, configure drives, resume.
-	sockPath := hypervisor.SocketPath(runDir)
+	// FC snapshot/load requires drives at the same paths as the source.
+	// Read-only layers are shared blobs (same path). COW changed path.
+	// Create a temp symlink from the source COW path -> clone COW so load succeeds.
+	symlinks := createDriveSymlinks(meta.StorageConfigs, storageConfigs)
+	defer cleanupSymlinks(symlinks)
 
+	sockPath := hypervisor.SocketPath(runDir)
 	withNetwork := len(networkConfigs) > 0
 	pid, err := fc.launchProcess(ctx, &hypervisor.VMRecord{
 		VM:     types.VM{NetworkConfigs: networkConfigs},
@@ -121,16 +123,10 @@ func (fc *Firecracker) cloneAfterExtract(ctx context.Context, vmID string, vmCfg
 		return nil, err
 	}
 
-	// Finalize record -> Running.
 	info := types.VM{
-		ID:             vmID,
-		State:          types.VMStateRunning,
-		Config:         *vmCfg,
-		StorageConfigs: storageConfigs,
-		NetworkConfigs: networkConfigs,
-		CreatedAt:      now,
-		UpdatedAt:      now,
-		StartedAt:      &now,
+		ID: vmID, State: types.VMStateRunning,
+		Config: *vmCfg, StorageConfigs: storageConfigs, NetworkConfigs: networkConfigs,
+		CreatedAt: now, UpdatedAt: now, StartedAt: &now,
 	}
 	if err := fc.DB.Update(ctx, func(idx *hypervisor.VMIndex) error {
 		r := idx.VMs[vmID]
@@ -169,13 +165,9 @@ func (fc *Firecracker) restoreAndResumeClone(
 		return fmt.Errorf("snapshot/load: %w", err)
 	}
 
-	// Re-configure drives after snapshot load.
-	// FC snapshot/load does NOT preserve drive config; drives must be re-attached.
 	if err = fc.reconfigureDrives(ctx, hc, storageConfigs); err != nil {
 		return fmt.Errorf("reconfigure drives: %w", err)
 	}
-
-	// Re-configure network interfaces for the clone (new TAP devices, new MACs).
 	if err = fc.reconfigureNetworks(ctx, hc, networkConfigs); err != nil {
 		return fmt.Errorf("reconfigure networks: %w", err)
 	}
@@ -186,56 +178,50 @@ func (fc *Firecracker) restoreAndResumeClone(
 	return nil
 }
 
-// rebuildFromSnapshot reconstructs StorageConfigs, BootConfig, and blob IDs
-// from the VM's image (looked up via vmCfg.Image) plus the new COW path.
-// FC only supports OCI (direct boot), so we always have a kernel+initrd+layers.
-func (fc *Firecracker) rebuildFromSnapshot(ctx context.Context, _ string, vmCfg *types.VMConfig, cowPath string) ([]*types.StorageConfig, *types.BootConfig, map[string]struct{}, error) {
-	// Look up the original VM that was snapshotted to find its storage layout.
-	// For clone, the snapshot already carried the COW; we need the read-only layers
-	// which are shared blobs on disk (referenced by the image).
-	// The caller (cmd layer) already resolved the image and passed storageConfigs
-	// via snapshotConfig.ImageBlobIDs. We reconstruct from the index.
-
-	// Search for any existing VM with the same image to get layer paths.
-	// This is a fallback; the primary path is through the image resolver at the cmd layer.
-	var storageConfigs []*types.StorageConfig
-	var bootCfg *types.BootConfig
-
-	if err := fc.DB.With(ctx, func(idx *hypervisor.VMIndex) error {
-		for _, rec := range idx.VMs {
-			if rec == nil || rec.Config.Image != vmCfg.Image {
-				continue
-			}
-			// Found a VM with the same image; reuse its read-only layers and boot config.
-			for _, sc := range rec.StorageConfigs {
-				if sc.RO {
-					storageConfigs = append(storageConfigs, &types.StorageConfig{
-						Path:   sc.Path,
-						RO:     true,
-						Serial: sc.Serial,
-					})
+// rebuildCloneStorage creates new StorageConfigs from snapshot metadata,
+// keeping read-only layer paths unchanged and updating the COW path.
+func rebuildCloneStorage(meta *snapshotMeta, cowPath string) []*types.StorageConfig {
+	var configs []*types.StorageConfig
+	for _, sc := range meta.StorageConfigs {
+		if sc.RO {
+			configs = append(configs, &types.StorageConfig{Path: sc.Path, RO: true, Serial: sc.Serial})
+		}
+	}
+	configs = append(configs, &types.StorageConfig{Path: cowPath, RO: false, Serial: CowSerial})
+	return configs
+}
+
+// createDriveSymlinks creates temporary symlinks from source drive paths to
+// clone drive paths so FC snapshot/load can find drives at their original locations.
+// Only creates symlinks for paths that actually changed (i.e., COW disk).
+func createDriveSymlinks(srcConfigs, dstConfigs []*types.StorageConfig) []string {
+	var symlinks []string
+	dstPaths := make(map[string]string) // srcPath → dstPath
+	for i, src := range srcConfigs {
+		if i < len(dstConfigs) && src.Path != dstConfigs[i].Path {
+			dstPaths[src.Path] = dstConfigs[i].Path
+		}
+	}
+	for srcPath, dstPath := range dstPaths {
+		// Only create if source path doesn't already exist (avoid overwriting real files).
+		if _, err := os.Stat(srcPath); err != nil {
+			// Ensure parent directory exists for the symlink.
+			if mkErr := os.MkdirAll(filepath.Dir(srcPath), 0o700); mkErr == nil {
+				if linkErr := os.Symlink(dstPath, srcPath); linkErr == nil {
+					symlinks = append(symlinks, srcPath)
 				}
 			}
-			if rec.BootConfig != nil {
-				b := *rec.BootConfig
-				bootCfg = &b
-			}
-			return nil
 		}
-		return fmt.Errorf("no VM with image %q found for layer reference", vmCfg.Image)
-	}); err != nil {
-		return nil, nil, nil, err
 	}
+	return symlinks
+}
 
-	// Append the new COW disk.
-	storageConfigs = append(storageConfigs, &types.StorageConfig{
-		Path:   cowPath,
-		RO:     false,
-		Serial: CowSerial,
-	})
-
-	blobIDs := hypervisor.ExtractBlobIDs(storageConfigs, bootCfg)
-	return storageConfigs, bootCfg, blobIDs, nil
+func cleanupSymlinks(paths []string) {
+	for _, p := range paths {
+		_ = os.Remove(p)
+		// Clean up parent dir if it was created for the symlink (best-effort).
+		_ = os.Remove(filepath.Dir(p))
+	}
 }
 
 // reconfigureDrives re-attaches drives after FC snapshot/load.

hypervisor/firecracker/firecracker.go

@@ -6,7 +6,6 @@ import (
 	"fmt"
 
 	"github.com/cocoonstack/cocoon/config"
-	"github.com/cocoonstack/cocoon/gc"
 	"github.com/cocoonstack/cocoon/hypervisor"
 	"github.com/cocoonstack/cocoon/lock/flock"
 	storejson "github.com/cocoonstack/cocoon/storage/json"
@@ -83,5 +82,3 @@ func (fc *Firecracker) Delete(ctx context.Context, refs []string, force bool) ([
 		})
 	})
 }
-
-func (fc *Firecracker) RegisterGC(_ *gc.Orchestrator) {}

hypervisor/firecracker/gc.go

@@ -0,0 +1,110 @@
+package firecracker
+
+import (
+	"context"
+	"errors"
+	"slices"
+	"time"
+
+	"github.com/cocoonstack/cocoon/gc"
+	"github.com/cocoonstack/cocoon/hypervisor"
+	"github.com/cocoonstack/cocoon/types"
+	"github.com/cocoonstack/cocoon/utils"
+)
+
+const creatingStateGCGrace = 24 * time.Hour
+
+type fcSnapshot struct {
+	blobIDs     map[string]struct{}
+	vmIDs       map[string]struct{}
+	staleCreate []string
+	runDirs     []string
+	logDirs     []string
+}
+
+func (s fcSnapshot) UsedBlobIDs() map[string]struct{} { return s.blobIDs }
+func (s fcSnapshot) ActiveVMIDs() map[string]struct{} { return s.vmIDs }
+
+// GCModule returns the GC module for cross-module blob pinning and orphan cleanup.
+func (fc *Firecracker) GCModule() gc.Module[fcSnapshot] {
+	return gc.Module[fcSnapshot]{
+		Name:   typ,
+		Locker: fc.Locker,
+		ReadDB: func(_ context.Context) (fcSnapshot, error) {
+			var snap fcSnapshot
+			cutoff := time.Now().Add(-creatingStateGCGrace)
+			if err := fc.DB.ReadRaw(func(idx *hypervisor.VMIndex) error {
+				snap.blobIDs = make(map[string]struct{})
+				snap.vmIDs = make(map[string]struct{})
+				for id, rec := range idx.VMs {
+					if rec == nil {
+						continue
+					}
+					snap.vmIDs[id] = struct{}{}
+					for hex := range rec.ImageBlobIDs {
+						snap.blobIDs[hex] = struct{}{}
+					}
+					if rec.State == types.VMStateCreating && rec.UpdatedAt.Before(cutoff) {
+						snap.staleCreate = append(snap.staleCreate, id)
+					}
+				}
+				return nil
+			}); err != nil {
+				return snap, err
+			}
+			var err error
+			if snap.runDirs, err = utils.ScanSubdirs(fc.conf.RunDir()); err != nil {
+				return snap, err
+			}
+			if snap.logDirs, err = utils.ScanSubdirs(fc.conf.LogDir()); err != nil {
+				return snap, err
+			}
+			return snap, nil
+		},
+		Resolve: func(snap fcSnapshot, _ map[string]any) []string {
+			reserved := map[string]struct{}{"db": {}}
+			runOrphans := utils.FilterUnreferenced(snap.runDirs, snap.vmIDs, reserved)
+			logOrphans := utils.FilterUnreferenced(snap.logDirs, snap.vmIDs, reserved)
+			candidates := slices.Concat(runOrphans, logOrphans, snap.staleCreate)
+			slices.Sort(candidates)
+			return slices.Compact(candidates)
+		},
+		Collect: func(ctx context.Context, ids []string) error {
+			var errs []error
+			for _, id := range ids {
+				runDir, logDir := fc.conf.VMRunDir(id), fc.conf.VMLogDir(id)
+				if rec, loadErr := fc.LoadRecord(ctx, id); loadErr == nil {
+					runDir, logDir = rec.RunDir, rec.LogDir
+				}
+				if err := hypervisor.RemoveVMDirs(runDir, logDir); err != nil {
+					errs = append(errs, err)
+				}
+			}
+			if err := fc.cleanStalePlaceholders(ctx, ids); err != nil {
+				errs = append(errs, err)
+			}
+			return errors.Join(errs...)
+		},
+	}
+}
+
+// RegisterGC registers the Firecracker GC module with the given Orchestrator.
+func (fc *Firecracker) RegisterGC(orch *gc.Orchestrator) {
+	gc.Register(orch, fc.GCModule())
+}
+
+func (fc *Firecracker) cleanStalePlaceholders(_ context.Context, ids []string) error {
+	if len(ids) == 0 {
+		return nil
+	}
+	cutoff := time.Now().Add(-creatingStateGCGrace)
+	return fc.DB.WriteRaw(func(idx *hypervisor.VMIndex) error {
+		utils.CleanStaleRecords(idx.VMs, idx.Names, ids,
+			func(r *hypervisor.VMRecord) string { return r.Config.Name },
+			func(r *hypervisor.VMRecord) bool {
+				return r.State == types.VMStateCreating && r.UpdatedAt.Before(cutoff)
+			},
+		)
+		return nil
+	})
+}

hypervisor/firecracker/relay.go

@@ -99,7 +99,13 @@ func relayBidirectional(ctx context.Context, a io.ReadWriter, b io.ReadWriteClos
 	case <-done:
 	case <-ctx.Done():
 	}
-	// Close the conn to unblock the other goroutine's io.Copy.
+	// Close the conn to unblock conn→master io.Copy.
+	// The master→conn io.Copy may remain blocked on PTY read until the
+	// guest writes to serial or FC exits. Use a timeout to avoid blocking
+	// the accept loop indefinitely.
 	_ = b.Close()
-	<-done // wait for the second goroutine
+	select {
+	case <-done:
+	case <-time.After(3 * time.Second): //nolint:mnd
+	}
 }

hypervisor/firecracker/snapshot.go

@@ -2,6 +2,7 @@ package firecracker
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"io"
 	"maps"
@@ -18,6 +19,7 @@ import (
 const (
 	snapshotVMStateFile = "vmstate"
 	snapshotMemFile     = "mem"
+	snapshotMetaFile    = "cocoon.json"
 )
 
 // Snapshot pauses the VM, captures its full state (CPU/device state via FC
@@ -88,6 +90,13 @@ func (fc *Firecracker) Snapshot(ctx context.Context, ref string) (*types.Snapsho
 		return nil, nil, fmt.Errorf("snapshot VM %s: %w", vmID, err)
 	}
 
+	// Save snapshot metadata so clones can reconstruct storage/boot config
+	// without depending on live VM records.
+	if metaErr := saveSnapshotMeta(tmpDir, rec.StorageConfigs, rec.BootConfig); metaErr != nil {
+		os.RemoveAll(tmpDir) //nolint:errcheck,gosec
+		return nil, nil, fmt.Errorf("save snapshot metadata: %w", metaErr)
+	}
+
 	// Generate snapshot ID and record it on the VM atomically.
 	snapID, genErr := utils.GenerateID()
 	if genErr != nil {
@@ -126,3 +135,31 @@ func (fc *Firecracker) Snapshot(ctx context.Context, ref string) (*types.Snapsho
 
 	return cfg, utils.TarDirStreamWithRemove(tmpDir), nil
 }
+
+// snapshotMeta is persisted as cocoon.json inside the snapshot tar.
+// It makes the snapshot self-contained: clones can reconstruct storage/boot
+// config without depending on live VM records or image backends.
+type snapshotMeta struct {
+	StorageConfigs []*types.StorageConfig `json:"storage_configs"`
+	BootConfig     *types.BootConfig      `json:"boot_config,omitempty"`
+}
+
+func saveSnapshotMeta(dir string, storageConfigs []*types.StorageConfig, boot *types.BootConfig) error {
+	data, err := json.Marshal(snapshotMeta{StorageConfigs: storageConfigs, BootConfig: boot})
+	if err != nil {
+		return fmt.Errorf("marshal: %w", err)
+	}
+	return os.WriteFile(filepath.Join(dir, snapshotMetaFile), data, 0o600)
+}
+
+func loadSnapshotMeta(dir string) (*snapshotMeta, error) {
+	data, err := os.ReadFile(filepath.Join(dir, snapshotMetaFile)) //nolint:gosec
+	if err != nil {
+		return nil, fmt.Errorf("read %s: %w", snapshotMetaFile, err)
+	}
+	var meta snapshotMeta
+	if err := json.Unmarshal(data, &meta); err != nil {
+		return nil, fmt.Errorf("decode %s: %w", snapshotMetaFile, err)
+	}
+	return &meta, nil
+}