diff options
| author | sangtran <tranquangsang12.7@gmail.com> | 2026-09-03 23:10:33 +0700 |
|---|---|---|
| committer | sangtran <tranquangsang12.7@gmail.com> | 2026-09-03 23:10:33 +0700 |
| commit | 02123df8cc15896f97c1580e49d8e2a991e37ddd (patch) | |
| tree | d0c00bc3f65e5f5493a9dff8034c80640cb12960 /workerpool | |
| signature | Verified sangtran <tranquangsang12.7@gmail.com> Key: 8DD79823D235E1BB ultimate trust | |
| download | quacker-main.tar.gz quacker-main.zip | |
Diffstat (limited to 'workerpool')
| -rw-r--r-- | workerpool/workerpool.go | 262 | ||||
| -rw-r--r-- | workerpool/workerpool_bench_test.go | 50 | ||||
| -rw-r--r-- | workerpool/workerpool_test.go | 608 |
3 files changed, 920 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. | ||
| 2 | package workerpool | ||
| 3 | |||
| 4 | import ( | ||
| 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. | ||
| 16 | var ErrQueueFull = errors.New("workerpool: task queue is full") | ||
| 17 | |||
| 18 | // Task represents a unit of work that can be executed by the worker pool. | ||
| 19 | type 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. | ||
| 25 | type 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. | ||
| 32 | type 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. | ||
| 42 | type Option func(*WorkerPoolConfig) | ||
| 43 | |||
| 44 | // WorkerPool manages a fixed-size pool of workers that process tasks concurrently. | ||
| 45 | type 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 | |||
| 55 | type 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. | ||
| 63 | type 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. | ||
| 70 | func 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 | |||
| 102 | func WithName(name string) Option { | ||
| 103 | return func(cfg *WorkerPoolConfig) { cfg.Name = name } | ||
| 104 | } | ||
| 105 | |||
| 106 | func WithNumWorkers(n int) Option { | ||
| 107 | return func(cfg *WorkerPoolConfig) { cfg.NumWorkers = n } | ||
| 108 | } | ||
| 109 | |||
| 110 | func WithTaskQueueSize(size int) Option { | ||
| 111 | return func(cfg *WorkerPoolConfig) { cfg.TaskQueueSize = size } | ||
| 112 | } | ||
| 113 | |||
| 114 | func WithObserver(observer WorkerPoolObserver) Option { | ||
| 115 | return func(cfg *WorkerPoolConfig) { cfg.Observer = observer } | ||
| 116 | } | ||
| 117 | |||
| 118 | func WithErrorHandler(handler func(workerID int, taskID string, err error)) Option { | ||
| 119 | return func(cfg *WorkerPoolConfig) { cfg.ErrorHandler = handler } | ||
| 120 | } | ||
| 121 | |||
| 122 | func 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. | ||
| 127 | func (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. | ||
| 160 | func (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. | ||
| 172 | func (w *WorkerPool[T]) Stop() { | ||
| 173 | w.once.Do(func() { | ||
| 174 | close(w.taskQueue) | ||
| 175 | }) | ||
| 176 | } | ||
| 177 | |||
| 178 | // Wait blocks until all workers finish. | ||
| 179 | func (w *WorkerPool[T]) Wait() error { | ||
| 180 | return w.eg.Wait() | ||
| 181 | } | ||
| 182 | |||
| 183 | // StopAndWait closes the queue and waits for all workers to drain. | ||
| 184 | func (w *WorkerPool[T]) StopAndWait() error { | ||
| 185 | w.Stop() | ||
| 186 | return w.Wait() | ||
| 187 | } | ||
| 188 | |||
| 189 | // GetMetrics returns a snapshot of current pool metrics. | ||
| 190 | func (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 | |||
| 204 | func (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 | |||
| 218 | func (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 | |||
| 247 | func (w *WorkerPool[T]) recordTaskSuccess(workerName string) { | ||
| 248 | w.metrics.mu.Lock() | ||
| 249 | defer w.metrics.mu.Unlock() | ||
| 250 | w.metrics.TasksProcessed[workerName]++ | ||
| 251 | } | ||
| 252 | |||
| 253 | func (w *WorkerPool[T]) recordTaskFailed() { | ||
| 254 | w.metrics.mu.Lock() | ||
| 255 | defer w.metrics.mu.Unlock() | ||
| 256 | w.metrics.TasksFailed++ | ||
| 257 | } | ||
| 258 | |||
| 259 | func safeObserve(fn func()) { | ||
| 260 | defer func() { recover() }() | ||
| 261 | fn() | ||
| 262 | } | ||
diff --git a/workerpool/workerpool_bench_test.go b/workerpool/workerpool_bench_test.go new file mode 100644 index 0000000..1fbc672 --- /dev/null +++ b/workerpool/workerpool_bench_test.go | |||
| @@ -0,0 +1,50 @@ | |||
| 1 | package workerpool | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "fmt" | ||
| 6 | "testing" | ||
| 7 | ) | ||
| 8 | |||
| 9 | type benchTask struct { | ||
| 10 | id string | ||
| 11 | } | ||
| 12 | |||
| 13 | func (t *benchTask) Execute(_ context.Context) error { return nil } | ||
| 14 | func (t *benchTask) GetID() string { return t.id } | ||
| 15 | |||
| 16 | func BenchmarkWorkerPool(b *testing.B) { | ||
| 17 | b.ReportAllocs() | ||
| 18 | |||
| 19 | pool, _ := NewWorkerPool[*benchTask](b.Context(), | ||
| 20 | WithNumWorkers(4), | ||
| 21 | WithTaskQueueSize(256), | ||
| 22 | ) | ||
| 23 | pool.Start() | ||
| 24 | |||
| 25 | task := &benchTask{id: "bench"} | ||
| 26 | |||
| 27 | for b.Loop() { | ||
| 28 | pool.Push(task) | ||
| 29 | } | ||
| 30 | |||
| 31 | pool.StopAndWait() | ||
| 32 | } | ||
| 33 | |||
| 34 | func BenchmarkWorkerPool_Push(b *testing.B) { | ||
| 35 | b.ReportAllocs() | ||
| 36 | |||
| 37 | pool, _ := NewWorkerPool[*benchTask](b.Context(), | ||
| 38 | WithNumWorkers(8), | ||
| 39 | WithTaskQueueSize(1024), | ||
| 40 | ) | ||
| 41 | pool.Start() | ||
| 42 | defer pool.StopAndWait() | ||
| 43 | |||
| 44 | b.RunParallel(func(pb *testing.PB) { | ||
| 45 | task := &benchTask{id: fmt.Sprintf("bench-%d", b.N)} | ||
| 46 | for pb.Next() { | ||
| 47 | pool.Push(task) | ||
| 48 | } | ||
| 49 | }) | ||
| 50 | } | ||
diff --git a/workerpool/workerpool_test.go b/workerpool/workerpool_test.go new file mode 100644 index 0000000..851874b --- /dev/null +++ b/workerpool/workerpool_test.go | |||
| @@ -0,0 +1,608 @@ | |||
| 1 | package workerpool | ||
| 2 | |||
| 3 | import ( | ||
| 4 | "context" | ||
| 5 | "errors" | ||
| 6 | "fmt" | ||
| 7 | "sync" | ||
| 8 | "testing" | ||
| 9 | "time" | ||
| 10 | ) | ||
| 11 | |||
| 12 | type mockTask struct { | ||
| 13 | id string | ||
| 14 | fn func(ctx context.Context) error | ||
| 15 | } | ||
| 16 | |||
| 17 | func (t *mockTask) Execute(ctx context.Context) error { | ||
| 18 | if t.fn != nil { | ||
| 19 | return t.fn(ctx) | ||
| 20 | } | ||
| 21 | return nil | ||
| 22 | } | ||
| 23 | |||
| 24 | func (t *mockTask) GetID() string { | ||
| 25 | return t.id | ||
| 26 | } | ||
| 27 | |||
| 28 | type taskResult struct { | ||
| 29 | taskID string | ||
| 30 | err error | ||
| 31 | } | ||
| 32 | |||
| 33 | type mockObserver struct { | ||
| 34 | mu sync.Mutex | ||
| 35 | workerStarts []int | ||
| 36 | workerStops []int | ||
| 37 | tasksDone []taskResult | ||
| 38 | t *testing.T | ||
| 39 | } | ||
| 40 | |||
| 41 | func (o *mockObserver) OnWorkerStart(workerID int) { | ||
| 42 | o.mu.Lock() | ||
| 43 | defer o.mu.Unlock() | ||
| 44 | o.workerStarts = append(o.workerStarts, workerID) | ||
| 45 | if o.t != nil { | ||
| 46 | o.t.Logf("observer: worker %d started", workerID) | ||
| 47 | } | ||
| 48 | } | ||
| 49 | |||
| 50 | func (o *mockObserver) OnWorkerStop(workerID int) { | ||
| 51 | o.mu.Lock() | ||
| 52 | defer o.mu.Unlock() | ||
| 53 | o.workerStops = append(o.workerStops, workerID) | ||
| 54 | if o.t != nil { | ||
| 55 | o.t.Logf("observer: worker %d stopped", workerID) | ||
| 56 | } | ||
| 57 | } | ||
| 58 | |||
| 59 | func (o *mockObserver) OnTaskDone(workerID int, taskID string, err error) { | ||
| 60 | o.mu.Lock() | ||
| 61 | defer o.mu.Unlock() | ||
| 62 | o.tasksDone = append(o.tasksDone, taskResult{taskID: taskID, err: err}) | ||
| 63 | if o.t != nil { | ||
| 64 | if err != nil { | ||
| 65 | o.t.Logf("observer: worker %d task %s failed: %v", workerID, taskID, err) | ||
| 66 | } else { | ||
| 67 | o.t.Logf("observer: worker %d task %s succeeded", workerID, taskID) | ||
| 68 | } | ||
| 69 | } | ||
| 70 | } | ||
| 71 | |||
| 72 | func TestNewWorkerPool(t *testing.T) { | ||
| 73 | t.Parallel() | ||
| 74 | |||
| 75 | tests := []struct { | ||
| 76 | name string | ||
| 77 | opts []Option | ||
| 78 | shouldErr bool | ||
| 79 | }{ | ||
| 80 | { | ||
| 81 | name: "valid config with defaults", | ||
| 82 | opts: []Option{}, | ||
| 83 | shouldErr: false, | ||
| 84 | }, | ||
| 85 | { | ||
| 86 | name: "valid config with all options", | ||
| 87 | opts: []Option{ | ||
| 88 | WithName("test-pool"), | ||
| 89 | WithNumWorkers(4), | ||
| 90 | WithTaskQueueSize(100), | ||
| 91 | }, | ||
| 92 | shouldErr: false, | ||
| 93 | }, | ||
| 94 | { | ||
| 95 | name: "invalid number of workers (zero)", | ||
| 96 | opts: []Option{WithNumWorkers(0)}, | ||
| 97 | shouldErr: true, | ||
| 98 | }, | ||
| 99 | { | ||
| 100 | name: "invalid number of workers (negative)", | ||
| 101 | opts: []Option{WithNumWorkers(-1)}, | ||
| 102 | shouldErr: true, | ||
| 103 | }, | ||
| 104 | { | ||
| 105 | name: "invalid task queue size (zero)", | ||
| 106 | opts: []Option{WithTaskQueueSize(0)}, | ||
| 107 | shouldErr: true, | ||
| 108 | }, | ||
| 109 | { | ||
| 110 | name: "invalid task queue size (negative)", | ||
| 111 | opts: []Option{WithTaskQueueSize(-1)}, | ||
| 112 | shouldErr: true, | ||
| 113 | }, | ||
| 114 | } | ||
| 115 | |||
| 116 | for _, test := range tests { | ||
| 117 | t.Run(test.name, func(t *testing.T) { | ||
| 118 | t.Parallel() | ||
| 119 | _, err := NewWorkerPool[*mockTask](t.Context(), test.opts...) | ||
| 120 | if (err != nil) != test.shouldErr { | ||
| 121 | t.Errorf("NewWorkerPool() error = %v, shouldErr %v", err, test.shouldErr) | ||
| 122 | } | ||
| 123 | }) | ||
| 124 | } | ||
| 125 | } | ||
| 126 | |||
| 127 | func TestWorkerPool_Observer(t *testing.T) { | ||
| 128 | t.Parallel() | ||
| 129 | ctx := t.Context() | ||
| 130 | |||
| 131 | obs := &mockObserver{t: t} | ||
| 132 | numWorkers := 2 | ||
| 133 | taskCount := 5 | ||
| 134 | |||
| 135 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 136 | WithName("observer-test"), | ||
| 137 | WithNumWorkers(numWorkers), | ||
| 138 | WithTaskQueueSize(10), | ||
| 139 | WithObserver(obs), | ||
| 140 | ) | ||
| 141 | if err != nil { | ||
| 142 | t.Fatalf("failed to create pool: %v", err) | ||
| 143 | } | ||
| 144 | |||
| 145 | pool.Start() | ||
| 146 | |||
| 147 | for i := range taskCount { | ||
| 148 | task := &mockTask{id: fmt.Sprintf("task-%d", i)} | ||
| 149 | if err := pool.Push(task); err != nil { | ||
| 150 | t.Errorf("failed to push task: %v", err) | ||
| 151 | } | ||
| 152 | } | ||
| 153 | |||
| 154 | pool.StopAndWait() | ||
| 155 | |||
| 156 | obs.mu.Lock() | ||
| 157 | defer obs.mu.Unlock() | ||
| 158 | |||
| 159 | if len(obs.workerStarts) != numWorkers { | ||
| 160 | t.Errorf("expected %d worker starts, got %d", numWorkers, len(obs.workerStarts)) | ||
| 161 | } | ||
| 162 | if len(obs.workerStops) != numWorkers { | ||
| 163 | t.Errorf("expected %d worker stops, got %d", numWorkers, len(obs.workerStops)) | ||
| 164 | } | ||
| 165 | |||
| 166 | if len(obs.tasksDone) != taskCount { | ||
| 167 | t.Errorf("expected %d tasks done, got %d", taskCount, len(obs.tasksDone)) | ||
| 168 | } | ||
| 169 | for _, r := range obs.tasksDone { | ||
| 170 | if r.err != nil { | ||
| 171 | t.Errorf("task %s should have succeeded, got error: %v", r.taskID, r.err) | ||
| 172 | } | ||
| 173 | } | ||
| 174 | } | ||
| 175 | |||
| 176 | func TestWorkerPool_ObserverTaskFailed(t *testing.T) { | ||
| 177 | t.Parallel() | ||
| 178 | ctx := t.Context() | ||
| 179 | |||
| 180 | obs := &mockObserver{t: t} | ||
| 181 | |||
| 182 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 183 | WithName("observer-fail-test"), | ||
| 184 | WithNumWorkers(1), | ||
| 185 | WithTaskQueueSize(10), | ||
| 186 | WithObserver(obs), | ||
| 187 | ) | ||
| 188 | if err != nil { | ||
| 189 | t.Fatalf("failed to create pool: %v", err) | ||
| 190 | } | ||
| 191 | |||
| 192 | pool.Start() | ||
| 193 | |||
| 194 | task := &mockTask{ | ||
| 195 | id: "fail-task", | ||
| 196 | fn: func(ctx context.Context) error { | ||
| 197 | return fmt.Errorf("intentional error") | ||
| 198 | }, | ||
| 199 | } | ||
| 200 | pool.Push(task) | ||
| 201 | |||
| 202 | pool.StopAndWait() | ||
| 203 | |||
| 204 | obs.mu.Lock() | ||
| 205 | defer obs.mu.Unlock() | ||
| 206 | |||
| 207 | failures := 0 | ||
| 208 | for _, r := range obs.tasksDone { | ||
| 209 | if r.err != nil { | ||
| 210 | failures++ | ||
| 211 | } | ||
| 212 | } | ||
| 213 | |||
| 214 | if failures != 1 { | ||
| 215 | t.Errorf("expected 1 failed task, got %d", failures) | ||
| 216 | } | ||
| 217 | if len(obs.tasksDone) > 0 && obs.tasksDone[0].taskID != "fail-task" { | ||
| 218 | t.Errorf("expected failed task id 'fail-task', got %s", obs.tasksDone[0].taskID) | ||
| 219 | } | ||
| 220 | } | ||
| 221 | |||
| 222 | func TestWorkerPool_ObserverTaskPanic(t *testing.T) { | ||
| 223 | t.Parallel() | ||
| 224 | ctx := t.Context() | ||
| 225 | |||
| 226 | panicCaught := make(chan string, 1) | ||
| 227 | |||
| 228 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 229 | WithName("observer-panic-test"), | ||
| 230 | WithNumWorkers(1), | ||
| 231 | WithTaskQueueSize(10), | ||
| 232 | WithPanicHandler(func(workerID int, taskID string, panicValue any) { | ||
| 233 | panicCaught <- taskID | ||
| 234 | }), | ||
| 235 | ) | ||
| 236 | if err != nil { | ||
| 237 | t.Fatalf("failed to create pool: %v", err) | ||
| 238 | } | ||
| 239 | |||
| 240 | pool.Start() | ||
| 241 | |||
| 242 | task := &mockTask{ | ||
| 243 | id: "panic-task", | ||
| 244 | fn: func(ctx context.Context) error { | ||
| 245 | panic("intentional panic") | ||
| 246 | }, | ||
| 247 | } | ||
| 248 | pool.Push(task) | ||
| 249 | |||
| 250 | pool.StopAndWait() | ||
| 251 | |||
| 252 | select { | ||
| 253 | case taskID := <-panicCaught: | ||
| 254 | if taskID != "panic-task" { | ||
| 255 | t.Errorf("expected panic from 'panic-task', got %s", taskID) | ||
| 256 | } | ||
| 257 | default: | ||
| 258 | t.Error("expected PanicHandler to be called") | ||
| 259 | } | ||
| 260 | } | ||
| 261 | |||
| 262 | func TestWorkerPool_ConcurrentPush(t *testing.T) { | ||
| 263 | t.Parallel() | ||
| 264 | ctx := t.Context() | ||
| 265 | var wg sync.WaitGroup | ||
| 266 | taskCount := 20 | ||
| 267 | |||
| 268 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 269 | WithName("concurrent-push"), | ||
| 270 | WithNumWorkers(3), | ||
| 271 | WithTaskQueueSize(taskCount), | ||
| 272 | ) | ||
| 273 | if err != nil { | ||
| 274 | t.Fatalf("failed to create pool: %v", err) | ||
| 275 | } | ||
| 276 | |||
| 277 | pool.Start() | ||
| 278 | defer pool.StopAndWait() | ||
| 279 | |||
| 280 | for i := range taskCount { | ||
| 281 | wg.Add(1) | ||
| 282 | go func(id int) { | ||
| 283 | defer wg.Done() | ||
| 284 | task := &mockTask{id: fmt.Sprintf("task-%d", id)} | ||
| 285 | if err := pool.Push(task); err != nil { | ||
| 286 | t.Errorf("push failed: %v", err) | ||
| 287 | } | ||
| 288 | }(i) | ||
| 289 | } | ||
| 290 | |||
| 291 | done := make(chan struct{}) | ||
| 292 | go func() { | ||
| 293 | wg.Wait() | ||
| 294 | close(done) | ||
| 295 | }() | ||
| 296 | |||
| 297 | select { | ||
| 298 | case <-done: | ||
| 299 | case <-time.After(5 * time.Second): | ||
| 300 | t.Fatal("timeout waiting for concurrent push") | ||
| 301 | } | ||
| 302 | } | ||
| 303 | |||
| 304 | func TestWorkerPool_TaskQueueFull(t *testing.T) { | ||
| 305 | t.Parallel() | ||
| 306 | ctx := t.Context() | ||
| 307 | |||
| 308 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 309 | WithName("queue-full"), | ||
| 310 | WithNumWorkers(1), | ||
| 311 | WithTaskQueueSize(2), | ||
| 312 | ) | ||
| 313 | if err != nil { | ||
| 314 | t.Fatalf("failed to create pool: %v", err) | ||
| 315 | } | ||
| 316 | |||
| 317 | for i := range 2 { | ||
| 318 | task := &mockTask{id: fmt.Sprintf("task-%d", i)} | ||
| 319 | if err := pool.Push(task); err != nil { | ||
| 320 | t.Errorf("unexpected error on push %d: %v", i, err) | ||
| 321 | } | ||
| 322 | } | ||
| 323 | |||
| 324 | task := &mockTask{id: "overflow-task"} | ||
| 325 | err = pool.Push(task) | ||
| 326 | if err == nil { | ||
| 327 | t.Fatal("expected error when queue is full, got nil") | ||
| 328 | } | ||
| 329 | if !errors.Is(err, ErrQueueFull) { | ||
| 330 | t.Fatalf("expected ErrQueueFull, got: %v", err) | ||
| 331 | } | ||
| 332 | } | ||
| 333 | |||
| 334 | func TestWorkerPool_MultipleStop(t *testing.T) { | ||
| 335 | t.Parallel() | ||
| 336 | ctx := t.Context() | ||
| 337 | |||
| 338 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 339 | WithName("multiple-stop"), | ||
| 340 | WithNumWorkers(2), | ||
| 341 | WithTaskQueueSize(10), | ||
| 342 | ) | ||
| 343 | if err != nil { | ||
| 344 | t.Fatalf("failed to create pool: %v", err) | ||
| 345 | } | ||
| 346 | |||
| 347 | pool.Start() | ||
| 348 | pool.Stop() | ||
| 349 | |||
| 350 | func() { | ||
| 351 | defer func() { | ||
| 352 | if r := recover(); r != nil { | ||
| 353 | t.Errorf("Stop() panicked on second call: %v", r) | ||
| 354 | } | ||
| 355 | }() | ||
| 356 | pool.Stop() | ||
| 357 | }() | ||
| 358 | |||
| 359 | pool.Wait() | ||
| 360 | } | ||
| 361 | |||
| 362 | func TestWorkerPool_MetricsThreadSafe(t *testing.T) { | ||
| 363 | t.Parallel() | ||
| 364 | ctx := t.Context() | ||
| 365 | var wg sync.WaitGroup | ||
| 366 | taskCount := 20 | ||
| 367 | |||
| 368 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 369 | WithName("metrics-safe"), | ||
| 370 | WithNumWorkers(10), | ||
| 371 | WithTaskQueueSize(taskCount), | ||
| 372 | ) | ||
| 373 | if err != nil { | ||
| 374 | t.Fatalf("failed to create pool: %v", err) | ||
| 375 | } | ||
| 376 | |||
| 377 | pool.Start() | ||
| 378 | defer pool.StopAndWait() | ||
| 379 | |||
| 380 | for i := range taskCount { | ||
| 381 | wg.Add(1) | ||
| 382 | go func(id int) { | ||
| 383 | defer wg.Done() | ||
| 384 | task := &mockTask{ | ||
| 385 | id: fmt.Sprintf("task-%d", id), | ||
| 386 | fn: func(ctx context.Context) error { | ||
| 387 | time.Sleep(time.Millisecond) | ||
| 388 | return nil | ||
| 389 | }, | ||
| 390 | } | ||
| 391 | pool.Push(task) | ||
| 392 | }(i) | ||
| 393 | } | ||
| 394 | |||
| 395 | concurrentReads := 100 | ||
| 396 | for range concurrentReads { | ||
| 397 | wg.Add(1) | ||
| 398 | go func() { | ||
| 399 | defer wg.Done() | ||
| 400 | for range 10 { | ||
| 401 | pool.GetMetrics() | ||
| 402 | } | ||
| 403 | }() | ||
| 404 | } | ||
| 405 | |||
| 406 | wg.Wait() | ||
| 407 | } | ||
| 408 | |||
| 409 | func TestWorkerPool_ContextCancellation(t *testing.T) { | ||
| 410 | t.Parallel() | ||
| 411 | ctx, cancel := context.WithCancel(t.Context()) | ||
| 412 | |||
| 413 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 414 | WithName("context-cancel"), | ||
| 415 | WithNumWorkers(2), | ||
| 416 | WithTaskQueueSize(10), | ||
| 417 | ) | ||
| 418 | if err != nil { | ||
| 419 | t.Fatalf("failed to create pool: %v", err) | ||
| 420 | } | ||
| 421 | |||
| 422 | pool.Start() | ||
| 423 | |||
| 424 | for i := range 5 { | ||
| 425 | task := &mockTask{id: fmt.Sprintf("task-%d", i)} | ||
| 426 | pool.Push(task) | ||
| 427 | } | ||
| 428 | |||
| 429 | cancel() | ||
| 430 | |||
| 431 | err = pool.Wait() | ||
| 432 | if err != nil && !errors.Is(err, context.Canceled) { | ||
| 433 | t.Errorf("expected context.Canceled, got: %v", err) | ||
| 434 | } | ||
| 435 | } | ||
| 436 | |||
| 437 | func TestWorkerPool_ErrorHandler(t *testing.T) { | ||
| 438 | t.Parallel() | ||
| 439 | ctx := t.Context() | ||
| 440 | |||
| 441 | errorHandlerCalled := make(chan struct{}) | ||
| 442 | |||
| 443 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 444 | WithName("error-handler-test"), | ||
| 445 | WithNumWorkers(1), | ||
| 446 | WithTaskQueueSize(10), | ||
| 447 | WithErrorHandler(func(workerID int, taskID string, err error) { | ||
| 448 | t.Logf("ErrorHandler called: worker=%d, task=%s, err=%v", workerID, taskID, err) | ||
| 449 | close(errorHandlerCalled) | ||
| 450 | }), | ||
| 451 | ) | ||
| 452 | if err != nil { | ||
| 453 | t.Fatalf("failed to create pool: %v", err) | ||
| 454 | } | ||
| 455 | |||
| 456 | pool.Start() | ||
| 457 | |||
| 458 | task := &mockTask{ | ||
| 459 | id: "error-task", | ||
| 460 | fn: func(ctx context.Context) error { | ||
| 461 | return fmt.Errorf("task error") | ||
| 462 | }, | ||
| 463 | } | ||
| 464 | pool.Push(task) | ||
| 465 | |||
| 466 | select { | ||
| 467 | case <-errorHandlerCalled: | ||
| 468 | case <-time.After(time.Second): | ||
| 469 | t.Fatal("ErrorHandler was not called within timeout") | ||
| 470 | } | ||
| 471 | |||
| 472 | pool.StopAndWait() | ||
| 473 | } | ||
| 474 | |||
| 475 | func TestWorkerPool_PanicHandler(t *testing.T) { | ||
| 476 | t.Parallel() | ||
| 477 | ctx := t.Context() | ||
| 478 | |||
| 479 | panicHandlerCalled := make(chan struct{}) | ||
| 480 | |||
| 481 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 482 | WithName("panic-handler-test"), | ||
| 483 | WithNumWorkers(1), | ||
| 484 | WithTaskQueueSize(10), | ||
| 485 | WithPanicHandler(func(workerID int, taskID string, panicValue any) { | ||
| 486 | t.Logf("PanicHandler called: worker=%d, task=%s, panic=%v", workerID, taskID, panicValue) | ||
| 487 | close(panicHandlerCalled) | ||
| 488 | }), | ||
| 489 | ) | ||
| 490 | if err != nil { | ||
| 491 | t.Fatalf("failed to create pool: %v", err) | ||
| 492 | } | ||
| 493 | |||
| 494 | pool.Start() | ||
| 495 | |||
| 496 | task := &mockTask{ | ||
| 497 | id: "panic-task", | ||
| 498 | fn: func(ctx context.Context) error { | ||
| 499 | panic("task panic") | ||
| 500 | }, | ||
| 501 | } | ||
| 502 | pool.Push(task) | ||
| 503 | |||
| 504 | select { | ||
| 505 | case <-panicHandlerCalled: | ||
| 506 | case <-time.After(time.Second): | ||
| 507 | t.Fatal("PanicHandler was not called within timeout") | ||
| 508 | } | ||
| 509 | |||
| 510 | pool.StopAndWait() | ||
| 511 | } | ||
| 512 | |||
| 513 | func TestWorkerPool_PushContextCancelled(t *testing.T) { | ||
| 514 | t.Parallel() | ||
| 515 | ctx, cancel := context.WithCancel(t.Context()) | ||
| 516 | |||
| 517 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 518 | WithName("push-ctx-cancelled"), | ||
| 519 | WithNumWorkers(1), | ||
| 520 | WithTaskQueueSize(1), | ||
| 521 | ) | ||
| 522 | if err != nil { | ||
| 523 | t.Fatalf("failed to create pool: %v", err) | ||
| 524 | } | ||
| 525 | |||
| 526 | task1 := &mockTask{id: "task-1"} | ||
| 527 | if err := pool.Push(task1); err != nil { | ||
| 528 | t.Fatalf("failed to fill queue: %v", err) | ||
| 529 | } | ||
| 530 | |||
| 531 | cancel() | ||
| 532 | |||
| 533 | task2 := &mockTask{id: "task-2"} | ||
| 534 | err = pool.Push(task2) | ||
| 535 | |||
| 536 | if err == nil { | ||
| 537 | t.Fatal("expected error when pushing to full queue with cancelled context") | ||
| 538 | } | ||
| 539 | |||
| 540 | if !errors.Is(err, ErrQueueFull) && !errors.Is(err, context.Canceled) { | ||
| 541 | t.Fatalf("expected ErrQueueFull or context.Canceled, got: %v", err) | ||
| 542 | } | ||
| 543 | } | ||
| 544 | |||
| 545 | func TestWorkerPool_NoObserver(t *testing.T) { | ||
| 546 | t.Parallel() | ||
| 547 | ctx := t.Context() | ||
| 548 | |||
| 549 | pool, err := NewWorkerPool[*mockTask](ctx, | ||
| 550 | WithName("no-observer"), | ||
| 551 | WithNumWorkers(2), | ||
| 552 | WithTaskQueueSize(10), | ||
| 553 | ) | ||
| 554 | if err != nil { | ||
| 555 | t.Fatalf("failed to create pool: %v", err) | ||
| 556 | } | ||
| 557 | |||
| 558 | pool.Start() | ||
| 559 | |||
| 560 | task1 := &mockTask{id: "success-task"} | ||
| 561 | pool.Push(task1) | ||
| 562 | |||
| 563 | task2 := &mockTask{ | ||
| 564 | id: "fail-task", | ||
| 565 | fn: func(ctx context.Context) error { | ||
| 566 | return fmt.Errorf("error") | ||
| 567 | }, | ||
| 568 | } | ||
| 569 | pool.Push(task2) | ||
| 570 | |||
| 571 | task3 := &mockTask{ | ||
| 572 | id: "panic-task", | ||
| 573 | fn: func(ctx context.Context) error { | ||
| 574 | panic("panic") | ||
| 575 | }, | ||
| 576 | } | ||
| 577 | pool.Push(task3) | ||
| 578 | |||
| 579 | pool.StopAndWait() | ||
| 580 | |||
| 581 | metrics := pool.GetMetrics() | ||
| 582 | if metrics.TasksFailed != 2 { | ||
| 583 | t.Errorf("expected 2 failed tasks, got %d", metrics.TasksFailed) | ||
| 584 | } | ||
| 585 | } | ||
| 586 | |||
| 587 | func TestWorkerPool_StartMultipleTimes(t *testing.T) { | ||
| 588 | t.Parallel() | ||
| 589 | |||
| 590 | pool, err := NewWorkerPool[*mockTask](t.Context(), | ||
| 591 | WithNumWorkers(1), | ||
| 592 | WithTaskQueueSize(1), | ||
| 593 | ) | ||
| 594 | if err != nil { | ||
| 595 | t.Fatalf("failed to create pool: %v", err) | ||
| 596 | } | ||
| 597 | |||
| 598 | pool.Start() | ||
| 599 | defer pool.StopAndWait() | ||
| 600 | |||
| 601 | defer func() { | ||
| 602 | if r := recover(); r == nil { | ||
| 603 | t.Error("expected panic when calling Start() multiple times") | ||
| 604 | } | ||
| 605 | }() | ||
| 606 | |||
| 607 | pool.Start() | ||
| 608 | } | ||
