diff --git a/cmd/atelet/internal/ategcs/objects.go b/cmd/atelet/internal/ategcs/objects.go index 20598b682..f72dffb5f 100644 --- a/cmd/atelet/internal/ategcs/objects.go +++ b/cmd/atelet/internal/ategcs/objects.go @@ -28,12 +28,20 @@ import ( "time" "github.com/agent-substrate/substrate/internal/ateerrors" + "github.com/agent-substrate/substrate/internal/startupsweep" "github.com/klauspost/compress/zstd" "go.opentelemetry.io/otel" ) var tracer = otel.Tracer("ategcs") +// uploadTempFilePrefix is the filename prefix for staging files created by +// sendBufferedZstd in os.TempDir(). S3/rustfs requires a seekable body with a +// known Content-Length, so the compressed payload is buffered to a temp file +// before upload. If atelet crashes mid-upload these files are never removed; +// SweepEntries registers them for cleanup at next startup. +const uploadTempFilePrefix = "substrate-upload-compress-" + type ObjectStorage interface { GetObject(ctx context.Context, bucket, object string) (io.ReadCloser, error) PutObject(ctx context.Context, bucket, object string, reader io.Reader) error @@ -184,7 +192,7 @@ func sendZstd(ctx context.Context, client ObjectStorage, gsURL string, content i // Used for backends (S3/rustfs) whose PutObject needs a seekable body to sign and // set Content-Length; the streaming counterpart is sendStreamingZstd. func sendBufferedZstd(ctx context.Context, client ObjectStorage, bucket, object string, content io.Reader, tStart time.Time) error { - tmpFile, err := os.CreateTemp("", "substrate-upload-compress-") + tmpFile, err := os.CreateTemp("", uploadTempFilePrefix) if err != nil { return fmt.Errorf("while creating temp compress file: %w", err) } @@ -296,6 +304,14 @@ func FetchLocalFileFromGCSWithZstd(ctx context.Context, client ObjectStorage, gs return nil } +// SweepEntries returns a registration function for the Sweeper that cleans up +// upload staging files left behind by a crash. +func SweepEntries() func(*startupsweep.Sweeper) { + return func(s *startupsweep.Sweeper) { + s.Register("GCS upload temp files", os.TempDir(), uploadTempFilePrefix+"*", os.Remove) + } +} + func fetchFromGCSWithZstd(ctx context.Context, client ObjectStorage, gsURL string, out io.Writer) (err error) { bucket, object, err := parseGCSURL(gsURL) if err != nil { diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 654a3c76f..6f96200e5 100644 --- a/cmd/atelet/main.go +++ b/cmd/atelet/main.go @@ -33,6 +33,7 @@ import ( "cloud.google.com/go/storage" "github.com/agent-substrate/substrate/cmd/atelet/internal/ategcs" "github.com/agent-substrate/substrate/internal/ateerrors" + "github.com/agent-substrate/substrate/internal/startupsweep" "github.com/agent-substrate/substrate/internal/ateinterceptors" "github.com/agent-substrate/substrate/internal/ateompath" "github.com/agent-substrate/substrate/internal/credbundle" @@ -121,6 +122,11 @@ func main() { } } + sweeper := startupsweep.New() + sweeper.Add(ategcs.SweepEntries()) + sweeper.Add(imagecache.SweepEntries(*imageCacheDir)) + sweeper.Sweep(ctx) + imageCache, err := imagecache.New(*imageCacheDir, imagecache.WithAuthenticator(gcpRegistryAuthn), imagecache.WithLocalhostRegistryReplacement(*localhostRegistryReplacement), diff --git a/internal/imagecache/imagecache.go b/internal/imagecache/imagecache.go index d474d0df3..9c8f419ec 100644 --- a/internal/imagecache/imagecache.go +++ b/internal/imagecache/imagecache.go @@ -61,6 +61,7 @@ import ( "sync" "time" + "github.com/agent-substrate/substrate/internal/startupsweep" "github.com/google/go-containerregistry/pkg/authn" "github.com/google/go-containerregistry/pkg/name" v1 "github.com/google/go-containerregistry/pkg/v1" @@ -189,14 +190,28 @@ func New(root string, opts ...Option) (*Store, error) { return nil, fmt.Errorf("while reading image cache version marker: %w", err) } - if err := s.sweepTempDirs(); err != nil { - return nil, err - } return s, nil } -func (s *Store) layersDir() string { return filepath.Join(s.root, "layers", "sha256") } -func (s *Store) manifestsDir() string { return filepath.Join(s.root, "manifests", "sha256") } +// SweepEntries returns a registration function for the Sweeper that cleans up +// temp dirs and manifest files left behind by a crash in the cache at root. +// ".tmp-*": an unpack in flight at crash. ".rm-*": a layer retirement renamed +// aside but not yet removed. ".*": writeRecord's temp files are +// "..json.tmp-"; finished records are ".json", so a leading +// dot alone identifies an interrupted manifest write. +func SweepEntries(root string) func(*startupsweep.Sweeper) { + return func(s *startupsweep.Sweeper) { + s.Register("image cache orphaned layer dirs", layersDirPath(root), ".tmp-*", RemoveAllWritable) + s.Register("image cache retired layer dirs", layersDirPath(root), retiredPrefix+"*", RemoveAllWritable) + s.Register("image cache orphaned manifest files", manifestsDirPath(root), ".*", os.Remove) + } +} + +func layersDirPath(root string) string { return filepath.Join(root, "layers", "sha256") } +func manifestsDirPath(root string) string { return filepath.Join(root, "manifests", "sha256") } + +func (s *Store) layersDir() string { return layersDirPath(s.root) } +func (s *Store) manifestsDir() string { return manifestsDirPath(s.root) } func (s *Store) layerDir(diffID v1.Hash) string { return filepath.Join(s.root, "layers", diffID.Algorithm, diffID.Hex) @@ -206,53 +221,6 @@ func (s *Store) recordPath(digest v1.Hash) string { return filepath.Join(s.root, "manifests", digest.Algorithm, digest.Hex+".json") } -// sweepTempDirs removes unpack temp dirs, retired layer dirs, and -// manifest-record temp files orphaned by a crash. A layer dir without the -// temp/retired prefix and a record without a leading dot are always -// complete (both are moved into place with a single rename), so this is -// the only recovery the pool needs. -func (s *Store) sweepTempDirs() error { - entries, err := os.ReadDir(s.layersDir()) - if err != nil { - return fmt.Errorf("while listing layer pool: %w", err) - } - swept := 0 - for _, e := range entries { - // ".tmp-": an unpack in flight at crash. ".rm-": a layer retirement - // renamed aside but not yet removed. Either way, unreferenced - // garbage. - if !strings.HasPrefix(e.Name(), ".tmp-") && !strings.HasPrefix(e.Name(), retiredPrefix) { - continue - } - p := filepath.Join(s.layersDir(), e.Name()) - slog.Info("Image cache sweeping orphaned layer dir", slog.String("dir", e.Name())) - if err := RemoveAllWritable(p); err != nil { - return fmt.Errorf("while sweeping orphaned layer dir %q: %w", p, err) - } - swept++ - } - if swept > 0 { - slog.Info("Image cache startup sweep removed orphaned layer dirs", slog.Int("count", swept)) - } - - // writeRecord's temp files are "..json.tmp-"; finished records - // are ".json", so a leading dot alone identifies an orphan. - records, err := os.ReadDir(s.manifestsDir()) - if err != nil { - return fmt.Errorf("while listing manifest records: %w", err) - } - for _, e := range records { - if !strings.HasPrefix(e.Name(), ".") { - continue - } - p := filepath.Join(s.manifestsDir(), e.Name()) - if err := os.Remove(p); err != nil { - return fmt.Errorf("while sweeping orphaned manifest temp file %q: %w", p, err) - } - } - return nil -} - // EnsureImage makes ref's image available in the pool and returns its config // and ordered layer directories. Digest refs hit the cache with no network // I/O; tag refs cost one HEAD request to resolve the tag to a manifest diff --git a/internal/imagecache/imagecache_test.go b/internal/imagecache/imagecache_test.go index f11748775..713adf3d9 100644 --- a/internal/imagecache/imagecache_test.go +++ b/internal/imagecache/imagecache_test.go @@ -27,6 +27,7 @@ import ( "slices" "testing" + "github.com/agent-substrate/substrate/internal/startupsweep" "github.com/google/go-containerregistry/pkg/name" "github.com/google/go-containerregistry/pkg/registry" v1 "github.com/google/go-containerregistry/pkg/v1" @@ -227,9 +228,12 @@ func TestNew_RecoveryAndVersioning(t *testing.T) { if err := os.WriteFile(record, []byte("{}"), 0o600); err != nil { t.Fatalf("planting record: %v", err) } + sw := startupsweep.New() + sw.Add(SweepEntries(root)) if _, err := New(root); err != nil { t.Fatalf("New(recovery): %v", err) } + sw.Sweep(context.Background()) if _, err := os.Stat(orphan); !os.IsNotExist(err) { t.Errorf("orphaned temp dir survived recovery") } diff --git a/internal/imagecache/retire_test.go b/internal/imagecache/retire_test.go index fe45c5a3c..94fcefda3 100644 --- a/internal/imagecache/retire_test.go +++ b/internal/imagecache/retire_test.go @@ -24,6 +24,7 @@ import ( "testing" "time" + "github.com/agent-substrate/substrate/internal/startupsweep" v1 "github.com/google/go-containerregistry/pkg/v1" ) @@ -101,9 +102,12 @@ func TestNewSweepsRetiredDirs(t *testing.T) { if err := os.Chmod(filepath.Join(retired, "fs", "ro"), 0o500); err != nil { t.Fatal(err) } + sw := startupsweep.New() + sw.Add(SweepEntries(root)) if _, err := New(root); err != nil { t.Fatalf("New (recovery): %v", err) } + sw.Sweep(context.Background()) if _, err := os.Stat(retired); !os.IsNotExist(err) { t.Errorf("retired dir not swept at startup: %v", err) } diff --git a/internal/startupsweep/sweep.go b/internal/startupsweep/sweep.go new file mode 100644 index 000000000..5b55f0fda --- /dev/null +++ b/internal/startupsweep/sweep.go @@ -0,0 +1,85 @@ +// 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 startupsweep provides a registry for cleaning up files and +// directories orphaned by a previous crash. Callers register (dir, glob, +// removeFn) tuples at init time; a single Sweep call removes all matches. +package startupsweep + +import ( + "context" + "errors" + "log/slog" + "os" + "path/filepath" +) + +type entry struct { + label string + dir string + pattern string + removeFn func(string) error +} + +// Sweeper accumulates glob-based cleanup tasks and runs them at startup. +type Sweeper struct { + entries []entry +} + +// New returns an empty Sweeper. +func New() *Sweeper { return &Sweeper{} } + +// Add calls fn with the Sweeper so the caller can register its entries. +func (s *Sweeper) Add(fn func(*Sweeper)) { fn(s) } + +// Register adds a cleanup entry: all paths matching filepath.Join(dir, pattern) +// will be passed to removeFn during Sweep. label is used in the summary log +// line emitted when at least one match is removed. Use os.Remove for plain +// files and os.RemoveAll (or a writable variant) for directories. +// Register is deliberately a no-op on a nil Sweeper. +func (s *Sweeper) Register(label, dir, pattern string, removeFn func(string) error) { + if s == nil { + return + } + s.entries = append(s.entries, entry{label, dir, pattern, removeFn}) +} + +// Sweep removes all registered orphaned paths. Each match is logged before +// removal. Failures are logged as warnings and do not abort the sweep; +// missing files are silently skipped. +func (s *Sweeper) Sweep(ctx context.Context) { + for _, e := range s.entries { + matches, err := filepath.Glob(filepath.Join(e.dir, e.pattern)) + if err != nil { + // filepath.Glob only errors on malformed patterns — treat as a bug. + slog.WarnContext(ctx, "Startup sweep: invalid glob pattern", + slog.String("dir", e.dir), slog.String("pattern", e.pattern), slog.Any("err", err)) + continue + } + swept := 0 + for _, m := range matches { + slog.InfoContext(ctx, "Startup sweep removing orphan", slog.String("path", m)) + if err := e.removeFn(m); err != nil && !errors.Is(err, os.ErrNotExist) { + slog.WarnContext(ctx, "Startup sweep failed to remove orphan", + slog.String("path", m), slog.Any("err", err)) + } else { + swept++ + } + } + if swept > 0 { + slog.InfoContext(ctx, "Startup sweep removed orphans", + slog.String("label", e.label), slog.Int("count", swept)) + } + } +} diff --git a/internal/startupsweep/sweep_test.go b/internal/startupsweep/sweep_test.go new file mode 100644 index 000000000..735eb1196 --- /dev/null +++ b/internal/startupsweep/sweep_test.go @@ -0,0 +1,149 @@ +// 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 startupsweep + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +// TestSweepRemovesMatchingFiles checks that only files matching +// the registered pattern are deleted. +func TestSweepRemovesMatchingFiles(t *testing.T) { + // Given + dir := t.TempDir() + orphan := filepath.Join(dir, ".tmp-orphan") + keep := filepath.Join(dir, "keep") + writeFile(t, orphan) + writeFile(t, keep) + + // When + sw := New() + sw.Register("test files", dir, ".tmp-*", os.Remove) + sw.Sweep(context.Background()) + + // Then + if _, err := os.Stat(orphan); !os.IsNotExist(err) { + t.Errorf("orphan file was not removed") + } + if _, err := os.Stat(keep); err != nil { + t.Errorf("non-matching file was removed: %v", err) + } +} + +// TestSweepRemovesMatchingDirs checks that matching directories +// are deleted correctly and that not matching dirs are kept. +func TestSweepRemovesMatchingDirs(t *testing.T) { + // Given + dir := t.TempDir() + orphan := filepath.Join(dir, ".tmp-orphan-dir") + keep := filepath.Join(dir, "keep-dir") + makeDir(t, orphan) + makeDir(t, keep) + + // When + sw := New() + sw.Register("test dirs", dir, ".tmp-*", os.RemoveAll) + sw.Sweep(context.Background()) + + // Then + if _, err := os.Stat(orphan); !os.IsNotExist(err) { + t.Errorf("orphan dir was not removed") + } + if _, err := os.Stat(keep); err != nil { + t.Errorf("non-matching dir was removed: %v", err) + } +} + +// TestSweepNoMatchIsNoop checks that no-op when there are no matching +// files or directories to be swept. +func TestSweepNoMatchIsNoop(t *testing.T) { + // Given + dir := t.TempDir() + keep := filepath.Join(dir, "keep") + writeFile(t, keep) + + // When + sw := New() + sw.Register("test", dir, ".tmp-*", os.Remove) + sw.Sweep(context.Background()) + + // Then + if _, err := os.Stat(keep); err != nil { + t.Errorf("non-matching file was removed: %v", err) + } +} + +// TestAddDelegatesRegistration validates that registration via Add +// removes matching files. +func TestAddDelegatesRegistration(t *testing.T) { + // Given + dir := t.TempDir() + orphan := filepath.Join(dir, ".tmp-via-add") + writeFile(t, orphan) + + register := func(s *Sweeper) { + s.Register("add test", dir, ".tmp-*", os.Remove) + } + + // When + sw := New() + sw.Add(register) + sw.Sweep(context.Background()) + + // Then + if _, err := os.Stat(orphan); !os.IsNotExist(err) { + t.Errorf("orphan registered via Add was not removed") + } +} + +// TestRegisterOnNilSweeperIsNoop checks that registering on a nil +// Sweeper does not panic. +func TestRegisterOnNilSweeperIsNoop(t *testing.T) { + // Given + var sw *Sweeper + + // Then + sw.Register("nil test", t.TempDir(), "*", os.Remove) // must not panic +} + +// TestSweepMissingDirIsNoop checks that sweeping a non-existent +// directory does not panic. +func TestSweepMissingDirIsNoop(t *testing.T) { + // Given + sw := New() + + // When + sw.Register("missing dir", "/nonexistent/dir/that/cannot/exist", "*.tmp", os.Remove) + + // Then + sw.Sweep(context.Background()) // must not panic or error +} + +func writeFile(t *testing.T, path string) { + t.Helper() + if err := os.WriteFile(path, []byte("x"), 0o600); err != nil { + t.Fatal(err) + } +} + +func makeDir(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(path, 0o700); err != nil { + t.Fatal(err) + } +}