Implement GetIndexState in Proxy

Signed-off-by: dragondriver <jiquan.long@zilliz.com>
This commit is contained in:
dragondriver 2021-02-04 11:08:36 +08:00 committed by yefu.chen
parent bc7ad02df4
commit b92ff69cea
4 changed files with 115 additions and 24 deletions

View File

@ -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)

View File

@ -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 {

View File

@ -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
}

View File

@ -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
}