From c77061fecda1dc92aa61d21e8fe8b980ad483995 Mon Sep 17 00:00:00 2001 From: yiming <1965768941@qq.com> Date: Wed, 30 Sep 2026 15:06:32 +0800 Subject: [PATCH] feat: add configurable per-task panic handler Add WithPanicHandler so callers can supply a custom callback invoked when a task panics inside a worker, replacing the default logger-based reporting. Added: - PanicHandler type func(task Task, recovered any, stack []byte) and the WithPanicHandler config option in config.go - Pool.panicHandler field, initialized from config in NewPool - Pool.handlePanic: delegates to the configured handler, or falls back to the previous "worker exits from panic" log line when none is set, and recovers a panicking handler so the worker goroutine survives Changed: - worker.runTask now reports recovered panics via pool.handlePanic instead of logging directly Tests: - Custom handler: receives the exact panicking task, the recovered value and a non-empty stack; the worker keeps serving subsequent tasks - Panicking handler is recovered and Wait() still returns Docs: - Document WithPanicHandler in the README options table --- config.go | 26 +++++++++++++- pool.go | 40 +++++++++++++++++----- pool_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++ readme.md | 1 + worker.go | 7 ++-- 5 files changed, 156 insertions(+), 13 deletions(-) diff --git a/config.go b/config.go index e5d14eb..0554d92 100644 --- a/config.go +++ b/config.go @@ -1,4 +1,4 @@ -package agilepool +package agilepool import ( "time" @@ -43,10 +43,25 @@ type Config struct { statsWindowSize int // number of windows for median calculation scalerPeriod time.Duration // scaler tick interval (e.g. 50ms) backlogDecayFactor float64 // queue backlog weight in scaler target (0-1) + panicHandler PanicHandler // custom task-panic handler (nil keeps default logging) } type ConfigOption func(*Config) +// PanicHandler is invoked when a task panics while running inside a worker. +// Setting one with WithPanicHandler replaces the pool's default behaviour of +// logging the recovered value and stack trace, letting callers route panics +// to their own error reporting. +// +// - task: the task whose execution panicked +// - recovered: the value passed to panic +// - stack: formatted call stack captured at the recovery point +// +// A nil handler keeps the default logging behaviour. A handler that panics is +// itself recovered and logged, so it can never crash the worker goroutine or +// skip task-completion bookkeeping; see Pool.handlePanic. +type PanicHandler func(task Task, recovered any, stack []byte) + func NewConfig(opts ...ConfigOption) *Config { config := &Config{ cleanPeriod: defaultCleanPeriod, @@ -148,3 +163,12 @@ func WithBacklogDecayFactor(factor float64) ConfigOption { } } } + +// WithPanicHandler installs a custom handler invoked when a task panics inside +// a worker, replacing the default logger-based reporting. Pass nil to keep the +// default behaviour. +func WithPanicHandler(handler PanicHandler) ConfigOption { + return func(c *Config) { + c.panicHandler = handler + } +} diff --git a/pool.go b/pool.go index 13fdcb8..76b8c2f 100644 --- a/pool.go +++ b/pool.go @@ -100,7 +100,8 @@ type Pool struct { consumeHist *histogram // consume count distribution per window exitHist *histogram // exit count distribution per window - hooks Hooks // lifecycle callbacks registered via SetHook + hooks Hooks // lifecycle callbacks registered via SetHook + panicHandler PanicHandler // custom task-panic handler (nil keeps default logging) } func NewPool(c *Config) *Pool { @@ -108,14 +109,15 @@ func NewPool(c *Config) *Pool { c = NewConfig() } p := &Pool{ - closePoolCn: make(chan struct{}), - config: c, - lock: &sync.Mutex{}, - muIdle: &sync.Mutex{}, // default value, overridden below based on config - logger: log.Default(), - capacity: c.workerNumCapacity, - taskQueue: make(chan Task, c.taskQueueSize), - taskBuf: newChunkedTaskBuffer(), + closePoolCn: make(chan struct{}), + config: c, + lock: &sync.Mutex{}, + muIdle: &sync.Mutex{}, // default value, overridden below based on config + logger: log.Default(), + capacity: c.workerNumCapacity, + taskQueue: make(chan Task, c.taskQueueSize), + taskBuf: newChunkedTaskBuffer(), + panicHandler: c.panicHandler, } // Select muIdle lock implementation based on config: SpinLock or MutexLock (sync.Mutex) @@ -625,3 +627,23 @@ func (p *Pool) dispatchHook(fn func(h Hooks)) { }() fn(p.hooks) } + +// handlePanic reports a task panic that was recovered by the worker's runTask +// defer. When a PanicHandler is configured it takes over; otherwise the +// recovered value and stack trace are written to the pool logger, preserving +// the historical default. A panicking handler is itself recovered and logged +// so it cannot crash the worker goroutine or skip the task-completion hook +// that runTask dispatches afterwards. +func (p *Pool) handlePanic(task Task, recovered any, stack []byte) { + handler := p.panicHandler + if handler == nil { + p.logger.Printf("worker exits from panic: %v\n%s\n", recovered, stack) + return + } + defer func() { + if r := recover(); r != nil { + p.logger.Printf("panic handler panicked: %v\n%s\n", r, Stack(1)) + } + }() + handler(task, recovered, stack) +} diff --git a/pool_test.go b/pool_test.go index 5c37e50..3e89ac6 100644 --- a/pool_test.go +++ b/pool_test.go @@ -5,6 +5,7 @@ import ( "errors" "io" "log" + "reflect" "sync" "sync/atomic" "testing" @@ -451,6 +452,100 @@ func TestAgilePoolTaskPanicDoesNotBreakPool(t *testing.T) { assert.Equal(t, int64(2), atomic.LoadInt64(&executed)) } +// TestAgilePoolCustomPanicHandler verifies that a handler registered via +// WithPanicHandler receives the panicking task, the recovered value and the +// captured stack trace, and that the worker keeps servicing subsequent tasks. +func TestAgilePoolCustomPanicHandler(t *testing.T) { + var ( + calls int64 + gotTask atomic.Value + gotValue atomic.Value + gotStack atomic.Value + ) + + agilePool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(1), + agilepool.WithTaskQueueSize(10), + agilepool.WithPanicHandler(func(task agilepool.Task, recovered any, stack []byte) { + atomic.AddInt64(&calls, 1) + gotTask.Store(task) + gotValue.Store(recovered) + gotStack.Store(string(stack)) + }), + )) + defer agilePool.Close() + + panicked := agilepool.TaskFunc(func() error { + panic("custom-handler-boom") + }) + agilePool.Submit(panicked) + + var executed int64 + agilePool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt64(&executed, 1) + return nil + })) + + done := make(chan struct{}) + go func() { + agilePool.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("pool.Wait() timed out after a task panic") + } + + assert.Equal(t, int64(1), atomic.LoadInt64(&calls)) + assert.Equal(t, "custom-handler-boom", gotValue.Load()) + assert.Equal(t, int64(1), atomic.LoadInt64(&executed)) + + // The exact task that panicked must be forwarded, not a wrapper. + assert.Equal(t, + reflect.ValueOf(panicked).Pointer(), + reflect.ValueOf(gotTask.Load()).Pointer(), + ) + + // A non-empty, formatted stack trace must be provided. + stack, _ := gotStack.Load().(string) + assert.Contains(t, stack, "worker.go") +} + +// TestAgilePoolPanicHandlerPanicIsRecovered guards Pool.handlePanic: a +// PanicHandler that itself panics must be recovered so the worker goroutine +// survives and Wait() still returns. +func TestAgilePoolPanicHandlerPanicIsRecovered(t *testing.T) { + agilePool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(1), + agilepool.WithTaskQueueSize(10), + agilepool.WithPanicHandler(func(agilepool.Task, any, []byte) { + panic("handler boom") + }), + )) + agilePool.SetLogger(log.New(io.Discard, "", 0)) + defer agilePool.Close() + + var executed int64 + agilePool.Submit(agilepool.TaskFunc(func() error { panic("task boom") })) + agilePool.Submit(agilepool.TaskFunc(func() error { + atomic.AddInt64(&executed, 1) + return nil + })) + + done := make(chan struct{}) + go func() { + agilePool.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("Wait() blocked after the panic handler panicked") + } + assert.Equal(t, int64(1), atomic.LoadInt64(&executed)) +} + // TestAgilePoolBatchWithScaler verifies that the scaler spawns workers to // drain tasks submitted in a burst, even when no worker goroutines existed // at submission time (replaces the old safety-net / race-stuck-task test diff --git a/readme.md b/readme.md index 07b6e68..28eb1f3 100644 --- a/readme.md +++ b/readme.md @@ -93,6 +93,7 @@ defer pool.Close() | `WithStatsWindowSize` | `10` | Number of sliding windows for median calculation. | | `WithScalerPeriod` | `10ms` | How often the scaler evaluates whether to spawn workers. | | `WithBacklogDecayFactor` | `0.3` | Weight (0–1) for backlog in the scaler target formula. Higher values make the scaler more aggressive at draining queues. | +| `WithPanicHandler` | `nil` | Custom `func(task Task, recovered any, stack []byte)` invoked when a task panics, replacing the default log line. Defaults to logging with `SetLogger`. | > **Note**: `WithTaskQueueSize` controls the capacity of the internal handoff channel. Tasks beyond that capacity are stored in the dynamically growing chunked buffer. diff --git a/worker.go b/worker.go index 2b269ee..2f7aabc 100644 --- a/worker.go +++ b/worker.go @@ -124,13 +124,14 @@ func (w *worker) runTask(task Task) { } // Capture task panics so the Completed hook can observe the recovered - // value. dispatchHook guards the hook call itself; w.pool.done() above - // runs afterwards regardless of what the hooks do. + // value. The configured PanicHandler (or the default logger) is invoked + // via w.pool.handlePanic; w.pool.done() above runs afterwards regardless + // of what the hooks or handler do. var recovered any defer func() { if p := recover(); p != nil { recovered = p - w.pool.logger.Printf("worker exits from panic: %v\n%s\n", p, Stack(1)) + w.pool.handlePanic(task, p, Stack(1)) } w.pool.dispatchHook(func(h Hooks) { h.DispatchTaskCompleted(hookCtx, recovered)