atelet: sweep crash-leftover temp files at startup

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 {

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),

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
+// ".<hex>.json.tmp-<rand>"; finished records are "<hex>.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 ".<hex>.json.tmp-<rand>"; finished records
-	// are "<hex>.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

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")
 	}

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)
 	}

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))
+		}
+	}
+}