feat: mark pulled clone sources and reclaim them under any name

docs/reconcile.md

@@ -80,12 +80,18 @@ published IP and macOS VNC-port annotations, removing an endpoint when the
 runtime no longer serves it; failed patches are retried on the next pass.
 
 At startup and every 10 minutes, the snapshot reclaimer walks the local
-snapshot store. It skips a `vk-` snapshot that a pod on this node runs, wakes
-or clones from. Any other `vk-` snapshot is removed with its fork snapshot once
-the registry's `:hibernate` tag is gone or names another snapshot. A copy the
-tag still names stays as the warm-wake cache, and a registry error keeps the
-copy until the next pass. A leftover `-hibernate-import` snapshot no pod on
-this node wakes from is removed.
+snapshot store. It skips any snapshot that a pod on this node runs, wakes or
+clones from. A clone source vk-cocoon pulled from the registry is imported with
+the description `vk-cocoon clone source`; once no pod on this node clones from
+it, it is removed whatever its name, since the registry still holds it. A `vk-`
+copy whose name carries a colon (a non-default tag or a digest, pulled before
+the description was set) is removed the same way. Any other `vk-` snapshot is
+removed with its fork snapshot once the registry's `:hibernate` tag is gone or
+names another snapshot. A copy the tag still names stays as the warm-wake
+cache, and a registry error keeps the copy until the next pass with a warning.
+A leftover `-hibernate-import` snapshot no pod on this node wakes from is
+removed. A snapshot outside the `vk-` names without the description, such as
+one an operator imported by hand, is never touched.
 
 ## VM event watcher
 

provider/cocoon/create.go

