Skip to content

Commit

Permalink
Return context in Stop() to notify user when everything has been stopped
Browse files Browse the repository at this point in the history
  • Loading branch information
Corentin Clabaut committed Jul 8, 2024
1 parent 3f439a7 commit 5e03cf1
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 4 deletions.
9 changes: 7 additions & 2 deletions pond.go
Original file line number Diff line number Diff line change
Expand Up @@ -318,8 +318,13 @@ 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)
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
Expand Down
4 changes: 2 additions & 2 deletions pond_blackbox_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {
Expand Down

0 comments on commit 5e03cf1

Please sign in to comment.