diff --git a/pond.go b/pond.go index e5e6192..f176537 100644 --- a/pond.go +++ b/pond.go @@ -318,8 +318,14 @@ func (p *WorkerPool) SubmitBefore(task func(), deadline time.Duration) { // 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. -func (p *WorkerPool) Stop() { - go p.stop(false) +// 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 diff --git a/pond_blackbox_test.go b/pond_blackbox_test.go index ca10b11..40a2aec 100644 --- a/pond_blackbox_test.go +++ b/pond_blackbox_test.go @@ -120,7 +120,7 @@ func TestSubmitAndStopWithoutWaiting(t *testing.T) { <-started // Stop without waiting for the rest of the tasks to start - pool.Stop() + ctx := pool.Stop() // Let the first task complete now completed <- true @@ -129,7 +129,7 @@ func TestSubmitAndStopWithoutWaiting(t *testing.T) { assertEqual(t, int32(1), atomic.LoadInt32(&doneCount)) // Make sure the exit lines in the worker pool are executed and covered - time.Sleep(6 * time.Millisecond) + <-ctx.Done() } func TestSubmitWithNilTask(t *testing.T) {