diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index d381a9985a..8f74a9e2e9 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/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) diff --git a/packages/orchestrator/pkg/template/build/builder.go b/packages/orchestrator/pkg/template/build/builder.go index 18a6fc24e0..e613cea397 100644 --- a/packages/orchestrator/pkg/template/build/builder.go +++ b/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( diff --git a/packages/orchestrator/pkg/template/build/layer/create_sandbox.go b/packages/orchestrator/pkg/template/build/layer/create_sandbox.go index a87904df2c..e0efea6739 100644 --- a/packages/orchestrator/pkg/template/build/layer/create_sandbox.go +++ b/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) diff --git a/packages/orchestrator/pkg/template/build/phases/base/builder.go b/packages/orchestrator/pkg/template/build/phases/base/builder.go index 5bbac9c963..d670056dfb 100644 --- a/packages/orchestrator/pkg/template/build/phases/base/builder.go +++ b/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) { diff --git a/packages/orchestrator/pkg/template/build/phases/base/provision.go b/packages/orchestrator/pkg/template/build/phases/base/provision.go index a803d20ef0..a52b0e59b8 100644 --- a/packages/orchestrator/pkg/template/build/phases/base/provision.go +++ b/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) diff --git a/packages/orchestrator/pkg/template/build/phases/finalize/builder.go b/packages/orchestrator/pkg/template/build/phases/finalize/builder.go index 59897e5535..998aca413f 100644 --- a/packages/orchestrator/pkg/template/build/phases/finalize/builder.go +++ b/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, diff --git a/packages/orchestrator/pkg/template/build/phases/steps/builder.go b/packages/orchestrator/pkg/template/build/phases/steps/builder.go index 0ffb48c412..0ff423e0b8 100644 --- a/packages/orchestrator/pkg/template/build/phases/steps/builder.go +++ b/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) diff --git a/packages/orchestrator/pkg/template/build/phases/steps/factory.go b/packages/orchestrator/pkg/template/build/phases/steps/factory.go index e78e7d8f68..c2261d1f86 100644 --- a/packages/orchestrator/pkg/template/build/phases/steps/factory.go +++ b/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, diff --git a/packages/orchestrator/pkg/template/build/phases/user/builder.go b/packages/orchestrator/pkg/template/build/phases/user/builder.go index 3d9242fbfe..121d2411c2 100644 --- a/packages/orchestrator/pkg/template/build/phases/user/builder.go +++ b/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"}, diff --git a/packages/orchestrator/pkg/template/build/sandboxtools/command.go b/packages/orchestrator/pkg/template/build/sandboxtools/command.go index 5c60440a09..c67677c358 100644 --- a/packages/orchestrator/pkg/template/build/sandboxtools/command.go +++ b/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.