Make client fail fast when they face an issue with API.

Author: clems4ever

internal/utils/workerpool.go

@@ -2,11 +2,15 @@ package utils
 
 import (
 	"context"
+	"errors"
 	"sync"
 
 	"golang.org/x/sync/semaphore"
 )
 
+// ErrTaskAborted error thrown when a taks is aborted
+var ErrTaskAborted = errors.New("Aborted")
+
 // Task is a representation of a task
 type Task struct {
 	fn   func() error
@@ -19,6 +23,7 @@ type WorkerPool struct {
 	taskQ   chan Task
 	wg      sync.WaitGroup
 	sem     *semaphore.Weighted
+	failed  bool
 }
 
 // NewWorkerPool instantiate a new worker pool
@@ -28,6 +33,7 @@ func NewWorkerPool(workers int) *WorkerPool {
 		workers: workers,
 		taskQ:   taskQ,
 		sem:     semaphore.NewWeighted(int64(workers)),
+		failed:  false,
 	}
 
 	wp.wg.Add(workers)
@@ -36,8 +42,18 @@ func NewWorkerPool(workers int) *WorkerPool {
 		go func() {
 			defer wp.wg.Done()
 			for task := range taskQ {
-				task.errC <- task.fn()
+				// If the worker pool has been aborted, we skip all tasks in the queue.
+				if wp.failed {
+					task.errC <- ErrTaskAborted
+					wp.sem.Release(1)
+					continue
+				}
+				err := task.fn()
 				wp.sem.Release(1)
+				if err != nil {
+					wp.failed = true
+				}
+				task.errC <- err
 			}
 		}()
 	}
@@ -47,10 +63,15 @@ func NewWorkerPool(workers int) *WorkerPool {
 
 // Exec enqueues one task into the worker pool
 func (wp *WorkerPool) Exec(f func() error) chan error {
+	errC := make(chan error, 1)
+	if wp.failed {
+		errC <- ErrTaskAborted
+		return errC
+	}
 	wp.sem.Acquire(context.Background(), 1)
 	t := Task{
 		fn:   f,
-		errC: make(chan error, 1),
+		errC: errC,
 	}
 	wp.taskQ <- t
 	return t.errC

internal/utils/workerpool_test.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -59,3 +60,36 @@ func TestWorkerPoolWithManyTasks(t *testing.T) {
 		<-f
 	}
 }
+
+func TestWorkerPoolAborted(t *testing.T) {
+	p := NewWorkerPool(5)
+	defer p.Close()
+
+	var errFailed = errors.New("failed")
+
+	futures := make([]chan error, 0)
+	for i := 0; i < 1000; i++ {
+		ix := i
+		fmt.Printf("task %d inserted\n", i)
+		f := p.Exec(func() error {
+			fmt.Printf("task %d started\n", ix)
+			time.Sleep(3 * time.Second)
+			if ix == 6 {
+				fmt.Printf("task failed!")
+				return errFailed
+			}
+			fmt.Printf("task %d ended\n", ix)
+			return nil
+		})
+		futures = append(futures, f)
+	}
+
+	errorsCount := 0
+	for _, f := range futures {
+		err := <-f
+		if err == errFailed {
+			errorsCount++
+		}
+	}
+	assert.Equal(t, 1, errorsCount)
+}