diff --git a/packages/orchestrator/cmd/template-gc/main.go b/packages/orchestrator/cmd/template-gc/main.go new file mode 100644 index 0000000000..21742b14db --- /dev/null +++ b/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) +} diff --git a/packages/orchestrator/internal/templategc/mark.go b/packages/orchestrator/internal/templategc/mark.go new file mode 100644 index 0000000000..02d9787a3c --- /dev/null +++ b/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 +} diff --git a/packages/orchestrator/internal/templategc/mark_test.go b/packages/orchestrator/internal/templategc/mark_test.go new file mode 100644 index 0000000000..95d43d1df9 --- /dev/null +++ b/packages/orchestrator/internal/templategc/mark_test.go @@ -0,0 +1,169 @@ +package templategc + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage" +) + +type stubReader struct { + closures map[uuid.UUID][]uuid.UUID + err error +} + +func (s stubReader) AncestorBuilds(_ context.Context, b uuid.UUID) ([]uuid.UUID, error) { + if s.err != nil { + return nil, s.err + } + + out := []uuid.UUID{b} + + return append(out, s.closures[b]...), nil +} + +func obj(buildID uuid.UUID, name string, size int64, mod time.Time) storage.ObjectInfo { + return storage.ObjectInfo{Path: buildID.String() + "/" + name, Size: size, LastModified: mod} +} + +// A build referenced only through a live root's header closure (never its own +// root) must be kept, not marked a candidate. +func TestMark_KeepsAncestorOfLiveRoot(t *testing.T) { + t.Parallel() + + now := time.Now() + old := now.Add(-100 * 24 * time.Hour) + + liveRoot := uuid.New() + ancestor := uuid.New() // shared base layer, old, not a root itself + + reader := stubReader{closures: map[uuid.UUID][]uuid.UUID{liveRoot: {ancestor}}} + + in := MarkInput{ + LiveRoots: []uuid.UUID{liveRoot}, + Objects: []storage.ObjectInfo{ + obj(liveRoot, "rootfs.ext4.header", 10, old), + obj(ancestor, "rootfs.ext4.header", 10, old), + obj(ancestor, "memfile", 1000, old), + }, + Now: now, + Grace: 60 * 24 * time.Hour, + } + + m, err := Mark(context.Background(), reader, in) + require.NoError(t, err) + + assert.Empty(t, m.Candidates, "an old ancestor referenced by a live root must not be a candidate") + assert.Contains(t, m.Keep, ancestor) + assert.Contains(t, m.Keep, liveRoot) +} + +// An unreferenced, old build is a candidate. +func TestMark_OrphanOldBuildIsCandidate(t *testing.T) { + t.Parallel() + + now := time.Now() + old := now.Add(-100 * 24 * time.Hour) + + liveRoot := uuid.New() + orphan := uuid.New() + + reader := stubReader{closures: map[uuid.UUID][]uuid.UUID{}} + + in := MarkInput{ + LiveRoots: []uuid.UUID{liveRoot}, + Objects: []storage.ObjectInfo{ + obj(liveRoot, "memfile", 1, old), + obj(orphan, "memfile", 2048, old), + obj(orphan, "rootfs.ext4.header", 20, old), + }, + Now: now, + Grace: 60 * 24 * time.Hour, + } + + m, err := Mark(context.Background(), reader, in) + require.NoError(t, err) + + require.Len(t, m.Candidates, 1) + assert.Equal(t, orphan, m.Candidates[0].BuildID) + assert.Equal(t, 2, m.Candidates[0].ObjectCount) + assert.Equal(t, int64(2068), m.Candidates[0].TotalBytes) +} + +// An unreferenced build with a fresh object is held back (grace), protecting the +// whole prefix — a fresh object may be an in-flight upload. +func TestMark_RecentObjectProtectsWholePrefix(t *testing.T) { + t.Parallel() + + now := time.Now() + old := now.Add(-100 * 24 * time.Hour) + fresh := now.Add(-1 * time.Hour) + + orphan := uuid.New() + + reader := stubReader{closures: map[uuid.UUID][]uuid.UUID{}} + + in := MarkInput{ + LiveRoots: nil, + Objects: []storage.ObjectInfo{ + obj(orphan, "memfile", 10, old), + obj(orphan, "memfile.header", 10, fresh), // one fresh object + }, + Now: now, + Grace: 60 * 24 * time.Hour, + } + + m, err := Mark(context.Background(), reader, in) + require.NoError(t, err) + + assert.Empty(t, m.Candidates) + require.Len(t, m.TooNew, 1) + assert.Equal(t, orphan, m.TooNew[0].BuildID) +} + +// If any live root's closure cannot be resolved, Mark fails closed. +func TestMark_FailsClosedOnUnreadableRoot(t *testing.T) { + t.Parallel() + + reader := stubReader{err: errors.New("header corrupt")} + + in := MarkInput{ + LiveRoots: []uuid.UUID{uuid.New()}, + Objects: []storage.ObjectInfo{obj(uuid.New(), "memfile", 1, time.Now().Add(-1000*time.Hour))}, + Now: time.Now(), + Grace: 60 * 24 * time.Hour, + } + + _, err := Mark(context.Background(), reader, in) + require.Error(t, err, "an unreadable live-root closure must fail closed, not produce candidates") +} + +// Keys not under a UUID prefix are surfaced as unknown, never deleted. +func TestMark_UnknownKeysSurfaced(t *testing.T) { + t.Parallel() + + now := time.Now() + old := now.Add(-100 * 24 * time.Hour) + + in := MarkInput{ + LiveRoots: nil, + Objects: []storage.ObjectInfo{ + {Path: "not-a-uuid/file", Size: 1, LastModified: old}, + {Path: "toplevel", Size: 1, LastModified: old}, + }, + Now: now, + Grace: 60 * 24 * time.Hour, + } + + m, err := Mark(context.Background(), stubReader{}, in) + require.NoError(t, err) + + assert.Empty(t, m.Candidates) + assert.ElementsMatch(t, []string{"not-a-uuid/file", "toplevel"}, m.UnknownKeys) +} diff --git a/packages/orchestrator/internal/templategc/reader.go b/packages/orchestrator/internal/templategc/reader.go new file mode 100644 index 0000000000..f438c1c97e --- /dev/null +++ b/packages/orchestrator/internal/templategc/reader.go @@ -0,0 +1,42 @@ +package templategc + +import ( + "context" + "fmt" + + "github.com/google/uuid" + + "github.com/e2b-dev/infra/packages/shared/pkg/storage" + "github.com/e2b-dev/infra/packages/shared/pkg/storage/header" +) + +// StorageAncestorReader reads a build's memfile and rootfs headers from template +// storage and unions the build IDs their block mappings reference. It is the +// production AncestorReader. A missing or unparseable header returns an error so +// Mark fails closed. +type StorageAncestorReader struct { + Storage storage.StorageProvider +} + +func (r StorageAncestorReader) AncestorBuilds(ctx context.Context, buildID uuid.UUID) ([]uuid.UUID, error) { + paths := storage.Paths{BuildID: buildID.String()} + + seen := map[uuid.UUID]struct{}{buildID: {}} + for _, hdrPath := range []string{paths.MemfileHeader(), paths.RootfsHeader()} { + h, _, err := header.LoadHeader(ctx, r.Storage, hdrPath) + if err != nil { + return nil, fmt.Errorf("load header %s: %w", hdrPath, err) + } + + for _, b := range h.Mapping.Builds() { + seen[b] = struct{}{} + } + } + + out := make([]uuid.UUID, 0, len(seen)) + for b := range seen { + out = append(out, b) + } + + return out, nil +} diff --git a/packages/shared/pkg/storage/objlister_test.go b/packages/shared/pkg/storage/objlister_test.go new file mode 100644 index 0000000000..3f92dfff38 --- /dev/null +++ b/packages/shared/pkg/storage/objlister_test.go @@ -0,0 +1,66 @@ +package storage + +import ( + "context" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestFSStorage_ListObjects(t *testing.T) { + t.Parallel() + + base := t.TempDir() + s := newFileSystemStorage(base, "", nil) + + // Two build prefixes, one with two objects. + require.NoError(t, os.MkdirAll(filepath.Join(base, "build-a"), 0o755)) + require.NoError(t, os.MkdirAll(filepath.Join(base, "build-b"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(base, "build-a", "memfile"), []byte("hello"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(base, "build-a", "rootfs.ext4.header"), []byte("hdr"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(base, "build-b", "memfile"), []byte("x"), 0o644)) + + objs, err := s.ListObjects(context.Background(), "") + require.NoError(t, err) + + paths := make(map[string]int64) + for _, o := range objs { + paths[o.Path] = o.Size + assert.False(t, o.LastModified.IsZero(), "LastModified must be populated for grace checks") + } + + assert.Equal(t, int64(5), paths["build-a/memfile"]) + assert.Equal(t, int64(3), paths["build-a/rootfs.ext4.header"]) + assert.Equal(t, int64(1), paths["build-b/memfile"]) + assert.Len(t, objs, 3) +} + +func TestFSStorage_ListObjects_PrefixScoped(t *testing.T) { + t.Parallel() + + base := t.TempDir() + s := newFileSystemStorage(base, "", nil) + require.NoError(t, os.MkdirAll(filepath.Join(base, "build-a"), 0o755)) + require.NoError(t, os.MkdirAll(filepath.Join(base, "build-b"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(base, "build-a", "memfile"), []byte("a"), 0o644)) + require.NoError(t, os.WriteFile(filepath.Join(base, "build-b", "memfile"), []byte("b"), 0o644)) + + objs, err := s.ListObjects(context.Background(), "build-a") + require.NoError(t, err) + require.Len(t, objs, 1) + assert.Equal(t, "build-a/memfile", objs[0].Path) +} + +// ListObjects helper returns ErrListUnsupported for a provider without the +// capability, so callers fail closed instead of seeing an empty bucket. +func TestListObjects_UnsupportedFailsClosed(t *testing.T) { + t.Parallel() + + _, err := ListObjects(context.Background(), noListProvider{}, "") + assert.ErrorIs(t, err, ErrListUnsupported) +} + +type noListProvider struct{ StorageProvider } diff --git a/packages/shared/pkg/storage/storage.go b/packages/shared/pkg/storage/storage.go index f517dfdede..a57e6624bb 100644 --- a/packages/shared/pkg/storage/storage.go +++ b/packages/shared/pkg/storage/storage.go @@ -39,6 +39,11 @@ var ErrObjectSoftDeleted = errors.New("object soft-deleted by storage index") // backend instead of assuming there is no tombstone. var ErrMetadataUnsupported = errors.New("blob does not support reading custom metadata") +// ErrListUnsupported means the storage backend cannot enumerate objects under a +// prefix. Callers (e.g. the template GC mark tool) must fail closed when they +// get this rather than treat the bucket as empty. +var ErrListUnsupported = errors.New("storage provider does not support listing objects") + // ObjectMetadataSoftDeleted is the storage-index soft-delete tombstone key. const ObjectMetadataSoftDeleted = storageopts.ObjectMetadataSoftDeleted @@ -149,6 +154,38 @@ type MetadataReader interface { Metadata(ctx context.Context) (ObjectMetadata, error) } +// ObjectInfo describes one object returned by ObjectLister. +type ObjectInfo struct { + // Path is the object's full key (e.g. "{buildID}/memfile.header"). + Path string + // Size is the object size in bytes. + Size int64 + // LastModified is the object's last-modified time, used by the GC grace + // window: a build prefix is only a deletion candidate when every object + // under it is older than the grace period. + LastModified time.Time +} + +// ObjectLister is an optional StorageProvider capability: enumerate every object +// under a prefix (paginated internally). Required by full-bucket operations like +// the template GC mark tool; backends that cannot list omit it and callers fail +// closed via ErrListUnsupported. +type ObjectLister interface { + ListObjects(ctx context.Context, prefix string) ([]ObjectInfo, error) +} + +// ListObjects lists every object under prefix, or returns ErrListUnsupported if +// the provider cannot enumerate — so callers can distinguish "bucket empty" from +// "couldn't list" and refuse to treat an unlistable bucket as having no objects. +func ListObjects(ctx context.Context, p StorageProvider, prefix string) ([]ObjectInfo, error) { + l, ok := p.(ObjectLister) + if !ok { + return nil, ErrListUnsupported + } + + return l.ListObjects(ctx, prefix) +} + // BlobCustomMetadata returns the blob's custom metadata, or ErrMetadataUnsupported // when the backend can't read it — so callers can distinguish "no tombstone" // from "couldn't check" and fail closed under enforcement. diff --git a/packages/shared/pkg/storage/storage_aws.go b/packages/shared/pkg/storage/storage_aws.go index 52c8d49ab1..45158443f5 100644 --- a/packages/shared/pkg/storage/storage_aws.go +++ b/packages/shared/pkg/storage/storage_aws.go @@ -580,3 +580,35 @@ func (m *awsPartUploader) Close() error { return err } + +// ListObjects enumerates every object under prefix, paginating through the full +// result set. Implements the ObjectLister optional capability. +func (s *awsStorage) ListObjects(ctx context.Context, prefix string) ([]ObjectInfo, error) { + paginator := s3.NewListObjectsV2Paginator(s.client, &s3.ListObjectsV2Input{ + Bucket: aws.String(s.bucketName), + Prefix: aws.String(prefix), + }) + + var out []ObjectInfo + for paginator.HasMorePages() { + pageCtx, cancel := context.WithTimeout(ctx, awsOperationTimeout) + list, err := paginator.NextPage(pageCtx) + cancel() + if err != nil { + return nil, err + } + + for _, obj := range list.Contents { + info := ObjectInfo{Path: aws.ToString(obj.Key)} + if obj.Size != nil { + info.Size = *obj.Size + } + if obj.LastModified != nil { + info.LastModified = *obj.LastModified + } + out = append(out, info) + } + } + + return out, nil +} diff --git a/packages/shared/pkg/storage/storage_fs.go b/packages/shared/pkg/storage/storage_fs.go index 4fc101ef06..3c0b0eaf0a 100644 --- a/packages/shared/pkg/storage/storage_fs.go +++ b/packages/shared/pkg/storage/storage_fs.go @@ -422,3 +422,47 @@ func (o *fsObject) OpenRangeReader(ctx context.Context, offsetU int64, length in return raw, SourceFS, nil } + +// ListObjects walks the filesystem under prefix and reports each regular file as +// an object. Implements the ObjectLister optional capability (used by tests and +// local dry runs). +func (s *fsStorage) ListObjects(_ context.Context, prefix string) ([]ObjectInfo, error) { + root := s.getPath(prefix) + + var out []ObjectInfo + err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error { + if err != nil { + if os.IsNotExist(err) { + return nil + } + + return err + } + if d.IsDir() { + return nil + } + + info, statErr := d.Info() + if statErr != nil { + return statErr + } + + rel, relErr := filepath.Rel(s.basePath, path) + if relErr != nil { + return relErr + } + + out = append(out, ObjectInfo{ + Path: filepath.ToSlash(rel), + Size: info.Size(), + LastModified: info.ModTime(), + }) + + return nil + }) + if err != nil { + return nil, err + } + + return out, nil +}