fix sandbox cleanup for template build

Author: dobrac

packages/orchestrator/internal/sandbox/sandbox.go

@@ -71,10 +71,9 @@ type RuntimeMetadata struct {
 }
 
 type Resources struct {
-	Slot     *network.Slot
-	rootfs   rootfs.Provider
-	memory   uffd.MemoryBackend
-	uffdExit chan error
+	Slot   *network.Slot
+	rootfs rootfs.Provider
+	memory uffd.MemoryBackend
 }
 
 type Metadata struct {
@@ -99,6 +98,8 @@ type Sandbox struct {
 	Checks *Checks
 
 	APIStoredConfig *orchestrator.SandboxConfig
+
+	exit *utils.SetOnce[struct{}]
 }
 
 func (s *Sandbox) LoggerMetadata() sbxlogger.SandboxMetadata {
@@ -131,6 +132,8 @@ func CreateSandbox(
 	childCtx, childSpan := tracer.Start(ctx, "create-sandbox")
 	defer childSpan.End()
 
+	exit := utils.NewSetOnce[struct{}]()
+
 	cleanup := NewCleanup()
 	defer func() {
 		if e != nil {
@@ -248,10 +251,9 @@ func CreateSandbox(
 	telemetry.ReportEvent(childCtx, "created fc process")
 
 	resources := &Resources{
-		Slot:     ips.slot,
-		rootfs:   rootfsProvider,
-		memory:   uffd.NewNoopMemory(memfileSize, memfile.BlockSize()),
-		uffdExit: make(chan error, 1),
+		Slot:   ips.slot,
+		rootfs: rootfsProvider,
+		memory: uffd.NewNoopMemory(memfileSize, memfile.BlockSize()),
 	}
 
 	metadata := &Metadata{
@@ -273,6 +275,8 @@ func CreateSandbox(
 		cleanup: cleanup,
 
 		APIStoredConfig: apiConfigToStore,
+
+		exit: exit,
 	}
 
 	checks, err := NewChecks(ctx, tracer, sbx, false)
@@ -282,9 +286,18 @@ func CreateSandbox(
 	sbx.Checks = checks
 
 	cleanup.AddPriority(func(ctx context.Context) error {
-		return sbx.Close(ctx, tracer)
+		// Stop the sandbox first if it is still running, otherwise do nothing
+		return sbx.Stop(ctx, tracer)
 	})
 
+	go func() {
+		// If the process exists, stop the sandbox properly
+		_, fcErr := fcHandle.Exit.Wait()
+		err := s.Stop(context.WithoutCancel(ctx), tracer)
+
+		exit.SetResult(struct{}{}, errors.Join(err, fcErr))
+	}()
+
 	return sbx, nil
 }
 
@@ -307,6 +320,8 @@ func ResumeSandbox(
 	childCtx, childSpan := tracer.Start(ctx, "resume-sandbox")
 	defer childSpan.End()
 
+	exit := utils.NewSetOnce[struct{}]()
+
 	cleanup := NewCleanup()
 	defer func() {
 		if e != nil {
@@ -379,18 +394,16 @@ func ResumeSandbox(
 		return nil, fmt.Errorf("failed to serve memory: %w", err)
 	}
 
+	// ==== END of resources initialization ====
 	uffdStartCtx, cancelUffdStartCtx := context.WithCancelCause(ctx)
 	defer cancelUffdStartCtx(fmt.Errorf("uffd finished starting"))
 
-	uffdExit := make(chan error, 1)
 	go func() {
-		uffdWaitErr := <-fcUffd.Exit()
-		uffdExit <- uffdWaitErr
+		_, uffdWaitErr := fcUffd.Exit().Wait()
 
 		cancelUffdStartCtx(fmt.Errorf("uffd process exited: %w", errors.Join(uffdWaitErr, context.Cause(uffdStartCtx))))
 	}()
 
-	// / ==== END of resources initialization ====
 	rootfsPath, err := rootfsOverlay.Path()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get rootfs path: %w", err)
@@ -452,10 +465,9 @@ func ResumeSandbox(
 	telemetry.ReportEvent(childCtx, "initialized FC")
 
 	resources := &Resources{
-		Slot:     ips.slot,
-		rootfs:   rootfsOverlay,
-		memory:   fcUffd,
-		uffdExit: uffdExit,
+		Slot:   ips.slot,
+		rootfs: rootfsOverlay,
+		memory: fcUffd,
 	}
 
 	metadata := &Metadata{
@@ -477,6 +489,8 @@ func ResumeSandbox(
 		cleanup: cleanup,
 
 		APIStoredConfig: apiConfigToStore,
+
+		exit: exit,
 	}
 
 	// Part of the sandbox as we need to stop Checks before pausing the sandbox
@@ -489,7 +503,8 @@ func ResumeSandbox(
 	sbx.Checks = checks
 
 	cleanup.AddPriority(func(ctx context.Context) error {
-		return sbx.Close(ctx, tracer)
+		// Stop the sandbox first if it is still running, otherwise do nothing
+		return sbx.Stop(ctx, tracer)
 	})
 
 	err = sbx.WaitForEnvd(
@@ -503,40 +518,38 @@ func ResumeSandbox(
 
 	go sbx.Checks.Start()
 
-	return sbx, nil
-}
+	go func() {
+		// Wait for either uffd or fc process to exit
+		select {
+		case <-fcUffd.Exit().Done:
+		case <-fcHandle.Exit.Done:
+		}
 
-func (s *Sandbox) Wait(ctx context.Context) error {
-	select {
-	case <-ctx.Done():
-		return ctx.Err()
-	case <-s.process.Exit.Done:
-		_, fcErr := s.process.Exit.Result()
-		stopErr := s.Stop(ctx)
-		uffdErr := <-s.uffdExit
+		err := s.Stop(context.WithoutCancel(ctx), tracer)
 
-		return errors.Join(fcErr, stopErr, uffdErr)
-	case uffdErr := <-s.uffdExit:
-		stopErr := s.Stop(ctx)
+		_, uffdWaitErr := fcUffd.Exit().Wait()
+		_, fcErr := fcHandle.Exit.Wait()
+		exit.SetResult(struct{}{}, errors.Join(err, fcErr, uffdWaitErr))
+	}()
 
-		_, fcErr := s.process.Exit.WaitWithContext(ctx)
+	return sbx, nil
+}
 
-		return errors.Join(uffdErr, stopErr, fcErr)
-	}
+func (s *Sandbox) Wait(ctx context.Context) error {
+	_, err := s.exit.WaitWithContext(ctx)
+	return err
 }
 
-// Stop starts the cleanup process for the sandbox.
-func (s *Sandbox) Stop(ctx context.Context) error {
+func (s *Sandbox) Close(ctx context.Context) error {
 	err := s.cleanup.Run(ctx)
 	if err != nil {
-		sbxlogger.I(s).Error("failed to stop sandbox", zap.Error(err))
-		return fmt.Errorf("failed to stop sandbox: %w", err)
+		return fmt.Errorf("failed to cleanup sandbox: %w", err)
 	}
-
 	return nil
 }
 
-func (s *Sandbox) Close(ctx context.Context, tracer trace.Tracer) error {
+// Stop kills the sandbox.
+func (s *Sandbox) Stop(ctx context.Context, tracer trace.Tracer) error {
 	_, span := tracer.Start(ctx, "sandbox-close")
 	defer span.End()
 
@@ -659,7 +672,7 @@ func (s *Sandbox) Pause(
 		originalRootfs.Header(),
 		&RootfsDiffCreator{
 			rootfs:   s.rootfs,
-			stopHook: s.Stop,
+			stopHook: s.Close,
 		},
 	)
 	if err != nil {
@@ -894,8 +907,8 @@ func (s *Sandbox) WaitForExit(
 		return fmt.Errorf("waiting for exit took too long")
 	case <-ctx.Done():
 		return nil
-	case <-s.process.Exit.Done:
-		_, err := s.process.Exit.Result()
+	case <-s.exit.Done:
+		_, err := s.exit.Result()
 		if err == nil {
 			return nil
 		}

packages/orchestrator/internal/sandbox/uffd/handler.go

@@ -16,6 +16,7 @@ import (
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/fdexit"
 	"github.com/e2b-dev/infra/packages/orchestrator/internal/sandbox/uffd/mapping"
 	"github.com/e2b-dev/infra/packages/shared/pkg/logger"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 const (
@@ -25,7 +26,7 @@ const (
 )
 
 type Uffd struct {
-	exitCh  chan error
+	exit    *utils.SetOnce[struct{}]
 	readyCh chan struct{}
 
 	fdExit *fdexit.FdExit
@@ -48,7 +49,7 @@ func New(memfile block.ReadonlyDevice, socketPath string, blockSize int64) (*Uff
 	}
 
 	return &Uffd{
-		exitCh:     make(chan error, 1),
+		exit:       utils.NewSetOnce[struct{}](),
 		readyCh:    make(chan struct{}, 1),
 		fdExit:     fdExit,
 		memfile:    trackedMemfile,
@@ -75,10 +76,9 @@ func (u *Uffd) Start(sandboxId string) error {
 		closeErr := u.lis.Close()
 		fdExitErr := u.fdExit.Close()
 
-		u.exitCh <- errors.Join(handleErr, closeErr, fdExitErr)
+		u.exit.SetResult(struct{}{}, errors.Join(handleErr, closeErr, fdExitErr))
 
 		close(u.readyCh)
-		close(u.exitCh)
 	}()
 
 	return nil
@@ -165,8 +165,8 @@ func (u *Uffd) Ready() chan struct{} {
 	return u.readyCh
 }
 
-func (u *Uffd) Exit() chan error {
-	return u.exitCh
+func (u *Uffd) Exit() *utils.SetOnce[struct{}] {
+	return u.exit
 }
 
 func (u *Uffd) TrackAndReturnNil() error {

packages/orchestrator/internal/sandbox/uffd/memory_backend.go

@@ -1,6 +1,10 @@
 package uffd
 
-import "github.com/bits-and-blooms/bitset"
+import (
+	"github.com/bits-and-blooms/bitset"
+
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
+)
 
 type MemoryBackend interface {
 	Disable() error
@@ -9,5 +13,5 @@ type MemoryBackend interface {
 	Start(sandboxId string) error
 	Stop() error
 	Ready() chan struct{}
-	Exit() chan error
+	Exit() *utils.SetOnce[struct{}]
 }

packages/orchestrator/internal/sandbox/uffd/noop.go

@@ -4,6 +4,7 @@ import (
 	"github.com/bits-and-blooms/bitset"
 
 	"github.com/e2b-dev/infra/packages/shared/pkg/storage/header"
+	"github.com/e2b-dev/infra/packages/shared/pkg/utils"
 )
 
 type NoopMemory struct {
@@ -48,6 +49,6 @@ func (m *NoopMemory) Ready() chan struct{} {
 	return ch
 }
 
-func (m *NoopMemory) Exit() chan error {
-	return make(chan error)
+func (m *NoopMemory) Exit() *utils.SetOnce[struct{}] {
+	return utils.NewSetOnce[struct{}]()
 }

packages/orchestrator/internal/server/sandboxes.go

@@ -138,7 +138,7 @@ func (s *server) Create(ctx context.Context, req *orchestrator.SandboxCreateRequ
 			sbxlogger.I(sbx).Error("failed to wait for sandbox, cleaning up", zap.Error(waitErr))
 		}
 
-		cleanupErr := sbx.Stop(ctx)
+		cleanupErr := sbx.Close(ctx)
 		if cleanupErr != nil {
 			sbxlogger.I(sbx).Error("failed to cleanup sandbox, will remove from cache", zap.Error(cleanupErr))
 		}
@@ -304,7 +304,7 @@ func (s *server) Delete(ctxConn context.Context, in *orchestrator.SandboxDeleteR
 	// Start the cleanup in a goroutine—the initial kill request should be send as the first thing in stop, and at this point you cannot route to the sandbox anymore.
 	// We don't wait for the whole cleanup to finish here.
 	go func() {
-		err := sbx.Stop(ctx)
+		err := sbx.Close(ctx)
 		if err != nil {
 			sbxlogger.I(sbx).Error("error stopping sandbox", logger.WithSandboxID(in.SandboxId), zap.Error(err))
 		}
@@ -372,7 +372,7 @@ func (s *server) Pause(ctx context.Context, in *orchestrator.SandboxPauseRequest
 			ctx, childSpan := s.tracer.Start(ctx, "sandbox-pause-stop")
 			defer childSpan.End()
 
-			err := sbx.Stop(ctx)
+			err := sbx.Close(ctx)
 			if err != nil {
 				sbxlogger.I(sbx).Error("error stopping sandbox after snapshot", logger.WithSandboxID(in.SandboxId), zap.Error(err))
 			}

packages/orchestrator/internal/template/build/core/filesystem/ext4.go

@@ -229,10 +229,15 @@ func ReadFile(ctx context.Context, tracer trace.Tracer, rootfsPath string, fileP
 	_, statSpan := tracer.Start(ctx, "ext4-read-file")
 	defer statSpan.End()
 
+	_, err := os.Stat(rootfsPath)
+	if err != nil {
+		return "1", fmt.Errorf("rootfs file does not exist: %w", err)
+	}
+
 	cmd := exec.Command("debugfs", "-R", fmt.Sprintf("cat \"%s\"", filePath), rootfsPath)
-	out, err := cmd.Output()
+	out, err := cmd.CombinedOutput()
 	if err != nil {
-		return "1", fmt.Errorf("error reading file: %w", err)
+		return "1", fmt.Errorf("error reading file %s: %w", filePath, err)
 	}
 
 	return string(out), nil

packages/orchestrator/internal/template/build/layer/layer_executor.go

@@ -92,13 +92,7 @@ func (lb *LayerExecutor) BuildLayer(
 	if err != nil {
 		return metadata.Template{}, err
 	}
-	defer sbx.Stop(ctx)
-	go func() {
-		err := sbx.Wait(context.WithoutCancel(ctx))
-		if err != nil {
-			lb.logger.Error("error waiting for sandbox", zap.Error(err))
-		}
-	}()
+	defer sbx.Close(ctx)
 
 	// Add to proxy so we can call envd commands
 	lb.sandboxes.Insert(sbx.Runtime.SandboxID, sbx)

packages/orchestrator/internal/template/build/phases/base/builder.go

@@ -284,7 +284,7 @@ func (bb *BaseBuilder) buildLayerFromOCI(
 	if err != nil {
 		return metadata.Template{}, fmt.Errorf("error creating sandbox: %w", err)
 	}
-	defer sourceSbx.Stop(ctx)
+	defer sourceSbx.Close(ctx)
 
 	err = sourceSbx.WaitForEnvd(
 		ctx,

packages/orchestrator/internal/template/build/phases/base/provision.go

@@ -99,7 +99,7 @@ func (bb *BaseBuilder) provisionSandbox(
 	if err != nil {
 		return fmt.Errorf("error creating sandbox: %w", err)
 	}
-	defer sbx.Stop(ctx)
+	defer sbx.Close(ctx)
 
 	err = sbx.WaitForExit(
 		ctx,

packages/shared/pkg/utils/set_once.go

@@ -44,6 +44,13 @@ func (s *SetOnce[T]) SetError(err error) error {
 	return s.setResult(result[T]{err: err})
 }
 
+func (s *SetOnce[T]) SetResult(value T, err error) error {
+	if err != nil {
+		return s.SetError(err)
+	}
+	return s.SetValue(value)
+}
+
 var ErrAlreadySet = fmt.Errorf("value already set")
 
 // SetResult internal method for setting the result only once.