worker/pool_test.go

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
	}
}