aboutsummaryrefslogtreecommitdiff
path: root/workerpool/workerpool.go
diff options
context:
space:
mode:
authorsangtran <tranquangsang12.7@gmail.com>2026-09-03 23:10:33 +0700
committersangtran <tranquangsang12.7@gmail.com>2026-09-03 23:10:33 +0700
commit02123df8cc15896f97c1580e49d8e2a991e37ddd (patch)
treed0c00bc3f65e5f5493a9dff8034c80640cb12960 /workerpool/workerpool.go
signatureVerified sangtran <tranquangsang12.7@gmail.com> Key: 8DD79823D235E1BB ultimate trust
downloadquacker-main.tar.gz
quacker-main.zip
Verified Moving all legacy quacker source code from github to cgitHEADmain
Diffstat (limited to 'workerpool/workerpool.go')
-rw-r--r--workerpool/workerpool.go262
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 @@
1// Package workerpool provides a generic worker pool implementation for concurrent task processing.
2package workerpool
3
4import (
5 "context"
6 "errors"
7 "fmt"
8 "maps"
9 "sync"
10 "sync/atomic"
11
12 "golang.org/x/sync/errgroup"
13)
14
15// ErrQueueFull is returned by Push when the task queue has no remaining capacity.
16var ErrQueueFull = errors.New("workerpool: task queue is full")
17
18// Task represents a unit of work that can be executed by the worker pool.
19type Task[T any] interface {
20 Execute(ctx context.Context) error
21 GetID() string
22}
23
24// WorkerPoolObserver defines callbacks for monitoring worker pool lifecycle events.
25type WorkerPoolObserver interface {
26 OnWorkerStart(workerID int)
27 OnWorkerStop(workerID int)
28 OnTaskDone(workerID int, taskID string, err error)
29}
30
31// WorkerPoolConfig holds configuration options for creating a WorkerPool instance.
32type WorkerPoolConfig struct {
33 Name string
34 NumWorkers int
35 TaskQueueSize int
36 Observer WorkerPoolObserver
37 ErrorHandler func(workerID int, taskID string, err error)
38 PanicHandler func(workerID int, taskID string, panicValue any)
39}
40
41// Option is a functional option for configuring a WorkerPool instance.
42type Option func(*WorkerPoolConfig)
43
44// WorkerPool manages a fixed-size pool of workers that process tasks concurrently.
45type WorkerPool[T Task[T]] struct {
46 cfg *WorkerPoolConfig
47 taskQueue chan T
48 metrics *metrics
49 eg *errgroup.Group
50 once sync.Once
51 ctx context.Context
52 started atomic.Bool
53}
54
55type metrics struct {
56 mu sync.RWMutex
57 TasksProcessed map[string]int
58 TasksFailed int
59 ActiveWorkers int
60}
61
62// MetricsSnapshot provides a point-in-time view of worker pool metrics.
63type MetricsSnapshot struct {
64 TasksProcessed map[string]int
65 TasksFailed int
66 ActiveWorkers int
67}
68
69// NewWorkerPool creates a new WorkerPool with the provided options.
70func NewWorkerPool[T Task[T]](ctx context.Context, opts ...Option) (*WorkerPool[T], error) {
71 cfg := &WorkerPoolConfig{
72 Name: "workerpool",
73 NumWorkers: 1,
74 TaskQueueSize: 1,
75 }
76
77 for _, opt := range opts {
78 opt(cfg)
79 }
80
81 if cfg.NumWorkers < 1 {
82 return nil, fmt.Errorf("workerpool: number of workers must be at least 1, got %d", cfg.NumWorkers)
83 }
84
85 if cfg.TaskQueueSize < 1 {
86 return nil, fmt.Errorf("workerpool: task queue size must be at least 1, got %d", cfg.TaskQueueSize)
87 }
88
89 eg, ctx := errgroup.WithContext(ctx)
90
91 return &WorkerPool[T]{
92 cfg: cfg,
93 taskQueue: make(chan T, cfg.TaskQueueSize),
94 ctx: ctx,
95 eg: eg,
96 metrics: &metrics{
97 TasksProcessed: make(map[string]int),
98 },
99 }, nil
100}
101
102func WithName(name string) Option {
103 return func(cfg *WorkerPoolConfig) { cfg.Name = name }
104}
105
106func WithNumWorkers(n int) Option {
107 return func(cfg *WorkerPoolConfig) { cfg.NumWorkers = n }
108}
109
110func WithTaskQueueSize(size int) Option {
111 return func(cfg *WorkerPoolConfig) { cfg.TaskQueueSize = size }
112}
113
114func WithObserver(observer WorkerPoolObserver) Option {
115 return func(cfg *WorkerPoolConfig) { cfg.Observer = observer }
116}
117
118func WithErrorHandler(handler func(workerID int, taskID string, err error)) Option {
119 return func(cfg *WorkerPoolConfig) { cfg.ErrorHandler = handler }
120}
121
122func WithPanicHandler(handler func(workerID int, taskID string, panicValue any)) Option {
123 return func(cfg *WorkerPoolConfig) { cfg.PanicHandler = handler }
124}
125
126// Start launches worker goroutines. It panics if called multiple times.
127func (w *WorkerPool[T]) Start() {
128 if !w.started.CompareAndSwap(false, true) {
129 panic("workerpool: Start() called multiple times")
130 }
131
132 for i := 0; i < w.cfg.NumWorkers; i++ {
133 workerID := i
134 workerName := fmt.Sprintf("%s.worker.%d", w.cfg.Name, workerID)
135 w.eg.Go(func() error {
136 w.metrics.mu.Lock()
137 w.metrics.ActiveWorkers++
138 w.metrics.mu.Unlock()
139
140 if w.cfg.Observer != nil {
141 safeObserve(func() { w.cfg.Observer.OnWorkerStart(workerID) })
142 }
143
144 defer func() {
145 w.metrics.mu.Lock()
146 w.metrics.ActiveWorkers--
147 w.metrics.mu.Unlock()
148
149 if w.cfg.Observer != nil {
150 safeObserve(func() { w.cfg.Observer.OnWorkerStop(workerID) })
151 }
152 }()
153
154 return w.worker(workerID, workerName)
155 })
156 }
157}
158
159// Push adds a task to the queue. Returns ErrQueueFull if full, or context error if cancelled.
160func (w *WorkerPool[T]) Push(task T) error {
161 select {
162 case w.taskQueue <- task:
163 return nil
164 case <-w.ctx.Done():
165 return fmt.Errorf("workerpool %q: push: %w", w.cfg.Name, w.ctx.Err())
166 default:
167 return fmt.Errorf("workerpool %q: capacity %d: %w", w.cfg.Name, cap(w.taskQueue), ErrQueueFull)
168 }
169}
170
171// Stop closes the task queue.
172func (w *WorkerPool[T]) Stop() {
173 w.once.Do(func() {
174 close(w.taskQueue)
175 })
176}
177
178// Wait blocks until all workers finish.
179func (w *WorkerPool[T]) Wait() error {
180 return w.eg.Wait()
181}
182
183// StopAndWait closes the queue and waits for all workers to drain.
184func (w *WorkerPool[T]) StopAndWait() error {
185 w.Stop()
186 return w.Wait()
187}
188
189// GetMetrics returns a snapshot of current pool metrics.
190func (w *WorkerPool[T]) GetMetrics() MetricsSnapshot {
191 w.metrics.mu.RLock()
192 defer w.metrics.mu.RUnlock()
193
194 tasksProcessed := make(map[string]int, len(w.metrics.TasksProcessed))
195 maps.Copy(tasksProcessed, w.metrics.TasksProcessed)
196
197 return MetricsSnapshot{
198 ActiveWorkers: w.metrics.ActiveWorkers,
199 TasksFailed: w.metrics.TasksFailed,
200 TasksProcessed: tasksProcessed,
201 }
202}
203
204func (w *WorkerPool[T]) worker(workerID int, workerName string) error {
205 for {
206 select {
207 case task, ok := <-w.taskQueue:
208 if !ok {
209 return nil
210 }
211 w.executeSafe(workerID, workerName, task)
212 case <-w.ctx.Done():
213 return w.ctx.Err()
214 }
215 }
216}
217
218func (w *WorkerPool[T]) executeSafe(workerID int, workerName string, task T) {
219 taskID := task.GetID()
220
221 defer func() {
222 if r := recover(); r != nil {
223 w.recordTaskFailed()
224 if w.cfg.PanicHandler != nil {
225 w.cfg.PanicHandler(workerID, taskID, r)
226 }
227 }
228 }()
229
230 if err := task.Execute(w.ctx); err != nil {
231 w.recordTaskFailed()
232 if w.cfg.Observer != nil {
233 safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, err) })
234 }
235 if w.cfg.ErrorHandler != nil {
236 w.cfg.ErrorHandler(workerID, taskID, err)
237 }
238 return
239 }
240
241 w.recordTaskSuccess(workerName)
242 if w.cfg.Observer != nil {
243 safeObserve(func() { w.cfg.Observer.OnTaskDone(workerID, taskID, nil) })
244 }
245}
246
247func (w *WorkerPool[T]) recordTaskSuccess(workerName string) {
248 w.metrics.mu.Lock()
249 defer w.metrics.mu.Unlock()
250 w.metrics.TasksProcessed[workerName]++
251}
252
253func (w *WorkerPool[T]) recordTaskFailed() {
254 w.metrics.mu.Lock()
255 defer w.metrics.mu.Unlock()
256 w.metrics.TasksFailed++
257}
258
259func safeObserve(fn func()) {
260 defer func() { recover() }()
261 fn()
262}