Get index info from master

Signed-off-by: Nbigsheeper <yihao.dai@zilliz.com>

configs/advanced/data_service.yaml

@@ -10,4 +10,4 @@ dataservice:
     # old name: segmentExpireDuration: 2000
     IDAssignExpiration: 2000 # ms
   insertChannelNumPerCollection: 4
-  dataNodeNum: 2
+  dataNodeNum: 1
\ No newline at end of file

internal/core/src/pb/milvus.pb.h

@@ -4073,6 +4073,7 @@ class DescribeSegmentResponse :
 
   enum : int {
     kStatusFieldNumber = 1,
+    kIndexDescriptionFieldNumber = 3,
     kIndexIDFieldNumber = 2,
   };
   // .milvus.proto.common.Status status = 1;
@@ -4083,6 +4084,14 @@ class DescribeSegmentResponse :
   ::milvus::proto::common::Status* mutable_status();
   void set_allocated_status(::milvus::proto::common::Status* status);
 
+  // .milvus.proto.milvus.IndexDescription index_description = 3;
+  bool has_index_description() const;
+  void clear_index_description();
+  const ::milvus::proto::milvus::IndexDescription& index_description() const;
+  ::milvus::proto::milvus::IndexDescription* release_index_description();
+  ::milvus::proto::milvus::IndexDescription* mutable_index_description();
+  void set_allocated_index_description(::milvus::proto::milvus::IndexDescription* index_description);
+
   // int64 indexID = 2;
   void clear_indexid();
   ::PROTOBUF_NAMESPACE_ID::int64 indexid() const;
@@ -4094,6 +4103,7 @@ class DescribeSegmentResponse :
 
   ::PROTOBUF_NAMESPACE_ID::internal::InternalMetadataWithArena _internal_metadata_;
   ::milvus::proto::common::Status* status_;
+  ::milvus::proto::milvus::IndexDescription* index_description_;
   ::PROTOBUF_NAMESPACE_ID::int64 indexid_;
   mutable ::PROTOBUF_NAMESPACE_ID::internal::CachedSize _cached_size_;
   friend struct ::TableStruct_milvus_2eproto;
@@ -10312,6 +10322,57 @@ inline void DescribeSegmentResponse::set_indexid(::PROTOBUF_NAMESPACE_ID::int64
   // @@protoc_insertion_point(field_set:milvus.proto.milvus.DescribeSegmentResponse.indexID)
 }
 
+// .milvus.proto.milvus.IndexDescription index_description = 3;
+inline bool DescribeSegmentResponse::has_index_description() const {
+  return this != internal_default_instance() && index_description_ != nullptr;
+}
+inline void DescribeSegmentResponse::clear_index_description() {
+  if (GetArenaNoVirtual() == nullptr && index_description_ != nullptr) {
+    delete index_description_;
+  }
+  index_description_ = nullptr;
+}
+inline const ::milvus::proto::milvus::IndexDescription& DescribeSegmentResponse::index_description() const {
+  const ::milvus::proto::milvus::IndexDescription* p = index_description_;
+  // @@protoc_insertion_point(field_get:milvus.proto.milvus.DescribeSegmentResponse.index_description)
+  return p != nullptr ? *p : *reinterpret_cast<const ::milvus::proto::milvus::IndexDescription*>(
+      &::milvus::proto::milvus::_IndexDescription_default_instance_);
+}
+inline ::milvus::proto::milvus::IndexDescription* DescribeSegmentResponse::release_index_description() {
+  // @@protoc_insertion_point(field_release:milvus.proto.milvus.DescribeSegmentResponse.index_description)
+  
+  ::milvus::proto::milvus::IndexDescription* temp = index_description_;
+  index_description_ = nullptr;
+  return temp;
+}
+inline ::milvus::proto::milvus::IndexDescription* DescribeSegmentResponse::mutable_index_description() {
+  
+  if (index_description_ == nullptr) {
+    auto* p = CreateMaybeMessage<::milvus::proto::milvus::IndexDescription>(GetArenaNoVirtual());
+    index_description_ = p;
+  }
+  // @@protoc_insertion_point(field_mutable:milvus.proto.milvus.DescribeSegmentResponse.index_description)
+  return index_description_;
+}
+inline void DescribeSegmentResponse::set_allocated_index_description(::milvus::proto::milvus::IndexDescription* index_description) {
+  ::PROTOBUF_NAMESPACE_ID::Arena* message_arena = GetArenaNoVirtual();
+  if (message_arena == nullptr) {
+    delete index_description_;
+  }
+  if (index_description) {
+    ::PROTOBUF_NAMESPACE_ID::Arena* submessage_arena = nullptr;
+    if (message_arena != submessage_arena) {
+      index_description = ::PROTOBUF_NAMESPACE_ID::internal::GetOwnedMessage(
+          message_arena, index_description, submessage_arena);
+    }
+    
+  } else {
+    
+  }
+  index_description_ = index_description;
+  // @@protoc_insertion_point(field_set_allocated:milvus.proto.milvus.DescribeSegmentResponse.index_description)
+}
+
 // -------------------------------------------------------------------
 
 // ShowSegmentRequest

