diff --git a/hook/dynamic/dynamic.go b/hook/dynamic/dynamic.go new file mode 100644 index 0000000..38b067d --- /dev/null +++ b/hook/dynamic/dynamic.go @@ -0,0 +1,169 @@ +// Package dynamic provides an opt-in lifecycle hook dispatcher whose +// callbacks may be registered while a pool is processing tasks. +// +// Each Dispatch method reads one immutable callback snapshot. A callback +// registered during a dispatch therefore applies to later dispatches, never +// to the one already in progress. Different lifecycle events for the same +// task can observe different snapshots. +package dynamic + +import ( + "context" + "log" + "sync" + "sync/atomic" + + agilepool "github.com/Yiming1997/agilePool/v2" +) + +// Hooks stores and dispatches lifecycle callbacks. Unlike hook.Hooks, it is +// safe to add callbacks while events are being dispatched. It implements +// agilepool.Hooks. +type Hooks struct { + mu sync.Mutex + snapshots atomic.Pointer[callbackSnapshot] + logger log.Logger +} + +// callbackSnapshot is one immutable generation of lifecycle callbacks. +// Dispatch reads one snapshot for the entire event, while registrations +// publish a new snapshot for later events. +type callbackSnapshot struct { + taskSubmitted []agilepool.TaskHook + taskEnqueued []agilepool.TaskHook + taskStarted []agilepool.TaskHook + taskCompleted []agilepool.TaskCompleteHook + poolClosed []agilepool.PoolHook +} + +var _ agilepool.Hooks = (*Hooks)(nil) + +// NewHooks returns an empty dynamic dispatcher that logs recovered callback +// panics through the standard logger. +func NewHooks() *Hooks { + h := &Hooks{logger: *log.Default()} + h.snapshots.Store(&callbackSnapshot{}) + return h +} + +// AddTaskSubmitted registers a callback for future task-submission events. +func (h *Hooks) AddTaskSubmitted(fn agilepool.TaskHook) { + h.mu.Lock() + current := h.snapshot() + next := *current + next.taskSubmitted = appendCallback(current.taskSubmitted, fn) + h.snapshots.Store(&next) + h.mu.Unlock() +} + +// AddTaskEnqueued registers a callback for future task-enqueue events. +func (h *Hooks) AddTaskEnqueued(fn agilepool.TaskHook) { + h.mu.Lock() + current := h.snapshot() + next := *current + next.taskEnqueued = appendCallback(current.taskEnqueued, fn) + h.snapshots.Store(&next) + h.mu.Unlock() +} + +// AddTaskStarted registers a callback for future task-start events. +func (h *Hooks) AddTaskStarted(fn agilepool.TaskHook) { + h.mu.Lock() + current := h.snapshot() + next := *current + next.taskStarted = appendCallback(current.taskStarted, fn) + h.snapshots.Store(&next) + h.mu.Unlock() +} + +// AddTaskCompleted registers a callback for future task-completion events. +func (h *Hooks) AddTaskCompleted(fn agilepool.TaskCompleteHook) { + h.mu.Lock() + current := h.snapshot() + next := *current + next.taskCompleted = appendCallback(current.taskCompleted, fn) + h.snapshots.Store(&next) + h.mu.Unlock() +} + +// AddPoolClosed registers a callback for future pool-close events. +func (h *Hooks) AddPoolClosed(fn agilepool.PoolHook) { + h.mu.Lock() + current := h.snapshot() + next := *current + next.poolClosed = appendCallback(current.poolClosed, fn) + h.snapshots.Store(&next) + h.mu.Unlock() +} + +// DispatchTaskSubmitted dispatches one stable snapshot of submission callbacks. +func (h *Hooks) DispatchTaskSubmitted(ctx context.Context) { + for _, fn := range h.snapshot().taskSubmitted { + h.invoke(func() { fn(ctx) }, "OnTaskSubmitted") + } +} + +// DispatchTaskEnqueued dispatches one stable snapshot of enqueue callbacks. +func (h *Hooks) DispatchTaskEnqueued(ctx context.Context) { + for _, fn := range h.snapshot().taskEnqueued { + h.invoke(func() { fn(ctx) }, "OnTaskEnqueued") + } +} + +// DispatchTaskStarted dispatches one stable snapshot of start callbacks. +func (h *Hooks) DispatchTaskStarted(ctx context.Context) { + for _, fn := range h.snapshot().taskStarted { + h.invoke(func() { fn(ctx) }, "OnTaskStarted") + } +} + +// DispatchTaskCompleted dispatches one stable snapshot of completion callbacks. +func (h *Hooks) DispatchTaskCompleted(ctx context.Context, recovered any) { + for _, fn := range h.snapshot().taskCompleted { + h.invoke(func() { fn(ctx, recovered) }, "OnTaskCompleted") + } +} + +// DispatchPoolClosed dispatches one stable snapshot of pool-close callbacks. +func (h *Hooks) DispatchPoolClosed(pool *agilepool.Pool) { + if pool == nil { + h.logger.Println("[ERROR] DispatchPoolClosed: received nil pointer") + return + } + for _, fn := range h.snapshot().poolClosed { + h.invoke(func() { fn(pool) }, "OnPoolClosed") + } +} + +// snapshot returns the currently published callback generation. The nil +// fallback keeps the zero value of Hooks safe; NewHooks always publishes an +// initial empty snapshot. +func (h *Hooks) snapshot() *callbackSnapshot { + if snapshot := h.snapshots.Load(); snapshot != nil { + return snapshot + } + return &callbackSnapshot{} +} + +// appendCallback returns a new callback slice. Plain append is deliberately +// avoided because it may reuse spare capacity and mutate the backing array +// of an already published snapshot, which lock-free dispatch depends on +// never happening. +func appendCallback[T any](callbacks []T, callback T) []T { + next := make([]T, len(callbacks)+1) + copy(next, callbacks) + next[len(callbacks)] = callback + return next +} + +// invoke runs one callback with per-callback panic recovery. A panicking +// callback must not prevent the remaining callbacks of the same event from +// running. +func (h *Hooks) invoke(fn func(), name string) { + defer func() { + if recovered := recover(); recovered != nil { + h.logger.Printf("hook %s panicked: %v\n%s \n", name, recovered, agilepool.Stack(2)) + } + }() + fn() +} diff --git a/hook/dynamic/dynamic_test.go b/hook/dynamic/dynamic_test.go new file mode 100644 index 0000000..df6fe68 --- /dev/null +++ b/hook/dynamic/dynamic_test.go @@ -0,0 +1,166 @@ +package dynamic + +import ( + "context" + "io" + "log" + "sync" + "sync/atomic" + "testing" + "time" + + agilepool "github.com/Yiming1997/agilePool/v2" +) + +func TestHooksDispatchesAllEvents(t *testing.T) { + hooks := NewHooks() + hooks.logger = *log.New(io.Discard, "", 0) + + var submitted, enqueued, started, completed, closed atomic.Int64 + hooks.AddTaskSubmitted(func(context.Context) { submitted.Add(1) }) + hooks.AddTaskEnqueued(func(context.Context) { enqueued.Add(1) }) + hooks.AddTaskStarted(func(context.Context) { started.Add(1) }) + hooks.AddTaskCompleted(func(context.Context, any) { completed.Add(1) }) + hooks.AddPoolClosed(func(*agilepool.Pool) { closed.Add(1) }) + + pool := agilepool.NewPool(agilepool.NewConfig()) + defer pool.Close() + + hooks.DispatchTaskSubmitted(context.Background()) + hooks.DispatchTaskEnqueued(context.Background()) + hooks.DispatchTaskStarted(context.Background()) + hooks.DispatchTaskCompleted(context.Background(), nil) + hooks.DispatchPoolClosed(pool) + + if got := submitted.Load(); got != 1 { + t.Fatalf("submitted callbacks = %d, want 1", got) + } + if got := enqueued.Load(); got != 1 { + t.Fatalf("enqueued callbacks = %d, want 1", got) + } + if got := started.Load(); got != 1 { + t.Fatalf("started callbacks = %d, want 1", got) + } + if got := completed.Load(); got != 1 { + t.Fatalf("completed callbacks = %d, want 1", got) + } + if got := closed.Load(); got != 1 { + t.Fatalf("closed callbacks = %d, want 1", got) + } +} + +func TestHooksRegistrationDuringDispatchAppliesToNextEvent(t *testing.T) { + hooks := NewHooks() + hooks.logger = *log.New(io.Discard, "", 0) + + var initialCalls, addedCalls atomic.Int64 + hooks.AddTaskStarted(func(context.Context) { + initialCalls.Add(1) + hooks.AddTaskStarted(func(context.Context) { + addedCalls.Add(1) + }) + }) + + hooks.DispatchTaskStarted(context.Background()) + if got := initialCalls.Load(); got != 1 { + t.Fatalf("initial callbacks after first dispatch = %d, want 1", got) + } + if got := addedCalls.Load(); got != 0 { + t.Fatalf("new callbacks after first dispatch = %d, want 0", got) + } + + hooks.DispatchTaskStarted(context.Background()) + if got := initialCalls.Load(); got != 2 { + t.Fatalf("initial callbacks after second dispatch = %d, want 2", got) + } + if got := addedCalls.Load(); got != 1 { + t.Fatalf("new callbacks after second dispatch = %d, want 1", got) + } +} + +func TestHooksCallbackPanicDoesNotStopRemainingCallbacks(t *testing.T) { + hooks := NewHooks() + hooks.logger = *log.New(io.Discard, "", 0) + + var calls atomic.Int64 + hooks.AddTaskCompleted(func(context.Context, any) { panic("boom") }) + hooks.AddTaskCompleted(func(context.Context, any) { calls.Add(1) }) + + hooks.DispatchTaskCompleted(context.Background(), nil) + if got := calls.Load(); got != 1 { + t.Fatalf("remaining callbacks = %d, want 1", got) + } +} + +func TestHooksCanRegisterWhilePoolRuns(t *testing.T) { + hooks := NewHooks() + hooks.logger = *log.New(io.Discard, "", 0) + + pool := agilepool.NewPool(agilepool.NewConfig()) + pool.SetLogger(log.New(io.Discard, "", 0)) + if err := pool.SetHook(hooks); err != nil { + t.Fatalf("SetHook() error = %v", err) + } + defer pool.Close() + + started := make(chan struct{}) + release := make(chan struct{}) + pool.Submit(agilepool.TaskFunc(func() error { + close(started) + <-release + return nil + })) + + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("task did not start") + } + + var completed atomic.Int64 + hooks.AddTaskCompleted(func(context.Context, any) { completed.Add(1) }) + close(release) + pool.Wait() + + if got := completed.Load(); got != 1 { + t.Fatalf("completed callbacks = %d, want 1", got) + } +} + +func TestHooksConcurrentRegistrationAndDispatch(t *testing.T) { + hooks := NewHooks() + hooks.logger = *log.New(io.Discard, "", 0) + + const ( + registrars = 4 + dispatchers = 4 + iterations = 100 + ) + + var registered atomic.Int64 + var callbacks sync.WaitGroup + for i := 0; i < registrars; i++ { + callbacks.Add(1) + go func() { + defer callbacks.Done() + for j := 0; j < iterations; j++ { + hooks.AddTaskEnqueued(func(context.Context) {}) + registered.Add(1) + } + }() + } + for i := 0; i < dispatchers; i++ { + callbacks.Add(1) + go func() { + defer callbacks.Done() + for j := 0; j < iterations; j++ { + hooks.DispatchTaskEnqueued(context.Background()) + } + }() + } + callbacks.Wait() + + if got, want := registered.Load(), int64(registrars*iterations); got != want { + t.Fatalf("registered callbacks = %d, want %d", got, want) + } +}