Files
worker/pool_test.go
2026-10-25 21:11:49 +08:00

948 lines
27 KiB
Go

package worker
import (
"bytes"
"errors"
"log"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestNewValidationAndNilTask(t *testing.T) {
for _, cfg := range [][2]int{{-1, 1}, {0, 0}, {2, 1}} {
if _, err := New(cfg[0], cfg[1]); !errors.Is(err, ErrInvalidConfig) {
t.Fatalf("New(%d,%d) error=%v, want ErrInvalidConfig", cfg[0], cfg[1], err)
}
}
p, err := New(0, 2)
if err != nil {
t.Fatal(err)
}
if cap(p.submit) != defaultSubmitBuffer {
t.Fatalf("default submit capacity=%d, want %d", cap(p.submit), defaultSubmitBuffer)
}
if err := p.Submit(nil); !errors.Is(err, ErrNilTask) {
t.Fatalf("Submit(nil) error=%v, want ErrNilTask", err)
}
p.Close()
if err := p.Submit(nil); !errors.Is(err, ErrNilTask) {
t.Fatalf("Submit(nil) after Close error=%v, want nil-task priority", err)
}
}
func TestSubmitBufferOptions(t *testing.T) {
for _, test := range []struct {
name string
options []Option
want int
invalid bool
}{
{name: "default", want: 128},
{name: "custom", options: []Option{WithSubmitBuffer(7)}, want: 7},
{name: "zero", options: []Option{WithSubmitBuffer(0)}, want: 0},
{name: "negative", options: []Option{WithSubmitBuffer(-1)}, invalid: true},
} {
t.Run(test.name, func(t *testing.T) {
p, err := New(0, 2, test.options...)
if test.invalid {
if !errors.Is(err, ErrInvalidConfig) {
t.Fatalf("New error=%v, want ErrInvalidConfig", err)
}
return
}
if err != nil {
t.Fatal(err)
}
if got := cap(p.submit); got != test.want {
t.Fatalf("submit buffer=%d, want %d", got, test.want)
}
p.Close()
})
}
}
func TestSchedulerStartsMinimumBeforeRun(t *testing.T) {
p := newTestPool(3, 5)
s := newScheduler(p)
for range p.min {
s.startWorker(nil, false)
}
if s.total != 3 || len(s.idle) != 3 {
t.Fatalf("initial scheduler total/idle=%d/%d, want 3/3", s.total, len(s.idle))
}
// The actors' minimum workers are real goroutines; stop this fixture cleanly.
s.closed = true
close(p.stopping)
s.pump()
for s.total > 0 {
event := waitEvent(t, p.events)
s.handleEvent(event)
}
}
func TestConcurrentSubmitRunsExactlyOnceAndNeverExceedsMaximum(t *testing.T) {
const count, maxWorkers = 120, 4
p, err := New(1, maxWorkers)
if err != nil {
t.Fatal(err)
}
var calls [count]atomic.Int32
var active, maxActive atomic.Int32
var submitters sync.WaitGroup
start := make(chan struct{})
results := make([]error, count)
for i := range count {
submitters.Add(1)
go func(i int) {
defer submitters.Done()
<-start
results[i] = p.Submit(func() {
n := active.Add(1)
for old := maxActive.Load(); n > old && !maxActive.CompareAndSwap(old, n); old = maxActive.Load() {
}
calls[i].Add(1)
active.Add(-1)
})
}(i)
}
close(start)
waitGroup(t, &submitters)
p.Close()
for i, result := range results {
if result != nil {
t.Fatalf("Submit(%d) error=%v, want accepted", i, result)
}
if got := calls[i].Load(); got != 1 {
t.Fatalf("task %d accepted but ran %d times", i, got)
}
}
if n := maxActive.Load(); n > maxWorkers {
t.Fatalf("maximum concurrent tasks=%d, worker limit=%d", n, maxWorkers)
}
}
func TestGrowsToMaximumBeforeUsingBoundedQueue(t *testing.T) {
p, err := New(0, 3, WithSubmitBuffer(0)) // one cached task, no ingress buffer
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)
}
if err := p.Submit(func() {}); err != nil {
t.Fatal(err)
}
blocked := make(chan error, 1)
attempted := make(chan struct{})
go func() { close(attempted); blocked <- p.Submit(func() {}) }()
waitSignal(t, attempted)
assertBlocked(t, blocked, "Submit beyond max workers and queue capacity")
close(gate)
if err := waitSignal(t, blocked); err != nil {
t.Fatalf("Submit after capacity release: %v", err)
}
p.Close()
}
func TestQueueCapacityAndMultipleBlockedSubmittersWakeIncrementally(t *testing.T) {
p, err := New(0, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer
if err != nil {
t.Fatal(err)
}
gateA, gateB := make(chan struct{}), make(chan struct{})
started := make(chan struct{}, 2)
if err := p.Submit(func() { started <- struct{}{}; <-gateA }); err != nil {
t.Fatal(err)
}
if err := p.Submit(func() { started <- struct{}{}; <-gateB }); err != nil {
t.Fatal(err)
}
waitSignal(t, started)
waitSignal(t, started)
baseStarted, baseGate := make(chan struct{}), make(chan struct{})
if err := p.Submit(func() { close(baseStarted); <-baseGate }); err != nil {
t.Fatal(err)
}
type submitResult struct {
index int
err error
}
results := make(chan submitResult, 3)
attempted := make(chan struct{}, 3)
blockedGates := []chan struct{}{make(chan struct{}), make(chan struct{}), make(chan struct{})}
blockedStarted := make(chan int, 3)
for i := range blockedGates {
index := i
go func() {
attempted <- struct{}{}
results <- submitResult{index: index, err: p.Submit(func() { blockedStarted <- index; <-blockedGates[index] })}
}()
}
for range 3 {
waitSignal(t, attempted)
}
for range 3 {
assertBlocked(t, results, "Submitter passed a full queue")
}
close(gateA)
waitSignal(t, baseStarted)
first := waitSignal(t, results)
if first.err != nil {
t.Fatalf("first released Submit: %v", first.err)
}
assertBlocked(t, results, "more than one Submit used one freed queue slot")
close(baseGate)
firstRun := waitSignal(t, blockedStarted)
second := waitSignal(t, results)
if second.err != nil {
t.Fatalf("second released Submit: %v", second.err)
}
assertBlocked(t, results, "third Submit used a still-full queue")
close(blockedGates[firstRun])
secondRun := waitSignal(t, blockedStarted)
close(blockedGates[secondRun])
third := waitSignal(t, results)
if third.err != nil {
t.Fatalf("third released Submit: %v", third.err)
}
thirdRun := waitSignal(t, blockedStarted)
close(blockedGates[thirdRun])
close(gateB)
p.Close()
}
func TestZeroQueueMinZeroAcceptsAfterWorkerFinishes(t *testing.T) {
p, err := New(0, 1, WithSubmitBuffer(0))
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)
second := make(chan error, 1)
attempted := make(chan struct{})
var secondRuns atomic.Int32
go func() {
close(attempted)
second <- p.Submit(func() { secondRuns.Add(1) })
}()
waitSignal(t, attempted)
assertBlocked(t, second, "max=1 accepted a task while its worker was busy")
close(gate)
if err := waitSignal(t, second); err != nil {
t.Fatal(err)
}
p.Close()
if got := secondRuns.Load(); got != 1 {
t.Fatalf("second task ran %d times, want once", got)
}
}
func TestQueueTasksAreDispatchedFIFO(t *testing.T) {
p, err := New(4, 4, WithSubmitBuffer(0)) // two cached tasks
if err != nil {
t.Fatal(err)
}
gates := []chan struct{}{make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{})}
startedWorkers := make(chan struct{}, 4)
for _, gate := range gates {
if err := p.Submit(func() { startedWorkers <- struct{}{}; <-gate }); err != nil {
t.Fatal(err)
}
}
for range 4 {
waitSignal(t, startedWorkers)
}
queueStarted := make(chan int, 2)
queueGate := make(chan struct{})
for i := 1; i <= 2; i++ {
index := i
if err := p.Submit(func() { queueStarted <- index; <-queueGate }); err != nil {
t.Fatal(err)
}
}
close(gates[0])
if got := waitSignal(t, queueStarted); got != 1 {
t.Fatalf("first queued task started=%d, want 1", got)
}
close(gates[1])
if got := waitSignal(t, queueStarted); got != 2 {
t.Fatalf("second queued task started=%d, want 2", got)
}
close(queueGate)
close(gates[2])
close(gates[3])
p.Close()
}
func TestStrictShrinkThresholdAndMinimumProtection(t *testing.T) {
// Exactly min/2 other idle workers is not enough to shrink.
s := newTestScheduler(2, 5)
current := s.addFakeWorker(workerBusy)
s.addFakeWorker(workerIdle)
s.addFakeWorker(workerBusy)
s.taskFinished(current)
if s.workers[current] != workerIdle || s.retiring != 0 {
t.Fatal("worker shrank when other idle count equaled min/2")
}
// Strictly more than min/2 permits shrink, but retiring workers are
// excluded from the effective total so repeated completions stop at min.
s = newTestScheduler(2, 6)
first := s.addFakeWorker(workerBusy)
s.addFakeWorker(workerIdle)
s.addFakeWorker(workerIdle)
second := s.addFakeWorker(workerBusy)
thirdBusy := s.addFakeWorker(workerBusy)
s.taskFinished(first)
if s.workers[first] != workerRetiring || s.retiring != 1 {
t.Fatal("worker did not shrink above the strict idle threshold")
}
s.taskFinished(second)
if s.workers[second] != workerRetiring || s.retiring != 2 {
t.Fatal("second worker should shrink while effective total remains above min")
}
s.taskFinished(thirdBusy)
if s.workers[thirdBusy] != workerRetiring || s.retiring != 3 || s.total-s.retiring != 2 {
t.Fatal("normal shrink lowered effective worker count below min")
}
lastWorker := <-s.idle
s.workers[lastWorker] = workerBusy
s.taskFinished(lastWorker)
if s.workers[lastWorker] != workerIdle || s.total-s.retiring != 2 {
t.Fatal("worker shrank below min after retiring workers were excluded")
}
}
func TestTaskPanicsAreLoggedAndSameWorkerContinues(t *testing.T) {
var output bytes.Buffer
previous := log.Writer()
log.SetOutput(&output)
defer log.SetOutput(previous)
w := &worker{inbox: make(chan func()), stop: make(chan struct{}), dead: make(chan struct{})}
events := make(chan workerEvent, 8)
go runWorker(w, events)
for _, task := range []func(){
func() { panic("task panic marker") },
func() { panic(nil) },
func() { defer func() { panic("deferred panic marker") }() },
func() {},
} {
w.inbox <- task
event := waitEvent(t, events)
if event.kind != taskFinished || event.worker != w {
t.Fatalf("event=%+v, want completion from same worker", event)
}
}
if !strings.Contains(output.String(), "task panic marker") || !strings.Contains(output.String(), "deferred panic marker") {
t.Fatalf("task panic logs missing markers: %q", output.String())
}
close(w.stop)
if event := waitEvent(t, events); event.kind != workerStopped || event.worker != w {
t.Fatalf("stop event=%+v, want normal stop from same worker", event)
}
}
type panicWriter struct {
always bool
calls atomic.Int32
}
func (w *panicWriter) Write(p []byte) (int, error) {
if w.calls.Add(1) > 0 && (w.always || w.calls.Load() == 1) {
panic("logger writer panic")
}
return len(p), nil
}
func TestLoggerWriterPanicDoesNotKillTaskWorker(t *testing.T) {
for _, always := range []bool{false, true} {
t.Run(map[bool]string{false: "once", true: "always"}[always], func(t *testing.T) {
previous := log.Writer()
writer := &panicWriter{always: always}
log.SetOutput(writer)
defer log.SetOutput(previous)
w := &worker{inbox: make(chan func()), stop: make(chan struct{}), dead: make(chan struct{})}
events := make(chan workerEvent, 4)
go runWorker(w, events)
w.inbox <- func() { panic("first task") }
if event := waitEvent(t, events); event.kind != taskFinished || event.worker != w {
t.Fatalf("first event=%+v", event)
}
w.inbox <- func() { panic("second task") }
if event := waitEvent(t, events); event.kind != taskFinished || event.worker != w {
t.Fatalf("second event=%+v", event)
}
w.inbox <- func() {}
if event := waitEvent(t, events); event.kind != taskFinished || event.worker != w {
t.Fatalf("final event=%+v", event)
}
if writer.calls.Load() == 0 {
t.Fatal("logger writer was not called")
}
close(w.stop)
waitEvent(t, events)
})
}
}
func TestWorkerOuterPanicRecoveryReportsExit(t *testing.T) {
var output bytes.Buffer
previous := log.Writer()
log.SetOutput(&output)
defer log.SetOutput(previous)
w := &worker{inbox: make(chan func()), stop: make(chan struct{}), dead: make(chan struct{})}
close(w.inbox) // test-only fault injection into the worker receive loop
events := make(chan workerEvent, 1)
go runWorker(w, events)
event := waitEvent(t, events)
if event.kind != workerPanicked || event.worker != w {
t.Fatalf("event=%+v, want worker panic exit", event)
}
if !strings.Contains(output.String(), "worker panic: worker inbox closed") {
t.Fatalf("worker panic was not logged: %q", output.String())
}
}
func TestAbnormalExitStateMachineKeepsQueueAndDoesNotRefillMinimum(t *testing.T) {
s := newTestScheduler(2, 3)
w := s.addFakeWorker(workerBusy)
survivor := s.addFakeWorker(workerBusy)
var runs atomic.Int32
done := make(chan struct{})
s.queue <- func() { runs.Add(1); close(done) }
s.handleEvent(workerEvent{worker: w, kind: workerPanicked})
waitSignal(t, done)
event := waitEvent(t, s.pool.events)
s.handleEvent(event)
if runs.Load() != 1 || len(s.queue) != 0 {
t.Fatalf("queue task runs=%d queue=%d, want once and empty", runs.Load(), len(s.queue))
}
if s.total != 2 || len(s.idle) != 1 || s.workers[survivor] != workerBusy {
t.Fatalf("post-replacement total/idle=%d/%d, survivor=%d", s.total, len(s.idle), s.workers[survivor])
}
// Retire the real replacement and account the fake surviving worker.
s.closed = true
close(s.pool.stopping)
s.pump()
s.handleEvent(waitEvent(t, s.pool.events))
s.workerExited(survivor, false)
s.finishCloseIfReady()
}
func TestAbnormalExitWithoutQueueDoesNotRefillButLaterSubmitWorks(t *testing.T) {
s := newTestScheduler(2, 3)
w := s.addFakeWorker(workerBusy)
s.handleEvent(workerEvent{worker: w, kind: workerPanicked})
if s.total != 0 || s.closed {
t.Fatalf("empty abnormal exit total/closed=%d/%v, want 0/open", s.total, s.closed)
}
done := make(chan struct{})
s.accept(func() { close(done) })
waitSignal(t, done)
s.handleEvent(waitEvent(t, s.pool.events))
if s.total != 1 || len(s.idle) != 1 {
t.Fatalf("later Submit total/idle=%d/%d, want one worker without min refill", s.total, len(s.idle))
}
s.closed = true
close(s.pool.stopping)
s.pump()
s.handleEvent(waitEvent(t, s.pool.events))
}
func TestAbnormalExitDuringCloseStillDrainsQueue(t *testing.T) {
s := newTestScheduler(0, 2)
w := s.addFakeWorker(workerBusy)
s.closed = true
close(s.pool.stopping)
started, gate, done := make(chan struct{}), make(chan struct{}), make(chan struct{})
s.queue <- func() { close(started); <-gate; close(done) }
s.handleEvent(workerEvent{worker: w, kind: workerPanicked})
waitSignal(t, started)
select {
case <-s.pool.done:
t.Fatal("Close completed while replacement task was running")
default:
}
close(gate)
waitSignal(t, done)
s.handleEvent(waitEvent(t, s.pool.events)) // task completion retires worker
s.handleEvent(waitEvent(t, s.pool.events)) // retirement confirmation
select {
case <-s.pool.done:
default:
t.Fatal("Close did not finish after accepted queue drained")
}
if s.total != 0 || len(s.queue) != 0 {
t.Fatalf("final total/queue=%d/%d, want 0/0", s.total, len(s.queue))
}
}
func TestBusyWorkerPanicPreservesQueuedTaskAndOpenPool(t *testing.T) {
p := newTestPool(0, 2)
s := newScheduler(p)
w1 := s.startWorker(nil, false)
_ = s.startWorker(nil, false)
started := make(chan string, 2)
gates := map[string]chan struct{}{"A": make(chan struct{}), "B": make(chan struct{})}
var runs [4]atomic.Int32
s.accept(func() { started <- "A"; <-gates["A"]; runs[0].Add(1) })
s.accept(func() { started <- "B"; <-gates["B"]; runs[1].Add(1) })
waitSignal(t, started)
waitSignal(t, started)
queuedDone := make(chan struct{})
s.queue <- func() { runs[2].Add(1); close(queuedDone) }
go s.run()
close(w1.inbox) // worker A fails after its current task reports completion
close(gates["A"])
waitSignal(t, queuedDone)
close(gates["B"])
postFailureDone := make(chan struct{})
if err := p.Submit(func() { runs[3].Add(1); close(postFailureDone) }); err != nil {
t.Fatalf("Submit after worker failure: %v", err)
}
waitSignal(t, postFailureDone)
if pClosed(p) {
t.Fatal("worker panic closed the pool")
}
p.Close()
for i := range runs {
if got := runs[i].Load(); got != 1 {
t.Errorf("task %d ran %d times, want once", i, got)
}
}
}
func TestIdleWorkerPanicDuringSubmitDoesNotLoseTask(t *testing.T) {
for range 8 {
p := newTestPool(0, 1)
s := newScheduler(p)
w := s.startWorker(nil, false)
close(w.inbox)
go s.run()
var runs atomic.Int32
done := make(chan struct{})
if err := p.Submit(func() { runs.Add(1); close(done) }); err != nil {
t.Fatalf("Submit during idle worker exit: %v", err)
}
waitSignal(t, done)
if pClosed(p) {
t.Fatal("worker panic closed the pool")
}
p.Close()
if got := runs.Load(); got != 1 {
t.Fatalf("task ran %d times, want once", got)
}
}
}
func TestBusyWorkerPanicDuringCloseDrainsQueuedTask(t *testing.T) {
p := newTestPool(0, 2)
s := newScheduler(p)
w1 := s.startWorker(nil, false)
_ = s.startWorker(nil, false)
started := make(chan string, 2)
gateA, gateB := make(chan struct{}), make(chan struct{})
var queuedRuns atomic.Int32
s.accept(func() { started <- "A"; <-gateA })
s.accept(func() { started <- "B"; <-gateB })
waitSignal(t, started)
waitSignal(t, started)
queuedDone := make(chan struct{})
s.queue <- func() { queuedRuns.Add(1); close(queuedDone) }
go s.run()
closeStarted := make(chan struct{})
go func() { p.Close(); close(closeStarted) }()
waitSignal(t, p.stopping)
close(w1.inbox)
close(gateA)
waitSignal(t, queuedDone)
select {
case <-closeStarted:
t.Fatal("Close returned before remaining worker finished")
default:
}
close(gateB)
waitSignal(t, closeStarted)
if got := queuedRuns.Load(); got != 1 {
t.Fatalf("queued task ran %d times, want once", got)
}
}
type gatedLogWriter struct {
entered chan struct{}
release chan struct{}
}
func (w *gatedLogWriter) Write(p []byte) (int, error) {
w.entered <- struct{}{}
<-w.release
return len(p), nil
}
func TestCloseWaitsForWorkerPanicLogsAndTerminalEvents(t *testing.T) {
writer := &gatedLogWriter{entered: make(chan struct{}, 2), release: make(chan struct{}, 2)}
previous := log.Writer()
log.SetOutput(writer)
defer log.SetOutput(previous)
p := newTestPool(0, 2)
s := newScheduler(p)
w1 := s.startWorker(nil, false)
w2 := s.startWorker(nil, false)
started := make(chan string, 2)
gateA, gateB, gateC := make(chan struct{}), make(chan struct{}), make(chan struct{})
var runs [3]atomic.Int32
s.accept(func() { started <- "A"; <-gateA; runs[0].Add(1) })
s.accept(func() { started <- "B"; <-gateB; runs[1].Add(1) })
waitSignal(t, started)
waitSignal(t, started)
cStarted, cDone := make(chan struct{}), make(chan struct{})
s.queue <- func() { runs[2].Add(1); close(cStarted); <-gateC; close(cDone) }
go s.run()
close(w1.inbox)
close(gateA)
waitSignal(t, writer.entered)
close(gateB)
waitSignal(t, cStarted) // the surviving worker continues queued work
close(w2.inbox)
close(gateC)
waitSignal(t, cDone)
closeDone := make(chan struct{})
go func() { p.Close(); close(closeDone) }()
waitSignal(t, p.stopping)
assertBlocked(t, closeDone, "Close returned before delayed panic logging and terminal events")
writer.release <- struct{}{}
waitSignal(t, writer.entered)
assertBlocked(t, closeDone, "Close returned before second worker terminal event")
writer.release <- struct{}{}
waitSignal(t, closeDone)
for i := range runs {
if got := runs[i].Load(); got != 1 {
t.Errorf("task %d ran %d times, want once", i, got)
}
}
}
func TestAcceptedPendingHandoffSurvivesDelayedExitAndClose(t *testing.T) {
writer := &gatedLogWriter{entered: make(chan struct{}, 1), release: make(chan struct{}, 1)}
previous := log.Writer()
log.SetOutput(writer)
defer log.SetOutput(previous)
p := newTestPool(0, 1) // zero queue capacity
s := newScheduler(p)
w := s.startWorker(nil, false)
close(w.inbox)
waitSignal(t, writer.entered)
var runs atomic.Int32
task := func() { runs.Add(1) }
if s.accept(task) {
t.Fatal("handoff to exited idle worker unexpectedly succeeded")
}
// This is the task already received by the actor: keep its idle-worker
// handoff slot reserved while the old worker finishes delayed logging.
s.pending = task
s.hasPending = true
go s.run()
closeDone := make(chan struct{})
go func() { p.Close(); close(closeDone) }()
waitSignal(t, p.stopping)
assertBlocked(t, closeDone, "Close returned before pending accepted task was dispatched")
writer.release <- struct{}{}
waitSignal(t, closeDone)
if got := runs.Load(); got != 1 {
t.Fatalf("pending task ran %d times, want once", got)
}
}
func TestCloseDrainsAcceptedTasksAndWakesBlockedSubmitters(t *testing.T) {
p, err := New(2, 2, WithSubmitBuffer(0)) // queue capacity one; zero ingress buffer
if err != nil {
t.Fatal(err)
}
gates := []chan struct{}{make(chan struct{}), make(chan struct{})}
started := make(chan struct{}, 2)
var completed atomic.Int32
for _, gate := range gates {
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() { completed.Add(1) }); err != nil {
t.Fatal(err)
}
blocked := make(chan error, 1)
attempted := make(chan struct{})
go func() { close(attempted); blocked <- p.Submit(func() { completed.Add(1000) }) }()
waitSignal(t, attempted)
assertBlocked(t, blocked, "saturated Submit returned before Close")
const closeCount = 4
closes := make(chan struct{}, closeCount)
closeStarted := make(chan struct{}, closeCount)
for range closeCount {
go func() { closeStarted <- struct{}{}; p.Close(); closes <- struct{}{} }()
}
for range closeCount {
waitSignal(t, closeStarted)
}
waitSignal(t, p.stopping)
if err := waitSignal(t, blocked); !errors.Is(err, ErrClosed) {
t.Fatalf("blocked Submit error=%v, want ErrClosed", err)
}
select {
case <-closes:
t.Fatal("Close returned while accepted gated work remained")
default:
}
close(gates[0])
close(gates[1])
for range closeCount {
waitSignal(t, closes)
}
if got := completed.Load(); got != 3 {
t.Fatalf("accepted tasks completed=%d, want 3", got)
}
if err := p.Submit(func() {}); !errors.Is(err, ErrClosed) {
t.Fatalf("Submit after Close=%v, want ErrClosed", err)
}
}
func TestCloseDrainsTasksAcceptedIntoSubmitBuffer(t *testing.T) {
const bufferedTasks = 2
p, err := New(1, 1, WithSubmitBuffer(bufferedTasks))
if err != nil {
t.Fatal(err)
}
gate := make(chan struct{})
started := make(chan int, bufferedTasks+1)
var calls [bufferedTasks + 1]atomic.Int32
if err := p.Submit(func() { started <- 0; <-gate; calls[0].Add(1) }); err != nil {
t.Fatal(err)
}
if got := waitSignal(t, started); got != 0 {
t.Fatalf("first task index=%d", got)
}
for i := 1; i <= bufferedTasks; i++ {
index := i
if err := p.Submit(func() { started <- index; calls[index].Add(1) }); err != nil {
t.Fatalf("buffered Submit(%d): %v", i, err)
}
}
closed := make(chan struct{})
go func() { p.Close(); close(closed) }()
waitSignal(t, p.stopping)
select {
case <-closed:
t.Fatal("Close returned while accepted work remained")
default:
}
close(gate)
waitSignal(t, closed)
for i := range calls {
if got := calls[i].Load(); got != 1 {
t.Errorf("task %d ran %d times, want once", i, got)
}
}
}
func TestFullSubmitBufferBlocksUntilSchedulerConsumes(t *testing.T) {
p, err := New(1, 1, WithSubmitBuffer(1))
if err != nil {
t.Fatal(err)
}
release := make(chan struct{})
started := make(chan int, 3)
var calls [3]atomic.Int32
if err := p.Submit(func() { started <- 0; <-release; calls[0].Add(1) }); err != nil {
t.Fatal(err)
}
waitSignal(t, started)
if err := p.Submit(func() { started <- 1; calls[1].Add(1) }); err != nil {
t.Fatal(err)
}
third := make(chan error, 1)
go func() { third <- p.Submit(func() { started <- 2; calls[2].Add(1) }) }()
assertBlocked(t, third, "Submit returned while its ingress buffer was full and worker was busy")
close(release)
if err := waitSignal(t, third); err != nil {
t.Fatalf("Submit after scheduler consumed ingress=%v", err)
}
p.Close()
for i := range calls {
if got := calls[i].Load(); got != 1 {
t.Errorf("task %d ran %d times, want once", i, got)
}
}
}
func TestCloseDrainsBufferedSubmissionsAfterWorkerPanic(t *testing.T) {
p := newTestPool(0, 1)
p.submit = make(chan func(), 3)
s := newScheduler(p)
w := s.startWorker(nil, false)
close(w.inbox)
var calls [3]atomic.Int32
for i := range calls {
index := i
if err := p.Submit(func() { calls[index].Add(1) }); err != nil {
t.Fatalf("buffered Submit(%d): %v", i, err)
}
}
go s.run()
closed := make(chan struct{})
go func() { p.Close(); close(closed) }()
waitSignal(t, closed)
for i := range calls {
if got := calls[i].Load(); got != 1 {
t.Errorf("buffered task %d ran %d times, want once", i, got)
}
}
}
func TestConcurrentSubmitAndCloseMatchesEveryResult(t *testing.T) {
const count = 100
p, err := New(2, 4)
if err != nil {
t.Fatal(err)
}
var calls [count]atomic.Int32
results := make([]error, count)
start := make(chan struct{})
var submits sync.WaitGroup
for i := range count {
submits.Add(1)
go func(i int) {
defer submits.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)
waitGroup(t, &submits)
waitSignal(t, closeDone)
for i, err := range results {
switch {
case err == nil && calls[i].Load() != 1:
t.Errorf("task %d accepted but ran %d times", i, calls[i].Load())
case errors.Is(err, ErrClosed) && calls[i].Load() != 0:
t.Errorf("task %d rejected but ran %d times", i, calls[i].Load())
case err != nil && !errors.Is(err, ErrClosed):
t.Errorf("task %d unexpected error: %v", i, err)
}
}
}
func TestReentrantSubmit(t *testing.T) {
p, err := New(1, 2)
if err != nil {
t.Fatal(err)
}
inner, outer := make(chan struct{}), make(chan struct{})
if err := p.Submit(func() {
if err := p.Submit(func() { close(inner) }); err != nil {
t.Errorf("nested Submit: %v", err)
}
<-inner
close(outer)
}); err != nil {
t.Fatal(err)
}
waitSignal(t, outer)
p.Close()
}
func newTestPool(min, max int) *Pool {
return &Pool{
submit: make(chan func()), closeReq: make(chan struct{}), stopping: make(chan struct{}),
done: make(chan struct{}), events: make(chan workerEvent, max), min: min, max: max, queueCap: max / 2,
}
}
func newTestScheduler(min, max int) *scheduler { return newScheduler(newTestPool(min, max)) }
func pClosed(p *Pool) bool {
select {
case <-p.stopping:
return true
default:
return false
}
}
func (s *scheduler) addFakeWorker(state workerState) *worker {
w := &worker{inbox: make(chan func(), 1), stop: make(chan struct{}), dead: make(chan struct{})}
s.workers[w] = state
s.total++
if state == workerIdle {
s.idle <- w
}
if state == workerRetiring {
s.retiring++
}
return w
}
func waitEvent(t *testing.T, events <-chan workerEvent) workerEvent {
t.Helper()
select {
case event := <-events:
return event
case <-time.After(3 * time.Second):
t.Fatal("timed out waiting for worker event")
return workerEvent{}
}
}
func waitSignal[T any](t *testing.T, ch <-chan T) T {
t.Helper()
select {
case value := <-ch:
return value
case <-time.After(3 * time.Second):
var zero T
t.Fatal("timed out waiting for signal")
return zero
}
}
func assertBlocked[T any](t *testing.T, ch <-chan T, message string) {
t.Helper()
select {
case value := <-ch:
t.Fatalf("%s: unexpectedly received %v", message, value)
case <-time.After(20 * time.Millisecond):
}
}
func waitGroup(t *testing.T, group *sync.WaitGroup) {
t.Helper()
done := make(chan struct{})
go func() { group.Wait(); close(done) }()
waitSignal(t, done)
}