Add cleanup for VM state

Author: JAORMX

image/cache.go

@@ -23,6 +23,14 @@ func NewCache(baseDir string) *Cache {
 	return &Cache{baseDir: baseDir}
 }
 
+// BaseDir returns the root directory used for cached rootfs entries.
+func (c *Cache) BaseDir() string {
+	if c == nil {
+		return ""
+	}
+	return c.baseDir
+}
+
 // Get returns the path to a cached rootfs for the given digest, and true
 // if it exists and appears valid. Returns ("", false) on a cache miss.
 func (c *Cache) Get(digest string) (string, bool) {

image/cache_test.go

@@ -22,6 +22,18 @@ func TestNewCache(t *testing.T) {
 	assert.Equal(t, baseDir, c.baseDir)
 }
 
+func TestCache_BaseDir(t *testing.T) {
+	t.Parallel()
+
+	baseDir := "/some/cache/dir"
+	c := NewCache(baseDir)
+
+	assert.Equal(t, baseDir, c.BaseDir())
+
+	var nilCache *Cache
+	assert.Equal(t, "", nilCache.BaseDir())
+}
+
 func TestCache_Has_EmptyCache(t *testing.T) {
 	t.Parallel()
 

options.go

@@ -79,6 +79,9 @@ type config struct {
 	imageCache            *image.Cache
 	imageFetcher          image.ImageFetcher // nil = default local-then-remote fallback
 	spawner               runner.Spawner     // nil = default runner.Spawn
+	cleanDataDir          bool
+	removeAll             func(string) error
+	stat                  func(string) (os.FileInfo, error)
 }
 
 func defaultConfig() *config {
@@ -92,6 +95,8 @@ func defaultConfig() *config {
 		preflight:   preflight.Default(),
 		imageCache:  image.NewCache(filepath.Join(dataDir, "cache")),
 		dataDir:     dataDir,
+		removeAll:   os.RemoveAll,
+		stat:        os.Stat,
 	}
 }
 
@@ -234,6 +239,13 @@ func WithDataDir(path string) Option {
 	})
 }
 
