1 files changed, 262 insertions, 0 deletions

diff --git a/workerpool/workerpool.go b/workerpool/workerpool.go
new file mode 100644
index 0000000..44a3edf
--- /dev/null
+++ b/workerpool/workerpool.go
@@ -0,0 +1,262 @@
+// Package workerpool provides a generic worker pool implementation for concurrent task processing.
+package workerpool
+
+import (
+        "context"
+        "errors"
+        "fmt"
+        "maps"
+        "sync"
+        "sync/atomic"
+
+        "golang.org/x/sync/errgroup"
+)
+
+// ErrQueueFull is returned by Push when the task queue has no remaining capacity.
+var ErrQueueFull = errors.New("workerpool: task queue is full")
+
+// Task represents a unit of work that can be executed by the worker pool.
+type Task[T any] interface {
+        Execute(ctx context.Context) error
+        GetID() string
+}
+
+// WorkerPoolObserver defines callbacks for monitoring worker pool lifecycle events.
+type WorkerPoolObserver interface {
+        OnWorkerStart(workerID int)
+        OnWorkerStop(workerID int)
+        OnTaskDone(workerID int, taskID string, err error)
+}
+
+// WorkerPoolConfig holds configuration options for creating a WorkerPool instance.
+type WorkerPoolConfig struct {
+        Name          string
+        NumWorkers    int
+        TaskQueueSize int
+        Observer      WorkerPoolObserver
+        ErrorHandler  func(workerID int, taskID string, err error)
+        PanicHandler  func(workerID int, taskID string, panicValue any)
+}
+
+// Option is a functional option for configuring a WorkerPool instance.
+type Option func(*WorkerPoolConfig)
+
+// WorkerPool manages a fixed-size pool of workers that process tasks concurrently.
+type WorkerPool[T Task[T]] struct {
+        cfg       *WorkerPoolConfig
+        taskQueue chan T
+        metrics   *metrics
+        eg        *errgroup.Group
+        once      sync.Once
+        ctx       context.Context
+        started   atomic.Bool
+}
+
+type metrics struct {
+        mu             sync.RWMutex
+        TasksProcessed map[string]int
+        TasksFailed    int
+        ActiveWorkers  int
+}
+
+// MetricsSnapshot provides a point-in-time view of worker pool metrics.
+type MetricsSnapshot struct {
+        TasksProcessed map[string]int
+        TasksFailed    int
+        ActiveWorkers  int
+}
+
+// NewWorkerPool creates a new WorkerPool with the provided options.
+func NewWorkerPool[T Task[T]](ctx context.Context, opts ...Option) (*WorkerPool[T], error) {
+        cfg := &WorkerPoolConfig{
+                Name:          "workerpool",
+                NumWorkers:    1,
+                TaskQueueSize: 1,
+        }
+
+        for _, opt := range opts {
+                opt(cfg)
+        }
+
+        if cfg.NumWorkers < 1 {
+                return nil, fmt.Errorf("workerpool: number of workers must be at least 1, got %d", cfg.NumWorkers)
+        }
+
+        if cfg.TaskQueueSize < 1 {
+                return nil, fmt.Errorf("workerpool: task queue size must be at least 1, got %d", cfg.TaskQueueSize)
+        }
+
+        eg, ctx := errgroup.WithContext(ctx)
+
+        return &WorkerPool[T]{
+                cfg:       cfg,
+                taskQueue: make(chan T, cfg.TaskQueueSize),
+                ctx:       ctx,
+                eg:        eg,
+                metrics: &metrics{
+                        TasksProcessed: make(map[string]int),
+                },
+        }, nil
+}
+
+func WithName(name string) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.Name = name }
+}
+
+func WithNumWorkers(n int) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.NumWorkers = n }
+}
+
+func WithTaskQueueSize(size int) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.TaskQueueSize = size }
+}
+
+func WithObserver(observer WorkerPoolObserver) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.Observer = observer }
+}
+
+func WithErrorHandler(handler func(workerID int, taskID string, err error)) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.ErrorHandler = handler }
+}
+
+func WithPanicHandler(handler func(workerID int, taskID string, panicValue any)) Option {
+        return func(cfg *WorkerPoolConfig) { cfg.PanicHandler = handler }
+}
+
+// Start launches worker goroutines. It panics if called multiple times.
+func (w *WorkerPool[T]) Start() {
+        if !w.started.CompareAndSwap(false, true) {
+                panic("workerpool: Start() called multiple times")
+        }
+
+        for i := 0; i < w.cfg.NumWorkers; i++ {
+                workerID := i
+                workerName := fmt.Sprintf("%s.worker.%d", w.cfg.Name, workerID)
+                w.eg.Go(func() error {
+                        w.metrics.mu.Lock()
+                        w.metrics.ActiveWorkers++
+                        w.metrics.mu.Unlock()
+
+                        if w.cfg.Observer != nil {
+                                safeObserve(func() { w.cfg.Observer.OnWorkerStart(workerID) })
+                        }
+
+                        defer func() {
+                                w.metrics.mu.Lock()
+                                w.metrics.ActiveWorkers--
+                                w.metrics.mu.Unlock()
+
+                                if w.cfg.Observer != nil {
+                                        safeObserve(func() { w.cfg.Observer.OnWorkerStop(workerID) })
+                                }
+                        }()
+
+                        return w.worker(workerID, workerName)
+                })
+        }
+}
+
+// Push adds a task to the queue. Returns ErrQueueFull if full, or context error if cancelled.
+func (w *WorkerPool[T]) Push(task T) error {
+        select {
+        case w.taskQueue <- task:
+                return nil
+        case <-w.ctx.Done():
+                return fmt.Errorf("workerpool %q: push: %w", w.cfg.Name, w.ctx.Err())
+        default:
+                return fmt.Errorf("workerpool %q: capacity %d: %w", w.cfg.Name, cap(w.taskQueue), ErrQueueFull)
+        }
+}
+
+// Stop closes the task queue.
+func (w *WorkerPool[T]) Stop() {
+        w.once.Do(func() {
+                close(w.taskQueue)
+        })
+}
+
+// Wait blocks until all workers finish.
+func (w *WorkerPool[T]) Wait() error {
+        return w.eg.Wait()
+}
+
+// StopAndWait closes the queue and waits for all workers to drain.
+func (w *WorkerPool[T]) StopAndWait() error {
+        w.Stop()
+        return w.Wait()
+}
+
+// GetMetrics returns a snapshot of current pool metrics.
+func (w *WorkerPool[T]) GetMetrics() MetricsSnapshot {
+        w.metrics.mu.RLock()
+        defer w.metrics.mu.RUnlock()
+
+        tasksProcessed := make(map[string]int, len(w.metrics.TasksProcessed))
+        maps.Copy(tasksProcessed, w.metrics.TasksProcessed)
+
+        return MetricsSnapshot{
+                ActiveWorkers:  w.metrics.ActiveWorkers,
+                TasksFailed:    w.metrics.TasksFailed,
+                TasksProcessed: tasksProcessed,
+        }
+}
+
+func (w *WorkerPool[T]) worker(workerID int, workerName string) error {
+        for {
+                select {
+                case task, ok := <-w.taskQueue:
+                        if !ok {
+                                return nil
+                        }
+                        w.executeSafe(workerID, workerName, task)
+                case <-w.ctx.Done():
+                        return w.ctx.Err()
+                }
+        }
+}
+
+func (w *WorkerPool[T]) executeSafe(workerID int, workerName string, task T) {
+        taskID := task.GetID()
+
+        defer func() {
+                if r := recover(); r != nil {
+                        w.recordTaskFailed()
+                        if w.cfg.PanicHandler != nil {
+                                w.cfg.PanicHandler(workerID, taskID, r)
+                        }
+                }
+        }()
+
+        if err := task.Execute(w.ctx); err != nil {
+                w.recordTaskFailed()
+                if w.cfg.Observer != nil {
+                        safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, err) })
+                }
+                if w.cfg.ErrorHandler != nil {
+                        w.cfg.ErrorHandler(workerID, taskID, err)
+                }
+                return
+        }
+
+        w.recordTaskSuccess(workerName)
+        if w.cfg.Observer != nil {
+                safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, nil) })
+        }
+}
+
+func (w *WorkerPool[T]) recordTaskSuccess(workerName string) {
+        w.metrics.mu.Lock()
+        defer w.metrics.mu.Unlock()
+        w.metrics.TasksProcessed[workerName]++
+}
+
+func (w *WorkerPool[T]) recordTaskFailed() {
+        w.metrics.mu.Lock()
+        defer w.metrics.mu.Unlock()
+        w.metrics.TasksFailed++
+}
+
+func safeObserve(fn func()) {
+        defer func() { recover() }()
+        fn()
+}