From 5bf542fc67eee72a6712d211b881229151c6234c Mon Sep 17 00:00:00 2001
From: billy <me@ybilly.com>
Date: Tue, 20 Oct 2026 21:29:45 +0800
Subject: [PATCH] Handle worker Goexit cleanup and discard orphaned queued
 tasks

---
 README.md    |   7 +-
 pool.go      |  36 ++++++-
 pool_test.go | 296 +++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 334 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index f18f33f..d70ec09 100644
--- a/README.md
+++ b/README.md
@@ -26,9 +26,10 @@ 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.
+Tasks should return normally or panic. `runtime.Goexit` terminates its worker
+without replacement. If the last worker exits this way, queued accepted tasks
+are discarded so `Close` can return; if other workers remain, they keep the
+queue. 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/pool.go b/pool.go
index 82a35bf..2aded5d 100644
--- a/pool.go
+++ b/pool.go
@@ -84,8 +84,10 @@ func (p *Pool) Submit(task func()) error {
 	}
 }
 
-// Close stops accepting tasks and waits for accepted tasks and workers to
-// finish. It is safe to call Close concurrently or more than once.
+// Close stops accepting tasks and waits for workers to exit. Tasks accepted
+// before Close normally finish before it returns. If a task calls
+// runtime.Goexit and that causes the last worker to exit, queued tasks are
+// discarded. It is safe to call Close concurrently or more than once.
 func (p *Pool) Close() {
 	p.mu.Lock()
 	p.closed = true
@@ -116,12 +118,19 @@ func (p *Pool) startWorkerLocked(task func(), withTask bool) {
 }
 
 func (p *Pool) run(w *worker) {
+	normalExit := false
+	defer func() {
+		if !normalExit {
+			p.workerExitedAbnormally()
+		}
+	}()
 	for {
 		task, ok := <-w.inbox
 		if !ok {
 			p.mu.Lock()
 			p.total--
 			p.workerExitedLocked()
+			normalExit = true
 			p.mu.Unlock()
 			return
 		}
@@ -138,6 +147,7 @@ func (p *Pool) run(w *worker) {
 		otherIdle := len(p.idle)
 		if p.closed || (p.total > p.min && otherIdle > p.min/2) {
 			p.total--
+			normalExit = true
 			if !p.closed {
 				p.submitCond.Signal()
 			} else {
@@ -180,6 +190,28 @@ func (p *Pool) workerExitedLocked() {
 	}
 }
 
+// workerExitedAbnormally accounts for a worker terminated by runtime.Goexit.
+// If other workers remain, they retain the FIFO queue. The last worker clears
+// the queue because no worker remains to consume it; all of this is protected
+// by mu so Submit and Close observe the same state transition.
+func (p *Pool) workerExitedAbnormally() {
+	p.mu.Lock()
+	p.total--
+	if p.total == 0 {
+		for i := range p.queue {
+			p.queue[i] = nil
+		}
+		p.queueHead = 0
+		p.queueTail = 0
+		p.queueLen = 0
+		p.submitCond.Broadcast()
+	} else {
+		p.submitCond.Signal()
+	}
+	p.workerExitedLocked()
+	p.mu.Unlock()
+}
+
 func runTask(task func()) {
 	defer func() {
 		if recovered := recover(); recovered != nil {
diff --git a/pool_test.go b/pool_test.go
index 190c7b1..909c171 100644
--- a/pool_test.go
+++ b/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 {