diff --git a/cmd/atelet/lifecycle_test.go b/cmd/atelet/lifecycle_test.go new file mode 100644 index 000000000..c285679c4 --- /dev/null +++ b/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) + } +} diff --git a/cmd/atelet/local_checkpoints.go b/cmd/atelet/local_checkpoints.go index b32a4bf15..ca7b96382 100644 --- a/cmd/atelet/local_checkpoints.go +++ b/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 } diff --git a/cmd/atelet/local_checkpoints_test.go b/cmd/atelet/local_checkpoints_test.go index c1a6ebba3..3c5391978 100644 --- a/cmd/atelet/local_checkpoints_test.go +++ b/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) + } + } } diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 3995f5668..7ad044ca6 100644 --- a/cmd/atelet/main.go +++ b/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()), )