From 327a81e06f288b88580f3232ae019dcf62363074 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Fri, 7 Jun 2024 19:58:24 +0100
Subject: [PATCH] Bloombuild tests: protect fakeBuilder against race conditions

Using atomic variables.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>
---
 pkg/bloombuild/planner/planner_test.go | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/pkg/bloombuild/planner/planner_test.go b/pkg/bloombuild/planner/planner_test.go
index 2aa2626ba29ea..64c6ef086dac8 100644
--- a/pkg/bloombuild/planner/planner_test.go
+++ b/pkg/bloombuild/planner/planner_test.go
@@ -16,6 +16,7 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/atomic"
 	"google.golang.org/grpc"
 
 	"github.com/grafana/loki/v3/pkg/bloombuild/protos"
@@ -809,15 +810,15 @@ func Test_processTenantTaskResults(t *testing.T) {
 }
 
 type fakeBuilder struct {
-	mx          sync.Mutex
+	mx          sync.Mutex // Protects tasks and currTaskIdx.
 	id          string
 	tasks       []*protos.Task
 	currTaskIdx int
 	grpc.ServerStream
 
-	returnError    bool
-	returnErrorMsg bool
-	wait           bool
+	returnError    atomic.Bool
+	returnErrorMsg atomic.Bool
+	wait           atomic.Bool
 	ctx            context.Context
 	ctxCancel      context.CancelFunc
 }
@@ -840,15 +841,15 @@ func (f *fakeBuilder) ReceivedTasks() []*protos.Task {
 }
 
 func (f *fakeBuilder) SetReturnError(b bool) {
-	f.returnError = b
+	f.returnError.Store(b)
 }
 
 func (f *fakeBuilder) SetReturnErrorMsg(b bool) {
-	f.returnErrorMsg = b
+	f.returnErrorMsg.Store(b)
 }
 
 func (f *fakeBuilder) SetWait(b bool) {
-	f.wait = b
+	f.wait.Store(b)
 }
 
 func (f *fakeBuilder) CancelContext(b bool) {
@@ -891,12 +892,12 @@ func (f *fakeBuilder) Recv() (*protos.BuilderToPlanner, error) {
 		}, nil
 	}
 
-	if f.returnError {
+	if f.returnError.Load() {
 		return nil, fmt.Errorf("fake error from %s", f.id)
 	}
 
 	// Wait until `wait` is false
-	for f.wait {
+	for f.wait.Load() {
 		time.Sleep(time.Second)
 	}
 
@@ -906,10 +907,12 @@ func (f *fakeBuilder) Recv() (*protos.BuilderToPlanner, error) {
 	}
 
 	var errMsg string
-	if f.returnErrorMsg {
+	if f.returnErrorMsg.Load() {
 		errMsg = fmt.Sprintf("fake error from %s", f.id)
 	}
 
+	f.mx.Lock()
+	defer f.mx.Unlock()
 	return &protos.BuilderToPlanner{
 		BuilderID: f.id,
 		Result: protos.ProtoTaskResult{