feat(orchestrator): run startup reclaim on boot

packages/orchestrator/pkg/cfg/service.go

@@ -3,11 +3,14 @@
 package cfg
 
 import (
+	"slices"
 	"strings"
 )
 
 type ServiceType string
 
+type Services []ServiceType
+
 const (
 	UnknownService  ServiceType = "orch-unknown"
 	Orchestrator    ServiceType = "orchestrator"
@@ -29,10 +32,10 @@ func ParseServiceType(s string) ServiceType {
 
 // GetServices parses the ORCHESTRATOR_SERVICES environment variable
 // and returns a slice of known ServiceTypes.
-func GetServices(config Config) []ServiceType {
+func GetServices(config Config) Services {
 	rawServiceNames := config.Services
 
-	var services []ServiceType
+	services := make(Services, 0, len(rawServiceNames))
 	for _, name := range rawServiceNames {
 		service := ParseServiceType(name)
 		if service != UnknownService {
@@ -43,9 +46,25 @@ func GetServices(config Config) []ServiceType {
 	return services
 }
 
+func (s Services) Has(service ServiceType) bool {
+	return slices.Contains(s, service)
+}
+
+func (s Services) RunsOrchestrator() bool {
+	return s.Has(Orchestrator)
+}
+
+func (s Services) RunsTemplateManager() bool {
+	return s.Has(TemplateManager)
+}
+
+func (s Services) UsesSandboxRuntime() bool {
+	return s.RunsOrchestrator() || s.RunsTemplateManager()
+}
+
 // GetServiceName returns a single string identifier for the given services.
 // If multiple services are present, they are joined with underscores.
-func GetServiceName(services []ServiceType) string {
+func GetServiceName(services Services) string {
 	if len(services) == 0 {
 		return string(UnknownService)
 	}

packages/orchestrator/pkg/cfg/service_test.go

@@ -0,0 +1,66 @@
+//go:build linux
+
+package cfg
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestServicesCapabilities(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name               string
+		services           Services
+		runsOrchestrator   bool
+		runsTemplateMgr    bool
+		usesSandboxRuntime bool
+	}{
+		{
+			name:               "orchestrator",
+			services:           Services{Orchestrator},
+			runsOrchestrator:   true,
+			usesSandboxRuntime: true,
+		},
+		{
+			name:               "template manager",
+			services:           Services{TemplateManager},
+			runsTemplateMgr:    true,
+			usesSandboxRuntime: true,
+		},
+		{
+			name:               "combined",
+			services:           Services{Orchestrator, TemplateManager},
+			runsOrchestrator:   true,
+			runsTemplateMgr:    true,
+			usesSandboxRuntime: true,
+		},
+		{
+			name:     "unknown only",
+			services: Services{UnknownService},
+		},
+		{
+			name: "empty",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			assert.Equal(t, tt.runsOrchestrator, tt.services.RunsOrchestrator())
+			assert.Equal(t, tt.runsTemplateMgr, tt.services.RunsTemplateManager())
+			assert.Equal(t, tt.usesSandboxRuntime, tt.services.UsesSandboxRuntime())
+		})
+	}
+}
+
+func TestGetServicesFiltersUnknownServices(t *testing.T) {
+	t.Parallel()
+
+	services := GetServices(Config{Services: []string{" orchestrator ", "unknown", "TEMPLATE-MANAGER"}})
+
+	assert.Equal(t, Services{Orchestrator, TemplateManager}, services)
+}

packages/orchestrator/pkg/factories/run.go

@@ -51,6 +51,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/server"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/service"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/service/machineinfo"
+	"github.com/e2b-dev/infra/packages/orchestrator/pkg/startupreclaim"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/template/constants"
 	tmplserver "github.com/e2b-dev/infra/packages/orchestrator/pkg/template/server"
 	"github.com/e2b-dev/infra/packages/orchestrator/pkg/volumes"
@@ -180,7 +181,7 @@ func run(config cfg.Config, opts Options) (success bool) {
 	// Check if the orchestrator crashed and restarted
 	// Skip this check in development mode
 	// We don't want to lock if the service is running with force stop; the subsequent start would fail.
-	if !env.IsDevelopment() && !config.ForceStop && slices.Contains(services, cfg.Orchestrator) {
+	if !env.IsDevelopment() && !config.ForceStop && services.RunsOrchestrator() {
 		fileLockName := config.OrchestratorLockPath
 		info, err := os.Stat(fileLockName)
 		if err == nil {
@@ -621,6 +622,18 @@ func run(config cfg.Config, opts Options) (success bool) {
 		closers = append(closers, closer{"egress proxy", egressSetup.Close})
 	}
 
+	// Sandbox-runtime reclaim must run before newStorage below: reclaim deletes
+	// leaked ns-* from /run/netns, and NewStorageLocal snapshots the remaining
+	// namespaces as foreign at construction.
+	if services.UsesSandboxRuntime() && !config.DisableStartupReclaim {
+		startupreclaim.Run(ctx, startupreclaim.Config{
+			NetworkConfig: config.NetworkConfig,
+			EgressProxy:   egressSetup.Proxy,
+			CgroupManager: cgroupManager,
+			StorageConfig: config.StorageConfig,
+		})
+	}
+
 	// device pool
 	devicePool, err := nbd.NewDevicePool(config.NBDPoolSize)
 	if err != nil {
@@ -735,7 +748,7 @@ func run(config cfg.Config, opts Options) (success bool) {
 	// template manager
 	var tmpl *tmplserver.ServerStore
 	var localUploadHandler *localupload.Handler
-	if slices.Contains(services, cfg.TemplateManager) {
+	if services.RunsTemplateManager() {
 		buildPersistence, uploadHandler, err := setupBuildStorage(ctx, limiter, config)
 		if err != nil {
 			logger.L().Fatal(ctx, "failed to setup build storage", zap.Error(err))