internal/dataservice/cluster.go

@@ -41,7 +41,7 @@ func newDataNodeCluster(finishCh chan struct{}) *dataNodeCluster {
 func (c *dataNodeCluster) Register(ip string, port int64, id int64) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if !c.checkDataNodeNotExist(ip, port) {
+	if c.checkDataNodeNotExist(ip, port) {
 		c.nodes = append(c.nodes, &dataNode{
 			id: id,
 			address: struct {
@@ -50,9 +50,9 @@ func (c *dataNodeCluster) Register(ip string, port int64, id int64) {
 			}{ip: ip, port: port},
 			channelNum: 0,
 		})
-	}
-	if len(c.nodes) == Params.DataNodeNum {
-		close(c.finishCh)
+		if len(c.nodes) == Params.DataNodeNum {
+			close(c.finishCh)
+		}
 	}
 }
 

internal/distributed/querynode/service.go

@@ -5,14 +5,17 @@ import (
 	"fmt"
 	"log"
 	"net"
+	"time"
 
 	"google.golang.org/grpc"
 
 	"github.com/zilliztech/milvus-distributed/internal/dataservice"
 	grpcdataserviceclient "github.com/zilliztech/milvus-distributed/internal/distributed/dataservice"
 	grpcindexserviceclient "github.com/zilliztech/milvus-distributed/internal/distributed/indexservice/client"
+	grpcmasterserviceclient "github.com/zilliztech/milvus-distributed/internal/distributed/masterservice"
 	grpcqueryserviceclient "github.com/zilliztech/milvus-distributed/internal/distributed/queryservice/client"
 	"github.com/zilliztech/milvus-distributed/internal/indexservice"
+	"github.com/zilliztech/milvus-distributed/internal/masterservice"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
@@ -36,6 +39,16 @@ func NewServer(ctx context.Context) *Server {
 		panic(err)
 	}
 
+	masterservice.Params.Init()
+	masterConnectTimeout := 10 * time.Second
+	masterClient, err := grpcmasterserviceclient.NewGrpcClient(masterservice.Params.Address, masterConnectTimeout)
+	if err != nil {
+		panic(err)
+	}
+	if err = server.node.SetMasterService(masterClient); err != nil {
+		panic(err)
+	}
+
 	indexservice.Params.Init()
 	indexClient := grpcindexserviceclient.NewClient(indexservice.Params.Address)
 	if err := server.node.SetIndexService(indexClient); err != nil {

internal/masterservice/meta_table.go

@@ -536,9 +536,12 @@ func (mt *metaTable) GetSegmentIndexInfoByID(segID typeutil.UniqueID, filedID in
 		return pb.SegmentIndexInfo{}, errors.Errorf("segment id %d not has any index", segID)
 	}
 
-	if filedID == -1 && idxName == "" { // return any index
+	if filedID == -1 && idxName == "" { // return default index
 		for _, seg := range *segIdxMap {
-			return seg, nil
+			info, ok := mt.indexID2Meta[seg.IndexID]
+			if ok && info.IndexName == Params.DefaultIndexName {
+				return seg, nil
+			}
 		}
 	} else {
 		for idxID, seg := range *segIdxMap {
@@ -705,3 +708,14 @@ func (mt *metaTable) GetIndexByName(collName string, fieldName string, indexName
 	}
 	return rst, nil
 }
+
+func (mt *metaTable) GetIndexByID(indexID typeutil.UniqueID) (*pb.IndexInfo, error) {
+	mt.ddLock.RLock()
+	mt.ddLock.RUnlock()
+
+	indexInfo, ok := mt.indexID2Meta[indexID]
+	if !ok {
+		return nil, errors.New("cannot find index, id =" + strconv.FormatInt(indexID, 10))
+	}
+	return &indexInfo, nil
+}

internal/masterservice/task.go

@@ -480,6 +480,14 @@ func (t *DescribeSegmentReqTask) Execute() error {
 		return err
 	}
 	t.Rsp.IndexID = segIdxInfo.IndexID
+	indexInfo, err := t.core.MetaTable.GetIndexByID(segIdxInfo.IndexID)
+	if err != nil {
+		return err
+	}
+	t.Rsp.IndexDescription = &milvuspb.IndexDescription{
+		IndexName: Params.DefaultIndexName, // TODO: set index name
+		Params:    indexInfo.IndexParams,
+	}
 	return nil
 }
 

internal/proto/milvus.proto

@@ -154,6 +154,7 @@ message DescribeSegmentRequest {
 message DescribeSegmentResponse {
   common.Status status = 1;
   int64 indexID = 2;
+  IndexDescription index_description = 3;
 }
 
 message ShowSegmentRequest {

internal/proxyservice/nodeid_allocator.go

 package proxyservice
 
 import (
-	"sync"
+	"context"
 
 	"github.com/zilliztech/milvus-distributed/internal/allocator"
 
@@ -17,21 +17,22 @@ type NodeIDAllocator interface {
 
 type NaiveNodeIDAllocatorImpl struct {
 	impl *allocator.IDAllocator
-	now  UniqueID
-	mtx  sync.Mutex
 }
 
 func (allocator *NaiveNodeIDAllocatorImpl) AllocOne() UniqueID {
-	allocator.mtx.Lock()
-	defer func() {
-		allocator.now++
-		allocator.mtx.Unlock()
-	}()
-	return allocator.now
+	id, err := allocator.impl.AllocOne()
+	if err != nil {
+		panic(err)
+	}
+	return id
 }
 
 func NewNodeIDAllocator() NodeIDAllocator {
+	impl, err := allocator.NewIDAllocator(context.Background(), Params.MasterAddress())
+	if err != nil {
+		panic(err)
+	}
 	return &NaiveNodeIDAllocatorImpl{
-		now: 0,
+		impl: impl,
 	}
 }

internal/querynode/query_node.go

@@ -16,7 +16,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 	"io"
 	"log"
 	"sync/atomic"
@@ -27,6 +26,7 @@ import (
 	"github.com/zilliztech/milvus-distributed/internal/msgstream/pulsarms"
 	"github.com/zilliztech/milvus-distributed/internal/msgstream/util"
 	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
+	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
 	queryPb "github.com/zilliztech/milvus-distributed/internal/proto/querypb"
 	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
@@ -69,9 +69,10 @@ type QueryNode struct {
 	closer io.Closer
 
 	// clients
-	queryClient QueryServiceInterface
-	indexClient IndexServiceInterface
-	dataClient  DataServiceInterface
+	masterClient MasterServiceInterface
+	queryClient  QueryServiceInterface
+	indexClient  IndexServiceInterface
+	dataClient   DataServiceInterface
 }
 
 func NewQueryNode(ctx context.Context, queryNodeID uint64) *QueryNode {
@@ -203,7 +204,7 @@ func (node *QueryNode) Start() error {
 	node.metaService = newMetaService(node.queryNodeLoopCtx, node.replica)
 	node.loadIndexService = newLoadIndexService(node.queryNodeLoopCtx, node.replica)
 	node.statsService = newStatsService(node.queryNodeLoopCtx, node.replica, node.loadIndexService.fieldStatsChan)
-	node.segManager = newSegmentManager(node.queryNodeLoopCtx, node.dataClient, node.indexClient, node.replica, node.dataSyncService.dmStream, node.loadIndexService.loadIndexReqChan)
+	node.segManager = newSegmentManager(node.queryNodeLoopCtx, node.masterClient, node.dataClient, node.indexClient, node.replica, node.dataSyncService.dmStream, node.loadIndexService.loadIndexReqChan)
 
 	// start services
 	go node.dataSyncService.start()
@@ -243,9 +244,17 @@ func (node *QueryNode) Stop() error {
 	return nil
 }
 
+func (node *QueryNode) SetMasterService(master MasterServiceInterface) error {
+	if master == nil {
+		return errors.New("null master service interface")
+	}
+	node.masterClient = master
+	return nil
+}
+
 func (node *QueryNode) SetQueryService(query QueryServiceInterface) error {
 	if query == nil {
-		return errors.New("query index service interface")
+		return errors.New("null query service interface")
 	}
 	node.queryClient = query
 	return nil

internal/querynode/segment_manager.go

@@ -12,6 +12,7 @@ import (
 	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
 	internalPb "github.com/zilliztech/milvus-distributed/internal/proto/internalpb2"
+	"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
 	"github.com/zilliztech/milvus-distributed/internal/storage"
 )
 
@@ -21,8 +22,9 @@ type segmentManager struct {
 	dmStream         msgstream.MsgStream
 	loadIndexReqChan chan []msgstream.TsMsg
 
-	dataClient  DataServiceInterface
-	indexClient IndexServiceInterface
+	masterClient MasterServiceInterface
+	dataClient   DataServiceInterface
+	indexClient  IndexServiceInterface
 
 	kv     kv.Base // minio kv
 	iCodec *storage.InsertCodec
@@ -39,6 +41,29 @@ func (s *segmentManager) seekSegment(positions []*internalPb.MsgPosition) error
 	return nil
 }
 
+func (s *segmentManager) getIndexInfo(collectionID UniqueID, segmentID UniqueID) (UniqueID, indexParam, error) {
+	req := &milvuspb.DescribeSegmentRequest{
+		Base: &commonpb.MsgBase{
+			MsgType: commonpb.MsgType_kDescribeSegment,
+		},
+		CollectionID: collectionID,
+		SegmentID:    segmentID,
+	}
+	response, err := s.masterClient.DescribeSegment(req)
+	if err != nil {
+		return 0, nil, err
+	}
+	if len(response.IndexDescription.Params) <= 0 {
+		return 0, nil, errors.New("null index param")
+	}
+
+	var targetIndexParam = make(map[string]string)
+	for _, param := range response.IndexDescription.Params {
+		targetIndexParam[param.Key] = param.Value
+	}
+	return response.IndexID, targetIndexParam, nil
+}
+
 func (s *segmentManager) loadSegment(collectionID UniqueID, partitionID UniqueID, segmentIDs []UniqueID, fieldIDs []int64) error {
 	// TODO: interim solution
 	if len(fieldIDs) == 0 {
@@ -52,7 +77,10 @@ func (s *segmentManager) loadSegment(collectionID UniqueID, partitionID UniqueID
 		}
 	}
 	for _, segmentID := range segmentIDs {
-		indexID := UniqueID(0) // TODO: get index id from master
+		indexID, indexParams, err := s.getIndexInfo(collectionID, segmentID)
+		if err != nil {
+			return err
+		}
 		paths, srcFieldIDs, err := s.getInsertBinlogPaths(segmentID)
 		if err != nil {
 			return err
@@ -76,11 +104,7 @@ func (s *segmentManager) loadSegment(collectionID UniqueID, partitionID UniqueID
 		if err != nil {
 			return err
 		}
-		iParam, err := s.getIndexParam()
-		if err != nil {
-			return err
-		}
-		err = s.loadIndex(segmentID, indexPaths, iParam)
+		err = s.loadIndex(segmentID, indexPaths, indexParams)
 		if err != nil {
 			// TODO: return or continue?
 			return err
@@ -236,12 +260,6 @@ func (s *segmentManager) getIndexPaths(indexID UniqueID) ([]string, error) {
 	return pathResponse.FilePaths[0].IndexFilePaths, nil
 }
 
-func (s *segmentManager) getIndexParam() (indexParam, error) {
-	var targetIndexParam indexParam
-	// TODO: get index param from master
-	return targetIndexParam, nil
-}
-
 func (s *segmentManager) loadIndex(segmentID UniqueID, indexPaths []string, indexParam indexParam) error {
 	// get vector field ids from schema to load index
 	vecFieldIDs, err := s.replica.getVecFieldsBySegmentID(segmentID)
@@ -288,7 +306,7 @@ func (s *segmentManager) sendLoadIndex(indexPaths []string,
 	s.loadIndexReqChan <- messages
 }
 
-func newSegmentManager(ctx context.Context, dataClient DataServiceInterface, indexClient IndexServiceInterface, replica collectionReplica, dmStream msgstream.MsgStream, loadIndexReqChan chan []msgstream.TsMsg) *segmentManager {
+func newSegmentManager(ctx context.Context, masterClient MasterServiceInterface, dataClient DataServiceInterface, indexClient IndexServiceInterface, replica collectionReplica, dmStream msgstream.MsgStream, loadIndexReqChan chan []msgstream.TsMsg) *segmentManager {
 	bucketName := Params.MinioBucketName
 	option := &miniokv.Option{
 		Address:           Params.MinioEndPoint,
@@ -309,8 +327,9 @@ func newSegmentManager(ctx context.Context, dataClient DataServiceInterface, ind
 		dmStream:         dmStream,
 		loadIndexReqChan: loadIndexReqChan,
 
-		dataClient:  dataClient,
-		indexClient: indexClient,
+		masterClient: masterClient,
+		dataClient:   dataClient,
+		indexClient:  indexClient,
 
 		kv:     minioKV,
 		iCodec: &storage.InsertCodec{},

internal/querynode/segment_manager_test.go

@@ -405,7 +405,7 @@ func TestSegmentManager_load_release_and_search(t *testing.T) {
 
 	ctx := node.queryNodeLoopCtx
 	node.loadIndexService = newLoadIndexService(ctx, node.replica)
-	node.segManager = newSegmentManager(ctx, nil, nil, node.replica, nil, node.loadIndexService.loadIndexReqChan)
+	node.segManager = newSegmentManager(ctx, nil, nil, nil, node.replica, nil, node.loadIndexService.loadIndexReqChan)
 	go node.loadIndexService.start()
 
 	collectionName := "collection0"

internal/querynode/type_def.go

@@ -3,6 +3,7 @@ package querynode
 import (
 	"github.com/zilliztech/milvus-distributed/internal/proto/datapb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
+	"github.com/zilliztech/milvus-distributed/internal/proto/milvuspb"
 	"github.com/zilliztech/milvus-distributed/internal/proto/querypb"
 	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
 )
@@ -17,6 +18,11 @@ type TimeRange struct {
 	timestampMax Timestamp
 }
 
+type MasterServiceInterface interface {
+	DescribeSegment(in *milvuspb.DescribeSegmentRequest) (*milvuspb.DescribeSegmentResponse, error)
+	DescribeIndex(in *milvuspb.DescribeIndexRequest) (*milvuspb.DescribeIndexResponse, error)
+}
+
 type QueryServiceInterface interface {
 	RegisterNode(req *querypb.RegisterNodeRequest) (*querypb.RegisterNodeResponse, error)
 }