feat(orchestrator): allow customizing number of pre-warmed NBDs (#2266)

* feat(orchestrator): allow customizing number of pre-warmed NBDs

* chore: add validation

* chore: use config instead of hardcoded value

Author: jakubno

packages/orchestrator/.env.local

@@ -20,4 +20,5 @@ PERSISTENT_VOLUME_MOUNTS=test-volume-type:./.data/test-volume
 REDIS_URL=localhost:6379
 SANDBOX_CACHE_DIR=./tmp/sandbox-cache-dir
 SNAPSHOT_CACHE_DIR=./tmp/snapshot-cache
+NBD_POOL_SIZE=16
 STORAGE_PROVIDER=Local

packages/orchestrator/benchmark_test.go

@@ -145,7 +145,7 @@ func BenchmarkBaseImageLaunch(b *testing.B) {
 		assert.NoError(b, err)
 	})
 
-	devicePool, err := nbd.NewDevicePool()
+	devicePool, err := nbd.NewDevicePool(config.NBDPoolSize)
 	require.NoError(b, err, "do you have the nbd kernel module installed?")
 	go func() {
 		devicePool.Populate(b.Context())

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

@@ -254,7 +254,18 @@ func doBuild(
 		return fmt.Errorf("build storage: %w", err)
 	}
 
-	devicePool, err := nbd.NewDevicePool()
+	blockMetrics, _ := blockmetrics.NewMetrics(noop.NewMeterProvider())
+
+	if os.Getenv("NODE_IP") == "" {
+		os.Setenv("NODE_IP", "127.0.0.1")
+	}
+
+	c, err := cfg.Parse()
+	if err != nil {
+		return fmt.Errorf("config: %w", err)
+	}
+
+	devicePool, err := nbd.NewDevicePool(c.NBDPoolSize)
 	if err != nil {
 		return fmt.Errorf("nbd pool: %w", err)
 	}
@@ -280,17 +291,6 @@ func doBuild(
 	}
 	defer dockerhubRepo.Close()
 
-	blockMetrics, _ := blockmetrics.NewMetrics(noop.NewMeterProvider())
-
-	if os.Getenv("NODE_IP") == "" {
-		os.Setenv("NODE_IP", "127.0.0.1")
-	}
-
-	c, err := cfg.Parse()
-	if err != nil {
-		return fmt.Errorf("config: %w", err)
-	}
-
 	templateCache, err := sbxtemplate.NewCache(c, featureFlags, persistenceTemplate, blockMetrics, peerclient.NopResolver())
 	if err != nil {
 		return fmt.Errorf("template cache: %w", err)

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

@@ -1013,7 +1013,7 @@ func run(ctx context.Context, buildID string, iterations int, coldStart, noPrefe
 	if verbose {
 		fmt.Println("🔧 Creating NBD device pool...")
 	}
-	devicePool, err := nbd.NewDevicePool()
+	devicePool, err := nbd.NewDevicePool(config.NBDPoolSize)
 	if err != nil {
 		return fmt.Errorf("nbd pool: %w", err)
 	}

packages/orchestrator/cmd/smoketest/smoke_test.go

@@ -187,7 +187,7 @@ func newTestInfra(t *testing.T, ctx context.Context) *testInfra {
 	require.NoError(t, err)
 
 	// NBD
-	devicePool, err := nbd.NewDevicePool()
+	devicePool, err := nbd.NewDevicePool(orcConfig.NBDPoolSize)
 	require.NoError(t, err)
 	go devicePool.Populate(ctx)
 	ti.closers = append(ti.closers, func(ctx context.Context) { devicePool.Close(ctx) })

packages/orchestrator/pkg/cfg/model.go

@@ -90,6 +90,7 @@ type Config struct {
 	RedisTLSCABase64           string            `env:"REDIS_TLS_CA_BASE64"`
 	RedisURL                   string            `env:"REDIS_URL"`
 	RedisPoolSize              int               `env:"REDIS_POOL_SIZE"               envDefault:"10"`
+	NBDPoolSize                int               `env:"NBD_POOL_SIZE"                 envDefault:"64"`
 	Services                   []string          `env:"ORCHESTRATOR_SERVICES"         envDefault:"orchestrator"`
 	PersistentVolumeMounts     map[string]string `env:"PERSISTENT_VOLUME_MOUNTS"`
 }

packages/orchestrator/pkg/factories/run.go

@@ -507,7 +507,7 @@ func run(config cfg.Config, opts Options) (success bool) {
 	}
 
 	// device pool
-	devicePool, err := nbd.NewDevicePool()
+	devicePool, err := nbd.NewDevicePool(config.NBDPoolSize)
 	if err != nil {
 		logger.L().Fatal(ctx, "failed to create device pool", zap.Error(err))
 	}

packages/orchestrator/pkg/sandbox/nbd/pool.go

@@ -20,9 +20,7 @@ import (
 	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
-// maxSlotsReady is the number of slots that are ready to be used.
 const (
-	maxSlotsReady                 = 64
 	waitOnNBDError                = 50 * time.Millisecond
 	devicePoolCloseReleaseTimeout = 10 * time.Minute
 )
@@ -82,7 +80,11 @@ type DevicePool struct {
 	slots chan DeviceSlot
 }
 
-func NewDevicePool() (*DevicePool, error) {
+func NewDevicePool(maxSlotsReady int) (*DevicePool, error) {
+	if maxSlotsReady <= 0 {
+		return nil, fmt.Errorf("maxSlotsReady must be > 0, got %d", maxSlotsReady)
+	}
+
 	maxDevices, err := getMaxDevices()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get max devices: %w", err)
@@ -95,7 +97,7 @@ func NewDevicePool() (*DevicePool, error) {
 	pool := &DevicePool{
 		done:      make(chan struct{}),
 		usedSlots: bitset.New(maxDevices),
-		slots:     make(chan DeviceSlot, int(math.Min(maxSlotsReady, float64(maxDevices)))),
+		slots:     make(chan DeviceSlot, int(math.Min(float64(maxSlotsReady), float64(maxDevices)))),
 	}
 
 	return pool, nil

packages/orchestrator/pkg/sandbox/nbd/testutils/nbd_device.go

@@ -13,7 +13,7 @@ import (
 func GetNBDDevice(ctx context.Context, backend block.Device, featureFlags *featureflags.Client) (nbd.DevicePath, *Cleaner, error) {
 	var cleaner Cleaner
 
-	devicePool, err := nbd.NewDevicePool()
+	devicePool, err := nbd.NewDevicePool(64)
 	if err != nil {
 		return "", &cleaner, fmt.Errorf("failed to create device pool: %w", err)
 	}