+// WithCleanDataDir removes any existing data directory contents before boot.
+// Use only when the data dir is VM-scoped; this will also clear the image cache
+// if it lives under the data dir.
+func WithCleanDataDir() Option {
+	return optionFunc(func(c *config) { c.cleanDataDir = true })
+}
+
 // WithRunnerPath sets the path to the propolis-runner binary.
 // When empty, the runner is found via $PATH or alongside the calling binary.
 func WithRunnerPath(path string) Option {

options_test.go

@@ -122,6 +122,16 @@ func TestWithDataDir(t *testing.T) {
 	assert.NotNil(t, cfg.imageCache) // imageCache should be recreated
 }
 
+func TestWithCleanDataDir(t *testing.T) {
+	t.Parallel()
+
+	cfg := defaultConfig()
+	assert.False(t, cfg.cleanDataDir)
+
+	WithCleanDataDir().apply(cfg)
+	assert.True(t, cfg.cleanDataDir)
+}
+
 func TestWithRootFSPath(t *testing.T) {
 	t.Parallel()
 

propolis.go

@@ -38,6 +38,11 @@ func Run(ctx context.Context, imageRef string, opts ...Option) (*VM, error) {
 	for _, opt := range opts {
 		opt.apply(cfg)
 	}
+	if cfg.cleanDataDir {
+		if err := cleanDataDir(cfg); err != nil {
+			return nil, err
+		}
+	}
 
 	if err := os.MkdirAll(cfg.dataDir, 0o700); err != nil {
 		return nil, fmt.Errorf("create data dir: %w", err)
@@ -158,6 +163,8 @@ func Run(ctx context.Context, imageRef string, opts ...Option) (*VM, error) {
 		dataDir:    cfg.dataDir,
 		rootfsPath: rootfs.Path,
 		ports:      cfg.ports,
+		cacheDir:   cacheDir(cfg),
+		removeAll:  cfg.removeAll,
 	}
 
 	// Best-effort state persistence for crash recovery.
@@ -192,6 +199,46 @@ func Run(ctx context.Context, imageRef string, opts ...Option) (*VM, error) {
 	return vm, nil
 }
 
+func cleanDataDir(cfg *config) error {
+	if cfg.dataDir == "" {
+		return nil
+	}
+	_, err := cfg.stat(cfg.dataDir)
+	if err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return fmt.Errorf("check data dir: %w", err)
+	}
+
+	var keep []string
+	cache := cacheDir(cfg)
+	if cache != "" && isWithin(cfg.dataDir, cache) {
+		keep = append(keep, cache)
+	}
+	if cfg.rootfsPath != "" && isWithin(cfg.dataDir, cfg.rootfsPath) {
+		keep = append(keep, cfg.rootfsPath)
+	}
+	if len(keep) > 0 {
+		if err := removeDataDirContentsExcept(cfg.removeAll, cfg.dataDir, keep); err != nil {
+			return fmt.Errorf("clean data dir contents: %w", err)
+		}
+		return nil
+	}
+
+	if err := cfg.removeAll(cfg.dataDir); err != nil {
+		return fmt.Errorf("clean data dir: %w", err)
+	}
+	return nil
+}
+
+func cacheDir(cfg *config) string {
+	if cfg == nil || cfg.imageCache == nil {
+		return ""
+	}
+	return cfg.imageCache.BaseDir()
+}
+
 func toRunnerPortForwards(ports []PortForward) []runner.PortForward {
 	out := make([]runner.PortForward, len(ports))
 	for i, p := range ports {

propolis_test.go

@@ -192,6 +192,44 @@ func TestRun_SpawnFailure(t *testing.T) {
 	assert.True(t, netProv.stopped)
 }
 
+func TestRun_WithCleanDataDir_RemovesStaleState(t *testing.T) {
+	t.Parallel()
+
+	dataDir := t.TempDir()
+	cacheDir := filepath.Join(dataDir, "cache")
+	stalePath := filepath.Join(dataDir, "stale.sock")
+
+	require.NoError(t, os.MkdirAll(cacheDir, 0o755))
+	require.NoError(t, os.WriteFile(filepath.Join(cacheDir, "marker"), []byte("cache"), 0o644))
+	require.NoError(t, os.WriteFile(stalePath, []byte("stale"), 0o644))
+
+	rootfsDir := filepath.Join(dataDir, "rootfs")
+	require.NoError(t, os.MkdirAll(rootfsDir, 0o755))
+
+	proc := &mockProcessHandle{pid: 1234, alive: true}
+	netProv := &mockNetProvider{sockPath: "/tmp/fake.sock"}
+
+	vm, err := Run(context.Background(), "test:latest",
+		WithDataDir(dataDir),
+		WithCleanDataDir(),
+		WithPreflightChecker(preflight.NewEmpty()),
+		WithRootFSPath(rootfsDir),
+		WithNetProvider(netProv),
+		WithSpawner(&mockSpawner{proc: proc}),
+	)
+	require.NoError(t, err)
+	require.NotNil(t, vm)
+
+	_, err = os.Stat(stalePath)
+	assert.True(t, os.IsNotExist(err))
+
+	_, err = os.Stat(cacheDir)
+	require.NoError(t, err)
+
+	_, err = os.Stat(rootfsDir)
+	require.NoError(t, err)
+}
+
 func TestRun_Success(t *testing.T) {
 	t.Parallel()
 

vm.go

@@ -7,6 +7,9 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"os"
+	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/stacklok/propolis/net"
@@ -22,6 +25,8 @@ type VM struct {
 	dataDir    string
 	rootfsPath string
 	ports      []PortForward
+	cacheDir   string
+	removeAll  func(string) error
 }
 
 // VMInfo contains status information about a VM.
@@ -80,14 +85,39 @@ func (vm *VM) Status(_ context.Context) (*VMInfo, error) {
 }
 
 // Remove stops the VM and cleans up its rootfs and state.
+// If the image cache lives under the data dir, its contents are preserved.
 func (vm *VM) Remove(ctx context.Context) error {
 	if vm.proc.IsAlive() {
 		if err := vm.Stop(ctx); err != nil {
 			return fmt.Errorf("stop before remove: %w", err)
 		}
 	}
-	// Note: we intentionally do NOT remove the image cache —
-	// only the VM-specific state and rootfs extraction.
+	if vm.removeAll == nil {
+		vm.removeAll = os.RemoveAll
+	}
+
+	if vm.rootfsPath != "" && !isWithin(vm.cacheDir, vm.rootfsPath) {
+		if err := vm.removeAll(vm.rootfsPath); err != nil {
+			return fmt.Errorf("remove rootfs: %w", err)
+		}
+	}
+
+	if vm.dataDir != "" {
+		var keep []string
+		if vm.cacheDir != "" && isWithin(vm.dataDir, vm.cacheDir) {
+			keep = append(keep, vm.cacheDir)
+		}
+		if len(keep) > 0 {
+			if err := removeDataDirContentsExcept(vm.removeAll, vm.dataDir, keep); err != nil {
+				return fmt.Errorf("remove data dir contents: %w", err)
+			}
+		} else {
+			if err := vm.removeAll(vm.dataDir); err != nil {
+				return fmt.Errorf("remove data dir: %w", err)
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -105,3 +135,44 @@ func (vm *VM) RootFSPath() string { return vm.rootfsPath }
 
 // Ports returns the configured port forwards.
 func (vm *VM) Ports() []PortForward { return vm.ports }
+
+func isWithin(base string, target string) bool {
+	if base == "" || target == "" {
+		return false
+	}
+	rel, err := filepath.Rel(base, target)
+	if err != nil {
+		return false
+	}
+	if rel == "." {
+		return true
+	}
+	if rel == ".." {
+		return false
+	}
+	return !strings.HasPrefix(rel, ".."+string(filepath.Separator))
+}
+
+func removeDataDirContentsExcept(removeAll func(string) error, dataDir string, keepPaths []string) error {
+	entries, err := os.ReadDir(dataDir)
+	if err != nil {
+		return fmt.Errorf("read data dir: %w", err)
+	}
+	keep := make(map[string]struct{}, len(keepPaths))
+	for _, path := range keepPaths {
+		if path == "" {
+			continue
+		}
+		keep[filepath.Clean(path)] = struct{}{}
+	}
+	for _, entry := range entries {
+		entryPath := filepath.Join(dataDir, entry.Name())
+		if _, ok := keep[filepath.Clean(entryPath)]; ok {
+			continue
+		}
+		if err := removeAll(entryPath); err != nil {
+			return fmt.Errorf("remove data dir entry %s: %w", entryPath, err)
+		}
+	}
+	return nil
+}

vm_test.go

@@ -6,6 +6,8 @@ package propolis
 import (
 	"context"
 	"fmt"
+	"os"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -160,6 +162,69 @@ func TestVM_Remove_StillRunning(t *testing.T) {
 	assert.True(t, netProv.stopped)
 }
 
+func TestVM_Remove_PreservesCacheContents(t *testing.T) {
+	t.Parallel()
+
+	dataDir := t.TempDir()
+	cacheDir := filepath.Join(dataDir, "cache")
+	rootfsDir := filepath.Join(cacheDir, "rootfs")
+	stalePath := filepath.Join(dataDir, "stale.sock")
+
+	require.NoError(t, os.MkdirAll(rootfsDir, 0o755))
+	require.NoError(t, os.WriteFile(filepath.Join(rootfsDir, "marker"), []byte("rootfs"), 0o644))
+	require.NoError(t, os.WriteFile(stalePath, []byte("stale"), 0o644))
+
+	vm := &VM{
+		name:       "test-vm",
+		proc:       &mockProcessHandle{pid: 42, alive: false},
+		dataDir:    dataDir,
+		rootfsPath: rootfsDir,
+		cacheDir:   cacheDir,
+		removeAll:  os.RemoveAll,
+	}
+
+	err := vm.Remove(context.Background())
+	require.NoError(t, err)
+
+	_, err = os.Stat(stalePath)
+	assert.True(t, os.IsNotExist(err))
+
+	_, err = os.Stat(rootfsDir)
+	require.NoError(t, err)
+
+	_, err = os.Stat(cacheDir)
+	require.NoError(t, err)
+}
+
+func TestVM_Remove_RemovesRootfsOutsideCache(t *testing.T) {
+	t.Parallel()
+
+	dataDir := t.TempDir()
+	cacheDir := filepath.Join(dataDir, "cache")
+	rootfsDir := t.TempDir()
+
+	require.NoError(t, os.MkdirAll(cacheDir, 0o755))
+	require.NoError(t, os.WriteFile(filepath.Join(rootfsDir, "marker"), []byte("rootfs"), 0o644))
+
+	vm := &VM{
+		name:       "test-vm",
+		proc:       &mockProcessHandle{pid: 42, alive: false},
+		dataDir:    dataDir,
+		rootfsPath: rootfsDir,
+		cacheDir:   cacheDir,
+		removeAll:  os.RemoveAll,
+	}
+
+	err := vm.Remove(context.Background())
+	require.NoError(t, err)
+
+	_, err = os.Stat(rootfsDir)
+	assert.True(t, os.IsNotExist(err))
+
+	_, err = os.Stat(cacheDir)
+	require.NoError(t, err)
+}
+
 func TestVM_Accessors(t *testing.T) {
 	t.Parallel()