Refactor QueryCoord (milvus-io#18836)

Signed-off-by: sunby <[email protected]> Co-authored-by: yah01 <[email protected]> Co-authored-by: Wei Liu <[email protected]> Co-authored-by: Congqi Xia <[email protected]> Signed-off-by: sunby <[email protected]> Co-authored-by: sunby <[email protected]> Co-authored-by: yah01 <[email protected]> Co-authored-by: Wei Liu <[email protected]> Co-authored-by: Congqi Xia <[email protected]>
zhengbuqian · Sep 15, 2022 · 626854c · 626854c
1 parent 92857a4
commit 626854c
Show file tree

Hide file tree

Showing 129 changed files with 20,662 additions and 860 deletions.
diff --git a/configs/milvus.yaml b/configs/milvus.yaml
@@ -179,6 +179,14 @@ queryCoord:
   overloadedMemoryThresholdPercentage: 90 # The threshold percentage that memory overload
   balanceIntervalSeconds: 60
   memoryUsageMaxDifferencePercentage: 30
+  checkInterval: 1000
+  channelTaskTimeout: 60000 # 1 minute
+  segmentTaskTimeout: 15000 # 15 seconds
+  distPullInterval: 500
+  loadTimeoutSeconds: 600
+  checkHandoffInterval: 5000
+
+
 
 # Related configuration of queryNode, used to run hybrid search between vector and scalar data.
 queryNode:

diff --git a/internal/core/src/pb/common.pb.cc b/internal/core/src/pb/common.pb.cc
diff --git a/internal/core/src/pb/common.pb.h b/internal/core/src/pb/common.pb.h
diff --git a/internal/distributed/querycoord/service.go b/internal/distributed/querycoord/service.go
@@ -39,7 +39,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
 	"github.com/milvus-io/milvus/internal/proto/querypb"
-	qc "github.com/milvus-io/milvus/internal/querycoord"
+	qc "github.com/milvus-io/milvus/internal/querycoordv2"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/dependency"
 	"github.com/milvus-io/milvus/internal/util/etcd"

diff --git a/internal/distributed/querynode/client/client.go b/internal/distributed/querynode/client/client.go
@@ -152,6 +152,20 @@ func (c *Client) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChanne
 	return ret.(*commonpb.Status), err
 }
 
