fix: make hook dispatch race-free

hook/hook.go

@@ -18,94 +18,126 @@ import (
 	"log"
 	"runtime/debug"
 	"sync"
+	"sync/atomic"
 
 	agilepool "github.com/Yiming1997/agilePool/v2"
 )
 
 // Hooks stores and dispatches lifecycle callbacks. It implements the
 // agilepool.Hooks interface.
 type Hooks struct {
-	mu            sync.RWMutex
+	// mu serializes registrations. Dispatch never takes this lock: callbacks
+	// are published as immutable snapshots and loaded atomically on the hot
+	// path.
+	mu        sync.Mutex
+	callbacks atomic.Pointer[callbackSnapshot]
+	logger    log.Logger
+}
+
+// callbackSnapshot is immutable after publication. Add methods create a new
+// snapshot before storing it, so a dispatch can safely range over a snapshot
+// while another goroutine registers a callback.
+type callbackSnapshot struct {
 	taskSubmitted []agilepool.TaskHook
 	taskEnqueued  []agilepool.TaskHook
 	taskStarted   []agilepool.TaskHook
 	taskCompleted []agilepool.TaskCompleteHook
 	poolClosed    []agilepool.PoolHook
-	logger        log.Logger
 }
 
 // NewHooks returns an empty dispatcher that logs recovered callback panics
 // through the standard logger.
 func NewHooks() *Hooks {
-	return &Hooks{
+	h := &Hooks{
 		logger: *log.Default(),
 	}
+	h.callbacks.Store(&callbackSnapshot{})
+	return h
 }
 
 // AddTaskSubmitted registers a callback for task submission.
 func (h *Hooks) AddTaskSubmitted(fn agilepool.TaskHook) {
 	h.mu.Lock()
-	h.taskSubmitted = append(h.taskSubmitted, fn)
+	current := h.currentCallbacks()
+	next := *current
+	next.taskSubmitted = appendCallback(current.taskSubmitted, fn)
+	h.callbacks.Store(&next)
 	h.mu.Unlock()
 }
 
 // AddTaskEnqueued registers a callback for task enqueue.
 func (h *Hooks) AddTaskEnqueued(fn agilepool.TaskHook) {
 	h.mu.Lock()
-	h.taskEnqueued = append(h.taskEnqueued, fn)
+	current := h.currentCallbacks()
+	next := *current
+	next.taskEnqueued = appendCallback(current.taskEnqueued, fn)
+	h.callbacks.Store(&next)
 	h.mu.Unlock()
 }
 
 // AddTaskStarted registers a callback for task start.
 func (h *Hooks) AddTaskStarted(fn agilepool.TaskHook) {
 	h.mu.Lock()
-	h.taskStarted = append(h.taskStarted, fn)
+	current := h.currentCallbacks()
+	next := *current
+	next.taskStarted = appendCallback(current.taskStarted, fn)
+	h.callbacks.Store(&next)
 	h.mu.Unlock()
 }
 
 // AddTaskCompleted registers a callback for task completion. The callback
 // receives the value a panicking task panicked with, or nil on normal exit.
 func (h *Hooks) AddTaskCompleted(fn agilepool.TaskCompleteHook) {
 	h.mu.Lock()
-	h.taskCompleted = append(h.taskCompleted, fn)
+	current := h.currentCallbacks()
+	next := *current
+	next.taskCompleted = appendCallback(current.taskCompleted, fn)
+	h.callbacks.Store(&next)
 	h.mu.Unlock()
 }
 
 // AddPoolClosed registers a callback for pool close.
 func (h *Hooks) AddPoolClosed(fn agilepool.PoolHook) {
 	h.mu.Lock()
-	h.poolClosed = append(h.poolClosed, fn)
+	current := h.currentCallbacks()
+	next := *current
+	next.poolClosed = appendCallback(current.poolClosed, fn)
+	h.callbacks.Store(&next)
 	h.mu.Unlock()
 }
 
 // DispatchTaskSubmitted dispatches submission callbacks. It must only be
 // called by the pool submission path.
 func (h *Hooks) DispatchTaskSubmitted(ctx context.Context) {
-	for _, fn := range h.taskSubmitted {
+	callbacks := h.currentCallbacks()
+	for _, fn := range callbacks.taskSubmitted {
 		h.invoke(func() { fn(ctx) }, "OnTaskSubmitted")
 	}
 }
 
 // DispatchTaskEnqueued dispatches enqueue callbacks. It must only be called
 // by the pool enqueue path.
 func (h *Hooks) DispatchTaskEnqueued(ctx context.Context) {
-	for _, fn := range h.taskEnqueued {
+	callbacks := h.currentCallbacks()
+	for _, fn := range callbacks.taskEnqueued {
 		h.invoke(func() { fn(ctx) }, "OnTaskEnqueued")
 	}
 }
 
 // DispatchTaskStarted dispatches start callbacks. It must only be called by
 // the worker execution path.
 func (h *Hooks) DispatchTaskStarted(ctx context.Context) {
-	for _, fn := range h.taskStarted {
+	callbacks := h.currentCallbacks()
+	for _, fn := range callbacks.taskStarted {
 		h.invoke(func() { fn(ctx) }, "OnTaskStarted")
 	}
 }
 
 // DispatchTaskCompleted dispatches completion callbacks. It must only be
 // called by the worker completion path.
 func (h *Hooks) DispatchTaskCompleted(ctx context.Context, recovered any) {
-	for _, fn := range h.taskCompleted {
+	callbacks := h.currentCallbacks()
+	for _, fn := range callbacks.taskCompleted {
 		h.invoke(func() { fn(ctx, recovered) }, "OnTaskCompleted")
 	}
 }
@@ -117,11 +149,29 @@ func (h *Hooks) DispatchPoolClosed(pool *agilepool.Pool) {
 		h.logger.Println("[ERROR] DispatchPoolClosed: received nil pointer")
 		return
 	}
-	for _, fn := range h.poolClosed {
+	callbacks := h.currentCallbacks()
+	for _, fn := range callbacks.poolClosed {
 		h.invoke(func() { fn(pool) }, "OnPoolClosed")
 	}
 }
 
+func (h *Hooks) currentCallbacks() *callbackSnapshot {
+	if callbacks := h.callbacks.Load(); callbacks != nil {
+		return callbacks
+	}
+	// Hooks' zero value remains usable. NewHooks is preferred because it also
+	// configures the standard logger, but a zero-value dispatcher can still
+	// safely register and dispatch callbacks.
+	return &callbackSnapshot{}
+}
+
+func appendCallback[T any](callbacks []T, fn T) []T {
+	next := make([]T, len(callbacks)+1)
+	copy(next, callbacks)
+	next[len(callbacks)] = fn
+	return next
+}
+
 func (h *Hooks) invoke(fn func(), name string) {
 	defer func() {
 		if recovered := recover(); recovered != nil {

hook/hook_test.go

@@ -0,0 +1,76 @@
+package hook
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+)
+
+// TestConcurrentRegistrationAndDispatch verifies that a callback registration
+// can overlap lifecycle dispatch without racing on the callback slices. Run
+// this test with -race; the old append-and-range implementation reports a
+// race here.
+func TestConcurrentRegistrationAndDispatch(t *testing.T) {
+	h := NewHooks()
+
+	const iterations = 1_000
+	start := make(chan struct{})
+	var registered atomic.Int32
+	var called atomic.Int32
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		<-start
+		for i := 0; i < iterations; i++ {
+			h.AddTaskStarted(func(context.Context) {
+				called.Add(1)
+			})
+			registered.Add(1)
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		<-start
+		for i := 0; i < iterations; i++ {
+			h.DispatchTaskStarted(context.Background())
+		}
+	}()
+
+	close(start)
+	wg.Wait()
+	// The concurrent dispatch loop is intentionally only a race exercise: it
+	// may finish before the first registration on a fast scheduler. Dispatch
+	// once after both goroutines finish to deterministically verify that the
+	// published callback snapshot is usable.
+	h.DispatchTaskStarted(context.Background())
+
+	if got := registered.Load(); got != iterations {
+		t.Fatalf("registered callbacks = %d, want %d", got, iterations)
+	}
+	if got := called.Load(); got == 0 {
+		t.Fatal("registered callbacks were never dispatched")
+	}
+}
+
+// TestCallbackCanRegisterAnotherCallback confirms dispatch does not hold the
+// registration mutex while invoking user code, preserving reentrant use.
+func TestCallbackCanRegisterAnotherCallback(t *testing.T) {
+	h := NewHooks()
+	var calls atomic.Int32
+
+	h.AddTaskStarted(func(context.Context) {
+		calls.Add(1)
+		h.AddTaskStarted(func(context.Context) { calls.Add(1) })
+	})
+
+	h.DispatchTaskStarted(context.Background())
+	h.DispatchTaskStarted(context.Background())
+
+	if got := calls.Load(); got != 3 {
+		t.Fatalf("callback calls = %d, want 3", got)
+	}
+}

hook_test.go

@@ -2,6 +2,7 @@ package agilepool
 
 import (
 	"context"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -162,3 +163,49 @@ func TestHookPoolClosedPanicDoesNotAbortClose(t *testing.T) {
 	}
 	p.Close() // must return normally despite the panicking hook
 }
+
+// TestSetHookConcurrentWithTaskDispatch verifies that replacing or disabling
+// a hook while Submit and worker dispatch paths are active is race-free. It
+// intentionally exercises Submitted, Enqueued, Started, and Completed.
+func TestSetHookConcurrentWithTaskDispatch(t *testing.T) {
+	p := NewPool(NewConfig(WithWorkerNumCapacity(8)))
+	defer p.Close()
+
+	const (
+		switches = 1_000
+		tasks    = 1_000
+	)
+
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		<-start
+		for i := 0; i < switches; i++ {
+			if err := p.SetHook(&panicHooks{}); err != nil {
+				t.Errorf("SetHook() error = %v", err)
+				return
+			}
+			if err := p.SetHook(nil); err != nil {
+				t.Errorf("SetHook(nil) error = %v", err)
+				return
+			}
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		<-start
+		for i := 0; i < tasks; i++ {
+			p.Submit(TaskFunc(func() error { return nil }))
+		}
+	}()
+
+	close(start)
+	wg.Wait()
+	if !waitPoolDone(p, 5*time.Second) {
+		t.Fatal("Wait() blocked after concurrent SetHook and task dispatch")
+	}
+}

pool.go

@@ -100,7 +100,16 @@ 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 holds an immutable snapshot wrapper. Dispatch only loads this
+	// pointer, while SetHook atomically publishes a replacement, so hook
+	// installation is safe while tasks are submitted or executed.
+	hooks atomic.Pointer[hookSnapshot]
+}
+
+// hookSnapshot wraps an interface so it can be atomically published. A
+// snapshot whose hooks field is nil explicitly disables lifecycle callbacks.
+type hookSnapshot struct {
+	hooks Hooks
 }
 
 func NewPool(c *Config) *Pool {
@@ -595,12 +604,11 @@ func (p *Pool) GetCapacity() int64 {
 	return p.capacity
 }
 
-// SetHook installs the lifecycle callback dispatcher. Register callbacks
-// before the pool starts processing tasks: the hook set is read without
-// synchronization by every dispatch path, so replacing it while tasks are in
-// flight is a data race. A nil set disables hook dispatch.
+// SetHook installs the lifecycle callback dispatcher. It is safe to call
+// while tasks are in flight: each lifecycle event uses the hook snapshot that
+// was current when its dispatch began. A nil set disables hook dispatch.
 func (p *Pool) SetHook(hooks Hooks) error {
-	p.hooks = hooks
+	p.hooks.Store(&hookSnapshot{hooks: hooks})
 	return nil
 }
 
@@ -610,13 +618,14 @@ func (p *Pool) SetHook(hooks Hooks) error {
 // mid-dispatch from crashing the submitting goroutine, a worker, or Close,
 // and from skipping pool bookkeeping such as wg.Done.
 func (p *Pool) dispatchHook(fn func(h Hooks)) {
-	if p.hooks == nil {
+	snapshot := p.hooks.Load()
+	if snapshot == nil || snapshot.hooks == nil {
 		return
 	}
 	defer func() {
 		if r := recover(); r != nil {
 			p.logger.Printf("hook dispatch panicked: %v\n%s\n", r, Stack(1))
 		}
 	}()
-	fn(p.hooks)
+	fn(snapshot.hooks)
 }