feat(pool): add EachCtx for cancellable fan-out (#352)

Each has no way to observe cancellation, so a fan-out scanner (full
port sweep, directory brute-force) keeps draining its whole queue even
after ctrl-c or -max-time fires. EachCtx takes a context and stops
feeding workers once it's cancelled, while still handing the ctx to
the callback so an in-flight item can bail out too. Each is now a thin
wrapper over EachCtx with context.Background(), so existing callers
are unaffected until they opt in.

Author: TBX3D

internal/pool/pool.go

@@ -16,12 +16,26 @@
 // holding it, the rest keep draining the queue instead of idling behind it.
 package pool
 
-import "sync"
+import (
+	"context"
+	"sync"
+)
 
 // Each runs fn for every item in items, concurrently, across at most workers
 // goroutines. order isn't preserved - fn must be safe to call from multiple
 // goroutines and guard any shared state itself. blocks until every item is done.
 func Each[T any](items []T, workers int, fn func(T)) {
+	EachCtx(context.Background(), items, workers, func(_ context.Context, item T) {
+		fn(item)
+	})
+}
+
+// EachCtx is Each with a context: once ctx is cancelled, workers stop pulling
+// new items off the queue (fn still gets called with ctx for the item it's
+// already holding, so it can bail out mid-item too). this is what lets a
+// fan-out scanner - a full port sweep, a directory brute-force - stop
+// promptly on ctrl-c / -max-time instead of draining the whole queue first.
+func EachCtx[T any](ctx context.Context, items []T, workers int, fn func(context.Context, T)) {
 	if len(items) == 0 {
 		return
 	}
@@ -46,10 +60,14 @@ func Each[T any](items []T, workers int, fn func(T)) {
 	for i := 0; i < workers; i++ {
 		go func() {
 			defer wg.Done()
-			// pull until the queue is drained; a worker that finishes its
-			// current item just grabs the next, which is the work-stealing.
+			// pull until the queue is drained or ctx is cancelled; a worker
+			// that finishes its current item just grabs the next, which is
+			// the work-stealing.
 			for item := range queue {
-				fn(item)
+				if ctx.Err() != nil {
+					return
+				}
+				fn(ctx, item)
 			}
 		}()
 	}

internal/pool/pool_test.go

@@ -13,9 +13,11 @@
 package pool
 
 import (
+	"context"
 	"sync"
 	"sync/atomic"
 	"testing"
+	"time"
 )
 
 // every item runs exactly once across a spread of sizes and worker counts,
@@ -143,3 +145,50 @@ func TestEachCapsAtItemCount(t *testing.T) {
 		t.Fatalf("peak concurrency %d exceeded item count %d", got, items)
 	}
 }
+
+// TestEachCtxStopsEarlyOnCancel proves the cancellation contract documented on
+// EachCtx (pool.go), for the -max-time / ctrl-c fan-out case.
+func TestEachCtxStopsEarlyOnCancel(t *testing.T) {
+	const (
+		items   = 1000
+		workers = 4
+	)
+	work := make([]int, items)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	var processed int64
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		EachCtx(ctx, work, workers, func(_ context.Context, _ int) {
+			// cancel partway through so most items are still queued when it fires.
+			if atomic.AddInt64(&processed, 1) == 20 {
+				cancel()
+			}
+			// give the other workers a chance to observe the cancellation
+			// before grabbing their next item, instead of racing straight
+			// through the buffered channel.
+			time.Sleep(time.Millisecond)
+		})
+	}()
+
+	done := make(chan struct{})
+	go func() {
+		wg.Wait()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("EachCtx did not return after cancellation; workers kept draining the queue")
+	}
+
+	got := atomic.LoadInt64(&processed)
+	if got >= items {
+		t.Fatalf("EachCtx processed all %d items despite cancellation (processed=%d): cancel had no effect", items, got)
+	}
+	t.Logf("EachCtx stopped after %d/%d items once ctx was cancelled", got, items)
+}