Repository: apigee/registry
Pull request #1114 (2023-03-21 17:24:16)

add continueOnError param to WorkerPool
Follow up for #917


Patch:
diff --git a/cmd/registry/tasks/tasks.go b/cmd/registry/tasks/tasks.go
index a2bcd44..67a3640 100644
--- a/cmd/registry/tasks/tasks.go
+++ b/cmd/registry/tasks/tasks.go
@@ -33,18 +33,24 @@ type Task interface {
 // The return value is the taskQueue and a wait function.
 // Do not directly close the taskQueue, use the wait function.
 // Clients should add new tasks to this taskQueue and call the wait function
-// when done. If a worker fails with an error, the task context will be
-// canceled, instructing the queue and workers to terminate. Tasks should
-// check ctx as appropriate. The first error encountered will be returned
-// from the wait function and all errors will be logged at Warn level.
-func WorkerPool(ctx context.Context, n int) (chan<- Task, func() error) {
-	eg, ctx := errgroup.WithContext(ctx)
+// when done. If a worker fails with an error and continueOnError is false,
+// the task context will be canceled, instructing the queue and workers to
+// terminate. Tasks should check ctx as appropriate. The first error encountered
+// will be returned from the wait function and all errors are logged at Warn level.
+func WorkerPool(ctx context.Context, n int, continueOnError bool) (chan<- Task, func() error) {
+	var eg *errgroup.Group
+	if continueOnError {
+		eg = &errgroup.Group{}
+	} else {
+		eg, ctx = errgroup.WithContext(ctx)
+	}
 	eg.SetLimit(n) // pool size
 	taskQueue := make(chan Task, 1024)
 	var wg sync.WaitGroup
 	wg.Add(1)
 
 	go func() {
+		defer wg.Done()
 		for task := range taskQueue {
 			select {
 			case <-ctx.Done():
@@ -61,7 +67,6 @@ func WorkerPool(ctx context.Context, n int) (chan<- Task, func() error) {
 				eg.Go(f) // blocks at n runners
 			}
 		}
-		wg.Done()
 	}()
 
 	wait := func() error {
@@ -77,11 +82,11 @@ func WorkerPool(ctx context.Context, n int) (chan<- Task, func() error) {
 	return taskQueue, wait
 }
 
-// WorkerPoolIgnoreError does the same as WorkerPool, but returns a wait
-// function that captures and logs the error at Error level. This is convenient
-// for defer.
+// WorkerPoolIgnoreError instantiates a WorkerPool with continueOnError=true and
+// returns a wait function that captures and logs the error at Error level. This
+// is convenient for defer.
 func WorkerPoolIgnoreError(ctx context.Context, n int) (chan<- Task, func()) {
-	wp, w := WorkerPool(ctx, n)
+	wp, w := WorkerPool(ctx, n, false)
 	wait := func() {
 		if err := w(); err != nil {
 			log.FromContext(ctx).WithError(err).Errorf("unhandled WorkerPool error")


Test patch:
diff --git a/cmd/registry/tasks/tasks_test.go b/cmd/registry/tasks/tasks_test.go
index fd9498f..ec3602b 100644
--- a/cmd/registry/tasks/tasks_test.go
+++ b/cmd/registry/tasks/tasks_test.go
@@ -21,15 +21,15 @@ import (
 	"testing"
 )
 
-func TestWorkerPool(t *testing.T) {
+func TestWorkerPoolContinueOnError(t *testing.T) {
 	ctx := context.Background()
 	jobs := 100
 	counter := new(atomicInt32)
-	want := 1000
+	total := 1000
 
 	func() {
-		taskQueue, wait := WorkerPool(ctx, jobs)
-		for i := 0; i < want; i++ {
+		taskQueue, wait := WorkerPool(ctx, jobs, true)
+		for i := 0; i < total; i++ {
 			taskQueue <- &incrTask{counter}
 		}
 		if err := wait(); err != nil {
@@ -38,8 +38,34 @@ func TestWorkerPool(t *testing.T) {
 	}()
 
 	got := counter.Load()
-	if got != int32(want) {
-		t.Errorf("want %d got: %d", want, got)
+	if got != int32(total) {
+		t.Errorf("want %d got: %d", total, got)
+	}
+}
+
+func TestWorkerPoolStopOnError(t *testing.T) {
+	ctx := context.Background()
+	jobs := 100
+	counter := new(atomicInt32)
+	total := 1000
+	errorAt := 500
+
+	func() {
+		taskQueue, wait := WorkerPool(ctx, jobs, false)
+		for i := 0; i < total; i++ {
+			if i == errorAt {
+				taskQueue <- &failTask{}
+			}
+			taskQueue <- &incrTask{counter}
+		}
+		if err := wait(); err == nil {
+			t.Log("expected error")
+		}
+	}()
+
+	got := counter.Load()
+	if got < int32(errorAt) || got >= int32(total) {
+		t.Errorf("want between %d and %d, got: %d", errorAt, total, got)
 	}
 }
 
@@ -56,19 +82,6 @@ func (t *incrTask) Run(ctx context.Context) error {
 	return nil
 }
 
-func TestWorkerPoolWithError(t *testing.T) {
-	ctx := context.Background()
-	jobs := 1
-
-	taskQueue, wait := WorkerPool(ctx, jobs)
-	for i := 0; i < 2; i++ {
-		taskQueue <- &failTask{}
-	}
-	if err := wait(); err == nil {
-		t.Error("want error, got nil")
-	}
-}
-
 func TestWorkerPoolIgnoreError(t *testing.T) {
 	ctx := context.Background()
 	jobs := 1