diff --git a/cmd/atelet/local_checkpoints.go b/cmd/atelet/local_checkpoints.go index b32a4bf15..6396e6cee 100644 --- a/cmd/atelet/local_checkpoints.go +++ b/cmd/atelet/local_checkpoints.go @@ -23,13 +23,13 @@ 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 the actor's local snapshots, except the one +// named by keep (pass "" to remove them 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 string) { +func pruneLocalCheckpointDir(ctx context.Context, dir, keep string) { entries, err := os.ReadDir(dir) if err != nil { if !os.IsNotExist(err) { @@ -38,6 +38,9 @@ func pruneLocalCheckpointDir(ctx context.Context, dir string) { 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)) @@ -45,5 +48,6 @@ func pruneLocalCheckpointDir(ctx context.Context, dir string) { } slog.InfoContext(ctx, "pruned local checkpoint", slog.String("path", path)) } + // Only removes the directory when it is empty, so a kept snapshot stays. _ = os.Remove(dir) } diff --git a/cmd/atelet/local_checkpoints_test.go b/cmd/atelet/local_checkpoints_test.go index c1a6ebba3..870318cfa 100644 --- a/cmd/atelet/local_checkpoints_test.go +++ b/cmd/atelet/local_checkpoints_test.go @@ -38,7 +38,7 @@ func TestPruneRemovesEverySnapshot(t *testing.T) { writeSnapshotDir(t, dir, "pause-2") writeSnapshotDir(t, dir, "pause-3") - pruneLocalCheckpointDir(context.Background(), dir) + pruneLocalCheckpointDir(context.Background(), dir, "") if _, err := os.Stat(dir); !os.IsNotExist(err) { t.Fatalf("dir still exists (err=%v), want removed entirely", err) @@ -46,5 +46,20 @@ func TestPruneRemovesEverySnapshot(t *testing.T) { } func TestPruneMissingDirIsNoop(t *testing.T) { - pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent")) + pruneLocalCheckpointDir(context.Background(), filepath.Join(t.TempDir(), "absent"), "") +} + +func TestPruneKeepsNamedSnapshot(t *testing.T) { + dir := t.TempDir() + writeSnapshotDir(t, dir, "pause-1") + writeSnapshotDir(t, dir, "pause-2") + + pruneLocalCheckpointDir(context.Background(), dir, "pause-2") + + if _, err := os.Stat(filepath.Join(dir, "pause-1")); !os.IsNotExist(err) { + t.Errorf("pause-1 still exists (err=%v), want pruned", err) + } + if _, err := os.Stat(filepath.Join(dir, "pause-2", "memory.img")); err != nil { + t.Errorf("pause-2 was pruned, want kept: %v", err) + } } diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 3f3cd1fd4..9211279ea 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -576,6 +576,21 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe actorUID := req.GetActorUid() actorRef := resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()} + // Fast-forward if the snapshot was already committed by a previous attempt, + // completing any remaining node teardown. + committed, err := s.checkpointAlreadyCommitted(ctx, req) + if err != nil { + return nil, err + } + if committed { + slog.InfoContext(ctx, "Checkpoint already committed to its destination; finishing its teardown", + slog.Any("actor", actorRef), slog.String("actor_uid", actorUID)) + // Finish teardown in case the previous attempt died before cleanup completed. + if err := s.finishCheckpoint(ctx, actorUID, req.GetSpec().GetVolumes()); err != nil { + return nil, err + } + return &ateletpb.CheckpointResponse{}, nil + } // Per-phase timing, recorded on the way out so a failed checkpoint still // reports the phases it completed. Phases left at zero never ran. tStart := time.Now() @@ -641,8 +656,7 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe }) dAteom = time.Since(tAteom) if err != nil { - // TODO: Ateom should classify checkpoint failures, and set "should-crash" - // in the metadata if the error is not retriable. + // Wrapping preserves any gRPC status ErrorInfo attached by ateom. op.failedPhase = ateattr.SnapshotPhaseAteomCheckpoint return nil, fmt.Errorf("while calling ateom.CheckpointWorkload: %w", err) } @@ -658,11 +672,9 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe sandboxRec.ActorTemplateName = req.GetActorTemplateName() sandboxRec.Scope = ateattr.SnapshotScopeValue(req.GetScope()) - // No earlier pause snapshot can ever be restored again, so remove them - // 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) + // Prune older local snapshots to free disk space, but keep the current + // checkpoint destination if an earlier attempt already moved files there. + pruneLocalCheckpoints(ctx, actorUID, req.GetLocalConfig().GetSnapshotName()) // Pruning stays outside the persist window: it collects superseded // snapshots on both paths, so timing it as part of an external upload would @@ -687,16 +699,25 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe } dPersist = time.Since(tPersist) - if err := s.unmountExternalVolumes(ctx, actorUID, req.GetSpec().GetVolumes()); err != nil { - return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonTerminalFileSystemError, ateerrors.ActorCrashedMetadata(), fmt.Errorf("while unmounting external volumes: %w", err)) + // Assign to named return err so deferred metrics record teardown failures. + if err = s.finishCheckpoint(ctx, actorUID, req.GetSpec().GetVolumes()); err != nil { + return nil, err + } + + return &ateletpb.CheckpointResponse{}, nil +} + +// finishCheckpoint unmounts external volumes and resets on-node actor directories. +func (s *AteomHerder) finishCheckpoint(ctx context.Context, actorUID string, volumes []*ateletpb.Volume) error { + if err := s.unmountExternalVolumes(ctx, actorUID, volumes); err != nil { + return ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonTerminalFileSystemError, ateerrors.ActorCrashedMetadata(), fmt.Errorf("while unmounting external volumes: %w", err)) } // Note: we do not crash the actor if resetting the directory fails. if err := resetActorDirs(actorUID); err != nil { - return nil, fmt.Errorf("while resetting actor dirs: %w", err) + return fmt.Errorf("while resetting actor dirs: %w", err) } - - return &ateletpb.CheckpointResponse{}, nil + return nil } func toAteomSnapshotScope(scope ateletpb.SnapshotScope) ateompb.SnapshotScope { @@ -711,19 +732,93 @@ func toAteomSnapshotScope(scope ateletpb.SnapshotScope) ateompb.SnapshotScope { } } +// checkpointAlreadyCommitted reports whether the requested snapshot is already +// committed to its destination. The manifest acts as the commit marker. +func (s *AteomHerder) checkpointAlreadyCommitted(ctx context.Context, req *ateletpb.CheckpointRequest) (bool, error) { + switch req.GetType() { + case ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL: + uri, err := resources.ParseSnapshotURI(req.GetExternalConfig().GetSnapshotUri()) + if err != nil { + return false, ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonInvalidObjectURL) + } + manifest, uploaded, err := s.fetchUploadedSnapshotManifest(ctx, uri) + if err != nil || !uploaded { + return false, err + } + return committedManifestAnswers(ctx, manifest, req), nil + + case ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL: + path := filepath.Join(ateompath.LocalSnapshotDir(req.GetActorUid(), req.GetLocalConfig().GetSnapshotName()), sandboxManifestName) + manifest, err := os.ReadFile(path) + switch { + case err == nil: + return committedManifestAnswers(ctx, manifest, req), nil + case errors.Is(err, os.ErrNotExist): + return false, nil + default: + return false, wrapFileSystemErr("while probing for an already-written local snapshot manifest", err) + } + + default: + return false, nil + } +} + +// committedManifestAnswers reports whether the snapshot manifest matches the requested scope. +func committedManifestAnswers(ctx context.Context, manifest []byte, req *ateletpb.CheckpointRequest) bool { + rec, err := unmarshalSandboxRecord(manifest) + if err != nil { + slog.WarnContext(ctx, "Snapshot manifest at this checkpoint's destination cannot be parsed; treating the checkpoint as not committed", + slog.String("actor_uid", req.GetActorUid()), slog.Any("err", err)) + return false + } + if want := ateattr.SnapshotScopeValue(req.GetScope()); rec.Scope != want { + slog.WarnContext(ctx, "Snapshot at this checkpoint's destination records a different scope; not treating it as this checkpoint's result", + slog.String("actor_uid", req.GetActorUid()), slog.String("manifest_scope", rec.Scope), slog.String("requested_scope", want)) + return false + } + return true +} + +// fetchUploadedSnapshotManifest returns the snapshot manifest at uri, and whether +// it exists. Missing objects return false without error; storage failures return an error. +func (s *AteomHerder) fetchUploadedSnapshotManifest(ctx context.Context, uri resources.SnapshotURI) ([]byte, bool, error) { + manifestURI, err := uri.ObjectURI(sandboxManifestName) + if err != nil { + return nil, false, ateerrors.CrashIfReason(ctx, fmt.Errorf("while addressing snapshot manifest in GCS: %w", err), ateerrors.ReasonInvalidObjectURL) + } + manifest, err := ategcs.FetchFromGCS(ctx, s.gcsClient, manifestURI) + if err != nil { + if errors.Is(err, ateerrors.ReasonFailedGetExternalObject) { + return nil, false, nil + } + return nil, false, fmt.Errorf("while probing for an already-uploaded snapshot manifest: %w", err) + } + return manifest, true, nil +} + func (s *AteomHerder) moveLocalCheckpoint(ctx context.Context, req *ateletpb.CheckpointRequest, checkpointDir string, rec *sandboxAssetsRecord) error { localCheckpointPath := ateompath.LocalSnapshotDir(req.GetActorUid(), req.GetLocalConfig().GetSnapshotName()) if err := os.MkdirAll(localCheckpointPath, 0o700); err != nil { return fmt.Errorf("while creating local checkpoint directory: %w", err) } - // Move exactly the files ateom reported. + // Move files reported by ateom. If a file is already at dst, it was moved + // by an earlier interrupted attempt. for _, fileName := range rec.SnapshotFiles { src := filepath.Join(checkpointDir, fileName) dst := filepath.Join(localCheckpointPath, fileName) recordSnapshotSize(ctx, fileName, src, req.GetActorTemplateNamespace(), req.GetActorTemplateName()) - if err := os.Rename(src, dst); err != nil { + err := os.Rename(src, dst) + if errors.Is(err, os.ErrNotExist) { + if _, statErr := os.Stat(dst); statErr == nil { + continue + } + // Gone from both sides: the snapshot cannot be assembled. + return wrapFileSystemErr(fmt.Sprintf("snapshot file %q is missing from both %s and %s", fileName, checkpointDir, localCheckpointPath), err) + } + if err != nil { return fmt.Errorf("failed to move %s to %s: %w", src, dst, err) } } @@ -733,7 +828,7 @@ func (s *AteomHerder) moveLocalCheckpoint(ctx context.Context, req *ateletpb.Che if err != nil { return fmt.Errorf("while marshaling snapshot manifest: %w", err) } - if err := os.WriteFile(filepath.Join(localCheckpointPath, sandboxManifestName), manifest, 0o600); err != nil { + if err := writeFileAtomic(filepath.Join(localCheckpointPath, sandboxManifestName), manifest, 0o600); err != nil { return fmt.Errorf("while writing snapshot manifest: %w", err) } @@ -844,7 +939,7 @@ 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()) + pruneLocalCheckpoints(ctx, req.GetActorUid(), "") return &ateletpb.UploadPausedCheckpointResponse{}, nil } @@ -854,11 +949,6 @@ func (s *AteomHerder) UploadPausedCheckpoint(ctx context.Context, req *ateletpb. // returns the sandbox class recorded in the snapshot manifest (empty when the // manifest was not read). Parameterized by localDir for tests. func (s *AteomHerder) uploadLocalCheckpointDir(ctx context.Context, req *ateletpb.UploadPausedCheckpointRequest, localDir string, uri resources.SnapshotURI) (string, error) { - manifestURI, err := uri.ObjectURI(sandboxManifestName) - if err != nil { - return "", fmt.Errorf("while addressing snapshot manifest in GCS: %w", err) - } - manifest, err := os.ReadFile(filepath.Join(localDir, sandboxManifestName)) if errors.Is(err, os.ErrNotExist) { // The local snapshot is gone. A previous invocation may have uploaded @@ -866,16 +956,16 @@ func (s *AteomHerder) uploadLocalCheckpointDir(ctx context.Context, req *ateletp // means the whole snapshot is committed and this retry already // succeeded. Absent on both sides, the paused actor's state is // unrecoverable. - _, fetchErr := ategcs.FetchFromGCS(ctx, s.gcsClient, manifestURI) - if fetchErr == nil { + _, uploaded, probeErr := s.fetchUploadedSnapshotManifest(ctx, uri) + if probeErr != nil { + return "", probeErr + } + if uploaded { slog.InfoContext(ctx, "Local snapshot already uploaded and pruned; nothing to do", slog.String("snapshot_uri", req.GetDestinationSnapshotUri())) return "", nil } - if errors.Is(fetchErr, ateerrors.ReasonFailedGetExternalObject) { - return "", ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonLocalSnapshotGone, ateerrors.ActorCrashedMetadata(), - fmt.Errorf("local snapshot %q is gone and no uploaded copy exists: %w", req.GetLocalSnapshotName(), fetchErr)) - } - return "", fmt.Errorf("while probing for an already-uploaded snapshot manifest: %w", fetchErr) + return "", ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonLocalSnapshotGone, ateerrors.ActorCrashedMetadata(), + fmt.Errorf("%w: local snapshot %q is gone and no uploaded copy exists", ateerrors.ReasonLocalSnapshotGone, req.GetLocalSnapshotName())) } if err != nil { return "", wrapFileSystemErr("while reading local snapshot manifest", err) diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index 2897eca67..a79b28e18 100644 --- a/cmd/atelet/main_test.go +++ b/cmd/atelet/main_test.go @@ -41,6 +41,7 @@ import ( "github.com/agent-substrate/substrate/internal/proto/ateompb" "github.com/agent-substrate/substrate/internal/resources" "github.com/agent-substrate/substrate/internal/serverboot" + "github.com/agent-substrate/substrate/internal/volume" "github.com/google/go-cmp/cmp" "github.com/klauspost/compress/zstd" "github.com/spf13/pflag" @@ -418,6 +419,17 @@ func validCheckpointRequest() *ateletpb.CheckpointRequest { } } +func validLocalCheckpointRequest(snapshotName string) *ateletpb.CheckpointRequest { + r := validCheckpointRequest() + r.Type = ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL + r.Config = &ateletpb.CheckpointRequest_LocalConfig{ + LocalConfig: &ateletpb.LocalCheckpointConfiguration{ + SnapshotName: snapshotName, + }, + } + return r +} + func validRestoreRequest() *ateletpb.RestoreRequest { return &ateletpb.RestoreRequest{ Atespace: "ate-demo", @@ -1570,11 +1582,15 @@ type recordingObjectStorage struct { mu sync.Mutex objects map[string][]byte putErr error + getErr error } func (r *recordingObjectStorage) GetObject(_ context.Context, bucket, object string) (io.ReadCloser, error) { r.mu.Lock() defer r.mu.Unlock() + if r.getErr != nil { + return nil, r.getErr + } b, ok := r.objects[bucket+"/"+object] if !ok { return nil, fmt.Errorf("%w: Bucket:%q, Object:%q", ateerrors.ReasonFailedGetExternalObject, bucket, object) @@ -1966,3 +1982,339 @@ func TestShouldHaveSnapshots(t *testing.T) { }) } } + +func useTempActorsDir(t *testing.T) { + t.Helper() + orig := ateompath.ActorsDir + t.Cleanup(func() { ateompath.ActorsDir = orig }) + ateompath.ActorsDir = t.TempDir() +} + +func TestCheckpointAlreadyCommitted(t *testing.T) { + ctx := context.Background() + const manifestKey = "bucket/root/snapshots/ate-demo/counter-1-snap/manifest.json" + + tests := []struct { + name string + req *ateletpb.CheckpointRequest + setup func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder + wantCommitted bool + wantErr bool + errTarget error + }{ + { + name: "external with an uploaded manifest", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{ + objects: map[string][]byte{manifestKey: []byte(`{"pauseImage":"pause:v1","scope":"full"}`)}, + }} + }, + wantCommitted: true, + }, + { + name: "external manifest recording a different scope", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{ + objects: map[string][]byte{manifestKey: []byte(`{"pauseImage":"pause:v1","scope":"data"}`)}, + }} + }, + wantCommitted: false, + }, + { + name: "external manifest with no scope recorded", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{ + objects: map[string][]byte{manifestKey: []byte(`{"pauseImage":"pause:v1"}`)}, + }} + }, + wantCommitted: false, + }, + { + name: "external manifest that cannot be parsed", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{ + objects: map[string][]byte{manifestKey: []byte("not json")}, + }} + }, + wantCommitted: false, + }, + { + name: "external with no manifest", + req: validCheckpointRequest(), + wantCommitted: false, + }, + { + name: "external probe failure is not read as uncommitted", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{getErr: errors.New("bucket unreachable")}} + }, + wantErr: true, + }, + { + name: "external with invalid snapshot URI", + req: func() *ateletpb.CheckpointRequest { + r := validCheckpointRequest() + r.Config = &ateletpb.CheckpointRequest_ExternalConfig{ + ExternalConfig: &ateletpb.ExternalCheckpointConfiguration{SnapshotUri: "invalid-uri"}, + } + return r + }(), + wantErr: true, + }, + { + name: "local with a written manifest", + req: validLocalCheckpointRequest("pause-snap-1"), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + writeLocalSnapshot(t, ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1"), + sandboxAssetsRecord{SandboxClass: "gvisor", PauseImage: testPauseImage, SnapshotFiles: []string{"checkpoint.img"}, Scope: ateattr.SnapshotScopeFull}, + map[string]string{"checkpoint.img": "img"}) + return &AteomHerder{} + }, + wantCommitted: true, + }, + { + name: "local snapshot recording a different scope", + req: validLocalCheckpointRequest("pause-snap-1"), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + writeLocalSnapshot(t, ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1"), + sandboxAssetsRecord{SandboxClass: "gvisor", PauseImage: testPauseImage, SnapshotFiles: []string{"durable-dir.tar"}, Scope: ateattr.SnapshotScopeData}, + map[string]string{"durable-dir.tar": "tar"}) + return &AteomHerder{} + }, + wantCommitted: false, + }, + { + name: "local with unparseable manifest", + req: validLocalCheckpointRequest("pause-snap-1"), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + dir := ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1") + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatalf("creating snapshot dir: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, sandboxManifestName), []byte("not json"), 0o600); err != nil { + t.Fatalf("writing manifest: %v", err) + } + return &AteomHerder{} + }, + wantCommitted: false, + }, + { + name: "local with manifest read error", + req: validLocalCheckpointRequest("pause-snap-1"), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + manifestDir := filepath.Join(ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1"), sandboxManifestName) + if err := os.MkdirAll(manifestDir, 0o700); err != nil { + t.Fatalf("creating manifest dir: %v", err) + } + return &AteomHerder{} + }, + wantErr: true, + errTarget: ateerrors.ReasonTerminalFileSystemError, + }, + { + name: "local with no snapshot dir", + req: validLocalCheckpointRequest("pause-snap-1"), + wantCommitted: false, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + useTempActorsDir(t) + s := &AteomHerder{gcsClient: &recordingObjectStorage{}} + if tc.setup != nil { + s = tc.setup(t, tc.req) + } + + got, err := s.checkpointAlreadyCommitted(ctx, tc.req) + if (err != nil) != tc.wantErr { + t.Fatalf("checkpointAlreadyCommitted err = %v, wantErr %v", err, tc.wantErr) + } + if tc.errTarget != nil && !errors.Is(err, tc.errTarget) { + t.Errorf("err = %v, want it tagged %v", err, tc.errTarget) + } + if got != tc.wantCommitted { + t.Errorf("committed = %v, want %v", got, tc.wantCommitted) + } + }) + } +} + +func TestCheckpointFastForwardsWhenAlreadyCommitted(t *testing.T) { + for _, tc := range []struct { + name string + req *ateletpb.CheckpointRequest + setup func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder + }{ + { + name: "external snapshot", + req: validCheckpointRequest(), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + return &AteomHerder{gcsClient: &recordingObjectStorage{ + objects: map[string][]byte{ + "bucket/root/snapshots/ate-demo/counter-1-snap/manifest.json": []byte(`{"pauseImage":"pause:v1","scope":"full"}`), + }, + }} + }, + }, + { + name: "local snapshot", + req: validLocalCheckpointRequest("pause-snap-1"), + setup: func(t *testing.T, req *ateletpb.CheckpointRequest) *AteomHerder { + writeLocalSnapshot(t, ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1"), + sandboxAssetsRecord{SandboxClass: "gvisor", PauseImage: testPauseImage, SnapshotFiles: []string{"checkpoint.img"}, Scope: ateattr.SnapshotScopeFull}, + map[string]string{"checkpoint.img": "img"}) + return &AteomHerder{} + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + useTempActorsDir(t) + req := tc.req + s := tc.setup(t, req) + + fakePlugin := &fakeWorkerPlugin{} + s.volumePlugins = map[string]volume.VolumePluginWorkerPlane{ + "mock-driver": fakePlugin, + } + req.Spec.Volumes = []*ateletpb.Volume{{ + Name: "vol-1", + Source: &ateletpb.Volume_External{ + External: &ateletpb.ExternalVolumeSource{ + StorageVolumeId: "mock-vol-1", + VolumeType: "mock-driver", + }, + }, + }} + + bundleDir := ateompath.OCIBundleDir(req.GetActorUid()) + if err := os.MkdirAll(bundleDir, 0o700); err != nil { + t.Fatalf("creating bundle dir: %v", err) + } + leftover := filepath.Join(bundleDir, "leftover") + if err := os.WriteFile(leftover, []byte("x"), 0o600); err != nil { + t.Fatalf("writing leftover: %v", err) + } + + resp, err := s.Checkpoint(context.Background(), req) + if err != nil { + t.Fatalf("Checkpoint: %v", err) + } + if resp == nil { + t.Fatal("Checkpoint returned a nil response") + } + + // Verify that fast-forwarding still resets actor directories and unmounts external volumes. + if _, err := os.Stat(leftover); !os.IsNotExist(err) { + t.Errorf("bundle dir still populated (err=%v), want the checkpoint teardown to have reset it", err) + } + if len(fakePlugin.unmounted) != 1 || fakePlugin.unmounted[0] != "mock-vol-1" { + t.Errorf("unmounted volumes = %v, want [mock-vol-1]", fakePlugin.unmounted) + } + }) + } +} + +func TestMoveLocalCheckpointResumesPartialMove(t *testing.T) { + ctx := context.Background() + useTempActorsDir(t) + + req := validLocalCheckpointRequest("pause-snap-1") + rec := &sandboxAssetsRecord{ + SandboxClass: "gvisor", + PauseImage: testPauseImage, + SnapshotFiles: []string{"checkpoint.img", "pages.img"}, + } + + checkpointDir := ateompath.CheckpointStateDir(req.GetActorUid()) + dstDir := ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1") + for dir, files := range map[string]map[string]string{ + checkpointDir: {"pages.img": "pages"}, + dstDir: {"checkpoint.img": "img"}, + } { + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatalf("creating %s: %v", dir, err) + } + for name, body := range files { + if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0o600); err != nil { + t.Fatalf("writing %s: %v", name, err) + } + } + } + + if err := (&AteomHerder{}).moveLocalCheckpoint(ctx, req, checkpointDir, rec); err != nil { + t.Fatalf("moveLocalCheckpoint: %v", err) + } + + for _, name := range append(rec.SnapshotFiles, sandboxManifestName) { + if _, err := os.Stat(filepath.Join(dstDir, name)); err != nil { + t.Errorf("%s missing from the snapshot dir: %v", name, err) + } + } +} + +func TestMoveLocalCheckpointLeavesOnlyTheCommittedSnapshot(t *testing.T) { + useTempActorsDir(t) + + req := validLocalCheckpointRequest("pause-snap-1") + rec := &sandboxAssetsRecord{SandboxClass: "gvisor", PauseImage: testPauseImage, SnapshotFiles: []string{"checkpoint.img"}} + + checkpointDir := ateompath.CheckpointStateDir(req.GetActorUid()) + if err := os.MkdirAll(checkpointDir, 0o700); err != nil { + t.Fatalf("creating checkpoint dir: %v", err) + } + if err := os.WriteFile(filepath.Join(checkpointDir, "checkpoint.img"), []byte("img"), 0o600); err != nil { + t.Fatalf("writing checkpoint.img: %v", err) + } + + if err := (&AteomHerder{}).moveLocalCheckpoint(context.Background(), req, checkpointDir, rec); err != nil { + t.Fatalf("moveLocalCheckpoint: %v", err) + } + + dstDir := ateompath.LocalSnapshotDir(req.GetActorUid(), "pause-snap-1") + entries, err := os.ReadDir(dstDir) + if err != nil { + t.Fatalf("reading snapshot dir: %v", err) + } + var got []string + for _, e := range entries { + got = append(got, e.Name()) + } + slices.Sort(got) + want := []string{"checkpoint.img", sandboxManifestName} + if !slices.Equal(got, want) { + t.Errorf("snapshot dir = %v, want exactly %v", got, want) + } + + manifest, err := os.ReadFile(filepath.Join(dstDir, sandboxManifestName)) + if err != nil { + t.Fatalf("reading manifest: %v", err) + } + if _, err := unmarshalSandboxRecord(manifest); err != nil { + t.Errorf("unmarshalSandboxRecord: %v, want the committed manifest to parse", err) + } +} + +func TestMoveLocalCheckpointFailsWhenFileGoneFromBothSides(t *testing.T) { + useTempActorsDir(t) + + req := validLocalCheckpointRequest("pause-snap-1") + checkpointDir := ateompath.CheckpointStateDir(req.GetActorUid()) + if err := os.MkdirAll(checkpointDir, 0o700); err != nil { + t.Fatalf("creating checkpoint dir: %v", err) + } + + rec := &sandboxAssetsRecord{SandboxClass: "gvisor", PauseImage: testPauseImage, SnapshotFiles: []string{"checkpoint.img"}} + err := (&AteomHerder{}).moveLocalCheckpoint(context.Background(), req, checkpointDir, rec) + if err == nil { + t.Fatal("moveLocalCheckpoint succeeded, want a failure: the snapshot cannot be assembled") + } + if !errors.Is(err, ateerrors.ReasonTerminalFileSystemError) { + t.Errorf("err = %v, want it tagged %v", err, ateerrors.ReasonTerminalFileSystemError) + } +} diff --git a/cmd/ateom-gvisor/checkpoint.go b/cmd/ateom-gvisor/checkpoint.go new file mode 100644 index 000000000..c6f0eee34 --- /dev/null +++ b/cmd/ateom-gvisor/checkpoint.go @@ -0,0 +1,225 @@ +//go:build linux + +// 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" + "fmt" + "log/slog" + "os" + "sort" + "strings" + "time" + + "github.com/agent-substrate/substrate/internal/ateerrors" + "github.com/agent-substrate/substrate/internal/ateompath" + "github.com/agent-substrate/substrate/internal/ateomstats" + "github.com/agent-substrate/substrate/internal/checkpointmarker" + "github.com/agent-substrate/substrate/internal/proto/ateompb" + "google.golang.org/grpc/codes" +) + +// Allow checkpointing even if the pod is shutting down. This will allow actors +// (or the harness) to suspend on shutdown. +func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.CheckpointWorkloadRequest) (*ateompb.CheckpointWorkloadResponse, error) { + s.lock.Lock() + defer s.lock.Unlock() + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + s.setActiveRPC(rpcCheckpointWorkload, cancel) + defer s.clearActiveRPC() + + attribution := ateomstats.ActorAttributionFromRequest(req) + + // Replay a previously completed checkpoint for this actor if available. + if rec, ok, err := checkpointmarker.Read(req.GetActorUid(), req.GetScope().String()); err != nil { + return nil, err + } else if ok { + slog.InfoContext(ctx, "Checkpoint already completed for this actor; replaying its result", + "actor", attribution.Ref, + "actorUID", req.GetActorUid(), + "snapshotFiles", rec.SnapshotFiles) + // Finish any pending workload termination, unless the ateom now holds a different actor. + if held := s.activeActor.Load(); held != nil && held.UID != req.GetActorUid() { + slog.WarnContext(ctx, "Not running the post-checkpoint teardown: this ateom now holds a different actor", + slog.String("id", req.GetActorUid()), slog.String("active_actor_uid", held.UID)) + return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil + } + if err := s.terminateWorkload(ctx, attribution.Ref, req.GetActorUid(), req.GetRunscPath(), req.GetSpec().GetContainers()); err != nil { + slog.WarnContext(ctx, "Failed to terminate workload while replaying checkpoint", + slog.String("actorUID", req.GetActorUid()), slog.Any("err", err)) + } + s.activeSession = nil + return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil + } + + if err := s.deactivateActorNetworking(ctx); err != nil { + return nil, err + } + + s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointing", attribution) + + // Contract with atelet: + // + // * After we exit, atelet will upload checkpoint to GCS + // * After we exit, atelet will tear down OCI bundles and reset the actor directory. + + // Checkpoint only saves state; no sizing is applied, so size is left zero. + rcmd := &runsc{ + path: req.GetRunscPath(), + actorUID: req.GetActorUid(), + } + + checkpointPath := ateompath.CheckpointStateDir(req.GetActorUid()) + // Start from a clean directory so retried attempts do not mix with stale + // or partially-written snapshot files from previous runs. + if err := os.RemoveAll(checkpointPath); err != nil { + return nil, fmt.Errorf("while clearing checkpoint directory: %w", err) + } + if err := os.MkdirAll(checkpointPath, 0o700); err != nil { + return nil, fmt.Errorf("while creating checkpoint directory: %w", err) + } + + // Always take durable-dir snapshot if at least one container has a durable-dir volume mount. + // TODO(dberkov): this is a temporary workaround until gVisor supports taking durable-dir snapshots in a single request with the process snapshot. + switch req.GetScope() { + case ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA: + var ddv []string + for _, ctr := range req.GetSpec().GetContainers() { + for _, m := range ctr.GetDurableDirVolumeMounts() { + ddv = append(ddv, m.GetMountPath()) + } + } + if len(ddv) == 0 { + return nil, fmt.Errorf("no durable-dir volumes found for DATA snapshot") + } + if err := rcmd.cmdFsCheckpoint(ctx, "pause", checkpointPath, ddv); err != nil { + return nil, classifyCheckpointFailure(ctx, rcmd, fmt.Errorf("while fscheckpointing durable-dir %q: %w", ddv[0], err)) + } + case ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL: + // Checkpoint pause container (root of the sandbox) + if err := rcmd.cmdCheckpoint(ctx, "pause", checkpointPath); err != nil { + return nil, classifyCheckpointFailure(ctx, rcmd, fmt.Errorf("while checkpointing pause: %w", err)) + } + default: + return nil, fmt.Errorf("unsupported snapshot scope: %v", req.GetScope()) + } + + // The sandbox is gone as of the checkpoint above, so the ateom is back to + // "available" from here on: there is nothing left to measure, and holding + // the attribution would let a later GetWorkloadStats report a checkpointed + // actor as though it were still running. + // + // Cleared here rather than at the end of the function because everything + // below is bookkeeping over a dead sandbox and can still fail (listing the + // snapshot files returns an error), which would otherwise leave the + // attribution behind. Conversely nothing above this point clears it: a + // checkpoint that failed may well have left the workload running, and + // reporting its usage is then the honest answer. + s.activeActor.Store(nil) + + // Report exactly the files runsc wrote so atelet ships precisely this set + // (checkpoint.img plus any pages images), rather than a hardcoded list. + snapshotFiles, err := listSnapshotFiles(checkpointPath) + if err != nil { + return nil, fmt.Errorf("while listing checkpoint files: %w", err) + } + + // Record checkpoint completion before answering. If writing the marker fails, + // log and continue since the snapshot files are already complete on disk. + if err := checkpointmarker.Write(req.GetActorUid(), req.GetScope().String(), snapshotFiles); err != nil { + slog.ErrorContext(ctx, "Failed to record the checkpoint completion marker; answering anyway, but a lost response can no longer be replayed", + "actor", attribution.Ref, + "actorUID", req.GetActorUid(), + "snapshotFiles", snapshotFiles, + "err", err) + } + + // Cleanup the containers after checkpointing. + // This is best-effort cleanup for actor containers that may have been left behind after checkpointing. + if err := s.terminateWorkload(ctx, attribution.Ref, attribution.UID, req.GetRunscPath(), req.GetSpec().GetContainers()); err != nil { + slog.WarnContext(ctx, "failed to terminate workload after checkpoint", + slog.String("actor", attribution.Ref.String()), + slog.String("actorUID", attribution.UID), + slog.Any("err", err)) + } + + s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointed", attribution) + s.activeSession = nil + + return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles}, nil +} + +// stateProbeTimeout bounds probing runsc state during failure classification. +const stateProbeTimeout = 15 * time.Second + +// classifyCheckpointFailure inspects runsc container state after a failure to +// distinguish retriable transient errors from unrecoverable errors (where the +// sandbox was destroyed and cannot be retried). +func classifyCheckpointFailure(ctx context.Context, rcmd *runsc, err error) error { + // Probe with a separate timeout so an expired caller context is not + // mistaken for a missing sandbox. + probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), stateProbeTimeout) + defer cancel() + + out, stateErr := rcmd.cmdStateOutput(probeCtx, "pause") + if stateErr == nil { + return err + } + if !sandboxNotFound(out) { + slog.WarnContext(ctx, "Checkpoint failed and the sandbox state could not be determined; leaving the failure retriable", + "actorUID", rcmd.actorUID, "stateErr", stateErr, "runscOutput", string(out), "err", err) + return err + } + slog.WarnContext(ctx, "Checkpoint failed and the sandbox is gone; the actor's state is unrecoverable", + "actorUID", rcmd.actorUID, "stateErr", stateErr, "err", err) + return ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonInvalidCheckpointResult, ateerrors.ActorCrashedMetadata(), + fmt.Errorf("%w: checkpoint failed and no sandbox remains to retry against: %w", ateerrors.ReasonInvalidCheckpointResult, err)) +} + +// sandboxNotFound checks if runsc output explicitly indicates the container does not exist. +func sandboxNotFound(runscOutput []byte) bool { + for line := range strings.Lines(string(runscOutput)) { + msg, ok := strings.CutPrefix(strings.TrimSpace(line), "error:") + if !ok { + continue + } + if strings.Contains(strings.ToLower(msg), "does not exist") { + return true + } + } + return false +} + +// listSnapshotFiles returns the (relative) names of regular files directly under dir. +func listSnapshotFiles(dir string) ([]string, error) { + entries, err := os.ReadDir(dir) + if err != nil { + return nil, err + } + var files []string + for _, e := range entries { + // ateom's own completion marker shares the directory but is + // bookkeeping, not snapshot content, so it never joins the set. + if e.Type().IsRegular() && e.Name() != ateompath.CheckpointDoneFileName { + files = append(files, e.Name()) + } + } + sort.Strings(files) + return files, nil +} diff --git a/cmd/ateom-gvisor/checkpoint_test.go b/cmd/ateom-gvisor/checkpoint_test.go new file mode 100644 index 000000000..96cb4b712 --- /dev/null +++ b/cmd/ateom-gvisor/checkpoint_test.go @@ -0,0 +1,284 @@ +//go:build linux + +// 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" + "errors" + "io" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "github.com/agent-substrate/substrate/internal/actorlog" + "github.com/agent-substrate/substrate/internal/ateompath" + "github.com/agent-substrate/substrate/internal/atunnel" + "github.com/agent-substrate/substrate/internal/checkpointmarker" + "github.com/agent-substrate/substrate/internal/proto/ateompb" + "github.com/agent-substrate/substrate/internal/resources" + "github.com/agent-substrate/substrate/internal/roottest" +) + +// testScope stands in for a CheckpointWorkloadRequest's stringified scope. +var testScope = ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL.String() + +// newCheckpointTestService builds a minimal AteomService for testing pre-sandbox checkpoint logic. +func newCheckpointTestService() *AteomService { + return &AteomService{ + lock: newCancelableMutex(), + atunnelIngress: &atunnel.Server{}, + atunnelEgress: &atunnel.Egress{}, + actorLogger: actorlog.NewActorLogger(io.Discard, false), + } +} + +// useTempActorsDir points the shared actor-state root at a temp directory and +// creates the actor's checkpoint dir. +func useTempActorsDir(t *testing.T, actorUID string) string { + t.Helper() + orig := ateompath.ActorsDir + t.Cleanup(func() { ateompath.ActorsDir = orig }) + ateompath.ActorsDir = t.TempDir() + + dir := ateompath.CheckpointStateDir(actorUID) + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatalf("creating checkpoint dir: %v", err) + } + return dir +} + +func TestListSnapshotFilesExcludesCompletionMarker(t *testing.T) { + const actorUID = "actor-1" + dir := useTempActorsDir(t, actorUID) + + for _, name := range []string{"checkpoint.img", "pages.img"} { + if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0o600); err != nil { + t.Fatalf("writing %s: %v", name, err) + } + } + if err := checkpointmarker.Write(actorUID, testScope, []string{"checkpoint.img", "pages.img"}); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + got, err := listSnapshotFiles(dir) + if err != nil { + t.Fatalf("listSnapshotFiles: %v", err) + } + want := []string{"checkpoint.img", "pages.img"} + if !slices.Equal(got, want) { + t.Errorf("listSnapshotFiles = %v, want %v", got, want) + } +} + +// If the worker has already been reassigned to a successor actor, replaying +// a completed checkpoint must skip teardown to avoid disrupting the new actor. +func TestCheckpointWorkloadReplaySkipsTeardownForAReassignedAteom(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + want := []string{"checkpoint.img", "pages.img"} + if err := checkpointmarker.Write(actorUID, testScope, want); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + // The ateom has moved on: it now runs a different actor. + successor := &resources.ActorAttribution{ + Ref: resources.ActorRef{Atespace: "ate-demo", Name: "counter-2"}, + UID: "actor-2", + } + s := newCheckpointTestService() + s.activeActor.Store(successor) + + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err != nil { + t.Fatalf("CheckpointWorkload: %v", err) + } + if !slices.Equal(resp.GetSnapshotFiles(), want) { + t.Errorf("SnapshotFiles = %v, want %v", resp.GetSnapshotFiles(), want) + } + + if got := s.activeActor.Load(); got != successor { + t.Errorf("activeActor = %v, want the successor %v left untouched", got, successor) + } +} + +// A marker from a differently-scoped checkpoint must not be replayed. +func TestCheckpointWorkloadDoesNotReplayADifferentScope(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + if err := checkpointmarker.Write(actorUID, ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA.String(), []string{"durable-dir.tar"}); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + s := newCheckpointTestService() + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err == nil { + t.Fatalf("CheckpointWorkload succeeded (files=%v), want a failure rather than a replay of the DATA snapshot", resp.GetSnapshotFiles()) + } + if slices.Contains(resp.GetSnapshotFiles(), "durable-dir.tar") { + t.Errorf("SnapshotFiles = %v, want the DATA marker's files not replayed", resp.GetSnapshotFiles()) + } +} + +func TestSandboxNotFound(t *testing.T) { + tests := []struct { + name string + out string + want bool + }{ + {"container absent", `error: loading container: container "pause" does not exist`, true}, + {"control server unresponsive", "error: connecting to control server: connection refused", false}, + {"runsc binary missing", "fork/exec /usr/bin/runsc: no such file or directory", false}, + {"probe timed out", "signal: killed", false}, + {"no output at all", "", false}, + { + "phrase in log line but not verdict", + `{"msg":"cgroup path \"/sys/fs/cgroup/runsc\" does not exist, skipping","level":"warning"} +error: connecting to control server: connection refused`, + false, + }, + { + "verdict after log noise", + `{"msg":"loading container","level":"info"} +error: loading container: container "pause" does not exist`, + true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := sandboxNotFound([]byte(tt.out)); got != tt.want { + t.Errorf("sandboxNotFound(%q) = %v, want %v", tt.out, got, tt.want) + } + }) + } +} + +// Retried checkpoints must clear stale checkpoint files before writing new images. +func TestCheckpointWorkloadClearsStaleCheckpointFiles(t *testing.T) { + const actorUID = "actor-1" + dir := useTempActorsDir(t, actorUID) + + stale := filepath.Join(dir, "pages.img") + if err := os.WriteFile(stale, []byte("half-written"), 0o600); err != nil { + t.Fatalf("writing stale image: %v", err) + } + + s := newCheckpointTestService() + if _, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }); err == nil { + t.Fatal("CheckpointWorkload succeeded, want a failure with no runsc to drive") + } + + if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { + t.Errorf("os.Stat(%q) = %v, want the previous attempt's image to be gone", stale, err) + } +} + +// When the ateom is still assigned to the checkpointed actor, replaying a +// completed checkpoint runs teardown and clears the active session. +func TestCheckpointWorkloadReplayCleansUpWhenNotReassigned(t *testing.T) { + roottest.Require(t, "teardown touches network namespaces and mounts") + + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + want := []string{"checkpoint.img", "pages.img"} + if err := checkpointmarker.Write(actorUID, testScope, want); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + current := &resources.ActorAttribution{ + Ref: resources.ActorRef{Atespace: "ate-demo", Name: "counter-1"}, + UID: actorUID, + } + s := newCheckpointTestService() + s.activeActor.Store(current) + s.activeSession = &workloadSession{} + + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err != nil { + t.Fatalf("CheckpointWorkload: %v", err) + } + if !slices.Equal(resp.GetSnapshotFiles(), want) { + t.Errorf("SnapshotFiles = %v, want %v", resp.GetSnapshotFiles(), want) + } + if s.activeSession != nil { + t.Errorf("activeSession = %v, want nil after replay cleanup", s.activeSession) + } +} + +func TestCheckpointWorkloadScopeValidation(t *testing.T) { + tests := []struct { + name string + scope ateompb.SnapshotScope + spec *ateompb.WorkloadSpec + errContains string + }{ + { + name: "data scope without durable volumes", + scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA, + spec: &ateompb.WorkloadSpec{}, + errContains: "no durable-dir volumes found", + }, + { + name: "unsupported scope", + scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED, + errContains: "unsupported snapshot scope", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + s := newCheckpointTestService() + _, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: tt.scope, + Spec: tt.spec, + }) + if err == nil || !strings.Contains(err.Error(), tt.errContains) { + t.Errorf("CheckpointWorkload err = %v, want error containing %q", err, tt.errContains) + } + }) + } +} diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index c555fcd56..ab244c5bb 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -26,7 +26,6 @@ import ( "os" "os/signal" "slices" - "sort" "strings" "sync" "sync/atomic" @@ -713,117 +712,6 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload return &ateompb.RunWorkloadResponse{}, nil } -// Allow checkpointing even if the pod is shutting down. This will allow actors -// (or the harness) to suspend on shutdown. -func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.CheckpointWorkloadRequest) (*ateompb.CheckpointWorkloadResponse, error) { - s.lock.Lock() - defer s.lock.Unlock() - - ctx, cancel := context.WithCancel(ctx) - defer cancel() - s.setActiveRPC(rpcCheckpointWorkload, cancel) - defer s.clearActiveRPC() - - if err := s.deactivateActorNetworking(ctx); err != nil { - return nil, err - } - - attribution := ateomstats.ActorAttributionFromRequest(req) - s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointing", attribution) - - // Contract with atelet: - // - // * After we exit, atelet will upload checkpoint to GCS - // * After we exit, atelet will tear down OCI bundles and reset the actor directory. - - // Checkpoint only saves state; no sizing is applied, so size is left zero. - rcmd := &runsc{ - path: req.GetRunscPath(), - actorUID: req.GetActorUid(), - } - - checkpointPath := ateompath.CheckpointStateDir(req.GetActorUid()) - if err := os.MkdirAll(checkpointPath, 0o700); err != nil { - return nil, fmt.Errorf("while creating checkpoint directory: %w", err) - } - - // Always take durable-dir snapshot if at least one container has a durable-dir volume mount. - // TODO(dberkov): this is a temporary workaround until gVisor supports taking durable-dir snapshots in a single request with the process snapshot. - switch req.GetScope() { - case ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA: - var ddv []string - for _, ctr := range req.GetSpec().GetContainers() { - for _, m := range ctr.GetDurableDirVolumeMounts() { - ddv = append(ddv, m.GetMountPath()) - } - } - if len(ddv) == 0 { - return nil, fmt.Errorf("no durable-dir volumes found for DATA snapshot") - } - if err := rcmd.cmdFsCheckpoint(ctx, "pause", checkpointPath, ddv); err != nil { - return nil, fmt.Errorf("while fscheckpointing durable-dir %q: %w", ddv[0], err) - } - case ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL: - // Checkpoint pause container (root of the sandbox) - if err := rcmd.cmdCheckpoint(ctx, "pause", checkpointPath); err != nil { - return nil, fmt.Errorf("while checkpointing pause: %w", err) - } - default: - return nil, fmt.Errorf("unsupported snapshot scope: %v", req.GetScope()) - } - - // The sandbox is gone as of the checkpoint above, so the ateom is back to - // "available" from here on: there is nothing left to measure, and holding - // the attribution would let a later GetWorkloadStats report a checkpointed - // actor as though it were still running. - // - // Cleared here rather than at the end of the function because everything - // below is bookkeeping over a dead sandbox and can still fail (listing the - // snapshot files returns an error), which would otherwise leave the - // attribution behind. Conversely nothing above this point clears it: a - // checkpoint that failed may well have left the workload running, and - // reporting its usage is then the honest answer. - s.activeActor.Store(nil) - - // Cleanup the containers after checkpointing. - // This is best-effort cleanup for actor containers that may have been left behind after checkpointing. - if err := s.terminateWorkload(ctx, attribution.Ref, attribution.UID, req.GetRunscPath(), req.GetSpec().GetContainers()); err != nil { - slog.WarnContext(ctx, "failed to terminate workload after checkpoint", - slog.String("actor", attribution.Ref.String()), - slog.String("actorUID", attribution.UID), - slog.Any("err", err)) - } - - // Report exactly the files runsc wrote so atelet ships precisely this set - // (checkpoint.img plus any pages images), rather than a hardcoded list. - snapshotFiles, err := listSnapshotFiles(checkpointPath) - if err != nil { - return nil, fmt.Errorf("while listing checkpoint files: %w", err) - } - - s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointed", attribution) - s.activeSession = nil - - return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles}, nil -} - -// listSnapshotFiles returns the (relative) names of regular files directly under -// dir, which atelet ships to object storage as the snapshot. -func listSnapshotFiles(dir string) ([]string, error) { - entries, err := os.ReadDir(dir) - if err != nil { - return nil, err - } - var files []string - for _, e := range entries { - if e.Type().IsRegular() { - files = append(files, e.Name()) - } - } - sort.Strings(files) - return files, nil -} - func (r *runsc) stopContainers(ctx context.Context, containers []*ateompb.Container) { for _, ctr := range containers { _ = r.cmdKill(ctx, ctr.GetName(), "SIGKILL") diff --git a/cmd/ateom-gvisor/runsc.go b/cmd/ateom-gvisor/runsc.go index 0d517a1b4..4e7fc44fb 100644 --- a/cmd/ateom-gvisor/runsc.go +++ b/cmd/ateom-gvisor/runsc.go @@ -316,19 +316,22 @@ func (r *runsc) cmdDelete(ctx context.Context, containerName string) error { return nil } -func (r *runsc) cmdState(ctx context.Context, containerName string) error { - reapLock.RLock() - defer reapLock.RUnlock() - - cmd := exec.CommandContext( - ctx, - r.path, +// stateArgs builds the argv for `runsc state `. +func (r *runsc) stateArgs(containerName string) []string { + return []string{ "-log-format", "json", "--alsologtostderr", "-root", ateompath.RunSCStateDir(r.actorUID), "state", containerName, - ) + } +} + +func (r *runsc) cmdState(ctx context.Context, containerName string) error { + reapLock.RLock() + defer reapLock.RUnlock() + + cmd := exec.CommandContext(ctx, r.path, r.stateArgs(containerName)...) cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr if err := cmd.Run(); err != nil { @@ -337,6 +340,18 @@ func (r *runsc) cmdState(ctx context.Context, containerName string) error { return nil } +// cmdStateOutput executes `runsc state` and captures combined stdout/stderr for failure inspection. +func (r *runsc) cmdStateOutput(ctx context.Context, containerName string) ([]byte, error) { + reapLock.RLock() + defer reapLock.RUnlock() + + out, err := exec.CommandContext(ctx, r.path, r.stateArgs(containerName)...).CombinedOutput() + if err != nil { + return out, fmt.Errorf("while running `runsc state`: %w", err) + } + return out, nil +} + // killArgs builds the argv for `runsc kill `. Factored out // so the argument construction can be unit-tested without executing runsc. func (r *runsc) killArgs(containerName, signal string) []string { diff --git a/cmd/ateom-microvm/checkpoint.go b/cmd/ateom-microvm/checkpoint.go index 8a4c8c75a..635bb6fd3 100644 --- a/cmd/ateom-microvm/checkpoint.go +++ b/cmd/ateom-microvm/checkpoint.go @@ -29,8 +29,10 @@ import ( "github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/ch" "github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/kata" + "github.com/agent-substrate/substrate/internal/ateerrors" "github.com/agent-substrate/substrate/internal/ateompath" "github.com/agent-substrate/substrate/internal/ateomstats" + "github.com/agent-substrate/substrate/internal/checkpointmarker" "github.com/agent-substrate/substrate/internal/imagecache" "github.com/agent-substrate/substrate/internal/proto/ateompb" "golang.org/x/sync/errgroup" @@ -70,12 +72,31 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec s.setActiveRPC(rpcCheckpointWorkload, cancel) defer s.clearActiveRPC() - if err := s.deactivateActorNetworking(ctx); err != nil { + attribution := ateomstats.ActorAttributionFromRequest(req) + actorUID := req.GetActorUid() + + // Replay a previously completed checkpoint for this actor if available. + if rec, ok, err := checkpointmarker.Read(actorUID, req.GetScope().String()); err != nil { return nil, err + } else if ok { + slog.InfoContext(ctx, "Checkpoint already completed for this actor; replaying its result", + slog.String("id", actorUID), slog.Any("snapshot_files", rec.SnapshotFiles)) + // Finish any pending workload termination, unless the ateom now holds a different actor. + if held := s.activeActor.Load(); held != nil && held.UID != actorUID { + slog.WarnContext(ctx, "Not running the post-checkpoint teardown: this ateom now holds a different actor", + slog.String("id", actorUID), slog.String("active_actor_uid", held.UID)) + return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil + } + if err := s.terminateWorkload(ctx, actorUID); err != nil { + slog.WarnContext(ctx, "Failed to terminate workload while replaying checkpoint", + slog.String("actorUID", actorUID), slog.Any("err", err)) + } + return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: rec.SnapshotFiles}, nil } - attribution := ateomstats.ActorAttributionFromRequest(req) - actorUID := req.GetActorUid() + if err := s.deactivateActorNetworking(ctx); err != nil { + return nil, err + } s.actorLogger.EmitLifecycleLog(ctx, "Actor checkpointing", attribution) @@ -111,6 +132,12 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec } client := ch.NewClient(chSocket) if _, err := client.WaitReady(ctx, 10*time.Second); err != nil { + // If the API socket is completely gone, the VMM no longer exists (unrecoverable). + // Otherwise, treat as a potentially retriable transient error. + if _, statErr := os.Stat(chSocket); errors.Is(statErr, os.ErrNotExist) { + return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonInvalidCheckpointResult, ateerrors.ActorCrashedMetadata(), + fmt.Errorf("%w: no guest remains to checkpoint: api-socket %q is gone: %w", ateerrors.ReasonInvalidCheckpointResult, chSocket, err)) + } return nil, fmt.Errorf("while waiting for CH api-socket: %w", err) } @@ -186,6 +213,12 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec return nil, fmt.Errorf("while listing snapshot files: %w", err) } + // Record checkpoint completion before teardown. In micro-VM, failing here is safe + // because the guest is only paused and can be retried from the top. + if err := checkpointmarker.Write(actorUID, req.GetScope().String(), snapshotFiles); err != nil { + return nil, err + } + // Tear down: the actor returns to "available". Best-effort; the snapshot is // already on disk for atelet to ship. tTeardown := time.Now() @@ -274,7 +307,9 @@ func listFiles(dir string) ([]string, error) { } var files []string for _, e := range entries { - if e.Type().IsRegular() { + // ateom's own completion marker shares the directory but is + // bookkeeping, not snapshot content, so it never joins the set. + if e.Type().IsRegular() && e.Name() != ateompath.CheckpointDoneFileName { files = append(files, e.Name()) } } diff --git a/cmd/ateom-microvm/checkpoint_test.go b/cmd/ateom-microvm/checkpoint_test.go new file mode 100644 index 000000000..d8ec30a5e --- /dev/null +++ b/cmd/ateom-microvm/checkpoint_test.go @@ -0,0 +1,233 @@ +//go:build linux + +// 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" + "io" + "os" + "path/filepath" + "slices" + "testing" + + "github.com/agent-substrate/substrate/internal/actorlog" + "github.com/agent-substrate/substrate/internal/ateompath" + "github.com/agent-substrate/substrate/internal/atunnel" + "github.com/agent-substrate/substrate/internal/checkpointmarker" + "github.com/agent-substrate/substrate/internal/proto/ateompb" + "github.com/agent-substrate/substrate/internal/resources" + "github.com/agent-substrate/substrate/internal/roottest" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// testScope stands in for a CheckpointWorkloadRequest's stringified scope. +var testScope = ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL.String() + +// newCheckpointTestService builds a minimal AteomService for testing pre-sandbox checkpoint logic. +func newCheckpointTestService() *AteomService { + return &AteomService{ + lock: newCancelableMutex(), + atunnelIngress: &atunnel.Server{}, + atunnelEgress: &atunnel.Egress{}, + actorLogger: actorlog.NewActorLogger(io.Discard, false), + running: make(map[string]*runningActor), + } +} + +func useTempActorsDir(t *testing.T, actorUID string) string { + t.Helper() + orig := ateompath.ActorsDir + t.Cleanup(func() { ateompath.ActorsDir = orig }) + ateompath.ActorsDir = t.TempDir() + + dir := ateompath.CheckpointStateDir(actorUID) + if err := os.MkdirAll(dir, 0o700); err != nil { + t.Fatalf("creating checkpoint dir: %v", err) + } + return dir +} + +func TestListFilesExcludesCompletionMarker(t *testing.T) { + const actorUID = "actor-1" + dir := useTempActorsDir(t, actorUID) + + for _, name := range []string{"snapshot.state", ateompath.CheckpointDoneFileName} { + if err := os.WriteFile(filepath.Join(dir, name), []byte("x"), 0o600); err != nil { + t.Fatalf("writing %s: %v", name, err) + } + } + + got, err := listFiles(dir) + if err != nil { + t.Fatalf("listFiles: %v", err) + } + want := []string{"snapshot.state"} + if !slices.Equal(got, want) { + t.Errorf("listFiles = %v, want %v", got, want) + } +} + +// If the worker has already been reassigned to a successor actor, replaying +// a completed checkpoint must skip teardown to avoid disrupting the new actor. +func TestCheckpointWorkloadReplaySkipsTeardownForAReassignedAteom(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + want := []string{"snapshot.mem", "snapshot.state"} + if err := checkpointmarker.Write(actorUID, ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL.String(), want); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + successor := &resources.ActorAttribution{ + Ref: resources.ActorRef{Atespace: "ate-demo", Name: "counter-2"}, + UID: "actor-2", + } + s := newCheckpointTestService() + s.activeActor.Store(successor) + s.guestStats.Store(&guestStatsTarget{actorUID: successor.UID}) + + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err != nil { + t.Fatalf("CheckpointWorkload: %v", err) + } + if !slices.Equal(resp.GetSnapshotFiles(), want) { + t.Errorf("SnapshotFiles = %v, want %v", resp.GetSnapshotFiles(), want) + } + + if got := s.activeActor.Load(); got != successor { + t.Errorf("activeActor = %v, want the successor %v left untouched", got, successor) + } + if got := s.guestStats.Load(); got == nil || got.actorUID != successor.UID { + t.Errorf("guestStats = %v, want the successor's target left untouched", got) + } +} + +// A marker from a differently-scoped checkpoint must not be replayed. +func TestCheckpointWorkloadDoesNotReplayADifferentScope(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + if err := checkpointmarker.Write(actorUID, ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA.String(), []string{"durable-dir.tar"}); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + s := newCheckpointTestService() + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err == nil { + t.Fatalf("CheckpointWorkload succeeded (files=%v), want a failure rather than a replay of the DATA snapshot", resp.GetSnapshotFiles()) + } + if slices.Contains(resp.GetSnapshotFiles(), "durable-dir.tar") { + t.Errorf("SnapshotFiles = %v, want the DATA marker's files not replayed", resp.GetSnapshotFiles()) + } +} + +// When the ateom is still assigned to the checkpointed actor, replaying a +// completed checkpoint runs teardown to clean up the actor's state. +func TestCheckpointWorkloadReplayCleansUpWhenNotReassigned(t *testing.T) { + roottest.Require(t, "teardown touches network namespaces and mounts") + + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + want := []string{"snapshot.mem", "snapshot.state"} + if err := checkpointmarker.Write(actorUID, testScope, want); err != nil { + t.Fatalf("checkpointmarker.Write: %v", err) + } + + current := &resources.ActorAttribution{ + Ref: resources.ActorRef{Atespace: "ate-demo", Name: "counter-1"}, + UID: actorUID, + } + s := newCheckpointTestService() + s.activeActor.Store(current) + s.guestStats.Store(&guestStatsTarget{actorUID: actorUID}) + s.running[actorUID] = &runningActor{} + + resp, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_FULL, + }) + if err != nil { + t.Fatalf("CheckpointWorkload: %v", err) + } + if !slices.Equal(resp.GetSnapshotFiles(), want) { + t.Errorf("SnapshotFiles = %v, want %v", resp.GetSnapshotFiles(), want) + } + + if got := s.activeActor.Load(); got != nil { + t.Errorf("activeActor = %v, want nil after replay cleanup", got) + } + if got := s.guestStats.Load(); got != nil { + t.Errorf("guestStats = %v, want nil after replay cleanup", got) + } + if _, exists := s.running[actorUID]; exists { + t.Errorf("running[%q] still exists, want removed after replay cleanup", actorUID) + } +} + +func TestCheckpointWorkloadScopePreconditions(t *testing.T) { + tests := []struct { + name string + scope ateompb.SnapshotScope + spec *ateompb.WorkloadSpec + wantCode codes.Code + }{ + { + name: "data scope without volumes fails precondition", + scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_DATA, + spec: &ateompb.WorkloadSpec{}, + wantCode: codes.FailedPrecondition, + }, + { + name: "unsupported scope returns invalid argument", + scope: ateompb.SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED, + wantCode: codes.InvalidArgument, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + s := newCheckpointTestService() + _, err := s.CheckpointWorkload(context.Background(), &ateompb.CheckpointWorkloadRequest{ + Atespace: "ate-demo", + ActorName: "counter-1", + ActorUid: actorUID, + Scope: tt.scope, + Spec: tt.spec, + }) + if gotCode := status.Code(err); gotCode != tt.wantCode { + t.Errorf("CheckpointWorkload code = %v, want %v (err: %v)", gotCode, tt.wantCode, err) + } + }) + } +} diff --git a/internal/ateompath/ateompath.go b/internal/ateompath/ateompath.go index 554781a0b..a78e1e8be 100644 --- a/internal/ateompath/ateompath.go +++ b/internal/ateompath/ateompath.go @@ -178,6 +178,18 @@ func CheckpointStateDir(actorUID string) string { ) } +// CheckpointDoneFileName is the filename of the checkpoint completion marker +// written into CheckpointStateDir upon successful checkpoint. +const CheckpointDoneFileName = "checkpoint-done.json" + +// CheckpointDoneFile returns the path to the actor's checkpoint completion marker file. +func CheckpointDoneFile(actorUID string) string { + return filepath.Join( + CheckpointStateDir(actorUID), + CheckpointDoneFileName, + ) +} + func LocalCheckpointsDir(actorUID string) string { return filepath.Join( ActorPath(actorUID), diff --git a/internal/checkpointmarker/checkpointmarker.go b/internal/checkpointmarker/checkpointmarker.go new file mode 100644 index 000000000..0f75cb201 --- /dev/null +++ b/internal/checkpointmarker/checkpointmarker.go @@ -0,0 +1,128 @@ +// 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 checkpointmarker reads and writes the per-actor checkpoint +// completion marker both ateom runtimes use for idempotent replays. +// +// Because checkpoints are destructive (taking the sandbox down), the marker +// records completed snapshot files so lost responses or retried calls can be +// replayed without re-running against a destroyed sandbox. +package checkpointmarker + +import ( + "encoding/json" + "errors" + "fmt" + "log/slog" + "os" + "path/filepath" + + "github.com/agent-substrate/substrate/internal/ateompath" +) + +// Record represents the marker written on checkpoint completion. +type Record struct { + SnapshotFiles []string `json:"snapshotFiles"` + // Scope is the checkpoint scope (e.g. FULL, DATA), ensuring a marker is only + // replayed for matching checkpoint requests. + Scope string `json:"scope"` +} + +// Write records the completed checkpoint for actorUID atomically (temp file +// plus rename and sync). An empty snapshotFiles slice is valid (e.g. CSI +// volumes snapshotted out-of-band) and still recorded. +func Write(actorUID, scope string, snapshotFiles []string) error { + if scope == "" { + return fmt.Errorf("refusing to record a checkpoint marker for actor %q with no scope", actorUID) + } + + data, err := json.Marshal(&Record{SnapshotFiles: snapshotFiles, Scope: scope}) + if err != nil { + return fmt.Errorf("while marshaling checkpoint marker: %w", err) + } + + path := ateompath.CheckpointDoneFile(actorUID) + tmp, err := os.CreateTemp(filepath.Dir(path), "."+ateompath.CheckpointDoneFileName+".tmp-") + if err != nil { + return fmt.Errorf("while creating checkpoint marker temp file: %w", err) + } + tmpName := tmp.Name() + defer func() { + tmp.Close() + os.Remove(tmpName) // no-op once the rename below succeeds + }() + + if _, err := tmp.Write(data); err != nil { + return fmt.Errorf("while writing checkpoint marker: %w", err) + } + // Flush bytes before rename to ensure file content is durable on disk. + if err := tmp.Sync(); err != nil { + return fmt.Errorf("while syncing checkpoint marker: %w", err) + } + if err := tmp.Close(); err != nil { + return fmt.Errorf("while closing checkpoint marker: %w", err) + } + + // Atomically publish the marker so readers never observe partial data. + if err := os.Rename(tmpName, path); err != nil { + return fmt.Errorf("while renaming checkpoint marker into place: %w", err) + } + + // Sync parent directory to persist the directory entry across crashes. + dir, err := os.Open(filepath.Dir(path)) + if err != nil { + return fmt.Errorf("while opening checkpoint marker directory to sync: %w", err) + } + defer dir.Close() + if err := dir.Sync(); err != nil { + return fmt.Errorf("while syncing checkpoint marker directory: %w", err) + } + return nil +} + +// Read returns the marker recorded for actorUID matching the requested scope. +// ok is false if no marker exists, the scope mismatches, or the marker is invalid. +func Read(actorUID, scope string) (_ *Record, ok bool, _ error) { + path := ateompath.CheckpointDoneFile(actorUID) + data, err := os.ReadFile(path) + if errors.Is(err, os.ErrNotExist) { + return nil, false, nil + } + if err != nil { + return nil, false, fmt.Errorf("while reading checkpoint marker: %w", err) + } + rec := &Record{} + if err := json.Unmarshal(data, rec); err != nil { + discardUnusable(path, actorUID, fmt.Errorf("while parsing checkpoint marker: %w", err)) + return nil, false, nil + } + if rec.Scope != scope { + slog.Info("Checkpoint marker records a different checkpoint; not replaying it", + slog.String("actor_uid", actorUID), slog.String("marker_scope", rec.Scope), slog.String("requested_scope", scope)) + return nil, false, nil + } + return rec, true, nil +} + +// discardUnusable removes corrupted marker files to prevent retry loops from +// repeatedly failing on unparseable state. +func discardUnusable(path, actorUID string, reason error) { + slog.Warn("Discarding unusable checkpoint marker; the checkpoint will be re-attempted", + slog.String("actor_uid", actorUID), slog.String("path", path), slog.Any("err", reason)) + if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) { + // Not fatal: a later successful checkpoint overwrites the marker by rename regardless. + slog.Warn("Failed to remove unusable checkpoint marker", + slog.String("actor_uid", actorUID), slog.String("path", path), slog.Any("err", err)) + } +} diff --git a/internal/checkpointmarker/checkpointmarker_test.go b/internal/checkpointmarker/checkpointmarker_test.go new file mode 100644 index 000000000..bcb04c320 --- /dev/null +++ b/internal/checkpointmarker/checkpointmarker_test.go @@ -0,0 +1,174 @@ +// 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 checkpointmarker + +import ( + "os" + "path/filepath" + "slices" + "testing" + + "github.com/agent-substrate/substrate/internal/ateompath" +) + +// useTempActorsDir points the shared actor-state root at a temp directory for +// the duration of the test, and creates the actor's checkpoint dir (ateom +// makes it before checkpointing). +func useTempActorsDir(t *testing.T, actorUID string) { + t.Helper() + orig := ateompath.ActorsDir + t.Cleanup(func() { ateompath.ActorsDir = orig }) + ateompath.ActorsDir = t.TempDir() + + if err := os.MkdirAll(ateompath.CheckpointStateDir(actorUID), 0o700); err != nil { + t.Fatalf("creating checkpoint dir: %v", err) + } +} + +// testScope stands in for a CheckpointWorkloadRequest's stringified scope. +const testScope = "SNAPSHOT_SCOPE_FULL" + +func TestWriteThenRead(t *testing.T) { + tests := []struct { + name string + files []string + }{ + {"with files", []string{"checkpoint.img", "pages.img", "pages_meta.img"}}, + {"empty file set", nil}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + if err := Write(actorUID, testScope, tt.files); err != nil { + t.Fatalf("Write: %v", err) + } + + rec, ok, err := Read(actorUID, testScope) + if err != nil { + t.Fatalf("Read: %v", err) + } + if !ok { + t.Fatal("Read reported no marker after Write") + } + if !slices.Equal(rec.SnapshotFiles, tt.files) { + t.Errorf("SnapshotFiles = %v, want %v", rec.SnapshotFiles, tt.files) + } + }) + } +} + +func TestWriteLeavesNoTempFiles(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + if err := Write(actorUID, testScope, []string{"checkpoint.img"}); err != nil { + t.Fatalf("Write: %v", err) + } + + entries, err := os.ReadDir(ateompath.CheckpointStateDir(actorUID)) + if err != nil { + t.Fatalf("ReadDir: %v", err) + } + if len(entries) != 1 || entries[0].Name() != ateompath.CheckpointDoneFileName { + var got []string + for _, e := range entries { + got = append(got, e.Name()) + } + t.Errorf("checkpoint dir contents = %v, want only %q", got, ateompath.CheckpointDoneFileName) + } +} + +func TestReadNoMarker(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + rec, ok, err := Read(actorUID, testScope) + if err != nil { + t.Fatalf("Read: %v", err) + } + if ok || rec != nil { + t.Errorf("Read = (%v, %v), want (nil, false)", rec, ok) + } +} + +// An unusable marker is discarded so the checkpoint can be re-attempted. +func TestReadDiscardsUnusableMarker(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + path := filepath.Join(ateompath.CheckpointStateDir(actorUID), ateompath.CheckpointDoneFileName) + if err := os.WriteFile(path, []byte("{not json"), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + rec, ok, err := Read(actorUID, testScope) + if err != nil { + t.Fatalf("Read: %v, want no error", err) + } + if ok || rec != nil { + t.Errorf("Read = (%v, %v), want (nil, false)", rec, ok) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Errorf("marker still on disk (err=%v), want removed", err) + } +} + +func TestWriteRejectsEmptyScope(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + if err := Write(actorUID, "", []string{"checkpoint.img"}); err == nil { + t.Fatal("Write succeeded, want an error") + } + path := filepath.Join(ateompath.CheckpointStateDir(actorUID), ateompath.CheckpointDoneFileName) + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Errorf("marker written (err=%v), want none", err) + } +} + +// Read rejects markers with a different or missing scope. +func TestReadRejectsMarkerFromADifferentScope(t *testing.T) { + tests := []struct { + name string + content string + }{ + {"different scope", `{"snapshotFiles":["durable-dir.tar"],"scope":"SNAPSHOT_SCOPE_DATA"}`}, + {"no scope", `{"snapshotFiles":["checkpoint.img"]}`}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + const actorUID = "actor-1" + useTempActorsDir(t, actorUID) + + path := filepath.Join(ateompath.CheckpointStateDir(actorUID), ateompath.CheckpointDoneFileName) + if err := os.WriteFile(path, []byte(tt.content), 0o600); err != nil { + t.Fatalf("WriteFile: %v", err) + } + + rec, ok, err := Read(actorUID, testScope) + if err != nil { + t.Fatalf("Read: %v, want no error", err) + } + if ok || rec != nil { + t.Errorf("Read = (%v, %v), want (nil, false)", rec, ok) + } + if _, err := os.Stat(path); err != nil { + t.Errorf("marker removed (err=%v), want it left in place", err) + } + }) + } +}