Handle worker Goexit cleanup and discard orphaned queued tasks

pool_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"errors"
 	"log"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -586,6 +587,284 @@ func TestReentrantSubmit(t *testing.T) {
 	p.Close()
 }
 
+func TestGoexitMaxOneCloseBeforeAndAfterWorkerExit(t *testing.T) {
+	t.Run("close while worker is running", func(t *testing.T) {
+		p, err := New(0, 1)
+		if err != nil {
+			t.Fatal(err)
+		}
+		gate, started, exited := make(chan struct{}), make(chan struct{}), make(chan struct{})
+		if err := p.Submit(func() {
+			defer close(exited)
+			close(started)
+			<-gate
+			runtime.Goexit()
+		}); err != nil {
+			t.Fatal(err)
+		}
+		waitSignal(t, started)
+		const closeCount = 3
+		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)
+		}
+		select {
+		case <-closeDone:
+			t.Fatal("Close returned before the gated task exited")
+		case <-time.After(10 * time.Millisecond):
+		}
+		close(gate)
+		waitSignal(t, exited)
+		for range closeCount {
+			waitSignal(t, closeDone)
+		}
+		waitPoolState(t, p, 0, 0)
+	})
+
+	t.Run("close after worker exit", func(t *testing.T) {
+		p, err := New(0, 1)
+		if err != nil {
+			t.Fatal(err)
+		}
+		exited := make(chan struct{})
+		if err := p.Submit(func() { defer close(exited); runtime.Goexit() }); err != nil {
+			t.Fatal(err)
+		}
+		waitSignal(t, exited)
+		waitPoolState(t, p, 0, 0)
+		const closeCount = 3
+		var closes sync.WaitGroup
+		for range closeCount {
+			closes.Add(1)
+			go func() { defer closes.Done(); p.Close() }()
+		}
+		done := make(chan struct{})
+		go func() { closes.Wait(); close(done) }()
+		waitSignal(t, done)
+	})
+}
+
+func TestGoexitKeepsQueuedTasksWhileAnotherWorkerLives(t *testing.T) {
+	p, err := New(2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	goexitGate, survivorGate := make(chan struct{}), make(chan struct{})
+	started := make(chan struct{}, 2)
+	goexitDone := make(chan struct{})
+	if err := p.Submit(func() { defer close(goexitDone); started <- struct{}{}; <-goexitGate; runtime.Goexit() }); err != nil {
+		t.Fatal(err)
+	}
+	if err := p.Submit(func() { started <- struct{}{}; <-survivorGate }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, started)
+	waitSignal(t, started)
+	queuedDone := make(chan struct{})
+	if err := p.Submit(func() { close(queuedDone) }); err != nil {
+		t.Fatal(err)
+	}
+	close(goexitGate)
+	waitSignal(t, goexitDone)
+	waitPoolState(t, p, 1, 0)
+	p.mu.Lock()
+	if p.queueLen != 1 || p.queue[p.queueHead] == nil {
+		p.mu.Unlock()
+		t.Fatal("queue was not retained while another worker remained")
+	}
+	p.mu.Unlock()
+	close(survivorGate)
+	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 1, 1) // no automatic replacement worker
+	p.Close()
+}
+
+func TestCloseWaitsThenDropsQueueAfterLastGoexit(t *testing.T) {
+	p, err := New(2, 2) // queue capacity one
+	if err != nil {
+		t.Fatal(err)
+	}
+	gateA, gateB := make(chan struct{}), make(chan struct{})
+	started := make(chan struct{}, 2)
+	exited := make(chan struct{}, 2)
+	for _, gate := range []chan struct{}{gateA, gateB} {
+		if err := p.Submit(func() {
+			defer func() { exited <- struct{}{} }()
+			started <- struct{}{}
+			<-gate
+			runtime.Goexit()
+		}); err != nil {
+			t.Fatal(err)
+		}
+	}
+	waitSignal(t, started)
+	waitSignal(t, started)
+	var queuedRuns atomic.Int32
+	if err := p.Submit(func() { queuedRuns.Add(1) }); err != nil {
+		t.Fatal(err)
+	}
+
+	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)
+	}
+	deadline := time.Now().Add(2 * time.Second)
+	for {
+		p.mu.Lock()
+		closed := p.closed
+		p.mu.Unlock()
+		if closed {
+			break
+		}
+		if time.Now().After(deadline) {
+			t.Fatal("Close did not mark the pool closed")
+		}
+		time.Sleep(time.Millisecond)
+	}
+	select {
+	case <-closeDone:
+		t.Fatal("Close returned while both gated tasks were still running")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	close(gateA)
+	waitSignal(t, exited)
+	waitPoolQueueState(t, p, 1, 0, 1)
+	p.mu.Lock()
+	if p.queue[p.queueHead] == nil {
+		p.mu.Unlock()
+		t.Fatal("queue was cleared before the last worker exited")
+	}
+	p.mu.Unlock()
+	select {
+	case <-closeDone:
+		t.Fatal("Close returned while the final worker task was still gated")
+	case <-time.After(10 * time.Millisecond):
+	}
+
+	close(gateB)
+	waitSignal(t, exited)
+	for range closeCount {
+		waitSignal(t, closeDone)
+	}
+	if got := queuedRuns.Load(); got != 0 {
+		t.Fatalf("queued task executed after last Goexit: %d times", got)
+	}
+	p.mu.Lock()
+	if p.total != 0 || len(p.idle) != 0 || p.queueLen != 0 || p.queueHead != 0 || p.queueTail != 0 {
+		p.mu.Unlock()
+		t.Fatalf("pool state after last Goexit: total=%d idle=%d queue=%d head=%d tail=%d", p.total, len(p.idle), p.queueLen, p.queueHead, p.queueTail)
+	}
+	for i, task := range p.queue {
+		if task != nil {
+			p.mu.Unlock()
+			t.Fatalf("queue slot %d retains a task after last Goexit", i)
+		}
+	}
+	p.mu.Unlock()
+	if err := p.Submit(func() {}); !errors.Is(err, ErrClosed) {
+		t.Fatalf("Submit after Close = %v, want ErrClosed", err)
+	}
+}
+
+func TestGoexitWakesSubmitAndAllowsWorkerRecreation(t *testing.T) {
+	p, err := New(0, 1) // zero-capacity queue
+	if err != nil {
+		t.Fatal(err)
+	}
+	gate, started, exited, taskDone := make(chan struct{}), make(chan struct{}), make(chan struct{}), make(chan struct{})
+	if err := p.Submit(func() { defer close(exited); close(started); <-gate; runtime.Goexit() }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, started)
+	attempted := make(chan struct{})
+	submitted := make(chan error, 1)
+	go func() { close(attempted); submitted <- p.Submit(func() { close(taskDone) }) }()
+	waitSignal(t, attempted)
+	select {
+	case err := <-submitted:
+		t.Fatalf("Submit returned while the only worker was busy: %v", err)
+	case <-time.After(10 * time.Millisecond):
+	}
+	close(gate)
+	waitSignal(t, exited)
+	if err := waitSignal(t, submitted); err != nil {
+		t.Fatalf("Submit after Goexit = %v", err)
+	}
+	waitSignal(t, taskDone)
+	waitPoolState(t, p, 1, 1)
+	p.Close()
+}
+
+func TestLastGoexitDropsQueueAndPoolCanBeReused(t *testing.T) {
+	p, err := New(4, 4) // queue capacity two; all workers are occupied
+	if err != nil {
+		t.Fatal(err)
+	}
+	gates := make([]chan struct{}, 4)
+	started := make(chan struct{}, 4)
+	exited := 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; exited <- struct{}{}; runtime.Goexit() }); err != nil {
+			t.Fatal(err)
+		}
+	}
+	for range 4 {
+		waitSignal(t, started)
+	}
+	var dropped atomic.Int32
+	for range 2 {
+		if err := p.Submit(func() { dropped.Add(1) }); err != nil {
+			t.Fatal(err)
+		}
+	}
+	for exitedWorkers := 0; exitedWorkers < 3; exitedWorkers++ {
+		close(gates[exitedWorkers])
+		waitSignal(t, exited)
+		waitPoolQueueState(t, p, 3-exitedWorkers, 0, 2) // queued work stays until the last worker exits
+	}
+	close(gates[3])
+	waitSignal(t, exited)
+	waitPoolState(t, p, 0, 0)
+	p.mu.Lock()
+	if p.queueLen != 0 || p.queueHead != 0 || p.queueTail != 0 {
+		p.mu.Unlock()
+		t.Fatalf("queue after last Goexit: len=%d head=%d tail=%d", p.queueLen, p.queueHead, p.queueTail)
+	}
+	for i, task := range p.queue {
+		if task != nil {
+			p.mu.Unlock()
+			t.Fatalf("queue slot %d still holds a task after last Goexit", i)
+		}
+	}
+	p.mu.Unlock()
+	if got := dropped.Load(); got != 0 {
+		t.Fatalf("queued tasks executed after last Goexit: %d", got)
+	}
+	reused := make(chan struct{})
+	if err := p.Submit(func() { close(reused) }); err != nil {
+		t.Fatalf("Submit after queue clear = %v", err)
+	}
+	waitSignal(t, reused)
+	p.Close()
+}
+
 func TestConcurrentSubmitAndClose(t *testing.T) {
 	const rounds = 10
 	const taskCount = 64
@@ -758,6 +1037,23 @@ func waitPoolState(t *testing.T, p *Pool, total, idle int) {
 	t.Fatalf("pool state workers/idle = %d/%d, want %d/%d", p.total, len(p.idle), total, idle)
 }
 
+func waitPoolQueueState(t *testing.T, p *Pool, total, idle, queueLen int) {
+	t.Helper()
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		p.mu.Lock()
+		gotTotal, gotIdle, gotQueue := p.total, len(p.idle), p.queueLen
+		p.mu.Unlock()
+		if gotTotal == total && gotIdle == idle && gotQueue == queueLen {
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	t.Fatalf("pool state workers/idle/queue = %d/%d/%d, want %d/%d/%d", p.total, len(p.idle), p.queueLen, total, idle, queueLen)
+}
+
 func waitSignal[T any](t *testing.T, ch <-chan T) T {
 	t.Helper()
 	select {