add a library to concurrently executes tasks (#197)

Author: DeimonDB

ctask/doer.go

@@ -0,0 +1,68 @@
+package ctask
+
+import (
+	"context"
+	"runtime"
+
+	"golang.org/x/sync/errgroup"
+)
+
+type DoOpt func(cfg *DoConfig)
+type DoConfig struct {
+	WorkerNum int
+}
+
+// Do execute tasks using the given executor function,
+// and return the results in the same order as the given tasks respectively.
+// It stops executing remaining tasks after any first error is encountered.
+//
+// The max number of goroutines can optionally be specified using the option WithWorkerNum.
+// By default, it is set to runtime.NumCPU()
+func Do[Task any, Result any](
+	ctx context.Context,
+	tasks []Task,
+	executor func(t Task) (Result, error),
+	opts ...DoOpt,
+) ([]Result, error) {
+	cfg := getConfigWithOptions(opts...)
+
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(int(cfg.WorkerNum))
+	results := make([]Result, len(tasks))
+	for idx, task := range tasks {
+		idx, task := idx, task // retain current loop values to be used in goroutine
+		g.Go(func() error {
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			default:
+				res, err := executor(task)
+				if err != nil {
+					return err
+				}
+				results[idx] = res
+				return nil
+			}
+		})
+	}
+	if err := g.Wait(); err != nil {
+		return nil, err
+	}
+	return results, nil
+}
+
+func getConfigWithOptions(opts ...DoOpt) DoConfig {
+	cfg := DoConfig{
+		WorkerNum: runtime.NumCPU(),
+	}
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+	return cfg
+}
+
+func WithWorkerNum(num int) DoOpt {
+	return func(cfg *DoConfig) {
+		cfg.WorkerNum = num
+	}
+}

ctask/doer_test.go

@@ -0,0 +1,90 @@
+package ctask
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestDo(t *testing.T) {
+	type T = int // task type
+	type R = int // result type
+
+	type args struct {
+		ctx      context.Context
+		tasks    []T
+		executor func(t T) (R, error)
+		opts     []DoOpt
+	}
+	tests := []struct {
+		name       string
+		args       args
+		want       []R
+		requireErr require.ErrorAssertionFunc
+	}{
+		{
+			name: "happy path",
+			args: args{
+				ctx:      context.Background(),
+				tasks:    []T{0, 1, 2, 3, 4, 5, 6},
+				executor: fibonacci,
+				opts:     nil,
+			},
+			want:       []R{1, 1, 2, 3, 5, 8, 13},
+			requireErr: require.NoError,
+		},
+		{
+			name: "empty slice",
+			args: args{
+				ctx:      context.Background(),
+				tasks:    []T{0, 1, 2, 3, 4, 5, 6},
+				executor: fibonacci,
+				opts:     nil,
+			},
+			want:       []R{1, 1, 2, 3, 5, 8, 13},
+			requireErr: require.NoError,
+		},
+		{
+			name: "error path & ensure tasks after error aren't executed (1000th fibonacci is too slow to be computed)",
+			args: args{
+				ctx:      context.Background(),
+				tasks:    []T{0, 1, 2, 1, -1, 1000},
+				executor: fibonacci,
+				opts:     []DoOpt{WithWorkerNum(1)},
+			},
+			want: nil,
+			requireErr: func(t require.TestingT, err error, i ...interface{}) {
+				require.Equal(t, errors.New("negative"), err)
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := Do(tt.args.ctx, tt.args.tasks, tt.args.executor, tt.args.opts...)
+			tt.requireErr(t, err)
+			require.Equal(t, tt.want, got)
+		})
+	}
+}
+
+func fibonacci(n int) (int, error) {
+	if n < 0 {
+		return 0, errors.New("negative")
+	}
+	if n < 2 {
+		return 1, nil
+	}
+	r1, err := fibonacci(n - 1)
+	if err != nil {
+		return 0, err
+	}
+
+	r2, err := fibonacci(n - 2)
+	if err != nil {
+		return 0, err
+	}
+
+	return r1 + r2, nil
+}

go.mod

@@ -18,6 +18,7 @@ require (
 	github.com/streadway/amqp v1.0.0
 	github.com/stretchr/testify v1.7.0
 	golang.org/x/exp v0.0.0-20220827204233-334a2380cb91
+	golang.org/x/sync v0.1.0
 	gorm.io/driver/postgres v1.2.3
 	gorm.io/gorm v1.22.4
 	gotest.tools v2.2.0+incompatible

go.sum

@@ -619,6 +619,8 @@ golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
+golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180823144017-11551d06cbcc/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=