Isolate worker panics and preserve queued tasks for surviving workers

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 {