+// UnsubDmChannel unsubscribes the channels about data manipulation.
+func (c *Client) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest) (*commonpb.Status, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client interface{}) (interface{}, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.(querypb.QueryNodeClient).UnsubDmChannel(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*commonpb.Status), err
+}
+
 // LoadSegments loads the segments to search.
 func (c *Client) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
 	ret, err := c.grpcClient.ReCall(ctx, func(client interface{}) (interface{}, error) {
@@ -305,3 +319,29 @@ func (c *Client) GetStatistics(ctx context.Context, request *querypb.GetStatisti
 	}
 	return ret.(*internalpb.GetStatisticsResponse), err
 }
+
+func (c *Client) GetDataDistribution(ctx context.Context, req *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
+	ret, err := c.grpcClient.Call(ctx, func(client interface{}) (interface{}, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.(querypb.QueryNodeClient).GetDataDistribution(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*querypb.GetDataDistributionResponse), err
+}
+
+func (c *Client) SyncDistribution(ctx context.Context, req *querypb.SyncDistributionRequest) (*commonpb.Status, error) {
+	ret, err := c.grpcClient.Call(ctx, func(client interface{}) (interface{}, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.(querypb.QueryNodeClient).SyncDistribution(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*commonpb.Status), err
+}
diff --git a/internal/distributed/querynode/service.go b/internal/distributed/querynode/service.go
@@ -268,6 +268,10 @@ func (s *Server) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChanne
 	return s.querynode.WatchDmChannels(ctx, req)
 }
 
+func (s *Server) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest) (*commonpb.Status, error) {
+	return s.querynode.UnsubDmChannel(ctx, req)
+}
+
 // LoadSegments loads the segments to search.
 func (s *Server) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest) (*commonpb.Status, error) {
 	// ignore ctx
@@ -321,3 +325,12 @@ func (s *Server) ShowConfigurations(ctx context.Context, req *internalpb.ShowCon
 func (s *Server) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
 	return s.querynode.GetMetrics(ctx, req)
 }
+
+// GetDataDistribution gets the distribution information of QueryNode.
+func (s *Server) GetDataDistribution(ctx context.Context, req *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
+	return s.querynode.GetDataDistribution(ctx, req)
+}
+
+func (s *Server) SyncDistribution(ctx context.Context, req *querypb.SyncDistributionRequest) (*commonpb.Status, error) {
+	return s.querynode.SyncDistribution(ctx, req)
+}
diff --git a/internal/distributed/querynode/service_test.go b/internal/distributed/querynode/service_test.go
@@ -48,6 +48,7 @@ type MockQueryNode struct {
 	StatsResp  *internalpb.GetStatisticsResponse
 	searchResp *internalpb.SearchResults
 	queryResp  *internalpb.RetrieveResults
+	distResp   *querypb.GetDataDistributionResponse
 }
 
 func (m *MockQueryNode) Init() error {
@@ -140,7 +141,16 @@ func (m *MockQueryNode) ShowConfigurations(ctx context.Context, req *internalpb.
 	return m.configResp, m.err
 }
 
-///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+func (m *MockQueryNode) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+func (m *MockQueryNode) GetDataDistribution(context.Context, *querypb.GetDataDistributionRequest) (*querypb.GetDataDistributionResponse, error) {
+	return m.distResp, m.err
+}
+func (m *MockQueryNode) SyncDistribution(context.Context, *querypb.SyncDistributionRequest) (*commonpb.Status, error) {
+	return m.status, m.err
+}
+
 type MockRootCoord struct {
 	types.RootCoord
 	initErr  error

diff --git a/internal/kv/etcd/etcd_kv.go b/internal/kv/etcd/etcd_kv.go
@@ -27,9 +27,10 @@ import (
 
 	"github.com/milvus-io/milvus/internal/common"
 
-	"github.com/milvus-io/milvus/internal/log"
 	clientv3 "go.etcd.io/etcd/client/v3"
 
+	"github.com/milvus-io/milvus/internal/log"
+
 	"go.uber.org/zap"
 )
 
@@ -284,7 +285,7 @@ func (kv *EtcdKV) LoadWithRevision(key string) ([]string, []string, int64, error
 	ctx, cancel := context.WithTimeout(context.TODO(), RequestTimeout)
 	defer cancel()
 	resp, err := kv.client.Get(ctx, key, clientv3.WithPrefix(),
-		clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend))
+		clientv3.WithSort(clientv3.SortByCreateRevision, clientv3.SortAscend))
 	if err != nil {
 		return nil, nil, 0, err
 	}

diff --git a/internal/log/global.go b/internal/log/global.go
@@ -103,7 +103,7 @@ func RatedWarn(cost float64, msg string, fields ...zap.Field) bool {
 // With creates a child logger and adds structured context to it.
 // Fields added to the child don't affect the parent, and vice versa.
 func With(fields ...zap.Field) *zap.Logger {
-	return L().With(fields...)
+	return L().With(fields...).WithOptions(zap.AddCallerSkip(-1))
 }
 
 // SetLevel alters the logging level.

diff --git a/internal/metastore/catalog.go b/internal/metastore/catalog.go
@@ -6,6 +6,7 @@ import (
 	"github.com/milvus-io/milvus/internal/metastore/model"
 	"github.com/milvus-io/milvus/internal/proto/datapb"
 	"github.com/milvus-io/milvus/internal/proto/milvuspb"
+	"github.com/milvus-io/milvus/internal/proto/querypb"
 	"github.com/milvus-io/milvus/internal/util/typeutil"
 )
 
@@ -93,3 +94,17 @@ type IndexCoordCatalog interface {
 	AlterSegmentIndexes(ctx context.Context, newSegIdxes []*model.SegmentIndex) error
 	DropSegmentIndex(ctx context.Context, collID, partID, segID, buildID typeutil.UniqueID) error
 }
+
+type QueryCoordCatalog interface {
+	SaveCollection(info *querypb.CollectionLoadInfo) error
+	SavePartition(info ...*querypb.PartitionLoadInfo) error
+	SaveReplica(replica *querypb.Replica) error
+	GetCollections() ([]*querypb.CollectionLoadInfo, error)
+	GetPartitions() (map[int64][]*querypb.PartitionLoadInfo, error)
+	GetReplicas() ([]*querypb.Replica, error)
+	ReleaseCollection(id int64) error
+	ReleasePartition(collection int64, partitions ...int64) error
+	ReleaseReplicas(collectionID int64) error
+	ReleaseReplica(collection, replica int64) error
+	RemoveHandoffEvent(segmentInfo *querypb.SegmentInfo) error
+}