diff --git a/hook/hook.go b/hook/hook.go index 4d1529b..737dbe8 100644 --- a/hook/hook.go +++ b/hook/hook.go @@ -18,6 +18,7 @@ import ( "log" "runtime/debug" "sync" + "sync/atomic" agilepool "github.com/Yiming1997/agilePool/v2" ) @@ -25,41 +26,62 @@ import ( // 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() } @@ -67,21 +89,28 @@ func (h *Hooks) AddTaskStarted(fn agilepool.TaskHook) { // 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") } } @@ -89,7 +118,8 @@ func (h *Hooks) DispatchTaskSubmitted(ctx context.Context) { // 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") } } @@ -97,7 +127,8 @@ func (h *Hooks) DispatchTaskEnqueued(ctx context.Context) { // 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") } } @@ -105,7 +136,8 @@ func (h *Hooks) DispatchTaskStarted(ctx context.Context) { // 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 { diff --git a/hook/hook_test.go b/hook/hook_test.go new file mode 100644 index 0000000..b20cf85 --- /dev/null +++ b/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) + } +} diff --git a/hook_test.go b/hook_test.go index 9e6fd4f..1d13c47 100644 --- a/hook_test.go +++ b/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") + } +} diff --git a/pool.go b/pool.go index 56fd2e8..d9064e8 100644 --- a/pool.go +++ b/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,7 +618,8 @@ 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() { @@ -618,5 +627,5 @@ func (p *Pool) dispatchHook(fn func(h Hooks)) { p.logger.Printf("hook dispatch panicked: %v\n%s\n", r, Stack(1)) } }() - fn(p.hooks) + fn(snapshot.hooks) }