refactor SetOnce[struct{}] to ErrorOnce

Author: dobrac

packages/orchestrator/internal/sandbox/fc/process.go

@@ -55,7 +55,7 @@ type Process struct {
 	kernelPath         string
 	files              *storage.SandboxFiles
 
-	Exit *utils.SetOnce[struct{}]
+	Exit *utils.ErrorOnce
 
 	client *apiClient
 }
@@ -110,7 +110,7 @@ func NewProcess(
 
 	return &Process{
 		Versions:              versions,
-		Exit:                  utils.NewSetOnce[struct{}](),
+		Exit:                  utils.NewErrorOnce(),
 		cmd:                   cmd,
 		firecrackerSocketPath: files.SandboxFirecrackerSocketPath(),
 		client:                newApiClient(files.SandboxFirecrackerSocketPath()),
@@ -170,7 +170,7 @@ func (p *Process) configure(
 			if errors.As(waitErr, &exitErr) {
 				// Check if the process was killed by a signal
 				if status, ok := exitErr.Sys().(syscall.WaitStatus); ok && status.Signaled() && (status.Signal() == syscall.SIGKILL || status.Signal() == syscall.SIGTERM) {
-					p.Exit.SetValue(struct{}{})
+					p.Exit.SetError(nil)
 
 					return
 				}
@@ -186,7 +186,7 @@ func (p *Process) configure(
 			return
 		}
 
-		p.Exit.SetValue(struct{}{})
+		p.Exit.SetError(nil)
 	}()
 
 	// Wait for the FC process to start so we can use FC API
@@ -435,7 +435,7 @@ func (p *Process) Stop() error {
 			}
 
 		// If the FC process exited, we can return.
-		case <-p.Exit.Done:
+		case <-p.Exit.Done():
 			return
 		}
 	}()

packages/orchestrator/internal/sandbox/sandbox.go

@@ -99,7 +99,7 @@ type Sandbox struct {
 
 	APIStoredConfig *orchestrator.SandboxConfig
 
-	exit *utils.SetOnce[struct{}]
+	exit *utils.ErrorOnce
 }
 
 func (s *Sandbox) LoggerMetadata() sbxlogger.SandboxMetadata {
@@ -132,7 +132,7 @@ func CreateSandbox(
 	childCtx, childSpan := tracer.Start(ctx, "create-sandbox")
 	defer childSpan.End()
 
-	exit := utils.NewSetOnce[struct{}]()
+	exit := utils.NewErrorOnce()
 
 	cleanup := NewCleanup()
 	defer func() {
@@ -292,10 +292,10 @@ func CreateSandbox(
 
 	go func() {
 		// If the process exists, stop the sandbox properly
-		_, fcErr := fcHandle.Exit.Wait()
+		fcErr := fcHandle.Exit.Wait()
 		err := sbx.Stop(context.WithoutCancel(ctx), tracer)
 
-		exit.SetResult(struct{}{}, errors.Join(err, fcErr))
+		exit.SetError(errors.Join(err, fcErr))
 	}()
 
 	return sbx, nil
@@ -320,7 +320,7 @@ func ResumeSandbox(
 	childCtx, childSpan := tracer.Start(ctx, "resume-sandbox")
 	defer childSpan.End()
 
-	exit := utils.NewSetOnce[struct{}]()
+	exit := utils.NewErrorOnce()
 
 	cleanup := NewCleanup()
 	defer func() {
@@ -399,7 +399,7 @@ func ResumeSandbox(
 	defer cancelUffdStartCtx(fmt.Errorf("uffd finished starting"))
 
 	go func() {
-		_, uffdWaitErr := fcUffd.Exit().Wait()
+		uffdWaitErr := fcUffd.Exit().Wait()
 
 		cancelUffdStartCtx(fmt.Errorf("uffd process exited: %w", errors.Join(uffdWaitErr, context.Cause(uffdStartCtx))))
 	}()
@@ -521,23 +521,22 @@ func ResumeSandbox(
 	go func() {
 		// Wait for either uffd or fc process to exit
 		select {
-		case <-fcUffd.Exit().Done:
-		case <-fcHandle.Exit.Done:
+		case <-fcUffd.Exit().Done():
+		case <-fcHandle.Exit.Done():
 		}
 
 		err := sbx.Stop(context.WithoutCancel(ctx), tracer)
 
-		_, uffdWaitErr := fcUffd.Exit().Wait()
-		_, fcErr := fcHandle.Exit.Wait()
-		exit.SetResult(struct{}{}, errors.Join(err, fcErr, uffdWaitErr))
+		uffdWaitErr := fcUffd.Exit().Wait()
+		fcErr := fcHandle.Exit.Wait()
+		exit.SetError(errors.Join(err, fcErr, uffdWaitErr))
 	}()
 
 	return sbx, nil
 }
 
 func (s *Sandbox) Wait(ctx context.Context) error {
-	_, err := s.exit.WaitWithContext(ctx)
-	return err
+	return s.exit.WaitWithContext(ctx)
 }
 
 func (s *Sandbox) Close(ctx context.Context) error {
@@ -565,7 +564,7 @@ func (s *Sandbox) Stop(ctx context.Context, tracer trace.Tracer) error {
 
 	// 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
+	<-s.process.Exit.Done()
 
 	uffdStopErr := s.Resources.memory.Stop()
 	if uffdStopErr != nil {
@@ -907,8 +906,8 @@ func (s *Sandbox) WaitForExit(
 		return fmt.Errorf("waiting for exit took too long")
 	case <-ctx.Done():
 		return nil
-	case <-s.exit.Done:
-		_, err := s.exit.Result()
+	case <-s.exit.Done():
+		err := s.exit.Error()
 		if err == nil {
 			return nil
 		}
@@ -942,8 +941,8 @@ func (s *Sandbox) WaitForEnvd(
 			syncCancel(fmt.Errorf("syncing took too long"))
 		case <-syncCtx.Done():
 			return
-		case <-s.process.Exit.Done:
-			_, err := s.process.Exit.Result()
+		case <-s.process.Exit.Done():
+			err := s.process.Exit.Error()
 
 			syncCancel(fmt.Errorf("fc process exited prematurely: %w", err))
 		}

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

@@ -26,7 +26,7 @@ const (
 )
 
 type Uffd struct {
-	exit    *utils.SetOnce[struct{}]
+	exit    *utils.ErrorOnce
 	readyCh chan struct{}
 
 	fdExit *fdexit.FdExit
@@ -49,7 +49,7 @@ func New(memfile block.ReadonlyDevice, socketPath string, blockSize int64) (*Uff
 	}
 
 	return &Uffd{
-		exit:       utils.NewSetOnce[struct{}](),
+		exit:       utils.NewErrorOnce(),
 		readyCh:    make(chan struct{}, 1),
 		fdExit:     fdExit,
 		memfile:    trackedMemfile,
@@ -76,7 +76,7 @@ func (u *Uffd) Start(sandboxId string) error {
 		closeErr := u.lis.Close()
 		fdExitErr := u.fdExit.Close()
 
-		u.exit.SetResult(struct{}{}, errors.Join(handleErr, closeErr, fdExitErr))
+		u.exit.SetError(errors.Join(handleErr, closeErr, fdExitErr))
 
 		close(u.readyCh)
 	}()
@@ -165,7 +165,7 @@ func (u *Uffd) Ready() chan struct{} {
 	return u.readyCh
 }
 
-func (u *Uffd) Exit() *utils.SetOnce[struct{}] {
+func (u *Uffd) Exit() *utils.ErrorOnce {
 	return u.exit
 }
 

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

@@ -13,5 +13,5 @@ type MemoryBackend interface {
 	Start(sandboxId string) error
 	Stop() error
 	Ready() chan struct{}
-	Exit() *utils.SetOnce[struct{}]
+	Exit() *utils.ErrorOnce
 }

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

@@ -13,7 +13,7 @@ type NoopMemory struct {
 
 	dirty *bitset.BitSet
 
-	exit *utils.SetOnce[struct{}]
+	exit *utils.ErrorOnce
 }
 
 func NewNoopMemory(size, blockSize int64) *NoopMemory {
@@ -26,7 +26,7 @@ func NewNoopMemory(size, blockSize int64) *NoopMemory {
 		size:      size,
 		blockSize: blockSize,
 		dirty:     dirty,
-		exit:      utils.NewSetOnce[struct{}](),
+		exit:      utils.NewErrorOnce(),
 	}
 }
 
@@ -43,7 +43,7 @@ func (m *NoopMemory) Start(sandboxId string) error {
 }
 
 func (m *NoopMemory) Stop() error {
-	return m.exit.SetValue(struct{}{})
+	return m.exit.SetSuccess()
 }
 
 func (m *NoopMemory) Ready() chan struct{} {
@@ -52,6 +52,6 @@ func (m *NoopMemory) Ready() chan struct{} {
 	return ch
 }
 
-func (m *NoopMemory) Exit() *utils.SetOnce[struct{}] {
+func (m *NoopMemory) Exit() *utils.ErrorOnce {
 	return m.exit
 }

packages/shared/pkg/utils/error_once.go

@@ -0,0 +1,55 @@
+package utils
+
+import "context"
+
+// ErrorOnce is a wrapper around SetOnce that can only be set with an error.
+// It's useful for cases where you only need to signal completion with a potential error,
+// without carrying any value.
+type ErrorOnce struct {
+	setOnce *SetOnce[struct{}]
+}
+
+// NewErrorOnce creates a new ErrorOnce instance.
+func NewErrorOnce() *ErrorOnce {
+	return &ErrorOnce{
+		setOnce: NewSetOnce[struct{}](),
+	}
+}
+
+// SetError sets the error once. Subsequent calls will return ErrAlreadySet.
+func (e *ErrorOnce) SetError(err error) error {
+	return e.setOnce.SetError(err)
+}
+
+// SetSuccess marks the operation as completed successfully (no error).
+// This is equivalent to SetError(nil).
+func (e *ErrorOnce) SetSuccess() error {
+	return e.setOnce.SetError(nil)
+}
+
+// Wait blocks until an error is set and returns it.
+// Returns nil if the operation completed successfully.
+func (e *ErrorOnce) Wait() error {
+	_, err := e.setOnce.Wait()
+	return err
+}
+
+// Error returns the error if one has been set, or ErrNotSet if not set yet.
+// Unlike Wait, this doesn't block.
+func (e *ErrorOnce) Error() error {
+	_, err := e.setOnce.Result()
+	return err
+}
+
+// WaitWithContext waits for an error to be set with context cancellation support.
+// Returns the set error or ctx.Err() if the context is cancelled first.
+func (e *ErrorOnce) WaitWithContext(ctx context.Context) error {
+	_, err := e.setOnce.WaitWithContext(ctx)
+	return err
+}
+
+// Done returns a channel that's closed when an error is set.
+// This allows using ErrorOnce in select statements.
+func (e *ErrorOnce) Done() <-chan struct{} {
+	return e.setOnce.Done
+}

packages/shared/pkg/utils/error_once_test.go

@@ -0,0 +1,52 @@
+package utils
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestErrorOnce(t *testing.T) {
+	errorOnce := NewErrorOnce()
+
+	// Test setting error
+	expectedErr := fmt.Errorf("test error")
+	err := errorOnce.SetError(expectedErr)
+	require.NoError(t, err)
+
+	// Wait should return the error
+	err = errorOnce.Wait()
+	require.Error(t, err)
+	assert.Equal(t, expectedErr, err)
+
+	// Trying to set again should return ErrAlreadySet
+	err = errorOnce.SetError(fmt.Errorf("another error"))
+	require.ErrorIs(t, err, ErrAlreadySet)
+
+	// Wait should still return the original error
+	err = errorOnce.Wait()
+	require.Error(t, err)
+	assert.Equal(t, expectedErr, err)
+}
+
+func TestErrorOnceSetSuccess(t *testing.T) {
+	errorOnce := NewErrorOnce()
+
+	// Test setting success (nil error)
+	err := errorOnce.SetSuccess()
+	require.NoError(t, err)
+
+	// Wait should return nil
+	err = errorOnce.Wait()
+	require.NoError(t, err)
+
+	// Trying to set error after success should return ErrAlreadySet
+	err = errorOnce.SetError(fmt.Errorf("test error"))
+	require.ErrorIs(t, err, ErrAlreadySet)
+
+	// Wait should still return nil
+	err = errorOnce.Wait()
+	require.NoError(t, err)
+}

packages/shared/pkg/utils/set_once.go

@@ -67,7 +67,7 @@ func (s *SetOnce[T]) setResult(r result[T]) error {
 		defer s.mu.Unlock()
 
 		if s.res != nil {
-			return fmt.Errorf("value already set")
+			return ErrAlreadySet
 		}
 
 		s.res = &r