path: root/workerpool/workerpool.go

// 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()
}