Implement GetIndexState in Proxy

Signed-off-by: Ndragondriver <jiquan.long@zilliz.com>

internal/proxynode/impl.go

@@ -442,8 +442,10 @@ func (node *NodeImpl) GetIndexState(request *milvuspb.IndexStateRequest) (*milvu
 	ctx, cancel := context.WithTimeout(context.Background(), reqTimeoutInterval)
 	defer cancel()
 	dipt := &GetIndexStateTask{
-		Condition:         NewTaskCondition(ctx),
-		IndexStateRequest: request,
+		Condition:             NewTaskCondition(ctx),
+		IndexStateRequest:     request,
+		indexServiceClient:    node.indexServiceClient,
+		masterClientInterface: node.masterClient,
 	}
 
 	err := node.sched.DdQueue.Enqueue(dipt)

internal/proxynode/interface.go

@@ -23,6 +23,7 @@ type MasterClient interface {
 	CreateIndex(in *milvuspb.CreateIndexRequest) (*commonpb.Status, error)
 	DescribeIndex(in *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error)
 	ShowSegments(in *milvuspb.ShowSegmentRequest) (*milvuspb.ShowSegmentResponse, error)
+	DescribeSegment(in *milvuspb.DescribeSegmentRequest) (*milvuspb.DescribeSegmentResponse, error)
 }
 
 type IndexServiceClient interface {

internal/proxynode/task.go

@@ -6,6 +6,8 @@ import (
 	"math"
 	"strconv"
 
+	"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
+
 	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 
 	"github.com/golang/protobuf/proto"
@@ -1285,8 +1287,9 @@ func (dit *DescribeIndexTask) PostExecute() error {
 type GetIndexStateTask struct {
 	Condition
 	*milvuspb.IndexStateRequest
-	indexServiceClient IndexServiceClient
-	result             *milvuspb.IndexStateResponse
+	indexServiceClient    IndexServiceClient
+	masterClientInterface MasterClient
+	result                *milvuspb.IndexStateResponse
 }
 
 func (dipt *GetIndexStateTask) OnEnqueue() error {
@@ -1336,17 +1339,98 @@ func (dipt *GetIndexStateTask) PreExecute() error {
 }
 
 func (dipt *GetIndexStateTask) Execute() error {
-	// TODO: use index service client
-	//var err error
-	//dipt.result, err = dipt.masterClient.GetIndexState(dipt.IndexStateRequest)
-	//return err
+	collectionName := dipt.CollectionName
+	collectionID, err := globalMetaCache.GetCollectionID(collectionName)
+	if err != nil { // err is not nil if collection not exists
+		return err
+	}
+
+	showPartitionRequest := &milvuspb.ShowPartitionRequest{
+		Base: &commonpb.MsgBase{
+			MsgType:   commonpb.MsgType_kShowPartitions,
+			MsgID:     dipt.Base.MsgID,
+			Timestamp: dipt.Base.Timestamp,
+			SourceID:  Params.ProxyID,
+		},
+		DbName:         dipt.DbName,
+		CollectionName: collectionName,
+		CollectionID:   collectionID,
+	}
+	partitions, err := dipt.masterClientInterface.ShowPartitions(showPartitionRequest)
+	if err != nil {
+		return err
+	}
+
+	for _, partitionID := range partitions.PartitionIDs {
+		showSegmentsRequest := &milvuspb.ShowSegmentRequest{
+			Base: &commonpb.MsgBase{
+				MsgType:   commonpb.MsgType_kShowSegment,
+				MsgID:     dipt.Base.MsgID,
+				Timestamp: dipt.Base.Timestamp,
+				SourceID:  Params.ProxyID,
+			},
+			CollectionID: collectionID,
+			PartitionID:  partitionID,
+		}
+		segments, err := dipt.masterClientInterface.ShowSegments(showSegmentsRequest)
+		if err != nil {
+			return err
+		}
+
+		getIndexStatesRequest := &indexpb.IndexStatesRequest{
+			IndexBuildIDs: make([]UniqueID, 0),
+		}
+		for _, segmentID := range segments.SegmentIDs {
+			describeSegmentRequest := &milvuspb.DescribeSegmentRequest{
+				Base: &commonpb.MsgBase{
+					MsgType:   commonpb.MsgType_kDescribeSegment,
+					MsgID:     dipt.Base.MsgID,
+					Timestamp: dipt.Base.Timestamp,
+					SourceID:  Params.ProxyID,
+				},
+				CollectionID: collectionID,
+				SegmentID:    segmentID,
+			}
+			segmentDesc, err := dipt.masterClientInterface.DescribeSegment(describeSegmentRequest)
+			if err != nil {
+				return err
+			}
+
+			getIndexStatesRequest.IndexBuildIDs = append(getIndexStatesRequest.IndexBuildIDs, segmentDesc.BuildID)
+		}
+
+		states, err := dipt.indexServiceClient.GetIndexStates(getIndexStatesRequest)
+		if err != nil {
+			return err
+		}
+
+		if states.Status.ErrorCode != commonpb.ErrorCode_SUCCESS {
+			dipt.result = &milvuspb.IndexStateResponse{
+				Status: states.Status,
+				State:  commonpb.IndexState_FAILED,
+			}
+			return nil
+		}
+
+		for _, state := range states.States {
+			if state.State != commonpb.IndexState_FINISHED {
+				dipt.result = &milvuspb.IndexStateResponse{
+					Status: states.Status,
+					State:  state.State,
+				}
+				return nil
+			}
+		}
+	}
+
 	dipt.result = &milvuspb.IndexStateResponse{
 		Status: &commonpb.Status{
-			ErrorCode: 0,
+			ErrorCode: commonpb.ErrorCode_SUCCESS,
 			Reason:    "",
 		},
 		State: commonpb.IndexState_FINISHED,
 	}
+
 	return nil
 }
 

internal/querynode/flow_graph_gc_node.go

@@ -21,27 +21,31 @@ func (gcNode *gcNode) Operate(in []*Msg) []*Msg {
 		// TODO: add error handling
 	}
 
-	gcMsg, ok := (*in[0]).(*gcMsg)
+	_, ok := (*in[0]).(*gcMsg)
 	if !ok {
 		log.Println("type assertion failed for gcMsg")
 		// TODO: add error handling
 	}
 
-	// drop collections
-	for _, collectionID := range gcMsg.gcRecord.collections {
-		err := gcNode.replica.removeCollection(collectionID)
-		if err != nil {
-			log.Println(err)
-		}
-	}
+	// Use `releasePartition` and `releaseCollection`,
+	// because if we drop collections or partitions here, query service doesn't know this behavior,
+	// which would lead the wrong result of `showCollections` or `showPartition`
 
-	// drop partitions
-	for _, partition := range gcMsg.gcRecord.partitions {
-		err := gcNode.replica.removePartition(partition.collectionID, partition.partitionID)
-		if err != nil {
-			log.Println(err)
-		}
-	}
+	//// drop collections
+	//for _, collectionID := range gcMsg.gcRecord.collections {
+	//	err := gcNode.replica.removeCollection(collectionID)
+	//	if err != nil {
+	//		log.Println(err)
+	//	}
+	//}
+	//
+	//// drop partitions
+	//for _, partition := range gcMsg.gcRecord.partitions {
+	//	err := gcNode.replica.removePartition(partition.collectionID, partition.partitionID)
+	//	if err != nil {
+	//		log.Println(err)
+	//	}
+	//}
 
 	return nil
 }