Delete local snapshots when an actor is terminated (#1246)

We were leaking snapshots after actor termination.

Note that this is a temporary fix: it only removes local snapshots in
one node. We should clean up the copies on any other
NodeVmsWithLocalSnapshots. This is fine *as of the day this was written*
because today NodeVmsWithLocalSnapshots has at most one item.

Related to https://github.com/agent-substrate/substrate/issues/668 and
#664


- [x] Tests pass
- [] Appropriate changes to documentation are included in the PR

cmd/atelet/lifecycle_test.go

@@ -0,0 +1,238 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"crypto/sha256"
+	"fmt"
+	"net"
+	"os"
+	"path/filepath"
+	"runtime"
+	"testing"
+
+	"github.com/agent-substrate/substrate/internal/ateompath"
+	"github.com/agent-substrate/substrate/internal/proto/ateletpb"
+	"github.com/agent-substrate/substrate/internal/proto/ateompb"
+	"google.golang.org/grpc"
+)
+
+// useTempNodeDirs roots atelet's on-node state in temp directories so a test
+// can drive the real filesystem layout. Not parallel-safe: the paths are
+// process-global.
+func useTempNodeDirs(t *testing.T) {
+	t.Helper()
+	root := t.TempDir()
+	origActors, origStatic := ateompath.ActorsDir, ateompath.StaticFilesDir
+	ateompath.ActorsDir = filepath.Join(root, "actors")
+	ateompath.StaticFilesDir = filepath.Join(root, "static-files")
+	t.Cleanup(func() {
+		ateompath.ActorsDir, ateompath.StaticFilesDir = origActors, origStatic
+	})
+}
+
+// fakeAteom is a fake ateom in a worker pod. It writes the files a
+// real checkpoint would leave in the checkpoint-state dir, and reads back
+// what a restore was handed.
+type fakeAteom struct {
+	ateompb.UnimplementedAteomServer
+	// snapshotFiles are written at checkpoint and reported back to atelet as
+	// the exact set the snapshot consists of.
+	snapshotFiles map[string]string
+	// restored holds the file contents staged into the restore-state dir by
+	// the most recent RestoreWorkload.
+	restored map[string]string
+}
+
+func (f *fakeAteom) RunWorkload(context.Context, *ateompb.RunWorkloadRequest) (*ateompb.RunWorkloadResponse, error) {
+	return &ateompb.RunWorkloadResponse{}, nil
+}
+
+func (f *fakeAteom) CheckpointWorkload(_ context.Context, req *ateompb.CheckpointWorkloadRequest) (*ateompb.CheckpointWorkloadResponse, error) {
+	dir := ateompath.CheckpointStateDir(req.GetActorUid())
+	names := make([]string, 0, len(f.snapshotFiles))
+	for name, body := range f.snapshotFiles {
+		if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600); err != nil {
+			return nil, err
+		}
+		names = append(names, name)
+	}
+	return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: names}, nil
+}
+
+func (f *fakeAteom) RestoreWorkload(_ context.Context, req *ateompb.RestoreWorkloadRequest) (*ateompb.RestoreWorkloadResponse, error) {
+	dir := ateompath.RestoreStateDir(req.GetActorUid())
+	f.restored = map[string]string{}
+	for name := range f.snapshotFiles {
+		body, err := os.ReadFile(filepath.Join(dir, name))
+		if err != nil {
+			return nil, err
+		}
+		f.restored[name] = string(body)
+	}
+	return &ateompb.RestoreWorkloadResponse{}, nil
+}
+
+func (f *fakeAteom) TerminateWorkload(context.Context, *ateompb.TerminateWorkloadRequest) (*ateompb.TerminateWorkloadResponse, error) {
+	return &ateompb.TerminateWorkloadResponse{}, nil
+}
+
+// serveFakeAteom serves ateom on a unix socket and points atelet's dialer at
+// it. The socket lives in its own short temp dir.
+func serveFakeAteom(t *testing.T, f *fakeAteom) {
+	t.Helper()
+	dir, err := os.MkdirTemp("", "ateom-")
+	if err != nil {
+		t.Fatalf("creating socket dir: %v", err)
+	}
+	t.Cleanup(func() { os.RemoveAll(dir) })
+
+	sock := filepath.Join(dir, "ateom.sock")
+	lis, err := net.Listen("unix", sock)
+	if err != nil {
+		t.Fatalf("listening on %q: %v", sock, err)
+	}
+	srv := grpc.NewServer()
+	ateompb.RegisterAteomServer(srv, f)
+	go func() { _ = srv.Serve(lis) }()
+	t.Cleanup(srv.Stop)
+
+	orig := ateomSocketPath
+	ateomSocketPath = func(string) string { return sock }
+	t.Cleanup(func() { ateomSocketPath = orig })
+}
+
+// TestLocalSnapshotGC walks an actor through
+// run -> pause -> resume -> terminate over atelet's RPC surface and ensures that
+// the local snapshot is garbage collected after the actor is terminated.
+func TestLocalSnapshotGC(t *testing.T) {
+	useTempNodeDirs(t)
+	ctx := t.Context()
+
+	const (
+		atespace     = "ate-demo"
+		actorName    = "counter"
+		actorUID     = "actor-uid-1"
+		ateomUID     = "ateom-uid-1"
+		snapshotName = "pause-snap-1"
+	)
+
+	ateom := &fakeAteom{snapshotFiles: map[string]string{"checkpoint.img": "guest-memory"}}
+	serveFakeAteom(t, ateom)
+
+	host := imageVolumeTestRegistry(t)
+	image := host + "/actor:v1"
+	pushTestImage(t, image, singleFileLayer(t, "bin/app", "app"))
+
+	// A single "runsc" asset served from a fake bucket: enough to exercise the
+	// content-addressed asset fetch without a gVisor release tarball.
+	runsc := []byte("runsc binary")
+	s := &AteomHerder{
+		ateomDialer:   newAteomDialer(1),
+		imageCache:    newImageVolumeStore(t),
+		anonGCSClient: fakeObjectStorage{data: runsc},
+	}
+	sandboxAssets := &ateletpb.SandboxAssets{
+		SandboxClass: "gvisor",
+		PauseImage:   image,
+		Assets: map[string]*ateletpb.ArchAssets{
+			runtime.GOARCH: {Files: map[string]*ateletpb.AssetFile{
+				runscAssetName: {
+					Url:    "gs://test-bucket/runsc",
+					Sha256: fmt.Sprintf("%x", sha256.Sum256(runsc)),
+				},
+			}},
+		},
+	}
+	spec := &ateletpb.WorkloadSpec{
+		Containers: []*ateletpb.Container{{Name: "app", Image: image, Command: []string{"/bin/app"}}},
+	}
+
+	if _, err := s.Run(ctx, &ateletpb.RunRequest{
+		Atespace:               atespace,
+		ActorName:              actorName,
+		ActorUid:               actorUID,
+		ActorTemplateNamespace: "default",
+		ActorTemplateName:      "counter",
+		TargetAteomUid:         ateomUID,
+		SandboxAssets:          sandboxAssets,
+		Spec:                   spec,
+	}); err != nil {
+		t.Fatalf("Run: %v", err)
+	}
+
+	// Pause: a local checkpoint, which leaves the snapshot on this node.
+	if _, err := s.Checkpoint(ctx, &ateletpb.CheckpointRequest{
+		Atespace:               atespace,
+		ActorName:              actorName,
+		ActorUid:               actorUID,
+		ActorTemplateNamespace: "default",
+		ActorTemplateName:      "counter",
+		TargetAteomUid:         ateomUID,
+		Spec:                   spec,
+		Scope:                  ateletpb.SnapshotScope_SNAPSHOT_SCOPE_FULL,
+		Type:                   ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL,
+		Config: &ateletpb.CheckpointRequest_LocalConfig{
+			LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotName: snapshotName},
+		},
+	}); err != nil {
+		t.Fatalf("Checkpoint: %v", err)
+	}
+	snapshotFile := filepath.Join(ateompath.LocalSnapshotDir(actorUID, snapshotName), "checkpoint.img")
+	if _, err := os.Stat(snapshotFile); err != nil {
+		t.Fatalf("pause did not write the local snapshot: %v", err)
+	}
+
+	// Resume: restores from that local snapshot.
+	if _, err := s.Restore(ctx, &ateletpb.RestoreRequest{
+		Atespace:               atespace,
+		ActorName:              actorName,
+		ActorUid:               actorUID,
+		ActorTemplateNamespace: "default",
+		ActorTemplateName:      "counter",
+		TargetAteomUid:         ateomUID,
+		Spec:                   spec,
+		Scope:                  ateletpb.SnapshotScope_SNAPSHOT_SCOPE_FULL,
+		Type:                   ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL,
+		Config: &ateletpb.RestoreRequest_LocalConfig{
+			LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotName: snapshotName},
+		},
+	}); err != nil {
+		t.Fatalf("Restore: %v", err)
+	}
+	if got := ateom.restored["checkpoint.img"]; got != "guest-memory" {
+		t.Fatalf("restore staged %q for ateom, want the pause snapshot's %q", got, "guest-memory")
+	}
+
+	// Terminate: the actor is gone, and so should its snapshot be.
+	if _, err := s.Terminate(ctx, &ateletpb.TerminateRequest{
+		Atespace:               atespace,
+		ActorName:              actorName,
+		ActorUid:               actorUID,
+		ActorTemplateNamespace: "default",
+		ActorTemplateName:      "counter",
+		TargetAteomUid:         ateomUID,
+		Spec:                   spec,
+	}); err != nil {
+		t.Fatalf("Terminate: %v", err)
+	}
+
+	localDir := ateompath.LocalCheckpointsDir(actorUID)
+	if _, err := os.Stat(localDir); !os.IsNotExist(err) {
+		leaked, _ := filepath.Glob(filepath.Join(localDir, "*", "*"))
+		t.Errorf("local checkpoint dir survived terminate (stat err = %v), leaked files: %v", err, leaked)
+	}
+}

