feat: persist synchronously-available layer sizes to env_builds via gRPC

packages/api/internal/db/layer_sizes.go

@@ -0,0 +1,32 @@
+package db
+
+import (
+	"github.com/google/uuid"
+
+	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
+)
+
+// LayerSizesParams maps the orchestrator's synchronously-available logical layer
+// sizes into params for SetEnvBuildLayerSizes. A zero value maps to NULL (e.g.
+// the memfile logical size is 0 for filesystem-only snapshots).
+func LayerSizesParams(buildID uuid.UUID, ls *orchestrator.LayerSizes) queries.SetEnvBuildLayerSizesParams {
+	p := queries.SetEnvBuildLayerSizesParams{BuildID: buildID}
+	if ls == nil {
+		return p
+	}
+
+	p.MemfileLogicalSizeBytes = sizePtr(ls.GetMemfileLogicalSize())
+
+	return p
+}
+
+func sizePtr(v uint64) *int64 {
+	if v == 0 {
+		return nil
+	}
+
+	i := int64(v)
+
+	return &i
+}

packages/api/internal/orchestrator/pause_instance.go

@@ -12,6 +12,7 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 
+	apidb "github.com/e2b-dev/infra/packages/api/internal/db"
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
@@ -55,7 +56,7 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 		zap.String("source_build_id", sbx.BuildID.String()),
 	)
 
-	err = snapshotInstance(ctx, node, sbx, result.TemplateID, result.BuildID.String(), filesystemOnly)
+	pauseResp, err := snapshotInstance(ctx, node, sbx, result.TemplateID, result.BuildID.String(), filesystemOnly)
 	if errors.Is(err, PauseQueueExhaustedError{}) {
 		telemetry.ReportCriticalError(ctx, "pause queue exhausted", err)
 
@@ -81,17 +82,23 @@ func (o *Orchestrator) pauseSandbox(ctx context.Context, node *nodemanager.Node,
 		return fmt.Errorf("error pausing sandbox: %w", err)
 	}
 
+	// Best-effort: the pause already succeeded, so don't fail it over layer-size
+	// bookkeeping.
+	if err := o.sqlcDB.SetEnvBuildLayerSizes(ctx, apidb.LayerSizesParams(result.BuildID, pauseResp.GetLayerSizes())); err != nil {
+		logger.L().Warn(ctx, "failed to persist build layer sizes", logger.WithBuildID(result.BuildID.String()), zap.Error(err))
+	}
+
 	o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sbx.SandboxID)
 
 	return nil
 }
 
