feat(orch): move reserved blocks setting from in-guest to host-side pre-boot hook (#2199)

* feat(orch): add PreBootFn hook to Factory.CreateSandbox for host-side rootfs modifications

Add a PreBootFn callback parameter to Factory.CreateSandbox() that runs
after the rootfs device is ready but before Firecracker boots the VM.
This enables host-side filesystem metadata changes (e.g. tune2fs) on the
rootfs path without depending on guest-side tooling.

Thread the hook through the layer CreateSandbox builder via a new
WithPreBootFn option.

* feat(orch): move reserved blocks setting from in-guest to host-side pre-boot hook

Replace SetReservedBlocksInGuest (tune2fs inside the VM) with a host-side
PreBootFn that calls SetReservedBlocksOnHost on the rootfs device path
before the guest kernel boots. This eliminates the dependency on e2fsprogs
being installed and compatible inside the guest image.

The previous in-guest approach was fragile across different build paths:
- fromImage builds: provisioning could install e2fsprogs, but older guest
  versions were incompatible with host-created ext4 features.
- fromTemplate builds: provisioning is skipped entirely, so e2fsprogs
  was only present if the parent template happened to have it.
- Cached builds: the base phase SetReservedBlocksOnHost was skipped on
  cache hits, leaving only the in-guest call which could fail.

Reserved blocks are now set at two points:
- Base phase: on the host rootfs file after provisioning, so root has
  protected disk space during user build steps.
- Finalize phase: via PreBootFn before every VM boot, regardless of
  build path or cache state, using the host's own tune2fs.

* refactor(orch): collect finalize sandbox options into a slice

Gather CreateSandbox options up front and spread them, rather than
always passing WithPreBootFn which could be nil.

* feat(orch): set reserved blocks on first uncached step layer via PreBootFn

Apply the host-side SetReservedBlocksOnHost pre-boot hook to the steps
phase as well, not just finalize. When the first uncached layer creates
a new sandbox, the PreBootFn sets reserved blocks before the guest boots.
This protects intermediate build sandboxes from disk-full conditions
during user RUN commands, not just the final template.

* refactor(orch): unify base phase reserved blocks to use PreBootFn hook

Move the inline SetReservedBlocksOnHost call in buildLayerFromOCI to use
the same WithPreBootFn mechanism as the steps and finalize phases. This
is a pure refactor — same file, same timing (before VM boot), just
consistent with the rest of the pipeline.

* refactor(orch): extract ReservedBlocksOptions helper to deduplicate PreBootFn setup

Consolidate the identical reserved blocks PreBootFn setup from base,
steps, and finalize phases into a single ReservedBlocksOptions() helper
in the layer package. Each call site now collapses from 5 lines to 1.

* fix(orch): guard finalize ReservedBlocksOptions with sourceLayer.Cached

Only set reserved blocks in finalize when the source layer was cached,
since otherwise a prior step phase already set them via its own
CreateSandbox PreBootFn.

Author: arkamar

packages/orchestrator/pkg/sandbox/sandbox.go

@@ -295,6 +295,12 @@ func NewFactory(
 	}
 }
 
