diff --git a/packages/orchestrator/pkg/sandbox/sandbox.go b/packages/orchestrator/pkg/sandbox/sandbox.go index 07a78a9f52..3b0e84c1e6 100644 --- a/packages/orchestrator/pkg/sandbox/sandbox.go +++ b/packages/orchestrator/pkg/sandbox/sandbox.go @@ -1720,20 +1720,33 @@ func (s *Sandbox) WaitForExit(ctx context.Context) error { ctx, span := tracer.Start(ctx, "sandbox-wait-for-exit") defer span.End() - timeout := time.Until(s.GetEndAt()) + for { + endAt := s.GetEndAt() + timeout := time.Until(endAt) + if timeout <= 0 { + return errors.New("sandbox TTL expired") + } - select { - case <-time.After(timeout): - return errors.New("waiting for exit took too long") - case <-ctx.Done(): - return nil - case <-s.exit.Done(): - err := s.exit.Error() - if err == nil { + timer := time.NewTimer(timeout) + select { + case <-timer.C: + // Re-check in case SetEndAt extended the deadline (e.g. via KeepAlive). + if s.GetEndAt().After(endAt) { + continue + } + return errors.New("waiting for exit took too long") + case <-ctx.Done(): + timer.Stop() return nil - } + case <-s.exit.Done(): + timer.Stop() + err := s.exit.Error() + if err == nil { + return nil + } - return fmt.Errorf("fc process exited prematurely: %w", err) + return fmt.Errorf("fc process exited prematurely: %w", err) + } } } diff --git a/packages/orchestrator/pkg/sandbox/wait_for_exit_test.go b/packages/orchestrator/pkg/sandbox/wait_for_exit_test.go new file mode 100644 index 0000000000..6ef98a67ed --- /dev/null +++ b/packages/orchestrator/pkg/sandbox/wait_for_exit_test.go @@ -0,0 +1,137 @@ +//go:build linux + +package sandbox + +import ( + "context" + "errors" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/e2b-dev/infra/packages/shared/pkg/utils" +) + +func newTestSandboxForExit(endAt time.Time) *Sandbox { + sbx := &Sandbox{ + Metadata: &Metadata{}, + exit: utils.NewErrorOnce(), + } + sbx.SetEndAt(endAt) + + return sbx +} + +// TestWaitForExit_AlreadyExpired verifies that WaitForExit returns an error +// immediately when endAt is already in the past. +func TestWaitForExit_AlreadyExpired(t *testing.T) { + t.Parallel() + + sbx := newTestSandboxForExit(time.Now().Add(-1 * time.Second)) + + err := sbx.WaitForExit(context.Background()) + + require.Error(t, err) +} + +// TestWaitForExit_ExitBeforeTTL verifies that WaitForExit returns nil when +// the sandbox exits cleanly before its TTL expires. +func TestWaitForExit_ExitBeforeTTL(t *testing.T) { + t.Parallel() + + sbx := newTestSandboxForExit(time.Now().Add(2 * time.Second)) + + go func() { + time.Sleep(50 * time.Millisecond) + sbx.exit.SetSuccess() + }() + + err := sbx.WaitForExit(context.Background()) + + require.NoError(t, err) +} + +// TestWaitForExit_ExitWithError verifies that WaitForExit returns a wrapped +// error when the Firecracker process exits with an error before the TTL. +func TestWaitForExit_ExitWithError(t *testing.T) { + t.Parallel() + + sbx := newTestSandboxForExit(time.Now().Add(2 * time.Second)) + + fcErr := errors.New("firecracker process killed") + go func() { + time.Sleep(50 * time.Millisecond) + sbx.exit.SetError(fcErr) + }() + + err := sbx.WaitForExit(context.Background()) + + require.Error(t, err) + assert.ErrorContains(t, err, "fc process exited prematurely") + assert.ErrorContains(t, err, fcErr.Error()) +} + +// TestWaitForExit_KeepAliveExtendsTTL verifies that a KeepAlive call that +// extends endAt via SetEndAt resets the node-side timer. Without this, the +// sandbox would be killed at the original TTL even though it was extended. +func TestWaitForExit_KeepAliveExtendsTTL(t *testing.T) { + t.Parallel() + + // Short initial TTL — would fire after 80 ms without the fix. + sbx := newTestSandboxForExit(time.Now().Add(80 * time.Millisecond)) + + done := make(chan error, 1) + go func() { + done <- sbx.WaitForExit(context.Background()) + }() + + // Simulate KeepAlive: extend endAt before the original TTL fires. + time.Sleep(40 * time.Millisecond) + sbx.SetEndAt(time.Now().Add(2 * time.Second)) + + // Wait past the original TTL; WaitForExit must NOT have returned yet. + time.Sleep(100 * time.Millisecond) + select { + case err := <-done: + t.Fatalf("WaitForExit returned early after KeepAlive extension: %v", err) + default: + // Good — still waiting. + } + + // Now signal a clean exit; WaitForExit should return nil. + sbx.exit.SetSuccess() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("WaitForExit did not return after exit signal") + } +} + +// TestWaitForExit_ContextCancelled verifies that WaitForExit returns nil when +// the context is cancelled before either TTL expiry or process exit. +func TestWaitForExit_ContextCancelled(t *testing.T) { + t.Parallel() + + sbx := newTestSandboxForExit(time.Now().Add(10 * time.Second)) + + ctx, cancel := context.WithCancel(context.Background()) + + done := make(chan error, 1) + go func() { + done <- sbx.WaitForExit(ctx) + }() + + time.Sleep(20 * time.Millisecond) + cancel() + + select { + case err := <-done: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("WaitForExit did not return after context cancellation") + } +} diff --git a/packages/orchestrator/pkg/server/sandboxes.go b/packages/orchestrator/pkg/server/sandboxes.go index 17f2ac24ee..21d4c9ec02 100644 --- a/packages/orchestrator/pkg/server/sandboxes.go +++ b/packages/orchestrator/pkg/server/sandboxes.go @@ -1041,9 +1041,12 @@ func (s *Server) setupSandboxLifecycle(ctx context.Context, sbx *sandbox.Sandbox ctx, childSpan := tracer.Start(context.WithoutCancel(ctx), "stop sandbox-lifecycle", trace.WithNewRoot()) defer childSpan.End() - waitErr := sbx.Wait(ctx) + waitErr := sbx.WaitForExit(ctx) if waitErr != nil { - sbxlogger.I(sbx).Error(ctx, "failed to wait for sandbox, cleaning up", zap.Error(waitErr)) + sbxlogger.I(sbx).Error(ctx, "sandbox exit wait failed, stopping", zap.Error(waitErr)) + if stopErr := sbx.Stop(ctx); stopErr != nil { + sbxlogger.I(sbx).Error(ctx, "failed to stop timed-out sandbox", zap.Error(stopErr)) + } } cleanupErr := sbx.Close(ctx)