Merge pull request #67 from CorentinClabaut/stopCtx

Return context in Stop() to notify user when everything has been stopped

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

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) {