+// PreBootFn is an optional callback invoked after the rootfs is ready but before
+// Firecracker boots. It receives the rootfs device path (e.g., a file path for
+// DirectProvider or /dev/nbdX for NBDProvider) and may modify the filesystem
+// on the host side.
+type PreBootFn func(ctx context.Context, rootfsPath string) error
+
 // CreateSandbox creates the sandbox.
 // IMPORTANT: You must Close() the sandbox after you are done with it.
 func (f *Factory) CreateSandbox(
@@ -306,6 +312,7 @@ func (f *Factory) CreateSandbox(
 	rootfsCachePath string,
 	processOptions fc.ProcessOptions,
 	apiConfigToStore *orchestrator.SandboxConfig,
+	preBootFn PreBootFn,
 ) (s *Sandbox, e error) {
 	ctx, span := tracer.Start(ctx, "create sandbox")
 	defer span.End()
@@ -380,6 +387,19 @@ func (f *Factory) CreateSandbox(
 		return nil, err
 	}
 
+	// Run the optional pre-boot hook before Firecracker starts.
+	// This allows host-side filesystem changes before the guest kernel takes charge.
+	if preBootFn != nil {
+		rootfsPath, pathErr := rootfsProvider.Path()
+		if pathErr != nil {
+			return nil, fmt.Errorf("failed to get rootfs path for pre-boot hook: %w", pathErr)
+		}
+
+		if hookErr := preBootFn(ctx, rootfsPath); hookErr != nil {
+			return nil, fmt.Errorf("pre-boot hook failed: %w", hookErr)
+		}
+	}
+
 	cgroupHandle, cgroupFD := createCgroup(ctx, f.cgroupManager, sandboxFiles.SandboxCgroupName(), cleanup)
 	defer releaseCgroupFD(ctx, cgroupHandle, runtime.SandboxID)
 

packages/orchestrator/pkg/template/build/builder.go

@@ -301,6 +301,7 @@ func runBuild(
 		commandExecutor,
 		index,
 		builder.metrics,
+		builder.featureFlags,
 		config.TemplateDefaultUser,
 		bc.Config.Force,
 	)
@@ -314,6 +315,7 @@ func runBuild(
 		commandExecutor,
 		index,
 		builder.metrics,
+		builder.featureFlags,
 	)
 
 	postProcessingBuilder := finalize.New(

packages/orchestrator/pkg/template/build/layer/create_sandbox.go

@@ -13,10 +13,12 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/fc"
 	sbxtemplate "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/template"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/config"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/core/filesystem"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/constants"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/units"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
 	"github.com/e2b-dev/infra/packages/shared/pkg/fc/models"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/id"
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
@@ -29,6 +31,7 @@ type CreateSandbox struct {
 
 	rootfsCachePath string
 	ioEngine        *string
+	preBootFn       sandbox.PreBootFn
 }
 
 const (
@@ -41,6 +44,7 @@ var _ SandboxCreator = (*CreateSandbox)(nil)
 type createSandboxOptions struct {
 	rootfsCachePath string
 	ioEngine        *string
+	preBootFn       sandbox.PreBootFn
 }
 
 type CreateSandboxOption func(*createSandboxOptions)
@@ -57,6 +61,31 @@ func WithRootfsCachePath(rootfsCachePath string) CreateSandboxOption {
 	}
 }
 
+// WithPreBootFn sets a callback that runs after the rootfs is ready but before
+// Firecracker boots. The callback receives the rootfs device path and can
+// modify filesystem on the host side.
+func WithPreBootFn(fn sandbox.PreBootFn) CreateSandboxOption {
+	return func(opts *createSandboxOptions) {
+		opts.preBootFn = fn
+	}
+}
+
+// ReservedBlocksOptions returns CreateSandboxOption(s) that set reserved blocks
+// on the rootfs before the guest boots, if the BuildReservedDiskSpaceMB feature
+// flag is greater than zero. Returns nil otherwise.
+func ReservedBlocksOptions(ctx context.Context, featureFlags *featureflags.Client, blockSize int64) []CreateSandboxOption {
+	reservedDiskSpaceMB := int64(featureFlags.IntFlag(ctx, featureflags.BuildReservedDiskSpaceMB))
+	if reservedDiskSpaceMB <= 0 {
+		return nil
+	}
+
+	return []CreateSandboxOption{
+		WithPreBootFn(func(ctx context.Context, rootfsPath string) error {
+			return filesystem.SetReservedBlocksOnHost(ctx, rootfsPath, reservedDiskSpaceMB, blockSize)
+		}),
+	}
+}
+
 func NewCreateSandbox(config *sandbox.Config, sandboxFactory *sandbox.Factory, timeout time.Duration, options ...CreateSandboxOption) *CreateSandbox {
 	opts := &createSandboxOptions{
 		rootfsCachePath: "",
@@ -72,6 +101,7 @@ func NewCreateSandbox(config *sandbox.Config, sandboxFactory *sandbox.Factory, t
 		rootfsCachePath: opts.rootfsCachePath,
 		sandboxFactory:  sandboxFactory,
 		ioEngine:        opts.ioEngine,
+		preBootFn:       opts.preBootFn,
 	}
 }
 
@@ -121,6 +151,7 @@ func (cs *CreateSandbox) Sandbox(
 			IoEngine:            cs.ioEngine,
 		},
 		nil,
+		cs.preBootFn,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("create sandbox: %w", err)

packages/orchestrator/pkg/template/build/phases/base/builder.go

@@ -252,21 +252,19 @@ func (bb *BaseBuilder) buildLayerFromOCI(
 		return metadata.Template{}, fmt.Errorf("error enlarging disk after provisioning: %w", err)
 	}
 
-	if reservedDiskSpaceMB := int64(bb.featureFlags.IntFlag(ctx, featureflags.BuildReservedDiskSpaceMB)); reservedDiskSpaceMB > 0 {
-		err = filesystem.SetReservedBlocksOnHost(ctx, rootfsPath, reservedDiskSpaceMB, bb.Config.RootfsBlockSize())
-		if err != nil {
-			return metadata.Template{}, fmt.Errorf("error setting reserved disk space: %w", err)
-		}
-	}
-
 	// Create sandbox for building template
 	userLogger.Debug(ctx, "Creating base sandbox template layer")
 
+	sandboxOptions := []layer.CreateSandboxOption{
+		layer.WithRootfsCachePath(rootfsPath),
+	}
+	sandboxOptions = append(sandboxOptions, layer.ReservedBlocksOptions(ctx, bb.featureFlags, bb.Config.RootfsBlockSize())...)
+
 	sandboxCreator := layer.NewCreateSandbox(
 		baseSbxConfig,
 		bb.sandboxFactory,
 		baseLayerTimeout,
-		layer.WithRootfsCachePath(rootfsPath),
+		sandboxOptions...,
 	)
 
 	actionExecutor := layer.NewFunctionAction(func(ctx context.Context, sbx *sandbox.Sandbox, meta metadata.Template) (metadata.Template, error) {

packages/orchestrator/pkg/template/build/phases/base/provision.go

@@ -139,6 +139,7 @@ func (bb *BaseBuilder) provisionSandbox(
 			Stderr: logsWriter,
 		},
 		nil,
+		nil, // no pre-boot hook for provisioning
 	)
 	if err != nil {
 		return fmt.Errorf("error creating sandbox: %w", err)

packages/orchestrator/pkg/template/build/phases/finalize/builder.go

@@ -174,12 +174,20 @@ func (ppb *PostProcessingBuilder) Build(
 	span.SetAttributes(attribute.String("io_engine", ioEngine))
 	ppb.logger.Debug(ctx, "using io engine", zap.String("io_engine", ioEngine))
 
+	// Collect sandbox creation options
+	sandboxOptions := []layer.CreateSandboxOption{
+		layer.WithIoEngine(ioEngine),
+	}
+	if sourceLayer.Cached {
+		sandboxOptions = append(sandboxOptions, layer.ReservedBlocksOptions(ctx, ppb.featureFlags, ppb.Config.RootfsBlockSize())...)
+	}
+
 	// Always restart the sandbox for the final layer to properly wire the rootfs path for the final template
 	sandboxCreator := layer.NewCreateSandbox(
 		sbxConfig,
 		ppb.sandboxFactory,
 		finalizeTimeout,
-		layer.WithIoEngine(ioEngine),
+		sandboxOptions...,
 	)
 
 	actionExecutor := layer.NewFunctionAction(ppb.postProcessingFn(userLogger))
@@ -219,16 +227,6 @@ func (ppb *PostProcessingBuilder) postProcessingFn(userLogger logger.Logger) lay
 				return
 			}
 
-			// Set reserved disk space for the guest OS before syncing
-			if reservedDiskSpaceMB := int64(ppb.featureFlags.IntFlag(ctx, featureflags.BuildReservedDiskSpaceMB)); reservedDiskSpaceMB > 0 {
-				err := sandboxtools.SetReservedBlocksInGuest(ctx, ppb.proxy, userLogger, sbx.Runtime.SandboxID, reservedDiskSpaceMB, ppb.Config.RootfsBlockSize())
-				if err != nil {
-					e = fmt.Errorf("error setting reserved disk space: %w", err)
-
-					return
-				}
-			}
-
 			// Ensure all changes are synchronized to disk so the sandbox can be restarted
 			err := sandboxtools.SyncChangesToDisk(
 				ctx,

packages/orchestrator/pkg/template/build/phases/steps/builder.go

@@ -24,6 +24,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/sandboxtools"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/storage/cache"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
@@ -47,6 +48,7 @@ type StepBuilder struct {
 	commandExecutor *commands.CommandExecutor
 	index           cache.Index
 	metrics         *metrics.BuildMetrics
+	featureFlags    *featureflags.Client
 }
 
 func New(
@@ -58,6 +60,7 @@ func New(
 	commandExecutor *commands.CommandExecutor,
 	index cache.Index,
 	metrics *metrics.BuildMetrics,
+	featureFlags *featureflags.Client,
 	step *templatemanager.TemplateStep,
 	stepNumber int,
 	defaultLoggingLevel zapcore.Level,
@@ -77,6 +80,7 @@ func New(
 		commandExecutor: commandExecutor,
 		index:           index,
 		metrics:         metrics,
+		featureFlags:    featureFlags,
 	}
 }
 
@@ -180,6 +184,7 @@ func (sb *StepBuilder) Build(
 			sbxConfig,
 			sb.sandboxFactory,
 			layerTimeout,
+			layer.ReservedBlocksOptions(ctx, sb.featureFlags, sb.Config.RootfsBlockSize())...,
 		)
 	} else {
 		sandboxCreator = layer.NewResumeSandbox(sbxConfig, sb.sandboxFactory, layerTimeout)

packages/orchestrator/pkg/template/build/phases/steps/factory.go

@@ -11,6 +11,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/metrics"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/phases"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/storage/cache"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
 
@@ -25,6 +26,7 @@ func CreateStepPhases(
 	commandExecutor *commands.CommandExecutor,
 	index cache.Index,
 	metrics *metrics.BuildMetrics,
+	featureFlags *featureflags.Client,
 ) []phases.BuilderPhase {
 	steps := make([]phases.BuilderPhase, 0, len(bc.Config.Steps))
 
@@ -39,6 +41,7 @@ func CreateStepPhases(
 				commandExecutor,
 				index,
 				metrics,
+				featureFlags,
 				step,
 				i+1, // stepNumber starts from 1
 				defaultLoggingLevel,

packages/orchestrator/pkg/template/build/phases/user/builder.go

@@ -15,6 +15,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/phases"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/phases/steps"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/storage/cache"
+	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
 	template_manager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
 )
@@ -34,6 +35,7 @@ func New(
 	commandExecutor *commands.CommandExecutor,
 	index cache.Index,
 	metrics *metrics.BuildMetrics,
+	featureFlags *featureflags.Client,
 	user string,
 	force *bool,
 ) *UserBuilder {
@@ -47,6 +49,7 @@ func New(
 			commandExecutor,
 			index,
 			metrics,
+			featureFlags,
 			&template_manager.TemplateStep{
 				Type:  "USER",
 				Args:  []string{user, "true"},

packages/orchestrator/pkg/template/build/sandboxtools/command.go

@@ -13,14 +13,12 @@ import (
 	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/trace"
-	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/proxy"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/build/core/rootfs"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/metadata"
-	"github.com/e2b-dev/infra/packages/orchestrator/pkg/units"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/envd/process/processconnect"
@@ -234,38 +232,6 @@ func logStream(ctx context.Context, logger logger.Logger, lvl zapcore.Level, id
 	}
 }
 
-// SetReservedBlocksInGuest sets the number of reserved filesystem blocks inside the sandbox.
-// Reserved blocks are only usable by root (uid 0), protecting the guest OS from disk-full conditions.
-// Requires e2fsprogs (tune2fs) installed in the guest image (standard on Debian-based images).
-func SetReservedBlocksInGuest(
-	ctx context.Context,
-	proxy *proxy.SandboxProxy,
-	logger logger.Logger,
-	sandboxID string,
-	reservedSpaceMB int64,
-	blockSize int64,
-) error {
-	if reservedSpaceMB <= 0 {
-		return nil
-	}
-
-	blocks := units.MBToBytes(reservedSpaceMB) / blockSize
-	tuneCmd := fmt.Sprintf("tune2fs -r %d /dev/vda", blocks)
-
-	return RunCommandWithLogger(
-		ctx,
-		proxy,
-		logger,
-		zap.DebugLevel,
-		"set-reserved-disk-space",
-		sandboxID,
-		tuneCmd,
-		metadata.Context{
-			User: "root",
-		},
-	)
-}
-
 // syncChangesToDisk synchronizes filesystem changes to the filesystem
 // This is useful to ensure that all changes made in the sandbox are written to disk
 // to be able to re-create the sandbox without resume.