From 02123df8cc15896f97c1580e49d8e2a991e37ddd Mon Sep 17 00:00:00 2001 From: sangtran Date: Thu, 3 Sep 2026 23:10:33 +0700 Subject: Moving all legacy quacker source code from github to cgit --- workerpool/workerpool.go | 262 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 262 insertions(+) create mode 100644 workerpool/workerpool.go (limited to 'workerpool/workerpool.go') 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() +} -- cgit v1.3.1