refactor(api): remove memory sandbox backends

.github/actions/start-services/action.yml

@@ -128,7 +128,6 @@ runs:
         # Client-proxy config
         API_INTERNAL_GRPC_ADDRESS: "localhost:5009"
         DEFAULT_PERSISTENT_VOLUME_TYPE: "test-volume-type"
-        SANDBOX_STORAGE_BACKEND: "redis"
         COMPRESS_ENABLED: ${{ inputs.compress_enabled }}
         COMPRESS_TYPE: ${{ inputs.compress_type }}
         COMPRESS_LEVEL: ${{ inputs.compress_level }}

iac/modules/job-api/jobs/api.hcl

@@ -185,8 +185,6 @@ job "api" {
         REDIS_TLS_CA_BASE64            = "${redis_tls_ca_base64}"
         REDIS_URL                      = "${redis_url}"
 
-        SANDBOX_STORAGE_BACKEND        = "${sandbox_storage_backend}"
-
 %{ if launch_darkly_api_key != "" }
         LAUNCH_DARKLY_API_KEY         = "${launch_darkly_api_key}"
 %{ endif }

iac/modules/job-api/main.tf

@@ -47,7 +47,6 @@ resource "nomad_job" "api" {
     clickhouse_connection_string            = var.clickhouse_connection_string
     loki_url                                = var.loki_url
     sandbox_access_token_hash_seed          = var.sandbox_access_token_hash_seed
-    sandbox_storage_backend                 = var.sandbox_storage_backend
     db_migrator_docker_image                = var.db_migrator_docker_image
     launch_darkly_api_key                   = trimspace(var.launch_darkly_api_key)
     default_persistent_volume_type          = var.default_persistent_volume_type

iac/modules/job-api/variables.tf

@@ -120,11 +120,6 @@ variable "sandbox_access_token_hash_seed" {
   sensitive = true
 }
 
-variable "sandbox_storage_backend" {
-  type    = string
-  default = "memory"
-}
-
 variable "redis_url" {
   type      = string
   sensitive = true

iac/provider-gcp/Makefile

@@ -40,7 +40,6 @@ tf_vars := \
 	$(call tfvar, GCP_ZONE) \
 	$(call tfvar, DOMAIN_NAME) \
 	$(call tfvar, PREFIX) \
-	$(call tfvar, SANDBOX_STORAGE_BACKEND) \
 	$(call tfvar, ORCHESTRATOR_ENABLED) \
 	$(call tfvar, ALLOW_SANDBOX_INTERNET) \
 	$(call tfvar, API_INTERNAL_GRPC_PORT) \

iac/provider-gcp/main.tf

@@ -258,7 +258,6 @@ module "nomad" {
   redis_cluster_url_secret_version                       = module.init.redis_cluster_url_secret_version
   redis_tls_ca_base64_secret_version                     = module.init.redis_tls_ca_base64_secret_version
   sandbox_access_token_hash_seed                         = random_password.sandbox_access_token_hash_seed.result
-  sandbox_storage_backend                                = var.sandbox_storage_backend
   db_max_open_connections                                = var.db_max_open_connections
   db_min_idle_connections                                = var.db_min_idle_connections
   auth_db_max_open_connections                           = var.auth_db_max_open_connections

iac/provider-gcp/nomad/main.tf

@@ -140,7 +140,6 @@ module "api" {
   clickhouse_connection_string            = local.clickhouse_connection_string
   loki_url                                = local.loki_url
   sandbox_access_token_hash_seed          = var.sandbox_access_token_hash_seed
-  sandbox_storage_backend                 = var.sandbox_storage_backend
   db_max_open_connections                 = var.db_max_open_connections
   db_min_idle_connections                 = var.db_min_idle_connections
   auth_db_max_open_connections            = var.auth_db_max_open_connections

iac/provider-gcp/nomad/variables.tf

@@ -113,11 +113,6 @@ variable "sandbox_access_token_hash_seed" {
   type = string
 }
 
-variable "sandbox_storage_backend" {
-  type    = string
-  default = "memory"
-}
-
 variable "db_max_open_connections" {
   type = number
 }

iac/provider-gcp/variables.tf

@@ -715,12 +715,6 @@ variable "loki_boot_disk_type" {
   default     = "pd-ssd"
 }
 
-variable "sandbox_storage_backend" {
-  description = "The sandbox storage backend to use. Valid values: 'memory', 'redis'."
-  type        = string
-  default     = ""
-}
-
 variable "db_max_open_connections" {
   type    = number
   default = 40

packages/api/go.mod

@@ -41,7 +41,6 @@ require (
 	github.com/launchdarkly/go-server-sdk/v7 v7.13.0
 	github.com/oapi-codegen/gin-middleware v1.0.2
 	github.com/oapi-codegen/runtime v1.4.0
-	github.com/orcaman/concurrent-map/v2 v2.0.1
 	github.com/posthog/posthog-go v0.0.0-20230801140217-d607812dee69
 	github.com/redis/go-redis/v9 v9.17.3
 	github.com/stretchr/testify v1.11.1
@@ -295,6 +294,7 @@ require (
 	github.com/opentracing-contrib/go-grpc v0.1.2 // indirect
 	github.com/opentracing-contrib/go-stdlib v1.1.0 // indirect
 	github.com/opentracing/opentracing-go v1.2.1-0.20220228012449-10b1cf09e00b // indirect
+	github.com/orcaman/concurrent-map/v2 v2.0.1 // indirect
 	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/paulmach/orb v0.11.1 // indirect
 	github.com/pb33f/jsonpath v0.8.2 // indirect

packages/api/internal/cfg/model.go

@@ -16,11 +16,6 @@ import (
 )
 
 const (
-	// SandboxStorageBackendMemory will use memory backend as a primary storage for sandbox data.
-	// It will also keep redis populated to allow for seamless migration to redis.
-	SandboxStorageBackendMemory = "memory"
-	SandboxStorageBackendRedis  = "redis"
-
 	// ServiceDiscoveryProviderNomad queries Nomad's HTTP API (the original / Nomad-based deploy).
 	ServiceDiscoveryProviderNomad = "nomad"
 	// ServiceDiscoveryProviderKubernetes queries the in-cluster K8s API (the K8s deploy).
@@ -92,10 +87,6 @@ type Config struct {
 
 	DefaultPersistentVolumeType string `env:"DEFAULT_PERSISTENT_VOLUME_TYPE"`
 
-	// SandboxStorageBackend selects the sandbox storage implementation.
-	// "redis" uses Redis directly; "populate_redis" uses in-memory with Redis shadow writes.
-	SandboxStorageBackend string `env:"SANDBOX_STORAGE_BACKEND" envDefault:"memory"`
-
 	DomainName string `env:"DOMAIN_NAME" envDefault:""`
 }
 
@@ -165,10 +156,6 @@ func Parse() (Config, error) {
 		config.AuthDBConnectionString = config.PostgresConnectionString
 	}
 
-	if !slices.Contains([]string{SandboxStorageBackendMemory, SandboxStorageBackendRedis}, config.SandboxStorageBackend) {
-		return config, fmt.Errorf("invalid sandbox storage backend: %s", config.SandboxStorageBackend)
-	}
-
 	if !slices.Contains([]string{ServiceDiscoveryProviderNomad, ServiceDiscoveryProviderKubernetes, ServiceDiscoveryProviderLocal}, config.ServiceDiscoveryProvider) {
 		return config, fmt.Errorf("invalid service discovery provider: %s", config.ServiceDiscoveryProvider)
 	}

packages/api/internal/cfg/model_test.go

@@ -51,13 +51,6 @@ func TestParse(t *testing.T) {
 		require.NoError(t, err)
 		assert.Equal(t, content, result.VolumesToken.SigningKey)
 	})
-
-	t.Run("test sandbox backend empty string", func(t *testing.T) {
-		t.Setenv("SANDBOX_STORAGE_BACKEND", "")
-		result, err := Parse()
-		require.NoError(t, err)
-		assert.Equal(t, SandboxStorageBackendMemory, result.SandboxStorageBackend)
-	})
 }
 
 // removeEnv was mostly copied from the implementation of t.Setenv

packages/api/internal/orchestrator/autoresume_test.go

@@ -9,19 +9,29 @@ import (
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.opentelemetry.io/otel/metric/noop"
 
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations"
-	sandboxmemory "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/memory"
+	redisreservations "github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations/redis"
+	sandboxredis "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
-func newTestAutoResumeOrchestrator() *Orchestrator {
+func newTestAutoResumeOrchestrator(t *testing.T) *Orchestrator {
+	t.Helper()
+
+	client := redis_utils.SetupInstance(t)
+	storage, err := sandboxredis.NewStorage(client, noop.NewMeterProvider())
+	require.NoError(t, err)
+	go storage.Start(t.Context())
+	t.Cleanup(func() { storage.Close(context.WithoutCancel(t.Context())) })
+
 	return &Orchestrator{
 		sandboxStore: sandbox.NewStore(
-			sandboxmemory.NewStorage(),
-			reservations.NewReservationStorage(),
+			storage,
+			redisreservations.NewReservationStorage(client, storage.Notifier()),
 			sandbox.Callbacks{
 				AddSandboxToRoutingTable: func(context.Context, sandbox.Sandbox) {},
 				AsyncNewlyCreatedSandbox: func(context.Context, sandbox.Sandbox, sandbox.CreationMetadata) {},
@@ -64,7 +74,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("running sandbox returns node ip immediately", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		registerNode(o, sbx, "10.0.0.1")
 
@@ -77,7 +87,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("running sandbox with empty ip returns error", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		registerNode(o, sbx, "")
 
@@ -90,7 +100,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("snapshotting sandbox waits and routes when transition finishes", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		addSandbox(t, o, sbx)
 		registerNode(o, sbx, "10.0.0.2")
@@ -118,7 +128,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("pausing sandbox returns still transitioning after retries", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		addSandbox(t, o, sbx)
 
@@ -138,10 +148,10 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 		assert.ErrorIs(t, err, ErrSandboxStillTransitioning)
 	})
 
-	t.Run("pausing sandbox wait failure returns internal error", func(t *testing.T) {
+	t.Run("pausing sandbox wait failure returns still transitioning", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		addSandbox(t, o, sbx)
 
@@ -157,13 +167,13 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 		_, handled, err := o.HandleExistingSandboxAutoResume(t.Context(), sbx.TeamID, sbx.SandboxID, pausingSandbox, time.Minute)
 		require.Error(t, err)
 		assert.False(t, handled)
-		assert.EqualError(t, err, "error waiting for sandbox to pause")
+		assert.ErrorIs(t, err, ErrSandboxStillTransitioning)
 	})
 
 	t.Run("pausing sandbox wait timeout returns failed precondition", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateRunning)
 		addSandbox(t, o, sbx)
 
@@ -183,7 +193,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("killing sandbox returns not found", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.StateKilling)
 
 		_, handled, err := o.HandleExistingSandboxAutoResume(t.Context(), sbx.TeamID, sbx.SandboxID, sbx, time.Minute)
@@ -195,7 +205,7 @@ func TestHandleExistingSandboxAutoResume(t *testing.T) {
 	t.Run("unknown sandbox state returns internal error", func(t *testing.T) {
 		t.Parallel()
 
-		o := newTestAutoResumeOrchestrator()
+		o := newTestAutoResumeOrchestrator(t)
 		sbx := testSandboxForAutoResume(sandbox.State("mystery"))
 
 		_, handled, err := o.HandleExistingSandboxAutoResume(t.Context(), sbx.TeamID, sbx.SandboxID, sbx, time.Minute)

packages/api/internal/orchestrator/create_instance_events_test.go

@@ -18,9 +18,10 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/placement"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations"
-	sandboxmemory "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/memory"
+	redisreservations "github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations/redis"
+	sandboxredis "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
@@ -55,9 +56,15 @@ func newOrchestratorWithCounter(t *testing.T) (*Orchestrator, *eventCounter) {
 
 	ec := &eventCounter{}
 
+	client := redis_utils.SetupInstance(t)
+	storage, err := sandboxredis.NewStorage(client, noop.NewMeterProvider())
+	require.NoError(t, err)
+	go storage.Start(t.Context())
+	t.Cleanup(func() { storage.Close(context.WithoutCancel(t.Context())) })
+
 	store := sandbox.NewStore(
-		sandboxmemory.NewStorage(),
-		reservations.NewReservationStorage(),
+		storage,
+		redisreservations.NewReservationStorage(client, storage.Notifier()),
 		sandbox.Callbacks{
 			AddSandboxToRoutingTable: func(context.Context, sandbox.Sandbox) {},
 			AsyncNewlyCreatedSandbox: ec.callback(),

packages/api/internal/orchestrator/create_instance_test.go

@@ -15,12 +15,13 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/placement"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations"
-	sandboxmemory "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/memory"
+	redisreservations "github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations/redis"
+	sandboxredis "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis"
 	teamtypes "github.com/e2b-dev/infra/packages/auth/pkg/types"
 	authqueries "github.com/e2b-dev/infra/packages/db/pkg/auth/queries"
 	"github.com/e2b-dev/infra/packages/db/queries"
 	"github.com/e2b-dev/infra/packages/shared/pkg/featureflags"
+	redis_utils "github.com/e2b-dev/infra/packages/shared/pkg/redis"
 	"github.com/e2b-dev/infra/packages/shared/pkg/smap"
 )
 
@@ -45,9 +46,15 @@ func testBuild() queries.EnvBuild {
 func newCreateSandboxTestOrchestrator(t *testing.T) (*Orchestrator, *nodemanager.Node) {
 	t.Helper()
 
+	client := redis_utils.SetupInstance(t)
+	storage, err := sandboxredis.NewStorage(client, noop.NewMeterProvider())
+	require.NoError(t, err)
+	go storage.Start(t.Context())
+	t.Cleanup(func() { storage.Close(context.WithoutCancel(t.Context())) })
+
 	store := sandbox.NewStore(
-		sandboxmemory.NewStorage(),
-		reservations.NewReservationStorage(),
+		storage,
+		redisreservations.NewReservationStorage(client, storage.Notifier()),
 		sandbox.Callbacks{
 			AddSandboxToRoutingTable: func(context.Context, sandbox.Sandbox) {},
 			AsyncNewlyCreatedSandbox: func(context.Context, sandbox.Sandbox, sandbox.CreationMetadata) {},
@@ -57,8 +64,8 @@ func newCreateSandboxTestOrchestrator(t *testing.T) (*Orchestrator, *nodemanager
 	meter := noop.NewMeterProvider().Meter("github.com/e2b-dev/infra/packages/api/internal/orchestrator")
 	counter, _ := meter.Int64Counter("test-created-sandboxes")
 
-	ffClient, err := featureflags.NewClientWithDatasource(ldtestdata.DataSource())
-	require.NoError(t, err)
+	ffClient, ffErr := featureflags.NewClientWithDatasource(ldtestdata.DataSource())
+	require.NoError(t, ffErr)
 
 	algo := placement.NewBestOfK(placement.DefaultBestOfKConfig()).(*placement.BestOfK)
 

packages/api/internal/orchestrator/orchestrator.go

@@ -21,10 +21,7 @@ import (
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/nodemanager"
 	"github.com/e2b-dev/infra/packages/api/internal/orchestrator/placement"
 	"github.com/e2b-dev/infra/packages/api/internal/sandbox"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations"
 	redisreservations "github.com/e2b-dev/infra/packages/api/internal/sandbox/reservations/redis"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/memory"
-	"github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/populate_redis"
 	redisbackend "github.com/e2b-dev/infra/packages/api/internal/sandbox/storage/redis"
 	sqlcdb "github.com/e2b-dev/infra/packages/db/client"
 	"github.com/e2b-dev/infra/packages/shared/pkg/env"
@@ -156,27 +153,9 @@ func New(
 		snapshotUpsertSem: snapshotUpsertSem,
 	}
 
-	var reservationStorage sandbox.ReservationStorage
-	var sandboxStorage sandbox.Storage
-
-	switch config.SandboxStorageBackend {
-	case cfg.SandboxStorageBackendMemory:
-		reservationStorage = reservations.NewReservationStorage()
-		sandboxStorage = populate_redis.NewStorage(memory.NewStorage(), redisStorage)
-		logger.L().Info(ctx, "Using populate_redis sandbox storage backend")
-
-		go redisbackend.NewCleaner(redisStorage).Start(ctx)
-	case cfg.SandboxStorageBackendRedis:
-		reservationStorage = redisreservations.NewReservationStorage(redisClient, redisStorage.Notifier())
-		sandboxStorage = redisStorage
-		logger.L().Info(ctx, "Using redis sandbox storage backend")
-	default:
-		return nil, fmt.Errorf("invalid sandbox storage backend: %s", config.SandboxStorageBackend)
-	}
-
 	o.sandboxStore = sandbox.NewStore(
-		sandboxStorage,
-		reservationStorage,
+		redisStorage,
+		redisreservations.NewReservationStorage(redisClient, redisStorage.Notifier()),
 		sandbox.Callbacks{
 			AddSandboxToRoutingTable: o.addSandboxToRoutingTable,
 			AsyncNewlyCreatedSandbox: o.handleNewlyCreatedSandbox,