diff --git a/cmd/atelet/durabletar_linux.go b/cmd/atelet/durabletar_linux.go new file mode 100644 index 0000000000..7927fa12b0 --- /dev/null +++ b/cmd/atelet/durabletar_linux.go @@ -0,0 +1,97 @@ +// 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" + "strings" + "sync/atomic" + + "github.com/agent-substrate/substrate/internal/tarutil" + "golang.org/x/sync/errgroup" +) + +// streamDurableTarEnv turns the streaming durable-dir capture on. Off by +// default: it moves the archive out of ateom's paused window (see +// CheckpointWorkloadRequest.skip_durable_dir_tar) and is worth measuring +// against the staged path before it becomes the only one. +const streamDurableTarEnv = "ATELET_STREAM_DURABLE_TAR" + +// streamDurableTarEnabled reports whether this atelet may take the durable-dir +// archive over from ateom. Callers must still check that the actor's runtime +// and checkpoint type allow it; see canStreamDurableDirTar. +func streamDurableTarEnabled() bool { + switch strings.ToLower(strings.TrimSpace(os.Getenv(streamDurableTarEnv))) { + case "1", "t", "true", "y", "yes", "on": + return true + } + return false +} + +// streamDurableDirTar archives srcDir and hands the archive to upload as it is +// produced, so the walk overlaps compression and the network PUT and the +// archive is never staged on disk. It returns the uncompressed archive size for +// the size metric, which is only meaningful when the error is nil. +// +// The archive must be identical to the one ateom would have written, or a +// snapshot taken this way would not restore the same tree: micro-VM archives +// all of the durable-dir mounts dir unfiltered (see ateom-microvm's +// tarDurableVolumes), which is what a nil SkipFunc reproduces here. +func streamDurableDirTar(ctx context.Context, srcDir string, upload func(io.Reader) error) (int64, error) { + pr, pw := io.Pipe() + counter := &countingWriter{w: pw} + + var g errgroup.Group + g.Go(func() error { + err := tarutil.CreateTo(ctx, counter, srcDir, nil) + // A nil error closes the write end normally, ending the upload's read + // at EOF; anything else surfaces at that read, so the upload aborts + // instead of committing a truncated object. + pw.CloseWithError(err) + return err + }) + + uploadErr := upload(pr) + // Unblocks the walk if the upload stopped reading early, so g.Wait cannot + // hang on a full pipe. + pr.CloseWithError(uploadErr) + tarErr := g.Wait() + + // uploadErr first: when the walk is what failed, the upload's read returned + // that same error and reports it with the upload's context attached, while + // tarErr in the reverse case is only io.ErrClosedPipe. + if uploadErr != nil { + return 0, uploadErr + } + if tarErr != nil { + return 0, tarErr + } + return counter.n.Load(), nil +} + +// countingWriter totals the bytes written through it. The archive never exists +// as a file, so this is the only place its size can be observed. +type countingWriter struct { + w io.Writer + n atomic.Int64 +} + +func (c *countingWriter) Write(p []byte) (int, error) { + n, err := c.w.Write(p) + c.n.Add(int64(n)) + return n, err +} diff --git a/cmd/atelet/durabletar_linux_test.go b/cmd/atelet/durabletar_linux_test.go new file mode 100644 index 0000000000..1c5042366e --- /dev/null +++ b/cmd/atelet/durabletar_linux_test.go @@ -0,0 +1,248 @@ +// 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 ( + "bytes" + "errors" + "io" + "os" + "path/filepath" + "slices" + "testing" + + "github.com/agent-substrate/substrate/internal/ateompath" + ateletpb "github.com/agent-substrate/substrate/internal/proto/ateletpb" + ateompb "github.com/agent-substrate/substrate/internal/proto/ateompb" + "github.com/agent-substrate/substrate/internal/resources" + "github.com/agent-substrate/substrate/internal/tarutil" + atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1" + "github.com/klauspost/compress/zstd" +) + +// durableDirFixture builds a small durable-dir tree and returns its path. +func durableDirFixture(t *testing.T) string { + t.Helper() + dir := t.TempDir() + if err := os.MkdirAll(filepath.Join(dir, "vol-a", "nested"), 0o755); err != nil { + t.Fatal(err) + } + // Big enough to outrun the pipe's internal handoff, so the walk and the + // consumer really do run concurrently rather than completing in one write. + if err := os.WriteFile(filepath.Join(dir, "vol-a", "big"), bytes.Repeat([]byte("substrate"), 1<<16), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "vol-a", "nested", "small"), []byte("hello"), 0o600); err != nil { + t.Fatal(err) + } + if err := os.Symlink("big", filepath.Join(dir, "vol-a", "link")); err != nil { + t.Fatal(err) + } + return dir +} + +// TestStreamDurableDirTarMatchesStagedArchive is the property the whole change +// rests on: a snapshot taken by streaming must restore the same tree as one +// ateom staged on disk, so the two archives have to be byte-identical. +func TestStreamDurableDirTarMatchesStagedArchive(t *testing.T) { + dir := durableDirFixture(t) + + staged := filepath.Join(t.TempDir(), "durable-dir.tar") + if err := tarutil.Create(t.Context(), staged, dir); err != nil { + t.Fatalf("staging archive: %v", err) + } + want, err := os.ReadFile(staged) + if err != nil { + t.Fatal(err) + } + + var got bytes.Buffer + size, err := streamDurableDirTar(t.Context(), dir, func(r io.Reader) error { + _, err := io.Copy(&got, r) + return err + }) + if err != nil { + t.Fatalf("streamDurableDirTar: %v", err) + } + if !bytes.Equal(got.Bytes(), want) { + t.Errorf("streamed archive differs from staged archive: got %d bytes, want %d", got.Len(), len(want)) + } + if size != int64(len(want)) { + t.Errorf("reported size = %d, want %d", size, len(want)) + } +} + +// TestStreamDurableDirTarUploadFailure checks the upload's error is what the +// caller sees, and that the walk does not deadlock once nobody is reading. +func TestStreamDurableDirTarUploadFailure(t *testing.T) { + dir := durableDirFixture(t) + sentinel := errors.New("object storage is down") + + // Reads one byte and quits, leaving the walk blocked on a pipe with no + // reader: only the CloseWithError on the read end lets it finish. + size, err := streamDurableDirTar(t.Context(), dir, func(r io.Reader) error { + if _, readErr := io.ReadFull(r, make([]byte, 1)); readErr != nil { + t.Errorf("reading first byte: %v", readErr) + } + return sentinel + }) + if !errors.Is(err, sentinel) { + t.Errorf("err = %v, want %v", err, sentinel) + } + if size != 0 { + t.Errorf("size = %d, want 0 on failure", size) + } +} + +// TestStreamDurableDirTarArchiveFailure checks a failed walk aborts the upload +// rather than committing whatever bytes it managed to produce. +func TestStreamDurableDirTarArchiveFailure(t *testing.T) { + missing := filepath.Join(t.TempDir(), "no-such-durable-dir") + + var uploadErr error + uploaded := false + if _, err := streamDurableDirTar(t.Context(), missing, func(r io.Reader) error { + _, uploadErr = io.Copy(io.Discard, r) + uploaded = uploadErr == nil + return uploadErr + }); err == nil { + t.Fatal("streamDurableDirTar succeeded on a missing source directory") + } + if uploaded { + t.Error("upload saw a clean EOF; a truncated archive would have been committed") + } + if !errors.Is(uploadErr, os.ErrNotExist) { + t.Errorf("upload read error = %v, want it to carry %v", uploadErr, os.ErrNotExist) + } +} + +func TestStreamDurableTarEnabled(t *testing.T) { + for _, tc := range []struct { + value string + want bool + }{ + {"1", true}, {"true", true}, {"TRUE", true}, {" on ", true}, {"yes", true}, + {"", false}, {"0", false}, {"false", false}, {"off", false}, + // Anything unrecognized stays on the staged path: the knob exists to + // try the new one, so a typo must not silently opt in. + {"onn", false}, {"enabled", false}, + } { + t.Run(tc.value, func(t *testing.T) { + t.Setenv(streamDurableTarEnv, tc.value) + if got := streamDurableTarEnabled(); got != tc.want { + t.Errorf("streamDurableTarEnabled(%q) = %v, want %v", tc.value, got, tc.want) + } + }) + } +} + +func TestCanStreamDurableDirTar(t *testing.T) { + externalReq := &ateletpb.CheckpointRequest{Type: ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL} + localReq := &ateletpb.CheckpointRequest{Type: ateletpb.CheckpointType_CHECKPOINT_TYPE_LOCAL} + durableSpec := &ateompb.WorkloadSpec{Containers: []*ateompb.Container{{ + Name: "app", + DurableDirVolumeMounts: []*ateompb.DurableDirVolumeMount{{VolumeName: "data", MountPath: "/data"}}, + }}} + plainSpec := &ateompb.WorkloadSpec{Containers: []*ateompb.Container{{Name: "app"}}} + microVM := &sandboxAssetsRecord{SandboxClass: string(atev1alpha1.SandboxClassMicroVM)} + + for _, tc := range []struct { + name string + env string + req *ateletpb.CheckpointRequest + spec *ateompb.WorkloadSpec + rec *sandboxAssetsRecord + want bool + }{ + {"eligible", "1", externalReq, durableSpec, microVM, true}, + {"knob off", "", externalReq, durableSpec, microVM, false}, + {"local checkpoint", "1", localReq, durableSpec, microVM, false}, + {"gvisor", "1", externalReq, durableSpec, &sandboxAssetsRecord{SandboxClass: string(atev1alpha1.SandboxClassGvisor)}, false}, + // The predicate ateom applies: a declared-but-unmounted durable volume + // makes ateom write no archive, so we must not invent one. + {"no durable mounts", "1", externalReq, plainSpec, microVM, false}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(streamDurableTarEnv, tc.env) + if got := canStreamDurableDirTar(tc.req, tc.spec, tc.rec); got != tc.want { + t.Errorf("canStreamDurableDirTar() = %v, want %v", got, tc.want) + } + }) + } +} + +// TestUploadSnapshotStreamsDurableDirTar covers the seam: the durable-dir tar +// is the one snapshot file with no copy in srcDir, and it still has to land as +// the same object, beside the files that do come off disk. +func TestUploadSnapshotStreamsDurableDirTar(t *testing.T) { + uri, err := resources.ParseSnapshotURI(pausedSnapshotURI) + if err != nil { + t.Fatalf("ParseSnapshotURI: %v", err) + } + + // Holds every file EXCEPT the durable-dir tar, exactly as ateom leaves it + // when it is told to skip that one. + srcDir := t.TempDir() + if err := os.WriteFile(filepath.Join(srcDir, "config.json"), []byte("cfg"), 0o600); err != nil { + t.Fatal(err) + } + durableDir := durableDirFixture(t) + + store := &recordingObjectStorage{} + s := &AteomHerder{gcsClient: store} + rec := &sandboxAssetsRecord{ + SandboxClass: string(atev1alpha1.SandboxClassMicroVM), + SnapshotFiles: []string{"config.json", ateompath.DurableDirTarFile}, + } + if err := s.uploadSnapshot(t.Context(), uri, srcDir, rec, "team-a", "tmpl", durableDir); err != nil { + t.Fatalf("uploadSnapshot: %v", err) + } + + want := []string{ + pausedSnapshotPath + "/config.json.zstd", + pausedSnapshotPath + "/durable-dir.tar.zstd", + pausedSnapshotPath + "/manifest.json", + } + if got := store.keys(); !slices.Equal(got, want) { + t.Fatalf("uploaded objects = %v, want %v", got, want) + } + + // Decompressing and extracting is what proves the streamed object is a + // usable archive rather than merely present. + dec, err := zstd.NewReader(bytes.NewReader(store.objects[pausedSnapshotPath+"/durable-dir.tar.zstd"])) + if err != nil { + t.Fatalf("opening zstd reader: %v", err) + } + defer dec.Close() + plain, err := io.ReadAll(dec) + if err != nil { + t.Fatalf("decompressing durable-dir tar: %v", err) + } + tarPath := filepath.Join(t.TempDir(), "durable-dir.tar") + if err := os.WriteFile(tarPath, plain, 0o600); err != nil { + t.Fatal(err) + } + restored := t.TempDir() + if err := tarutil.Extract(tarPath, restored); err != nil { + t.Fatalf("extracting streamed archive: %v", err) + } + if got, err := os.ReadFile(filepath.Join(restored, "vol-a", "nested", "small")); err != nil || string(got) != "hello" { + t.Errorf("restored vol-a/nested/small = %q, %v; want %q", got, err, "hello") + } +} + +// Compile-time check that the counter is a plain io.Writer, so tarutil writes +// through it unchanged. +var _ io.Writer = (*countingWriter)(nil) diff --git a/cmd/atelet/durabletar_other.go b/cmd/atelet/durabletar_other.go new file mode 100644 index 0000000000..6c3705b935 --- /dev/null +++ b/cmd/atelet/durabletar_other.go @@ -0,0 +1,34 @@ +// 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. + +//go:build !linux + +package main + +import ( + "context" + "errors" + "io" +) + +// tarutil archives the ownership, device nodes, and xattrs a durable-dir +// restore depends on, none of which are portable, so it is Linux-only. atelet +// runs on Linux; these keep the package building for local development +// elsewhere. + +func streamDurableTarEnabled() bool { return false } + +func streamDurableDirTar(_ context.Context, _ string, _ func(io.Reader) error) (int64, error) { + return 0, errors.New("streaming durable-dir capture is only implemented on Linux") +} diff --git a/cmd/atelet/internal/ategcs/objects.go b/cmd/atelet/internal/ategcs/objects.go index feea7057c2..9a24ca6854 100644 --- a/cmd/atelet/internal/ategcs/objects.go +++ b/cmd/atelet/internal/ategcs/objects.go @@ -93,6 +93,27 @@ func SendBytesToGCS(ctx context.Context, client ObjectStorage, gsURL string, con return nil } +// SendReaderToGCSWithZstd is SendLocalFileToGCSWithZstd for content that was +// never written to a file — an archive generated on the fly, say. Whether the +// content really streams depends on the backend: GCS pipes it straight into the +// PUT, while S3/rustfs still stages the compressed bytes in a temp file because +// its SDK needs a seekable body (see sendZstd). +// +// A file source is worth strictly more than a stream here, so this is the wrong +// entry point for content that already sits on disk: sendZstd hands a file to +// PutSparseFile, which compresses and uploads its ranges in parallel and skips +// its holes, and a non-seekable reader can have neither. Reach for this only +// where producing that file is itself the cost being avoided. +func SendReaderToGCSWithZstd(ctx context.Context, client ObjectStorage, gsURL string, content io.Reader) error { + ctx, span := tracer.Start(ctx, "sendReaderToGCSWithZstd") + defer span.End() + + if err := sendZstd(ctx, client, gsURL, content); err != nil { + return fmt.Errorf("in sendZstd: %w", err) + } + return nil +} + func SendLocalFileToGCSWithZstd(ctx context.Context, client ObjectStorage, gsURL string, localFilePath string) (err error) { ctx, span := tracer.Start(ctx, "sendLocalFileToGCSWithZstd") defer span.End() diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 433e71b2a6..e3fc00f885 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -568,7 +568,16 @@ func recordSnapshotSize(ctx context.Context, file, path, templateAtespace, templ slog.String("file", file), slog.String("path", path), slog.Any("err", err)) return } - snapshotSizeBytes.Record(ctx, fi.Size(), metric.WithAttributes( + emitSnapshotSize(ctx, file, fi.Size(), templateAtespace, templateName) +} + +// emitSnapshotSize is recordSnapshotSize for an image whose size is already +// known — one generated straight into its upload, with no file to stat. +func emitSnapshotSize(ctx context.Context, file string, size int64, templateAtespace, templateName string) { + if snapshotSizeBytes == nil { + return + } + snapshotSizeBytes.Record(ctx, size, metric.WithAttributes( semconv.FileNameKey.String(file), ateattr.TemplateAtespaceKey.String(templateAtespace), ateattr.TemplateNameKey.String(templateName), @@ -634,6 +643,12 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe return nil, status.Errorf(codes.InvalidArgument, "invalid workload spec: %v", err) } + // Decided before the call, because it changes what ateom writes. + streamDurableFrom := "" + if canStreamDurableDirTar(req, spec, sandboxRec) { + streamDurableFrom = ateompath.DurableDirVolumeMountsDir(actorUID) + } + tAteom := time.Now() resp, err := client.CheckpointWorkload(ctx, &ateompb.CheckpointWorkloadRequest{ Atespace: actorRef.Atespace, @@ -645,6 +660,7 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe Spec: spec, Scope: toAteomSnapshotScope(req.GetScope()), ActorUid: actorUID, + SkipDurableDirTar: streamDurableFrom != "", }) dAteom = time.Since(tAteom) if err != nil { @@ -655,6 +671,15 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe } sandboxRec.SnapshotFiles = resp.GetSnapshotFiles() + if streamDurableFrom != "" { + // ateom lists what it wrote, and it deliberately did not write this one. + // The manifest still has to name it: it is an object of this snapshot + // like any other, and restore looks it up by name. Sorted because + // ateom's own list comes from a directory read, so this keeps the + // manifest identical to the one the staged path writes. + sandboxRec.SnapshotFiles = append(sandboxRec.SnapshotFiles, ateompath.DurableDirTarFile) + slices.Sort(sandboxRec.SnapshotFiles) + } if len(sandboxRec.SnapshotFiles) == 0 && shouldHaveSnapshots(req) { return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonInvalidCheckpointResult, ateerrors.ActorCrashedMetadata(), errors.New("ateom reported no snapshot files for checkpoint")) } @@ -682,7 +707,7 @@ func (s *AteomHerder) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRe switch req.GetType() { case ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL: // TODO(#362): Because we do not cache the external snapshot files when upload fails, we have to mark the Actor as CRASHED. - if err := s.uploadExternalCheckpoint(ctx, req, checkpointDir, sandboxRec); err != nil { + if err := s.uploadExternalCheckpoint(ctx, req, checkpointDir, sandboxRec, streamDurableFrom); err != nil { dPersist = time.Since(tPersist) op.failedPhase = ateattr.SnapshotPhasePersist return nil, ateerrors.NewGRPCError(ctx, codes.DataLoss, ateerrors.ReasonFaileSaveSnapshot, ateerrors.ActorCrashedMetadata(), fmt.Errorf("%w: while uploading external snapshot: %w", ateerrors.ReasonFaileSaveSnapshot, err)) @@ -765,12 +790,49 @@ func shouldHaveSnapshots(req *ateletpb.CheckpointRequest) bool { return false } -func (s *AteomHerder) uploadExternalCheckpoint(ctx context.Context, req *ateletpb.CheckpointRequest, checkpointDir string, rec *sandboxAssetsRecord) error { +// canStreamDurableDirTar reports whether this checkpoint may archive the +// durable-dir volumes here, streaming straight into object storage, instead of +// having ateom stage the archive on disk for us to read back. +// +// The staged path writes the whole archive, fsyncs it, and reads it again +// before a single byte is uploaded, which at half a gibibyte costs more than +// the upload itself. Streaming overlaps the two and touches no disk. +// +// The conditions are what make the two archives interchangeable: +// +// - External checkpoints only. A local checkpoint's files stay on the node, +// so there is no upload to overlap the archive with, and moveLocalCheckpoint +// would find the file missing. +// - Micro-VM only. gVisor's archive drops the .gvisor.* files its runtime +// leaves in the durable dir, a rule that lives in ateom-gvisor; the +// archive written here would keep them. +// - A container must mount a durable-dir volume. This is deliberately the +// spec ateom is about to receive and deliberately the same predicate ateom +// applies to it (hasDurableVolumes), because the two decisions have to +// agree: disagreeing either loses the archive or invents one for a +// directory ateom would have left alone. +func canStreamDurableDirTar(req *ateletpb.CheckpointRequest, spec *ateompb.WorkloadSpec, rec *sandboxAssetsRecord) bool { + return streamDurableTarEnabled() && + req.GetType() == ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL && + atev1alpha1.SandboxClass(rec.SandboxClass) == atev1alpha1.SandboxClassMicroVM && + hasDurableDirMounts(spec) +} + +func hasDurableDirMounts(spec *ateompb.WorkloadSpec) bool { + for _, c := range spec.GetContainers() { + if len(c.GetDurableDirVolumeMounts()) > 0 { + return true + } + } + return false +} + +func (s *AteomHerder) uploadExternalCheckpoint(ctx context.Context, req *ateletpb.CheckpointRequest, checkpointDir string, rec *sandboxAssetsRecord, streamDurableFrom string) error { uri, err := resources.ParseSnapshotURI(req.GetExternalConfig().GetSnapshotUri()) if err != nil { return err } - return s.uploadSnapshot(ctx, uri, checkpointDir, rec, req.GetActorTemplateAtespace(), req.GetActorTemplateName()) + return s.uploadSnapshot(ctx, uri, checkpointDir, rec, req.GetActorTemplateAtespace(), req.GetActorTemplateName(), streamDurableFrom) } // uploadSnapshot uploads rec's snapshot files from srcDir to uri (each @@ -779,13 +841,35 @@ func (s *AteomHerder) uploadExternalCheckpoint(ctx context.Context, req *ateletp // assume every file it lists is already present. A crash mid-upload thus // leaves only orphaned files, never a manifest pointing at files that never // landed; retries overwrite the deterministic object names. -func (s *AteomHerder) uploadSnapshot(ctx context.Context, uri resources.SnapshotURI, srcDir string, rec *sandboxAssetsRecord, templateAtespace, templateName string) error { +// +// A non-empty streamDurableFrom names the durable-dir mounts directory to +// archive on the fly in place of reading DurableDirTarFile out of srcDir, where +// it was never written; see canStreamDurableDirTar. +func (s *AteomHerder) uploadSnapshot(ctx context.Context, uri resources.SnapshotURI, srcDir string, rec *sandboxAssetsRecord, templateAtespace, templateName, streamDurableFrom string) error { g, gCtx := errgroup.WithContext(ctx) for _, fileName := range rec.SnapshotFiles { + objectName := fileName + ".zstd" + if streamDurableFrom != "" && fileName == ateompath.DurableDirTarFile { + g.Go(func() error { + objectURI, err := uri.ObjectURI(objectName) + if err != nil { + return fmt.Errorf("while addressing %s in GCS: %w", fileName, err) + } + size, err := streamDurableDirTar(gCtx, streamDurableFrom, func(r io.Reader) error { + return ategcs.SendReaderToGCSWithZstd(gCtx, s.gcsClient, objectURI, r) + }) + if err != nil { + return fmt.Errorf("while archiving %s to GCS: %w", fileName, err) + } + emitSnapshotSize(ctx, fileName, size, templateAtespace, templateName) + return nil + }) + continue + } local := filepath.Join(srcDir, fileName) recordSnapshotSize(ctx, fileName, local, templateAtespace, templateName) g.Go(func() error { - objectURI, err := uri.ObjectURI(fileName + ".zstd") + objectURI, err := uri.ObjectURI(objectName) if err != nil { return fmt.Errorf("while addressing %s in GCS: %w", fileName, err) } @@ -917,7 +1001,9 @@ func (s *AteomHerder) uploadLocalCheckpointDir(ctx context.Context, req *ateletp } } - return rec.SandboxClass, s.uploadSnapshot(ctx, uri, localDir, rec, req.GetActorTemplateAtespace(), req.GetActorTemplateName()) + // Never streamed: this uploads a checkpoint already staged on disk by an + // earlier pause, so there is nothing left to overlap. + return rec.SandboxClass, s.uploadSnapshot(ctx, uri, localDir, rec, req.GetActorTemplateAtespace(), req.GetActorTemplateName(), "") } // narrowFullCaptureToData rewrites rec so a FULL capture uploads as a DATA diff --git a/cmd/ateom-gvisor/durable_test.go b/cmd/ateom-gvisor/durable_test.go new file mode 100644 index 0000000000..7ddad3871c --- /dev/null +++ b/cmd/ateom-gvisor/durable_test.go @@ -0,0 +1,40 @@ +//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 ( + "testing" + + ateompb "github.com/agent-substrate/substrate/internal/proto/ateompb" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// TestCheckpointWorkloadRejectsSkipDurableDirTar pins the rejection, and pins +// that it happens before the sandbox is touched: a zero-value service has no +// actor to check point, so reaching any of the real work would not return +// InvalidArgument. +func TestCheckpointWorkloadRejectsSkipDurableDirTar(t *testing.T) { + s := &AteomService{} + _, err := s.CheckpointWorkload(t.Context(), &ateompb.CheckpointWorkloadRequest{ + ActorUid: "actor-1", + SkipDurableDirTar: true, + }) + if status.Code(err) != codes.InvalidArgument { + t.Fatalf("CheckpointWorkload err = %v (code %v), want InvalidArgument", err, status.Code(err)) + } +} diff --git a/cmd/ateom-gvisor/main.go b/cmd/ateom-gvisor/main.go index ff7835cc27..b9733b7917 100644 --- a/cmd/ateom-gvisor/main.go +++ b/cmd/ateom-gvisor/main.go @@ -735,6 +735,15 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload // 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) { + // Rejected rather than ignored, and before the sandbox is touched: honoring + // it means atelet archives the durable dir instead, and atelet does not know + // to drop the .gvisor.* files this runtime leaves there (see + // tarDurableVolumes). Silently archiving them anyway would restore one + // sandbox's internals into the next. + if req.GetSkipDurableDirTar() { + return nil, status.Error(codes.InvalidArgument, "skip_durable_dir_tar is not supported by the gVisor runtime") + } + s.lock.Lock() defer s.lock.Unlock() diff --git a/cmd/ateom-microvm/checkpoint.go b/cmd/ateom-microvm/checkpoint.go index 8a4c8c75a7..1b935487a0 100644 --- a/cmd/ateom-microvm/checkpoint.go +++ b/cmd/ateom-microvm/checkpoint.go @@ -154,7 +154,11 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec return err }) } - if durable { + // Skipped when atelet has taken the durable-dir tar over to stream it + // straight into object storage; the directory is left in place for it, and + // terminateWorkload below is what makes a later archive just as coherent as + // one taken here. + if durable && !req.GetSkipDurableDirTar() { g.Go(func() error { t := time.Now() if err := tarDurableVolumes(gctx, ateompath.DurableDirVolumeMountsDir(actorUID), checkpointDir); err != nil { diff --git a/internal/proto/ateompb/ateom.pb.go b/internal/proto/ateompb/ateom.pb.go index 2d4407bca7..48112d9b9f 100644 --- a/internal/proto/ateompb/ateom.pb.go +++ b/internal/proto/ateompb/ateom.pb.go @@ -1107,9 +1107,20 @@ type CheckpointWorkloadRequest struct { // atelet fetched it to (see RunWorkloadRequest). Empty for gVisor. RuntimeAssetPaths map[string]string `protobuf:"bytes,9,rep,name=runtime_asset_paths,json=runtimeAssetPaths,proto3" json:"runtime_asset_paths,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"` // What content to include in the checkpoint. - Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` - unknownFields protoimpl.UnknownFields - sizeCache protoimpl.SizeCache + Scope SnapshotScope `protobuf:"varint,10,opt,name=scope,proto3,enum=ateom.SnapshotScope" json:"scope,omitempty"` + // skip_durable_dir_tar tells ateom not to archive the actor's durable-dir + // volumes, and to leave the directory in place for atelet to archive itself. + // atelet sets this when it will stream the archive straight into object + // storage instead of staging it on disk, which is only safe because ateom + // terminates the workload before returning: nothing writes the directory + // afterwards, so an archive taken later is as coherent as one taken inside + // the paused window. + // + // Every other snapshot file is unaffected; ateom still reports the set it + // wrote and atelet appends the durable-dir tar to it. + SkipDurableDirTar bool `protobuf:"varint,11,opt,name=skip_durable_dir_tar,json=skipDurableDirTar,proto3" json:"skip_durable_dir_tar,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache } func (x *CheckpointWorkloadRequest) Reset() { @@ -1212,6 +1223,13 @@ func (x *CheckpointWorkloadRequest) GetScope() SnapshotScope { return SnapshotScope_SNAPSHOT_SCOPE_UNSPECIFIED } +func (x *CheckpointWorkloadRequest) GetSkipDurableDirTar() bool { + if x != nil { + return x.SkipDurableDirTar + } + return false +} + type CheckpointWorkloadResponse struct { state protoimpl.MessageState `protogen:"open.v1"` // snapshot_files lists the files ateom wrote into the checkpoint directory @@ -1912,7 +1930,7 @@ const file_ateom_proto_rawDesc = "" + "\rHTTPGetAction\x12\x12\n" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x12\n" + "\x04port\x18\x02 \x01(\x05R\x04port\"\x15\n" + - "\x13RunWorkloadResponse\"\xa1\x04\n" + + "\x13RunWorkloadResponse\"\xd2\x04\n" + "\x19CheckpointWorkloadRequest\x12\x1a\n" + "\batespace\x18\x01 \x01(\tR\batespace\x12\x1d\n" + "\n" + @@ -1926,7 +1944,8 @@ const file_ateom_proto_rawDesc = "" + "\fsnapshot_uri\x18\b \x01(\tR\vsnapshotUri\x12g\n" + "\x13runtime_asset_paths\x18\t \x03(\v27.ateom.CheckpointWorkloadRequest.RuntimeAssetPathsEntryR\x11runtimeAssetPaths\x12*\n" + "\x05scope\x18\n" + - " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x1aD\n" + + " \x01(\x0e2\x14.ateom.SnapshotScopeR\x05scope\x12/\n" + + "\x14skip_durable_dir_tar\x18\v \x01(\bR\x11skipDurableDirTar\x1aD\n" + "\x16RuntimeAssetPathsEntry\x12\x10\n" + "\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" + "\x05value\x18\x02 \x01(\tR\x05value:\x028\x01\"C\n" + diff --git a/internal/proto/ateompb/ateom.proto b/internal/proto/ateompb/ateom.proto index 2a3a97960e..8e7c80fb8a 100644 --- a/internal/proto/ateompb/ateom.proto +++ b/internal/proto/ateompb/ateom.proto @@ -285,6 +285,18 @@ message CheckpointWorkloadRequest { // What content to include in the checkpoint. SnapshotScope scope = 10; + + // skip_durable_dir_tar tells ateom not to archive the actor's durable-dir + // volumes, and to leave the directory in place for atelet to archive itself. + // atelet sets this when it will stream the archive straight into object + // storage instead of staging it on disk, which is only safe because ateom + // terminates the workload before returning: nothing writes the directory + // afterwards, so an archive taken later is as coherent as one taken inside + // the paused window. + // + // Every other snapshot file is unaffected; ateom still reports the set it + // wrote and atelet appends the durable-dir tar to it. + bool skip_durable_dir_tar = 11; } message CheckpointWorkloadResponse { diff --git a/internal/tarutil/tarutil.go b/internal/tarutil/tarutil.go index 0a89ad746b..0ed8aaabea 100644 --- a/internal/tarutil/tarutil.go +++ b/internal/tarutil/tarutil.go @@ -109,8 +109,32 @@ func CreateFiltered(ctx context.Context, tarPath, srcDir string, skip SkipFunc) } defer f.Close() + if err := writeArchive(ctx, f, srcDir, skip); err != nil { + return err + } + // Handed to atelet for upload as soon as we return, so flush to disk rather + // than trusting the page cache to outlive us. + if err := f.Sync(); err != nil { + return fmt.Errorf("syncing tar %q: %w", tarPath, err) + } + return nil +} + +// CreateTo is Create writing to an arbitrary destination instead of a file, for +// callers that pipe the archive somewhere (an object-storage upload) rather +// than staging it on disk. It does not close w. +// +// Errors from w surface here, so a consumer that stops reading early must make +// its own reason for stopping the more informative one. +func CreateTo(ctx context.Context, w io.Writer, srcDir string, skip SkipFunc) error { + return writeArchive(ctx, w, srcDir, skip) +} + +// writeArchive streams srcDir to w through a pooled buffer, flushed before +// returning so the whole archive has reached w by then. +func writeArchive(ctx context.Context, w io.Writer, srcDir string, skip SkipFunc) error { bw := tarWriterPool.Get().(*bufio.Writer) - bw.Reset(f) + bw.Reset(w) defer func() { bw.Reset(nil) tarWriterPool.Put(bw) @@ -120,17 +144,12 @@ func CreateFiltered(ctx context.Context, tarPath, srcDir string, skip SkipFunc) return err } if err := tw.Close(); err != nil { - return fmt.Errorf("closing tar %q: %w", tarPath, err) + return fmt.Errorf("closing tar of %q: %w", srcDir, err) } - // The buffer has to reach the file before the sync below, or the sync - // durably persists a truncated archive. + // Everything buffered has to reach w before we report success: the caller + // may sync, or finish an upload, the moment we return. if err := bw.Flush(); err != nil { - return fmt.Errorf("flushing tar %q: %w", tarPath, err) - } - // Durable-dir tars are handed to atelet for upload as soon as we return, so - // flush to disk rather than trusting the page cache to outlive us. - if err := f.Sync(); err != nil { - return fmt.Errorf("syncing tar %q: %w", tarPath, err) + return fmt.Errorf("flushing tar of %q: %w", srcDir, err) } return nil } diff --git a/internal/tarutil/tarutil_test.go b/internal/tarutil/tarutil_test.go index f85447ccfe..5674c8abd4 100644 --- a/internal/tarutil/tarutil_test.go +++ b/internal/tarutil/tarutil_test.go @@ -18,6 +18,7 @@ package tarutil import ( "archive/tar" + "bytes" "errors" "net" "os" @@ -316,6 +317,64 @@ func TestCreateFilteredSkipsSubtree(t *testing.T) { } } +// TestCreateToMatchesCreate pins the two entry points on one archive: callers +// that stream the tar somewhere instead of staging it on disk must produce the +// same bytes, or a snapshot taken one way would not restore like the other. +func TestCreateToMatchesCreate(t *testing.T) { + src := t.TempDir() + if err := os.MkdirAll(filepath.Join(src, "vol", "nested"), 0o755); err != nil { + t.Fatalf("mkdir: %v", err) + } + // Larger than the pooled stream buffer, so the flush CreateTo owes its + // caller actually has something left to flush. + if err := os.WriteFile(filepath.Join(src, "vol", "big"), []byte(strings.Repeat("substrate", 1<<14)), 0o644); err != nil { + t.Fatalf("writing big: %v", err) + } + if err := os.WriteFile(filepath.Join(src, "vol", "nested", "small"), []byte("hello"), 0o600); err != nil { + t.Fatalf("writing small: %v", err) + } + if err := os.Symlink("big", filepath.Join(src, "vol", "link")); err != nil { + t.Fatalf("symlink: %v", err) + } + + tarPath := filepath.Join(t.TempDir(), "staged.tar") + if err := Create(t.Context(), tarPath, src); err != nil { + t.Fatalf("Create: %v", err) + } + want, err := os.ReadFile(tarPath) + if err != nil { + t.Fatalf("reading staged tar: %v", err) + } + + var got bytes.Buffer + if err := CreateTo(t.Context(), &got, src, nil); err != nil { + t.Fatalf("CreateTo: %v", err) + } + if !bytes.Equal(got.Bytes(), want) { + t.Errorf("CreateTo wrote %d bytes, Create wrote %d; archives differ", got.Len(), len(want)) + } +} + +// TestCreateToReportsWriterErrors checks a destination that stops accepting +// bytes fails the archive rather than silently truncating it. +func TestCreateToReportsWriterErrors(t *testing.T) { + src := t.TempDir() + if err := os.WriteFile(filepath.Join(src, "f"), []byte(strings.Repeat("x", 1<<20)), 0o644); err != nil { + t.Fatalf("writing f: %v", err) + } + + sentinel := errors.New("destination closed") + err := CreateTo(t.Context(), errWriter{err: sentinel}, src, nil) + if !errors.Is(err, sentinel) { + t.Errorf("CreateTo err = %v, want it to carry %v", err, sentinel) + } +} + +// errWriter fails every write, standing in for an upload that died mid-archive. +type errWriter struct{ err error } + +func (w errWriter) Write([]byte) (int, error) { return 0, w.err } + // TestRoundTripSpecialModeBits pins setuid/setgid/sticky across a round trip. // FileMode.Perm() silently drops them, so without this the archive would record // bits that extraction throws away — and a setgid data directory would come back