From 1b9501ced51cd8975c175d226def6f50418ad5ac Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Mon, 1 Jun 2026 16:57:51 -0700 Subject: [PATCH 1/6] feat(orchestrator): add draining to orchestrator --- packages/orchestrator/pkg/factories/run.go | 58 +++++-- .../orchestrator/pkg/sandbox/fc/process.go | 87 +++++++--- .../pkg/sandbox/fc/process_test.go | 37 +++++ packages/orchestrator/pkg/sandbox/map.go | 117 +++++++++++++- packages/orchestrator/pkg/sandbox/map_test.go | 136 ++++++++++++++++ .../pkg/sandbox/network/network.go | 36 +++-- packages/orchestrator/pkg/sandbox/sandbox.go | 42 ++--- packages/orchestrator/pkg/server/main.go | 153 +++++++++++++++++- packages/orchestrator/pkg/server/sandboxes.go | 18 ++- packages/orchestrator/pkg/server/utils.go | 63 ++++++++ .../orchestrator/pkg/server/utils_test.go | 57 +++++++ 11 files changed, 725 insertions(+), 79 deletions(-) create mode 100644 packages/orchestrator/pkg/sandbox/fc/process_test.go create mode 100644 packages/orchestrator/pkg/sandbox/map_test.go create mode 100644 packages/orchestrator/pkg/server/utils_test.go diff --git a/packages/orchestrator/pkg/factories/run.go b/packages/orchestrator/pkg/factories/run.go index 513df32b44..5370716254 100644 --- a/packages/orchestrator/pkg/factories/run.go +++ b/packages/orchestrator/pkg/factories/run.go @@ -119,6 +119,16 @@ func (e serviceDoneError) Error() string { return fmt.Sprintf("service %s finished", e.name) } +func isServiceDoneError(err error) bool { + var sde serviceDoneError + + return errors.As(err, &sde) +} + +func isIgnorableSyncError(err error) bool { + return errors.Is(err, syscall.EINVAL) +} + // Run starts the orchestrator, blocking until shutdown. // Returns true on clean shutdown. func Run(opts Options) bool { @@ -234,7 +244,7 @@ func run(config cfg.Config, opts Options) (success bool) { // there's a panic. defer func(g *errgroup.Group) { err := g.Wait() - if err != nil { + if err != nil && !isServiceDoneError(err) { log.Printf("error while shutting down: %v", err) success = false } @@ -275,7 +285,7 @@ func run(config cfg.Config, opts Options) (success bool) { })) defer func(l logger.Logger) { err := l.Sync() - if err != nil { + if err != nil && !isIgnorableSyncError(err) { log.Printf("error while shutting down logger: %v", err) success = false } @@ -293,7 +303,7 @@ func run(config cfg.Config, opts Options) (success bool) { ) defer func(l logger.Logger) { err := l.Sync() - if err != nil { + if err != nil && !isIgnorableSyncError(err) { log.Printf("error while shutting down sandbox logger: %v", err) success = false } @@ -311,7 +321,7 @@ func run(config cfg.Config, opts Options) (success bool) { ) defer func(l logger.Logger) { err := l.Sync() - if err != nil { + if err != nil && !isIgnorableSyncError(err) { log.Printf("error while shutting down sandbox logger: %v", err) success = false } @@ -575,8 +585,8 @@ func run(config cfg.Config, opts Options) (success bool) { if err != nil { logger.L().Fatal(ctx, "failed to create orchestrator server", zap.Error(err)) } - closers = append(closers, closer{"orchestrator server", func(context.Context) error { - return orchestratorService.Close() + closers = append(closers, closer{"orchestrator server", func(closeCtx context.Context) error { + return orchestratorService.Close(closeCtx) }}) // template manager sandbox logger @@ -591,8 +601,7 @@ func run(config cfg.Config, opts Options) (success bool) { ) closers = append(closers, closer{ "template manager sandbox logger", func(context.Context) error { - // Sync returns EINVAL when path is /dev/stdout (for example) - if err := tmplSbxLoggerExternal.Sync(); err != nil && !errors.Is(err, syscall.EINVAL) { + if err := tmplSbxLoggerExternal.Sync(); err != nil && !isIgnorableSyncError(err) { return err } @@ -782,6 +791,36 @@ func run(config cfg.Config, opts Options) (success bool) { } } + forceStopSandboxes := func() { + forceErr := orchestratorService.ForceStopSandboxes(context.WithoutCancel(ctx)) + if forceErr != nil { + logger.L().Error(ctx, "forced sandbox shutdown failed", zap.Error(forceErr)) + success = false + } + } + + closers = append(closers, closer{"sandbox drain", func(context.Context) error { + logger.L().Info(ctx, "Starting sandbox drain phase", + zap.Bool("forced", config.ForceStop), + zap.Int("sandbox_count", sandboxes.Count()), + ) + + if config.ForceStop { + forceStopSandboxes() + + return nil + } + + err := orchestratorService.DrainSandboxes(closeCtx) + if err != nil { + logger.L().Warn(ctx, "sandbox drain phase did not complete gracefully; forcing sandbox shutdown", zap.Error(err)) + + forceStopSandboxes() + } + + return nil + }}) + slices.Reverse(closers) for _, closer := range closers { clog := globalLogger.With(zap.String("service", closer.name), zap.Bool("forced", config.ForceStop)) @@ -793,8 +832,7 @@ func run(config cfg.Config, opts Options) (success bool) { } logger.L().Info(ctx, "Waiting for services to finish") - var sde serviceDoneError - if err := g.Wait(); err != nil && !errors.As(err, &sde) { + if err := g.Wait(); err != nil && !isServiceDoneError(err) { logger.L().Error(ctx, "service group error", zap.Error(err)) success = false } diff --git a/packages/orchestrator/pkg/sandbox/fc/process.go b/packages/orchestrator/pkg/sandbox/fc/process.go index 0074176a48..eae4bb4db0 100644 --- a/packages/orchestrator/pkg/sandbox/fc/process.go +++ b/packages/orchestrator/pkg/sandbox/fc/process.go @@ -661,21 +661,25 @@ func (p *Process) Stop(ctx context.Context) error { logger.L().Warn(ctx, "failed to remove fc metrics FIFO", zap.Error(removeErr), logger.WithSandboxID(p.files.SandboxID)) } - // Check if process has already exited. + pid := p.cmd.Process.Pid + + // Check if process has already exited. The parent exiting is not enough for + // cleanup: descendants in the same process group can still hold resources. select { case <-p.Exit.Done(): logger.L().Info(ctx, "fc process already exited", logger.WithSandboxID(p.files.SandboxID)) - - return nil + if !processGroupExists(pid) { + return nil + } default: } // this function should never fail b/c a previous context was canceled. ctx = context.WithoutCancel(ctx) - err := p.cmd.Process.Signal(syscall.SIGTERM) + err := signalProcessGroup(pid, syscall.SIGTERM) if err != nil { - if errors.Is(err, os.ErrProcessDone) { + if errors.Is(err, os.ErrProcessDone) && !processGroupExists(pid) { logger.L().Info(ctx, "fc process already exited", logger.WithSandboxID(p.files.SandboxID)) return nil @@ -684,40 +688,79 @@ func (p *Process) Stop(ctx context.Context) error { logger.L().Warn(ctx, "failed to send SIGTERM to fc process", zap.Error(err), logger.WithSandboxID(p.files.SandboxID)) } - go func() { + termDeadline := time.NewTimer(10 * time.Second) + defer termDeadline.Stop() + poll := time.NewTicker(50 * time.Millisecond) + defer poll.Stop() + + for processGroupExists(pid) { select { - // Wait 10 sec for the FC process to exit, if it doesn't, send SIGKILL. - case <-time.After(10 * time.Second): - // Check process status right before Kill — the pre-SIGTERM status - // captured above is 10s stale and no longer useful here. - status, stateErr := getProcessStatus(p.cmd.Process.Pid) + case <-termDeadline.C: + status, stateErr := getProcessStatus(pid) if errors.Is(stateErr, process.ErrorProcessNotRunning) { - // Process already exited, no need to send SIGKILL. - return + logger.L().Info(ctx, "fc parent process exited before SIGKILL; checking process group", logger.WithSandboxID(p.files.SandboxID)) } else if stateErr != nil { logger.L().Warn(ctx, "failed to get fc process status before SIGKILL", zap.Error(stateErr), logger.WithSandboxID(p.files.SandboxID)) } - err := p.cmd.Process.Kill() - if err == nil { - logger.L().Info(ctx, "sent SIGKILL to fc process because it was not responding to SIGTERM for 10 seconds", + killErr := signalProcessGroup(pid, syscall.SIGKILL) + if killErr == nil { + logger.L().Info(ctx, "sent SIGKILL to fc process group because it was not responding to SIGTERM for 10 seconds", zap.Strings("status", status), logger.WithSandboxID(p.files.SandboxID), ) } - if err != nil && !errors.Is(err, os.ErrProcessDone) { - logger.L().Warn(ctx, "failed to send SIGKILL to fc process", zap.Error(err), logger.WithSandboxID(p.files.SandboxID)) + if killErr != nil && !errors.Is(killErr, os.ErrProcessDone) { + logger.L().Warn(ctx, "failed to send SIGKILL to fc process", zap.Error(killErr), logger.WithSandboxID(p.files.SandboxID)) } - // If the FC process exited, we can return. - case <-p.Exit.Done(): - return + killDeadline := time.NewTimer(time.Second) + for processGroupExists(pid) { + select { + case <-killDeadline.C: + return fmt.Errorf("fc process group %d still exists after SIGKILL", pid) + case <-poll.C: + } + } + killDeadline.Stop() + + return nil + case <-poll.C: } - }() + } return nil } +func signalProcessGroup(pid int, signal syscall.Signal) error { + if pid <= 0 { + return os.ErrProcessDone + } + + // Firecracker is launched with Setsid, so the process PID is also the process + // group ID. Signal the group so unshare/bash/ip descendants cannot keep the VM + // mount namespace or Firecracker process alive after shutdown. + if err := syscall.Kill(-pid, signal); err != nil { + if errors.Is(err, syscall.ESRCH) { + return os.ErrProcessDone + } + + return err + } + + return nil +} + +func processGroupExists(pid int) bool { + if pid <= 0 { + return false + } + + err := syscall.Kill(-pid, 0) + + return err == nil || errors.Is(err, syscall.EPERM) +} + func (p *Process) Pause(ctx context.Context) error { ctx, childSpan := tracer.Start(ctx, "pause-fc") defer childSpan.End() diff --git a/packages/orchestrator/pkg/sandbox/fc/process_test.go b/packages/orchestrator/pkg/sandbox/fc/process_test.go new file mode 100644 index 0000000000..245c757044 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/fc/process_test.go @@ -0,0 +1,37 @@ +//go:build linux + +package fc + +import ( + "context" + "os/exec" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestProcessGroupExistsForSetsidChild(t *testing.T) { + t.Parallel() + + ctx, cancel := context.WithTimeout(t.Context(), time.Minute) + defer cancel() + + cmd := exec.CommandContext(ctx, "sleep", "60") + cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true} + require.NoError(t, cmd.Start()) + + pid := cmd.Process.Pid + t.Cleanup(func() { + _ = syscall.Kill(-pid, syscall.SIGKILL) + _ = cmd.Wait() + }) + + require.True(t, processGroupExists(pid)) + require.NoError(t, signalProcessGroup(pid, syscall.SIGKILL)) + require.Error(t, cmd.Wait()) + require.Eventually(t, func() bool { + return !processGroupExists(pid) + }, time.Second, 10*time.Millisecond) +} diff --git a/packages/orchestrator/pkg/sandbox/map.go b/packages/orchestrator/pkg/sandbox/map.go index b14aa209c0..da38efe849 100644 --- a/packages/orchestrator/pkg/sandbox/map.go +++ b/packages/orchestrator/pkg/sandbox/map.go @@ -9,6 +9,8 @@ import ( "net" "sync" + "go.uber.org/zap" + "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/smap" ) @@ -25,14 +27,49 @@ type MapSubscriber interface { OnNetworkRelease(ctx context.Context, sbx *Sandbox) } -// Map holds sandboxes that are live (running) together with a IP-to-sandbox index -// The two maps are managed independently. +type SandboxState string + +const ( + SandboxStateRunning SandboxState = "running" + SandboxStateStopping SandboxState = "stopping" +) + +type lifecycleEntry struct { + sandbox *Sandbox + stateLock sync.RWMutex + state SandboxState +} + +func newLifecycleEntry(sbx *Sandbox, state SandboxState) *lifecycleEntry { + return &lifecycleEntry{ + sandbox: sbx, + state: state, + } +} + +func (e *lifecycleEntry) setState(state SandboxState) { + e.stateLock.Lock() + defer e.stateLock.Unlock() + + e.state = state +} + +func (e *lifecycleEntry) getState() SandboxState { + e.stateLock.RLock() + defer e.stateLock.RUnlock() + + return e.state +} + +// Map holds sandboxes that are live (running), known active lifecycles, +// together with a IP-to-sandbox index. The indexes are managed independently. // // AssignNetwork/NetworkReleased manage the IP map, // MarkRunning/MarkStopping manage the live set. type Map struct { - live *smap.Map[*Sandbox] - network *smap.Map[*Sandbox] + live *smap.Map[*Sandbox] + lifecycles *smap.Map[*lifecycleEntry] + network *smap.Map[*Sandbox] subs []MapSubscriber subsLock sync.RWMutex @@ -40,11 +77,16 @@ type Map struct { func NewSandboxesMap() *Map { return &Map{ - live: smap.New[*Sandbox](), - network: smap.New[*Sandbox](), + live: smap.New[*Sandbox](), + lifecycles: smap.New[*lifecycleEntry](), + network: smap.New[*Sandbox](), } } +func sandboxLifecycleKey(sandboxID, lifecycleID string) string { + return fmt.Sprintf("%s/%s", sandboxID, lifecycleID) +} + func (m *Map) Subscribe(subscriber MapSubscriber) { m.subsLock.Lock() defer m.subsLock.Unlock() @@ -73,6 +115,35 @@ func (m *Map) Get(sandboxID string) (*Sandbox, bool) { return m.live.Get(sandboxID) } +func (m *Map) LifecycleItems() []*Sandbox { + entries := m.lifecycles.Items() + sandboxes := make([]*Sandbox, 0, len(entries)) + for _, entry := range entries { + sandboxes = append(sandboxes, entry.sandbox) + } + + return sandboxes +} + +func (m *Map) LifecycleItemsByState(states ...SandboxState) []*Sandbox { + stateSet := make(map[SandboxState]struct{}, len(states)) + for _, state := range states { + stateSet[state] = struct{}{} + } + + entries := m.lifecycles.Items() + sandboxes := make([]*Sandbox, 0, len(entries)) + for _, entry := range entries { + if _, ok := stateSet[entry.getState()]; !ok { + continue + } + + sandboxes = append(sandboxes, entry.sandbox) + } + + return sandboxes +} + // GetByHostPort looks up a sandbox by its host IP address parsed from hostPort. func (m *Map) GetByHostPort(hostPort string) (*Sandbox, error) { reqIP, _, err := net.SplitHostPort(hostPort) @@ -100,12 +171,25 @@ func (m *Map) AssignNetwork(ctx context.Context, sbx *Sandbox) { ) } +func (m *Map) TrackLifecycle(ctx context.Context, sbx *Sandbox, state SandboxState) { + m.lifecycles.Insert(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID), newLifecycleEntry(sbx, state)) + + logger.L().Info(ctx, "sandbox lifecycle tracked", + logger.WithSandboxID(sbx.Runtime.SandboxID), + logger.WithLifecycleID(sbx.LifecycleID), + logger.WithSandboxIP(sbx.Slot.HostIPString()), + zap.String("state", string(state)), + ) +} + // MarkRunning makes the sandbox visible to Get/Items/Count and notifies OnInsert subscribers. func (m *Map) MarkRunning(ctx context.Context, sbx *Sandbox) { if !m.live.InsertIfAbsent(sbx.Runtime.SandboxID, sbx) { return } + m.TrackLifecycle(ctx, sbx, SandboxStateRunning) + m.trigger(ctx, func(ctx context.Context, s MapSubscriber) { s.OnInsert(ctx, sbx) }) @@ -126,6 +210,7 @@ func (m *Map) MarkRunning(ctx context.Context, sbx *Sandbox) { // Returns true if the sandbox was successfully removed. func (m *Map) MarkStopping(ctx context.Context, sandboxID, lifecycleID string) bool { stopped := false + m.markLifecycleState(sandboxID, lifecycleID, SandboxStateStopping) m.live.RemoveCb(sandboxID, func(_ string, sbx *Sandbox, exists bool) bool { if !exists { @@ -150,6 +235,26 @@ func (m *Map) MarkStopping(ctx context.Context, sandboxID, lifecycleID string) b return stopped } +func (m *Map) MarkStopped(ctx context.Context, sbx *Sandbox) { + m.lifecycles.Remove(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID)) + + logger.L().Info(ctx, "sandbox lifecycle stopped", + logger.WithSandboxID(sbx.Runtime.SandboxID), + logger.WithLifecycleID(sbx.LifecycleID), + logger.WithSandboxIP(sbx.Slot.HostIPString()), + ) +} + +func (m *Map) markLifecycleState(sandboxID, lifecycleID string, state SandboxState) { + key := sandboxLifecycleKey(sandboxID, lifecycleID) + entry, ok := m.lifecycles.Get(key) + if !ok { + return + } + + entry.setState(state) +} + // NetworkReleased unregisters a sandbox's IP and notifies OnNetworkRelease // subscribers after a successful removal. func (m *Map) NetworkReleased(ctx context.Context, ip string) { diff --git a/packages/orchestrator/pkg/sandbox/map_test.go b/packages/orchestrator/pkg/sandbox/map_test.go new file mode 100644 index 0000000000..54358363e1 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/map_test.go @@ -0,0 +1,136 @@ +//go:build linux + +package sandbox + +import ( + "sync" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox/network" +) + +func TestMapMarkRunningTracksLifecycle(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + + sandboxes.MarkRunning(t.Context(), sbx) + require.Len(t, sandboxes.Items(), 1) + require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) +} + +func TestMapLifecycleItemsRemainAfterMarkStopping(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + + sandboxes.MarkRunning(t.Context(), sbx) + require.Len(t, sandboxes.Items(), 1) + require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) + + marked := sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID) + require.True(t, marked) + require.Empty(t, sandboxes.Items()) + require.Len(t, sandboxes.LifecycleItems(), 1) + require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateStopping), 1) + + sandboxes.MarkStopped(t.Context(), sbx) + require.Empty(t, sandboxes.LifecycleItems()) +} + +func TestSandboxCloseMarksLifecycleStopped(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + sbx.cleanup = NewCleanup() + sbx.sandboxes = sandboxes + + sandboxes.MarkRunning(t.Context(), sbx) + require.Len(t, sandboxes.LifecycleItems(), 1) + + require.NoError(t, sbx.Close(t.Context())) + require.Empty(t, sandboxes.LifecycleItems()) +} + +func TestMapLifecycleItemsAllowDuplicateSandboxIDs(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + oldSbx := testMapSandbox(t, "lifecycle-old") + newSbx := testMapSandbox(t, "lifecycle-new") + + sandboxes.TrackLifecycle(t.Context(), oldSbx, SandboxStateStopping) + sandboxes.TrackLifecycle(t.Context(), newSbx, SandboxStateRunning) + + require.Len(t, sandboxes.LifecycleItems(), 2) + require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateStopping), 1) + require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) +} + +func TestMapLifecycleStateUpdateAfterRemovalDoesNotResurrect(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + + sandboxes.TrackLifecycle(t.Context(), sbx, SandboxStateRunning) + entry, ok := sandboxes.lifecycles.Get(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID)) + require.True(t, ok) + + sandboxes.MarkStopped(t.Context(), sbx) + entry.setState(SandboxStateStopping) + + require.Empty(t, sandboxes.LifecycleItems()) +} + +func TestMapConcurrentMarkStoppingAndStoppedDoesNotResurrectLifecycle(t *testing.T) { + t.Parallel() + + for range 1000 { + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + sandboxes.MarkRunning(t.Context(), sbx) + + start := make(chan struct{}) + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + <-start + sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID) + }() + + go func() { + defer wg.Done() + <-start + sandboxes.MarkStopped(t.Context(), sbx) + }() + + close(start) + wg.Wait() + + require.Empty(t, sandboxes.LifecycleItems()) + } +} + +func testMapSandbox(t *testing.T, lifecycleID string) *Sandbox { + t.Helper() + + slot, err := network.NewSlot("test", 1, network.Config{}, network.NoopEgressProxy{}) + require.NoError(t, err) + + return &Sandbox{ + LifecycleID: lifecycleID, + Metadata: &Metadata{ + Config: NewConfig(Config{}), + Runtime: RuntimeMetadata{SandboxID: "sandbox-1"}, + }, + Resources: &Resources{Slot: slot}, + } +} diff --git a/packages/orchestrator/pkg/sandbox/network/network.go b/packages/orchestrator/pkg/sandbox/network/network.go index 7e39467f83..2164470318 100644 --- a/packages/orchestrator/pkg/sandbox/network/network.go +++ b/packages/orchestrator/pkg/sandbox/network/network.go @@ -328,24 +328,26 @@ func (s *Slot) RemoveNetwork() error { } } - // Delete NFS proxy redirect rule - err = tables.Delete("nat", "PREROUTING", - "--in-interface", s.VethName(), "--protocol", "tcp", - "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "2049", - "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.NFSProxyPort)), - ) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting sandbox NFS proxy redirect rule: %w", err)) - } + if tables != nil { + // Delete NFS proxy redirect rule + err = tables.Delete("nat", "PREROUTING", + "--in-interface", s.VethName(), "--protocol", "tcp", + "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "2049", + "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.NFSProxyPort)), + ) + if err != nil { + errs = append(errs, fmt.Errorf("error deleting sandbox NFS proxy redirect rule: %w", err)) + } - // Delete portmapper redirect rule - err = tables.Delete("nat", "PREROUTING", - "--in-interface", s.VethName(), "--protocol", "tcp", - "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "111", - "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.PortmapperPort)), - ) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting sandbox portmapper redirect rule: %w", err)) + // Delete portmapper redirect rule + err = tables.Delete("nat", "PREROUTING", + "--in-interface", s.VethName(), "--protocol", "tcp", + "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "111", + "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.PortmapperPort)), + ) + if err != nil { + errs = append(errs, fmt.Errorf("error deleting sandbox portmapper redirect rule: %w", err)) + } } err = netns.DeleteNamed(s.NamespaceID()) diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index b24ce64c06..45fdc2df28 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -236,6 +236,8 @@ type Sandbox struct { files *storage.SandboxFiles cleanup *Cleanup + sandboxes *Map + featureFlags *featureflags.Client process *fc.Process @@ -477,10 +479,11 @@ func (f *Factory) CreateSandbox( Metadata: metadata, cgroupHandle: cgroupHandle, - Template: template, - config: f.config, - files: sandboxFiles, - process: fcHandle, + Template: template, + config: f.config, + files: sandboxFiles, + process: fcHandle, + sandboxes: f.Sandboxes, cleanup: cleanup, featureFlags: f.featureFlags, @@ -822,10 +825,11 @@ func (f *Factory) ResumeSandbox( Metadata: metadata, cgroupHandle: cgroupHandle, - Template: t, - config: f.config, - files: sandboxFiles, - process: fcHandle, + Template: t, + config: f.config, + files: sandboxFiles, + process: fcHandle, + sandboxes: f.Sandboxes, cleanup: cleanup, featureFlags: f.featureFlags, @@ -964,6 +968,10 @@ func (s *Sandbox) Wait(ctx context.Context) error { func (s *Sandbox) Close(ctx context.Context) error { err := s.cleanup.Run(ctx) + if s.sandboxes != nil { + s.sandboxes.MarkStopped(context.WithoutCancel(ctx), s) + } + if err != nil { return fmt.Errorf("failed to cleanup sandbox: %w", err) } @@ -994,9 +1002,11 @@ func (s *Sandbox) doStop(ctx context.Context) error { errs = append(errs, fmt.Errorf("failed to stop FC: %w", fcStopErr)) } - // The process exited, we can continue with the rest of the cleanup. - // We could use select with ctx.Done() to wait for cancellation, but if the process is not exited the whole cleanup will be in a bad state and will result in unexpected behavior. - <-s.process.Exit.Done() + // The process should exit before the rest of cleanup, but memory shutdown + // must still run if the wait context is canceled so UFFD can exit. + if waitErr := s.process.Exit.WaitWithContext(ctx); waitErr != nil { + errs = append(errs, fmt.Errorf("failed waiting for FC exit: %w", waitErr)) + } uffdStopErr := s.Resources.memory.Stop() if uffdStopErr != nil { @@ -1388,15 +1398,7 @@ func getNetworkSlot( ctx, span := tracer.Start(ctx, "clean network-slot") defer span.End() - // We can run this cleanup asynchronously, as it is not important for the sandbox lifecycle - go func(ctx context.Context) { - returnErr := networkPool.Return(ctx, slot, networkReleased, network.ReturnDelay) - if returnErr != nil { - logger.L().Error(ctx, "failed to return network slot", zap.Error(returnErr)) - } - }(context.WithoutCancel(ctx)) - - return nil + return networkPool.Return(ctx, slot, networkReleased, network.ReturnDelay) }) return slot, nil diff --git a/packages/orchestrator/pkg/server/main.go b/packages/orchestrator/pkg/server/main.go index 37c3a0f513..aa53e65070 100644 --- a/packages/orchestrator/pkg/server/main.go +++ b/packages/orchestrator/pkg/server/main.go @@ -4,6 +4,7 @@ package server import ( "context" + "errors" "fmt" "sync" "time" @@ -38,6 +39,19 @@ const uploadedBuildsTTL = 1 * time.Hour // MaxStartingInstancesPerNode feature flag and resize the semaphore. const startingSandboxesLimitRefreshInterval = 30 * time.Second +const sandboxDrainPollInterval = 5 * time.Second + +func sandboxDrainLogInterval(elapsed time.Duration) time.Duration { + switch { + case elapsed < time.Minute: + return 5 * time.Second + case elapsed < time.Hour: + return time.Minute + default: + return 15 * time.Minute + } +} + type Server struct { orchestrator.UnimplementedSandboxServiceServer orchestrator.UnimplementedChunkServiceServer @@ -61,6 +75,9 @@ type Server struct { done chan struct{} closeOnce sync.Once + + sandboxStartMu sync.RWMutex + sandboxLifecycleWG sync.WaitGroup } type ServiceConfig struct { @@ -156,16 +173,148 @@ func New(ctx context.Context, cfg ServiceConfig) (*Server, error) { return server, nil } -func (s *Server) Close() error { +func (s *Server) Close(ctx context.Context) error { + s.startDraining(ctx) + s.uploadedBuilds.Stop() + + return nil +} + +func (s *Server) startDraining(ctx context.Context) { s.closeOnce.Do(func() { + logger.L().Info(ctx, "orchestrator server entering sandbox drain mode", + zap.Int("live_sandboxes", s.sandboxFactory.Sandboxes.Count()), + ) close(s.done) }) +} - s.uploadedBuilds.Stop() +func (s *Server) DrainSandboxes(ctx context.Context) error { + s.startDraining(ctx) + if err := s.waitSandboxStarts(ctx); err != nil { + return err + } + + live := s.sandboxFactory.Sandboxes.Count() + logger.L().Info(ctx, "starting graceful sandbox drain", zap.Int("live_sandboxes", live)) + if live == 0 { + logger.L().Info(ctx, "graceful sandbox drain complete", zap.Int("live_sandboxes", live)) + + return s.waitSandboxLifecycles(ctx) + } + + ticker := time.NewTicker(sandboxDrainPollInterval) + defer ticker.Stop() + startedAt := time.Now() + lastLoggedAt := startedAt + + for { + select { + case <-ctx.Done(): + remaining := s.sandboxFactory.Sandboxes.Count() + logger.L().Warn(ctx, "graceful sandbox drain timed out", + zap.Int("remaining_sandboxes", remaining), + zap.Error(ctx.Err()), + ) + + return ctx.Err() + case <-ticker.C: + now := time.Now() + remaining := s.sandboxFactory.Sandboxes.Count() + if remaining == 0 { + logger.L().Info(ctx, "graceful sandbox drain complete", zap.Int("live_sandboxes", remaining)) + + return s.waitSandboxLifecycles(ctx) + } + + elapsed := now.Sub(startedAt) + if now.Sub(lastLoggedAt) >= sandboxDrainLogInterval(elapsed) { + logger.L().Info(ctx, "waiting for sandbox drain", + zap.Int("remaining_sandboxes", remaining), + zap.Duration("elapsed", elapsed), + ) + lastLoggedAt = now + } + } + } +} + +func (s *Server) ForceStopSandboxes(ctx context.Context) error { + s.startDraining(ctx) + if err := s.waitSandboxStarts(ctx); err != nil { + return err + } + + sandboxes := s.sandboxFactory.Sandboxes.LifecycleItems() + logger.L().Warn(ctx, "starting forced sandbox shutdown", zap.Int("sandbox_count", len(sandboxes))) + if len(sandboxes) == 0 { + return s.waitSandboxLifecycles(ctx) + } + + var wg sync.WaitGroup + errCh := make(chan error, len(sandboxes)) + + for _, sbx := range sandboxes { + wg.Go(func() { + sbxLog := logger.L().With( + logger.WithSandboxID(sbx.Runtime.SandboxID), + logger.WithLifecycleID(sbx.LifecycleID), + logger.WithSandboxIP(sbx.Slot.HostIPString()), + ) + sbxLog.Warn(ctx, "force stopping sandbox during orchestrator shutdown") + + marked := s.sandboxFactory.Sandboxes.MarkStopping(ctx, sbx.Runtime.SandboxID, sbx.LifecycleID) + if !marked { + sbxLog.Info(ctx, "sandbox was already removed from live map before force stop") + } + + if err := sbx.Stop(ctx); err != nil { + errCh <- fmt.Errorf("stop sandbox %s/%s: %w", sbx.Runtime.SandboxID, sbx.LifecycleID, err) + sbxLog.Error(ctx, "failed to force stop sandbox", zap.Error(err)) + } + + sbxLog.Info(ctx, "forced sandbox stop requested") + }) + } + + wg.Wait() + close(errCh) + + var errs []error + for err := range errCh { + errs = append(errs, err) + } + + if err := s.waitSandboxLifecycles(ctx); err != nil { + errs = append(errs, err) + } + + if err := errors.Join(errs...); err != nil { + logger.L().Error(ctx, "forced sandbox shutdown finished with errors", zap.Error(err)) + + return err + } + + logger.L().Info(ctx, "forced sandbox shutdown complete") return nil } +func (s *Server) waitSandboxLifecycles(ctx context.Context) error { + done := make(chan struct{}) + go func() { + s.sandboxLifecycleWG.Wait() + close(done) + }() + + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for sandbox lifecycle cleanup: %w", ctx.Err()) + case <-done: + return nil + } +} + func (s *Server) refreshStartingSandboxesLimit(ctx context.Context) { ticker := time.NewTicker(startingSandboxesLimitRefreshInterval) defer ticker.Stop() diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index b1783f1381..40fce97c9f 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -90,6 +90,11 @@ func (s *Server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ telemetry.WithEnvdVersion(req.GetSandbox().GetEnvdVersion()), ) + if err := s.enterSandboxStart(ctx, "sandbox-create"); err != nil { + return nil, err + } + defer s.leaveSandboxStart() + // setup launch darkly ctx = featureflags.AddToContext( ctx, @@ -187,6 +192,10 @@ func (s *Server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ SandboxType: sandbox.SandboxTypeSandbox, } + if err := s.rejectIfDraining(ctx, "sandbox-create-before-start"); err != nil { + return nil, err + } + sbx, err := s.sandboxFactory.ResumeSandbox( ctx, template, @@ -621,6 +630,11 @@ func (s *Server) Checkpoint(ctx context.Context, in *orchestrator.SandboxCheckpo Build(), ) + if err := s.enterSandboxStart(ctx, "sandbox-checkpoint"); err != nil { + return nil, err + } + defer s.leaveSandboxStart() + sbx, ok := s.sandboxFactory.Sandboxes.Get(in.GetSandboxId()) if !ok { telemetry.ReportCriticalError(ctx, "sandbox not found", nil, telemetry.WithSandboxID(in.GetSandboxId())) @@ -906,7 +920,7 @@ func (s *Server) uploadSnapshotAsync(ctx context.Context, sbx *sandbox.Sandbox, // setupSandboxLifecycle sets up the cleanup goroutine for a sandbox. func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox) { - go func() { + s.sandboxLifecycleWG.Go(func() { ctx, childSpan := tracer.Start(context.WithoutCancel(ctx), "stop sandbox-lifecycle", trace.WithNewRoot()) defer childSpan.End() @@ -926,7 +940,7 @@ func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox } sbxlogger.E(sbx).Info(ctx, "Sandbox stopped") - }() + }) } // stopSandboxAsync stops the sandbox in a background goroutine. diff --git a/packages/orchestrator/pkg/server/utils.go b/packages/orchestrator/pkg/server/utils.go index 3ef314b3e6..f6d6e91f6d 100644 --- a/packages/orchestrator/pkg/server/utils.go +++ b/packages/orchestrator/pkg/server/utils.go @@ -4,14 +4,77 @@ package server import ( "context" + "fmt" + "time" + "go.uber.org/zap" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/telemetry" ) +const sandboxStartWaitPollInterval = 50 * time.Millisecond + +func (s *Server) rejectIfDraining(ctx context.Context, operation string) error { + select { + case <-s.done: + logger.L().Info(ctx, "rejecting sandbox operation during orchestrator drain", zap.String("operation", operation)) + + return status.Error(codes.Unavailable, "orchestrator is draining") + default: + return nil + } +} + +func (s *Server) enterSandboxStart(ctx context.Context, operation string) error { + if err := s.rejectIfDraining(ctx, operation); err != nil { + return err + } + + s.sandboxStartMu.RLock() + if err := s.rejectIfDraining(ctx, operation); err != nil { + s.sandboxStartMu.RUnlock() + + return err + } + + return nil +} + +func (s *Server) leaveSandboxStart() { + s.sandboxStartMu.RUnlock() +} + +func (s *Server) waitSandboxStarts(ctx context.Context) error { + logger.L().Info(ctx, "waiting for in-flight sandbox start operations to finish") + + ticker := time.NewTicker(sandboxStartWaitPollInterval) + defer ticker.Stop() + + for { + if s.sandboxStartMu.TryLock() { + logger.L().Info(ctx, "in-flight sandbox start gate acquired") + s.sandboxStartMu.Unlock() + logger.L().Info(ctx, "in-flight sandbox start operations finished") + + return nil + } + + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for in-flight sandbox start operations: %w", ctx.Err()) + case <-ticker.C: + } + } +} + func (s *Server) waitForAcquire(ctx context.Context) error { + if err := s.rejectIfDraining(ctx, "wait-for-acquire"); err != nil { + return err + } + ctx, cancel := context.WithTimeout(ctx, acquireTimeout) defer cancel() diff --git a/packages/orchestrator/pkg/server/utils_test.go b/packages/orchestrator/pkg/server/utils_test.go new file mode 100644 index 0000000000..04bb6dabc1 --- /dev/null +++ b/packages/orchestrator/pkg/server/utils_test.go @@ -0,0 +1,57 @@ +//go:build linux + +package server + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +func TestWaitSandboxStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { + t.Parallel() + + s := &Server{done: make(chan struct{})} + + s.sandboxStartMu.RLock() + defer s.sandboxStartMu.RUnlock() + + waitCtx, cancel := context.WithCancel(t.Context()) + waitErr := make(chan error, 1) + go func() { + waitErr <- s.waitSandboxStarts(waitCtx) + }() + + // Give waitSandboxStarts a chance to observe the held read lock. The old + // implementation left a queued writer here, which blocked future RLock calls. + time.Sleep(2 * sandboxStartWaitPollInterval) + cancel() + + select { + case err := <-waitErr: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("waitSandboxStarts did not return after cancellation") + } + + close(s.done) + + enterErr := make(chan error, 1) + go func() { + enterErr <- s.enterSandboxStart(t.Context(), "test") + }() + + select { + case err := <-enterErr: + if err == nil { + s.leaveSandboxStart() + } + require.Equal(t, codes.Unavailable, status.Code(err)) + case <-time.After(time.Second): + t.Fatal("enterSandboxStart blocked instead of rejecting while draining") + } +} From 49d0b3af8b58c3c78768e13173044436ec4c1d4c Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Tue, 19 May 2026 16:39:11 -0700 Subject: [PATCH 2/6] fix(orchestrator): harden shutdown network cleanup --- .../pkg/sandbox/network/firewall.go | 12 +- .../pkg/sandbox/network/network.go | 117 ++++++++++++------ .../pkg/sandbox/network/network_test.go | 49 ++++++++ .../orchestrator/pkg/sandbox/network/pool.go | 22 +++- 4 files changed, 159 insertions(+), 41 deletions(-) create mode 100644 packages/orchestrator/pkg/sandbox/network/network_test.go diff --git a/packages/orchestrator/pkg/sandbox/network/firewall.go b/packages/orchestrator/pkg/sandbox/network/firewall.go index 5889667fdc..08d9d5f35b 100644 --- a/packages/orchestrator/pkg/sandbox/network/firewall.go +++ b/packages/orchestrator/pkg/sandbox/network/firewall.go @@ -3,6 +3,7 @@ package network import ( + "errors" "fmt" "net/netip" "slices" @@ -109,7 +110,16 @@ func NewFirewall(tapIf string, orchestratorInternalIP string, extraAllowedCIDRs } func (fw *Firewall) Close() error { - return fw.conn.CloseLasting() + fw.conn.DelTable(&nftables.Table{ + Name: tableName, + Family: nftables.TableFamilyINet, + }) + deleteErr := fw.conn.Flush() + if errors.Is(deleteErr, unix.ENOENT) { + deleteErr = nil + } + + return errors.Join(deleteErr, fw.conn.CloseLasting()) } // tapIfaceMatch returns expressions that match packets from the tap interface. diff --git a/packages/orchestrator/pkg/sandbox/network/network.go b/packages/orchestrator/pkg/sandbox/network/network.go index 2164470318..14a386f9b7 100644 --- a/packages/orchestrator/pkg/sandbox/network/network.go +++ b/packages/orchestrator/pkg/sandbox/network/network.go @@ -7,6 +7,7 @@ import ( "errors" "fmt" "net" + "os" "runtime" "strconv" @@ -14,11 +15,67 @@ import ( "github.com/vishvananda/netlink" "github.com/vishvananda/netns" "go.uber.org/zap" + "golang.org/x/sys/unix" "github.com/e2b-dev/infra/packages/shared/pkg/logger" ) -func (s *Slot) CreateNetwork(ctx context.Context) error { +type notExistError interface { + IsNotExist() bool +} + +type multiUnwrapError interface { + Unwrap() []error +} + +func ignoreExpectedAbsent(err error, isExpected func(error) bool) bool { + if err == nil { + return true + } + + var joined multiUnwrapError + if errors.As(err, &joined) { + for _, child := range joined.Unwrap() { + if !ignoreExpectedAbsent(child, isExpected) { + return false + } + } + + return true + } + + return isExpected(err) +} + +func isIPTablesNotExist(err error) bool { + var notExist notExistError + + return errors.As(err, ¬Exist) && notExist.IsNotExist() +} + +func isRouteNotExist(err error) bool { + return errors.Is(err, unix.ESRCH) || errors.Is(err, unix.ENOENT) +} + +func isLinkNotExist(err error) bool { + var linkNotFound netlink.LinkNotFoundError + + return errors.As(err, &linkNotFound) || errors.Is(err, unix.ENODEV) || errors.Is(err, unix.ENOENT) +} + +func isNamespaceNotExist(err error) bool { + return os.IsNotExist(err) || errors.Is(err, unix.ENOENT) +} + +func appendUnlessExpectedAbsentf(errs *[]error, err error, isExpected func(error) bool, format string) { + if ignoreExpectedAbsent(err, isExpected) { + return + } + + *errs = append(*errs, fmt.Errorf(format, err)) +} + +func (s *Slot) CreateNetwork(ctx context.Context) (retErr error) { // Prevent thread changes so we can safely manipulate with namespaces runtime.LockOSThread() defer runtime.UnlockOSThread() @@ -29,10 +86,19 @@ func (s *Slot) CreateNetwork(ctx context.Context) error { return fmt.Errorf("cannot get current (host) namespace: %w", err) } + cleanupNeeded := false defer func() { - err = netns.Set(hostNS) - if err != nil { - logger.L().Error(ctx, "error resetting network namespace back to the host namespace", zap.Error(err)) + restoreErr := netns.Set(hostNS) + if restoreErr != nil { + logger.L().Error(ctx, "error resetting network namespace back to the host namespace", zap.Error(restoreErr)) + } + + if retErr != nil && cleanupNeeded { + if restoreErr != nil { + retErr = errors.Join(retErr, fmt.Errorf("error resetting network namespace back to the host namespace before cleanup: %w", restoreErr)) + } else if cleanupErr := s.RemoveNetwork(); cleanupErr != nil { + retErr = errors.Join(retErr, fmt.Errorf("error cleaning up partially created network: %w", cleanupErr)) + } } err = hostNS.Close() @@ -46,6 +112,7 @@ func (s *Slot) CreateNetwork(ctx context.Context) error { if err != nil { return fmt.Errorf("cannot create new namespace: %w", err) } + cleanupNeeded = true defer ns.Close() @@ -273,20 +340,14 @@ func (s *Slot) RemoveNetwork() error { } else { // Delete host forwarding rules err = tables.Delete("filter", "FORWARD", "-i", s.VethName(), "-o", defaultGateway, "-j", "ACCEPT") - if err != nil { - errs = append(errs, fmt.Errorf("error deleting host forwarding rule to default gateway: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting host forwarding rule to default gateway: %w") err = tables.Delete("filter", "FORWARD", "-i", defaultGateway, "-o", s.VethName(), "-j", "ACCEPT") - if err != nil { - errs = append(errs, fmt.Errorf("error deleting host forwarding rule from default gateway: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting host forwarding rule from default gateway: %w") // Delete host postrouting rules err = tables.Delete("nat", "POSTROUTING", "-s", s.HostCIDR(), "-o", defaultGateway, "-j", "MASQUERADE") - if err != nil { - errs = append(errs, fmt.Errorf("error deleting host postrouting rule: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting host postrouting rule: %w") // Delete hyperloop proxy redirect rule err = tables.Delete( @@ -294,15 +355,11 @@ func (s *Slot) RemoveNetwork() error { "-p", "tcp", "-d", s.config.OrchestratorInSandboxIPAddress, "--dport", "80", "-j", "REDIRECT", "--to-port", s.hyperloopPort, ) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting sandbox hyperloop proxy redirect rule: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting sandbox hyperloop proxy redirect rule: %w") // Delete changes made by egress proxy err = s.egressProxy.OnSlotDelete(s, tables) - if err != nil { - errs = append(errs, err) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "%w") } // Delete routing from host to FC namespace @@ -310,9 +367,7 @@ func (s *Slot) RemoveNetwork() error { Gw: s.VpeerIP(), Dst: s.HostNet(), }) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting route from host to FC: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isRouteNotExist, "error deleting route from host to FC: %w") // Delete veth device // We explicitly delete the veth device from the host namespace because even though deleting @@ -320,12 +375,10 @@ func (s *Slot) RemoveNetwork() error { // the same name immediately after deleting the namespace. veth, err := netlink.LinkByName(s.VethName()) if err != nil { - errs = append(errs, fmt.Errorf("error finding veth: %w", err)) + appendUnlessExpectedAbsentf(&errs, err, isLinkNotExist, "error finding veth: %w") } else { err = netlink.LinkDel(veth) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting veth device: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isLinkNotExist, "error deleting veth device: %w") } if tables != nil { @@ -335,9 +388,7 @@ func (s *Slot) RemoveNetwork() error { "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "2049", "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.NFSProxyPort)), ) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting sandbox NFS proxy redirect rule: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting sandbox NFS proxy redirect rule: %w") // Delete portmapper redirect rule err = tables.Delete("nat", "PREROUTING", @@ -345,15 +396,11 @@ func (s *Slot) RemoveNetwork() error { "--destination", s.config.OrchestratorInSandboxIPAddress, "--dport", "111", "--jump", "REDIRECT", "--to-port", strconv.Itoa(int(s.config.PortmapperPort)), ) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting sandbox portmapper redirect rule: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isIPTablesNotExist, "error deleting sandbox portmapper redirect rule: %w") } err = netns.DeleteNamed(s.NamespaceID()) - if err != nil { - errs = append(errs, fmt.Errorf("error deleting namespace: %w", err)) - } + appendUnlessExpectedAbsentf(&errs, err, isNamespaceNotExist, "error deleting namespace: %w") return errors.Join(errs...) } diff --git a/packages/orchestrator/pkg/sandbox/network/network_test.go b/packages/orchestrator/pkg/sandbox/network/network_test.go new file mode 100644 index 0000000000..cfd582c3b2 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/network/network_test.go @@ -0,0 +1,49 @@ +//go:build linux + +package network + +import ( + "errors" + "fmt" + "os" + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/sys/unix" +) + +type fakeNotExistError struct { + notExist bool +} + +func (e fakeNotExistError) Error() string { + return "fake iptables error" +} + +func (e fakeNotExistError) IsNotExist() bool { + return e.notExist +} + +func TestIgnoreExpectedAbsentHandlesWrappedAndJoinedErrors(t *testing.T) { + t.Parallel() + + wrapped := fmt.Errorf("wrapped: %w", fakeNotExistError{notExist: true}) + joined := errors.Join(wrapped, fakeNotExistError{notExist: true}) + + require.True(t, ignoreExpectedAbsent(joined, isIPTablesNotExist)) + require.False(t, ignoreExpectedAbsent(errors.Join(joined, errors.New("boom")), isIPTablesNotExist)) + require.False(t, ignoreExpectedAbsent(fakeNotExistError{notExist: false}, isIPTablesNotExist)) +} + +func TestExpectedAbsentClassifiers(t *testing.T) { + t.Parallel() + + require.True(t, isRouteNotExist(fmt.Errorf("route delete failed: %w", unix.ESRCH))) + require.False(t, isRouteNotExist(unix.EPERM)) + + require.True(t, isLinkNotExist(fmt.Errorf("link delete failed: %w", unix.ENODEV))) + require.False(t, isLinkNotExist(unix.EPERM)) + + require.True(t, isNamespaceNotExist(&os.PathError{Op: "remove", Path: "/var/run/netns/missing", Err: unix.ENOENT})) + require.False(t, isNamespaceNotExist(unix.EPERM)) +} diff --git a/packages/orchestrator/pkg/sandbox/network/pool.go b/packages/orchestrator/pkg/sandbox/network/pool.go index f0d37b40a2..e5a1218eaf 100644 --- a/packages/orchestrator/pkg/sandbox/network/pool.go +++ b/packages/orchestrator/pkg/sandbox/network/pool.go @@ -108,15 +108,13 @@ func NewPool(newSlotsPoolSize, reusedSlotsPoolSize int, slotStorage Storage, con newSlots := make(chan *Slot, newSlotsPoolSize-1) reusedSlots := make(chan *Slot, reusedSlotsPoolSize) - pool := &Pool{ + return &Pool{ config: config, done: make(chan struct{}), newSlots: newSlots, reusedSlots: reusedSlots, slotStorage: slotStorage, } - - return pool } func (p *Pool) createNetworkSlot(ctx context.Context) (*Slot, error) { @@ -153,8 +151,22 @@ func (p *Pool) Populate(ctx context.Context) { continue } - newSlotsAvailableCounter.Add(ctx, 1) - p.newSlots <- slot + select { + case <-p.done: + if err := p.cleanup(context.WithoutCancel(ctx), slot); err != nil { + logger.L().Error(ctx, "[network slot pool]: failed to cleanup created slot while closing", zap.Error(err), zap.Int("slot_index", slot.Idx)) + } + + return + case <-ctx.Done(): + if err := p.cleanup(context.WithoutCancel(ctx), slot); err != nil { + logger.L().Error(ctx, "[network slot pool]: failed to cleanup created slot after context cancellation", zap.Error(err), zap.Int("slot_index", slot.Idx)) + } + + return + case p.newSlots <- slot: + newSlotsAvailableCounter.Add(ctx, 1) + } } } } From e6e9091196c73b522bcadd85da2c0a93c45850ad Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Tue, 2 Jun 2026 00:50:28 -0700 Subject: [PATCH 3/6] feat(orchestrator): gate template builds during drain --- packages/orchestrator/pkg/factories/run.go | 24 +++-- packages/orchestrator/pkg/server/main.go | 8 +- .../pkg/template/server/create_template.go | 15 ++++ .../pkg/template/server/delete_template.go | 17 +++- .../orchestrator/pkg/template/server/drain.go | 89 +++++++++++++++++++ .../pkg/template/server/drain_test.go | 80 +++++++++++++++++ .../orchestrator/pkg/template/server/main.go | 11 +++ .../server/upload_layer_files_template.go | 3 + 8 files changed, 228 insertions(+), 19 deletions(-) create mode 100644 packages/orchestrator/pkg/template/server/drain.go create mode 100644 packages/orchestrator/pkg/template/server/drain_test.go diff --git a/packages/orchestrator/pkg/factories/run.go b/packages/orchestrator/pkg/factories/run.go index 5370716254..4a4f4121ee 100644 --- a/packages/orchestrator/pkg/factories/run.go +++ b/packages/orchestrator/pkg/factories/run.go @@ -782,6 +782,8 @@ func run(config cfg.Config, opts Options) (success bool) { } } + orchestratorService.StartDraining(ctx) + // Wait for services to be drained before closing them if tmpl != nil { err := tmpl.Wait(closeCtx) @@ -799,27 +801,21 @@ func run(config cfg.Config, opts Options) (success bool) { } } - closers = append(closers, closer{"sandbox drain", func(context.Context) error { - logger.L().Info(ctx, "Starting sandbox drain phase", - zap.Bool("forced", config.ForceStop), - zap.Int("sandbox_count", sandboxes.Count()), - ) - - if config.ForceStop { - forceStopSandboxes() - - return nil - } + logger.L().Info(ctx, "Starting sandbox drain phase", + zap.Bool("forced", config.ForceStop), + zap.Int("sandbox_count", sandboxes.Count()), + ) + if config.ForceStop { + forceStopSandboxes() + } else { err := orchestratorService.DrainSandboxes(closeCtx) if err != nil { logger.L().Warn(ctx, "sandbox drain phase did not complete gracefully; forcing sandbox shutdown", zap.Error(err)) forceStopSandboxes() } - - return nil - }}) + } slices.Reverse(closers) for _, closer := range closers { diff --git a/packages/orchestrator/pkg/server/main.go b/packages/orchestrator/pkg/server/main.go index aa53e65070..1f66835e80 100644 --- a/packages/orchestrator/pkg/server/main.go +++ b/packages/orchestrator/pkg/server/main.go @@ -174,13 +174,13 @@ func New(ctx context.Context, cfg ServiceConfig) (*Server, error) { } func (s *Server) Close(ctx context.Context) error { - s.startDraining(ctx) + s.StartDraining(ctx) s.uploadedBuilds.Stop() return nil } -func (s *Server) startDraining(ctx context.Context) { +func (s *Server) StartDraining(ctx context.Context) { s.closeOnce.Do(func() { logger.L().Info(ctx, "orchestrator server entering sandbox drain mode", zap.Int("live_sandboxes", s.sandboxFactory.Sandboxes.Count()), @@ -190,7 +190,7 @@ func (s *Server) startDraining(ctx context.Context) { } func (s *Server) DrainSandboxes(ctx context.Context) error { - s.startDraining(ctx) + s.StartDraining(ctx) if err := s.waitSandboxStarts(ctx); err != nil { return err } @@ -240,7 +240,7 @@ func (s *Server) DrainSandboxes(ctx context.Context) error { } func (s *Server) ForceStopSandboxes(ctx context.Context) error { - s.startDraining(ctx) + s.StartDraining(ctx) if err := s.waitSandboxStarts(ctx); err != nil { return err } diff --git a/packages/orchestrator/pkg/template/server/create_template.go b/packages/orchestrator/pkg/template/server/create_template.go index 7aeb1aeb90..a4351ec8b0 100644 --- a/packages/orchestrator/pkg/template/server/create_template.go +++ b/packages/orchestrator/pkg/template/server/create_template.go @@ -106,6 +106,20 @@ func (s *ServerStore) TemplateCreate(ctx context.Context, templateRequest *templ } logs := buildlogger.NewLogEntryLogger() + if err := s.enterBuildStart(ctx, "template-create"); err != nil { + return nil, err + } + buildStartReleased := false + releaseBuildStart := func() { + if buildStartReleased { + return + } + + s.leaveBuildStart() + buildStartReleased = true + } + defer releaseBuildStart() + buildInfo, err := s.buildCache.Create(template.TeamID, metadata.BuildID, logs) if err != nil { return nil, fmt.Errorf("error while creating build cache: %w", err) @@ -123,6 +137,7 @@ func (s *ServerStore) TemplateCreate(ctx context.Context, templateRequest *templ s.wg.Add(1) s.activeBuilds.Add(1) + releaseBuildStart() go func(ctx context.Context) { defer s.wg.Done() defer s.activeBuilds.Add(-1) diff --git a/packages/orchestrator/pkg/template/server/delete_template.go b/packages/orchestrator/pkg/template/server/delete_template.go index 2165fd68f6..4b7b591e39 100644 --- a/packages/orchestrator/pkg/template/server/delete_template.go +++ b/packages/orchestrator/pkg/template/server/delete_template.go @@ -26,7 +26,22 @@ func (s *ServerStore) TemplateBuildDelete(ctx context.Context, in *templatemanag )) defer childSpan.End() + if err := s.enterBuildStart(ctx, "template-delete"); err != nil { + return nil, err + } + deleteStartReleased := false + releaseDeleteStart := func() { + if deleteStartReleased { + return + } + + s.leaveBuildStart() + deleteStartReleased = true + } + defer releaseDeleteStart() + s.wg.Add(1) + releaseDeleteStart() defer s.wg.Done() if in.GetTemplateID() == "" || in.GetBuildID() == "" { @@ -48,5 +63,5 @@ func (s *ServerStore) TemplateBuildDelete(ctx context.Context, in *templatemanag return nil, err } - return nil, nil + return &emptypb.Empty{}, nil } diff --git a/packages/orchestrator/pkg/template/server/drain.go b/packages/orchestrator/pkg/template/server/drain.go new file mode 100644 index 0000000000..78aef43f9f --- /dev/null +++ b/packages/orchestrator/pkg/template/server/drain.go @@ -0,0 +1,89 @@ +//go:build linux + +package server + +import ( + "context" + "fmt" + "time" + + "go.uber.org/zap" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + "github.com/e2b-dev/infra/packages/shared/pkg/logger" +) + +const buildStartWaitPollInterval = 50 * time.Millisecond + +func (s *ServerStore) StartDraining(ctx context.Context) { + s.drainOnce.Do(func() { + if s.drainDone == nil { + s.drainDone = make(chan struct{}) + } + + s.log().Info(ctx, "template server entering drain mode", zap.Int64("active_builds", s.activeBuilds.Load())) + close(s.drainDone) + }) +} + +func (s *ServerStore) rejectIfDraining(ctx context.Context, operation string) error { + select { + case <-s.drainDone: + s.log().Info(ctx, "rejecting template operation during orchestrator drain", zap.String("operation", operation)) + + return status.Error(codes.Unavailable, "orchestrator is draining") + default: + return nil + } +} + +func (s *ServerStore) enterBuildStart(ctx context.Context, operation string) error { + if err := s.rejectIfDraining(ctx, operation); err != nil { + return err + } + + s.buildStartMu.RLock() + if err := s.rejectIfDraining(ctx, operation); err != nil { + s.buildStartMu.RUnlock() + + return err + } + + return nil +} + +func (s *ServerStore) leaveBuildStart() { + s.buildStartMu.RUnlock() +} + +func (s *ServerStore) waitBuildStarts(ctx context.Context) error { + s.log().Info(ctx, "waiting for in-flight template build start operations to finish") + + ticker := time.NewTicker(buildStartWaitPollInterval) + defer ticker.Stop() + + for { + if s.buildStartMu.TryLock() { + s.log().Info(ctx, "in-flight template build start gate acquired") + s.buildStartMu.Unlock() + s.log().Info(ctx, "in-flight template build start operations finished") + + return nil + } + + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for in-flight template build start operations: %w", ctx.Err()) + case <-ticker.C: + } + } +} + +func (s *ServerStore) log() logger.Logger { + if s.logger != nil { + return s.logger + } + + return logger.L() +} diff --git a/packages/orchestrator/pkg/template/server/drain_test.go b/packages/orchestrator/pkg/template/server/drain_test.go new file mode 100644 index 0000000000..56192a503b --- /dev/null +++ b/packages/orchestrator/pkg/template/server/drain_test.go @@ -0,0 +1,80 @@ +//go:build linux + +package server + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + + templatemanager "github.com/e2b-dev/infra/packages/shared/pkg/grpc/template-manager" + "github.com/e2b-dev/infra/packages/shared/pkg/logger" +) + +func TestWaitBuildStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { + t.Parallel() + + s := &ServerStore{ + logger: logger.NewNopLogger(), + drainDone: make(chan struct{}), + } + + s.buildStartMu.RLock() + defer s.buildStartMu.RUnlock() + + waitCtx, cancel := context.WithCancel(t.Context()) + waitErr := make(chan error, 1) + go func() { + waitErr <- s.waitBuildStarts(waitCtx) + }() + + time.Sleep(2 * buildStartWaitPollInterval) + cancel() + + select { + case err := <-waitErr: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("waitBuildStarts did not return after cancellation") + } + + s.StartDraining(t.Context()) + + enterErr := make(chan error, 1) + go func() { + enterErr <- s.enterBuildStart(t.Context(), "test") + }() + + select { + case err := <-enterErr: + if err == nil { + s.leaveBuildStart() + } + require.Equal(t, codes.Unavailable, status.Code(err)) + case <-time.After(time.Second): + t.Fatal("enterBuildStart blocked instead of rejecting while draining") + } +} + +func TestTemplateBuildDeleteRejectsAfterDrainStarts(t *testing.T) { + t.Parallel() + + s := &ServerStore{ + logger: logger.NewNopLogger(), + drainDone: make(chan struct{}), + wg: &sync.WaitGroup{}, + } + s.StartDraining(t.Context()) + + got, err := s.TemplateBuildDelete(t.Context(), &templatemanager.TemplateBuildDeleteRequest{ + TemplateID: "template-id", + BuildID: "build-id", + }) + require.Equal(t, codes.Unavailable, status.Code(err)) + require.Nil(t, got) +} diff --git a/packages/orchestrator/pkg/template/server/main.go b/packages/orchestrator/pkg/template/server/main.go index 2013af14fa..8a3820e770 100644 --- a/packages/orchestrator/pkg/template/server/main.go +++ b/packages/orchestrator/pkg/template/server/main.go @@ -47,6 +47,9 @@ type ServerStore struct { wg *sync.WaitGroup // wait group for running builds activeBuilds atomic.Int64 // counter for active builds (for debugging) + drainOnce sync.Once + drainDone chan struct{} + buildStartMu sync.RWMutex closers []closeable } @@ -123,6 +126,7 @@ func New( templateStorage: templatePersistence, buildStorage: buildPersistence, wg: &sync.WaitGroup{}, + drainDone: make(chan struct{}), closers: closers, } @@ -130,6 +134,8 @@ func New( } func (s *ServerStore) Close(ctx context.Context) error { + s.StartDraining(ctx) + select { case <-ctx.Done(): return errors.New("force exit, not waiting for builds to finish") @@ -150,6 +156,11 @@ func (s *ServerStore) Close(ctx context.Context) error { } func (s *ServerStore) Wait(ctx context.Context) error { + s.StartDraining(ctx) + if err := s.waitBuildStarts(ctx); err != nil { + return err + } + select { case <-ctx.Done(): return errors.New("force exit, not waiting for builds to finish") diff --git a/packages/orchestrator/pkg/template/server/upload_layer_files_template.go b/packages/orchestrator/pkg/template/server/upload_layer_files_template.go index 2776500077..a483175561 100644 --- a/packages/orchestrator/pkg/template/server/upload_layer_files_template.go +++ b/packages/orchestrator/pkg/template/server/upload_layer_files_template.go @@ -17,6 +17,9 @@ const signedUrlExpiration = time.Minute * 30 func (s *ServerStore) InitLayerFileUpload(ctx context.Context, in *templatemanager.InitLayerFileUploadRequest) (*templatemanager.InitLayerFileUploadResponse, error) { ctx, childSpan := tracer.Start(ctx, "template-create") defer childSpan.End() + if err := s.rejectIfDraining(ctx, "template-layer-file-upload"); err != nil { + return nil, err + } // default to scope by template ID cacheScope := in.GetTemplateID() From bb8eddf567de3827a9a933334fa4f0f76e6b3280 Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Wed, 3 Jun 2026 13:00:28 -0700 Subject: [PATCH 4/6] refactor(orchestrator): unify sandbox lifecycle drain --- .../orchestrator/pkg/sandbox/factory_test.go | 74 ++++++++++++ packages/orchestrator/pkg/sandbox/map.go | 108 ++++++------------ packages/orchestrator/pkg/sandbox/map_test.go | 56 ++++++--- packages/orchestrator/pkg/sandbox/sandbox.go | 107 +++++++++++++++++ packages/orchestrator/pkg/server/main.go | 97 ++++++++-------- packages/orchestrator/pkg/server/sandboxes.go | 4 +- packages/orchestrator/pkg/server/utils.go | 20 ++++ .../orchestrator/pkg/server/utils_test.go | 68 +++++++++++ 8 files changed, 403 insertions(+), 131 deletions(-) create mode 100644 packages/orchestrator/pkg/sandbox/factory_test.go diff --git a/packages/orchestrator/pkg/sandbox/factory_test.go b/packages/orchestrator/pkg/sandbox/factory_test.go new file mode 100644 index 0000000000..e99b77b6e3 --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/factory_test.go @@ -0,0 +1,74 @@ +//go:build linux + +package sandbox + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestFactoryStartDrainingRejectsNewStarts(t *testing.T) { + t.Parallel() + + factory := testFactory() + factory.StartDraining(t.Context()) + + require.ErrorIs(t, factory.enterSandboxStart(), ErrFactoryDraining) +} + +func TestFactoryWaitSandboxStartsWaitsUntilStartLeaves(t *testing.T) { + t.Parallel() + + factory := testFactory() + require.NoError(t, factory.enterSandboxStart()) + + waitCtx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + + done := make(chan error, 1) + go func() { + done <- factory.WaitSandboxStarts(waitCtx) + }() + + select { + case err := <-done: + require.Failf(t, "WaitSandboxStarts returned before start left", "err: %v", err) + case <-time.After(25 * time.Millisecond): + } + + factory.leaveSandboxStart() + require.NoError(t, <-done) +} + +func TestFactoryWaitSandboxStartsReturnsContextError(t *testing.T) { + t.Parallel() + + factory := testFactory() + require.NoError(t, factory.enterSandboxStart()) + defer factory.leaveSandboxStart() + + waitCtx, cancel := context.WithCancel(t.Context()) + cancel() + + require.ErrorIs(t, factory.WaitSandboxStarts(waitCtx), context.Canceled) +} + +func TestFactoryTryWaitSandboxStartsDoesNotBlock(t *testing.T) { + t.Parallel() + + factory := testFactory() + require.NoError(t, factory.enterSandboxStart()) + defer factory.leaveSandboxStart() + + require.False(t, factory.TryWaitSandboxStarts(t.Context())) +} + +func testFactory() *Factory { + return &Factory{ + Sandboxes: NewSandboxesMap(), + drainDone: make(chan struct{}), + } +} diff --git a/packages/orchestrator/pkg/sandbox/map.go b/packages/orchestrator/pkg/sandbox/map.go index da38efe849..6142189424 100644 --- a/packages/orchestrator/pkg/sandbox/map.go +++ b/packages/orchestrator/pkg/sandbox/map.go @@ -9,8 +9,6 @@ import ( "net" "sync" - "go.uber.org/zap" - "github.com/e2b-dev/infra/packages/shared/pkg/logger" "github.com/e2b-dev/infra/packages/shared/pkg/smap" ) @@ -27,40 +25,6 @@ type MapSubscriber interface { OnNetworkRelease(ctx context.Context, sbx *Sandbox) } -type SandboxState string - -const ( - SandboxStateRunning SandboxState = "running" - SandboxStateStopping SandboxState = "stopping" -) - -type lifecycleEntry struct { - sandbox *Sandbox - stateLock sync.RWMutex - state SandboxState -} - -func newLifecycleEntry(sbx *Sandbox, state SandboxState) *lifecycleEntry { - return &lifecycleEntry{ - sandbox: sbx, - state: state, - } -} - -func (e *lifecycleEntry) setState(state SandboxState) { - e.stateLock.Lock() - defer e.stateLock.Unlock() - - e.state = state -} - -func (e *lifecycleEntry) getState() SandboxState { - e.stateLock.RLock() - defer e.stateLock.RUnlock() - - return e.state -} - // Map holds sandboxes that are live (running), known active lifecycles, // together with a IP-to-sandbox index. The indexes are managed independently. // @@ -68,18 +32,22 @@ func (e *lifecycleEntry) getState() SandboxState { // MarkRunning/MarkStopping manage the live set. type Map struct { live *smap.Map[*Sandbox] - lifecycles *smap.Map[*lifecycleEntry] + lifecycles *smap.Map[*Sandbox] network *smap.Map[*Sandbox] + lifecycleMu sync.Mutex + lifecycleChanged chan struct{} + subs []MapSubscriber subsLock sync.RWMutex } func NewSandboxesMap() *Map { return &Map{ - live: smap.New[*Sandbox](), - lifecycles: smap.New[*lifecycleEntry](), - network: smap.New[*Sandbox](), + live: smap.New[*Sandbox](), + lifecycles: smap.New[*Sandbox](), + network: smap.New[*Sandbox](), + lifecycleChanged: make(chan struct{}), } } @@ -116,32 +84,33 @@ func (m *Map) Get(sandboxID string) (*Sandbox, bool) { } func (m *Map) LifecycleItems() []*Sandbox { - entries := m.lifecycles.Items() - sandboxes := make([]*Sandbox, 0, len(entries)) - for _, entry := range entries { - sandboxes = append(sandboxes, entry.sandbox) + items := m.lifecycles.Items() + sandboxes := make([]*Sandbox, 0, len(items)) + for _, sbx := range items { + sandboxes = append(sandboxes, sbx) } return sandboxes } -func (m *Map) LifecycleItemsByState(states ...SandboxState) []*Sandbox { - stateSet := make(map[SandboxState]struct{}, len(states)) - for _, state := range states { - stateSet[state] = struct{}{} - } +func (m *Map) WaitLifecycles(ctx context.Context) error { + for { + m.lifecycleMu.Lock() + if m.lifecycles.Count() == 0 { + m.lifecycleMu.Unlock() - entries := m.lifecycles.Items() - sandboxes := make([]*Sandbox, 0, len(entries)) - for _, entry := range entries { - if _, ok := stateSet[entry.getState()]; !ok { - continue + return nil } - sandboxes = append(sandboxes, entry.sandbox) - } + changed := m.lifecycleChanged + m.lifecycleMu.Unlock() - return sandboxes + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for sandbox lifecycle cleanup: %w", ctx.Err()) + case <-changed: + } + } } // GetByHostPort looks up a sandbox by its host IP address parsed from hostPort. @@ -171,14 +140,16 @@ func (m *Map) AssignNetwork(ctx context.Context, sbx *Sandbox) { ) } -func (m *Map) TrackLifecycle(ctx context.Context, sbx *Sandbox, state SandboxState) { - m.lifecycles.Insert(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID), newLifecycleEntry(sbx, state)) +func (m *Map) trackLifecycle(ctx context.Context, sbx *Sandbox) { + m.lifecycleMu.Lock() + m.lifecycles.Insert(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID), sbx) + m.notifyLifecycleChangeLocked() + m.lifecycleMu.Unlock() logger.L().Info(ctx, "sandbox lifecycle tracked", logger.WithSandboxID(sbx.Runtime.SandboxID), logger.WithLifecycleID(sbx.LifecycleID), logger.WithSandboxIP(sbx.Slot.HostIPString()), - zap.String("state", string(state)), ) } @@ -188,7 +159,7 @@ func (m *Map) MarkRunning(ctx context.Context, sbx *Sandbox) { return } - m.TrackLifecycle(ctx, sbx, SandboxStateRunning) + m.trackLifecycle(ctx, sbx) m.trigger(ctx, func(ctx context.Context, s MapSubscriber) { s.OnInsert(ctx, sbx) @@ -210,7 +181,6 @@ func (m *Map) MarkRunning(ctx context.Context, sbx *Sandbox) { // Returns true if the sandbox was successfully removed. func (m *Map) MarkStopping(ctx context.Context, sandboxID, lifecycleID string) bool { stopped := false - m.markLifecycleState(sandboxID, lifecycleID, SandboxStateStopping) m.live.RemoveCb(sandboxID, func(_ string, sbx *Sandbox, exists bool) bool { if !exists { @@ -236,7 +206,10 @@ func (m *Map) MarkStopping(ctx context.Context, sandboxID, lifecycleID string) b } func (m *Map) MarkStopped(ctx context.Context, sbx *Sandbox) { + m.lifecycleMu.Lock() m.lifecycles.Remove(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID)) + m.notifyLifecycleChangeLocked() + m.lifecycleMu.Unlock() logger.L().Info(ctx, "sandbox lifecycle stopped", logger.WithSandboxID(sbx.Runtime.SandboxID), @@ -245,14 +218,9 @@ func (m *Map) MarkStopped(ctx context.Context, sbx *Sandbox) { ) } -func (m *Map) markLifecycleState(sandboxID, lifecycleID string, state SandboxState) { - key := sandboxLifecycleKey(sandboxID, lifecycleID) - entry, ok := m.lifecycles.Get(key) - if !ok { - return - } - - entry.setState(state) +func (m *Map) notifyLifecycleChangeLocked() { + close(m.lifecycleChanged) + m.lifecycleChanged = make(chan struct{}) } // NetworkReleased unregisters a sandbox's IP and notifies OnNetworkRelease diff --git a/packages/orchestrator/pkg/sandbox/map_test.go b/packages/orchestrator/pkg/sandbox/map_test.go index 54358363e1..b8b2f18a42 100644 --- a/packages/orchestrator/pkg/sandbox/map_test.go +++ b/packages/orchestrator/pkg/sandbox/map_test.go @@ -3,8 +3,10 @@ package sandbox import ( + "context" "sync" "testing" + "time" "github.com/stretchr/testify/require" @@ -19,7 +21,7 @@ func TestMapMarkRunningTracksLifecycle(t *testing.T) { sandboxes.MarkRunning(t.Context(), sbx) require.Len(t, sandboxes.Items(), 1) - require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) + require.Len(t, sandboxes.LifecycleItems(), 1) } func TestMapLifecycleItemsRemainAfterMarkStopping(t *testing.T) { @@ -30,13 +32,12 @@ func TestMapLifecycleItemsRemainAfterMarkStopping(t *testing.T) { sandboxes.MarkRunning(t.Context(), sbx) require.Len(t, sandboxes.Items(), 1) - require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) + require.Len(t, sandboxes.LifecycleItems(), 1) marked := sandboxes.MarkStopping(t.Context(), sbx.Runtime.SandboxID, sbx.LifecycleID) require.True(t, marked) require.Empty(t, sandboxes.Items()) require.Len(t, sandboxes.LifecycleItems(), 1) - require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateStopping), 1) sandboxes.MarkStopped(t.Context(), sbx) require.Empty(t, sandboxes.LifecycleItems()) @@ -64,28 +65,57 @@ func TestMapLifecycleItemsAllowDuplicateSandboxIDs(t *testing.T) { oldSbx := testMapSandbox(t, "lifecycle-old") newSbx := testMapSandbox(t, "lifecycle-new") - sandboxes.TrackLifecycle(t.Context(), oldSbx, SandboxStateStopping) - sandboxes.TrackLifecycle(t.Context(), newSbx, SandboxStateRunning) + sandboxes.MarkRunning(t.Context(), oldSbx) + require.True(t, sandboxes.MarkStopping(t.Context(), oldSbx.Runtime.SandboxID, oldSbx.LifecycleID)) + sandboxes.MarkRunning(t.Context(), newSbx) require.Len(t, sandboxes.LifecycleItems(), 2) - require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateStopping), 1) - require.Len(t, sandboxes.LifecycleItemsByState(SandboxStateRunning), 1) } -func TestMapLifecycleStateUpdateAfterRemovalDoesNotResurrect(t *testing.T) { +func TestMapWaitLifecyclesReturnsWhenEmpty(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + + require.NoError(t, sandboxes.WaitLifecycles(t.Context())) +} + +func TestMapWaitLifecyclesWaitsUntilStopped(t *testing.T) { t.Parallel() sandboxes := NewSandboxesMap() sbx := testMapSandbox(t, "lifecycle-1") + sandboxes.MarkRunning(t.Context(), sbx) - sandboxes.TrackLifecycle(t.Context(), sbx, SandboxStateRunning) - entry, ok := sandboxes.lifecycles.Get(sandboxLifecycleKey(sbx.Runtime.SandboxID, sbx.LifecycleID)) - require.True(t, ok) + waitCtx, cancel := context.WithTimeout(t.Context(), time.Second) + defer cancel() + + done := make(chan error, 1) + go func() { + done <- sandboxes.WaitLifecycles(waitCtx) + }() + + select { + case err := <-done: + require.Failf(t, "WaitLifecycles returned before lifecycle stopped", "err: %v", err) + case <-time.After(25 * time.Millisecond): + } sandboxes.MarkStopped(t.Context(), sbx) - entry.setState(SandboxStateStopping) + require.NoError(t, <-done) +} - require.Empty(t, sandboxes.LifecycleItems()) +func TestMapWaitLifecyclesReturnsContextError(t *testing.T) { + t.Parallel() + + sandboxes := NewSandboxesMap() + sbx := testMapSandbox(t, "lifecycle-1") + sandboxes.MarkRunning(t.Context(), sbx) + + waitCtx, cancel := context.WithCancel(t.Context()) + cancel() + + require.ErrorIs(t, sandboxes.WaitLifecycles(waitCtx), context.Canceled) } func TestMapConcurrentMarkStoppingAndStoppedDoesNotResurrectLifecycle(t *testing.T) { diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index 45fdc2df28..c1ceb70ced 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -50,6 +50,9 @@ var ( waitForEnvdDurationHistogram = utils.Must(telemetry.GetHistogram(meter, telemetry.WaitForEnvdDurationHistogramName)) ) +// ErrFactoryDraining is returned when a sandbox start is attempted after the factory entered drain mode. +var ErrFactoryDraining = errors.New("sandbox factory is draining") + var SandboxHttpTransport = otelhttp.NewTransport( &http.Transport{ DisableKeepAlives: true, @@ -57,6 +60,8 @@ var SandboxHttpTransport = otelhttp.NewTransport( }, ) +const factoryStartWaitPollInterval = 50 * time.Millisecond + // Http client that should be used for requests to sandboxes. var sandboxHttpClient = http.Client{ Timeout: 10 * time.Second, @@ -293,6 +298,10 @@ type Factory struct { hostStatsDelivery hoststats.Delivery cgroupManager cgroup.Manager egressProxy network.EgressProxy + + drainOnce sync.Once + drainDone chan struct{} + sandboxStartMu sync.RWMutex } func NewFactory( @@ -314,6 +323,96 @@ func NewFactory( hostStatsDelivery: hostStatsDelivery, cgroupManager: cgroupManager, egressProxy: egressProxy, + drainDone: make(chan struct{}), + } +} + +func (f *Factory) StartDraining(ctx context.Context) { + if f == nil || f.drainDone == nil { + return + } + + f.drainOnce.Do(func() { + trackedSandboxes := 0 + if f.Sandboxes != nil { + trackedSandboxes = len(f.Sandboxes.LifecycleItems()) + } + + logger.L().Info(ctx, "sandbox factory entering drain mode", zap.Int("tracked_sandboxes", trackedSandboxes)) + close(f.drainDone) + }) +} + +func (f *Factory) enterSandboxStart() error { + if err := f.rejectIfDraining(); err != nil { + return err + } + + f.sandboxStartMu.RLock() + if err := f.rejectIfDraining(); err != nil { + f.sandboxStartMu.RUnlock() + + return err + } + + return nil +} + +func (f *Factory) leaveSandboxStart() { + f.sandboxStartMu.RUnlock() +} + +func (f *Factory) WaitSandboxStarts(ctx context.Context) error { + logger.L().Info(ctx, "waiting for in-flight sandbox factory start operations to finish") + + ticker := time.NewTicker(factoryStartWaitPollInterval) + defer ticker.Stop() + + for { + if f.sandboxStartMu.TryLock() { + logger.L().Info(ctx, "in-flight sandbox factory start gate acquired") + f.sandboxStartMu.Unlock() + logger.L().Info(ctx, "in-flight sandbox factory start operations finished") + + return nil + } + + select { + case <-ctx.Done(): + return fmt.Errorf("waiting for in-flight sandbox factory start operations: %w", ctx.Err()) + case <-ticker.C: + } + } +} + +// TryWaitSandboxStarts returns whether no sandbox factory starts are in flight. +func (f *Factory) TryWaitSandboxStarts(ctx context.Context) bool { + if f == nil { + return true + } + + if !f.sandboxStartMu.TryLock() { + logger.L().Warn(ctx, "in-flight sandbox factory start operations still running") + + return false + } + + f.sandboxStartMu.Unlock() + logger.L().Info(ctx, "in-flight sandbox factory start operations finished") + + return true +} + +func (f *Factory) rejectIfDraining() error { + if f == nil || f.drainDone == nil { + return nil + } + + select { + case <-f.drainDone: + return ErrFactoryDraining + default: + return nil } } @@ -339,6 +438,10 @@ func (f *Factory) CreateSandbox( ctx, span := tracer.Start(ctx, "create sandbox") defer span.End() defer handleSpanError(span, &e) + if err := f.enterSandboxStart(); err != nil { + return nil, err + } + defer f.leaveSandboxStart() execCtx, execSpan := startExecutionSpan(ctx) @@ -591,6 +694,10 @@ func (f *Factory) ResumeSandbox( ctx, span := tracer.Start(ctx, "resume sandbox") defer span.End() defer handleSpanError(span, &e) + if err := f.enterSandboxStart(); err != nil { + return nil, err + } + defer f.leaveSandboxStart() execCtx, execSpan := startExecutionSpan(ctx) diff --git a/packages/orchestrator/pkg/server/main.go b/packages/orchestrator/pkg/server/main.go index 1f66835e80..29f4c1c7e4 100644 --- a/packages/orchestrator/pkg/server/main.go +++ b/packages/orchestrator/pkg/server/main.go @@ -76,8 +76,7 @@ type Server struct { done chan struct{} closeOnce sync.Once - sandboxStartMu sync.RWMutex - sandboxLifecycleWG sync.WaitGroup + sandboxStartMu sync.RWMutex } type ServiceConfig struct { @@ -191,6 +190,9 @@ func (s *Server) StartDraining(ctx context.Context) { func (s *Server) DrainSandboxes(ctx context.Context) error { s.StartDraining(ctx) + // The sandbox factory is shared by API sandboxes and template-build sandboxes. + // Drain it before waiting so the lifecycle snapshot cannot miss a new build sandbox. + s.sandboxFactory.StartDraining(ctx) if err := s.waitSandboxStarts(ctx); err != nil { return err } @@ -241,48 +243,62 @@ func (s *Server) DrainSandboxes(ctx context.Context) error { func (s *Server) ForceStopSandboxes(ctx context.Context) error { s.StartDraining(ctx) - if err := s.waitSandboxStarts(ctx); err != nil { - return err - } + // The sandbox factory is shared by API sandboxes and template-build sandboxes. + // Drain it before waiting so no new starts can enter while shutdown proceeds. + s.sandboxFactory.StartDraining(ctx) + stopped := make(map[string]struct{}) + var errs []error - sandboxes := s.sandboxFactory.Sandboxes.LifecycleItems() - logger.L().Warn(ctx, "starting forced sandbox shutdown", zap.Int("sandbox_count", len(sandboxes))) - if len(sandboxes) == 0 { - return s.waitSandboxLifecycles(ctx) - } + forceStop := func(sandboxes []*sandbox.Sandbox) { + var wg sync.WaitGroup + errCh := make(chan error, len(sandboxes)) - var wg sync.WaitGroup - errCh := make(chan error, len(sandboxes)) + for _, sbx := range sandboxes { + key := fmt.Sprintf("%s/%s", sbx.Runtime.SandboxID, sbx.LifecycleID) + if _, ok := stopped[key]; ok { + continue + } + stopped[key] = struct{}{} - for _, sbx := range sandboxes { - wg.Go(func() { - sbxLog := logger.L().With( - logger.WithSandboxID(sbx.Runtime.SandboxID), - logger.WithLifecycleID(sbx.LifecycleID), - logger.WithSandboxIP(sbx.Slot.HostIPString()), - ) - sbxLog.Warn(ctx, "force stopping sandbox during orchestrator shutdown") + wg.Go(func() { + sbxLog := logger.L().With( + logger.WithSandboxID(sbx.Runtime.SandboxID), + logger.WithLifecycleID(sbx.LifecycleID), + logger.WithSandboxIP(sbx.Slot.HostIPString()), + ) + sbxLog.Warn(ctx, "force stopping sandbox during orchestrator shutdown") - marked := s.sandboxFactory.Sandboxes.MarkStopping(ctx, sbx.Runtime.SandboxID, sbx.LifecycleID) - if !marked { - sbxLog.Info(ctx, "sandbox was already removed from live map before force stop") - } + marked := s.sandboxFactory.Sandboxes.MarkStopping(ctx, sbx.Runtime.SandboxID, sbx.LifecycleID) + if !marked { + sbxLog.Info(ctx, "sandbox was already removed from live map before force stop") + } - if err := sbx.Stop(ctx); err != nil { - errCh <- fmt.Errorf("stop sandbox %s/%s: %w", sbx.Runtime.SandboxID, sbx.LifecycleID, err) - sbxLog.Error(ctx, "failed to force stop sandbox", zap.Error(err)) - } + if err := sbx.Stop(ctx); err != nil { + errCh <- fmt.Errorf("stop sandbox %s/%s: %w", sbx.Runtime.SandboxID, sbx.LifecycleID, err) + sbxLog.Error(ctx, "failed to force stop sandbox", zap.Error(err)) + } + + sbxLog.Info(ctx, "forced sandbox stop requested") + }) + } + + wg.Wait() + close(errCh) - sbxLog.Info(ctx, "forced sandbox stop requested") - }) + for err := range errCh { + errs = append(errs, err) + } } - wg.Wait() - close(errCh) + sandboxes := s.sandboxFactory.Sandboxes.LifecycleItems() + logger.L().Warn(ctx, "starting forced sandbox shutdown", zap.Int("sandbox_count", len(sandboxes))) + forceStop(sandboxes) - var errs []error - for err := range errCh { + if err := s.waitSandboxStarts(ctx); err != nil { errs = append(errs, err) + } else { + newSandboxes := s.sandboxFactory.Sandboxes.LifecycleItems() + forceStop(newSandboxes) } if err := s.waitSandboxLifecycles(ctx); err != nil { @@ -301,18 +317,7 @@ func (s *Server) ForceStopSandboxes(ctx context.Context) error { } func (s *Server) waitSandboxLifecycles(ctx context.Context) error { - done := make(chan struct{}) - go func() { - s.sandboxLifecycleWG.Wait() - close(done) - }() - - select { - case <-ctx.Done(): - return fmt.Errorf("waiting for sandbox lifecycle cleanup: %w", ctx.Err()) - case <-done: - return nil - } + return s.sandboxFactory.Sandboxes.WaitLifecycles(ctx) } func (s *Server) refreshStartingSandboxesLimit(ctx context.Context) { diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index 40fce97c9f..285eac49fa 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -920,7 +920,7 @@ func (s *Server) uploadSnapshotAsync(ctx context.Context, sbx *sandbox.Sandbox, // setupSandboxLifecycle sets up the cleanup goroutine for a sandbox. func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox) { - s.sandboxLifecycleWG.Go(func() { + go func() { ctx, childSpan := tracer.Start(context.WithoutCancel(ctx), "stop sandbox-lifecycle", trace.WithNewRoot()) defer childSpan.End() @@ -940,7 +940,7 @@ func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox } sbxlogger.E(sbx).Info(ctx, "Sandbox stopped") - }) + }() } // stopSandboxAsync stops the sandbox in a background goroutine. diff --git a/packages/orchestrator/pkg/server/utils.go b/packages/orchestrator/pkg/server/utils.go index f6d6e91f6d..a10668b6b6 100644 --- a/packages/orchestrator/pkg/server/utils.go +++ b/packages/orchestrator/pkg/server/utils.go @@ -58,6 +58,9 @@ func (s *Server) waitSandboxStarts(ctx context.Context) error { logger.L().Info(ctx, "in-flight sandbox start gate acquired") s.sandboxStartMu.Unlock() logger.L().Info(ctx, "in-flight sandbox start operations finished") + if s.sandboxFactory != nil { + return s.sandboxFactory.WaitSandboxStarts(ctx) + } return nil } @@ -70,6 +73,23 @@ func (s *Server) waitSandboxStarts(ctx context.Context) error { } } +func (s *Server) tryWaitSandboxStarts(ctx context.Context) bool { + if !s.sandboxStartMu.TryLock() { + logger.L().Warn(ctx, "in-flight sandbox start operations still running") + + return false + } + + s.sandboxStartMu.Unlock() + logger.L().Info(ctx, "in-flight sandbox start operations finished") + + if s.sandboxFactory != nil { + return s.sandboxFactory.TryWaitSandboxStarts(ctx) + } + + return true +} + func (s *Server) waitForAcquire(ctx context.Context) error { if err := s.rejectIfDraining(ctx, "wait-for-acquire"); err != nil { return err diff --git a/packages/orchestrator/pkg/server/utils_test.go b/packages/orchestrator/pkg/server/utils_test.go index 04bb6dabc1..46984802fe 100644 --- a/packages/orchestrator/pkg/server/utils_test.go +++ b/packages/orchestrator/pkg/server/utils_test.go @@ -10,6 +10,8 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" + + "github.com/e2b-dev/infra/packages/orchestrator/pkg/sandbox" ) func TestWaitSandboxStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { @@ -55,3 +57,69 @@ func TestWaitSandboxStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { t.Fatal("enterSandboxStart blocked instead of rejecting while draining") } } + +func TestTryWaitSandboxStartsDoesNotBlock(t *testing.T) { + t.Parallel() + + s := &Server{done: make(chan struct{})} + s.sandboxStartMu.RLock() + defer s.sandboxStartMu.RUnlock() + + require.False(t, s.tryWaitSandboxStarts(t.Context())) +} + +func TestForceStopSandboxesWaitsForInFlightStarts(t *testing.T) { + t.Parallel() + + s := forceStopTestServer() + s.sandboxStartMu.RLock() + locked := true + defer func() { + if locked { + s.sandboxStartMu.RUnlock() + } + }() + + done := make(chan error, 1) + go func() { + done <- s.ForceStopSandboxes(t.Context()) + }() + + select { + case err := <-done: + require.Failf(t, "ForceStopSandboxes returned before start left", "err: %v", err) + case <-time.After(2 * sandboxStartWaitPollInterval): + } + + s.sandboxStartMu.RUnlock() + locked = false + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("ForceStopSandboxes did not return after start left") + } +} + +func TestForceStopSandboxesReturnsInFlightStartContextError(t *testing.T) { + t.Parallel() + + s := forceStopTestServer() + s.sandboxStartMu.RLock() + defer s.sandboxStartMu.RUnlock() + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + require.ErrorIs(t, s.ForceStopSandboxes(ctx), context.Canceled) +} + +func forceStopTestServer() *Server { + return &Server{ + done: make(chan struct{}), + sandboxFactory: &sandbox.Factory{ + Sandboxes: sandbox.NewSandboxesMap(), + }, + } +} From 114563f88e155d4f704a27b1625ab683349c8701 Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Wed, 3 Jun 2026 13:22:02 -0700 Subject: [PATCH 5/6] refactor(orchestrator): return drain gate release callbacks --- .../orchestrator/pkg/sandbox/factory_test.go | 23 ++++------- packages/orchestrator/pkg/sandbox/sandbox.go | 39 ++++++------------- packages/orchestrator/pkg/server/sandboxes.go | 10 +++-- packages/orchestrator/pkg/server/utils.go | 29 ++++---------- .../orchestrator/pkg/server/utils_test.go | 20 +++------- .../pkg/template/server/create_template.go | 12 +----- .../pkg/template/server/delete_template.go | 12 +----- .../orchestrator/pkg/template/server/drain.go | 12 +++--- .../pkg/template/server/drain_test.go | 10 +++-- 9 files changed, 56 insertions(+), 111 deletions(-) diff --git a/packages/orchestrator/pkg/sandbox/factory_test.go b/packages/orchestrator/pkg/sandbox/factory_test.go index e99b77b6e3..f987112ee6 100644 --- a/packages/orchestrator/pkg/sandbox/factory_test.go +++ b/packages/orchestrator/pkg/sandbox/factory_test.go @@ -16,14 +16,16 @@ func TestFactoryStartDrainingRejectsNewStarts(t *testing.T) { factory := testFactory() factory.StartDraining(t.Context()) - require.ErrorIs(t, factory.enterSandboxStart(), ErrFactoryDraining) + _, err := factory.enterSandboxStart() + require.ErrorIs(t, err, ErrFactoryDraining) } func TestFactoryWaitSandboxStartsWaitsUntilStartLeaves(t *testing.T) { t.Parallel() factory := testFactory() - require.NoError(t, factory.enterSandboxStart()) + release, err := factory.enterSandboxStart() + require.NoError(t, err) waitCtx, cancel := context.WithTimeout(t.Context(), time.Second) defer cancel() @@ -39,7 +41,7 @@ func TestFactoryWaitSandboxStartsWaitsUntilStartLeaves(t *testing.T) { case <-time.After(25 * time.Millisecond): } - factory.leaveSandboxStart() + release() require.NoError(t, <-done) } @@ -47,8 +49,9 @@ func TestFactoryWaitSandboxStartsReturnsContextError(t *testing.T) { t.Parallel() factory := testFactory() - require.NoError(t, factory.enterSandboxStart()) - defer factory.leaveSandboxStart() + release, err := factory.enterSandboxStart() + require.NoError(t, err) + defer release() waitCtx, cancel := context.WithCancel(t.Context()) cancel() @@ -56,16 +59,6 @@ func TestFactoryWaitSandboxStartsReturnsContextError(t *testing.T) { require.ErrorIs(t, factory.WaitSandboxStarts(waitCtx), context.Canceled) } -func TestFactoryTryWaitSandboxStartsDoesNotBlock(t *testing.T) { - t.Parallel() - - factory := testFactory() - require.NoError(t, factory.enterSandboxStart()) - defer factory.leaveSandboxStart() - - require.False(t, factory.TryWaitSandboxStarts(t.Context())) -} - func testFactory() *Factory { return &Factory{ Sandboxes: NewSandboxesMap(), diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index c1ceb70ced..d72b2eb9f9 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -343,19 +343,20 @@ func (f *Factory) StartDraining(ctx context.Context) { }) } -func (f *Factory) enterSandboxStart() error { +func (f *Factory) enterSandboxStart() (func(), error) { if err := f.rejectIfDraining(); err != nil { - return err + return nil, err } f.sandboxStartMu.RLock() + release := sync.OnceFunc(f.leaveSandboxStart) if err := f.rejectIfDraining(); err != nil { - f.sandboxStartMu.RUnlock() + release() - return err + return nil, err } - return nil + return release, nil } func (f *Factory) leaveSandboxStart() { @@ -385,24 +386,6 @@ func (f *Factory) WaitSandboxStarts(ctx context.Context) error { } } -// TryWaitSandboxStarts returns whether no sandbox factory starts are in flight. -func (f *Factory) TryWaitSandboxStarts(ctx context.Context) bool { - if f == nil { - return true - } - - if !f.sandboxStartMu.TryLock() { - logger.L().Warn(ctx, "in-flight sandbox factory start operations still running") - - return false - } - - f.sandboxStartMu.Unlock() - logger.L().Info(ctx, "in-flight sandbox factory start operations finished") - - return true -} - func (f *Factory) rejectIfDraining() error { if f == nil || f.drainDone == nil { return nil @@ -438,10 +421,11 @@ func (f *Factory) CreateSandbox( ctx, span := tracer.Start(ctx, "create sandbox") defer span.End() defer handleSpanError(span, &e) - if err := f.enterSandboxStart(); err != nil { + releaseSandboxStart, err := f.enterSandboxStart() + if err != nil { return nil, err } - defer f.leaveSandboxStart() + defer releaseSandboxStart() execCtx, execSpan := startExecutionSpan(ctx) @@ -694,10 +678,11 @@ func (f *Factory) ResumeSandbox( ctx, span := tracer.Start(ctx, "resume sandbox") defer span.End() defer handleSpanError(span, &e) - if err := f.enterSandboxStart(); err != nil { + releaseSandboxStart, err := f.enterSandboxStart() + if err != nil { return nil, err } - defer f.leaveSandboxStart() + defer releaseSandboxStart() execCtx, execSpan := startExecutionSpan(ctx) diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index 285eac49fa..8850f94b85 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -90,10 +90,11 @@ func (s *Server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ telemetry.WithEnvdVersion(req.GetSandbox().GetEnvdVersion()), ) - if err := s.enterSandboxStart(ctx, "sandbox-create"); err != nil { + releaseSandboxStart, err := s.enterSandboxStart(ctx, "sandbox-create") + if err != nil { return nil, err } - defer s.leaveSandboxStart() + defer releaseSandboxStart() // setup launch darkly ctx = featureflags.AddToContext( @@ -630,10 +631,11 @@ func (s *Server) Checkpoint(ctx context.Context, in *orchestrator.SandboxCheckpo Build(), ) - if err := s.enterSandboxStart(ctx, "sandbox-checkpoint"); err != nil { + releaseSandboxStart, err := s.enterSandboxStart(ctx, "sandbox-checkpoint") + if err != nil { return nil, err } - defer s.leaveSandboxStart() + defer releaseSandboxStart() sbx, ok := s.sandboxFactory.Sandboxes.Get(in.GetSandboxId()) if !ok { diff --git a/packages/orchestrator/pkg/server/utils.go b/packages/orchestrator/pkg/server/utils.go index a10668b6b6..b3ffd9b1a6 100644 --- a/packages/orchestrator/pkg/server/utils.go +++ b/packages/orchestrator/pkg/server/utils.go @@ -5,6 +5,7 @@ package server import ( "context" "fmt" + "sync" "time" "go.uber.org/zap" @@ -28,19 +29,20 @@ func (s *Server) rejectIfDraining(ctx context.Context, operation string) error { } } -func (s *Server) enterSandboxStart(ctx context.Context, operation string) error { +func (s *Server) enterSandboxStart(ctx context.Context, operation string) (func(), error) { if err := s.rejectIfDraining(ctx, operation); err != nil { - return err + return nil, err } s.sandboxStartMu.RLock() + release := sync.OnceFunc(s.leaveSandboxStart) if err := s.rejectIfDraining(ctx, operation); err != nil { - s.sandboxStartMu.RUnlock() + release() - return err + return nil, err } - return nil + return release, nil } func (s *Server) leaveSandboxStart() { @@ -73,23 +75,6 @@ func (s *Server) waitSandboxStarts(ctx context.Context) error { } } -func (s *Server) tryWaitSandboxStarts(ctx context.Context) bool { - if !s.sandboxStartMu.TryLock() { - logger.L().Warn(ctx, "in-flight sandbox start operations still running") - - return false - } - - s.sandboxStartMu.Unlock() - logger.L().Info(ctx, "in-flight sandbox start operations finished") - - if s.sandboxFactory != nil { - return s.sandboxFactory.TryWaitSandboxStarts(ctx) - } - - return true -} - func (s *Server) waitForAcquire(ctx context.Context) error { if err := s.rejectIfDraining(ctx, "wait-for-acquire"); err != nil { return err diff --git a/packages/orchestrator/pkg/server/utils_test.go b/packages/orchestrator/pkg/server/utils_test.go index 46984802fe..b650322d85 100644 --- a/packages/orchestrator/pkg/server/utils_test.go +++ b/packages/orchestrator/pkg/server/utils_test.go @@ -44,30 +44,22 @@ func TestWaitSandboxStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { enterErr := make(chan error, 1) go func() { - enterErr <- s.enterSandboxStart(t.Context(), "test") + release, err := s.enterSandboxStart(t.Context(), "test") + if err == nil { + release() + } + + enterErr <- err }() select { case err := <-enterErr: - if err == nil { - s.leaveSandboxStart() - } require.Equal(t, codes.Unavailable, status.Code(err)) case <-time.After(time.Second): t.Fatal("enterSandboxStart blocked instead of rejecting while draining") } } -func TestTryWaitSandboxStartsDoesNotBlock(t *testing.T) { - t.Parallel() - - s := &Server{done: make(chan struct{})} - s.sandboxStartMu.RLock() - defer s.sandboxStartMu.RUnlock() - - require.False(t, s.tryWaitSandboxStarts(t.Context())) -} - func TestForceStopSandboxesWaitsForInFlightStarts(t *testing.T) { t.Parallel() diff --git a/packages/orchestrator/pkg/template/server/create_template.go b/packages/orchestrator/pkg/template/server/create_template.go index a4351ec8b0..dc3adcf7e3 100644 --- a/packages/orchestrator/pkg/template/server/create_template.go +++ b/packages/orchestrator/pkg/template/server/create_template.go @@ -106,18 +106,10 @@ func (s *ServerStore) TemplateCreate(ctx context.Context, templateRequest *templ } logs := buildlogger.NewLogEntryLogger() - if err := s.enterBuildStart(ctx, "template-create"); err != nil { + releaseBuildStart, err := s.enterBuildStart(ctx, "template-create") + if err != nil { return nil, err } - buildStartReleased := false - releaseBuildStart := func() { - if buildStartReleased { - return - } - - s.leaveBuildStart() - buildStartReleased = true - } defer releaseBuildStart() buildInfo, err := s.buildCache.Create(template.TeamID, metadata.BuildID, logs) diff --git a/packages/orchestrator/pkg/template/server/delete_template.go b/packages/orchestrator/pkg/template/server/delete_template.go index 4b7b591e39..7437788053 100644 --- a/packages/orchestrator/pkg/template/server/delete_template.go +++ b/packages/orchestrator/pkg/template/server/delete_template.go @@ -26,18 +26,10 @@ func (s *ServerStore) TemplateBuildDelete(ctx context.Context, in *templatemanag )) defer childSpan.End() - if err := s.enterBuildStart(ctx, "template-delete"); err != nil { + releaseDeleteStart, err := s.enterBuildStart(ctx, "template-delete") + if err != nil { return nil, err } - deleteStartReleased := false - releaseDeleteStart := func() { - if deleteStartReleased { - return - } - - s.leaveBuildStart() - deleteStartReleased = true - } defer releaseDeleteStart() s.wg.Add(1) diff --git a/packages/orchestrator/pkg/template/server/drain.go b/packages/orchestrator/pkg/template/server/drain.go index 78aef43f9f..944506133d 100644 --- a/packages/orchestrator/pkg/template/server/drain.go +++ b/packages/orchestrator/pkg/template/server/drain.go @@ -5,6 +5,7 @@ package server import ( "context" "fmt" + "sync" "time" "go.uber.org/zap" @@ -38,19 +39,20 @@ func (s *ServerStore) rejectIfDraining(ctx context.Context, operation string) er } } -func (s *ServerStore) enterBuildStart(ctx context.Context, operation string) error { +func (s *ServerStore) enterBuildStart(ctx context.Context, operation string) (func(), error) { if err := s.rejectIfDraining(ctx, operation); err != nil { - return err + return nil, err } s.buildStartMu.RLock() + release := sync.OnceFunc(s.leaveBuildStart) if err := s.rejectIfDraining(ctx, operation); err != nil { - s.buildStartMu.RUnlock() + release() - return err + return nil, err } - return nil + return release, nil } func (s *ServerStore) leaveBuildStart() { diff --git a/packages/orchestrator/pkg/template/server/drain_test.go b/packages/orchestrator/pkg/template/server/drain_test.go index 56192a503b..ecfce031cf 100644 --- a/packages/orchestrator/pkg/template/server/drain_test.go +++ b/packages/orchestrator/pkg/template/server/drain_test.go @@ -47,14 +47,16 @@ func TestWaitBuildStartsCanceledDoesNotBlockDrainingRejection(t *testing.T) { enterErr := make(chan error, 1) go func() { - enterErr <- s.enterBuildStart(t.Context(), "test") + release, err := s.enterBuildStart(t.Context(), "test") + if err == nil { + release() + } + + enterErr <- err }() select { case err := <-enterErr: - if err == nil { - s.leaveBuildStart() - } require.Equal(t, codes.Unavailable, status.Code(err)) case <-time.After(time.Second): t.Fatal("enterBuildStart blocked instead of rejecting while draining") From 3c8efb1a71dacd1258aaddd43d73cef4d49b1005 Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Thu, 4 Jun 2026 21:18:30 -0700 Subject: [PATCH 6/6] fix(iac): Increase kill_timeout so orchestrator has time to drain --- iac/modules/job-orchestrator/jobs/orchestrator.hcl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/iac/modules/job-orchestrator/jobs/orchestrator.hcl b/iac/modules/job-orchestrator/jobs/orchestrator.hcl index 703b58e132..2d6470ea1a 100644 --- a/iac/modules/job-orchestrator/jobs/orchestrator.hcl +++ b/iac/modules/job-orchestrator/jobs/orchestrator.hcl @@ -60,6 +60,9 @@ job "orchestrator-${latest_orchestrator_job_id}" { attempts = 0 } + # Matches the Nomad client max_kill_timeout. + kill_timeout = "24h" + resources { memory = 1024 memory_max = -1