Reintroduce pkg/util/concurrency, in place of dskit/concurrency (#4614)

aknuds1 · web-flow · commit 66df8876f549 · 2022-01-13T15:16:24.000+01:00
Signed-off-by: Arve Knudsen &lt;arve.knudsen@gmail.com&gt;
diff --git a/pkg/alertmanager/alertstore/bucketclient/bucket_client.go b/pkg/alertmanager/alertstore/bucketclient/bucket_client.go
@@ -9,13 +9,13 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/gogo/protobuf/proto"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/runutil"
 	"github.com/pkg/errors"
 	"github.com/thanos-io/thanos/pkg/objstore"
 
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertspb"
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 )
 
 const (
diff --git a/pkg/alertmanager/alertstore/objectclient/store.go b/pkg/alertmanager/alertstore/objectclient/store.go
@@ -9,12 +9,12 @@ import (
 	"sync"
 
 	"github.com/go-kit/log"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/runutil"
 	"github.com/pkg/errors"
 
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertspb"
 	"github.com/cortexproject/cortex/pkg/chunk"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 )
 
 // Object Alert Storage Schema
diff --git a/pkg/alertmanager/api.go b/pkg/alertmanager/api.go
@@ -12,7 +12,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/pkg/errors"
 	"github.com/prometheus/alertmanager/config"
 	"github.com/prometheus/alertmanager/template"
@@ -22,6 +21,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/alertmanager/alertspb"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
diff --git a/pkg/alertmanager/multitenant.go b/pkg/alertmanager/multitenant.go
@@ -15,7 +15,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/kv"
 	"github.com/grafana/dskit/services"
@@ -38,6 +37,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring/client"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
diff --git a/pkg/alertmanager/multitenant_test.go b/pkg/alertmanager/multitenant_test.go
@@ -20,7 +20,6 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/kv/consul"
 	"github.com/grafana/dskit/services"
@@ -47,6 +46,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring"
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	"github.com/cortexproject/cortex/pkg/util/test"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
diff --git a/pkg/compactor/blocks_cleaner.go b/pkg/compactor/blocks_cleaner.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/services"
 	"github.com/oklog/ulid"
 	"github.com/pkg/errors"
@@ -21,6 +20,7 @@ import (
 	cortex_tsdb "github.com/cortexproject/cortex/pkg/storage/tsdb"
 	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 )
 
diff --git a/pkg/compactor/compactor_test.go b/pkg/compactor/compactor_test.go
@@ -18,7 +18,6 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/kv/consul"
 	"github.com/grafana/dskit/services"
@@ -39,6 +38,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ring"
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	cortex_tsdb "github.com/cortexproject/cortex/pkg/storage/tsdb"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	cortex_testutil "github.com/cortexproject/cortex/pkg/util/test"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
diff --git a/pkg/frontend/frontend_test.go b/pkg/frontend/frontend_test.go
@@ -14,7 +14,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/gorilla/mux"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/services"
 	otgrpc "github.com/opentracing-contrib/go-grpc"
@@ -30,6 +29,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/frontend/transport"
 	"github.com/cortexproject/cortex/pkg/frontend/v1/frontendv1pb"
 	querier_worker "github.com/cortexproject/cortex/pkg/querier/worker"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 )
 
 const (
diff --git a/pkg/ingester/ingester_v2.go b/pkg/ingester/ingester_v2.go
@@ -13,7 +13,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/services"
 	"github.com/oklog/ulid"
 	"github.com/pkg/errors"
@@ -40,6 +39,7 @@ import (
 	cortex_tsdb "github.com/cortexproject/cortex/pkg/storage/tsdb"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	"github.com/cortexproject/cortex/pkg/util/extract"
 	logutil "github.com/cortexproject/cortex/pkg/util/log"
 	util_math "github.com/cortexproject/cortex/pkg/util/math"
diff --git a/pkg/querier/tenantfederation/merge_queryable.go b/pkg/querier/tenantfederation/merge_queryable.go
@@ -6,7 +6,6 @@ import (
 	"sort"
 	"strings"
 
-	"github.com/grafana/dskit/concurrency"
 	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
@@ -15,6 +14,7 @@ import (
 	"github.com/weaveworks/common/user"
 
 	"github.com/cortexproject/cortex/pkg/tenant"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
 )
 
diff --git a/pkg/ring/basic_lifecycler_delegates_test.go b/pkg/ring/basic_lifecycler_delegates_test.go
@@ -7,12 +7,13 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/services"
 	"github.com/grafana/dskit/test"
 	"github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 )
 
 func TestLeaveOnStoppingDelegate(t *testing.T) {
diff --git a/pkg/ruler/ruler.go b/pkg/ruler/ruler.go
@@ -14,7 +14,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/grafana/dskit/flagext"
 	"github.com/grafana/dskit/kv"
 	"github.com/grafana/dskit/services"
@@ -34,6 +33,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/ruler/rulestore"
 	"github.com/cortexproject/cortex/pkg/tenant"
 	"github.com/cortexproject/cortex/pkg/util"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 	"github.com/cortexproject/cortex/pkg/util/grpcclient"
 	util_log "github.com/cortexproject/cortex/pkg/util/log"
 	"github.com/cortexproject/cortex/pkg/util/validation"
diff --git a/pkg/storegateway/bucket_index_metadata_fetcher_test.go b/pkg/storegateway/bucket_index_metadata_fetcher_test.go
@@ -9,7 +9,6 @@ import (
 	"time"
 
 	"github.com/go-kit/log"
-	"github.com/grafana/dskit/concurrency"
 	"github.com/oklog/ulid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/testutil"
@@ -22,6 +21,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
 	cortex_testutil "github.com/cortexproject/cortex/pkg/storage/tsdb/testutil"
+	"github.com/cortexproject/cortex/pkg/util/concurrency"
 )
 
 func TestBucketIndexMetadataFetcher_Fetch(t *testing.T) {
diff --git a/pkg/util/concurrency/buffer.go b/pkg/util/concurrency/buffer.go
diff --git a/pkg/util/concurrency/runner.go b/pkg/util/concurrency/runner.go
@@ -4,10 +4,10 @@ import (
 	"context"
 	"sync"
 
+	"github.com/grafana/dskit/multierror"
 	"golang.org/x/sync/errgroup"
 
-	"github.com/grafana/dskit/internal/math"
-	"github.com/grafana/dskit/multierror"
+	"github.com/cortexproject/cortex/pkg/util/math"
 )
 
 // ForEachUser runs the provided userFunc for each userIDs up to concurrency concurrent workers.
diff --git a/pkg/util/concurrency/runner_test.go b/pkg/util/concurrency/runner_test.go
@@ -0,0 +1,168 @@
+package concurrency
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
+)
+
+func TestForEachUser(t *testing.T) {
+	var (
+		ctx = context.Background()
+
+		// Keep track of processed users.
+		processedMx sync.Mutex
+		processed   []string
+	)
+
+	input := []string{"a", "b", "c"}
+
+	err := ForEachUser(ctx, input, 2, func(ctx context.Context, user string) error {
+		processedMx.Lock()
+		defer processedMx.Unlock()
+		processed = append(processed, user)
+		return nil
+	})
+
+	require.NoError(t, err)
+	assert.ElementsMatch(t, input, processed)
+}
+
+func TestForEachUser_ShouldContinueOnErrorButReturnIt(t *testing.T) {
+	var (
+		ctx = context.Background()
+
+		// Keep the processed users count.
+		processed atomic.Int32
+	)
+
+	input := []string{"a", "b", "c"}
+
+	err := ForEachUser(ctx, input, 2, func(ctx context.Context, user string) error {
+		if processed.CAS(0, 1) {
+			return errors.New("the first request is failing")
+		}
+
+		// Wait 1s and increase the number of processed jobs, unless the context get canceled earlier.
+		select {
+		case <-time.After(time.Second):
+			processed.Add(1)
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+
+		return nil
+	})
+
+	require.EqualError(t, err, "the first request is failing")
+
+	// Since we expect it continues on error, the number of processed users should be equal to the input length.
+	assert.Equal(t, int32(len(input)), processed.Load())
+}
+
+func TestForEachUser_ShouldReturnImmediatelyOnNoUsersProvided(t *testing.T) {
+	require.NoError(t, ForEachUser(context.Background(), nil, 2, func(ctx context.Context, user string) error {
+		return nil
+	}))
+}
+
+func TestForEach(t *testing.T) {
+	var (
+		ctx = context.Background()
+
+		// Keep track of processed jobs.
+		processedMx sync.Mutex
+		processed   []string
+	)
+
+	jobs := []string{"a", "b", "c"}
+
+	err := ForEach(ctx, CreateJobsFromStrings(jobs), 2, func(ctx context.Context, job interface{}) error {
+		processedMx.Lock()
+		defer processedMx.Unlock()
+		processed = append(processed, job.(string))
+		return nil
+	})
+
+	require.NoError(t, err)
+	assert.ElementsMatch(t, jobs, processed)
+}
+
+func TestForEach_ShouldBreakOnFirstError_ContextCancellationHandled(t *testing.T) {
+	var (
+		ctx = context.Background()
+
+		// Keep the processed jobs count.
+		processed atomic.Int32
+	)
+
+	err := ForEach(ctx, []interface{}{"a", "b", "c"}, 2, func(ctx context.Context, job interface{}) error {
+		if processed.CAS(0, 1) {
+			return errors.New("the first request is failing")
+		}
+
+		// Wait 1s and increase the number of processed jobs, unless the context get canceled earlier.
+		select {
+		case <-time.After(time.Second):
+			processed.Add(1)
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+
+		return nil
+	})
+
+	require.EqualError(t, err, "the first request is failing")
+
+	// Since we expect the first error interrupts the workers, we should only see
+	// 1 job processed (the one which immediately returned error).
+	assert.Equal(t, int32(1), processed.Load())
+}
+
+func TestForEach_ShouldBreakOnFirstError_ContextCancellationUnhandled(t *testing.T) {
+	var (
+		ctx = context.Background()
+
+		// Keep the processed jobs count.
+		processed atomic.Int32
+	)
+
+	// waitGroup to await the start of the first two jobs
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	err := ForEach(ctx, []interface{}{"a", "b", "c"}, 2, func(ctx context.Context, job interface{}) error {
+		wg.Done()
+
+		if processed.CAS(0, 1) {
+			// wait till two jobs have been started
+			wg.Wait()
+			return errors.New("the first request is failing")
+		}
+
+		// Wait till context is cancelled to add processed jobs.
+		<-ctx.Done()
+		processed.Add(1)
+
+		return nil
+	})
+
+	require.EqualError(t, err, "the first request is failing")
+
+	// Since we expect the first error interrupts the workers, we should only
+	// see 2 job processed (the one which immediately returned error and the
+	// job with "b").
+	assert.Equal(t, int32(2), processed.Load())
+}
+
+func TestForEach_ShouldReturnImmediatelyOnNoJobsProvided(t *testing.T) {
+	require.NoError(t, ForEach(context.Background(), nil, 2, func(ctx context.Context, job interface{}) error {
+		return nil
+	}))
+}
diff --git a/vendor/github.com/grafana/dskit/internal/math/math.go b/vendor/github.com/grafana/dskit/internal/math/math.go
diff --git a/vendor/modules.txt b/vendor/modules.txt
