fix: [Cherry-pick] sync action load segment with lack collection inde…

…x info list (#28956) relate: #28779 #28637 pr: #28788 Signed-off-by: aoiasd <[email protected]>
milvus-io · Dec 7, 2023 · 8502037 · 8502037
1 parent f670055
commit 8502037
Show file tree

Hide file tree

Showing 7 changed files with 370 additions and 321 deletions.
diff --git a/internal/proto/query_coord.proto b/internal/proto/query_coord.proto
@@ -586,6 +586,7 @@ message SyncDistributionRequest {
   LoadMetaInfo load_meta = 6;
   int64 replicaID = 7;
   int64 version = 8;
+  repeated index.IndexInfo index_info_list = 9;
 }
 
 message ResourceGroup {

diff --git a/internal/proto/querypb/query_coord.pb.go b/internal/proto/querypb/query_coord.pb.go
diff --git a/internal/querycoordv2/observers/leader_observer.go b/internal/querycoordv2/observers/leader_observer.go
@@ -216,6 +216,14 @@ func (o *LeaderObserver) sync(ctx context.Context, replicaID int64, leaderView *
 		log.Warn("sync distribution failed, cannot get schema of collection", zap.Error(err))
 		return false
 	}
+
+	// Get collection index info
+	indexInfo, err := o.broker.DescribeIndex(ctx, leaderView.CollectionID)
+	if err != nil {
+		log.Warn("fail to get index info of collection", zap.Error(err))
+		return false
+	}
+
 	partitions, err := utils.GetPartitions(o.meta.CollectionManager, leaderView.CollectionID)
 	if err != nil {
 		log.Warn("sync distribution failed, cannot get partitions of collection", zap.Error(err))
@@ -236,7 +244,8 @@ func (o *LeaderObserver) sync(ctx context.Context, replicaID int64, leaderView *
 			CollectionID: leaderView.CollectionID,
 			PartitionIDs: partitions,
 		},
-		Version: time.Now().UnixNano(),
+		Version:       time.Now().UnixNano(),
+		IndexInfoList: indexInfo,
 	}
 	ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.SegmentTaskTimeout.GetAsDuration(time.Millisecond))
 	defer cancel()

diff --git a/internal/querycoordv2/observers/leader_observer_test.go b/internal/querycoordv2/observers/leader_observer_test.go
@@ -18,6 +18,7 @@ package observers
 
 import (
 	"context"
+	"fmt"
 	"testing"
 	"time"
 
@@ -31,6 +32,7 @@ import (
 	etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
 	"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
+	"github.com/milvus-io/milvus/internal/proto/indexpb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/querycoordv2/meta"
 	. "github.com/milvus-io/milvus/internal/querycoordv2/params"
@@ -120,6 +122,11 @@ func (suite *LeaderObserverTestSuite) TestSyncLoadedSegments() {
 	suite.broker.EXPECT().GetCollectionSchema(mock.Anything, int64(1)).Return(schema, nil)
 	suite.broker.EXPECT().GetSegmentInfo(mock.Anything, int64(1)).Return(
 		&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{info}}, nil)
+	// will cause sync failed once
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("mock error")).Once()
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return([]*indexpb.IndexInfo{
+		{IndexName: "test"},
+	}, nil)
 	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return(
 		channels, segments, nil)
 	observer.target.UpdateCollectionNextTarget(int64(1))
@@ -154,7 +161,8 @@ func (suite *LeaderObserverTestSuite) TestSyncLoadedSegments() {
 				CollectionID: 1,
 				PartitionIDs: []int64{1},
 			},
-			Version: version,
+			Version:       version,
+			IndexInfoList: []*indexpb.IndexInfo{{IndexName: "test"}},
 		}
 	}
 
@@ -212,6 +220,9 @@ func (suite *LeaderObserverTestSuite) TestIgnoreSyncLoadedSegments() {
 		&datapb.GetSegmentInfoResponse{Infos: []*datapb.SegmentInfo{info}}, nil)
 	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return(
 		channels, segments, nil)
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return([]*indexpb.IndexInfo{
+		{IndexName: "test"},
+	}, nil)
 	observer.target.UpdateCollectionNextTarget(int64(1))
 	observer.target.UpdateCollectionCurrentTarget(1)
 	observer.target.UpdateCollectionNextTarget(int64(1))
@@ -246,7 +257,8 @@ func (suite *LeaderObserverTestSuite) TestIgnoreSyncLoadedSegments() {
 				CollectionID: 1,
 				PartitionIDs: []int64{1},
 			},
-			Version: version,
+			Version:       version,
+			IndexInfoList: []*indexpb.IndexInfo{{IndexName: "test"}},
 		}
 	}
 	called := atomic.NewBool(false)
@@ -344,6 +356,7 @@ func (suite *LeaderObserverTestSuite) TestSyncLoadedSegmentsWithReplicas() {
 	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return(
 		channels, segments, nil)
 	suite.broker.EXPECT().GetCollectionSchema(mock.Anything, int64(1)).Return(schema, nil)
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return([]*indexpb.IndexInfo{{IndexName: "test"}}, nil)
 	observer.target.UpdateCollectionNextTarget(int64(1))
 	observer.target.UpdateCollectionCurrentTarget(1)
 	observer.dist.SegmentDistManager.Update(1, utils.CreateTestSegment(1, 1, 1, 1, 1, "test-insert-channel"))
@@ -380,7 +393,8 @@ func (suite *LeaderObserverTestSuite) TestSyncLoadedSegmentsWithReplicas() {
 				CollectionID: 1,
 				PartitionIDs: []int64{1},
 			},
-			Version: version,
+			Version:       version,
+			IndexInfoList: []*indexpb.IndexInfo{{IndexName: "test"}},
 		}
 	}
 	called := atomic.NewBool(false)
