From e2568e9ad7c63220f3895d9df8ea3f934fa405df Mon Sep 17 00:00:00 2001 From: NekoPunch Date: Sun, 2 Aug 2026 17:04:16 -0700 Subject: [PATCH] atelet: prune local pause snapshots Pause snapshots accumulated until the node disk filled; the control plane only ever references the latest one. Pause now keeps just the snapshot it wrote, suspend removes them all. Snapshot prefixes are validated as single path segments at the RPC boundary: a nested prefix would nest on write but be pruned as its top-level directory, deleting the fresh snapshot. --- cmd/atelet/local_checkpoints.go | 54 +++++++++++++++++++ cmd/atelet/local_checkpoints_test.go | 78 ++++++++++++++++++++++++++++ cmd/atelet/main.go | 12 +++-- cmd/atelet/main_test.go | 16 ++++++ internal/resources/validate.go | 14 +++++ internal/resources/validate_test.go | 25 +++++++++ 6 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 cmd/atelet/local_checkpoints.go create mode 100644 cmd/atelet/local_checkpoints_test.go diff --git a/cmd/atelet/local_checkpoints.go b/cmd/atelet/local_checkpoints.go new file mode 100644 index 000000000..5426cbc2b --- /dev/null +++ b/cmd/atelet/local_checkpoints.go @@ -0,0 +1,54 @@ +// 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" + "log/slog" + "os" + "path/filepath" + + "github.com/agent-substrate/substrate/internal/ateompath" +) + +// pruneLocalCheckpoints removes the actor's local pause snapshots except keep; +// empty keep removes all. Best-effort: failures are logged, never fatal. +func pruneLocalCheckpoints(ctx context.Context, actorUID, keep string) { + pruneLocalCheckpointDir(ctx, ateompath.LocalCheckpointsDir(actorUID), keep) +} + +func pruneLocalCheckpointDir(ctx context.Context, dir, keep string) { + 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)) + } + return + } + for _, entry := range entries { + if keep != "" && entry.Name() == keep { + continue + } + 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)) + continue + } + slog.InfoContext(ctx, "pruned local checkpoint", slog.String("path", path)) + } + if keep == "" { + _ = os.Remove(dir) + } +} diff --git a/cmd/atelet/local_checkpoints_test.go b/cmd/atelet/local_checkpoints_test.go new file mode 100644 index 000000000..c7503559f --- /dev/null +++ b/cmd/atelet/local_checkpoints_test.go @@ -0,0 +1,78 @@ +// 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" + "os" + "path/filepath" + "testing" +) + +func writeSnapshotDir(t *testing.T, dir, prefix string) { + t.Helper() + p := filepath.Join(dir, prefix) + if err := os.MkdirAll(p, 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(p, "memory.img"), []byte("x"), 0o600); err != nil { + t.Fatal(err) + } +} + +func listDir(t *testing.T, dir string) []string { + t.Helper() + entries, err := os.ReadDir(dir) + if os.IsNotExist(err) { + return nil + } + if err != nil { + t.Fatal(err) + } + var names []string + for _, e := range entries { + names = append(names, e.Name()) + } + return names +} + +func TestPruneKeepsOnlyNamedSnapshot(t *testing.T) { + dir := t.TempDir() + writeSnapshotDir(t, dir, "pause-1") + writeSnapshotDir(t, dir, "pause-2") + writeSnapshotDir(t, dir, "pause-3") + + pruneLocalCheckpointDir(context.Background(), dir, "pause-3") + + if got := listDir(t, dir); len(got) != 1 || got[0] != "pause-3" { + t.Fatalf("dir = %v, want only pause-3", got) + } +} + +func TestPruneEmptyKeepRemovesEverything(t *testing.T) { + dir := t.TempDir() + writeSnapshotDir(t, dir, "pause-1") + writeSnapshotDir(t, dir, "pause-2") + + pruneLocalCheckpointDir(context.Background(), dir, "") + + if _, err := os.Stat(dir); !os.IsNotExist(err) { + t.Fatalf("dir still exists (err=%v), want removed entirely", err) + } +} + +func TestPruneMissingDirIsNoop(t *testing.T) { + pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent"), "anything") +} diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 7b5464699..61e7a9e77 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -398,10 +398,14 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe if err := s.uploadExternalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while uploading external snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err)) } + // The durable snapshot supersedes local pause snapshots. + pruneLocalCheckpoints(ctx, actorUID, "") case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: if err := s.moveLocalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while moving to local snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err)) } + // Without pruning, every pause leaves a snapshot behind until the disk fills. + pruneLocalCheckpoints(ctx, actorUID, req.GetLocalConfig().GetSnapshotPrefix()) default: return nil, fmt.Errorf("unexpected checkpoint type: %v", req.GetType()) } @@ -1049,8 +1053,8 @@ func validateCheckpointRequest(req *ateletpb.CheckpointRequest) error { return err } case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: - if req.GetLocalConfig().GetSnapshotPrefix() == "" { - return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String()) + if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil { + return err } default: return fmt.Errorf("invalid checkpoint type: %v", req.GetType()) @@ -1101,8 +1105,8 @@ func validateRestoreRequest(req *ateletpb.RestoreRequest) error { return err } case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: - if req.GetLocalConfig().GetSnapshotPrefix() == "" { - return fmt.Errorf("snapshot prefix must be non-empty for type %s", req.GetType().String()) + if err := resources.ValidateLocalSnapshotPrefix(req.GetLocalConfig().GetSnapshotPrefix()); err != nil { + return err } default: return fmt.Errorf("invalid checkpoint type: %v", req.GetType()) diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index 0378c71c2..fedfb9ef1 100644 --- a/cmd/atelet/main_test.go +++ b/cmd/atelet/main_test.go @@ -277,6 +277,14 @@ func TestValidateCheckpointRequest(t *testing.T) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}} }), true}, + {"nested local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}} + }), true}, + {"traversal local snapshot prefix", makeReq(func(r *ateletpb.CheckpointRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.CheckpointRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}} + }), true}, {"unspecified snapshot type", makeReq(func(r *ateletpb.CheckpointRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true}, {"unspecified snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true}, {"invalid snapshot scope", makeReq(func(r *ateletpb.CheckpointRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true}, @@ -323,6 +331,14 @@ func TestValidateRestoreRequest(t *testing.T) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ""}} }), true}, + {"nested local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: "pause/2"}} + }), true}, + {"traversal local snapshot prefix", makeReq(func(r *ateletpb.RestoreRequest) { + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.RestoreRequest_LocalConfig{LocalConfig: &ateletpb.LocalCheckpointConfiguration{SnapshotPrefix: ".."}} + }), true}, {"unspecified snapshot type", makeReq(func(r *ateletpb.RestoreRequest) { r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_UNSPECIFIED }), true}, {"unspecified snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED }), true}, {"invalid snapshot scope", makeReq(func(r *ateletpb.RestoreRequest) { r.Scope = ateletpb.SnapshotScope(23) }), true}, diff --git a/internal/resources/validate.go b/internal/resources/validate.go index 09c985df7..302f8dedd 100644 --- a/internal/resources/validate.go +++ b/internal/resources/validate.go @@ -173,6 +173,20 @@ func ValidateSnapshotURIPrefix(prefix string) error { return nil } +// ValidateLocalSnapshotPrefix ensures a local snapshot prefix is a single +// path segment: it is joined onto the actor's checkpoint directory and +// compared against that directory's entries when pruning, so a nested or +// relative prefix would escape the directory or delete the wrong snapshots. +func ValidateLocalSnapshotPrefix(prefix string) error { + if prefix == "" { + return fmt.Errorf("snapshot prefix must be non-empty") + } + if prefix == "." || prefix == ".." || strings.ContainsAny(prefix, `/\`) { + return fmt.Errorf("invalid snapshot prefix %q: must be a single path segment", prefix) + } + return nil +} + // ValidateWorker checks that the worker message is well-formed. func ValidateWorker(worker *ateapipb.Worker, fldPath *field.Path) field.ErrorList { var errs field.ErrorList diff --git a/internal/resources/validate_test.go b/internal/resources/validate_test.go index b25de1d55..f8a7fad32 100644 --- a/internal/resources/validate_test.go +++ b/internal/resources/validate_test.go @@ -238,6 +238,31 @@ func TestValidateSnapshotURIPrefix(t *testing.T) { } } +func TestValidateLocalSnapshotPrefix(t *testing.T) { + tests := []struct { + name string + prefix string + wantErr bool + }{ + {"valid", "pause", false}, + {"valid with dash and digits", "pause-2", false}, + {"empty", "", true}, + {"dot", ".", true}, + {"dotdot", "..", true}, + {"nested", "pause/2", true}, + {"absolute", "/pause", true}, + {"traversal", "../other-actor", true}, + {"backslash", `pause\2`, true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if err := ValidateLocalSnapshotPrefix(tt.prefix); (err != nil) != tt.wantErr { + t.Errorf("ValidateLocalSnapshotPrefix(%q) err = %v, wantErr %v", tt.prefix, err, tt.wantErr) + } + }) + } +} + func TestValidateWorker(t *testing.T) { tests := []struct { name string