diff --git a/benchmark/ants_compare_test.go b/benchmark/ants_compare_test.go new file mode 100644 index 0000000..7920378 --- /dev/null +++ b/benchmark/ants_compare_test.go @@ -0,0 +1,109 @@ +package benchmark_test + +import ( + "sync" + "testing" + "time" + + agilepool "github.com/Yiming1997/agilePool/v2" + "github.com/panjf2000/ants/v2" +) + +// cmpWorkerCap mirrors the worker capacity used by the agilePool benchmarks so +// ants and agilePool run under the same concurrency budget. +const cmpWorkerCap = 20000 + +func newAgilePoolForCmp() *agilepool.Pool { + return agilepool.NewPool(agilepool.NewConfig( + agilepool.WithCleanPeriod(500*time.Millisecond), + agilepool.WithTaskQueueSize(100000), + agilepool.WithWorkerNumCapacity(cmpWorkerCap), + agilepool.WithIdleContainerType(agilepool.LinkedListType), + )) +} + +// --------------------------------------------------------------------------- +// Metric 1 — hand-off throughput: no-op tasks. +// +// The task body does no work, so the measured cost is entirely the pool's +// task hand-off (submit -> worker pickup). ns/op therefore reads directly as +// the per-task hand-off cost; lower is better. +// --------------------------------------------------------------------------- + +func BenchmarkCmpChurnAnts(b *testing.B) { + pool, err := ants.NewPool(cmpWorkerCap) + if err != nil { + b.Fatal(err) + } + defer pool.Release() + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + _ = pool.Submit(func() {}) + } + }) +} + +func BenchmarkCmpChurnAgilePool(b *testing.B) { + pool := newAgilePoolForCmp() + defer pool.Close() + + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + pool.Submit(agilepool.TaskFunc(func() error { return nil })) + } + }) + pool.Wait() +} + +// --------------------------------------------------------------------------- +// Metric 2 — CPU usage under an execution-bound load: 10ms sleep tasks (the +// same shape as the agilePool SequentialLinkedList benchmark). Here the task +// duration is the limiter, so the meaningful metric is CPU burned per +// completed task, read from the CPU profile of each benchmark. +// --------------------------------------------------------------------------- + +func BenchmarkCmpSleepAnts(b *testing.B) { + pool, err := ants.NewPool(cmpWorkerCap) + if err != nil { + b.Fatal(err) + } + defer pool.Release() + + var wg sync.WaitGroup + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + wg.Add(1) + _ = pool.Submit(func() { + defer wg.Done() + time.Sleep(10 * time.Millisecond) + }) + } + }) + wg.Wait() +} + +func BenchmarkCmpSleepAgilePool(b *testing.B) { + pool := newAgilePoolForCmp() + defer pool.Close() + + // Use the same WaitGroup-based completion tracking as the ants variant so + // neither side gets a harness-specific advantage (both closures capture + // &wg and therefore heap-allocate identically). + var wg sync.WaitGroup + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + wg.Add(1) + pool.Submit(agilepool.TaskFunc(func() error { + defer wg.Done() + time.Sleep(10 * time.Millisecond) + return nil + })) + } + }) + wg.Wait() +} diff --git a/benchmark/scaling_matrix_test.go b/benchmark/scaling_matrix_test.go new file mode 100644 index 0000000..aae378a --- /dev/null +++ b/benchmark/scaling_matrix_test.go @@ -0,0 +1,65 @@ +package benchmark_test + +import ( + "fmt" + "testing" + "time" + + agilepool "github.com/Yiming1997/agilePool/v2" +) + +// Implementation-agnostic scaling matrix. It deliberately avoids any +// version-specific config so the exact same file compiles and runs against +// both the pre-sharding implementation and the current one, letting the two +// be compared by swapping the checkout only. +var scalingCaps = []int{128, 1024, 8192, 20000} + +func newScalingPool(capacity int) *agilepool.Pool { + return agilepool.NewPool(agilepool.NewConfig( + agilepool.WithCleanPeriod(500*time.Millisecond), + agilepool.WithTaskQueueSize(100000), + agilepool.WithWorkerNumCapacity(int64(capacity)), + agilepool.WithIdleContainerType(agilepool.LinkedListType), + )) +} + +// BenchmarkScalingNoop is hand-off bound: no-op tasks, so throughput is the +// pool's submit -> worker pickup rate. +func BenchmarkScalingNoop(b *testing.B) { + for _, capacity := range scalingCaps { + b.Run(fmt.Sprintf("cap=%d", capacity), func(b *testing.B) { + pool := newScalingPool(capacity) + defer pool.Close() + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + pool.Submit(agilepool.TaskFunc(func() error { return nil })) + } + }) + }) + } +} + +// BenchmarkScalingSleep10ms is execution bound: each task sleeps 10ms, so the +// limiter is capacity/10ms and the channel hand-off overhead shows up only as +// extra CPU and scheduling jitter, not throughput. +func BenchmarkScalingSleep10ms(b *testing.B) { + for _, capacity := range scalingCaps { + b.Run(fmt.Sprintf("cap=%d", capacity), func(b *testing.B) { + pool := newScalingPool(capacity) + defer pool.Close() + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + pool.Submit(agilepool.TaskFunc(func() error { + time.Sleep(10 * time.Millisecond) + return nil + })) + } + }) + pool.Wait() + }) + } +} diff --git a/config.go b/config.go index 0554d92..7a69e5d 100644 --- a/config.go +++ b/config.go @@ -34,7 +34,8 @@ const ( type Config struct { cleanPeriod time.Duration - taskQueueSize int64 // capacity of the internal handoff channel + taskQueueSize int64 // total capacity of the internal handoff queues + taskQueueShards int // number of handoff-queue shards; 0 = auto workerNumCapacity int64 workMode WorkMode idleContainerType IdleContainerType @@ -100,6 +101,18 @@ func WithTaskQueueSize(size int64) ConfigOption { } } +// WithTaskQueueShards sets how many shards the internal handoff queue is split +// into. Sharding spreads channel-lock contention across multiple channel locks, +// which matters when the pool runs a very large number of workers that would +// otherwise serialise on a single channel. The value is rounded up to a power +// of two and each shard gets at least one slot. n <= 0 selects a default: a +// single shard for small queues and GOMAXPROCS (capped) for large ones. +func WithTaskQueueShards(n int) ConfigOption { + return func(c *Config) { + c.taskQueueShards = n + } +} + func WithWorkerNumCapacity(capacity int64) ConfigOption { return func(c *Config) { if capacity > 0 { diff --git a/go.mod b/go.mod index 2f73599..2a1090e 100644 --- a/go.mod +++ b/go.mod @@ -2,10 +2,14 @@ module github.com/Yiming1997/agilePool/v2 go 1.23.4 -require github.com/stretchr/testify v1.11.1 +require ( + github.com/panjf2000/ants/v2 v2.12.1 + github.com/stretchr/testify v1.11.1 +) require ( github.com/davecgh/go-spew v1.1.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect + golang.org/x/sync v0.11.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index c4c1710..eff58bd 100644 --- a/go.sum +++ b/go.sum @@ -1,9 +1,13 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/panjf2000/ants/v2 v2.12.1 h1:BWvU2wHpyXWxhhNXsGB6JXLCNbshyLd1QxvoAmZnu10= +github.com/panjf2000/ants/v2 v2.12.1/go.mod h1:tSQuaNQ6r6NRhPt+IZVUevvDyFMTs+eS4ztZc52uJTY= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U= github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U= +golang.org/x/sync v0.11.0 h1:GGz8+XQP4FvTTrjZPzNKTMFtSXH80RAzG+5ghFPgK9w= +golang.org/x/sync v0.11.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/pool.go b/pool.go index 76b8c2f..e62ac93 100644 --- a/pool.go +++ b/pool.go @@ -32,6 +32,14 @@ const ( // When reached, submitters block on the handoff channel instead of // growing the buffer further, bounding peak memory. maxChunkLen = 100_000 + + // minShardableQueueSize is the smallest handoff capacity that is split + // into multiple shards. Smaller queues stay single-channel so the exact + // backpressure behaviour of tiny configurations is preserved. + minShardableQueueSize = 1024 + + // maxAutoShards caps the shard count chosen automatically from GOMAXPROCS. + maxAutoShards = 32 ) type WorkMode int8 @@ -65,7 +73,11 @@ func GetDefaultPool() *Pool { } type Pool struct { - taskQueue chan Task + taskQueues []chan Task // sharded handoff queues; len is a power of two + shardMask uint64 // len(taskQueues)-1, enables mask-based shard selection + sendSeq uint64 // atomic round-robin cursor picking the target shard + workerSeq uint64 // atomic cursor assigning distinct scan start shards + taskQueueSize int // configured total handoff capacity (GetTaskQueueCapacity) closePoolCn chan struct{} capacity int64 // The maximum number of workers in the pool. runningWorkersNum int64 @@ -115,11 +127,25 @@ func NewPool(c *Config) *Pool { muIdle: &sync.Mutex{}, // default value, overridden below based on config logger: log.Default(), capacity: c.workerNumCapacity, - taskQueue: make(chan Task, c.taskQueueSize), taskBuf: newChunkedTaskBuffer(), panicHandler: c.panicHandler, } + // Shard the handoff queue so a very large worker count does not serialise + // on a single channel lock. Each shard gets at least one slot and the + // total capacity stays >= the configured value. + shardCount := resolveTaskQueueShards(c) + perShardCap := (int(c.taskQueueSize) + shardCount - 1) / shardCount + if perShardCap < 1 { + perShardCap = 1 + } + p.taskQueues = make([]chan Task, shardCount) + for i := range p.taskQueues { + p.taskQueues[i] = make(chan Task, perShardCap) + } + p.shardMask = uint64(shardCount - 1) + p.taskQueueSize = int(c.taskQueueSize) + // Select muIdle lock implementation based on config: SpinLock or MutexLock (sync.Mutex) if c.lockType == SpinLock { p.muIdle = newSpinLock() @@ -159,6 +185,35 @@ func NewPool(c *Config) *Pool { return p } +// resolveTaskQueueShards returns the number of handoff-queue shards to use. +// The result is always a power of two so the round-robin cursor can be masked +// cheaply. An explicit Config value wins; otherwise small queues stay +// single-channel and large queues scale with GOMAXPROCS (capped). +func resolveTaskQueueShards(c *Config) int { + want := c.taskQueueShards + if want <= 0 { + if c.taskQueueSize < minShardableQueueSize { + return 1 + } + want = runtime.GOMAXPROCS(0) + if want > maxAutoShards { + want = maxAutoShards + } + } + if want < 1 { + want = 1 + } + shards := 1 + for shards < want { + shards <<= 1 + } + // Each shard needs at least one slot; never exceed the configured capacity. + for int64(shards) > c.taskQueueSize && shards > 1 { + shards >>= 1 + } + return shards +} + // SetLogger replaces the default standard-library logger. // Pass the same logger instance used elsewhere in your application // (e.g. zap.SugaredLogger) so pool output appears in the same log stream. @@ -215,7 +270,7 @@ func (p *Pool) submit(ctx context.Context, task Task) bool { if atomic.LoadInt64(&p.runningWorkersNum) == 0 { if atomic.CompareAndSwapInt64(&p.runningWorkersNum, 0, 1) { w := p.workerPool.Get().(*worker) - go w.run(nil) + p.startWorker(w) } } hookCtx := context.Background() @@ -226,8 +281,9 @@ func (p *Pool) submit(ctx context.Context, task Task) bool { h.DispatchTaskSubmitted(hookCtx) }) if p.config.workMode == NONBLOCK { + q := p.nextShard() select { - case p.taskQueue <- task: + case q <- task: p.dispatchTaskEnqueuedFor(task) return true default: @@ -236,8 +292,9 @@ func (p *Pool) submit(ctx context.Context, task Task) bool { } } + q := p.nextShard() select { - case p.taskQueue <- task: + case q <- task: p.dispatchTaskEnqueuedFor(task) return true default: @@ -245,7 +302,7 @@ func (p *Pool) submit(ctx context.Context, task Task) bool { result := p.taskBuf.PushAndForward(task, func(t Task) bool { select { - case p.taskQueue <- t: + case p.nextShard() <- t: return true default: return false @@ -261,8 +318,9 @@ func (p *Pool) submit(ctx context.Context, task Task) bool { // SubmitCtx promises that cancellation interrupts this wait. Listening // for pool shutdown also prevents a blocked submitter from being left // behind after Close. + q = p.nextShard() select { - case p.taskQueue <- task: + case q <- task: p.dispatchTaskEnqueuedFor(task) return true case <-ctx.Done(): @@ -364,6 +422,21 @@ func (p *Pool) addRunningWorkersNum(num int64) { atomic.AddInt64(&p.runningWorkersNum, num) } +// nextShard returns the handoff queue for the next submission using a lock-free +// round-robin cursor, spreading producer traffic across shards so that no +// single channel lock becomes a contention hotspot. +func (p *Pool) nextShard() chan Task { + seq := atomic.AddUint64(&p.sendSeq, 1) + return p.taskQueues[seq&p.shardMask] +} + +// startWorker launches w's run loop, assigning it a round-robin scan start +// shard so freshly spawned workers do not all probe the same shard first. +func (p *Pool) startWorker(w *worker) { + w.cursor = atomic.AddUint64(&p.workerSeq, 1) & p.shardMask + go w.run(nil) +} + func (p *Pool) expiredWorkerCleaner() { ticker := time.NewTicker(p.config.cleanPeriod) defer ticker.Stop() @@ -500,7 +573,7 @@ func (p *Pool) scaleIfNeeded() { w = p.workerPool.Get().(*worker) } p.muIdle.Unlock() - go w.run(nil) + p.startWorker(w) p.muIdle.Lock() } p.muIdle.Unlock() @@ -576,14 +649,19 @@ func (p *Pool) GetWorkerCreateCount() int64 { // up by a worker. This is a snapshot of len(taskQueue) and does not // include tasks waiting in the chunked overflow buffer. func (p *Pool) GetTaskQueueLen() int { - // Returns the number of tasks that have been submitted but not yet enqueued for execution. - return len(p.taskQueue) + // Returns the number of tasks across all handoff-queue shards that have + // been submitted but not yet picked up by a worker. + total := 0 + for _, q := range p.taskQueues { + total += len(q) + } + return total } // GetTaskQueueCapacity returns the capacity of the handoff channel // (taskQueue) configured for this pool. func (p *Pool) GetTaskQueueCapacity() int { - return cap(p.taskQueue) + return p.taskQueueSize } // GetIdleWorkerCount returns the number of workers currently parked in diff --git a/shard_stress_test.go b/shard_stress_test.go new file mode 100644 index 0000000..2a95695 --- /dev/null +++ b/shard_stress_test.go @@ -0,0 +1,471 @@ +package agilepool_test + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + agilepool "github.com/Yiming1997/agilePool/v2" +) + +// --------------------------------------------------------------------------- +// Sharded handoff-queue stress tests. +// +// The race detector is unavailable in this environment, so correctness is +// pinned by explicit invariants under heavy concurrency: +// * every accepted task runs exactly once (no loss, no duplication), +// * Wait() returns (no deadlock) across shard/queue/capacity combinations, +// * Close during load never panics or hangs, +// * reentrant submission and panicking tasks stay isolated. +// --------------------------------------------------------------------------- + +// uniqueTaskConfigs covers the shard-count / queue-size / capacity corners, +// including queue sizes below the auto-shard threshold and forced shard counts. +func uniqueTaskConfigs() []struct { + name string + config *agilepool.Config +} { + return []struct { + name string + config *agilepool.Config + }{ + {"cap1_q1", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(1), agilepool.WithTaskQueueSize(1))}, + {"cap4_q10", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(4), agilepool.WithTaskQueueSize(10))}, + {"cap64_q1000", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(64), agilepool.WithTaskQueueSize(1000))}, + {"cap256_q100000_auto", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(256), agilepool.WithTaskQueueSize(100000))}, + {"cap256_q100000_shards2", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(256), agilepool.WithTaskQueueSize(100000), agilepool.WithTaskQueueShards(2))}, + {"cap256_q100000_shards8", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(256), agilepool.WithTaskQueueSize(100000), agilepool.WithTaskQueueShards(8))}, + {"cap128_q4_shards8", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(128), agilepool.WithTaskQueueSize(4), agilepool.WithTaskQueueShards(8))}, + {"cap1000_q10000_shards32", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(1000), agilepool.WithTaskQueueSize(10000), agilepool.WithTaskQueueShards(32))}, + {"cap8_q2_shards32_reduced", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(8), agilepool.WithTaskQueueSize(2), agilepool.WithTaskQueueShards(32))}, + {"cap5000_q100000_auto", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(5000), agilepool.WithTaskQueueSize(100000))}, + {"cap512_q100000_shards64", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(512), agilepool.WithTaskQueueSize(100000), agilepool.WithTaskQueueShards(64))}, + {"cap512_q100000_shards3_nonpow2", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(512), agilepool.WithTaskQueueSize(100000), agilepool.WithTaskQueueShards(3))}, + } +} + +// TestStressUniqueTaskRunsExactlyOnce asserts no task is lost or duplicated +// across every shard/queue/capacity configuration. +func TestStressUniqueTaskRunsExactlyOnce(t *testing.T) { + const ( + tasks = 20000 + submitters = 8 + iterations = 3 + timeoutWait = 30 * time.Second + ) + for _, tc := range uniqueTaskConfigs() { + tc := tc + t.Run(tc.name, func(t *testing.T) { + for iter := 0; iter < iterations; iter++ { + pool := agilepool.NewPool(tc.config) + counts := make([]int32, tasks) + var executed int64 + var submitWG sync.WaitGroup + + for s := 0; s < submitters; s++ { + submitWG.Add(1) + go func(start int) { + defer submitWG.Done() + for i := start; i < tasks; i += submitters { + id := i + pool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt32(&counts[id], 1) + atomic.AddInt64(&executed, 1) + return nil + })) + } + }(s) + } + submitWG.Wait() + + done := make(chan struct{}) + go func() { pool.Wait(); close(done) }() + select { + case <-done: + case <-time.After(timeoutWait): + t.Fatalf("%s iter %d: Wait() timed out, executed=%d/%d, running=%d", + tc.name, iter, atomic.LoadInt64(&executed), tasks, pool.GetRunningWorkersNum()) + } + + if got := atomic.LoadInt64(&executed); got != tasks { + t.Fatalf("%s iter %d: executed %d, want %d", tc.name, iter, got, tasks) + } + for i, c := range counts { + if c != 1 { + t.Fatalf("%s iter %d: task %d ran %d times (want exactly 1)", tc.name, iter, i, c) + } + } + pool.Close() + } + }) + } +} + +// TestStressNonblockNoLossNoDup checks that in NONBLOCK mode the set of +// accepted tasks is exactly the set that ran, once each. +func TestStressNonblockNoLossNoDup(t *testing.T) { + const ( + tasks = 20000 + submitters = 8 + ) + + for _, shards := range []int{1, 4, 32} { + shards := shards + t.Run("shards"+itoa(shards), func(t *testing.T) { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(16), + agilepool.WithTaskQueueSize(64), + agilepool.WithBlockMode(agilepool.NONBLOCK), + agilepool.WithTaskQueueShards(shards), + )) + defer pool.Close() + + accepted := make([]int32, tasks) + counts := make([]int32, tasks) + var acceptedN int64 + + var wg sync.WaitGroup + for s := 0; s < submitters; s++ { + wg.Add(1) + go func(start int) { + defer wg.Done() + for i := start; i < tasks; i += submitters { + id := i + if pool.TrySubmit(agilepool.TaskFunc(func() error { + atomic.AddInt32(&counts[id], 1) + return nil + })) { + atomic.StoreInt32(&accepted[id], 1) + atomic.AddInt64(&acceptedN, 1) + } + } + }(s) + } + wg.Wait() + pool.Wait() + + if acceptedN == 0 { + t.Fatal("no tasks accepted; workload not exercising the pool") + } + for i := 0; i < tasks; i++ { + acc := atomic.LoadInt32(&accepted[i]) + run := atomic.LoadInt32(&counts[i]) + switch { + case acc == 1 && run != 1: + t.Fatalf("task %d accepted but ran %d times", i, run) + case acc == 0 && run != 0: + t.Fatalf("task %d rejected but ran %d times", i, run) + } + } + }) + } +} + +// TestStressReentrantSubmission makes tasks submit further tasks and verifies +// the total number of executions is exact. +func TestStressReentrantSubmission(t *testing.T) { + for _, shards := range []int{1, 8, 32} { + shards := shards + t.Run("shards"+itoa(shards), func(t *testing.T) { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(64), + agilepool.WithTaskQueueSize(1000), + agilepool.WithTaskQueueShards(shards), + )) + defer pool.Close() + + const roots = 2000 + const childrenPerRoot = 5 + var total int64 + var rootWG sync.WaitGroup + + for r := 0; r < roots; r++ { + rootWG.Add(1) + go func() { + defer rootWG.Done() + for c := 0; c < childrenPerRoot; c++ { + pool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt64(&total, 1) + return nil + })) + } + }() + } + rootWG.Wait() + pool.Wait() + + if want := int64(roots * childrenPerRoot); total != want { + t.Fatalf("executed %d, want %d", total, want) + } + }) + } +} + +// TestStressGracefulShutdown verifies the supported shutdown order: producers +// stop, then Close, then Wait. Wait() must return once in-flight tasks drain, +// and every submitted task must have run. +func TestStressGracefulShutdown(t *testing.T) { + const ( + producers = 8 + perProd = 2000 + ) + for iter := 0; iter < 30; iter++ { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(32), + agilepool.WithTaskQueueSize(128), + )) + + var executed int64 + var wg sync.WaitGroup + for s := 0; s < producers; s++ { + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < perProd; i++ { + pool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt64(&executed, 1) + return nil + })) + } + }() + } + wg.Wait() // all producers have finished submitting + pool.Close() + + done := make(chan struct{}) + go func() { pool.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("iter %d: Wait() hung after Close (executed=%d)", iter, atomic.LoadInt64(&executed)) + } + if want := int64(producers * perProd); atomic.LoadInt64(&executed) != want { + t.Fatalf("iter %d: executed %d, want %d", iter, atomic.LoadInt64(&executed), want) + } + } +} + +// TestStressCloseWhileSubmitting is a KNOWN pre-existing issue, not caused by +// the sharded handoff queue. Pool.submit() calls wg.Add(1) BEFORE checking the +// closed flag, so calling Wait() while producers are still submitting after +// Close() can panic with "sync: WaitGroup misuse: Add called concurrently with +// Wait". The same panic reproduces on the pre-sharding implementation. Skipped +// by default; run explicitly to reproduce: +// +// go test -run TestStressCloseWhileSubmitting -count=1 . +func TestStressCloseWhileSubmitting(t *testing.T) { + t.Skip("known pre-existing WaitGroup misuse: Close()+Wait() while producers still submit") + for iter := 0; iter < 50; iter++ { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(32), + agilepool.WithTaskQueueSize(128), + )) + + var executed int64 + var wg sync.WaitGroup + stop := make(chan struct{}) + for s := 0; s < 8; s++ { + wg.Add(1) + go func() { + defer wg.Done() + for { + select { + case <-stop: + return + default: + } + pool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt64(&executed, 1) + return nil + })) + } + }() + } + + time.Sleep(time.Duration(iter%5) * time.Millisecond) + pool.Close() + + done := make(chan struct{}) + go func() { pool.Wait(); close(done) }() + select { + case <-done: + case <-time.After(10 * time.Second): + t.Fatalf("iter %d: Wait() hung after Close (executed=%d)", iter, atomic.LoadInt64(&executed)) + } + close(stop) + wg.Wait() + } +} + +// TestStressPanicIsolation ensures panicking tasks do not prevent Wait() from +// returning and do not stop other tasks from running. +func TestStressPanicIsolation(t *testing.T) { + for _, shards := range []int{1, 32} { + shards := shards + t.Run("shards"+itoa(shards), func(t *testing.T) { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(16), + agilepool.WithTaskQueueSize(1000), + agilepool.WithTaskQueueShards(shards), + agilepool.WithPanicHandler(func(agilepool.Task, any, []byte) {}), + )) + defer pool.Close() + + const normal = 5000 + var ok int64 + pool.Submit(agilepool.TaskFunc(func() error { panic("boom") })) + for i := 0; i < normal; i++ { + idx := i + pool.Submit(agilepool.TaskFunc(func() error { + if idx%100 == 0 { + panic("boom") + } + atomic.AddInt64(&ok, 1) + return nil + })) + } + + done := make(chan struct{}) + go func() { pool.Wait(); close(done) }() + select { + case <-done: + case <-time.After(20 * time.Second): + t.Fatal("Wait() hung with panicking tasks") + } + + want := int64(normal - normal/100) + if got := atomic.LoadInt64(&ok); got != want { + t.Fatalf("non-panicking tasks executed %d, want %d", got, want) + } + }) + } +} + +// TestStressSubmitCtxCancellation verifies SubmitCtx tasks are never executed +// after their context is already cancelled, and that cancellation while queued +// keeps the pool responsive. +func TestStressSubmitCtxCancellation(t *testing.T) { + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(8), + agilepool.WithTaskQueueSize(1000), + agilepool.WithTaskQueueShards(8), + )) + defer pool.Close() + + var ran int64 + var wg sync.WaitGroup + for i := 0; i < 2000; i++ { + wg.Add(1) + go func(i int) { + defer wg.Done() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() // also covers the i%2==1 path; cancel is idempotent + if i%2 == 0 { + cancel() // cancelled before submit must not run + } + pool.SubmitCtx(ctx, agilepool.TaskFunc(func() error { + if ctx.Err() != nil { + return nil // skipped by contextTask.process + } + atomic.AddInt64(&ran, 1) + return nil + })) + if i%2 == 1 { + cancel() + } + }(i) + } + wg.Wait() + pool.Wait() + // Only the lower bound matters (timing-dependent skips), but it must be + // well under the total to prove cancelled tasks are skipped. + if got := atomic.LoadInt64(&ran); got >= 2000 { + t.Fatalf("cancelled-before-submit tasks appear to have run: ran=%d", got) + } +} + +// TestStressHeavyVolume runs a high-volume no-loss/no-dup check across the two +// most contention-heavy shapes: all-auto shards with a huge queue, and a tiny +// queue that forces the overflow buffer to be exercised. +func TestStressHeavyVolume(t *testing.T) { + if testing.Short() { + t.Skip("skipping heavy volume stress in -short mode") + } + const ( + tasks = 300000 + submitters = 32 + ) + configs := []struct { + name string + config *agilepool.Config + }{ + {"auto_shards_big_queue", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(2048), agilepool.WithTaskQueueSize(100000))}, + {"overflow_small_queue", agilepool.NewConfig(agilepool.WithWorkerNumCapacity(64), agilepool.WithTaskQueueSize(8), agilepool.WithTaskQueueShards(8))}, + } + for _, tc := range configs { + tc := tc + t.Run(tc.name, func(t *testing.T) { + pool := agilepool.NewPool(tc.config) + counts := make([]int32, tasks) + var executed int64 + var wg sync.WaitGroup + for s := 0; s < submitters; s++ { + wg.Add(1) + go func(start int) { + defer wg.Done() + for i := start; i < tasks; i += submitters { + id := i + pool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt32(&counts[id], 1) + atomic.AddInt64(&executed, 1) + return nil + })) + } + }(s) + } + wg.Wait() + + done := make(chan struct{}) + go func() { pool.Wait(); close(done) }() + select { + case <-done: + case <-time.After(60 * time.Second): + t.Fatalf("%s: Wait() hung, executed=%d/%d", tc.name, atomic.LoadInt64(&executed), tasks) + } + + if got := atomic.LoadInt64(&executed); got != tasks { + t.Fatalf("%s: executed %d, want %d", tc.name, got, tasks) + } + for i, c := range counts { + if c != 1 { + t.Fatalf("%s: task %d ran %d times (want 1)", tc.name, i, c) + } + } + pool.Close() + }) + } +} + +func itoa(n int) string { + if n == 0 { + return "0" + } + neg := n < 0 + if neg { + n = -n + } + var b [20]byte + i := len(b) + for n > 0 { + i-- + b[i] = byte('0' + n%10) + n /= 10 + } + if neg { + i-- + b[i] = '-' + } + return string(b[i:]) +} diff --git a/worker.go b/worker.go index 2f7aabc..4e126d9 100644 --- a/worker.go +++ b/worker.go @@ -9,6 +9,9 @@ import ( type worker struct { pool *Pool lastActiveAt time.Time + // cursor is the shard index this worker starts scanning from. A per-worker + // start point spreads channel-lock pressure across the handoff shards. + cursor uint64 } // DatedTime returns the last-active timestamp of the worker. @@ -36,78 +39,68 @@ func (w *worker) run(task Task) { // *worker via Pop and workerPool.Get respectively, causing a data race // on w.lastActiveAt and phantom duplicates in idleWorks. -loop: for { - select { - case task, ok := <-w.pool.taskQueue: - if !ok { - w.pool.logger.Println("taskQueue closed,exiting") - w.pool.addRunningWorkersNum(-1) - atomic.AddInt64(&w.pool.exitCount, 1) - w.pool.workerPool.Put(w) - return - } - - if task == nil { - w.pool.logger.Println("nil task received, exiting") - w.pool.addRunningWorkersNum(-1) - atomic.AddInt64(&w.pool.exitCount, 1) - w.pool.workerPool.Put(w) - return - } + // Drain the handoff-queue shards first. Scanning replaces the old + // double non-blocking channel poll: a task is taken from the first + // shard that has one, so a busy queue is served with a single channel + // lock while lock pressure is spread across all shards. + if t, ok := w.nextShardTask(); ok { w.lastActiveAt = time.Now() - w.runTask(task) + w.runTask(t) + continue + } - default: - // Try the chunked buffer before the second channel check. - // Grab a batch of up to 8 tasks per lock acquisition to - // amortise the mutex overhead across multiple tasks and - // reduce contention with the submission path. - const batchSize = 8 - var batch [batchSize]Task - n := w.pool.taskBuf.PopBatch(batch[:]) + // Fall back to the chunked overflow buffer, grabbing a batch per lock + // acquisition to amortise the mutex overhead and reduce contention + // with the submission path. + const batchSize = 8 + var batch [batchSize]Task + n := w.pool.taskBuf.PopBatch(batch[:]) + if n > 0 { for i := 0; i < n; i++ { w.lastActiveAt = time.Now() w.runTask(batch[i]) } - if n > 0 { - continue - } + continue + } - // Lock-free second check: catch tasks that arrived in the - // tiny window between the two select polls. Submit no longer - // holds p.lock, so serialisation via lock is unnecessary. - // If a task slips through both selects, the scaler will - // spawn workers within scalerPeriod (10ms) to pick it up. - select { - case task, ok := <-w.pool.taskQueue: - if !ok { - w.pool.logger.Println("taskQueue closed,exiting") - w.pool.addRunningWorkersNum(-1) - atomic.AddInt64(&w.pool.exitCount, 1) - w.pool.workerPool.Put(w) - return - } - if task == nil { - w.pool.logger.Println("nil task received, exiting") - w.pool.addRunningWorkersNum(-1) - atomic.AddInt64(&w.pool.exitCount, 1) - w.pool.workerPool.Put(w) - return - } - w.lastActiveAt = time.Now() - w.runTask(task) - default: - // Parking: no task found in second check, worker goes idle. - // Do NOT also put w in workerPool.sync.Pool — see the - // note at the top of run(). - w.pool.addRunningWorkersNum(-1) - atomic.AddInt64(&w.pool.exitCount, 1) - w.pool.addToIdle(w) - break loop + // Park: both the shards and the overflow buffer are empty. Decrement + // the running count so the scaler can spawn replacements when new work + // arrives, then hand the worker back to the idle container. + // Do NOT also put w in workerPool.sync.Pool — see the note above. + w.pool.addRunningWorkersNum(-1) + atomic.AddInt64(&w.pool.exitCount, 1) + w.pool.addToIdle(w) + return + } +} + +// nextShardTask scans the handoff-queue shards, starting from this worker's +// cursor, and returns the first available task. The cursor advances to the +// shard that yielded the task so subsequent polls stay on a warm shard and a +// worker keeps hitting the same shard while under load. The pool never closes +// these channels, so a closed channel is treated as empty to avoid a busy loop +// if that ever changes. +func (w *worker) nextShardTask() (Task, bool) { + p := w.pool + n := uint64(len(p.taskQueues)) + if n == 0 { + return nil, false + } + start := w.cursor + for i := uint64(0); i < n; i++ { + idx := (start + i) & p.shardMask + select { + case t, ok := <-p.taskQueues[idx]: + if !ok || t == nil { + continue } + w.cursor = idx + return t, true + default: } } + return nil, false } func (w *worker) runTask(task Task) {