Group context

.github/workflows/main.yml

@@ -10,7 +10,7 @@ jobs:
     name: Test
     strategy:
       matrix:
-        go-version: [1.15.x, 1.16.x, 1.17.x]
+        go-version: [1.15.x, 1.16.x, 1.17.x, 1.18.x]
         os: [ubuntu-latest, macos-latest, windows-latest]
     runs-on: ${{ matrix.os }}
     steps:

go.mod

@@ -1,3 +1,3 @@
 module github.com/alitto/pond
 
-go 1.17
+go 1.18

group.go

@@ -0,0 +1,71 @@
+package pond
+
+import (
+	"context"
+	"fmt"
+	"sync"
+)
+
+type TaskGroup = Group[func()]
+
+type TaskGroupWithContext = Group[func() error]
+
+// Group represents a group of related tasks
+type Group[T func() | func() error] struct {
+	pool      *WorkerPool
+	waitGroup sync.WaitGroup
+	ctx       context.Context
+	cancel    context.CancelFunc
+	errOnce   sync.Once
+	err       error
+}
+
+// Submit adds a task to this group and sends it to the worker pool to be executed
+func (g *Group[T]) Submit(task T) {
+	g.waitGroup.Add(1)
+
+	g.pool.Submit(func() {
+		defer g.waitGroup.Done()
+
+		// If context has already been cancelled, skip task execution
+		if g.ctx != nil {
+			select {
+			case <-g.ctx.Done():
+				return
+			default:
+			}
+		}
+
+		switch any(task).(type) {
+		case func():
+			any(task).(func())()
+		case func() error:
+			// don't actually ignore errors
+			err := any(task).(func() error)()
+			if err != nil {
+				fmt.Printf("Error not nil %v\n", err)
+				g.errOnce.Do(func() {
+					g.err = err
+					if g.cancel != nil {
+						g.cancel()
+					}
+				})
+			} else {
+				fmt.Printf("Error is nil %v\n", err)
+			}
+		}
+	})
+}
+
+// Wait waits until all the tasks in this group have completed
+func (g *Group[T]) Wait() error {
+
+	// Wait for all tasks to complete
+	g.waitGroup.Wait()
+
+	// If context was provided, cancel it to signal all running tasks to stop
+	if g.cancel != nil {
+		g.cancel()
+	}
+	return g.err
+}

pond.go

@@ -487,6 +487,19 @@ func (p *WorkerPool) Group() *TaskGroup {
 	}
 }
 
+// GroupContext creates a new task group and an associated Context derived from ctx.
+//
+// The derived Context is canceled the first time a function submitted to the group
+// returns a non-nil error or the first time Wait returns, whichever occurs first.
+func (p *WorkerPool) GroupContext(ctx context.Context) (*TaskGroupWithContext, context.Context) {
+	ctx, cancel := context.WithCancel(ctx)
+	return &TaskGroupWithContext{
+		pool:   p,
+		ctx:    ctx,
+		cancel: cancel,
+	}, ctx
+}
+
 // worker launches a worker goroutine
 func worker(context context.Context, firstTask func(), tasks <-chan func(), idleWorkerCount *int32, exitHandler func(), taskExecutor func(func())) {
 
@@ -528,25 +541,3 @@ func worker(context context.Context, firstTask func(), tasks <-chan func(), idle
 		}
 	}
 }
-
-// TaskGroup represents a group of related tasks
-type TaskGroup struct {
-	pool      *WorkerPool
-	waitGroup sync.WaitGroup
-}
-
-// Submit adds a task to this group and sends it to the worker pool to be executed
-func (g *TaskGroup) Submit(task func()) {
-	g.waitGroup.Add(1)
-	g.pool.Submit(func() {
-		defer g.waitGroup.Done()
-		task()
-	})
-}
-
-// Wait waits until all the tasks in this group have completed
-func (g *TaskGroup) Wait() {
-
-	// Wait for all tasks to complete
-	g.waitGroup.Wait()
-}

pond_blackbox_test.go

@@ -2,6 +2,7 @@ package pond_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"sync"
 	"sync/atomic"
@@ -478,6 +479,41 @@ func TestGroupSubmit(t *testing.T) {
 	assertEqual(t, int32(taskCount), atomic.LoadInt32(&doneCount))
 }
 
+func TestGroupContext(t *testing.T) {
+
+	pool := pond.New(5, 100)
+	assertEqual(t, 0, pool.RunningWorkers())
+
+	expectedErr := errors.New("Something went wrong")
+
+	// Submit a group of tasks
+	var doneCount int32
+	group, ctx := pool.GroupContext(context.Background())
+	for i := 0; i < 20; i++ {
+		n := i
+		group.Submit(func() error {
+
+			if n == 2 {
+				time.Sleep(5 * time.Millisecond)
+				return expectedErr
+			}
+
+			select {
+			case <-time.After(time.Duration(n) * time.Millisecond):
+				atomic.AddInt32(&doneCount, 1)
+			case <-ctx.Done():
+			}
+
+			return nil
+		})
+	}
+
+	err := group.Wait()
+	assertEqual(t, expectedErr, err)
+
+	assertEqual(t, int32(5), atomic.LoadInt32(&doneCount))
+}
+
 func TestPoolWithCustomStrategy(t *testing.T) {
 
 	pool := pond.New(3, 3, pond.Strategy(pond.RatedResizer(2)))