@@ -412,7 +426,9 @@ func (suite *LeaderObserverTestSuite) TestSyncRemovedSegments() {
 
 	schema := utils.CreateTestSchema()
 	suite.broker.EXPECT().GetCollectionSchema(mock.Anything, int64(1)).Return(schema, nil)
-
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return([]*indexpb.IndexInfo{
+		{IndexName: "test"},
+	}, nil)
 	channels := []*datapb.VchannelInfo{
 		{
 			CollectionID: 1,
@@ -450,7 +466,8 @@ func (suite *LeaderObserverTestSuite) TestSyncRemovedSegments() {
 				CollectionID: 1,
 				PartitionIDs: []int64{1},
 			},
-			Version: version,
+			Version:       version,
+			IndexInfoList: []*indexpb.IndexInfo{{IndexName: "test"}},
 		}
 	}
 	ch := make(chan struct{})
@@ -494,6 +511,9 @@ func (suite *LeaderObserverTestSuite) TestIgnoreSyncRemovedSegments() {
 	suite.broker.EXPECT().GetCollectionSchema(mock.Anything, int64(1)).Return(schema, nil)
 	suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, int64(1)).Return(
 		channels, segments, nil)
+	suite.broker.EXPECT().DescribeIndex(mock.Anything, mock.Anything).Return([]*indexpb.IndexInfo{
+		{IndexName: "test"},
+	}, nil)
 	observer.target.UpdateCollectionNextTarget(int64(1))
 
 	observer.dist.ChannelDistManager.Update(2, utils.CreateTestChannel(1, 2, 1, "test-insert-channel"))
@@ -519,7 +539,8 @@ func (suite *LeaderObserverTestSuite) TestIgnoreSyncRemovedSegments() {
 				CollectionID: 1,
 				PartitionIDs: []int64{1},
 			},
-			Version: version,
+			Version:       version,
+			IndexInfoList: []*indexpb.IndexInfo{{IndexName: "test"}},
 		}
 	}
 	called := atomic.NewBool(false)

diff --git a/internal/querycoordv2/observers/target_observer.go b/internal/querycoordv2/observers/target_observer.go
@@ -353,6 +353,13 @@ func (ob *TargetObserver) sync(ctx context.Context, replicaID int64, leaderView
 		return false
 	}
 
+	// Get collection index info
+	indexInfo, err := ob.broker.DescribeIndex(ctx, leaderView.CollectionID)
+	if err != nil {
+		log.Warn("fail to get index info of collection", zap.Error(err))
+		return false
+	}
+
 	req := &querypb.SyncDistributionRequest{
 		Base: commonpbutil.NewMsgBase(
 			commonpbutil.WithMsgType(commonpb.MsgType_SyncDistribution),
@@ -367,7 +374,8 @@ func (ob *TargetObserver) sync(ctx context.Context, replicaID int64, leaderView
 			CollectionID: leaderView.CollectionID,
 			PartitionIDs: partitions,
 		},
-		Version: time.Now().UnixNano(),
+		Version:       time.Now().UnixNano(),
+		IndexInfoList: indexInfo,
 	}
 	ctx, cancel := context.WithTimeout(ctx, paramtable.Get().QueryCoordCfg.SegmentTaskTimeout.GetAsDuration(time.Millisecond))
 	defer cancel()

diff --git a/internal/querycoordv2/observers/target_observer_test.go b/internal/querycoordv2/observers/target_observer_test.go
@@ -213,7 +213,7 @@ func (suite *TargetObserverSuite) TestTriggerUpdateTarget() {
 	}, 7*time.Second, 1*time.Second)
 }
 
-func (suite *TargetObserverSuite) TearDownSuite() {
+func (suite *TargetObserverSuite) TearDownTest() {
 	suite.kv.Close()
 	suite.observer.Stop()
 }

diff --git a/internal/querynodev2/services.go b/internal/querynodev2/services.go
@@ -1360,15 +1360,16 @@ func (node *QueryNode) SyncDistribution(ctx context.Context, req *querypb.SyncDi
 						commonpbutil.WithMsgType(commonpb.MsgType_LoadSegments),
 						commonpbutil.WithMsgID(req.Base.GetMsgID()),
 					),
-					Infos:        []*querypb.SegmentLoadInfo{action.GetInfo()},
-					Schema:       req.GetSchema(),
-					LoadMeta:     req.GetLoadMeta(),
-					CollectionID: req.GetCollectionID(),
-					ReplicaID:    req.GetReplicaID(),
-					DstNodeID:    action.GetNodeID(),
-					Version:      action.GetVersion(),
-					NeedTransfer: false,
-					LoadScope:    querypb.LoadScope_Delta,
+					Infos:         []*querypb.SegmentLoadInfo{action.GetInfo()},
+					Schema:        req.GetSchema(),
+					LoadMeta:      req.GetLoadMeta(),
+					CollectionID:  req.GetCollectionID(),
+					ReplicaID:     req.GetReplicaID(),
+					DstNodeID:     action.GetNodeID(),
+					Version:       action.GetVersion(),
+					NeedTransfer:  false,
+					LoadScope:     querypb.LoadScope_Delta,
+					IndexInfoList: req.GetIndexInfoList(),
 				})
 			})
 		case querypb.SyncType_UpdateVersion: