Refactor the Create/Resume functions to avoid globals/reduce args (#1258)

Co-authored-by: Jakub Novak <jakub@e2b.dev>

Author: djeebus

packages/orchestrator/cmd/build-template/main.go

@@ -155,13 +155,15 @@ func buildTemplate(
 	if err != nil {
 		zap.L().Fatal("failed to create build metrics", zap.Error(err))
 	}
+
+	sandboxFactory := sandbox.NewFactory(networkPool, devicePool, featureFlags, true)
+
 	builder := build.NewBuilder(
 		logger,
+		sandboxFactory,
 		persistenceTemplate,
 		persistenceBuild,
 		artifactRegistry,
-		devicePool,
-		networkPool,
 		sandboxProxy,
 		sandboxes,
 		templateCache,

packages/orchestrator/internal/config/config.go

@@ -13,7 +13,6 @@ import (
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
-	globalconfig "github.com/e2b-dev/infra/packages/orchestrator/internal/config"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/block"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/build"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/fc"
@@ -24,6 +23,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/metadata"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
+	featureflags "github.com/e2b-dev/infra/packages/shared/pkg/feature-flags"
 	"github.com/e2b-dev/infra/packages/shared/pkg/grpc/orchestrator"
 	sbxlogger "github.com/e2b-dev/infra/packages/shared/pkg/logger/sandbox"
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage"
@@ -117,12 +117,32 @@ type networkSlotRes struct {
 	err  error
 }
 
+type Factory struct {
+	networkPool  *network.Pool
+	devicePool   *nbd.DevicePool
+	featureFlags *featureflags.Client
+
+	defaultAllowInternetAccess bool
+}
+
+func NewFactory(
+	networkPool *network.Pool,
+	devicePool *nbd.DevicePool,
+	featureFlags *featureflags.Client,
+	defaultAllowInternetAccess bool,
+) *Factory {
+	return &Factory{
+		networkPool:                networkPool,
+		devicePool:                 devicePool,
+		featureFlags:               featureFlags,
+		defaultAllowInternetAccess: defaultAllowInternetAccess,
+	}
+}
+
 // CreateSandbox creates the sandbox.
 // IMPORTANT: You must Close() the sandbox after you are done with it.
-func CreateSandbox(
+func (f *Factory) CreateSandbox(
 	ctx context.Context,
-	networkPool *network.Pool,
-	devicePool *nbd.DevicePool,
 	config Config,
 	runtime RuntimeMetadata,
 	fcVersions fc.FirecrackerVersions,
@@ -148,12 +168,13 @@ func CreateSandbox(
 		}
 	}()
 
-	allowInternet := globalconfig.AllowSandboxInternet
+	// TODO: Temporarily set this based on global config, should be removed later (it should be passed as a parameter in build)
+	allowInternet := f.defaultAllowInternetAccess
 	if config.AllowInternetAccess != nil {
 		allowInternet = *config.AllowInternetAccess
 	}
 
-	ipsCh := getNetworkSlotAsync(ctx, networkPool, cleanup, allowInternet)
+	ipsCh := getNetworkSlotAsync(ctx, f.networkPool, cleanup, allowInternet)
 	defer func() {
 		// Ensure the slot is received from chan so the slot is cleaned up properly in cleanup
 		<-ipsCh
@@ -172,7 +193,7 @@ func CreateSandbox(
 		rootfsProvider, err = rootfs.NewNBDProvider(
 			rootFS,
 			sandboxFiles.SandboxCacheRootfsPath(),
-			devicePool,
+			f.devicePool,
 		)
 	} else {
 		rootfsProvider, err = rootfs.NewDirectProvider(
@@ -305,17 +326,14 @@ func endSpan(span trace.Span, err error) {
 
 // ResumeSandbox resumes the sandbox from already saved template or snapshot.
 // IMPORTANT: You must Close() the sandbox after you are done with it.
-func ResumeSandbox(
+func (f *Factory) ResumeSandbox(
 	ctx context.Context,
-	networkPool *network.Pool,
 	t template.Template,
 	config Config,
 	runtime RuntimeMetadata,
 	traceID string,
 	startedAt time.Time,
 	endAt time.Time,
-	devicePool *nbd.DevicePool,
-	useClickhouseMetrics bool,
 	apiConfigToStore *orchestrator.SandboxConfig,
 ) (s *Sandbox, e error) {
 	ctx, span := tracer.Start(ctx, "resume sandbox")
@@ -334,12 +352,14 @@ func ResumeSandbox(
 		}
 	}()
 
-	allowInternet := globalconfig.AllowSandboxInternet
+	// TODO: Temporarily set this based on global config, should be removed later
+	//  (it should be passed as a non nil parameter from API)
+	allowInternet := f.defaultAllowInternetAccess
 	if config.AllowInternetAccess != nil {
 		allowInternet = *config.AllowInternetAccess
 	}
 
-	ipsCh := getNetworkSlotAsync(ctx, networkPool, cleanup, allowInternet)
+	ipsCh := getNetworkSlotAsync(ctx, f.networkPool, cleanup, allowInternet)
 	defer func() {
 		// Ensure the slot is received from chan before ResumeSandbox returns so the slot is cleaned up properly in cleanup
 		<-ipsCh
@@ -360,7 +380,7 @@ func ResumeSandbox(
 	rootfsOverlay, err := rootfs.NewNBDProvider(
 		readonlyRootfs,
 		sandboxFiles.SandboxCacheRootfsPath(),
-		devicePool,
+		f.devicePool,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create rootfs overlay: %w", err)
@@ -508,6 +528,11 @@ func ResumeSandbox(
 		exit: exit,
 	}
 
+	useClickhouseMetrics, flagErr := f.featureFlags.BoolFlag(ctx, featureflags.MetricsWriteFlagName)
+	if flagErr != nil {
+		zap.L().Error("soft failing during metrics write feature flag receive", zap.Error(flagErr))
+	}
+
 	// Part of the sandbox as we need to stop Checks before pausing the sandbox
 	// This is to prevent race condition of reporting unhealthy sandbox
 	sbx.Checks = NewChecks(sbx, useClickhouseMetrics)

packages/orchestrator/internal/sandbox/sandbox.go

@@ -27,6 +27,7 @@ import (
 type server struct {
 	orchestrator.UnimplementedSandboxServiceServer
 
+	sandboxFactory    *sandbox.Factory
 	info              *service.ServiceInfo
 	sandboxes         *smap.Map[*sandbox.Sandbox]
 	proxy             *proxy.SandboxProxy
@@ -56,22 +57,21 @@ type ServiceConfig struct {
 	TemplateCache    *template.Cache
 	Info             *service.ServiceInfo
 	Proxy            *proxy.SandboxProxy
+	SandboxFactory   *sandbox.Factory
 	Sandboxes        *smap.Map[*sandbox.Sandbox]
 	Persistence      storage.StorageProvider
 	FeatureFlags     *featureflags.Client
 	SbxEventsService events.EventsService[event.SandboxEvent]
 }
 
-func New(
-	ctx context.Context,
-	cfg ServiceConfig,
-) (*Service, error) {
+func New(cfg ServiceConfig) *Service {
 	srv := &Service{
 		info:        cfg.Info,
 		proxy:       cfg.Proxy,
 		persistence: cfg.Persistence,
 	}
 	srv.server = &server{
+		sandboxFactory:    cfg.SandboxFactory,
 		info:              cfg.Info,
 		proxy:             srv.proxy,
 		sandboxes:         cfg.Sandboxes,
@@ -96,5 +96,5 @@ func New(
 
 	orchestrator.RegisterSandboxServiceServer(cfg.GRPC.GRPCServer(), srv.server)
 
-	return srv, nil
+	return srv
 }

packages/orchestrator/internal/server/main.go

@@ -87,11 +87,6 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 	}
 	defer s.startingSandboxes.Release(1)
 
-	metricsWriteFlag, flagErr := s.featureFlags.BoolFlag(ctx, featureflags.MetricsWriteFlagName)
-	if flagErr != nil {
-		zap.L().Error("soft failing during metrics write feature flag receive", zap.Error(flagErr))
-	}
-
 	template, err := s.templateCache.GetTemplate(
 		ctx,
 		req.GetSandbox().GetBuildId(),
@@ -104,9 +99,8 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		return nil, fmt.Errorf("failed to get template snapshot data: %w", err)
 	}
 
-	sbx, err := sandbox.ResumeSandbox(
+	sbx, err := s.sandboxFactory.ResumeSandbox(
 		ctx,
-		s.networkPool,
 		template,
 		sandbox.Config{
 			BaseTemplateID: req.Sandbox.BaseTemplateId,
@@ -133,8 +127,6 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 		childSpan.SpanContext().TraceID().String(),
 		req.StartTime.AsTime(),
 		req.EndTime.AsTime(),
-		s.devicePool,
-		metricsWriteFlag,
 		req.Sandbox,
 	)
 	if err != nil {

packages/orchestrator/internal/server/sandboxes.go

@@ -13,8 +13,6 @@ import (
 
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/proxy"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox"
-	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/nbd"
-	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/network"
 	sbxtemplate "github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/template"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/buildcontext"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/template/build/commands"
@@ -42,10 +40,9 @@ var tracer = otel.Tracer("github.com/e2b-dev/infra/packages/orchestrator/interna
 type Builder struct {
 	logger *zap.Logger
 
+	sandboxFactory   *sandbox.Factory
 	templateStorage  storage.StorageProvider
 	buildStorage     storage.StorageProvider
-	devicePool       *nbd.DevicePool
-	networkPool      *network.Pool
 	artifactRegistry artifactsregistry.ArtifactsRegistry
 	proxy            *proxy.SandboxProxy
 	sandboxes        *smap.Map[*sandbox.Sandbox]
@@ -55,11 +52,10 @@ type Builder struct {
 
 func NewBuilder(
 	logger *zap.Logger,
+	sandboxFactory *sandbox.Factory,
 	templateStorage storage.StorageProvider,
 	buildStorage storage.StorageProvider,
 	artifactRegistry artifactsregistry.ArtifactsRegistry,
-	devicePool *nbd.DevicePool,
-	networkPool *network.Pool,
 	proxy *proxy.SandboxProxy,
 	sandboxes *smap.Map[*sandbox.Sandbox],
 	templateCache *sbxtemplate.Cache,
@@ -70,12 +66,11 @@ func NewBuilder(
 		templateStorage:  templateStorage,
 		buildStorage:     buildStorage,
 		artifactRegistry: artifactRegistry,
-		devicePool:       devicePool,
-		networkPool:      networkPool,
 		proxy:            proxy,
 		sandboxes:        sandboxes,
 		templateCache:    templateCache,
 		metrics:          buildMetrics,
+		sandboxFactory:   sandboxFactory,
 	}
 }
 
@@ -189,10 +184,9 @@ func runBuild(
 ) (*Result, error) {
 	index := cache.NewHashIndex(bc.CacheScope, builder.buildStorage, builder.templateStorage)
 
-	layerExecutor := layer.NewLayerExecutor(bc,
+	layerExecutor := layer.NewLayerExecutor(
+		bc,
 		builder.logger,
-		builder.networkPool,
-		builder.devicePool,
 		builder.templateCache,
 		builder.proxy,
 		builder.sandboxes,
@@ -206,12 +200,11 @@ func runBuild(
 		builder.logger,
 		builder.proxy,
 		builder.templateStorage,
-		builder.devicePool,
-		builder.networkPool,
 		builder.artifactRegistry,
 		layerExecutor,
 		index,
 		builder.metrics,
+		builder.sandboxFactory,
 	)
 
 	commandExecutor := commands.NewCommandExecutor(
@@ -222,6 +215,7 @@ func runBuild(
 
 	stepBuilders := steps.CreateStepPhases(
 		bc,
+		builder.sandboxFactory,
 		builder.logger,
 		builder.proxy,
 		layerExecutor,
@@ -232,6 +226,7 @@ func runBuild(
 
 	postProcessingBuilder := finalize.New(
 		bc,
+		builder.sandboxFactory,
 		builder.templateStorage,
 		builder.proxy,
 		layerExecutor,

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

@@ -20,9 +20,10 @@ import (
 
 // CreateSandbox creates sandboxes for new templates
 type CreateSandbox struct {
-	config     sandbox.Config
-	timeout    time.Duration
-	fcVersions fc.FirecrackerVersions
+	config         sandbox.Config
+	timeout        time.Duration
+	fcVersions     fc.FirecrackerVersions
+	sandboxFactory *sandbox.Factory
 
 	rootfsCachePath string
 }
@@ -33,12 +34,24 @@ const (
 
 var _ SandboxCreator = (*CreateSandbox)(nil)
 
-func NewCreateSandbox(config sandbox.Config, timeout time.Duration, fcVersions fc.FirecrackerVersions) *CreateSandbox {
-	return &CreateSandbox{config: config, timeout: timeout, fcVersions: fcVersions, rootfsCachePath: ""}
+func NewCreateSandbox(config sandbox.Config, sandboxFactory *sandbox.Factory, timeout time.Duration, fcVersions fc.FirecrackerVersions) *CreateSandbox {
+	return &CreateSandbox{
+		config:          config,
+		sandboxFactory:  sandboxFactory,
+		timeout:         timeout,
+		fcVersions:      fcVersions,
+		rootfsCachePath: "",
+	}
 }
 
-func NewCreateSandboxFromCache(config sandbox.Config, timeout time.Duration, fcVersions fc.FirecrackerVersions, rootfsCachePath string) *CreateSandbox {
-	return &CreateSandbox{config: config, timeout: timeout, fcVersions: fcVersions, rootfsCachePath: rootfsCachePath}
+func NewCreateSandboxFromCache(config sandbox.Config, sandboxFactory *sandbox.Factory, timeout time.Duration, fcVersions fc.FirecrackerVersions, rootfsCachePath string) *CreateSandbox {
+	return &CreateSandbox{
+		config:          config,
+		timeout:         timeout,
+		fcVersions:      fcVersions,
+		rootfsCachePath: rootfsCachePath,
+		sandboxFactory:  sandboxFactory,
+	}
 }
 
 func (cs *CreateSandbox) Sandbox(
@@ -65,10 +78,8 @@ func (cs *CreateSandbox) Sandbox(
 	}
 
 	// In case of a new sandbox, base template ID is now used as the potentially exported template base ID.
-	sbx, err := sandbox.CreateSandbox(
+	sbx, err := cs.sandboxFactory.CreateSandbox(
 		ctx,
-		layerExecutor.networkPool,
-		layerExecutor.devicePool,
 		cs.config,
 		sandbox.RuntimeMetadata{
 			TemplateID:  layerExecutor.Config.TemplateID,