From 4310d9fab22aaea0b3856a4f8542c8efc5bdaa1b Mon Sep 17 00:00:00 2001 From: yiming <1965768941@qq.com> Date: Sat, 17 Oct 2026 10:40:28 +0800 Subject: [PATCH] fix: make Pool.SetLogger safe during concurrent task execution SetLogger wrote the Logger interface field (16 bytes) non-atomically while worker goroutines read it in handlePanic when reporting a recovered task panic. TSan reports this as a data race: Read at ... (*Pool).handlePanic pool.go:718 Previous write at ... (*Pool).SetLogger pool.go:221 Store the logger in an atomic.Pointer[Logger] and read it through getLogger(), which also falls back to the standard-library logger after SetLogger(nil) instead of leaving a nil interface behind. Tests: - TestAgilePoolSetLoggerConcurrentWithTaskPanics reproduces the race under -race against the previous implementation and also pins the visibility contract (the logger installed last receives the next panic report). - TestAgilePoolSetLoggerNilRestoresDefault covers the nil fallback path. --- logger_test.go | 114 +++++++++++++++++++++++++++++++++++++++++++++++++ pool.go | 30 ++++++++++--- 2 files changed, 138 insertions(+), 6 deletions(-) create mode 100644 logger_test.go diff --git a/logger_test.go b/logger_test.go new file mode 100644 index 0000000..f9c278e --- /dev/null +++ b/logger_test.go @@ -0,0 +1,114 @@ +package agilepool_test + +import ( + "bytes" + "log" + "runtime" + "strings" + "sync" + "sync/atomic" + "testing" + + agilepool "github.com/Yiming1997/agilePool/v2" +) + +// countingLogger is a Logger that only counts writes, keeping the concurrent +// probe free of extra synchronization. +type countingLogger struct { + calls atomic.Int64 +} + +func (l *countingLogger) Printf(string, ...interface{}) { l.calls.Add(1) } +func (l *countingLogger) Println(...interface{}) { l.calls.Add(1) } + +// lockedBuffer is an io.Writer safe for the goroutine a worker logs from. +type lockedBuffer struct { + mu sync.Mutex + buf bytes.Buffer +} + +func (b *lockedBuffer) Write(p []byte) (int, error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.Write(p) +} + +func (b *lockedBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buf.String() +} + +// TestAgilePoolSetLoggerConcurrentWithTaskPanics is the regression test for the +// data race between Pool.SetLogger and the panic-reporting path: a worker +// goroutine reads the logger while the caller swaps it. Run with -race, this +// test fails whenever SetLogger writes the field non-atomically; without -race +// it also pins the visibility contract, i.e. a logger installed after the last +// SetLogger call must receive the next panic report. +func TestAgilePoolSetLoggerConcurrentWithTaskPanics(t *testing.T) { + const iterations = 4000 + + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(64), + )) + defer pool.Close() + pool.SetLogger(&countingLogger{}) // keep the default logger quiet + + panicTask := agilepool.TaskFunc(func() error { panic("setlogger race probe") }) + + var submitWG sync.WaitGroup + submitWG.Add(1) + go func() { + defer submitWG.Done() + for i := 0; i < iterations; i++ { + pool.Submit(panicTask) + if i%32 == 0 { + runtime.Gosched() + } + } + }() + + // Swap the logger while workers are reporting the panics above. + replacement := &countingLogger{} + for i := 0; i < iterations; i++ { + pool.SetLogger(replacement) + runtime.Gosched() + } + submitWG.Wait() + pool.Wait() + + // Visibility: the most recently installed logger must be the one the next + // panic report goes to. + final := &countingLogger{} + pool.SetLogger(final) + pool.Submit(panicTask) + pool.Wait() + if got := final.calls.Load(); got != 1 { + t.Fatalf("logger installed by SetLogger saw %d panic reports, want 1", got) + } +} + +// TestAgilePoolSetLoggerNilRestoresDefault checks that SetLogger(nil) falls +// back to the standard-library logger instead of leaving a nil interface that +// would panic on the next panic report. +func TestAgilePoolSetLoggerNilRestoresDefault(t *testing.T) { + var buf lockedBuffer + previous := log.Writer() + log.SetOutput(&buf) + defer log.SetOutput(previous) + + pool := agilepool.NewPool(agilepool.NewConfig( + agilepool.WithWorkerNumCapacity(4), + )) + defer pool.Close() + + pool.SetLogger(&countingLogger{}) + pool.SetLogger(nil) + + pool.Submit(agilepool.TaskFunc(func() error { panic("nil logger fallback") })) + pool.Wait() + + if got := buf.String(); !strings.Contains(got, "worker exits from panic") { + t.Fatalf("standard logger did not receive the panic report; output = %q", got) + } +} diff --git a/pool.go b/pool.go index e62ac93..42f306b 100644 --- a/pool.go +++ b/pool.go @@ -88,7 +88,10 @@ type Pool struct { config *Config lock *sync.Mutex wg sync.WaitGroup - logger Logger + // logger is read by every goroutine that reports a panic (workers, the + // submitting goroutine, Close), so it is stored atomically: SetLogger may + // be called while the pool is running. + logger atomic.Pointer[Logger] // workerCreateCount counts the total allocations from sync.Pool.New // over the pool lifetime. For the number of currently active workers, // use GetRunningWorkersNum(). @@ -125,11 +128,12 @@ func NewPool(c *Config) *Pool { config: c, lock: &sync.Mutex{}, muIdle: &sync.Mutex{}, // default value, overridden below based on config - logger: log.Default(), capacity: c.workerNumCapacity, taskBuf: newChunkedTaskBuffer(), panicHandler: c.panicHandler, } + defaultLogger := Logger(log.Default()) + p.logger.Store(&defaultLogger) // Shard the handoff queue so a very large worker count does not serialise // on a single channel lock. Each shard gets at least one slot and the @@ -217,8 +221,22 @@ func resolveTaskQueueShards(c *Config) int { // SetLogger replaces the default standard-library logger. // Pass the same logger instance used elsewhere in your application // (e.g. zap.SugaredLogger) so pool output appears in the same log stream. +// +// It is safe to call SetLogger while the pool is running tasks: the logger is +// read atomically by every logging path (task panics, hook panics and panic +// handlers). Passing nil restores the standard-library logger. func (p *Pool) SetLogger(l Logger) { - p.logger = l + p.logger.Store(&l) +} + +// getLogger returns the currently configured logger, falling back to the +// standard-library logger when none is installed (including after +// SetLogger(nil)). +func (p *Pool) getLogger() Logger { + if lp := p.logger.Load(); lp != nil && *lp != nil { + return *lp + } + return log.Default() } func (p *Pool) Submit(task Task) { @@ -700,7 +718,7 @@ func (p *Pool) dispatchHook(fn func(h Hooks)) { } defer func() { if r := recover(); r != nil { - p.logger.Printf("hook dispatch panicked: %v\n%s\n", r, Stack(1)) + p.getLogger().Printf("hook dispatch panicked: %v\n%s\n", r, Stack(1)) } }() fn(p.hooks) @@ -715,12 +733,12 @@ func (p *Pool) dispatchHook(fn func(h Hooks)) { 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) + p.getLogger().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)) + p.getLogger().Printf("panic handler panicked: %v\n%s\n", r, Stack(1)) } }() handler(task, recovered, stack)