1.fix data race in ClientMock.MockGet() (#5248)

2.fix compactor unit test anomalous error

Signed-off-by: yiyang5055 <[email protected]>

Author: yiyang5055

pkg/compactor/compactor_test.go

@@ -1215,8 +1215,8 @@ func TestCompactor_ShouldCompactOnlyShardsOwnedByTheInstanceOnShardingEnabledWit
 			bucketClient.MockGet(userID+"/"+blockID+"/meta.json", mockBlockMetaJSONWithTime(blockID, userID, blockTimes["startTime"], blockTimes["endTime"]), nil)
 			bucketClient.MockGet(userID+"/"+blockID+"/deletion-mark.json", "", nil)
 			bucketClient.MockGet(userID+"/"+blockID+"/no-compact-mark.json", "", nil)
-			bucketClient.MockGetTimes(userID+"/"+blockID+"/visit-mark.json", "", nil, 1)
 			bucketClient.MockGet(userID+"/"+blockID+"/visit-mark.json", string(visitMarkerFileContent), nil)
+			bucketClient.MockGetRequireUpload(userID+"/"+blockID+"/visit-mark.json", string(visitMarkerFileContent), nil)
 			bucketClient.MockUpload(userID+"/"+blockID+"/visit-mark.json", nil)
 			blockDirectory = append(blockDirectory, userID+"/"+blockID)
 
@@ -1243,6 +1243,7 @@ func TestCompactor_ShouldCompactOnlyShardsOwnedByTheInstanceOnShardingEnabledWit
 	for i := 1; i <= 4; i++ {
 		cfg := prepareConfig()
 		cfg.ShardingEnabled = true
+		cfg.CompactionInterval = 15 * time.Second
 		cfg.ShardingStrategy = util.ShardingStrategyShuffle
 		cfg.ShardingRing.InstanceID = fmt.Sprintf("compactor-%d", i)
 		cfg.ShardingRing.InstanceAddr = fmt.Sprintf("127.0.0.%d", i)
@@ -1280,7 +1281,7 @@ func TestCompactor_ShouldCompactOnlyShardsOwnedByTheInstanceOnShardingEnabledWit
 
 	// Wait until a run has been completed on each compactor
 	for _, c := range compactors {
-		cortex_testutil.Poll(t, 60*time.Second, 1.0, func() interface{} {
+		cortex_testutil.Poll(t, 60*time.Second, 2.0, func() interface{} {
 			return prom_testutil.ToFloat64(c.compactionRunsCompleted)
 		})
 	}

pkg/storage/bucket/client_mock.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"sync"
 	"time"
 
 	"github.com/stretchr/testify/mock"
@@ -16,10 +17,14 @@ var errObjectDoesNotExist = errors.New("object does not exist")
 // ClientMock mocks objstore.Bucket
 type ClientMock struct {
 	mock.Mock
+	uploaded sync.Map
 }
 
 // Upload mocks objstore.Bucket.Upload()
 func (m *ClientMock) Upload(ctx context.Context, name string, r io.Reader) error {
+	if _, ok := m.uploaded.Load(name); ok {
+		m.uploaded.Store(name, true)
+	}
 	args := m.Called(ctx, name, r)
 	return args.Error(0)
 }
@@ -30,6 +35,8 @@ func (m *ClientMock) MockUpload(name string, err error) {
 
 // Delete mocks objstore.Bucket.Delete()
 func (m *ClientMock) Delete(ctx context.Context, name string) error {
+	m.uploaded.Delete(name)
+
 	args := m.Called(ctx, name)
 	return args.Error(0)
 }
@@ -70,7 +77,20 @@ func (m *ClientMock) MockIterWithCallback(prefix string, objects []string, err e
 
 // Get mocks objstore.Bucket.Get()
 func (m *ClientMock) Get(ctx context.Context, name string) (io.ReadCloser, error) {
+	if val, ok := m.uploaded.Load(name); ok {
+		uploaded := val.(bool)
+		if !uploaded {
+			return nil, errObjectDoesNotExist
+		}
+	}
+
 	args := m.Called(ctx, name)
+
+	// Allow to mock the Get() with a function which is called each time.
+	if fn, ok := args.Get(0).(func(ctx context.Context, name string) (io.ReadCloser, error)); ok {
+		return fn(ctx, name)
+	}
+
 	val, err := args.Get(0), args.Error(1)
 	if val == nil {
 		return nil, err
@@ -90,9 +110,8 @@ func (m *ClientMock) MockGet(name, content string, err error) {
 		// Since we return an ReadCloser and it can be consumed only once,
 		// each time the mocked Get() is called we do create a new one, so
 		// that getting the same mocked object twice works as expected.
-		mockedGet := m.On("Get", mock.Anything, name)
-		mockedGet.Run(func(args mock.Arguments) {
-			mockedGet.Return(io.NopCloser(bytes.NewReader([]byte(content))), err)
+		m.On("Get", mock.Anything, name).Return(func(_ context.Context, _ string) (io.ReadCloser, error) {
+			return io.NopCloser(bytes.NewReader([]byte(content))), err
 		})
 	} else {
 		m.On("Exists", mock.Anything, name).Return(false, err)
@@ -101,6 +120,13 @@ func (m *ClientMock) MockGet(name, content string, err error) {
 	}
 }
 
+// MockGetRequireUpload is a convenient method to mock Get() return resulst after upload,
+// otherwise return errObjectDoesNotExist
+func (m *ClientMock) MockGetRequireUpload(name, content string, err error) {
+	m.uploaded.Store(name, false)
+	m.MockGet(name, content, err)
+}
+
 // MockGetTimes is a convenient method to mock Get() and Exists() to run x time
 func (m *ClientMock) MockGetTimes(name, content string, err error, times int) {
 	if content != "" {

pkg/storage/bucket/client_test.go

@@ -2,6 +2,8 @@ package bucket
 
 import (
 	"context"
+	"io"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -92,3 +94,30 @@ func TestNewClient(t *testing.T) {
 		})
 	}
 }
+
+func TestClientMock_MockGet(t *testing.T) {
+	expected := "body"
+
+	m := ClientMock{}
+	m.MockGet("test", expected, nil)
+
+	// Run many goroutines all requesting the same mocked object and
+	// ensure there's no race.
+	wg := sync.WaitGroup{}
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			reader, err := m.Get(context.Background(), "test")
+			require.NoError(t, err)
+			actual, err := io.ReadAll(reader)
+			require.NoError(t, err)
+			require.Equal(t, []byte(expected), actual)
+
+			require.NoError(t, reader.Close())
+		}()
+	}
+
+	wg.Wait()
+}