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.
This commit is contained in:
Tigah
2026-07-22 12:56:27 -07:00
committed by GitHub
parent 75597fc90d
commit 7e6f6fe2e6
2 changed files with 71 additions and 4 deletions
+22 -4
View File
@@ -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)
}
}()
}
+49
View File
@@ -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)
}