fix: make Pool.SetLogger safe during concurrent task execution

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)
+	}
+}

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)