From 7ca844ab99dd2e0efed29aa2e859ab9b25e2ebfb Mon Sep 17 00:00:00 2001 From: SimFG Date: Wed, 28 Sep 2022 19:28:54 +0800 Subject: [PATCH] Fix the data race in the indexcoord unittest (#19515) Signed-off-by: SimFG Signed-off-by: SimFG --- internal/indexcoord/index_coord_mock.go | 13 +++++++++++++ internal/indexcoord/index_coord_test.go | 12 ++++++++---- 2 files changed, 21 insertions(+), 4 deletions(-) diff --git a/internal/indexcoord/index_coord_mock.go b/internal/indexcoord/index_coord_mock.go index 5b11297f63..3450045334 100644 --- a/internal/indexcoord/index_coord_mock.go +++ b/internal/indexcoord/index_coord_mock.go @@ -19,6 +19,7 @@ package indexcoord import ( "context" "math/rand" + "sync" "time" clientv3 "go.etcd.io/etcd/client/v3" @@ -348,6 +349,8 @@ func NewRootCoordMock() *RootCoordMock { type DataCoordMock struct { types.DataCoord + FuncLock sync.RWMutex + CallInit func() error CallStart func() error CallGetComponentStates func(ctx context.Context) (*internalpb.ComponentStates, error) @@ -370,7 +373,17 @@ func (dcm *DataCoordMock) GetComponentStates(ctx context.Context) (*internalpb.C return dcm.CallGetComponentStates(ctx) } +func (dcm *DataCoordMock) SetFunc(f func()) { + dcm.FuncLock.Lock() + defer dcm.FuncLock.Unlock() + + f() +} + func (dcm *DataCoordMock) GetSegmentInfo(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) { + dcm.FuncLock.RLock() + defer dcm.FuncLock.RUnlock() + return dcm.CallGetSegmentInfo(ctx, req) } func (dcm *DataCoordMock) AcquireSegmentLock(ctx context.Context, req *datapb.AcquireSegmentLockRequest) (*commonpb.Status, error) { diff --git a/internal/indexcoord/index_coord_test.go b/internal/indexcoord/index_coord_test.go index bdb0e31709..fb24a6f917 100644 --- a/internal/indexcoord/index_coord_test.go +++ b/internal/indexcoord/index_coord_test.go @@ -211,13 +211,17 @@ func TestIndexCoord(t *testing.T) { assert.NoError(t, err) assert.Equal(t, resp.Status.ErrorCode, commonpb.ErrorCode_UnexpectedError) dcm.CallGetFlushedSegment = originFunc1 - dcm.CallGetSegmentInfo = func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) { - return nil, errors.New("mock error") - } + dcm.SetFunc(func() { + dcm.CallGetSegmentInfo = func(ctx context.Context, req *datapb.GetSegmentInfoRequest) (*datapb.GetSegmentInfoResponse, error) { + return nil, errors.New("mock error") + } + }) resp, err = ic.DescribeIndex(ctx, req) assert.NoError(t, err) assert.Equal(t, resp.Status.ErrorCode, commonpb.ErrorCode_UnexpectedError) - dcm.CallGetSegmentInfo = originFunc2 + dcm.SetFunc(func() { + dcm.CallGetSegmentInfo = originFunc2 + }) }) t.Run("FlushedSegmentWatcher", func(t *testing.T) {