Distributor: Prevent data race on worker pool

pkg/distributor/distributor.go

@@ -522,13 +522,23 @@ func New(cfg Config, clientConfig ingester_client.Config, limits *validation.Ove
 	subservices = append(subservices, d.ingesterPool, d.activeUsers)
 
 	if cfg.ReusableIngesterPushWorkers > 0 {
+		var wpMu sync.Mutex // Prevents data races on the pool's jobs channel.
 		wp := concurrency.NewReusableGoroutinesPool(cfg.ReusableIngesterPushWorkers)
-		d.doBatchPushWorkers = wp.Go
+		d.doBatchPushWorkers = func(f func()) {
+			wpMu.Lock()
+			defer wpMu.Unlock()
+			wp.Go(f)
+		}
 		// Closing the pool doesn't stop the workload it's running, we're doing this just to avoid leaking goroutines in tests.
 		subservices = append(subservices, services.NewBasicService(
 			nil,
 			func(ctx context.Context) error { <-ctx.Done(); return nil },
-			func(_ error) error { wp.Close(); return nil },
+			func(_ error) error {
+				wpMu.Lock()
+				defer wpMu.Unlock()
+				wp.Close()
+				return nil
+			},
 		))
 	}
 

pkg/distributor/distributor_test.go

@@ -315,6 +315,48 @@ func TestDistributor_PushWithDoBatchWorkers(t *testing.T) {
 	require.GreaterOrEqual(t, counter.Load(), int64(3))
 }
 
+func TestDistributor_PushWithDoBatchWorkers_DataRace(t *testing.T) {
+	limits := prepareDefaultLimits()
+	limits.IngestionRate = 20
+	limits.IngestionBurstSize = 20
+
+	ds, _, _, _ := prepare(t, prepConfig{
+		numIngesters:    3,
+		happyIngesters:  3,
+		numDistributors: 1,
+		limits:          limits,
+		configure: func(cfg *Config) {
+			// 2 workers, so 1 push would need to spawn a new goroutine.
+			cfg.ReusableIngesterPushWorkers = 2
+		},
+	})
+	require.Len(t, ds, 1)
+	distributor := ds[0]
+
+	require.NotNil(t, distributor.doBatchPushWorkers)
+
+	reqCt := 10
+	errs := make(chan error, reqCt)
+	for i := 0; i < reqCt; i++ {
+		go func() {
+			request := makeWriteRequest(123456789000, 3, 5, false, false, "foo")
+			ctx := user.InjectOrgID(context.Background(), "user")
+			_, err := distributor.Push(ctx, request)
+			errs <- err
+		}()
+	}
+	time.Sleep(10 * time.Millisecond)
+	distributor.StopAsync()
+	errCt := 0
+	for i := 0; i < reqCt; i++ {
+		if <-errs != nil {
+			errCt++
+		}
+	}
+	require.NotZero(t, errCt, "expected at least one error")
+	require.Less(t, errCt, reqCt, "expected not all requests to fail")
+}
+
 func TestDistributor_ContextCanceledRequest(t *testing.T) {
 	now := time.Now()
 	mtime.NowForce(now)