feat: Better parallelism

pkg/preparation/internal/worker/group.go

@@ -0,0 +1,128 @@
+package worker
+
+import (
+	"context"
+	"sync"
+)
+
+// Group runs a collection of tasks that return [TaskError]. Non-fatal errors
+// are accumulated; a fatal error from any task cancels the context derived by
+// [WithContext] so siblings watching that context can exit early. It's
+// analogous to [golang.org/x/sync/errgroup.Group], with two deliberate
+// differences: (1) it distinguishes fatal from non-fatal errors, and (2) it
+// accumulates every task's errors rather than keeping only the first.
+//
+// If the group is cancelled, the cause will be the fatal error.
+type Group struct {
+	ctx    context.Context
+	cancel context.CancelCauseFunc
+	wg     sync.WaitGroup
+	sem    chan struct{}
+
+	mu  sync.Mutex
+	err taskError
+}
+
+// WithContext returns a new [Group] and a derived context. The context is
+// cancelled the first time a task reports a fatal error, or the first time
+// [Group.Wait] returns, whichever occurs first.
+func WithContext(ctx context.Context) (*Group, context.Context) {
+	ctx, cancel := context.WithCancelCause(ctx)
+	return &Group{ctx: ctx, cancel: cancel}, ctx
+}
+
+// SetLimit limits the number of concurrently-running tasks to n. Must be
+// called before any call to [Group.Go]. A value <= 0 removes any limit.
+func (g *Group) SetLimit(n int) {
+	if n <= 0 {
+		g.sem = nil
+		return
+	}
+	g.sem = make(chan struct{}, n)
+}
+
+// Go runs the given task in a new goroutine. If a concurrency limit is set, Go
+// blocks until a slot is available.  The task's [TaskError] is folded into the
+// group's accumulated result; if the task returns a fatal error, the derived
+// context is cancelled.
+//
+// If the group's context is cancelled when Go is called or while waiting for a
+// slot, the task is not started and Go returns immediately.
+func (g *Group) Go(task func() TaskError) {
+	if g.sem != nil {
+		select {
+		case g.sem <- struct{}{}:
+		case <-g.ctx.Done():
+			return
+		}
+	}
+	g.launch(task)
+}
+
+// TryGo runs the given task in a new goroutine if a slot is available. It
+// returns true if the task was started, false otherwise. If no concurrency
+// limit is set, TryGo always starts the task and returns true.
+func (g *Group) TryGo(task func() TaskError) bool {
+	if g.sem != nil {
+		select {
+		case g.sem <- struct{}{}:
+		default:
+			return false
+		}
+	}
+	g.launch(task)
+	return true
+}
+
+func (g *Group) launch(task func() TaskError) {
+	g.wg.Add(1)
+	go func() {
+		defer func() {
+			if g.sem != nil {
+				<-g.sem
+			}
+			g.wg.Done()
+		}()
+		select {
+		case <-g.ctx.Done():
+			// Group has cancelled before this task started. Skip it silently:
+			// the real fatal is already in the accumulator, and an un-started
+			// task has nothing to report. Tasks already executing when cancel
+			// fires still run to completion and report whatever they want.
+			return
+		default:
+		}
+		g.collect(task())
+	}()
+}
+
+// Wait blocks until all goroutines started with [Group.Go] or [Group.TryGo]
+// have returned, then returns the accumulated [TaskError]. It returns nil iff
+// no task reported a fatal error and no task reported any non-fatal errors.
+func (g *Group) Wait() TaskError {
+	g.wg.Wait()
+
+	// Clean up resources: a `cancel` must always be called eventually. We've just
+	// `Wait()`ed for all tasks to finish, so this won't stop anything.
+	g.cancel(nil)
+
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	if g.err.isEmpty() {
+		return nil
+	}
+	result := g.err
+	return &result
+}
+
+func (g *Group) collect(result TaskError) {
+	if result == nil {
+		return
+	}
+	g.mu.Lock()
+	g.err.add(result)
+	g.mu.Unlock()
+	if result.FatalError() != nil {
+		g.cancel(result.FatalError())
+	}
+}

pkg/preparation/internal/worker/group_test.go

