milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: data race in querynode task test (#31019) 

issue: https://github.com/milvus-io/milvus/issues/31022

Signed-off-by: sunby <sunbingyi1992@gmail.com>
pull/31000/head
Bingyi Sun 2024-03-05 16:26:59 +08:00 committed by GitHub
parent 3b66c17279
commit e3cce11dd9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 72 additions and 56 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

48
internal/querycoordv2/task/task_test.go
View File

@ -200,14 +200,16 @@ func (suite *TaskSuite) TestSubscribeChannelTask() {
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection). suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).
Return(&milvuspb.DescribeCollectionResponse{ RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestSubscribeChannelTask", Name: "TestSubscribeChannelTask",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
for channel, segment := range suite.growingSegments { for channel, segment := range suite.growingSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment). suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).
Return(&datapb.GetSegmentInfoResponse{ Return(&datapb.GetSegmentInfoResponse{
@ -389,14 +391,16 @@ func (suite *TaskSuite) TestLoadSegmentTask() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestLoadSegmentTask", Name: "TestLoadSegmentTask",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,
@ -487,14 +491,16 @@ func (suite *TaskSuite) TestLoadSegmentTaskNotIndex() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestLoadSegmentTaskNotIndex", Name: "TestLoadSegmentTaskNotIndex",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,
@ -585,14 +591,16 @@ func (suite *TaskSuite) TestLoadSegmentTaskFailed() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestLoadSegmentTaskNotIndex", Name: "TestLoadSegmentTaskNotIndex",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
for _, segment := range suite.loadSegments { for _, segment := range suite.loadSegments {
suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{ suite.broker.EXPECT().GetSegmentInfo(mock.Anything, segment).Return(&datapb.GetSegmentInfoResponse{
Infos: []*datapb.SegmentInfo{ Infos: []*datapb.SegmentInfo{
@ -785,14 +793,16 @@ func (suite *TaskSuite) TestMoveSegmentTask() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestMoveSegmentTask", Name: "TestMoveSegmentTask",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,
@ -957,14 +967,16 @@ func (suite *TaskSuite) TestTaskCanceled() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestSubscribeChannelTask", Name: "TestSubscribeChannelTask",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,
@ -1046,14 +1058,16 @@ func (suite *TaskSuite) TestSegmentTaskStale() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestSegmentTaskStale", Name: "TestSegmentTaskStale",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,
@ -1226,14 +1240,16 @@ func (suite *TaskSuite) TestLeaderTaskSet() {
} }
// Expect // Expect
suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).Return(&milvuspb.DescribeCollectionResponse{ suite.broker.EXPECT().DescribeCollection(mock.Anything, suite.collection).RunAndReturn(func(ctx context.Context, i int64) (*milvuspb.DescribeCollectionResponse, error) {
return &milvuspb.DescribeCollectionResponse{
Schema: &schemapb.CollectionSchema{ Schema: &schemapb.CollectionSchema{
Name: "TestLoadSegmentTask", Name: "TestLoadSegmentTask",
Fields: []*schemapb.FieldSchema{ Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector}, {FieldID: 100, Name: "vec", DataType: schemapb.DataType_FloatVector},
}, },
}, },
}, nil) }, nil
})
suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{ suite.broker.EXPECT().DescribeIndex(mock.Anything, suite.collection).Return([]*indexpb.IndexInfo{
{ {
CollectionID: suite.collection, CollectionID: suite.collection,