Keep pool running through worker failures until explicit close

README.md

@@ -26,12 +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. A panic in worker scheduling code is
+Only `Close` shuts down the pool. A task panic is logged and the same worker
-logged and terminates that worker without replacement. `runtime.Goexit` from a
+continues. A worker scheduling panic or `runtime.Goexit` terminates that
-task has the same worker-exit behavior. If the last worker exits abnormally,
+worker; the pool starts a replacement as needed, and accepted queued tasks
-queued accepted tasks are discarded so `Close` can return; if other workers
+remain available even while `Close` is draining. The task that panicked or
-remain, they keep the queue. Call `Close` from outside pool tasks so it does
+called `runtime.Goexit` is not retried. Call `Close` from outside pool tasks so
-not wait for itself.
+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.

pool.go

@@ -69,8 +69,13 @@ func (p *Pool) Submit(task func()) error {
 		if n := len(p.idle); n > 0 {
 			w := p.idle[n-1]
 			p.idle = p.idle[:n-1]
-			w.inbox <- task
+			if trySend(w, task) {
-			return nil
+				return nil
+			}
+			// A worker can have published itself idle immediately before its
+			// closed inbox is observed. Keep task and let the normal state
+			// machine assign it again after that worker is accounted for.
+			continue
 		}
 		if p.total < p.max {
 			p.startWorkerLocked(task, true)
@@ -85,21 +90,20 @@ 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 runtime.Goexit or a
+// before Close finish before it returns, even if workers exit abnormally and
-// worker panic exits the last worker, queued tasks are discarded. It is safe
+// need replacement. It is safe to call Close concurrently or more than once.
-// to call Close concurrently or more than once.
 func (p *Pool) Close() {
 	p.mu.Lock()
+	defer p.mu.Unlock()
 	p.closed = true
 	for _, w := range p.idle {
-		close(w.inbox)
+		safeClose(w)
 	}
 	p.idle = nil
 	p.submitCond.Broadcast()
 	for p.total != 0 {
 		p.closeCond.Wait()
 	}
-	p.mu.Unlock()
 }
 
 // All pool state (worker count, idle set, queue indexes and contents, and
@@ -117,43 +121,49 @@ func (p *Pool) startWorkerLocked(task func(), withTask bool) {
 	go p.run(w)
 }
 
+func trySend(w *worker, task func()) (sent bool) {
+	defer func() {
+		if recover() != nil {
+			sent = false
+		}
+	}()
+	w.inbox <- task
+	return true
+}
+
+func safeClose(w *worker) {
+	defer func() { _ = recover() }()
+	close(w.inbox)
+}
+
 func (p *Pool) run(w *worker) {
-	normalExit := false
+	exitAccounted := false
 	defer func() {
 		if recovered := recover(); recovered != nil {
 			safeLogf("worker panic: %v", recovered)
 		}
-		if !normalExit {
+		if !exitAccounted {
-			p.workerExitedAbnormally()
+			p.accountWorkerExit(w, &exitAccounted)
 		}
 	}()
 	for {
 		task, ok := <-w.inbox
 		if !ok {
-			p.finishClosedWorker(&normalExit)
+			p.accountWorkerExit(w, &exitAccounted)
 			return
 		}
 		runTask(task)
-		if p.finishTask(w, &normalExit) {
+		if p.finishTask(w, &exitAccounted) {
 			return
 		}
 	}
 }
 
-// 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 {
+func (p *Pool) finishTask(w *worker, exitAccounted *bool) bool {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 	if p.queueLen > 0 {
@@ -166,7 +176,7 @@ func (p *Pool) finishTask(w *worker, normalExit *bool) bool {
 	otherIdle := len(p.idle)
 	if p.closed || (p.total > p.min && otherIdle > p.min/2) {
 		p.total--
-		*normalExit = true
+		*exitAccounted = true
 		if !p.closed {
 			p.submitCond.Signal()
 		} else {
@@ -207,26 +217,34 @@ func (p *Pool) workerExitedLocked() {
 	}
 }
 
-// workerExitedAbnormally accounts for a worker terminated by runtime.Goexit.
+// accountWorkerExit decrements a worker and, while still holding mu, starts a
-// If other workers remain, they retain the FIFO queue. The last worker clears
+// replacement for the oldest accepted task or restores the minimum worker count.
-// the queue because no worker remains to consume it; all of this is protected
+// This applies to panic, Goexit, and a closed inbox. It never retries the task
-// by mu so Submit and Close observe the same state transition.
+// that caused the worker to exit and never discards queued tasks.
-func (p *Pool) workerExitedAbnormally() {
+func (p *Pool) accountWorkerExit(w *worker, exitAccounted *bool) {
 	p.mu.Lock()
-	p.total--
+	defer p.mu.Unlock()
-	if p.total == 0 {
+	for i, idleWorker := range p.idle {
-		for i := range p.queue {
+		if idleWorker == w {
-			p.queue[i] = nil
+			copy(p.idle[i:], p.idle[i+1:])
+			last := len(p.idle) - 1
+			p.idle[last] = nil
+			p.idle = p.idle[:last]
+			break
 		}
-		p.queueHead = 0
-		p.queueTail = 0
-		p.queueLen = 0
-		p.submitCond.Broadcast()
-	} else {
-		p.submitCond.Signal()
 	}
+	p.total--
+	*exitAccounted = true
+	if p.queueLen > 0 {
+		task := p.dequeueLocked()
+		p.startWorkerLocked(task, true)
+	} else {
+		if !p.closed && p.total < p.min {
+			p.startWorkerLocked(nil, false)
+		}
+	}
+	p.submitCond.Signal()
 	p.workerExitedLocked()
-	p.mu.Unlock()
 }
 
 func safeLogf(format string, args ...any) {

pool_test.go

@@ -523,7 +523,7 @@ func TestLoggerWriterPanicDoesNotKillTaskWorker(t *testing.T) {
 	}
 }
 
-func TestWorkerPanicLeavesQueueForSurvivingWorker(t *testing.T) {
+func TestWorkerPanicReplacesWorkerAndPreservesQueue(t *testing.T) {
 	p, err := New(2, 2)
 	if err != nil {
 		t.Fatal(err)
@@ -555,32 +555,176 @@ func TestWorkerPanicLeavesQueueForSurvivingWorker(t *testing.T) {
 	}
 	close(victim.inbox) // fault injection while the selected worker is gated
 	close(victimGate)
-	waitPoolState(t, p, 1, 0)
+	waitSignal(t, queuedDone)
-	p.mu.Lock()
+	waitPoolState(t, p, 2, 1)
-	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 {
+	replacement := p.idle[0]
+	if replacement == victim || replacement == survivor {
 		p.mu.Unlock()
-		t.Fatal("surviving worker was replaced after another worker panicked")
+		t.Fatal("worker panic did not create a fresh replacement")
 	}
 	p.mu.Unlock()
+	close(survivorGate)
+	waitPoolState(t, p, 2, 2)
 	p.Close()
 	if output.String() == "" {
 		t.Fatal("worker panic log was lost")
 	}
 }
 
-func TestLastWorkerPanicDropsQueueAndCloseReturns(t *testing.T) {
+func TestUnexpectedClosedIdleInboxRestoresMinimum(t *testing.T) {
+	p, err := New(1, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p.mu.Lock()
+	dead := p.idle[0]
+	p.mu.Unlock()
+	close(dead.inbox)
+	deadline := time.Now().Add(2 * time.Second)
+	for time.Now().Before(deadline) {
+		p.mu.Lock()
+		replaced := p.total == 1 && len(p.idle) == 1 && p.idle[0] != dead
+		p.mu.Unlock()
+		if replaced {
+			break
+		}
+		time.Sleep(time.Millisecond)
+	}
+	p.mu.Lock()
+	replaced := p.total == 1 && len(p.idle) == 1 && p.idle[0] != dead
+	p.mu.Unlock()
+	if !replaced {
+		t.Fatal("pool retained an idle worker with a closed inbox")
+	}
+	done := make(chan struct{})
+	if err := p.Submit(func() { close(done) }); err != nil {
+		t.Fatal(err)
+	}
+	waitSignal(t, done)
+	waitPoolState(t, p, 1, 1)
+	p.Close()
+}
+
+func TestAllOriginalWorkersExitPoolRemainsOpenAndDrainsQueue(t *testing.T) {
+	p, err := New(2, 2)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p.mu.Lock()
+	workerA := p.idle[len(p.idle)-1]
+	p.mu.Unlock()
+	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)
+	}
+	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
+	queuedDone := make(chan struct{})
+	if err := p.Submit(func() { queuedRuns.Add(1); close(queuedDone) }); err != nil {
+		t.Fatal(err)
+	}
+	close(workerA.inbox)
+	close(workerB.inbox)
+	close(gateA)
+	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 2, 1)
+	close(gateB)
+	waitPoolState(t, p, 2, 2)
+	if got := queuedRuns.Load(); got != 1 {
+		t.Fatalf("queued task ran %d times, want once", got)
+	}
+	p.mu.Lock()
+	closed := p.closed
+	p.mu.Unlock()
+	if closed {
+		t.Fatal("worker panics closed the pool without Close")
+	}
+	continued := make(chan struct{})
+	if err := p.Submit(func() { close(continued) }); err != nil {
+		t.Fatalf("Submit after worker panics = %v", err)
+	}
+	waitSignal(t, continued)
+	waitPoolState(t, p, 2, 2)
+	p.Close()
+}
+
+func TestSubmitHandlesClosedIdleInboxBeforeWorkerCleanup(t *testing.T) {
+	p, err := New(1, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p.mu.Lock()
+	dead := p.idle[0]
+	submitted := make(chan error, 1)
+	done := make(chan struct{})
+	attempted := make(chan struct{})
+	go func() {
+		close(attempted)
+		submitted <- p.Submit(func() { close(done) })
+	}()
+	waitSignal(t, attempted)
+	// Keep Submit queued on mu first, then make the published idle inbox dead.
+	time.Sleep(10 * time.Millisecond)
+	close(dead.inbox)
+	p.mu.Unlock()
+	if err := waitSignal(t, submitted); err != nil {
+		t.Fatalf("Submit after stale idle handoff = %v", err)
+	}
+	waitSignal(t, done)
+	waitPoolState(t, p, 1, 1)
+	p.mu.Lock()
+	if p.closed || p.idle[0] == dead {
+		p.mu.Unlock()
+		t.Fatal("Submit closed the pool or retained the dead worker")
+	}
+	p.mu.Unlock()
+	p.Close()
+}
+
+func TestCloseHandlesClosedIdleInboxBeforeWorkerCleanup(t *testing.T) {
+	p, err := New(1, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	p.mu.Lock()
+	dead := p.idle[0]
+	closeStarted, closeDone := make(chan struct{}), make(chan struct{})
+	go func() {
+		close(closeStarted)
+		p.Close()
+		close(closeDone)
+	}()
+	waitSignal(t, closeStarted)
+	// Queue Close on mu before waking the worker's closed-inbox cleanup.
+	time.Sleep(10 * time.Millisecond)
+	close(dead.inbox)
+	p.mu.Unlock()
+	waitSignal(t, closeDone)
+	p.mu.Lock()
+	if !p.closed || p.total != 0 || len(p.idle) != 0 {
+		p.mu.Unlock()
+		t.Fatalf("pool after Close: closed=%v total=%d idle=%d", p.closed, p.total, len(p.idle))
+	}
+	p.mu.Unlock()
+	if !p.mu.TryLock() {
+		t.Fatal("Close returned while still holding the pool mutex")
+	}
+	p.mu.Unlock()
+}
+
+func TestWorkerPanicDrainsQueueAndCloseWaits(t *testing.T) {
 	p, err := New(2, 2)
 	if err != nil {
 		t.Fatal(err)
@@ -606,8 +750,9 @@ func TestLastWorkerPanicDropsQueueAndCloseReturns(t *testing.T) {
 	}
 	waitSignal(t, started)
 	waitSignal(t, started)
+	queuedStarted, queuedGate, queuedDone := make(chan struct{}), make(chan struct{}), make(chan struct{})
 	var queuedRuns atomic.Int32
-	if err := p.Submit(func() { queuedRuns.Add(1) }); err != nil {
+	if err := p.Submit(func() { close(queuedStarted); <-queuedGate; queuedRuns.Add(1); close(queuedDone) }); err != nil {
 		t.Fatal(err)
 	}
 	const closeCount = 3
@@ -622,18 +767,20 @@ func TestLastWorkerPanicDropsQueueAndCloseReturns(t *testing.T) {
 	close(workerA.inbox)
 	close(workerB.inbox)
 	close(gateA)
-	waitPoolQueueState(t, p, 1, 0, 1)
+	waitSignal(t, queuedStarted)
 	select {
 	case <-closeDone:
-		t.Fatal("Close returned before the last worker exited")
+		t.Fatal("Close returned before the queued task finished")
 	case <-time.After(10 * time.Millisecond):
 	}
 	close(gateB)
+	close(queuedGate)
+	waitSignal(t, queuedDone)
 	for range closeCount {
 		waitSignal(t, closeDone)
 	}
-	if got := queuedRuns.Load(); got != 0 {
+	if got := queuedRuns.Load(); got != 1 {
-		t.Fatalf("queued task ran after last worker panic: %d times", got)
+		t.Fatalf("queued task ran after worker panics %d times, want once", got)
 	}
 	p.mu.Lock()
 	if p.total != 0 || len(p.idle) != 0 || p.queueLen != 0 || p.queueHead != 0 || p.queueTail != 0 {
@@ -858,7 +1005,7 @@ func TestGoexitMaxOneCloseBeforeAndAfterWorkerExit(t *testing.T) {
 	})
 }
 
-func TestGoexitKeepsQueuedTasksWhileAnotherWorkerLives(t *testing.T) {
+func TestGoexitReplacesWorkerAndKeepsQueuedTasks(t *testing.T) {
 	p, err := New(2, 2)
 	if err != nil {
 		t.Fatal(err)
@@ -880,20 +1027,20 @@ func TestGoexitKeepsQueuedTasksWhileAnotherWorkerLives(t *testing.T) {
 	}
 	close(goexitGate)
 	waitSignal(t, goexitDone)
-	waitPoolState(t, p, 1, 0)
+	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 2, 1)
 	p.mu.Lock()
-	if p.queueLen != 1 || p.queue[p.queueHead] == nil {
+	if p.queueLen != 0 {
 		p.mu.Unlock()
-		t.Fatal("queue was not retained while another worker remained")
+		t.Fatalf("queue length after replacement task = %d, want 0", p.queueLen)
 	}
 	p.mu.Unlock()
 	close(survivorGate)
-	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 2, 2)
-	waitPoolState(t, p, 1, 1) // no automatic replacement worker
 	p.Close()
 }
 
-func TestCloseWaitsThenDropsQueueAfterLastGoexit(t *testing.T) {
+func TestCloseWaitsAndDrainsQueueAfterGoexit(t *testing.T) {
 	p, err := New(2, 2) // queue capacity one
 	if err != nil {
 		t.Fatal(err)
@@ -913,8 +1060,9 @@ func TestCloseWaitsThenDropsQueueAfterLastGoexit(t *testing.T) {
 	}
 	waitSignal(t, started)
 	waitSignal(t, started)
+	queuedStarted, queuedGate, queuedDone := make(chan struct{}), make(chan struct{}), make(chan struct{})
 	var queuedRuns atomic.Int32
-	if err := p.Submit(func() { queuedRuns.Add(1) }); err != nil {
+	if err := p.Submit(func() { close(queuedStarted); <-queuedGate; queuedRuns.Add(1); close(queuedDone) }); err != nil {
 		t.Fatal(err)
 	}
 
@@ -952,36 +1100,37 @@ func TestCloseWaitsThenDropsQueueAfterLastGoexit(t *testing.T) {
 
 	close(gateA)
 	waitSignal(t, exited)
-	waitPoolQueueState(t, p, 1, 0, 1)
+	waitSignal(t, queuedStarted)
-	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")
+		t.Fatal("Close returned while the replacement task was still gated")
 	case <-time.After(10 * time.Millisecond):
 	}
 
 	close(gateB)
 	waitSignal(t, exited)
+	select {
+	case <-closeDone:
+		t.Fatal("Close returned before the queued task finished")
+	case <-time.After(10 * time.Millisecond):
+	}
+	close(queuedGate)
+	waitSignal(t, queuedDone)
 	for range closeCount {
 		waitSignal(t, closeDone)
 	}
-	if got := queuedRuns.Load(); got != 0 {
+	if got := queuedRuns.Load(); got != 1 {
-		t.Fatalf("queued task executed after last Goexit: %d times", got)
+		t.Fatalf("queued task executed after Goexit %d times, want once", 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)
+		t.Fatalf("pool state after Close: 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)
+			t.Fatalf("queue slot %d retains a task after Close", i)
 		}
 	}
 	p.mu.Unlock()
@@ -1019,7 +1168,7 @@ func TestGoexitWakesSubmitAndAllowsWorkerRecreation(t *testing.T) {
 	p.Close()
 }
 
-func TestLastGoexitDropsQueueAndPoolCanBeReused(t *testing.T) {
+func TestAllGoexitWorkersAreReplacedAndPoolCanBeReused(t *testing.T) {
 	p, err := New(4, 4) // queue capacity two; all workers are occupied
 	if err != nil {
 		t.Fatal(err)
@@ -1038,34 +1187,39 @@ func TestLastGoexitDropsQueueAndPoolCanBeReused(t *testing.T) {
 	for range 4 {
 		waitSignal(t, started)
 	}
-	var dropped atomic.Int32
+	var queuedRuns [2]atomic.Int32
-	for range 2 {
+	queuedDone := make(chan struct{}, len(queuedRuns))
-		if err := p.Submit(func() { dropped.Add(1) }); err != nil {
+	for i := range queuedRuns {
+		if err := p.Submit(func() { queuedRuns[i].Add(1); queuedDone <- struct{}{} }); err != nil {
 			t.Fatal(err)
 		}
 	}
-	for exitedWorkers := 0; exitedWorkers < 3; exitedWorkers++ {
+	close(gates[0])
-		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)
+	for range queuedRuns {
+		waitSignal(t, queuedDone)
+	}
+	for i := 1; i < len(gates); i++ {
+		close(gates[i])
+		waitSignal(t, exited)
+	}
+	waitPoolState(t, p, 4, 4)
 	p.mu.Lock()
-	if p.queueLen != 0 || p.queueHead != 0 || p.queueTail != 0 {
+	if p.closed || 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)
+		t.Fatalf("pool after worker Goexit: closed=%v queue len/head/tail=%d/%d/%d", p.closed, 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)
+			t.Fatalf("queue slot %d retains a completed task", i)
 		}
 	}
 	p.mu.Unlock()
-	if got := dropped.Load(); got != 0 {
+	for i := range queuedRuns {
-		t.Fatalf("queued tasks executed after last Goexit: %d", got)
+		if got := queuedRuns[i].Load(); got != 1 {
+			t.Fatalf("queued task %d ran %d times, want once", i, got)
+		}
 	}
 	reused := make(chan struct{})
 	if err := p.Submit(func() { close(reused) }); err != nil {