@@ -0,0 +1,78 @@
+package worker_test
+
+import (
+	"errors"
+	"sync/atomic"
+	"testing"
+
+	"github.com/storacha/guppy/pkg/preparation/internal/worker"
+	"github.com/stretchr/testify/require"
+)
+
+func TestGroupLaunchSkipsIfCancelled(t *testing.T) {
+	t.Run("task is skipped if group ctx is already cancelled before launch", func(t *testing.T) {
+		fatalErr := errors.New("fatal")
+
+		g, gctx := worker.WithContext(t.Context())
+
+		// First task cancels the group by returning a fatal.
+		g.Go(func() worker.TaskError {
+			return worker.NewFatalError(fatalErr)
+		})
+
+		// Wait for cancellation to propagate before enqueuing the second
+		// task, so the second task's goroutine sees a cancelled ctx on entry.
+		<-gctx.Done()
+
+		var ranSecond atomic.Bool
+		g.Go(func() worker.TaskError {
+			ranSecond.Store(true)
+			return worker.NewFatalError(errors.New("should not be reported"))
+		})
+
+		res := g.Wait()
+		require.False(t, ranSecond.Load(), "task enqueued after cancellation should be skipped")
+		require.NotNil(t, res)
+		require.True(t, res.IsFatal())
+		require.ErrorIs(t, res.FatalError(), fatalErr)
+		require.NotContains(t, res.FatalError().Error(), "should not be reported")
+	})
+
+	t.Run("in-flight task can still report errors after cancellation", func(t *testing.T) {
+		firstFatal := errors.New("first fatal")
+		secondFatal := errors.New("second fatal from in-flight task")
+
+		g, gctx := worker.WithContext(t.Context())
+
+		started := make(chan struct{})
+		release := make(chan struct{})
+
+		// In-flight task: signals it has started, waits to be released, then
+		// returns a fatal. It's already executing when the group cancels.
+		g.Go(func() worker.TaskError {
+			close(started)
+			<-release
+			return worker.NewFatalError(secondFatal)
+		})
+
+		// Wait for the in-flight task to be running.
+		<-started
+
+		// Fire a fatal from another task to cancel the group.
+		g.Go(func() worker.TaskError {
+			return worker.NewFatalError(firstFatal)
+		})
+
+		// Wait for gctx to be cancelled.
+		<-gctx.Done()
+
+		// Release the in-flight task so it can complete.
+		close(release)
+
+		res := g.Wait()
+		require.NotNil(t, res)
+		require.True(t, res.IsFatal())
+		require.ErrorIs(t, res.FatalError(), firstFatal, "first fatal should be recorded")
+		require.ErrorIs(t, res.FatalError(), secondFatal, "in-flight task's fatal should also be recorded")
+	})
+}

pkg/preparation/internal/worker/types.go

@@ -0,0 +1,133 @@
+package worker
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+)
+
+// TaskError is the error type returned by a [Task]. A TaskError may carry a
+// fatal error (which causes the worker to cancel dispatch of remaining tasks)
+// and/or a set of non-fatal errors (which are collected and reported after
+// all tasks complete).
+type TaskError interface {
+	error
+	NonFatalErrors() []error
+	FatalError() error
+	IsFatal() bool
+}
+
+// NewFatalError returns a [TaskError] carrying the given error as a fatal
+// error. If err is nil, it returns nil.
+func NewFatalError(err error) TaskError {
+	if err == nil {
+		return nil
+	}
+	return &taskError{fatalError: err}
+}
+
+// NewNonFatalError returns a [TaskError] carrying the given errors as
+// non-fatal errors. Any nil entries are dropped; if the result would be empty,
+// it returns nil.
+func NewNonFatalError(errs ...error) TaskError {
+	var filtered []error
+	for _, err := range errs {
+		if err != nil {
+			filtered = append(filtered, err)
+		}
+	}
+	if len(filtered) == 0 {
+		return nil
+	}
+	return &taskError{nonFatalErrors: filtered}
+}
+
+type taskError struct {
+	nonFatalErrors []error
+	fatalError     error
+}
+
+// Task is a worker task. It returns a [TaskError] to signal the result:
+//   - return nil for success,
+//   - return [NewFatalError] (or any TaskError whose IsFatal is true) to abort
+//     further dispatch,
+//   - return [NewNonFatalError] to report errors that should be collected but
+//     not abort dispatch.
+type Task func(context.Context) TaskError
+
+func (e *taskError) NonFatalErrors() []error {
+	return e.nonFatalErrors
+}
+
+func (e *taskError) FatalError() error {
+	return e.fatalError
+}
+
+func (e *taskError) IsFatal() bool {
+	return e.fatalError != nil
+}
+
+func (e *taskError) Error() string {
+	nonFatalString := ""
+	if len(e.nonFatalErrors) > 0 {
+		nonFatalString = "non-fatal errors:"
+	}
+	for _, err := range e.nonFatalErrors {
+		nonFatalString += fmt.Sprintf("\n- %s", err)
+	}
+
+	fatalString := ""
+	if e.fatalError != nil {
+		fatalString = fmt.Sprintf("fatal error: %s", e.fatalError)
+	}
+
+	return fmt.Sprintf("worker encountered %s", strings.Join([]string{nonFatalString, fatalString}, "\n"))
+}
+
+func (e *taskError) Unwrap() []error {
+	errs := make([]error, len(e.nonFatalErrors))
+	copy(errs, e.nonFatalErrors)
+	if e.fatalError != nil {
+		errs = append(errs, e.fatalError)
+	}
+	return errs
+}
+
+// add folds another [TaskError] into the receiver in place.
+func (e *taskError) add(other TaskError) {
+	if other == nil {
+		return
+	}
+	e.nonFatalErrors = append(e.nonFatalErrors, other.NonFatalErrors()...)
+	if other.FatalError() != nil {
+		e.fatalError = errors.Join(e.fatalError, other.FatalError())
+	}
+}
+
+// isEmpty reports whether the taskError carries neither a fatal error nor any
+// non-fatal errors.
+func (e *taskError) isEmpty() bool {
+	return e.fatalError == nil && len(e.nonFatalErrors) == 0
+}
+
+// Join combines two [TaskError]s. Non-fatal errors are concatenated; fatal
+// errors are joined with [errors.Join]. Either operand may be nil.
+func Join(a TaskError, b TaskError) TaskError {
+	switch {
+	case a == nil && b == nil:
+		return nil
+	case a == nil:
+		return b
+	case b == nil:
+		return a
+	}
+	joined := &taskError{
+		nonFatalErrors: append(a.NonFatalErrors(), b.NonFatalErrors()...),
+		fatalError:     errors.Join(a.FatalError(), b.FatalError()),
+	}
+	if joined.isEmpty() {
+		return nil
+	}
+	return joined
+}