-func snapshotInstance(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox, templateID, buildID string, filesystemOnly bool) error {
+func snapshotInstance(ctx context.Context, node *nodemanager.Node, sbx sandbox.Sandbox, templateID, buildID string, filesystemOnly bool) (*orchestrator.SandboxPauseResponse, error) {
 	childCtx, childSpan := tracer.Start(ctx, "snapshot-instance")
 	defer childSpan.End()
 
 	client, childCtx := node.GetSandboxDeleteCtx(childCtx, sbx.SandboxID, sbx.ExecutionID)
-	_, err := client.Sandbox.Pause(
+	resp, err := client.Sandbox.Pause(
 		childCtx, &orchestrator.SandboxPauseRequest{
 			SandboxId:      sbx.SandboxID,
 			TemplateId:     templateID,
@@ -103,19 +110,19 @@ func snapshotInstance(ctx context.Context, node *nodemanager.Node, sbx sandbox.S
 	if err == nil {
 		telemetry.ReportEvent(ctx, "Paused sandbox")
 
-		return nil
+		return resp, nil
 	}
 
 	st, ok := status.FromError(err)
 	if !ok {
-		return err
+		return nil, err
 	}
 
 	if st.Code() == codes.ResourceExhausted {
-		return PauseQueueExhaustedError{}
+		return nil, PauseQueueExhaustedError{}
 	}
 
-	return fmt.Errorf("failed to pause sandbox '%s': %w", sbx.SandboxID, err)
+	return nil, fmt.Errorf("failed to pause sandbox '%s': %w", sbx.SandboxID, err)
 }
 
 func (o *Orchestrator) WaitForStateChange(ctx context.Context, teamID uuid.UUID, sandboxID string) error {

packages/api/internal/orchestrator/snapshot_template.go

@@ -7,13 +7,16 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"go.uber.org/zap"
 
+	apidb "github.com/e2b-dev/infra/packages/api/internal/db"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/consts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
+	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/storageopts"
 	"github.com/e2b-dev/infra/packages/shared/pkg/telemetry"
 )
@@ -88,7 +91,7 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U
 	// kills the sandbox itself; RemoveSandbox is still needed to clean up
 	// API-side state (store, routing, analytics).
 	client, childCtx := node.GetClient(ctx)
-	_, err = client.Sandbox.Checkpoint(childCtx, &orchestrator.SandboxCheckpointRequest{
+	checkpointResp, err := client.Sandbox.Checkpoint(childCtx, &orchestrator.SandboxCheckpointRequest{
 		SandboxId: sbx.SandboxID,
 		BuildId:   upsertResult.BuildID.String(),
 		Metadata:  map[string]string{storageopts.ObjectMetadataTemplateID: snapshotTemplateEnvID},
@@ -119,6 +122,12 @@ func (o *Orchestrator) CreateSnapshotTemplate(ctx context.Context, teamID uuid.U
 		return SnapshotTemplateResult{}, fmt.Errorf("error updating build status: %w", err)
 	}
 
+	// Best-effort: the checkpoint already succeeded, so don't fail it over
+	// layer-size bookkeeping.
+	if err := o.sqlcDB.SetEnvBuildLayerSizes(ctx, apidb.LayerSizesParams(upsertResult.BuildID, checkpointResp.GetLayerSizes())); err != nil {
+		logger.L().Warn(ctx, "failed to persist build layer sizes", logger.WithBuildID(upsertResult.BuildID.String()), zap.Error(err))
+	}
+
 	o.snapshotCache.Invalidate(context.WithoutCancel(ctx), sandboxID)
 
 	telemetry.ReportEvent(ctx, "Snapshot template completed")

packages/api/internal/template-manager/template_manager_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/google/uuid"
 
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	templatemanagergrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
@@ -26,7 +27,7 @@ func (f fakeTemplateManagerClient) SetStatus(context.Context, uuid.UUID, types.B
 	return f.setStatusError
 }
 
-func (f fakeTemplateManagerClient) SetFinished(context.Context, uuid.UUID, int64, string, string, string) error {
+func (f fakeTemplateManagerClient) SetFinished(context.Context, uuid.UUID, int64, string, string, string, *orchestrator.LayerSizes) error {
 	return f.setFinishedError
 }
 

packages/api/internal/template-manager/template_status.go

@@ -10,8 +10,10 @@ import (
 	"github.com/google/uuid"
 	"go.uber.org/zap"
 
+	apidb "github.com/e2b-dev/infra/packages/api/internal/db"
 	"github.com/e2b-dev/infra/packages/db/pkg/types"
 	"github.com/e2b-dev/infra/packages/db/queries"
+	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	templatemanagergrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
@@ -84,7 +86,7 @@ func (tm *TemplateManager) BuildStatusSync(ctx context.Context, buildID uuid.UUI
 
 type templateManagerClient interface {
 	SetStatus(ctx context.Context, buildID uuid.UUID, statusGroup types.BuildStatusGroup, reason *templatemanagergrpc.TemplateBuildStatusReason) error
-	SetFinished(ctx context.Context, buildID uuid.UUID, rootfsSize int64, envdVersion, kernelVersion, firecrackerVersion string) error
+	SetFinished(ctx context.Context, buildID uuid.UUID, rootfsSize int64, envdVersion, kernelVersion, firecrackerVersion string, layerSizes *orchestrator.LayerSizes) error
 	GetStatus(ctx context.Context, buildId uuid.UUID, templateID string, clusterID uuid.UUID, nodeID string) (*templatemanagergrpc.TemplateBuildStatusResponse, error)
 }
 
@@ -207,6 +209,7 @@ func (c *PollBuildStatus) dispatchBasedOnStatus(ctx context.Context, status *tem
 			meta.GetEnvdVersionKey(),
 			meta.GetKernelVersion(),
 			meta.GetFirecrackerVersion(),
+			meta.GetLayerSizes(),
 		)
 		if err != nil {
 			return false, fmt.Errorf("error when finishing build: %w", err)
@@ -307,7 +310,7 @@ func (tm *TemplateManager) SetStatus(ctx context.Context, buildID uuid.UUID, sta
 	return err
 }
 
-func (tm *TemplateManager) SetFinished(ctx context.Context, buildID uuid.UUID, rootfsSize int64, envdVersion, kernelVersion, firecrackerVersion string) error {
+func (tm *TemplateManager) SetFinished(ctx context.Context, buildID uuid.UUID, rootfsSize int64, envdVersion, kernelVersion, firecrackerVersion string, layerSizes *orchestrator.LayerSizes) error {
 	// first do database update to prevent race condition while calling status
 	// TODO(ENG-3469): Switch to types.BuildStatusReady once all consumers are migrated.
 	err := tm.sqlcDB.FinishTemplateBuild(ctx, queries.FinishTemplateBuildParams{
@@ -318,6 +321,9 @@ func (tm *TemplateManager) SetFinished(ctx context.Context, buildID uuid.UUID, r
 		FirecrackerVersion: firecrackerVersion,
 		BuildID:            buildID,
 	})
+	if err == nil {
+		err = tm.sqlcDB.SetEnvBuildLayerSizes(ctx, apidb.LayerSizesParams(buildID, layerSizes))
+	}
 
 	tm.buildCache.Invalidate(ctx, buildID)
 

packages/db/migrations/20260627130000_add_layer_size_columns.sql

@@ -0,0 +1,11 @@
+-- +goose Up
+-- Memfile logical (virtual device) size captured synchronously at snapshot time.
+-- The rootfs logical size is already stored as total_disk_size_mb. Mapped and
+-- diff sizes (for both rootfs and memfile) are kept in the data objects' custom
+-- metadata instead, so each size category lives entirely in one place.
+ALTER TABLE public.env_builds
+  ADD COLUMN IF NOT EXISTS memfile_logical_size_bytes bigint;
+
+-- +goose Down
+ALTER TABLE public.env_builds
+  DROP COLUMN IF EXISTS memfile_logical_size_bytes;

packages/db/queries/builds/set_env_build_layer_sizes.sql

@@ -0,0 +1,6 @@
+-- name: SetEnvBuildLayerSizes :exec
+-- Persists the synchronously-available memfile logical size for a build. Nullable
+-- so a filesystem-only snapshot (no memfile) can be written as NULL.
+UPDATE "public"."env_builds"
+SET memfile_logical_size_bytes = sqlc.narg(memfile_logical_size_bytes)
+WHERE id = @build_id;