-
-
Notifications
You must be signed in to change notification settings - Fork 62
/
pond.go
556 lines (453 loc) · 14.4 KB
/
pond.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
package pond
import (
"context"
"errors"
"fmt"
"runtime/debug"
"sync"
"sync/atomic"
"time"
)
const (
// defaultIdleTimeout defines the default idle timeout to use when not explicitly specified
// via the IdleTimeout() option
defaultIdleTimeout = 5 * time.Second
)
var (
// ErrSubmitOnStoppedPool is thrown when attempting to submit a task to a pool that has been stopped
ErrSubmitOnStoppedPool = errors.New("worker pool has been stopped and is no longer accepting tasks")
)
// defaultPanicHandler is the default panic handler
func defaultPanicHandler(panic interface{}) {
fmt.Printf("Worker exits from a panic: %v\nStack trace: %s\n", panic, string(debug.Stack()))
}
// ResizingStrategy represents a pool resizing strategy
type ResizingStrategy interface {
Resize(runningWorkers, minWorkers, maxWorkers int) bool
}
// Option represents an option that can be passed when instantiating a worker pool to customize it
type Option func(*WorkerPool)
// IdleTimeout allows to change the idle timeout for a worker pool
func IdleTimeout(idleTimeout time.Duration) Option {
return func(pool *WorkerPool) {
pool.idleTimeout = idleTimeout
}
}
// MinWorkers allows to change the minimum number of workers of a worker pool
func MinWorkers(minWorkers int) Option {
return func(pool *WorkerPool) {
pool.minWorkers = minWorkers
}
}
// Strategy allows to change the strategy used to resize the pool
func Strategy(strategy ResizingStrategy) Option {
return func(pool *WorkerPool) {
pool.strategy = strategy
}
}
// PanicHandler allows to change the panic handler function of a worker pool
func PanicHandler(panicHandler func(interface{})) Option {
return func(pool *WorkerPool) {
pool.panicHandler = panicHandler
}
}
// Context configures a parent context on a worker pool to stop all workers when it is cancelled
func Context(parentCtx context.Context) Option {
return func(pool *WorkerPool) {
pool.context, pool.contextCancel = context.WithCancel(parentCtx)
}
}
// WorkerPool models a pool of workers
type WorkerPool struct {
// Atomic counters, should be placed first so alignment is guaranteed
// for atomic operations.
workerCount int32
idleWorkerCount int32
waitingTaskCount uint64
submittedTaskCount uint64
successfulTaskCount uint64
failedTaskCount uint64
// Configurable settings
maxWorkers int
maxCapacity int
minWorkers int
idleTimeout time.Duration
strategy ResizingStrategy
panicHandler func(interface{})
context context.Context
contextCancel context.CancelFunc
// Private properties
tasks chan func()
tasksCloseOnce sync.Once
workersWaitGroup sync.WaitGroup
tasksWaitGroup sync.WaitGroup
mutex sync.Mutex
stopped int32
}
// New creates a worker pool with that can scale up to the given maximum number of workers (maxWorkers).
// The maxCapacity parameter determines the number of tasks that can be submitted to this pool without blocking,
// because it defines the size of the buffered channel used to receive tasks.
// The options parameter can take a list of functions to customize configuration values on this worker pool.
func New(maxWorkers, maxCapacity int, options ...Option) *WorkerPool {
// Instantiate the pool
pool := &WorkerPool{
maxWorkers: maxWorkers,
maxCapacity: maxCapacity,
idleTimeout: defaultIdleTimeout,
strategy: Eager(),
panicHandler: defaultPanicHandler,
}
// Apply all options
for _, opt := range options {
opt(pool)
}
// Make sure options are consistent
if pool.maxWorkers <= 0 {
pool.maxWorkers = 1
}
if pool.minWorkers > pool.maxWorkers {
pool.minWorkers = pool.maxWorkers
}
if pool.maxCapacity < 0 {
pool.maxCapacity = 0
}
if pool.idleTimeout < 0 {
pool.idleTimeout = defaultIdleTimeout
}
// Initialize base context (if not already set)
if pool.context == nil {
Context(context.Background())(pool)
}
// Create tasks channel
pool.tasks = make(chan func(), pool.maxCapacity)
// Start purger goroutine
pool.workersWaitGroup.Add(1)
go pool.purge()
// Start minWorkers workers
if pool.minWorkers > 0 {
for i := 0; i < pool.minWorkers; i++ {
pool.maybeStartWorker(nil)
}
}
return pool
}
// RunningWorkers returns the current number of running workers
func (p *WorkerPool) RunningWorkers() int {
return int(atomic.LoadInt32(&p.workerCount))
}
// IdleWorkers returns the current number of idle workers
func (p *WorkerPool) IdleWorkers() int {
return int(atomic.LoadInt32(&p.idleWorkerCount))
}
// MinWorkers returns the minimum number of worker goroutines
func (p *WorkerPool) MinWorkers() int {
return p.minWorkers
}
// MaxWorkers returns the maximum number of worker goroutines
func (p *WorkerPool) MaxWorkers() int {
return p.maxWorkers
}
// MaxCapacity returns the maximum number of tasks that can be waiting in the queue
// at any given time (queue size)
func (p *WorkerPool) MaxCapacity() int {
return p.maxCapacity
}
// Strategy returns the configured pool resizing strategy
func (p *WorkerPool) Strategy() ResizingStrategy {
return p.strategy
}
// SubmittedTasks returns the total number of tasks submitted since the pool was created
func (p *WorkerPool) SubmittedTasks() uint64 {
return atomic.LoadUint64(&p.submittedTaskCount)
}
// WaitingTasks returns the current number of tasks in the queue that are waiting to be executed
func (p *WorkerPool) WaitingTasks() uint64 {
return atomic.LoadUint64(&p.waitingTaskCount)
}
// SuccessfulTasks returns the total number of tasks that have successfully completed their exection
// since the pool was created
func (p *WorkerPool) SuccessfulTasks() uint64 {
return atomic.LoadUint64(&p.successfulTaskCount)
}
// FailedTasks returns the total number of tasks that completed with panic since the pool was created
func (p *WorkerPool) FailedTasks() uint64 {
return atomic.LoadUint64(&p.failedTaskCount)
}
// CompletedTasks returns the total number of tasks that have completed their exection either successfully
// or with panic since the pool was created
func (p *WorkerPool) CompletedTasks() uint64 {
return p.SuccessfulTasks() + p.FailedTasks()
}
// Stopped returns true if the pool has been stopped and is no longer accepting tasks, and false otherwise.
func (p *WorkerPool) Stopped() bool {
return atomic.LoadInt32(&p.stopped) == 1
}
// Submit sends a task to this worker pool for execution. If the queue is full,
// it will wait until the task is dispatched to a worker goroutine.
func (p *WorkerPool) Submit(task func()) {
p.submit(task, true)
}
// TrySubmit attempts to send a task to this worker pool for execution. If the queue is full,
// it will not wait for a worker to become idle. It returns true if it was able to dispatch
// the task and false otherwise.
func (p *WorkerPool) TrySubmit(task func()) bool {
return p.submit(task, false)
}
func (p *WorkerPool) submit(task func(), mustSubmit bool) (submitted bool) {
if task == nil {
return
}
if p.Stopped() {
// Pool is stopped and caller must submit the task
if mustSubmit {
panic(ErrSubmitOnStoppedPool)
}
return
}
// Increment submitted and waiting task counters as soon as we receive a task
atomic.AddUint64(&p.submittedTaskCount, 1)
atomic.AddUint64(&p.waitingTaskCount, 1)
p.tasksWaitGroup.Add(1)
defer func() {
if !submitted {
// Task was not sumitted to the pool, decrement submitted and waiting task counters
atomic.AddUint64(&p.submittedTaskCount, ^uint64(0))
atomic.AddUint64(&p.waitingTaskCount, ^uint64(0))
p.tasksWaitGroup.Done()
}
}()
// Start a worker as long as we haven't reached the limit
if submitted = p.maybeStartWorker(task); submitted {
return
}
if !mustSubmit {
// Attempt to dispatch to an idle worker without blocking
select {
case p.tasks <- task:
submitted = true
return
default:
// Channel is full and can't wait for an idle worker, so need to exit
return
}
}
// Submit the task to the tasks channel and wait for it to be picked up by a worker
p.tasks <- task
submitted = true
return
}
// SubmitAndWait sends a task to this worker pool for execution and waits for it to complete
// before returning
func (p *WorkerPool) SubmitAndWait(task func()) {
if task == nil {
return
}
done := make(chan struct{})
p.Submit(func() {
defer close(done)
task()
})
<-done
}
// SubmitBefore attempts to send a task for execution to this worker pool but aborts it
// if the task did not start before the given deadline.
func (p *WorkerPool) SubmitBefore(task func(), deadline time.Duration) {
if task == nil {
return
}
timer := time.NewTimer(deadline)
p.Submit(func() {
select {
case <-timer.C:
// Deadline was reached, abort the task
default:
// Deadline not reached, execute the task
defer timer.Stop()
task()
}
})
}
// Stop causes this pool to stop accepting new tasks and signals all workers to exit.
// Tasks being executed by workers will continue until completion (unless the process is terminated).
// Tasks in the queue will not be executed.
// This method returns a context object that is cancelled when the pool has stopped completely.
func (p *WorkerPool) Stop() context.Context {
ctx, cancel := context.WithCancel(context.Background())
go func() {
p.stop(false)
cancel()
}()
return ctx
}
// StopAndWait causes this pool to stop accepting new tasks and then waits for all tasks in the queue
// to complete before returning.
func (p *WorkerPool) StopAndWait() {
p.stop(true)
}
// StopAndWaitFor stops this pool and waits until either all tasks in the queue are completed
// or the given deadline is reached, whichever comes first.
func (p *WorkerPool) StopAndWaitFor(deadline time.Duration) {
// Launch goroutine to detect when worker pool has stopped gracefully
workersDone := make(chan struct{})
go func() {
p.stop(true)
workersDone <- struct{}{}
}()
// Wait until either all workers have exited or the deadline is reached
select {
case <-workersDone:
return
case <-time.After(deadline):
p.contextCancel()
return
}
}
func (p *WorkerPool) stop(waitForQueuedTasksToComplete bool) {
// Mark pool as stopped
atomic.StoreInt32(&p.stopped, 1)
if waitForQueuedTasksToComplete {
// Wait for all queued tasks to complete
p.tasksWaitGroup.Wait()
}
// Reset worker count
p.resetWorkerCount()
// Terminate all workers & purger goroutine
p.contextCancel()
// Wait for all workers & purger goroutine to exit
p.workersWaitGroup.Wait()
// close tasks channel (only once, in case multiple concurrent calls to StopAndWait are made)
p.tasksCloseOnce.Do(func() {
close(p.tasks)
})
}
// purge represents the work done by the purger goroutine
func (p *WorkerPool) purge() {
defer p.workersWaitGroup.Done()
idleTicker := time.NewTicker(p.idleTimeout)
defer idleTicker.Stop()
for {
select {
// Timed out waiting for any activity to happen, attempt to stop an idle worker
case <-idleTicker.C:
p.maybeStopIdleWorker()
// Pool context was cancelled, exit
case <-p.context.Done():
return
}
}
}
// maybeStopIdleWorker attempts to stop an idle worker by sending it a nil task
func (p *WorkerPool) maybeStopIdleWorker() bool {
if decremented := p.decrementWorkerCount(); !decremented {
return false
}
// Send a nil task to stop an idle worker
p.tasks <- nil
return true
}
// maybeStartWorker attempts to create a new worker goroutine to run the given task.
// If the worker pool has reached the maximum number of workers or there are idle workers,
// it will not create a new one.
func (p *WorkerPool) maybeStartWorker(firstTask func()) bool {
if incremented := p.incrementWorkerCount(); !incremented {
return false
}
if firstTask == nil {
// Worker starts idle
atomic.AddInt32(&p.idleWorkerCount, 1)
}
// Launch worker goroutine
go worker(p.context, &p.workersWaitGroup, firstTask, p.tasks, p.executeTask, &p.tasksWaitGroup)
return true
}
// executeTask executes the given task and updates task-related counters
func (p *WorkerPool) executeTask(task func(), isFirstTask bool) {
defer func() {
if panic := recover(); panic != nil {
// Increment failed task count
atomic.AddUint64(&p.failedTaskCount, 1)
// Invoke panic handler
p.panicHandler(panic)
// Increment idle count
atomic.AddInt32(&p.idleWorkerCount, 1)
}
p.tasksWaitGroup.Done()
}()
// Decrement idle count
if !isFirstTask {
atomic.AddInt32(&p.idleWorkerCount, -1)
}
// Decrement waiting task count
atomic.AddUint64(&p.waitingTaskCount, ^uint64(0))
// Execute task
task()
// Increment successful task count
atomic.AddUint64(&p.successfulTaskCount, 1)
// Increment idle count
atomic.AddInt32(&p.idleWorkerCount, 1)
}
func (p *WorkerPool) incrementWorkerCount() bool {
p.mutex.Lock()
defer p.mutex.Unlock()
runningWorkerCount := p.RunningWorkers()
// Reached max workers, do not create a new one
if runningWorkerCount >= p.maxWorkers {
return false
}
// Idle workers available, do not create a new one
if runningWorkerCount >= p.minWorkers && runningWorkerCount > 0 && p.IdleWorkers() > 0 {
return false
}
// Execute the resizing strategy to determine if we should create more workers
if resize := p.strategy.Resize(runningWorkerCount, p.minWorkers, p.maxWorkers); !resize {
return false
}
// Increment worker count
atomic.AddInt32(&p.workerCount, 1)
// Increment wait group
p.workersWaitGroup.Add(1)
return true
}
func (p *WorkerPool) decrementWorkerCount() bool {
p.mutex.Lock()
defer p.mutex.Unlock()
if p.IdleWorkers() <= 0 || p.RunningWorkers() <= p.minWorkers || p.Stopped() {
return false
}
// Decrement worker count
atomic.AddInt32(&p.workerCount, -1)
// Decrement idle count
atomic.AddInt32(&p.idleWorkerCount, -1)
return true
}
func (p *WorkerPool) resetWorkerCount() {
p.mutex.Lock()
defer p.mutex.Unlock()
// Reset worker count
atomic.StoreInt32(&p.workerCount, 0)
// Reset idle count
atomic.StoreInt32(&p.idleWorkerCount, 0)
}
// Group creates a new task group
func (p *WorkerPool) Group() *TaskGroup {
return &TaskGroup{
pool: p,
}
}
// GroupContext creates a new task group and an associated Context derived from ctx.
//
// The derived Context is canceled the first time a function submitted to the group
// returns a non-nil error or the first time Wait returns, whichever occurs first.
func (p *WorkerPool) GroupContext(ctx context.Context) (*TaskGroupWithContext, context.Context) {
if ctx == nil {
panic("a non-nil context needs to be specified when using GroupContext")
}
ctx, cancel := context.WithCancel(ctx)
return &TaskGroupWithContext{
TaskGroup: TaskGroup{
pool: p,
},
ctx: ctx,
cancel: cancel,
}, ctx
}