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

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.