@@ -27,6 +27,8 @@ import (
 const (
 	// importDetachTimeout bounds a shared import flight, which outlives every caller that joins it.
 	importDetachTimeout = 30 * time.Minute
+	// cloneSourceDescription marks a local snapshot vk-cocoon pulled as a clone source; the reclaimer removes only marked ones outside the vk- names.
+	cloneSourceDescription = "vk-cocoon clone source"
 
 	// cpuPeriodUs is passed with the quota so the math cannot drift from cocoon's default cpu.max period.
 	cpuPeriodUs = 100000
@@ -433,7 +435,7 @@ func (p *Provider) ensureSnapshot(ctx context.Context, repo, tag, local string)
 			return snapshot, nil
 		}
 		pullStart := time.Now()
-		if pullErr := p.Puller.PullSnapshot(shared, repo, tag, local); pullErr != nil {
+		if pullErr := p.Puller.PullSnapshot(shared, repo, tag, local, cloneSourceDescription); pullErr != nil {
 			return nil, pullErr
 		}
 		metrics.SnapshotPullDuration.Observe(time.Since(pullStart).Seconds())

provider/cocoon/create_test.go

@@ -1214,6 +1214,9 @@ func TestEnsureSnapshotColdBurstSingleflight(t *testing.T) {
 	if len(rt.snapshotImports) != 1 {
 		t.Errorf("snapshot imports = %v, want exactly one", rt.snapshotImports)
 	}
+	if !slices.Equal(rt.snapshotImportDescriptions, []string{cloneSourceDescription}) {
+		t.Errorf("import descriptions = %q, want the clone-source mark", rt.snapshotImportDescriptions)
+	}
 	if got := reg.manifests.Load(); got != 1 {
 		t.Errorf("manifest fetches = %d, want 1", got)
 	}
@@ -2619,11 +2622,12 @@ type fakeRuntime struct {
 	startErr   error
 	events     []vm.VMEvent
 
-	mu              sync.Mutex
-	snapshotImports []string
-	imageImports    []string
-	importHook      func()
-	ensureImageHook func()
+	mu                         sync.Mutex
+	snapshotImports            []string
+	snapshotImportDescriptions []string
+	imageImports               []string
+	importHook                 func()
+	ensureImageHook            func()
 
 	onRemove func()
 
@@ -2765,9 +2769,10 @@ func (f *fakeRuntime) SnapshotList(_ context.Context) ([]vm.Snapshot, error) {
 	return out, nil
 }
 
-func (f *fakeRuntime) SnapshotImport(ctx context.Context, name string) (io.WriteCloser, func() error, error) {
+func (f *fakeRuntime) SnapshotImport(ctx context.Context, name, description string) (io.WriteCloser, func() error, error) {
 	f.mu.Lock()
 	f.snapshotImports = append(f.snapshotImports, name)
+	f.snapshotImportDescriptions = append(f.snapshotImportDescriptions, description)
 	hook := f.importHook
 	f.mu.Unlock()
 	if hook != nil {

provider/cocoon/gc.go

@@ -32,15 +32,16 @@ func (p *Provider) reclaimLocalSnapshots(ctx context.Context) {
 	inUse := p.snapshotNamesInUse()
 	for _, s := range snapshots {
 		vmName, isImport := strings.CutSuffix(s.Name, meta.HibernateImportSuffix)
-		if !strings.HasPrefix(vmName, "vk-") || inUse[vmName] {
+		pulledCloneSource := s.Description == cloneSourceDescription
+		if inUse[vmName] || !pulledCloneSource && !strings.HasPrefix(vmName, "vk-") {
 			continue
 		}
 		if isImport {
 			logger.Infof(ctx, "reclaim local snapshot %s: no pod on this node wakes %s", s.Name, vmName)
 			p.removeSnapshotDetached(ctx, s.Name)
 			continue
 		}
-		if strings.Contains(vmName, ":") {
+		if pulledCloneSource || strings.Contains(vmName, ":") {
 			logger.Infof(ctx, "reclaim local snapshot %s: no pod on this node clones from it", s.Name)
 			p.removeSnapshotDetached(ctx, s.Name)
 			continue

provider/cocoon/gc_test.go

@@ -63,6 +63,39 @@ func TestTheReclaimerDropsOnlyWhatNoPodOnTheNodeCanWakeFrom(t *testing.T) {
 	}
 }
 
+func TestTheReclaimerDropsMarkedCloneSourcesUnderAnyName(t *testing.T) {
+	const pulled = "simular/osworld-ref:ref12"
+	cloner := meta.VMSpec{VMName: "vk-ns-clone-0-9a1f3c", Image: pulled}
+	tests := []struct {
+		name  string
+		local vm.Snapshot
+		pod   meta.VMSpec
+		want  []string
+	}{
+		{"a marked copy no pod clones from", vm.Snapshot{Name: pulled, Description: cloneSourceDescription}, meta.VMSpec{}, []string{pulled}},
+		{"a marked copy a pod clones from", vm.Snapshot{Name: pulled, Description: cloneSourceDescription}, cloner, nil},
+		{"an unmarked local snapshot under the same name", vm.Snapshot{Name: pulled}, meta.VMSpec{}, nil},
+		{"an unmarked default-tag local snapshot", vm.Snapshot{Name: "simular/golden"}, meta.VMSpec{}, nil},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			rt := &fakeRuntime{snapshots: map[string]*vm.Snapshot{tt.local.Name: &tt.local}}
+			p := newTestProvider(t)
+			p.Runtime = rt
+			p.Registry = registryDown(t)
+			if tt.pod.VMName != "" {
+				p.trackPod(newPodWithSpec(tt.pod), nil)
+			}
+
+			p.reclaimLocalSnapshots(t.Context())
+
+			if got := slices.Sorted(slices.Values(rt.snapshotRemoveCalls)); !slices.Equal(got, tt.want) {
+				t.Errorf("removed %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func tagGone(t *testing.T) oci.Registry {
 	r := newWakeVerifyRegistry(t, "SNAP-1")
 	r.tagExists = false

provider/cocoon/update.go

@@ -634,7 +634,7 @@ func (p *Provider) resolveWakeSource(ctx context.Context, namespace, vmName stri
 	}
 	importName := vmName + meta.HibernateImportSuffix
 	pullStart := time.Now()
-	if pullErr := p.Puller.PullSnapshot(ctx, vmName, meta.HibernateSnapshotTag, importName); pullErr != nil {
+	if pullErr := p.Puller.PullSnapshot(ctx, vmName, meta.HibernateSnapshotTag, importName, ""); pullErr != nil {
 		metrics.SnapshotPullTotal.WithLabelValues("failed").Inc()
 		return wakeSource{}, fmt.Errorf("pull hibernation snapshot %s: %w", vmName, pullErr)
 	}

snapshots/puller.go

@@ -29,8 +29,8 @@ type Puller struct {
 	Transfer TransferConfig
 }
 
-// PullSnapshot fetches and imports a snapshot from the registry. localName defaults to name.
-func (p *Puller) PullSnapshot(ctx context.Context, name, tag, localName string) error {
+// PullSnapshot fetches and imports a snapshot from the registry under localName (default name) with description.
+func (p *Puller) PullSnapshot(ctx context.Context, name, tag, localName, description string) error {
 	raw, _, err := p.Registry.GetManifest(ctx, name, tag)
 	if err != nil {
 		return fmt.Errorf("get snapshot manifest %s:%s: %w", name, tag, err)
@@ -45,7 +45,7 @@ func (p *Puller) PullSnapshot(ctx context.Context, name, tag, localName string)
 		defer pullGate.Release(1)
 	}
 
-	importer, wait, err := p.Runtime.SnapshotImport(ctx, localName)
+	importer, wait, err := p.Runtime.SnapshotImport(ctx, localName, description)
 	if err != nil {
 		return fmt.Errorf("open cocoon snapshot import: %w", err)
 	}

vm/cocoon_cli.go

@@ -245,11 +245,15 @@ func (c *CocoonCLI) SnapshotList(ctx context.Context) ([]Snapshot, error) {
 }
 
 // SnapshotImport spawns `cocoon snapshot import` and returns its stdin pipe, removing a stale same-name snapshot up front for idempotency.
-func (c *CocoonCLI) SnapshotImport(ctx context.Context, name string) (io.WriteCloser, func() error, error) {
+func (c *CocoonCLI) SnapshotImport(ctx context.Context, name, description string) (io.WriteCloser, func() error, error) {
 	if err := c.SnapshotRemoveIfExists(ctx, name); err != nil {
 		return nil, nil, err
 	}
-	cmd := c.command(ctx, "snapshot", "import", "--name", name)
+	args := []string{"snapshot", "import", "--name", name}
+	if description != "" {
+		args = append(args, "--description", description)
+	}
+	cmd := c.command(ctx, args...)
 	return startCmdPipe(ctx, cmd, cmd.StdinPipe, "cocoon snapshot import")
 }
 

vm/cocoon_cli_test.go

@@ -470,6 +470,39 @@ exit 1
 	}
 }
 
+func TestSnapshotImportPassesTheDescriptionOnlyWhenSet(t *testing.T) {
+	t.Parallel()
+	for _, tt := range []struct {
+		description string
+		want        string
+	}{
+		{"", "snapshot import --name snap-a"},
+		{"vk-cocoon clone source", "snapshot import --name snap-a --description vk-cocoon clone source"},
+	} {
+		dir := t.TempDir()
+		bin := filepath.Join(dir, "cocoon")
+		argsFile := filepath.Join(dir, "args")
+		writeFakeCocoon(t, bin, "#!/bin/sh\n[ \"$1\" = snapshot ] && [ \"$2\" = rm ] && exit 0\necho \"$*\" > "+argsFile+"\ncat > /dev/null\n")
+		w, wait, err := NewCocoonCLI(bin).SnapshotImport(t.Context(), "snap-a", tt.description)
+		if err != nil {
+			t.Fatalf("SnapshotImport: %v", err)
+		}
+		if err := w.Close(); err != nil {
+			t.Fatalf("close stdin: %v", err)
+		}
+		if err := wait(); err != nil {
+			t.Fatalf("wait: %v", err)
+		}
+		got, err := os.ReadFile(argsFile)
+		if err != nil {
+			t.Fatalf("read args: %v", err)
+		}
+		if strings.TrimSpace(string(got)) != tt.want {
+			t.Errorf("description %q: argv %q, want %q", tt.description, strings.TrimSpace(string(got)), tt.want)
+		}
+	}
+}
+
 func TestWatchEventsSkipsUndecodableLine(t *testing.T) {
 	script := filepath.Join(t.TempDir(), "cocoon")
 	payload := `#!/bin/sh

vm/runtime.go

@@ -71,6 +71,7 @@ type Snapshot struct {
 	Image       string `json:"image"`
 	ImageDigest string `json:"image_digest"` // resolved image digest (e.g. "sha256:abc...")
 	Hypervisor  string `json:"hypervisor"`
+	Description string `json:"description"`
 }
 
 // RestoreMode maps to `cocoon vm clone --restore-mode`; RestoreMmap needs a CH build with mmap restore support.
@@ -138,7 +139,7 @@ type Runtime interface {
 	SnapshotRemoveIfExists(ctx context.Context, name string) error
 	Snapshot(ctx context.Context, name string) (*Snapshot, error)
 	SnapshotList(ctx context.Context) ([]Snapshot, error)
-	SnapshotImport(ctx context.Context, name string) (io.WriteCloser, func() error, error)
+	SnapshotImport(ctx context.Context, name, description string) (io.WriteCloser, func() error, error)
 	SnapshotExport(ctx context.Context, vmName string) (io.ReadCloser, func() error, error)
 	EnsureImage(ctx context.Context, image string, force bool) error
 	// Image probes local presence via `cocoon image inspect`; absence is ErrImageNotFound.