feat(orchestrator): template-gc dry-run mark tool (storage ObjectLister + mark engine)

packages/orchestrator/cmd/template-gc/main.go

@@ -0,0 +1,136 @@
+// Command template-gc is the DRY-RUN mark phase of template-storage garbage
+// collection. It reads a list of live-root build IDs (exported separately from
+// the control plane), expands each root's header-mapping ancestor closure,
+// scans the template bucket, and writes two manifests: the builds to KEEP and
+// the build prefixes that are unreferenced and older than the grace period and
+// therefore safe to delete.
+//
+// It NEVER deletes anything. Physical deletion is a separate, operator-gated
+// step that consumes the candidate manifest after quarantine and revalidation.
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/e2b-dev/infra/packages/orchestrator/cmd/internal/cmdutil"
+	"github.com/e2b-dev/infra/packages/orchestrator/internal/templategc"
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+)
+
+func main() {
+	storagePath := flag.String("storage", "", "template storage: local path or gs://bucket")
+	rootsFile := flag.String("roots", "", "path to a JSON file: array of live-root build IDs (exported separately)")
+	graceDays := flag.Int("grace-days", 60, "a build prefix is a candidate only if every object under it is older than this many days")
+	outDir := flag.String("out", ".", "directory to write keep.json and candidates.json manifests")
+	flag.Parse()
+
+	if *storagePath == "" || *rootsFile == "" {
+		fmt.Fprintln(os.Stderr, "both -storage and -roots are required")
+		os.Exit(2)
+	}
+
+	ctx := context.Background()
+	if err := run(ctx, *storagePath, *rootsFile, *graceDays, *outDir); err != nil {
+		fmt.Fprintf(os.Stderr, "template-gc mark failed (no objects were deleted): %v\n", err)
+		os.Exit(1)
+	}
+}
+
+func run(ctx context.Context, storagePath, rootsFile string, graceDays int, outDir string) error {
+	roots, err := readRoots(rootsFile)
+	if err != nil {
+		return fmt.Errorf("read roots: %w", err)
+	}
+
+	spec, err := cmdutil.TemplateSpec(storagePath)
+	if err != nil {
+		return fmt.Errorf("resolve storage spec: %w", err)
+	}
+
+	provider, err := storage.NewProvider(ctx, spec)
+	if err != nil {
+		return fmt.Errorf("storage provider: %w", err)
+	}
+
+	objects, err := storage.ListObjects(ctx, provider, "")
+	if err != nil {
+		return fmt.Errorf("list bucket (fail closed; cannot treat an unlistable bucket as empty): %w", err)
+	}
+
+	manifest, err := templategc.Mark(ctx, templategc.StorageAncestorReader{Storage: provider}, templategc.MarkInput{
+		LiveRoots: roots,
+		Objects:   objects,
+		Now:       time.Now(),
+		Grace:     time.Duration(graceDays) * 24 * time.Hour,
+	})
+	if err != nil {
+		return err
+	}
+
+	if err := writeJSON(outDir, "keep.json", manifest.Keep); err != nil {
+		return err
+	}
+	if err := writeJSON(outDir, "candidates.json", manifest.Candidates); err != nil {
+		return err
+	}
+	if err := writeJSON(outDir, "too-new.json", manifest.TooNew); err != nil {
+		return err
+	}
+	if err := writeJSON(outDir, "unknown-keys.json", manifest.UnknownKeys); err != nil {
+		return err
+	}
+
+	var candBytes int64
+	for _, c := range manifest.Candidates {
+		candBytes += c.TotalBytes
+	}
+
+	fmt.Printf("mark complete (DRY RUN, nothing deleted)\n")
+	fmt.Printf("  live roots:      %d\n", len(roots))
+	fmt.Printf("  keep builds:     %d\n", len(manifest.Keep))
+	fmt.Printf("  candidates:      %d (%d bytes)\n", len(manifest.Candidates), candBytes)
+	fmt.Printf("  held (too new):  %d\n", len(manifest.TooNew))
+	fmt.Printf("  unknown keys:    %d\n", len(manifest.UnknownKeys))
+	fmt.Printf("  manifests in:    %s\n", outDir)
+
+	return nil
+}
+
+func readRoots(path string) ([]uuid.UUID, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	var raw []string
+	if err := json.Unmarshal(data, &raw); err != nil {
+		return nil, fmt.Errorf("roots file must be a JSON array of build-ID strings: %w", err)
+	}
+
+	roots := make([]uuid.UUID, 0, len(raw))
+	for _, s := range raw {
+		id, err := uuid.Parse(s)
+		if err != nil {
+			return nil, fmt.Errorf("invalid root build ID %q: %w", s, err)
+		}
+		roots = append(roots, id)
+	}
+
+	return roots, nil
+}
+
+func writeJSON(dir, name string, v any) error {
+	data, err := json.MarshalIndent(v, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	return os.WriteFile(dir+"/"+name, data, 0o644)
+}

packages/orchestrator/internal/templategc/mark.go

@@ -0,0 +1,148 @@
+// Package templategc implements the mark phase of template-storage garbage
+// collection. It is dry-run only: it computes which build prefixes are
+// unreferenced and old enough to be deletion candidates, and NEVER deletes
+// anything. Physical deletion is a separate, operator-gated step.
+package templategc
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"time"
+
+	"github.com/google/uuid"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
+)
+
+// AncestorReader resolves a build's header-mapping ancestor closure (the build
+// itself plus every build its memfile/rootfs headers reference). An error means
+// the dependency set is unknown; the marker then fails closed.
+type AncestorReader interface {
+	AncestorBuilds(ctx context.Context, buildID uuid.UUID) ([]uuid.UUID, error)
+}
+
+// Manifest is the result of a mark run. It is advisory: Candidates lists the
+// build prefixes that are safe to delete under the inputs given, but nothing is
+// deleted here.
+type Manifest struct {
+	// Keep is every build ID still reachable from a live root's header closure.
+	Keep []uuid.UUID
+	// Candidates are build prefixes not in Keep and whose every object is older
+	// than the grace period — safe to delete.
+	Candidates []BuildPrefix
+	// TooNew are unreferenced builds held back because at least one object under
+	// the prefix is newer than the grace cutoff (possibly an in-flight upload).
+	TooNew []BuildPrefix
+	// UnknownKeys are object keys that do not sit under a UUID build prefix;
+	// surfaced for human review, never deleted.
+	UnknownKeys []string
+}
+
+// BuildPrefix is a build's object-storage prefix and its aggregate stats.
+type BuildPrefix struct {
+	BuildID      uuid.UUID
+	ObjectCount  int
+	TotalBytes   int64
+	NewestObject time.Time
+}
+
+// MarkInput is everything the pure Mark computation needs.
+type MarkInput struct {
+	// LiveRoots are the build IDs that must be kept: the union of DB active
+	// templates' builds, in-progress builds, and runtime-sandbox-pinned builds.
+	// The caller collects these from all sources and fails the whole run if any
+	// source could not be read.
+	LiveRoots []uuid.UUID
+	// Objects is the full, paginated listing of the template bucket.
+	Objects []storage.ObjectInfo
+	// Now and Grace define the age cutoff: a build is only a candidate when every
+	// object under its prefix predates Now.Add(-Grace).
+	Now   time.Time
+	Grace time.Duration
+}
+
+// Mark computes the deletion manifest. It fails closed: if any live root's
+// ancestor closure cannot be resolved, it returns an error and no candidates, so
+// an unreadable header can never cause a reachable build to be marked deletable.
+func Mark(ctx context.Context, reader AncestorReader, in MarkInput) (*Manifest, error) {
+	keep := make(map[uuid.UUID]struct{})
+	for _, root := range in.LiveRoots {
+		closure, err := reader.AncestorBuilds(ctx, root)
+		if err != nil {
+			return nil, fmt.Errorf("resolve ancestor closure of live root %s: %w", root, err)
+		}
+
+		for _, b := range closure {
+			keep[b] = struct{}{}
+		}
+	}
+
+	cutoff := in.Now.Add(-in.Grace)
+
+	// Group objects by build prefix (the first path segment, which must be a UUID).
+	type agg struct {
+		count  int
+		bytes  int64
+		newest time.Time
+	}
+	byBuild := make(map[uuid.UUID]*agg)
+	var unknown []string
+
+	for _, obj := range in.Objects {
+		seg, _, ok := strings.Cut(obj.Path, "/")
+		if !ok {
+			unknown = append(unknown, obj.Path)
+
+			continue
+		}
+
+		id, err := uuid.Parse(seg)
+		if err != nil {
+			unknown = append(unknown, obj.Path)
+
+			continue
+		}
+
+		a := byBuild[id]
+		if a == nil {
+			a = &agg{}
+			byBuild[id] = a
+		}
+		a.count++
+		a.bytes += obj.Size
+		if obj.LastModified.After(a.newest) {
+			a.newest = obj.LastModified
+		}
+	}
+
+	m := &Manifest{UnknownKeys: unknown}
+	for b := range keep {
+		m.Keep = append(m.Keep, b)
+	}
+
+	for id, a := range byBuild {
+		if _, kept := keep[id]; kept {
+			continue
+		}
+
+		bp := BuildPrefix{
+			BuildID:      id,
+			ObjectCount:  a.count,
+			TotalBytes:   a.bytes,
+			NewestObject: a.newest,
+		}
+
+		// Protect the whole prefix if any object is newer than the cutoff: a
+		// fresh object may be an in-flight upload not yet reflected in the roots.
+		if a.newest.After(cutoff) {
+			m.TooNew = append(m.TooNew, bp)
+
+			continue
+		}
+
+		m.Candidates = append(m.Candidates, bp)
+	}
+
+	return m, nil
+}