diff --git a/examples/dynamic_size/go.mod b/examples/dynamic_size/go.mod index 882c813..83a60be 100644 --- a/examples/dynamic_size/go.mod +++ b/examples/dynamic_size/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/dynamic_size -go 1.19 +go 1.18 require git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/examples/fixed_size/go.mod b/examples/fixed_size/go.mod index eeadf23..747c279 100644 --- a/examples/fixed_size/go.mod +++ b/examples/fixed_size/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/fixed_size -go 1.19 +go 1.18 require git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/examples/group_context/go.mod b/examples/group_context/go.mod index 5a71afe..ea9eb7e 100644 --- a/examples/group_context/go.mod +++ b/examples/group_context/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/group_context -go 1.19 +go 1.18 require git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/examples/pool_context/go.mod b/examples/pool_context/go.mod index 26db644..309e015 100644 --- a/examples/pool_context/go.mod +++ b/examples/pool_context/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/pool_context -go 1.19 +go 1.18 require git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/examples/prometheus/go.mod b/examples/prometheus/go.mod index 9952b92..49188de 100644 --- a/examples/prometheus/go.mod +++ b/examples/prometheus/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/fixed_size -go 1.19 +go 1.18 require ( git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/examples/task_group/go.mod b/examples/task_group/go.mod index caa31dc..9db2a55 100644 --- a/examples/task_group/go.mod +++ b/examples/task_group/go.mod @@ -1,6 +1,6 @@ module git.company.lan/gopkg/pond/examples/task_group -go 1.19 +go 1.18 require git.company.lan/gopkg/pond v0.0.0-00010101000000-000000000000 diff --git a/go.mod b/go.mod index 850a033..fb42eb8 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,3 @@ module git.company.lan/gopkg/pond -go 1.19 +go 1.18 diff --git a/pond.go b/pond.go index e5e6192..6bd1a6f 100644 --- a/pond.go +++ b/pond.go @@ -364,13 +364,13 @@ func (p *WorkerPool) stop(waitForQueuedTasksToComplete bool) { // 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) }) + + // Wait for all workers & purger goroutine to exit + p.workersWaitGroup.Wait() } // purge represents the work done by the purger goroutine @@ -420,7 +420,7 @@ func (p *WorkerPool) maybeStartWorker(firstTask func()) bool { } // Launch worker goroutine - go worker(p.context, &p.workersWaitGroup, firstTask, p.tasks, p.executeTask) + go worker(p.context, &p.workersWaitGroup, firstTask, p.tasks, p.executeTask, &p.tasksWaitGroup) return true } diff --git a/pond_blackbox_test.go b/pond_blackbox_test.go index 0234909..282d46f 100644 --- a/pond_blackbox_test.go +++ b/pond_blackbox_test.go @@ -542,6 +542,47 @@ func TestSubmitWithContext(t *testing.T) { assertEqual(t, int32(0), atomic.LoadInt32(&doneCount)) } +func TestSubmitWithContextCancelWithIdleTasks(t *testing.T) { + + ctx, cancel := context.WithCancel(context.Background()) + + pool := pond.New(1, 5, pond.Context(ctx)) + + var doneCount, taskCount int32 + + // Submit a long-running, cancellable task + pool.Submit(func() { + atomic.AddInt32(&taskCount, 1) + select { + case <-ctx.Done(): + return + case <-time.After(10 * time.Minute): + atomic.AddInt32(&doneCount, 1) + return + } + }) + + // Submit a long-running, cancellable task + pool.Submit(func() { + atomic.AddInt32(&taskCount, 1) + select { + case <-ctx.Done(): + return + case <-time.After(10 * time.Minute): + atomic.AddInt32(&doneCount, 1) + return + } + }) + + // Cancel the context + cancel() + + pool.StopAndWait() + + assertEqual(t, int32(1), atomic.LoadInt32(&taskCount)) + assertEqual(t, int32(0), atomic.LoadInt32(&doneCount)) +} + func TestConcurrentStopAndWait(t *testing.T) { pool := pond.New(1, 5) diff --git a/pond_test.go b/pond_test.go index 76a4876..e498673 100644 --- a/pond_test.go +++ b/pond_test.go @@ -39,6 +39,9 @@ func TestPurgeAfterPoolStopped(t *testing.T) { pool.SubmitAndWait(func() { atomic.AddInt32(&doneCount, 1) }) + + time.Sleep(10 * time.Millisecond) + assertEqual(t, int32(1), atomic.LoadInt32(&doneCount)) assertEqual(t, 1, pool.RunningWorkers()) @@ -59,6 +62,8 @@ func TestPurgeDuringSubmit(t *testing.T) { atomic.AddInt32(&doneCount, 1) }) + time.Sleep(10 * time.Millisecond) + assertEqual(t, 1, pool.IdleWorkers()) // Stop an idle worker right before submitting another task diff --git a/worker.go b/worker.go index 1677c27..c312bde 100644 --- a/worker.go +++ b/worker.go @@ -6,7 +6,7 @@ import ( ) // worker represents a worker goroutine -func worker(context context.Context, waitGroup *sync.WaitGroup, firstTask func(), tasks <-chan func(), taskExecutor func(func(), bool)) { +func worker(context context.Context, waitGroup *sync.WaitGroup, firstTask func(), tasks <-chan func(), taskExecutor func(func(), bool), taskWaitGroup *sync.WaitGroup) { // If provided, execute the first task immediately, before listening to the tasks channel if firstTask != nil { @@ -20,16 +20,33 @@ func worker(context context.Context, waitGroup *sync.WaitGroup, firstTask func() for { select { case <-context.Done(): - // Pool context was cancelled, exit + // Pool context was cancelled, empty tasks channel and exit + drainTasks(tasks, taskWaitGroup) return case task, ok := <-tasks: - if task == nil || !ok { - // We have received a signal to exit - return - } + // Prioritize context.Done statement (https://stackoverflow.com/questions/46200343/force-priority-of-go-select-statement) + select { + case <-context.Done(): + if task != nil && ok { + // We have received a task, ignore it + taskWaitGroup.Done() + } + default: + if task == nil || !ok { + // We have received a signal to exit + return + } - // We have received a task, execute it - taskExecutor(task, false) + // We have received a task, execute it + taskExecutor(task, false) + } } } } + +// drainPendingTasks discards queued tasks and decrements the corresponding wait group +func drainTasks(tasks <-chan func(), tasksWaitGroup *sync.WaitGroup) { + for _ = range tasks { + tasksWaitGroup.Done() + } +}