Fix snapshot compression cleanup races

Author: sjmiller609

lib/forkvm/copy.go

@@ -42,7 +42,7 @@ func CopyGuestDirectory(srcDir, dstDir string) error {
 		if d.IsDir() && shouldSkipDirectory(relPath) {
 			return filepath.SkipDir
 		}
-		if shouldSkipRegularFile(relPath) {
+		if !d.IsDir() && shouldSkipRegularFile(relPath) {
 			return nil
 		}
 

lib/forkvm/copy_test.go

@@ -43,3 +43,17 @@ func TestCopyGuestDirectory(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "metadata.json", linkTarget)
 }
+
+func TestCopyGuestDirectory_DoesNotSkipTmpSuffixedDirectories(t *testing.T) {
+	src := filepath.Join(t.TempDir(), "src")
+	dst := filepath.Join(t.TempDir(), "dst")
+
+	tmpDir := filepath.Join(src, "snapshots", "snapshot-latest", "memory-ranges.lz4.tmp")
+	require.NoError(t, os.MkdirAll(tmpDir, 0755))
+	require.NoError(t, os.WriteFile(filepath.Join(tmpDir, "nested.txt"), []byte("nested"), 0644))
+
+	require.NoError(t, CopyGuestDirectory(src, dst))
+
+	assert.DirExists(t, filepath.Join(dst, "snapshots", "snapshot-latest", "memory-ranges.lz4.tmp"))
+	assert.FileExists(t, filepath.Join(dst, "snapshots", "snapshot-latest", "memory-ranges.lz4.tmp", "nested.txt"))
+}

lib/instances/delete.go

@@ -35,6 +35,14 @@ func (m *manager) deleteInstance(
 	stored := &meta.StoredMetadata
 	log.DebugContext(ctx, "loaded instance", "instance_id", id, "state", inst.State)
 
+	target, err := m.cancelAndWaitCompressionJob(ctx, m.snapshotJobKeyForInstance(id))
+	if err != nil {
+		return fmt.Errorf("wait for instance compression to stop: %w", err)
+	}
+	if target != nil {
+		m.recordSnapshotCompressionPreemption(ctx, snapshotCompressionPreemptionDeleteInstance, *target)
+	}
+
 	// 2. Get network allocation BEFORE killing VMM (while we can still query it)
 	var networkAlloc *network.Allocation
 	if inst.NetworkEnabled {

lib/instances/metrics.go

@@ -41,6 +41,8 @@ const (
 	snapshotCompressionPreemptionRestoreSnapshot snapshotCompressionPreemptionOperation = "restore_snapshot"
 	snapshotCompressionPreemptionForkSnapshot    snapshotCompressionPreemptionOperation = "fork_snapshot"
 	snapshotCompressionPreemptionCreateSnapshot  snapshotCompressionPreemptionOperation = "create_snapshot"
+	snapshotCompressionPreemptionDeleteInstance  snapshotCompressionPreemptionOperation = "delete_instance"
+	snapshotCompressionPreemptionDeleteSnapshot  snapshotCompressionPreemptionOperation = "delete_snapshot"
 )
 
 // Metrics holds the metrics instruments for instance operations.

lib/instances/snapshot.go

@@ -213,7 +213,13 @@ func (m *manager) createSnapshot(ctx context.Context, id string, req CreateSnaps
 }
 
 func (m *manager) deleteSnapshot(ctx context.Context, snapshotID string) error {
-	_ = ctx
+	target, err := m.cancelAndWaitCompressionJob(ctx, m.snapshotJobKeyForSnapshot(snapshotID))
+	if err != nil {
+		return fmt.Errorf("wait for snapshot compression to stop: %w", err)
+	}
+	if target != nil {
+		m.recordSnapshotCompressionPreemption(ctx, snapshotCompressionPreemptionDeleteSnapshot, *target)
+	}
 	if err := m.snapshotStore().Delete(snapshotID); err != nil {
 		if errors.Is(err, snapshotstore.ErrNotFound) {
 			return ErrSnapshotNotFound

lib/instances/snapshot_compression_test.go

@@ -1,9 +1,13 @@
 package instances
 
 import (
+	"context"
 	"errors"
+	"os"
 	"testing"
+	"time"
 
+	"github.com/kernel/hypeman/lib/hypervisor"
 	snapshotstore "github.com/kernel/hypeman/lib/snapshot"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -206,3 +210,111 @@ func TestValidateCreateSnapshotRequestRejectsStoppedCompression(t *testing.T) {
 	require.Error(t, err)
 	assert.True(t, errors.Is(err, ErrInvalidRequest))
 }
+
+func TestDeleteInstanceCancelsCompressionJob(t *testing.T) {
+	t.Parallel()
+
+	mgr, _ := setupTestManager(t)
+	ctx := context.Background()
+	const instanceID = "delete-instance-compression"
+
+	require.NoError(t, mgr.ensureDirectories(instanceID))
+	now := time.Now()
+	require.NoError(t, mgr.saveMetadata(&metadata{StoredMetadata: StoredMetadata{
+		Id:             instanceID,
+		Name:           instanceID,
+		DataDir:        mgr.paths.InstanceDir(instanceID),
+		HypervisorType: hypervisor.TypeCloudHypervisor,
+		CreatedAt:      now,
+		StoppedAt:      &now,
+	}}))
+
+	target := installCancelableCompressionJob(mgr, compressionTarget{
+		Key:         mgr.snapshotJobKeyForInstance(instanceID),
+		OwnerID:     instanceID,
+		SnapshotDir: mgr.paths.InstanceSnapshotLatest(instanceID),
+		Source:      snapshotCompressionSourceStandby,
+		Policy: snapshotstore.SnapshotCompressionConfig{
+			Enabled:   true,
+			Algorithm: snapshotstore.SnapshotCompressionAlgorithmZstd,
+			Level:     intPtr(1),
+		},
+	})
+
+	require.NoError(t, mgr.DeleteInstance(ctx, instanceID))
+	assertCompressionJobCanceled(t, mgr, target)
+	_, err := os.Stat(mgr.paths.InstanceDir(instanceID))
+	require.Error(t, err)
+	assert.True(t, os.IsNotExist(err))
+}
+
+func TestDeleteSnapshotCancelsCompressionJob(t *testing.T) {
+	t.Parallel()
+
+	mgr, _ := setupTestManager(t)
+	ctx := context.Background()
+	const snapshotID = "delete-snapshot-compression"
+
+	snapshotDir := mgr.paths.SnapshotGuestDir(snapshotID)
+	require.NoError(t, os.MkdirAll(snapshotDir, 0o755))
+	require.NoError(t, mgr.saveSnapshotRecord(&snapshotRecord{
+		Snapshot: Snapshot{
+			Id:        snapshotID,
+			Name:      snapshotID,
+			Kind:      SnapshotKindStandby,
+			CreatedAt: time.Now(),
+		},
+	}))
+
+	target := installCancelableCompressionJob(mgr, compressionTarget{
+		Key:         mgr.snapshotJobKeyForSnapshot(snapshotID),
+		SnapshotID:  snapshotID,
+		SnapshotDir: snapshotDir,
+		Source:      snapshotCompressionSourceSnapshot,
+		Policy: snapshotstore.SnapshotCompressionConfig{
+			Enabled:   true,
+			Algorithm: snapshotstore.SnapshotCompressionAlgorithmLz4,
+			Level:     intPtr(0),
+		},
+	})
+
+	require.NoError(t, mgr.DeleteSnapshot(ctx, snapshotID))
+	assertCompressionJobCanceled(t, mgr, target)
+	_, err := os.Stat(mgr.paths.SnapshotDir(snapshotID))
+	require.Error(t, err)
+	assert.True(t, os.IsNotExist(err))
+}
+
+func installCancelableCompressionJob(mgr *manager, target compressionTarget) *compressionTarget {
+	ctx, cancel := context.WithCancel(context.Background())
+	done := make(chan struct{})
+
+	mgr.compressionMu.Lock()
+	mgr.compressionJobs[target.Key] = &compressionJob{
+		cancel: cancel,
+		done:   done,
+		target: target,
+	}
+	mgr.compressionMu.Unlock()
+
+	go func() {
+		<-ctx.Done()
+		mgr.compressionMu.Lock()
+		delete(mgr.compressionJobs, target.Key)
+		mgr.compressionMu.Unlock()
+		close(done)
+	}()
+
+	return &target
+}
+
+func assertCompressionJobCanceled(t *testing.T, mgr *manager, target *compressionTarget) {
+	t.Helper()
+
+	require.Eventually(t, func() bool {
+		mgr.compressionMu.Lock()
+		defer mgr.compressionMu.Unlock()
+		_, ok := mgr.compressionJobs[target.Key]
+		return !ok
+	}, time.Second, 10*time.Millisecond)
+}