Add grpc interface for replica Search/Query in QueryNode (#16197)

Resolves #16195
Add Search and Query grpc interface in query proto and types/types.go
Signed-off-by: NCongqi Xia <congqi.xia@zilliz.com>

internal/distributed/querynode/client/client.go

@@ -241,6 +241,34 @@ func (c *Client) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmen
 	return ret.(*commonpb.Status), err
 }
 
+// Search performs replica search tasks in QueryNode.
+func (c *Client) Search(ctx context.Context, req *querypb.SearchRequest) (*milvuspb.SearchResults, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client interface{}) (interface{}, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.(querypb.QueryNodeClient).Search(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*milvuspb.SearchResults), err
+}
+
+// Query performs replica query tasks in QueryNode.
+func (c *Client) Query(ctx context.Context, req *querypb.QueryRequest) (*milvuspb.QueryResults, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client interface{}) (interface{}, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.(querypb.QueryNodeClient).Query(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*milvuspb.QueryResults), err
+}
+
 // GetSegmentInfo gets the information of the specified segments in QueryNode.
 func (c *Client) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
 	ret, err := c.grpcClient.ReCall(ctx, func(client interface{}) (interface{}, error) {

internal/distributed/querynode/client/client_test.go

@@ -48,6 +48,8 @@ func Test_NewClient(t *testing.T) {
 	err = client.Register()
 	assert.Nil(t, err)
 
+	ctx, cancel := context.WithCancel(ctx)
+
 	checkFunc := func(retNotNil bool) {
 		retCheck := func(notNil bool, ret interface{}, err error) {
 			if notNil {
@@ -97,6 +99,12 @@ func Test_NewClient(t *testing.T) {
 
 		r13, err := client.WatchDeltaChannels(ctx, nil)
 		retCheck(retNotNil, r13, err)
+
+		r14, err := client.Search(ctx, nil)
+		retCheck(retNotNil, r14, err)
+
+		r15, err := client.Query(ctx, nil)
+		retCheck(retNotNil, r15, err)
 	}
 
 	client.grpcClient = &mock.ClientBase{
@@ -132,6 +140,15 @@ func Test_NewClient(t *testing.T) {
 	client.grpcClient.SetNewGrpcClientFunc(newFunc3)
 
 	checkFunc(true)
+
+	// ctx canceled
+	client.grpcClient = &mock.ClientBase{
+		GetGrpcClientErr: nil,
+	}
+	client.grpcClient.SetNewGrpcClientFunc(newFunc1)
+	cancel() // make context canceled
+	checkFunc(false)
+
 	err = client.Stop()
 	assert.Nil(t, err)
 }

internal/distributed/querynode/service.go

@@ -304,6 +304,16 @@ func (s *Server) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfo
 	return s.querynode.GetSegmentInfo(ctx, req)
 }
 
+// Search performs search of streaming/historical replica on QueryNode.
+func (s *Server) Search(ctx context.Context, req *querypb.SearchRequest) (*milvuspb.SearchResults, error) {
+	return s.querynode.Search(ctx, req)
+}
+
+// Query performs query of streaming/historical replica on QueryNode.
+func (s *Server) Query(ctx context.Context, req *querypb.QueryRequest) (*milvuspb.QueryResults, error) {
+	return s.querynode.Query(ctx, req)
+}
+
 // GetMetrics gets the metrics information of QueryNode.
 func (s *Server) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
 	return s.querynode.GetMetrics(ctx, req)

internal/distributed/querynode/service_test.go

@@ -43,6 +43,8 @@ type MockQueryNode struct {
 	strResp    *milvuspb.StringResponse
 	infoResp   *querypb.GetSegmentInfoResponse
 	metricResp *milvuspb.GetMetricsResponse
+	searchResp *milvuspb.SearchResults
+	queryResp  *milvuspb.QueryResults
 }
 
 func (m *MockQueryNode) Init() error {
@@ -113,6 +115,14 @@ func (m *MockQueryNode) GetMetrics(ctx context.Context, req *milvuspb.GetMetrics
 	return m.metricResp, m.err
 }
 
+func (m *MockQueryNode) Search(ctx context.Context, req *querypb.SearchRequest) (*milvuspb.SearchResults, error) {
+	return m.searchResp, m.err
+}
+
+func (m *MockQueryNode) Query(ctx context.Context, req *querypb.QueryRequest) (*milvuspb.QueryResults, error) {
+	return m.queryResp, m.err
+}
+
 func (m *MockQueryNode) SetEtcdClient(client *clientv3.Client) {
 }
 
@@ -307,6 +317,20 @@ func Test_NewServer(t *testing.T) {
 		assert.Equal(t, commonpb.ErrorCode_Success, resp.Status.ErrorCode)
 	})
 
+	t.Run("Search", func(t *testing.T) {
+		req := &querypb.SearchRequest{}
+		resp, err := server.Search(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("Query", func(t *testing.T) {
+		req := &querypb.QueryRequest{}
+		resp, err := server.Query(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
 	err = server.Stop()
 	assert.Nil(t, err)
 }

internal/proto/query_coord.proto

@@ -47,6 +47,9 @@ service QueryNode {
   rpc ReleaseSegments(ReleaseSegmentsRequest) returns (common.Status) {}
   rpc GetSegmentInfo(GetSegmentInfoRequest) returns (GetSegmentInfoResponse) {}
 
+  rpc Search(SearchRequest) returns (milvus.SearchResults) {}
+  rpc Query(QueryRequest) returns (milvus.QueryResults) {}
+
   // https://wiki.lfaidata.foundation/display/MIL/MEP+8+--+Add+metrics+for+proxy
   rpc GetMetrics(milvus.GetMetricsRequest) returns (milvus.GetMetricsResponse) {}
 }
@@ -233,6 +236,18 @@ message ReleaseSegmentsRequest {
   repeated int64 segmentIDs = 6;
 }
 
+message SearchRequest {
+  milvus.SearchRequest req = 1;
+  string dml_channel = 2;
+  repeated int64 segmentIDs = 3;
+}
+
+message QueryRequest {
+  milvus.QueryRequest req = 1;
+  string dml_channel = 2;
+  repeated int64 segmentIDs = 3;
+}
+
 //----------------request auto triggered by QueryCoord-----------------
 message HandoffSegmentsRequest {
   common.MsgBase base = 1;
@@ -355,4 +370,5 @@ message ReplicaInfo {    // ReplicaGroup
 message ShardReplica {
     int64 leader = 1;
     string dm_channel_name = 2;
-}
\ No newline at end of file
+}
+

internal/querycoord/mock_querynode_client_test.go

@@ -157,3 +157,11 @@ func (client *queryNodeClientMock) GetSegmentInfo(ctx context.Context, req *quer
 func (client *queryNodeClientMock) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {
 	return client.grpcClient.GetMetrics(ctx, req)
 }
+
+func (client *queryNodeClientMock) Search(ctx context.Context, req *querypb.SearchRequest) (*milvuspb.SearchResults, error) {
+	return client.grpcClient.Search(ctx, req)
+}
+
+func (client *queryNodeClientMock) Query(ctx context.Context, req *querypb.QueryRequest) (*milvuspb.QueryResults, error) {
+	return client.grpcClient.Query(ctx, req)
+}

internal/querynode/impl.go

@@ -18,6 +18,7 @@ package querynode
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"go.uber.org/zap"
@@ -554,6 +555,16 @@ func (node *QueryNode) isHealthy() bool {
 	return code == internalpb.StateCode_Healthy
 }
 
+// Search performs replica search tasks.
+func (node *QueryNode) Search(ctx context.Context, req *queryPb.SearchRequest) (*milvuspb.SearchResults, error) {
+	return nil, errors.New("not implemented")
+}
+
+// Query performs replica query tasks.
+func (node *QueryNode) Query(ctx context.Context, req *queryPb.QueryRequest) (*milvuspb.QueryResults, error) {
+	return nil, errors.New("not implemented")
+}
+
 // GetMetrics return system infos of the query node, such as total memory, memory usage, cpu usage ...
 // TODO(dragondriver): cache the Metrics and set a retention to the cache
 func (node *QueryNode) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) {

internal/querynode/impl_test.go

@@ -25,6 +25,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 
 	"github.com/milvus-io/milvus/internal/proto/commonpb"
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
@@ -575,3 +576,25 @@ func TestImpl_ReleaseSegments(t *testing.T) {
 	})
 	wg.Wait()
 }
+
+func TestImpl_Search(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	node, err := genSimpleQueryNode(ctx)
+	require.NoError(t, err)
+
+	_, err = node.Search(ctx, nil)
+	assert.Error(t, err)
+}
+
+func TestImpl_Query(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	node, err := genSimpleQueryNode(ctx)
+	require.NoError(t, err)
+
+	_, err = node.Query(ctx, nil)
+	assert.Error(t, err)
+}

internal/types/types.go

@@ -1130,6 +1130,9 @@ type QueryNode interface {
 	ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error)
 	GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error)
 
+	Search(ctx context.Context, req *querypb.SearchRequest) (*milvuspb.SearchResults, error)
+	Query(ctx context.Context, req *querypb.QueryRequest) (*milvuspb.QueryResults, error)
+
 	// GetMetrics gets the metrics about QueryNode.
 	GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error)
 }

internal/util/mock/grpcclient.go

@@ -119,9 +119,7 @@ func (c *ClientBase) Call(ctx context.Context, caller func(client interface{}) (
 }
 
 func (c *ClientBase) ReCall(ctx context.Context, caller func(client interface{}) (interface{}, error)) (interface{}, error) {
-	if !funcutil.CheckCtxValid(ctx) {
-		return nil, ctx.Err()
-	}
+	// omit ctx check in mock first time to let each function has failed context
 	ret, err := c.callOnce(ctx, caller)
 	if err == nil {
 		return ret, nil

internal/util/mock/querynode_client.go

@@ -78,6 +78,13 @@ func (m *QueryNodeClient) GetSegmentInfo(ctx context.Context, in *querypb.GetSeg
 	return &querypb.GetSegmentInfoResponse{}, m.Err
 }
 
+func (m *QueryNodeClient) Search(ctx context.Context, in *querypb.SearchRequest, opts ...grpc.CallOption) (*milvuspb.SearchResults, error) {
+	return &milvuspb.SearchResults{}, m.Err
+}
+
+func (m *QueryNodeClient) Query(ctx context.Context, in *querypb.QueryRequest, opts ...grpc.CallOption) (*milvuspb.QueryResults, error) {
+	return &milvuspb.QueryResults{}, m.Err
+}
 func (m *QueryNodeClient) GetMetrics(ctx context.Context, in *milvuspb.GetMetricsRequest, opts ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) {
 	return &milvuspb.GetMetricsResponse{}, m.Err
 }