feat: add unified worker pool for recursive downloads

internal/pool/pool.go

@@ -0,0 +1,125 @@
+// Package pool provides a protocol-agnostic recursive worker pool.
+//
+// It replaces the three near-identical semaphore+WaitGroup
+// implementations that existed in the FTP, SFTP, and WebDAV packages
+// with a single generic type.  Each protocol supplies its own task
+// type and handler function; the pool takes care of concurrency
+// limiting, goroutine lifecycle, and error collection.
+//
+// Tasks may submit additional tasks from inside the handler (tree
+// recursion).  Submit never blocks the caller—each task gets its
+// own goroutine that waits for a semaphore slot—so recursive
+// submissions cannot deadlock even when all pool slots are occupied.
+package pool
+
+import (
+	"context"
+	"sync"
+)
+
+// ctxKey is an unexported type for context value keys in this package.
+type ctxKey struct{}
+
+// Pool manages a bounded set of goroutines that execute tasks of
+// type T.  Tasks are submitted via Submit and executed by the
+// handler function provided to New.  The pool is safe for
+// concurrent use from multiple goroutines.
+type Pool[T any] struct {
+	handler func(ctx context.Context, task T)
+
+	sem    chan struct{}
+	wg     sync.WaitGroup
+	errMu  sync.Mutex
+	errs   []error
+	ctx    context.Context
+	cancel context.CancelFunc
+}
+
+// New creates a pool with max parallelism workers.  handler is
+// called for every submitted task, running in its own goroutine
+// bounded by the semaphore.  The returned pool is ready to use.
+//
+// ctx is the parent context for all handler goroutines.  When ctx
+// is cancelled, waiting and running handlers observe the
+// cancellation via the context passed to them.
+//
+// The context passed to handler carries a reference to the pool
+// itself, retrievable via Get.  This allows handlers to submit
+// child tasks without an explicit closure capture.
+func New[T any](ctx context.Context, maxParallelism int, handler func(ctx context.Context, task T)) *Pool[T] {
+	if maxParallelism < 1 {
+		maxParallelism = 1
+	}
+
+	ctx, cancel := context.WithCancel(ctx)
+
+	return &Pool[T]{
+		handler: handler,
+		sem:     make(chan struct{}, maxParallelism),
+		cancel:  cancel,
+		ctx:     ctx,
+	}
+}
+
+// Get retrieves the pool of type *Pool[T] from the context.
+// It returns nil when the context does not carry a pool (for
+// example, when called outside a handler goroutine).
+func Get[T any](ctx context.Context) *Pool[T] {
+	v, _ := ctx.Value(ctxKey{}).(*Pool[T])
+	return v
+}
+
+// Submit enqueues a task for execution.  Each submitted task gets
+// its own goroutine that will wait for a semaphore slot before
+// running the handler.  Submit may be called from inside a handler
+// (for recursive tasks).
+//
+// Submit is a no-op if the pool has already been waited on.
+func (p *Pool[T]) Submit(task T) {
+	// After Wait returns the context is cancelled; skip work that
+	// can never run.
+	select {
+	case <-p.ctx.Done():
+		return
+	default:
+	}
+
+	p.wg.Add(1)
+	go func() {
+		defer p.wg.Done()
+		select {
+		case p.sem <- struct{}{}:
+			defer func() { <-p.sem }()
+			handlerCtx := context.WithValue(p.ctx, ctxKey{}, p)
+			p.handler(handlerCtx, task)
+		case <-p.ctx.Done():
+			return
+		}
+	}()
+}
+
+// Wait blocks until all submitted tasks (including tasks submitted
+// by handlers) have completed.  It returns all errors collected
+// during execution.  After Wait returns the pool must not be used.
+func (p *Pool[T]) Wait() []error {
+	p.wg.Wait()
+	p.cancel()
+
+	p.errMu.Lock()
+	defer p.errMu.Unlock()
+	if len(p.errs) == 0 {
+		return nil
+	}
+	return p.errs
+}
+
+// AddError records a task-level error.  It is safe to call from
+// any goroutine, typically from inside a handler.
+func (p *Pool[T]) AddError(err error) {
+	if err == nil {
+		return
+	}
+	p.errMu.Lock()
+	p.errs = append(p.errs, err)
+	p.errMu.Unlock()
+}