mirror of
https://github.com/lunchcat/sif.git
synced 2026-07-28 14:37:01 -07:00
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:
+22
-4
@@ -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)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user