From 24c7ab3a653d825dcb1b49e63b1eae97ea47fb5a Mon Sep 17 00:00:00 2001 From: billy Date: Tue, 20 Oct 2026 22:04:17 +0800 Subject: [PATCH] Keep pool running through worker failures until explicit close --- README.md | 12 +-- pool.go | 96 +++++++++++-------- pool_test.go | 264 ++++++++++++++++++++++++++++++++++++++++----------- 3 files changed, 272 insertions(+), 100 deletions(-) diff --git a/README.md b/README.md index 7cbb40a..e8f8ff0 100644 --- a/README.md +++ b/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 -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. +Only `Close` shuts down the pool. A task panic is logged and the same worker +continues. A worker scheduling panic or `runtime.Goexit` terminates that +worker; the pool starts a replacement as needed, and accepted queued tasks +remain available even while `Close` is draining. The task that panicked or +called `runtime.Goexit` is not retried. 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 79d740a..8ee2045 100644 --- a/pool.go +++ b/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 - return nil + if trySend(w, task) { + 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 -// worker panic exits the last worker, queued tasks are discarded. It is safe -// to call Close concurrently or more than once. +// before Close finish before it returns, even if workers exit abnormally and +// need replacement. It is safe 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 { - p.workerExitedAbnormally() + if !exitAccounted { + 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. -// 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() { +// accountWorkerExit decrements a worker and, while still holding mu, starts a +// replacement for the oldest accepted task or restores the minimum worker count. +// This applies to panic, Goexit, and a closed inbox. It never retries the task +// that caused the worker to exit and never discards queued tasks. +func (p *Pool) accountWorkerExit(w *worker, exitAccounted *bool) { p.mu.Lock() - p.total-- - if p.total == 0 { - for i := range p.queue { - p.queue[i] = nil + defer p.mu.Unlock() + for i, idleWorker := range p.idle { + if idleWorker == w { + 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) { diff --git a/pool_test.go b/pool_test.go index 68418f6..09e65b9 100644 --- a/pool_test.go +++ b/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 {