create WaitGroupWithStopOnError option to cancel other tasks and NewWaitGroupWithContext function to create WaitGroup with custom context

Author: mrsoftware

README.md

@@ -184,7 +184,52 @@ if err := wg.Wait(); err != nil {
     // oh, something bad happened in one of routines above.
 } 
 ```
+## Stop all tasks on first error
+```go
+import (
+    "github.com/mrsoftware/errors"
+)
+
+// in this example we are using ants goroutine pool.
+wg := errors.NewWaitGroup(errors.WaitGroupWithStopOnError()) 
+
+ctx := wg.Context()
+
+wg.Do(func() error {
+    return callingHttpClient(ctx)
+})
+
+wg.Do(func() error {
+    return callingHttpClient(ctx)
+})
+ 
+// if one of above task failed, context will cancel and other task will stop (the task must ba aware of context cancellation like http pkg do)
 
+if err := wg.Wait(); err != nil {
+    // oh, something bad happened in one of routines above.
+} 
+```
+**or you can use NewWaitGroupWithContext method:**
+```go
+import (
+    "github.com/mrsoftware/errors"
+)
+
+// in this example we are using ants goroutine pool.
+ctx, wg := errors.NewWaitGroupWithContext(context.Background(), errors.WaitGroupWithStopOnError()) 
+
+wg.Do(func() error {
+    return callingHttpClient(ctx)
+})
+
+wg.Do(func() error {
+    return callingHttpClient(ctx)
+})
+ 
+if err := wg.Wait(); err != nil {
+    // oh, something bad happened in one of routines above.
+} 
+```
 
 for mode details, check the [documentation](https://godoc.org/github.com/mrsoftware/errors)
 

waitGroup.go

@@ -1,19 +1,30 @@
 package errors
 
 import (
+	"context"
 	"sync"
 )
 
 // WaitGroup is sync.WaitGroup with error support.
 type WaitGroup struct {
-	noCopy  noCopy
-	options *WaitGroupOptions
-	errors  MultiError
-	gch     chan struct{}
+	noCopy     noCopy
+	options    *WaitGroupOptions
+	errors     MultiError
+	gch        chan struct{}
+	ctx        context.Context
+	cancel     context.CancelCauseFunc
+	cancelOnce sync.Once
 }
 
 // NewWaitGroup create new WaitGroup.
 func NewWaitGroup(options ...WaitGroupOption) *WaitGroup {
+	_, wg := NewWaitGroupWithContext(context.Background(), options...)
+
+	return wg
+}
+
+// NewWaitGroupWithContext create new WaitGroup with custom context.
+func NewWaitGroupWithContext(ctx context.Context, options ...WaitGroupOption) (context.Context, *WaitGroup) {
 	ops := &WaitGroupOptions{
 		Wg:         &sync.WaitGroup{},
 		TaskRunner: func(task func()) { go task() },
@@ -28,13 +39,27 @@ func NewWaitGroup(options ...WaitGroupOption) *WaitGroup {
 		gch = make(chan struct{}, ops.TaskLimit)
 	}
 
-	return &WaitGroup{options: ops, gch: gch}
+	ctx, cancel := context.WithCancelCause(context.Background())
+
+	return ctx, &WaitGroup{options: ops, gch: gch, ctx: ctx, cancel: cancel}
+}
+
+// Context of current waitGroup.
+func (g *WaitGroup) Context() context.Context {
+	return g.ctx
+}
+
+// Stop send cancel signal to all tasks.
+func (g *WaitGroup) Stop(err error) {
+	g.cancelOnce.Do(func() { g.cancel(err) })
 }
 
 // Wait is sync.WaitGroup.Wait.
-func (g *WaitGroup) Wait() error {
+func (g *WaitGroup) Wait() (err error) {
 	g.options.Wg.Wait()
 
+	defer func() { g.Stop(err) }()
+
 	if g.errors.Len() == 0 {
 		return nil
 	}
@@ -55,6 +80,10 @@ func (g *WaitGroup) Done(err error) {
 		return
 	}
 
+	if g.options.StopOnError {
+		g.Stop(err)
+	}
+
 	g.errors.Add(err)
 }
 
@@ -84,9 +113,10 @@ type noCopy struct{}
 
 // WaitGroupOptions for WaitGroup.
 type WaitGroupOptions struct {
-	Wg         *sync.WaitGroup
-	TaskLimit  int
-	TaskRunner WaitGroupTaskRunner
+	Wg          *sync.WaitGroup
+	TaskLimit   int
+	TaskRunner  WaitGroupTaskRunner
+	StopOnError bool
 }
 
 type WaitGroupOption func(group *WaitGroupOptions)
@@ -114,3 +144,10 @@ func WaitGroupWithTaskRunner(runner WaitGroupTaskRunner) WaitGroupOption {
 		g.TaskRunner = runner
 	}
 }
+
+// WaitGroupWithStopOnError used if you want to stop all tasks on first error.
+func WaitGroupWithStopOnError() WaitGroupOption {
+	return func(g *WaitGroupOptions) {
+		g.StopOnError = true
+	}
+}

waitGroup_test.go

@@ -1,6 +1,7 @@
 package errors
 
 import (
+	"context"
 	"errors"
 	"sync"
 	"sync/atomic"
@@ -131,6 +132,30 @@ func TestGroup(t *testing.T) {
 		assert.Nil(t, err)
 		assert.True(t, isUsedCustomRunner)
 	})
+
+	t.Run("set StopOnError options", func(t *testing.T) {
+		error1 := errors.New("error 1")
+
+		ctx, wg := NewWaitGroupWithContext(context.Background(), WaitGroupWithStopOnError())
+
+		wg.Do(func() error { return error1 })
+
+		// sample long-running and context aware task.
+		wg.Do(func() error {
+			for {
+				select {
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+			}
+		})
+
+		err := wg.Wait()
+
+		expected := NewMultiError(error1, context.Canceled)
+		assert.ElementsMatch(t, expected.errors, err.(*MultiError).errors)
+		assert.Equal(t, ctx, wg.Context())
+	})
 }
 
 // all below test cases are copied from sync/waitgroup_test.go and transformed to group.