cmd/atelet/local_checkpoints.go

@@ -16,6 +16,8 @@ package main
 
 import (
 	"context"
+	"errors"
+	"fmt"
 	"log/slog"
 	"os"
 	"path/filepath"
@@ -23,27 +25,36 @@ import (
 	"github.com/agent-substrate/substrate/internal/ateompath"
 )
 
-// pruneLocalCheckpoints removes every local snapshot of the actor.
-// Best-effort: failures are logged, never fatal.
-func pruneLocalCheckpoints(ctx context.Context, actorUID string) {
-	pruneLocalCheckpointDir(ctx, ateompath.LocalCheckpointsDir(actorUID))
+// pruneLocalCheckpoints removes every local snapshot of the actor. A missing
+// directory is not an error, so retries are safe.
+func pruneLocalCheckpoints(ctx context.Context, actorUID string) error {
+	return pruneLocalCheckpointDir(ctx, ateompath.LocalCheckpointsDir(actorUID))
 }
 
-func pruneLocalCheckpointDir(ctx context.Context, dir string) {
+func pruneLocalCheckpointDir(ctx context.Context, dir string) error {
 	entries, err := os.ReadDir(dir)
 	if err != nil {
-		if !os.IsNotExist(err) {
-			slog.WarnContext(ctx, "failed to list local checkpoints for pruning", slog.String("dir", dir), slog.Any("err", err))
+		if os.IsNotExist(err) {
+			return nil
 		}
-		return
+		return fmt.Errorf("while listing local checkpoints in %s: %w", dir, err)
 	}
+	// Every entry is attempted: one undeletable snapshot must not strand the
+	// others on disk.
+	var errs []error
 	for _, entry := range entries {
 		path := filepath.Join(dir, entry.Name())
 		if err := os.RemoveAll(path); err != nil {
-			slog.WarnContext(ctx, "failed to prune local checkpoint", slog.String("path", path), slog.Any("err", err))
+			errs = append(errs, fmt.Errorf("while pruning local checkpoint %s: %w", path, err))
 			continue
 		}
 		slog.InfoContext(ctx, "pruned local checkpoint", slog.String("path", path))
 	}
-	_ = os.Remove(dir)
+	if len(errs) > 0 {
+		return errors.Join(errs...)
+	}
+	if err := os.Remove(dir); err != nil && !os.IsNotExist(err) {
+		return fmt.Errorf("while removing local checkpoints dir %s: %w", dir, err)
+	}
+	return nil
 }

cmd/atelet/local_checkpoints_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 )
 
@@ -38,7 +39,9 @@ func TestPruneRemovesEverySnapshot(t *testing.T) {
 	writeSnapshotDir(t, dir, "pause-2")
 	writeSnapshotDir(t, dir, "pause-3")
 
-	pruneLocalCheckpointDir(context.Background(), dir)
+	if err := pruneLocalCheckpointDir(context.Background(), dir); err != nil {
+		t.Fatalf("pruneLocalCheckpointDir() = %v, want nil", err)
+	}
 
 	if _, err := os.Stat(dir); !os.IsNotExist(err) {
 		t.Fatalf("dir still exists (err=%v), want removed entirely", err)
@@ -46,5 +49,40 @@ func TestPruneRemovesEverySnapshot(t *testing.T) {
 }
 
 func TestPruneMissingDirIsNoop(t *testing.T) {
-	pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent"))
+	if err := pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent")); err != nil {
+		t.Fatalf("pruneLocalCheckpointDir() = %v, want nil", err)
+	}
+}
+
+// An undeletable snapshot must be reported — Terminate turns that error into a
+// failed RPC so the delete workflow retries — without stranding the snapshots
+// that could have been removed.
+func TestPruneReportsFailureAndStillRemovesTheRest(t *testing.T) {
+	if os.Geteuid() == 0 {
+		t.Skip("root ignores directory permissions, so no snapshot can be made undeletable")
+	}
+	dir := t.TempDir()
+	writeSnapshotDir(t, dir, "pause-1")
+	writeSnapshotDir(t, dir, "pause-2")
+
+	// A snapshot dir with no write bit: its files cannot be unlinked.
+	stuck := filepath.Join(dir, "pause-stuck")
+	writeSnapshotDir(t, dir, "pause-stuck")
+	if err := os.Chmod(stuck, 0o500); err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = os.Chmod(stuck, 0o700) })
+
+	err := pruneLocalCheckpointDir(context.Background(), dir)
+	if err == nil {
+		t.Fatal("pruneLocalCheckpointDir() = nil, want an error naming the undeletable snapshot")
+	}
+	if !strings.Contains(err.Error(), "pause-stuck") {
+		t.Errorf("pruneLocalCheckpointDir() = %v, want it to name pause-stuck", err)
+	}
+	for _, name := range []string{"pause-1", "pause-2"} {
+		if _, err := os.Stat(filepath.Join(dir, name)); !os.IsNotExist(err) {
+			t.Errorf("%s still exists (err=%v), want removed", name, err)
+		}
+	}
 }

cmd/atelet/main.go

@@ -660,7 +660,11 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe
 	// all: the actor's current state was just captured by CheckpointWorkload,
 	// and the control plane tracks only a single local snapshot, which this
 	// checkpoint either overwrites (pause) or clears (suspend).
-	pruneLocalCheckpoints(ctx, actorUID)
+	//
+	// Best-effort: if this fail, the actor's terminate prunes again.
+	if err := pruneLocalCheckpoints(ctx, actorUID); err != nil {
+		slog.WarnContext(ctx, "failed to prune superseded local checkpoints", slog.Any("actor", actorRef), slog.Any("err", err))
+	}
 
 	// Pruning stays outside the persist window: it collects superseded
 	// snapshots on both paths, so timing it as part of an external upload would
@@ -842,7 +846,9 @@ func (s *AteomHerder) UploadPausedCheckpoint(ctx context.Context, req *ateletpb.
 
 	// The uploaded snapshot supersedes every local pause snapshot of this
 	// actor; free the node's disk (best-effort, like Checkpoint).
-	pruneLocalCheckpoints(ctx, req.GetActorUid())
+	if err := pruneLocalCheckpoints(ctx, req.GetActorUid()); err != nil {
+		slog.WarnContext(ctx, "failed to prune uploaded local checkpoints", slog.String("actorUID", req.GetActorUid()), slog.Any("err", err))
+	}
 
 	return &ateletpb.UploadPausedCheckpointResponse{}, nil
 }
@@ -1264,6 +1270,15 @@ func (s *AteomHerder) Terminate(ctx context.Context, req *ateletpb.TerminateRequ
 		return nil, fmt.Errorf("failed to unmount external volumes during terminate (actor: %s, actorUID: %s): %w", actorRef, actorUID, err)
 	}
 
+	// The actor is gone, so no pause snapshot of it can ever be restored again.
+	// TODO(#664): this only removes local snapshots in one node. We should clean
+	// up the copies on any other NodeVmsWithLocalSnapshots. This is fine *as of
+	// the day this was written* because today NodeVmsWithLocalSnapshots has at
+	// most one item.
+	if err := pruneLocalCheckpoints(ctx, actorUID); err != nil {
+		return nil, fmt.Errorf("failed to prune local checkpoints during terminate (actor: %s, actorUID: %s): %w", actorRef, actorUID, err)
+	}
+
 	// Reset actor directories on the node
 	if err := resetActorDirs(actorUID); err != nil {
 		return nil, fmt.Errorf("failed to reset actor directories during terminate (actor: %s, actorUID: %s): %w", actorRef, actorUID, err)
@@ -1812,6 +1827,11 @@ func newAteomDialer(size int) *AteomDialer {
 	}
 }
 
+// ateomSocketPath resolves a pod UID to the ateom socket atelet dials. A
+// variable because the real path is rooted at the node's BasePath, which a
+// test cannot serve on.
+var ateomSocketPath = ateompath.AteomSocketPath
+
 func (d *AteomDialer) DialAteomPod(ctx context.Context, podUID string) (*grpc.ClientConn, error) {
 	key := podUID
 
@@ -1821,7 +1841,7 @@ func (d *AteomDialer) DialAteomPod(ctx context.Context, podUID string) (*grpc.Cl
 	}
 
 	conn, err := grpc.NewClient(
-		"unix://"+ateompath.AteomSocketPath(podUID),
+		"unix://"+ateomSocketPath(podUID),
 		grpc.WithTransportCredentials(insecure.NewCredentials()),
 		grpc.WithStatsHandler(otelgrpc.NewClientHandler()),
 	)