diff --git a/group_done_cancel_test.go b/group_done_cancel_test.go new file mode 100644 index 0000000..fadf6c3 --- /dev/null +++ b/group_done_cancel_test.go @@ -0,0 +1,32 @@ +package pond + +import ( + "context" + "testing" + "time" +) + +func TestTaskGroupDoneClosesOnContextCancellation(t *testing.T) { + pool := NewPool(1) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + release := make(chan struct{}) + t.Cleanup(func() { + close(release) + pool.StopAndWait() + }) + group := pool.NewGroupContext(ctx) + started := make(chan struct{}) + group.Submit(func() { + close(started) + <-release + }) + <-started + done := group.Done() + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Done did not close when the group context was cancelled") + } +} diff --git a/internal/future/composite.go b/internal/future/composite.go index d88230e..c7d0e60 100644 --- a/internal/future/composite.go +++ b/internal/future/composite.go @@ -67,6 +67,16 @@ func (f *CompositeFuture[V]) Done(count int) <-chan struct{} { ch: ch, }) + go func() { + select { + case <-f.ctx.Done(): + f.mutex.Lock() + defer f.mutex.Unlock() + f.notifyListeners() + case <-ch: + } + }() + return ch }