feat: add unified worker pool for recursive downloads

2026-06-29 15:30:31 +02:00
parent 1672698936
commit 11aed3771d
6 changed files with 563 additions and 320 deletions
@@ -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()
+}
@@ -0,0 +1,210 @@
+//go:build linux || freebsd
+// +build linux freebsd
+
+package pool
+
+import (
+	"context"
+	"sort"
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+func TestPoolBasic(t *testing.T) {
+	var count atomic.Int64
+	p := New(context.Background(), 4, func(_ context.Context, task int) {
+		count.Add(1)
+	})
+
+	for i := 0; i < 10; i++ {
+		p.Submit(i)
+	}
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+	if got := count.Load(); got != 10 {
+		t.Fatalf("processed %d tasks, want 10", got)
+	}
+}
+
+func TestPoolConcurrencyLimit(t *testing.T) {
+	const parallel = 3
+	var running atomic.Int64
+	var maxRunning atomic.Int64
+
+	p := New(context.Background(), parallel, func(_ context.Context, task int) {
+		cur := running.Add(1)
+		for {
+			old := maxRunning.Load()
+			if cur <= old || maxRunning.CompareAndSwap(old, cur) {
+				break
+			}
+		}
+		// Simulate work so concurrent goroutines overlap.
+		for i := 0; i < 1000; i++ {
+		}
+		running.Add(-1)
+	})
+
+	for i := 0; i < 20; i++ {
+		p.Submit(i)
+	}
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+	if got := maxRunning.Load(); got > parallel {
+		t.Fatalf("max concurrent goroutines %d, want <= %d", got, parallel)
+	}
+}
+
+func TestPoolRecursive(t *testing.T) {
+	// Simulate a tree: each task submits two children until depth 3.
+	// Expected total tasks: 1 + 2 + 4 + 8 = 15.
+	var count atomic.Int64
+
+	type task struct {
+		depth int
+	}
+
+	p := New(context.Background(), 4, func(ctx context.Context, tk task) {
+		count.Add(1)
+		if tk.depth > 0 {
+			Get[task](ctx).Submit(task{depth: tk.depth - 1})
+			Get[task](ctx).Submit(task{depth: tk.depth - 1})
+		}
+	})
+
+	p.Submit(task{depth: 3})
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+	if got := count.Load(); got != 15 {
+		t.Fatalf("processed %d tasks, want 15", got)
+	}
+}
+
+func TestPoolErrorCollection(t *testing.T) {
+	p := New(context.Background(), 2, func(ctx context.Context, task int) {
+		if task%2 == 0 {
+			Get[int](ctx).AddError(&testError{task})
+		}
+	})
+
+	for i := 0; i < 6; i++ {
+		p.Submit(i)
+	}
+
+	errs := p.Wait()
+	if len(errs) != 3 {
+		t.Fatalf("collected %d errors, want 3", len(errs))
+	}
+}
+
+func TestGetOutsideHandler(t *testing.T) {
+	// Get returns nil when called outside a handler goroutine.
+	if got := Get[int](context.Background()); got != nil {
+		t.Fatalf("Get outside handler returned %v, want nil", got)
+	}
+}
+
+func TestPoolContextCancel(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	p := New(ctx, 4, func(taskCtx context.Context, task int) {
+		// Each task blocks until context is cancelled.
+		<-taskCtx.Done()
+	})
+
+	for i := 0; i < 5; i++ {
+		p.Submit(i)
+	}
+
+	// Cancel after a short delay so the tasks start.
+	go func() {
+		// Give goroutines time to start and acquire semaphore.
+		cancel()
+	}()
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+}
+
+func TestPoolDeterministic(t *testing.T) {
+	// Verify that all submitted tasks are processed exactly once.
+	const n = 100
+	var mu sync.Mutex
+	seen := make(map[int]int)
+
+	p := New(context.Background(), 8, func(_ context.Context, task int) {
+		mu.Lock()
+		seen[task]++
+		mu.Unlock()
+	})
+
+	for i := 0; i < n; i++ {
+		p.Submit(i)
+	}
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+
+	if len(seen) != n {
+		t.Fatalf("saw %d unique tasks, want %d", len(seen), n)
+	}
+	for k, v := range seen {
+		if v != 1 {
+			t.Errorf("task %d processed %d times, want 1", k, v)
+		}
+	}
+}
+
+func TestPoolSortedResults(t *testing.T) {
+	// Collect results from concurrent workers and verify we get them all.
+	const n = 50
+	var mu sync.Mutex
+	results := make([]int, 0, n)
+
+	p := New(context.Background(), 4, func(_ context.Context, task int) {
+		mu.Lock()
+		results = append(results, task)
+		mu.Unlock()
+	})
+
+	for i := 0; i < n; i++ {
+		p.Submit(i)
+	}
+
+	errs := p.Wait()
+	if len(errs) != 0 {
+		t.Fatalf("unexpected errors: %v", errs)
+	}
+
+	sort.Ints(results)
+	if len(results) != n {
+		t.Fatalf("got %d results, want %d", len(results), n)
+	}
+	for i, v := range results {
+		if v != i {
+			t.Fatalf("results[%d] = %d, want %d", i, v, i)
+		}
+	}
+}
+
+type testError struct {
+	task int
+}
+
+func (e *testError) Error() string {
+	return "error in task"
+}