diff --git a/worker.go b/worker.go index 00a8780..c312bde 100644 --- a/worker.go +++ b/worker.go @@ -21,9 +21,7 @@ func worker(context context.Context, waitGroup *sync.WaitGroup, firstTask func() select { case <-context.Done(): // Pool context was cancelled, empty tasks channel and exit - for _ = range tasks { - taskWaitGroup.Done() - } + drainTasks(tasks, taskWaitGroup) return case task, ok := <-tasks: // Prioritize context.Done statement (https://stackoverflow.com/questions/46200343/force-priority-of-go-select-statement) @@ -45,3 +43,10 @@ func worker(context context.Context, waitGroup *sync.WaitGroup, firstTask func() } } } + +// drainPendingTasks discards queued tasks and decrements the corresponding wait group +func drainTasks(tasks <-chan func(), tasksWaitGroup *sync.WaitGroup) { + for _ = range tasks { + tasksWaitGroup.Done() + } +}