From 44cadc6d26ac23bc9e96b6fceffd2a80ff541dbe Mon Sep 17 00:00:00 2001
From: billy <me@ybilly.com>
Date: Tue, 20 Oct 2026 21:47:55 +0800
Subject: [PATCH] Isolate worker panics and preserve queued tasks for surviving
 workers

---
 README.md    |  10 ++-
 pool.go      |  84 ++++++++++++-------
 pool_test.go | 229 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 286 insertions(+), 37 deletions(-)

diff --git a/README.md b/README.md
index d70ec09..7cbb40a 100644
--- a/README.md
+++ b/README.md
@@ -26,10 +26,12 @@ 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. `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.
+Tasks should return normally or panic. A panic in worker scheduling code is
+logged and terminates that worker without replacement. `runtime.Goexit` from a
+task has the same worker-exit behavior. If the last worker exits abnormally,
+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 2aded5d..79d740a 100644
--- a/pool.go
+++ b/pool.go
@@ -85,9 +85,9 @@ func (p *Pool) Submit(task func()) error {
 }
 
 // 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.
+// before Close normally finish before it returns. If runtime.Goexit or a
+// worker panic exits the last worker, 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
@@ -120,6 +120,9 @@ func (p *Pool) startWorkerLocked(task func(), withTask bool) {
 func (p *Pool) run(w *worker) {
 	normalExit := false
 	defer func() {
+		if recovered := recover(); recovered != nil {
+			safeLogf("worker panic: %v", recovered)
+		}
 		if !normalExit {
 			p.workerExitedAbnormally()
 		}
@@ -127,41 +130,55 @@ func (p *Pool) run(w *worker) {
 	for {
 		task, ok := <-w.inbox
 		if !ok {
-			p.mu.Lock()
-			p.total--
-			p.workerExitedLocked()
-			normalExit = true
-			p.mu.Unlock()
+			p.finishClosedWorker(&normalExit)
 			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--
-			normalExit = true
-			if !p.closed {
-				p.submitCond.Signal()
-			} else {
-				p.workerExitedLocked()
-			}
-			p.mu.Unlock()
+		if p.finishTask(w, &normalExit) {
 			return
 		}
-		p.idle = append(p.idle, w)
-		p.submitCond.Signal()
-		p.mu.Unlock()
 	}
 }
 
+// finishClosedWorker accounts for a normal Close of an idle worker.
+func (p *Pool) finishClosedWorker(normalExit *bool) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.total--
+	*normalExit = true
+	p.workerExitedLocked()
+}
+
+// finishTask commits the post-task transition under mu. If queue work exists,
+// it first hands off the head item and only then removes it from the ring, so
+// a closed inbox panic cannot lose an accepted task. The deferred unlock runs
+// before run's panic cleanup acquires the mutex.
+func (p *Pool) finishTask(w *worker, normalExit *bool) bool {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.queueLen > 0 {
+		next := p.queue[p.queueHead]
+		w.inbox <- next
+		p.dequeueLocked()
+		p.submitCond.Signal()
+		return false
+	}
+	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 {
+			p.workerExitedLocked()
+		}
+		return true
+	}
+	p.idle = append(p.idle, w)
+	p.submitCond.Signal()
+	return false
+}
+
 // 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()) {
@@ -212,10 +229,15 @@ func (p *Pool) workerExitedAbnormally() {
 	p.mu.Unlock()
 }
 
+func safeLogf(format string, args ...any) {
+	defer func() { _ = recover() }()
+	log.Printf(format, args...)
+}
+
 func runTask(task func()) {
 	defer func() {
 		if recovered := recover(); recovered != nil {
-			log.Printf("worker task panic: %v", recovered)
+			safeLogf("worker task panic: %v", recovered)
 		}
 	}()
 	task()
diff --git a/pool_test.go b/pool_test.go
index 909c171..68418f6 100644
--- a/pool_test.go
+++ b/pool_test.go
@@ -425,20 +425,230 @@ func TestPanicIsLoggedAndWorkerContinues(t *testing.T) {
 	previousOutput := log.Writer()
 	log.SetOutput(&output)
 	defer log.SetOutput(previousOutput)
+	p.mu.Lock()
+	workerBefore := p.idle[0]
+	p.mu.Unlock()
 	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() { panic(nil) }); err != nil {
+		log.SetOutput(previousOutput)
+		t.Fatal(err)
+	}
+	if err := p.Submit(func() {
+		defer func() { panic("expected deferred 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)
+	waitPoolState(t, p, 1, 1)
+	p.mu.Lock()
+	if p.total != 1 || len(p.idle) != 1 || p.idle[0] != workerBefore {
+		p.mu.Unlock()
+		t.Fatal("task panic replaced or lost the original worker")
+	}
+	p.mu.Unlock()
 	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)
+	got := output.String()
+	for _, marker := range []string{"expected test panic", "expected deferred panic"} {
+		if !strings.Contains(got, marker) {
+			t.Errorf("panic log %q does not contain %q", got, marker)
+		}
+	}
+}
+
+type panicLogWriter struct {
+	panicAlways bool
+	panicked    atomic.Bool
+	calls       atomic.Int32
+}
+
+func (w *panicLogWriter) Write(data []byte) (int, error) {
+	w.calls.Add(1)
+	if w.panicAlways || w.panicked.CompareAndSwap(false, true) {
+		panic("test logger writer panic")
+	}
+	return len(data), nil
+}
+
+func TestLoggerWriterPanicDoesNotKillTaskWorker(t *testing.T) {
+	for _, panicAlways := range []bool{false, true} {
+		name := "once"
+		if panicAlways {
+			name = "always"
+		}
+		t.Run(name, func(t *testing.T) {
+			p, err := New(1, 1)
+			if err != nil {
+				t.Fatal(err)
+			}
+			previousOutput := log.Writer()
+			writer := &panicLogWriter{panicAlways: panicAlways}
+			log.SetOutput(writer)
+			defer log.SetOutput(previousOutput)
+			p.mu.Lock()
+			workerBefore := p.idle[0]
+			p.mu.Unlock()
+			done := make(chan struct{})
+			for _, task := range []func(){
+				func() { panic("first task panic") },
+				func() { panic("second task panic") },
+				func() { close(done) },
+			} {
+				if err := p.Submit(task); err != nil {
+					t.Fatal(err)
+				}
+			}
+			waitSignal(t, done)
+			waitPoolState(t, p, 1, 1)
+			p.mu.Lock()
+			if p.idle[0] != workerBefore {
+				p.mu.Unlock()
+				t.Fatal("task panic replaced the original worker")
+			}
+			p.mu.Unlock()
+			p.Close()
+			if writer.calls.Load() == 0 {
+				t.Fatal("test logger writer was not called")
+			}
+		})
+	}
+}
+
+func TestWorkerPanicLeavesQueueForSurvivingWorker(t *testing.T) {
+	p, err := New(2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var output bytes.Buffer
+	previousOutput := log.Writer()
+	log.SetOutput(&output)
+	defer log.SetOutput(previousOutput)
+
+	p.mu.Lock()
+	victim := p.idle[len(p.idle)-1]
+	p.mu.Unlock()
+	victimGate, victimStarted := make(chan struct{}), make(chan struct{})
+	if err := p.Submit(func() { close(victimStarted); <-victimGate }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, victimStarted)
+	p.mu.Lock()
+	survivor := p.idle[len(p.idle)-1]
+	p.mu.Unlock()
+	survivorGate, survivorStarted := make(chan struct{}), make(chan struct{})
+	if err := p.Submit(func() { close(survivorStarted); <-survivorGate }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, survivorStarted)
+	queuedDone := make(chan struct{})
+	if err := p.Submit(func() { close(queuedDone) }); err != nil {
+		t.Fatal(err)
+	}
+	close(victim.inbox) // fault injection while the selected worker is gated
+	close(victimGate)
+	waitPoolState(t, p, 1, 0)
+	p.mu.Lock()
+	if p.queueLen != 1 || p.queue[p.queueHead] == nil {
+		p.mu.Unlock()
+		t.Fatal("worker panic removed queued work while another worker remained")
+	}
+	p.mu.Unlock()
+	if !strings.Contains(output.String(), "worker panic: send on closed channel") {
+		t.Fatalf("worker panic was not logged: %q", output.String())
+	}
+	close(survivorGate)
+	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 1, 1)
+	p.mu.Lock()
+	if p.idle[0] != survivor {
+		p.mu.Unlock()
+		t.Fatal("surviving worker was replaced after another worker panicked")
+	}
+	p.mu.Unlock()
+	p.Close()
+	if output.String() == "" {
+		t.Fatal("worker panic log was lost")
+	}
+}
+
+func TestLastWorkerPanicDropsQueueAndCloseReturns(t *testing.T) {
+	p, err := New(2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	var output bytes.Buffer
+	previousOutput := log.Writer()
+	log.SetOutput(&output)
+	defer log.SetOutput(previousOutput)
+
+	p.mu.Lock()
+	workerA := p.idle[len(p.idle)-1]
+	p.mu.Unlock()
+	started := make(chan struct{}, 2)
+	gateA, gateB := make(chan struct{}), make(chan struct{})
+	if err := p.Submit(func() { started <- struct{}{}; <-gateA }); err != nil {
+		t.Fatal(err)
+	}
+	p.mu.Lock()
+	workerB := p.idle[len(p.idle)-1]
+	p.mu.Unlock()
+	if err := p.Submit(func() { started <- struct{}{}; <-gateB }); 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 = 3
+	closeStarted, closeDone := make(chan struct{}, closeCount), make(chan struct{}, closeCount)
+	for range closeCount {
+		go func() { closeStarted <- struct{}{}; p.Close(); closeDone <- struct{}{} }()
+	}
+	for range closeCount {
+		waitSignal(t, closeStarted)
+	}
+	waitPoolClosed(t, p)
+	close(workerA.inbox)
+	close(workerB.inbox)
+	close(gateA)
+	waitPoolQueueState(t, p, 1, 0, 1)
+	select {
+	case <-closeDone:
+		t.Fatal("Close returned before the last worker exited")
+	case <-time.After(10 * time.Millisecond):
+	}
+	close(gateB)
+	for range closeCount {
+		waitSignal(t, closeDone)
+	}
+	if got := queuedRuns.Load(); got != 0 {
+		t.Fatalf("queued task ran after last worker panic: %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 worker panics: 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 task after worker panics", i)
+		}
+	}
+	p.mu.Unlock()
+	if !strings.Contains(output.String(), "worker panic: send on closed channel") {
+		t.Fatalf("worker-level panics not logged: %q", output.String())
 	}
 }
 
@@ -1054,6 +1264,21 @@ func waitPoolQueueState(t *testing.T, p *Pool, total, idle, queueLen int) {
 	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 waitPoolClosed(t *testing.T, p *Pool) {
+	t.Helper()
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		p.mu.Lock()
+		closed := p.closed
+		p.mu.Unlock()
+		if closed {
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+	t.Fatal("pool did not enter the closed state")
+}
+
 func waitSignal[T any](t *testing.T, ch <-chan T) T {
 	t.Helper()
 	select {