Keep pool running through worker failures until explicit close

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)
-	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()
+	waitSignal(t, queuedDone)
+	waitPoolState(t, p, 2, 1)
 	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 {
-		t.Fatalf("queued task ran after last worker panic: %d times", got)
+	if got := queuedRuns.Load(); got != 1 {
+		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, 1, 1) // no automatic replacement worker
+	waitPoolState(t, p, 2, 2)
 	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)
-	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()
+	waitSignal(t, queuedStarted)
 	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 {
-		t.Fatalf("queued task executed after last Goexit: %d times", got)
+	if got := queuedRuns.Load(); got != 1 {
+		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
-	for range 2 {
-		if err := p.Submit(func() { dropped.Add(1) }); err != nil {
+	var queuedRuns [2]atomic.Int32
+	queuedDone := make(chan struct{}, len(queuedRuns))
+	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[exitedWorkers])
-		waitSignal(t, exited)
-		waitPoolQueueState(t, p, 3-exitedWorkers, 0, 2) // queued work stays until the last worker exits
-	}
-	close(gates[3])
+	close(gates[0])
 	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 {
-		t.Fatalf("queued tasks executed after last Goexit: %d", got)
+	for i := range queuedRuns {
+		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 {