commit aa93168583db2a86e6dc7a74c7cb33a722336a83 Author: billy Date: Tue Oct 20 20:26:43 2026 +0800 Add dynamic worker pool with bounded queue and concurrency tests diff --git a/README.md b/README.md new file mode 100644 index 0000000..f18f33f --- /dev/null +++ b/README.md @@ -0,0 +1,34 @@ +# worker + +`worker` is a small Go pool for `func()` tasks. Create it with a minimum and +maximum worker count, then submit tasks: + +```go +pool, err := worker.New(2, 8) +if err != nil { + return err +} +defer pool.Close() + +if err := pool.Submit(func() { doWork() }); err != nil { + return err +} +``` + +The pool creates the minimum workers immediately and grows on demand up to the +maximum. Its task queue has a fixed capacity of `maxWorkers/2`; once all workers +and queue slots are occupied, `Submit` waits for room. `Close` stops new +submissions and waits for accepted tasks to finish. A successful `Submit` means +the task was accepted; it does not mean the task has started or finished. With +`maxWorkers == 1`, the queue has zero capacity and submissions wait for a worker. + +After a task finishes, a worker exits immediately only when `total > min` and +the number of *other* idle workers is greater than `min/2`. There is no idle +timeout. A task panic is logged and that worker continues with later tasks. + +Tasks should return normally or panic. Do not call `runtime.Goexit` from a task: +it terminates the worker and may leave `Close` waiting indefinitely. Call +`Close` from outside pool tasks so it does not wait for itself. + +A task may submit another task to the same pool. If all workers block submitting +to a full queue, the caller must resolve that deadlock in its task logic. diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..e98b055 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module git.ybilly.com/billy/worker + +go 1.27.1 diff --git a/pool.go b/pool.go new file mode 100644 index 0000000..82a35bf --- /dev/null +++ b/pool.go @@ -0,0 +1,190 @@ +// Package worker provides a small dynamically sized pool for func tasks. +package worker + +import ( + "errors" + "log" + "sync" +) + +var ( + ErrClosed = errors.New("worker pool is closed") + ErrNilTask = errors.New("worker task is nil") + ErrInvalidConfig = errors.New("invalid worker pool configuration") +) + +// Pool runs submitted tasks using a bounded number of workers. +type Pool struct { + mu sync.Mutex + submitCond *sync.Cond + closeCond *sync.Cond + min int + max int + total int + idle []*worker + queue []func() + queueHead int + queueTail int + queueLen int + queueCap int + closed bool +} + +type worker struct { + inbox chan func() +} + +// New creates a pool with minWorkers persistent workers and at most +// maxWorkers workers. The task queue has capacity maxWorkers/2. +func New(minWorkers, maxWorkers int) (*Pool, error) { + if maxWorkers <= 0 || minWorkers < 0 || minWorkers > maxWorkers { + return nil, ErrInvalidConfig + } + p := &Pool{ + min: minWorkers, + max: maxWorkers, + queue: make([]func(), maxWorkers/2), + queueCap: maxWorkers / 2, + } + p.submitCond = sync.NewCond(&p.mu) + p.closeCond = sync.NewCond(&p.mu) + for range minWorkers { + p.startWorkerLocked(nil, false) + } + return p, nil +} + +// Submit adds task to the pool. It blocks if all workers and the bounded +// queue are full. A successful return means the pool accepted the task. +func (p *Pool) Submit(task func()) error { + if task == nil { + return ErrNilTask + } + p.mu.Lock() + defer p.mu.Unlock() + for { + if p.closed { + return ErrClosed + } + if n := len(p.idle); n > 0 { + w := p.idle[n-1] + p.idle = p.idle[:n-1] + w.inbox <- task + return nil + } + if p.total < p.max { + p.startWorkerLocked(task, true) + return nil + } + if p.queueLen < p.queueCap { + p.enqueueLocked(task) + return nil + } + p.submitCond.Wait() + } +} + +// Close stops accepting tasks and waits for accepted tasks and workers to +// finish. It is safe to call Close concurrently or more than once. +func (p *Pool) Close() { + p.mu.Lock() + p.closed = true + for _, w := range p.idle { + close(w.inbox) + } + p.idle = nil + p.submitCond.Broadcast() + for p.total != 0 { + p.closeCond.Wait() + } + p.mu.Unlock() +} + +// All pool state (worker count, idle set, queue indexes and contents, and +// closed state) is guarded by mu. submitCond is only for Submit capacity +// changes; closeCond is only for Close waiting on all workers to exit. Each +// worker has a one-slot inbox representing an idle worker's reserved handoff. +func (p *Pool) startWorkerLocked(task func(), withTask bool) { + w := &worker{inbox: make(chan func(), 1)} + p.total++ + if withTask { + w.inbox <- task + } else { + p.idle = append(p.idle, w) + } + go p.run(w) +} + +func (p *Pool) run(w *worker) { + for { + task, ok := <-w.inbox + if !ok { + p.mu.Lock() + p.total-- + p.workerExitedLocked() + p.mu.Unlock() + return + } + runTask(task) + + p.mu.Lock() + if p.queueLen > 0 { + next := p.dequeueLocked() + w.inbox <- next + p.submitCond.Signal() + p.mu.Unlock() + continue + } + otherIdle := len(p.idle) + if p.closed || (p.total > p.min && otherIdle > p.min/2) { + p.total-- + if !p.closed { + p.submitCond.Signal() + } else { + p.workerExitedLocked() + } + p.mu.Unlock() + return + } + p.idle = append(p.idle, w) + p.submitCond.Signal() + p.mu.Unlock() + } +} + +// enqueueLocked and dequeueLocked are called with mu held. The queue is a +// fixed-capacity FIFO ring; clearing the dequeued slot releases its closure. +func (p *Pool) enqueueLocked(task func()) { + p.queue[p.queueTail] = task + p.queueTail++ + if p.queueTail == p.queueCap { + p.queueTail = 0 + } + p.queueLen++ +} + +func (p *Pool) dequeueLocked() func() { + task := p.queue[p.queueHead] + p.queue[p.queueHead] = nil + p.queueHead++ + if p.queueHead == p.queueCap { + p.queueHead = 0 + } + p.queueLen-- + return task +} + +func (p *Pool) workerExitedLocked() { + if p.closed && p.total == 0 { + p.closeCond.Broadcast() + } +} + +func runTask(task func()) { + defer func() { + if recovered := recover(); recovered != nil { + log.Printf("worker task panic: %v", recovered) + } + }() + task() +} diff --git a/pool_bench_test.go b/pool_bench_test.go new file mode 100644 index 0000000..4124dd7 --- /dev/null +++ b/pool_bench_test.go @@ -0,0 +1,19 @@ +package worker + +import "testing" + +func BenchmarkQueueRingSteady(b *testing.B) { + p, err := New(0, 6) + if err != nil { + b.Fatal(err) + } + task := func() {} + b.ReportAllocs() + b.ResetTimer() + for i := 0; i < b.N; i++ { + p.mu.Lock() + p.enqueueLocked(task) + p.dequeueLocked() + p.mu.Unlock() + } +} diff --git a/pool_test.go b/pool_test.go new file mode 100644 index 0000000..190c7b1 --- /dev/null +++ b/pool_test.go @@ -0,0 +1,771 @@ +package worker + +import ( + "bytes" + "errors" + "log" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestNewAndNilTask(t *testing.T) { + for _, config := range [][2]int{{-1, 1}, {0, 0}, {2, 1}} { + if _, err := New(config[0], config[1]); !errors.Is(err, ErrInvalidConfig) { + t.Fatalf("New(%d, %d) error = %v", config[0], config[1], err) + } + } + p, err := New(0, 2) + if err != nil { + t.Fatal(err) + } + if err := p.Submit(nil); !errors.Is(err, ErrNilTask) { + t.Fatalf("Submit(nil) error = %v", err) + } + p.Close() +} + +func TestMinimumWorkersStartImmediately(t *testing.T) { + p, err := New(3, 5) + if err != nil { + t.Fatal(err) + } + p.mu.Lock() + if p.total != 3 || len(p.idle) != 3 { + t.Fatalf("initial workers/idle = %d/%d, want 3/3", p.total, len(p.idle)) + } + p.mu.Unlock() + p.Close() +} + +func TestConcurrentSubmitRunsEveryTaskOnceAndRespectsMaximum(t *testing.T) { + const taskCount = 160 + p, err := New(1, 4) + if err != nil { + t.Fatal(err) + } + var running, maxRunning atomic.Int32 + var calls [taskCount]atomic.Int32 + var submitters sync.WaitGroup + start := make(chan struct{}) + for i := range taskCount { + submitters.Add(1) + go func(i int) { + defer submitters.Done() + err := p.Submit(func() { + current := running.Add(1) + for old := maxRunning.Load(); current > old && !maxRunning.CompareAndSwap(old, current); old = maxRunning.Load() { + } + calls[i].Add(1) + <-start + running.Add(-1) + }) + if err != nil { + t.Errorf("Submit(%d): %v", i, err) + } + }(i) + } + doneSubmit := make(chan struct{}) + go func() { submitters.Wait(); close(doneSubmit) }() + close(start) + select { + case <-doneSubmit: + case <-time.After(5 * time.Second): + t.Fatal("concurrent submitters did not finish") + } + p.Close() + if got := maxRunning.Load(); got > 4 { + t.Fatalf("maximum concurrent tasks = %d, want <= 4", got) + } + for i := range taskCount { + if got := calls[i].Load(); got != 1 { + t.Fatalf("task %d ran %d times, want once", i, got) + } + } +} + +func TestQueueCapacityBlocksAndUnblocksSubmit(t *testing.T) { + p, err := New(1, 2) // queue capacity is one + if err != nil { + t.Fatal(err) + } + firstGate, secondGate := make(chan struct{}), make(chan struct{}) + started := make(chan struct{}, 2) + if err := p.Submit(func() { started <- struct{}{}; <-firstGate }); err != nil { + t.Fatal(err) + } + if err := p.Submit(func() { started <- struct{}{}; <-secondGate }); err != nil { + t.Fatal(err) + } + waitSignal(t, started) + waitSignal(t, started) + if err := p.Submit(func() {}); err != nil { + t.Fatal(err) + } + accepted := make(chan error, 1) + go func() { accepted <- p.Submit(func() {}) }() + select { + case err := <-accepted: + t.Fatalf("Submit returned before queue space opened: %v", err) + case <-time.After(30 * time.Millisecond): + } + close(firstGate) + select { + case err := <-accepted: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("blocked Submit did not resume") + } + close(secondGate) + p.Close() +} + +func TestWorkersGrowOnlyToMaximum(t *testing.T) { + p, err := New(1, 3) + if err != nil { + t.Fatal(err) + } + gate := make(chan struct{}) + started := make(chan struct{}, 3) + for range 3 { + if err := p.Submit(func() { started <- struct{}{}; <-gate }); err != nil { + t.Fatal(err) + } + } + for range 3 { + waitSignal(t, started) + } + p.mu.Lock() + if p.total != 3 { + p.mu.Unlock() + t.Fatalf("workers at maximum load = %d, want 3", p.total) + } + p.mu.Unlock() + if err := p.Submit(func() {}); err != nil { // fills the max/2 queue + t.Fatal(err) + } + blocked := make(chan error, 1) + go func() { blocked <- p.Submit(func() {}) }() + select { + case err := <-blocked: + t.Fatalf("Submit beyond maximum workers and queue returned early: %v", err) + case <-time.After(30 * time.Millisecond): + } + close(gate) + select { + case err := <-blocked: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("Submit did not resume after capacity became available") + } + p.Close() +} + +func TestQueueRingWrapsInFIFOOrderAndClearsSlots(t *testing.T) { + p, err := New(0, 4) + if err != nil { + t.Fatal(err) + } + var order []int + f1 := func() { order = append(order, 1) } + f2 := func() { order = append(order, 2) } + f3 := func() { order = append(order, 3) } + p.mu.Lock() + p.enqueueLocked(f1) + p.enqueueLocked(f2) + if p.queueLen != p.queueCap { + p.mu.Unlock() + t.Fatalf("queue length = %d, capacity = %d", p.queueLen, p.queueCap) + } + first := p.dequeueLocked() + p.enqueueLocked(f3) // tail wraps to the beginning of the fixed ring + if p.queueHead != 1 || p.queueTail != 1 || p.queueLen != 2 { + p.mu.Unlock() + t.Fatalf("ring head/tail/len = %d/%d/%d, want 1/1/2", p.queueHead, p.queueTail, p.queueLen) + } + second, third := p.dequeueLocked(), p.dequeueLocked() + if p.queueLen != 0 || p.queue[0] != nil || p.queue[1] != nil { + p.mu.Unlock() + t.Fatal("dequeued task references remain in ring slots") + } + p.mu.Unlock() + first() + second() + third() + if len(order) != 3 || order[0] != 1 || order[1] != 2 || order[2] != 3 { + t.Fatalf("dequeue order = %v, want [1 2 3]", order) + } + p.Close() +} + +func TestSubmitPreservesFIFOAcrossQueueRingWraps(t *testing.T) { + const taskCount = 12 + p, err := New(4, 4) + if err != nil { + t.Fatal(err) + } + otherGate := make(chan struct{}) + otherStarted := make(chan struct{}, 3) + for range 3 { + if err := p.Submit(func() { otherStarted <- struct{}{}; <-otherGate }); err != nil { + t.Fatal(err) + } + } + for range 3 { + waitSignal(t, otherStarted) + } + consumerGate := make(chan struct{}) + consumerStarted := make(chan struct{}) + if err := p.Submit(func() { close(consumerStarted); <-consumerGate }); err != nil { + t.Fatal(err) + } + waitSignal(t, consumerStarted) + + startTask := make(chan int, taskCount) + taskGates := make([]chan struct{}, taskCount) + order := make([]int, 0, taskCount) + newTask := func(i int) func() { + taskGates[i] = make(chan struct{}) + return func() { + order = append(order, i) + startTask <- i + <-taskGates[i] + } + } + if err := p.Submit(newTask(0)); err != nil { + t.Fatal(err) + } + if err := p.Submit(newTask(1)); err != nil { + t.Fatal(err) + } + + submitted := make(chan int, taskCount-2) + submitterDone := make(chan struct{}) + go func() { + defer close(submitterDone) + for i := 2; i < taskCount; i++ { + if err := p.Submit(newTask(i)); err != nil { + submitted <- -1 + return + } + submitted <- i + } + }() + close(consumerGate) + + for i := 0; i < taskCount; i++ { + if got := waitSignal(t, startTask); got != i { + t.Fatalf("task execution order[%d] = %d, want %d", i, got, i) + } + if i+2 < taskCount { + if got := waitSignal(t, submitted); got != i+2 { + t.Fatalf("submitted task after queue slot release = %d, want %d", got, i+2) + } + } + close(taskGates[i]) + } + waitSignal(t, submitterDone) + close(otherGate) + p.Close() + if len(order) != taskCount { + t.Fatalf("executed %d tasks, want %d", len(order), taskCount) + } + for i, got := range order { + if got != i { + t.Fatalf("execution order = %v, want FIFO [0..%d]", order, taskCount-1) + } + } +} + +func TestZeroCapacityQueueDoesNotIndexOrWrap(t *testing.T) { + p, err := New(0, 1) + if err != nil { + t.Fatal(err) + } + if p.queueCap != 0 || len(p.queue) != 0 || p.queueLen != 0 { + t.Fatalf("zero-capacity queue state = cap:%d backing:%d len:%d", p.queueCap, len(p.queue), p.queueLen) + } + p.Close() +} + +func TestMultipleBlockedSubmittersWakeOnEachCapacityChange(t *testing.T) { + p, err := New(2, 2) + if err != nil { + t.Fatal(err) + } + gateA, gateB, gateC, gateD, gateE := make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{}) + started := make(chan struct{}, 5) + for _, gate := range []chan struct{}{gateA, gateB} { + if err := p.Submit(func() { started <- struct{}{}; <-gate }); err != nil { + t.Fatal(err) + } + } + waitSignal(t, started) + waitSignal(t, started) + if err := p.Submit(func() { started <- struct{}{}; <-gateC }); err != nil { // fill queue + t.Fatal(err) + } + results := make(chan error, 2) + attempted := make(chan struct{}, 2) + for _, gate := range []chan struct{}{gateD, gateE} { + go func(gate chan struct{}) { + attempted <- struct{}{} + results <- p.Submit(func() { started <- struct{}{}; <-gate }) + }(gate) + } + waitSignal(t, attempted) + waitSignal(t, attempted) + select { + case err := <-results: + t.Fatalf("Submit returned while the workers and queue were full: %v", err) + case <-time.After(30 * time.Millisecond): + } + // Release each busy worker separately. Each dequeue frees one slot and + // wakes one blocked submitter; the admitted task then occupies that slot. + close(gateA) + if err := waitSignal(t, results); err != nil { + t.Fatal(err) + } + waitSignal(t, started) // queue task C began + close(gateB) + if err := waitSignal(t, results); err != nil { + t.Fatal(err) + } + waitSignal(t, started) // admitted task D began + close(gateC) + waitSignal(t, started) // admitted task E began + close(gateD) + close(gateE) + p.Close() +} + +func TestUnbufferedMaxOneMinZero(t *testing.T) { + p, err := New(0, 1) + if err != nil { + t.Fatal(err) + } + gate := make(chan struct{}) + started := make(chan struct{}, 2) + if err := p.Submit(func() { started <- struct{}{}; <-gate }); err != nil { + t.Fatal(err) + } + waitSignal(t, started) + second := make(chan error, 1) + go func() { second <- p.Submit(func() { started <- struct{}{} }) }() + select { + case err := <-second: + t.Fatalf("unbuffered Submit returned while busy: %v", err) + case <-time.After(30 * time.Millisecond): + } + close(gate) + select { + case err := <-second: + if err != nil { + t.Fatal(err) + } + case <-time.After(2 * time.Second): + t.Fatal("Submit did not resume") + } + waitSignal(t, started) + p.Close() +} + +func TestShrinkUsesStrictIdleThresholdAndKeepsMinimum(t *testing.T) { + p, err := New(2, 4) + if err != nil { + t.Fatal(err) + } + gates := make([]chan struct{}, 4) + started := make(chan struct{}, 4) + for i := range gates { + gates[i] = make(chan struct{}) + } + for _, gate := range gates { + if err := p.Submit(func() { started <- struct{}{}; <-gate }); err != nil { + t.Fatal(err) + } + } + for range 4 { + waitSignal(t, started) + } + close(gates[0]) + waitPoolState(t, p, 4, 1) + close(gates[1]) + waitPoolState(t, p, 4, 2) // other idle == min/2, so no exit + p.mu.Lock() + if p.total != 4 { + t.Fatalf("workers at strict threshold = %d, want 4", p.total) + } + p.mu.Unlock() + close(gates[2]) + waitPoolState(t, p, 3, 2) + close(gates[3]) + waitPoolState(t, p, 2, 2) + p.mu.Lock() + if p.total < p.min { + t.Fatalf("workers %d fell below min %d", p.total, p.min) + } + p.mu.Unlock() + p.Close() +} + +func TestPanicIsLoggedAndWorkerContinues(t *testing.T) { + p, err := New(1, 1) + if err != nil { + t.Fatal(err) + } + var output bytes.Buffer + previousOutput := log.Writer() + log.SetOutput(&output) + defer log.SetOutput(previousOutput) + completed := make(chan struct{}) + if err := p.Submit(func() { panic("expected test panic") }); err != nil { + log.SetOutput(previousOutput) + t.Fatal(err) + } + if err := p.Submit(func() { close(completed) }); err != nil { + log.SetOutput(previousOutput) + t.Fatal(err) + } + waitSignal(t, completed) + p.Close() + log.SetOutput(previousOutput) + if got := output.String(); !strings.Contains(got, "expected test panic") { + t.Fatalf("panic log %q does not contain panic value", got) + } +} + +func TestCloseDrainsTasksRejectsAndWakesBlockedSubmit(t *testing.T) { + p, err := New(1, 2) + if err != nil { + t.Fatal(err) + } + gate := make(chan struct{}) + started := make(chan struct{}, 3) + var completed atomic.Int32 + for range 2 { + if err := p.Submit(func() { started <- struct{}{}; <-gate; completed.Add(1) }); err != nil { + t.Fatal(err) + } + } + waitSignal(t, started) + waitSignal(t, started) + if err := p.Submit(func() { started <- struct{}{}; completed.Add(1) }); err != nil { + t.Fatal(err) + } + const submitterCount = 8 + blockedSubmit := make(chan error, submitterCount) + submitAttempted := make(chan struct{}, submitterCount) + for range submitterCount { + go func() { + submitAttempted <- struct{}{} + blockedSubmit <- p.Submit(func() {}) + }() + } + for range submitterCount { + waitSignal(t, submitAttempted) + } + select { + case err := <-blockedSubmit: + t.Fatalf("Submit did not block on full queue: %v", err) + case <-time.After(30 * time.Millisecond): + } + const closeCount = 4 + closeStarted := make(chan struct{}, closeCount) + closeDone := make(chan struct{}, closeCount) + for range closeCount { + go func() { + closeStarted <- struct{}{} + p.Close() + closeDone <- struct{}{} + }() + } + for range closeCount { + waitSignal(t, closeStarted) + } + for range submitterCount { + select { + case err := <-blockedSubmit: + if !errors.Is(err, ErrClosed) { + t.Fatalf("blocked Submit error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Close did not wake every blocked Submit") + } + } + if err := p.Submit(func() {}); !errors.Is(err, ErrClosed) { + t.Fatalf("Submit after Close error = %v", err) + } + select { + case <-closeDone: + t.Fatal("Close returned before running tasks finished") + case <-time.After(30 * time.Millisecond): + } + close(gate) + waitSignal(t, started) + for range closeCount { + select { + case <-closeDone: + case <-time.After(2 * time.Second): + t.Fatal("Close did not return after accepted tasks completed") + } + } + if got := completed.Load(); got != 3 { + t.Fatalf("completed tasks = %d, want 3", got) + } + p.mu.Lock() + workers := p.total + p.mu.Unlock() + if workers != 0 { + t.Fatalf("workers after concurrent Close = %d, want 0", workers) + } +} + +func TestConcurrentCloseWaitsForRunningTask(t *testing.T) { + p, err := New(1, 2) + if err != nil { + t.Fatal(err) + } + gate, started := make(chan struct{}), make(chan struct{}) + if err := p.Submit(func() { close(started); <-gate }); err != nil { + t.Fatal(err) + } + waitSignal(t, started) + const closeCount = 4 + closeStarted := make(chan struct{}, closeCount) + closeDone := make(chan struct{}, closeCount) + for range 4 { + go func() { + closeStarted <- struct{}{} + p.Close() + closeDone <- struct{}{} + }() + } + for range closeCount { + waitSignal(t, closeStarted) + } + select { + case <-closeDone: + t.Fatal("Close returned while a task was still running") + case <-time.After(30 * time.Millisecond): + } + close(gate) + for range closeCount { + waitSignal(t, closeDone) + } + p.mu.Lock() + if p.total != 0 { + p.mu.Unlock() + t.Fatalf("workers after concurrent Close = %d, want 0", p.total) + } + p.mu.Unlock() +} + +func TestReentrantSubmit(t *testing.T) { + p, err := New(1, 2) + if err != nil { + t.Fatal(err) + } + innerDone, outerDone := make(chan struct{}), make(chan struct{}) + if err := p.Submit(func() { + if err := p.Submit(func() { close(innerDone) }); err != nil { + t.Errorf("reentrant Submit: %v", err) + } + <-innerDone + close(outerDone) + }); err != nil { + t.Fatal(err) + } + waitSignal(t, outerDone) + p.Close() +} + +func TestConcurrentSubmitAndClose(t *testing.T) { + const rounds = 10 + const taskCount = 64 + for round := range rounds { + p, err := New(2, 4) + if err != nil { + t.Fatal(err) + } + var calls [taskCount]atomic.Int32 + results := make([]error, taskCount) + start := make(chan struct{}) + var submitters sync.WaitGroup + for i := range taskCount { + submitters.Add(1) + go func(i int) { + defer submitters.Done() + <-start + results[i] = p.Submit(func() { calls[i].Add(1) }) + }(i) + } + closeDone := make(chan struct{}) + go func() { <-start; p.Close(); close(closeDone) }() + close(start) + allSubmitted := make(chan struct{}) + go func() { submitters.Wait(); close(allSubmitted) }() + select { + case <-allSubmitted: + case <-time.After(3 * time.Second): + t.Fatalf("round %d: Submit calls did not finish", round) + } + select { + case <-closeDone: + case <-time.After(3 * time.Second): + t.Fatalf("round %d: Close did not finish", round) + } + for i, result := range results { + got := calls[i].Load() + switch { + case result == nil && got != 1: + t.Fatalf("round %d task %d: Submit succeeded, execution count = %d", round, i, got) + case errors.Is(result, ErrClosed) && got != 0: + t.Fatalf("round %d task %d: Submit returned ErrClosed, execution count = %d", round, i, got) + case result != nil && !errors.Is(result, ErrClosed): + t.Fatalf("round %d task %d: unexpected Submit error %v", round, i, result) + } + } + p.mu.Lock() + workers := p.total + p.mu.Unlock() + if workers != 0 { + t.Fatalf("round %d: workers after Close = %d, want 0", round, workers) + } + } +} + +func TestSaturatedSubmitAndClose(t *testing.T) { + const rounds = 10 + const taskCount = 64 + for round := range rounds { + p, err := New(4, 4) + if err != nil { + t.Fatal(err) + } + gate := make(chan struct{}) + workersStarted := make(chan struct{}, 4) + for range 4 { + if err := p.Submit(func() { workersStarted <- struct{}{}; <-gate }); err != nil { + t.Fatal(err) + } + } + for range 4 { + waitSignal(t, workersStarted) + } + var calls [taskCount]atomic.Int32 + type submitResult struct { + index int + err error + } + results := make(chan submitResult, taskCount) + start := make(chan struct{}) + attempted := make(chan struct{}, taskCount) + var submitters sync.WaitGroup + for i := range taskCount { + submitters.Add(1) + go func(i int) { + defer submitters.Done() + attempted <- struct{}{} + <-start + results <- submitResult{index: i, err: p.Submit(func() { calls[i].Add(1) })} + }(i) + } + for range taskCount { + waitSignal(t, attempted) + } + close(start) + seen := make([]bool, taskCount) + acceptedTasks := make([]bool, taskCount) + accepted := 0 + for accepted < p.queueCap { + result := waitSignal(t, results) + if seen[result.index] { + t.Fatalf("round %d task %d: duplicate Submit result", round, result.index) + } + seen[result.index] = true + if result.err != nil { + t.Fatalf("round %d: Submit returned %v before Close", round, result.err) + } + acceptedTasks[result.index] = true + accepted++ + } + closeDone := make(chan struct{}) + go func() { p.Close(); close(closeDone) }() + for received := accepted; received < taskCount; received++ { + result := waitSignal(t, results) + if seen[result.index] { + t.Fatalf("round %d task %d: duplicate Submit result", round, result.index) + } + seen[result.index] = true + if !errors.Is(result.err, ErrClosed) { + t.Fatalf("round %d task %d: Submit error = %v, want ErrClosed", round, result.index, result.err) + } + } + submitters.Wait() + select { + case <-closeDone: + t.Fatalf("round %d: Close returned before blocked worker tasks finished", round) + case <-time.After(10 * time.Millisecond): + } + close(gate) + select { + case <-closeDone: + case <-time.After(3 * time.Second): + t.Fatalf("round %d: Close did not finish", round) + } + for i := range taskCount { + if !seen[i] { + t.Fatalf("round %d task %d: missing Submit result", round, i) + } + got := calls[i].Load() + switch { + case acceptedTasks[i] && got != 1: + t.Fatalf("round %d task %d: Submit succeeded, execution count = %d", round, i, got) + case !acceptedTasks[i] && got != 0: + t.Fatalf("round %d task %d: Submit returned ErrClosed, execution count = %d", round, i, got) + } + } + p.mu.Lock() + workers := p.total + p.mu.Unlock() + if workers != 0 { + t.Fatalf("round %d: workers after Close = %d, want 0", round, workers) + } + } +} + +func waitPoolState(t *testing.T, p *Pool, total, idle int) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + p.mu.Lock() + gotTotal, gotIdle := p.total, len(p.idle) + p.mu.Unlock() + if gotTotal == total && gotIdle == idle { + return + } + time.Sleep(time.Millisecond) + } + p.mu.Lock() + defer p.mu.Unlock() + t.Fatalf("pool state workers/idle = %d/%d, want %d/%d", p.total, len(p.idle), total, idle) +} + +func waitSignal[T any](t *testing.T, ch <-chan T) T { + t.Helper() + select { + case value := <-ch: + return value + case <-time.After(2 * time.Second): + var zero T + t.Fatal("timed out waiting for signal") + return zero + } +}