chore: import upstream snapshot with attribution

internal/distributed/querynode/api_testonly.go

@@ -0,0 +1,24 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build test
+// +build test
+
+package grpcquerynode
+
+func (s *Server) GetServerIDForTestOnly() int64 {
+	return s.serverID.Load()
+}

internal/distributed/querynode/client/client.go

@@ -0,0 +1,449 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcquerynodeclient
+
+import (
+	"context"
+	"fmt"
+
+	"google.golang.org/grpc"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
+	"github.com/milvus-io/milvus/internal/distributed/utils"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/internal/util/grpcclient"
+	"github.com/milvus-io/milvus/internal/util/sessionutil"
+	"github.com/milvus-io/milvus/pkg/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+var Params *paramtable.ComponentParam = paramtable.Get()
+
+// Client is the grpc client of QueryNode.
+type Client struct {
+	grpcClient grpcclient.GrpcClient[querypb.QueryNodeClient]
+	addr       string
+	nodeID     int64
+}
+
+// NewClient creates a new QueryNode client.
+func NewClient(ctx context.Context, addr string, nodeID int64) (types.QueryNodeClient, error) {
+	if addr == "" {
+		return nil, merr.WrapErrParameterInvalidMsg("addr is empty")
+	}
+	sess := sessionutil.NewSession(context.Background())
+	if sess == nil {
+		err := merr.WrapErrServiceUnavailable("new session error, maybe can not connect to etcd")
+		mlog.Debug(ctx, "QueryNodeClient NewClient failed", mlog.Err(err))
+		return nil, err
+	}
+	config := &paramtable.Get().QueryNodeGrpcClientCfg
+	client := &Client{
+		addr:       addr,
+		grpcClient: grpcclient.NewClientBase[querypb.QueryNodeClient](config, "milvus.proto.query.QueryNode"),
+		nodeID:     nodeID,
+	}
+	// node shall specify node id
+	client.grpcClient.SetRole(fmt.Sprintf("%s-%d", typeutil.QueryNodeRole, nodeID))
+	client.grpcClient.SetGetAddrFunc(client.getAddr)
+	client.grpcClient.SetNewGrpcClientFunc(client.newGrpcClient)
+	client.grpcClient.SetNodeID(nodeID)
+	client.grpcClient.SetSession(sess)
+
+	if Params.InternalTLSCfg.InternalTLSEnabled.GetAsBool() {
+		client.grpcClient.EnableEncryption()
+		cp, err := utils.CreateCertPoolforClient(Params.InternalTLSCfg.InternalTLSCaPemPath.GetValue(), "QueryNode")
+		if err != nil {
+			mlog.Error(ctx, "Failed to create cert pool for QueryNode client")
+			return nil, err
+		}
+		client.grpcClient.SetInternalTLSCertPool(cp)
+		client.grpcClient.SetInternalTLSServerName(Params.InternalTLSCfg.InternalTLSSNI.GetValue())
+	}
+	return client, nil
+}
+
+// Close close QueryNode's grpc client
+func (c *Client) Close() error {
+	return c.grpcClient.Close()
+}
+
+func (c *Client) newGrpcClient(cc *grpc.ClientConn) querypb.QueryNodeClient {
+	return querypb.NewQueryNodeClient(cc)
+}
+
+func (c *Client) getAddr() (string, error) {
+	return c.addr, nil
+}
+
+func wrapGrpcCall[T any](ctx context.Context, c *Client, call func(grpcClient querypb.QueryNodeClient) (*T, error)) (*T, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client querypb.QueryNodeClient) (any, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return call(client)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(*T), err
+}
+
+// GetComponentStates gets the component states of QueryNode.
+func (c *Client) GetComponentStates(ctx context.Context, _ *milvuspb.GetComponentStatesRequest, _ ...grpc.CallOption) (*milvuspb.ComponentStates, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*milvuspb.ComponentStates, error) {
+		return client.GetComponentStates(ctx, &milvuspb.GetComponentStatesRequest{})
+	})
+}
+
+// GetTimeTickChannel gets the time tick channel of QueryNode.
+func (c *Client) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest, _ ...grpc.CallOption) (*milvuspb.StringResponse, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*milvuspb.StringResponse, error) {
+		return client.GetTimeTickChannel(ctx, &internalpb.GetTimeTickChannelRequest{})
+	})
+}
+
+// GetStatisticsChannel gets the statistics channel of QueryNode.
+func (c *Client) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest, _ ...grpc.CallOption) (*milvuspb.StringResponse, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*milvuspb.StringResponse, error) {
+		return client.GetStatisticsChannel(ctx, &internalpb.GetStatisticsChannelRequest{})
+	})
+}
+
+// WatchDmChannels watches the channels about data manipulation.
+func (c *Client) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChannelsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.WatchDmChannels(ctx, req)
+	})
+}
+
+// UnsubDmChannel unsubscribes the channels about data manipulation.
+func (c *Client) UnsubDmChannel(ctx context.Context, req *querypb.UnsubDmChannelRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.UnsubDmChannel(ctx, req)
+	})
+}
+
+// LoadSegments loads the segments to search.
+func (c *Client) LoadSegments(ctx context.Context, req *querypb.LoadSegmentsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.LoadSegments(ctx, req)
+	})
+}
+
+// ReleaseCollection releases the data of the specified collection in QueryNode.
+func (c *Client) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.ReleaseCollection(ctx, req)
+	})
+}
+
+// LoadPartitions updates partitions meta info in QueryNode.
+func (c *Client) LoadPartitions(ctx context.Context, req *querypb.LoadPartitionsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.LoadPartitions(ctx, req)
+	})
+}
+
+// ReleasePartitions releases the data of the specified partitions in QueryNode.
+func (c *Client) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.ReleasePartitions(ctx, req)
+	})
+}
+
+// ReleaseSegments releases the data of the specified segments in QueryNode.
+func (c *Client) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.ReleaseSegments(ctx, req)
+	})
+}
+
+// Search performs replica search tasks in QueryNode.
+func (c *Client) Search(ctx context.Context, req *querypb.SearchRequest, _ ...grpc.CallOption) (*internalpb.SearchResults, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.SearchResults, error) {
+		return client.Search(ctx, req)
+	})
+}
+
+func (c *Client) SearchSegments(ctx context.Context, req *querypb.SearchRequest, _ ...grpc.CallOption) (*internalpb.SearchResults, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.SearchResults, error) {
+		return client.SearchSegments(ctx, req)
+	})
+}
+
+// Query performs replica query tasks in QueryNode.
+func (c *Client) Query(ctx context.Context, req *querypb.QueryRequest, _ ...grpc.CallOption) (*internalpb.RetrieveResults, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.RetrieveResults, error) {
+		return client.Query(ctx, req)
+	})
+}
+
+func (c *Client) QueryStream(ctx context.Context, req *querypb.QueryRequest, _ ...grpc.CallOption) (querypb.QueryNode_QueryStreamClient, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client querypb.QueryNodeClient) (any, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+
+		return client.QueryStream(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(querypb.QueryNode_QueryStreamClient), nil
+}
+
+func (c *Client) QuerySegments(ctx context.Context, req *querypb.QueryRequest, _ ...grpc.CallOption) (*internalpb.RetrieveResults, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.RetrieveResults, error) {
+		return client.QuerySegments(ctx, req)
+	})
+}
+
+func (c *Client) QueryStreamSegments(ctx context.Context, req *querypb.QueryRequest, _ ...grpc.CallOption) (querypb.QueryNode_QueryStreamSegmentsClient, error) {
+	ret, err := c.grpcClient.ReCall(ctx, func(client querypb.QueryNodeClient) (any, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+
+		return client.QueryStreamSegments(ctx, req)
+	})
+	if err != nil || ret == nil {
+		return nil, err
+	}
+	return ret.(querypb.QueryNode_QueryStreamSegmentsClient), nil
+}
+
+// GetSegmentInfo gets the information of the specified segments in QueryNode.
+func (c *Client) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest, _ ...grpc.CallOption) (*querypb.GetSegmentInfoResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.GetSegmentInfoResponse, error) {
+		return client.GetSegmentInfo(ctx, req)
+	})
+}
+
+// SyncReplicaSegments syncs replica node segments information to shard leaders.
+func (c *Client) SyncReplicaSegments(ctx context.Context, req *querypb.SyncReplicaSegmentsRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.SyncReplicaSegments(ctx, req)
+	})
+}
+
+// ShowConfigurations gets specified configurations para of QueryNode
+func (c *Client) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest, _ ...grpc.CallOption) (*internalpb.ShowConfigurationsResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.ShowConfigurationsResponse, error) {
+		return client.ShowConfigurations(ctx, req)
+	})
+}
+
+// GetMetrics gets the metrics information of QueryNode.
+func (c *Client) GetMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest, _ ...grpc.CallOption) (*milvuspb.GetMetricsResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*milvuspb.GetMetricsResponse, error) {
+		return client.GetMetrics(ctx, req)
+	})
+}
+
+func (c *Client) GetStatistics(ctx context.Context, request *querypb.GetStatisticsRequest, _ ...grpc.CallOption) (*internalpb.GetStatisticsResponse, error) {
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.GetStatisticsResponse, error) {
+		return client.GetStatistics(ctx, request)
+	})
+}
+
+func (c *Client) GetDataDistribution(ctx context.Context, req *querypb.GetDataDistributionRequest, _ ...grpc.CallOption) (*querypb.GetDataDistributionResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.GetDataDistributionResponse, error) {
+		return client.GetDataDistribution(ctx, req)
+	})
+}
+
+func (c *Client) SyncDistribution(ctx context.Context, req *querypb.SyncDistributionRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID))
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.SyncDistribution(ctx, req)
+	})
+}
+
+// Delete is used to forward delete message between delegator and workers.
+func (c *Client) Delete(ctx context.Context, req *querypb.DeleteRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.Delete(ctx, req)
+	})
+}
+
+// DeleteBatch is the API to apply same delete data into multiple segments.
+// it's basically same as `Delete` but cost less memory pressure.
+func (c *Client) DeleteBatch(ctx context.Context, req *querypb.DeleteBatchRequest, _ ...grpc.CallOption) (*querypb.DeleteBatchResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.DeleteBatchResponse, error) {
+		return client.DeleteBatch(ctx, req)
+	})
+}
+
+func (c *Client) UpdateSchema(ctx context.Context, req *querypb.UpdateSchemaRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.UpdateSchema(ctx, req)
+	})
+}
+
+func (c *Client) RunAnalyzer(ctx context.Context, req *querypb.RunAnalyzerRequest, _ ...grpc.CallOption) (*milvuspb.RunAnalyzerResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*milvuspb.RunAnalyzerResponse, error) {
+		return client.RunAnalyzer(ctx, req)
+	})
+}
+
+func (c *Client) DropIndex(ctx context.Context, req *querypb.DropIndexRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.DropIndex(ctx, req)
+	})
+}
+
+func (c *Client) UpdateIndex(ctx context.Context, req *querypb.UpdateIndexRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.UpdateIndex(ctx, req)
+	})
+}
+
+func (c *Client) ValidateAnalyzer(ctx context.Context, req *querypb.ValidateAnalyzerRequest, _ ...grpc.CallOption) (*querypb.ValidateAnalyzerResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.ValidateAnalyzerResponse, error) {
+		return client.ValidateAnalyzer(ctx, req)
+	})
+}
+
+func (c *Client) GetHighlight(ctx context.Context, req *querypb.GetHighlightRequest, _ ...grpc.CallOption) (*querypb.GetHighlightResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.GetHighlightResponse, error) {
+		return client.GetHighlight(ctx, req)
+	})
+}
+
+func (c *Client) SyncFileResource(ctx context.Context, req *internalpb.SyncFileResourceRequest, _ ...grpc.CallOption) (*commonpb.Status, error) {
+	req = typeutil.Clone(req)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*commonpb.Status, error) {
+		return client.SyncFileResource(ctx, req)
+	})
+}
+
+func (c *Client) ClearReadTaskQueue(ctx context.Context, req *internalpb.ClearReadTaskQueueRequest, _ ...grpc.CallOption) (*internalpb.ClearReadTaskQueueResponse, error) {
+	req = typeutil.Clone(req)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*internalpb.ClearReadTaskQueueResponse, error) {
+		return client.ClearReadTaskQueue(ctx, req)
+	})
+}
+
+func (c *Client) ComputePhraseMatchSlop(ctx context.Context, req *querypb.ComputePhraseMatchSlopRequest, _ ...grpc.CallOption) (*querypb.ComputePhraseMatchSlopResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(c.nodeID),
+	)
+	return wrapGrpcCall(ctx, c, func(client querypb.QueryNodeClient) (*querypb.ComputePhraseMatchSlopResponse, error) {
+		return client.ComputePhraseMatchSlop(ctx, req)
+	})
+}

internal/distributed/querynode/client/client_test.go

@@ -0,0 +1,172 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcquerynodeclient
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc"
+
+	"github.com/milvus-io/milvus/internal/util/mock"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+)
+
+func Test_NewClient(t *testing.T) {
+	paramtable.Init()
+
+	ctx := context.Background()
+	client, err := NewClient(ctx, "", 1)
+	assert.Nil(t, client)
+	assert.Error(t, err)
+
+	client, err = NewClient(ctx, "test", 2)
+	assert.NoError(t, err)
+	assert.NotNil(t, client)
+
+	ctx, cancel := context.WithCancel(ctx)
+
+	checkFunc := func(retNotNil bool) {
+		retCheck := func(notNil bool, ret any, err error) {
+			if notNil {
+				assert.NotNil(t, ret)
+				assert.NoError(t, err)
+			} else {
+				assert.Nil(t, ret)
+				assert.Error(t, err)
+			}
+		}
+
+		r1, err := client.GetComponentStates(ctx, nil)
+		retCheck(retNotNil, r1, err)
+
+		r2, err := client.GetTimeTickChannel(ctx, nil)
+		retCheck(retNotNil, r2, err)
+
+		r3, err := client.GetStatisticsChannel(ctx, nil)
+		retCheck(retNotNil, r3, err)
+
+		r6, err := client.WatchDmChannels(ctx, nil)
+		retCheck(retNotNil, r6, err)
+
+		r7, err := client.LoadSegments(ctx, nil)
+		retCheck(retNotNil, r7, err)
+
+		r8, err := client.ReleaseCollection(ctx, nil)
+		retCheck(retNotNil, r8, err)
+
+		r8, err = client.LoadPartitions(ctx, nil)
+		retCheck(retNotNil, r8, err)
+
+		r9, err := client.ReleasePartitions(ctx, nil)
+		retCheck(retNotNil, r9, err)
+
+		r10, err := client.ReleaseSegments(ctx, nil)
+		retCheck(retNotNil, r10, err)
+
+		r11, err := client.GetSegmentInfo(ctx, nil)
+		retCheck(retNotNil, r11, err)
+
+		r12, err := client.GetMetrics(ctx, nil)
+		retCheck(retNotNil, r12, err)
+
+		r14, err := client.Search(ctx, nil)
+		retCheck(retNotNil, r14, err)
+
+		r15, err := client.Query(ctx, nil)
+		retCheck(retNotNil, r15, err)
+
+		r16, err := client.SyncReplicaSegments(ctx, nil)
+		retCheck(retNotNil, r16, err)
+
+		r17, err := client.GetStatistics(ctx, nil)
+		retCheck(retNotNil, r17, err)
+
+		r18, err := client.ShowConfigurations(ctx, nil)
+		retCheck(retNotNil, r18, err)
+
+		r19, err := client.QuerySegments(ctx, nil)
+		retCheck(retNotNil, r19, err)
+
+		r20, err := client.SearchSegments(ctx, nil)
+		retCheck(retNotNil, r20, err)
+
+		r21, err := client.DeleteBatch(ctx, nil)
+		retCheck(retNotNil, r21, err)
+
+		r22, err := client.RunAnalyzer(ctx, nil)
+		retCheck(retNotNil, r22, err)
+
+		r23, err := client.ValidateAnalyzer(ctx, nil)
+		retCheck(retNotNil, r23, err)
+
+		r24, err := client.GetHighlight(ctx, nil)
+		retCheck(retNotNil, r24, err)
+
+		// stream rpc
+		client, err := client.QueryStream(ctx, nil)
+		retCheck(retNotNil, client, err)
+	}
+
+	client.(*Client).grpcClient = &mock.GRPCClientBase[querypb.QueryNodeClient]{
+		GetGrpcClientErr: errors.New("dummy"),
+	}
+
+	newFunc1 := func(cc *grpc.ClientConn) querypb.QueryNodeClient {
+		return &mock.GrpcQueryNodeClient{Err: nil}
+	}
+	client.(*Client).grpcClient.SetNewGrpcClientFunc(newFunc1)
+
+	checkFunc(false)
+
+	client.(*Client).grpcClient = &mock.GRPCClientBase[querypb.QueryNodeClient]{
+		GetGrpcClientErr: nil,
+	}
+
+	newFunc2 := func(cc *grpc.ClientConn) querypb.QueryNodeClient {
+		return &mock.GrpcQueryNodeClient{Err: errors.New("dummy")}
+	}
+
+	client.(*Client).grpcClient.SetNewGrpcClientFunc(newFunc2)
+
+	checkFunc(false)
+
+	client.(*Client).grpcClient = &mock.GRPCClientBase[querypb.QueryNodeClient]{
+		GetGrpcClientErr: nil,
+	}
+
+	newFunc3 := func(cc *grpc.ClientConn) querypb.QueryNodeClient {
+		return &mock.GrpcQueryNodeClient{Err: nil}
+	}
+	client.(*Client).grpcClient.SetNewGrpcClientFunc(newFunc3)
+
+	checkFunc(true)
+
+	// ctx canceled
+	client.(*Client).grpcClient = &mock.GRPCClientBase[querypb.QueryNodeClient]{
+		GetGrpcClientErr: nil,
+	}
+	client.(*Client).grpcClient.SetNewGrpcClientFunc(newFunc1)
+	cancel() // make context canceled
+	checkFunc(false)
+
+	err = client.Close()
+	assert.NoError(t, err)
+}

internal/distributed/querynode/service.go

@@ -0,0 +1,489 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcquerynode
+
+import (
+	"context"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/cockroachdb/errors"
+	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
+	clientv3 "go.etcd.io/etcd/client/v3"
+	"go.uber.org/atomic"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
+	mix "github.com/milvus-io/milvus/internal/distributed/mixcoord/client"
+	"github.com/milvus-io/milvus/internal/distributed/utils"
+	"github.com/milvus-io/milvus/internal/flushcommon/broker"
+	qn "github.com/milvus-io/milvus/internal/querynodev2"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/internal/util/componentutil"
+	"github.com/milvus-io/milvus/internal/util/dependency"
+	_ "github.com/milvus-io/milvus/internal/util/grpcclient"
+	"github.com/milvus-io/milvus/pkg/v3/mlog"
+	"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/tracer"
+	"github.com/milvus-io/milvus/pkg/v3/util/etcd"
+	"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/interceptor"
+	"github.com/milvus-io/milvus/pkg/v3/util/netutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+	"github.com/milvus-io/milvus/pkg/v3/util/retry"
+	"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
+	"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
+)
+
+// UniqueID is an alias for type typeutil.UniqueID, used as a unique identifier for the request.
+type UniqueID = typeutil.UniqueID
+
+// Server is the grpc server of QueryNode.
+type Server struct {
+	querynode   types.QueryNodeComponent
+	grpcWG      sync.WaitGroup
+	ctx         context.Context
+	cancel      context.CancelFunc
+	grpcErrChan chan error
+
+	serverID atomic.Int64
+
+	grpcServer *grpc.Server
+	listener   *netutil.NetListener
+
+	etcdCli  *clientv3.Client
+	mixCoord *syncutil.Future[types.MixCoordClient]
+}
+
+// lazyBinlogSaver implements segments.BinlogSaver with lazy initialization.
+// It waits for MixCoordClient to be ready on first SaveBinlogPaths call.
+type lazyBinlogSaver struct {
+	mixCoordFuture *syncutil.Future[types.MixCoordClient]
+	mu             sync.Mutex
+	broker         broker.Broker
+}
+
+func (s *lazyBinlogSaver) SaveBinlogPaths(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error {
+	s.mu.Lock()
+	if s.broker == nil {
+		client, err := s.mixCoordFuture.GetWithContext(ctx)
+		if err != nil {
+			s.mu.Unlock()
+			return errors.Wrap(err, "failed to get MixCoordClient for SaveBinlogPaths")
+		}
+		s.broker = broker.NewCoordBroker(client, paramtable.GetNodeID())
+	}
+	b := s.broker
+	s.mu.Unlock()
+	return b.SaveBinlogPaths(ctx, req)
+}
+
+func (s *Server) GetStatistics(ctx context.Context, request *querypb.GetStatisticsRequest) (*internalpb.GetStatisticsResponse, error) {
+	return s.querynode.GetStatistics(ctx, request)
+}
+
+func (s *Server) GetQueryNode() types.QueryNodeComponent {
+	return s.querynode
+}
+
+// NewServer create a new QueryNode grpc server.
+func NewServer(ctx context.Context, factory dependency.Factory) (*Server, error) {
+	ctx1, cancel := context.WithCancel(ctx) //nolint:gosec
+
+	s := &Server{
+		ctx:         ctx1,
+		cancel:      cancel,
+		querynode:   qn.NewQueryNode(ctx, factory),
+		grpcErrChan: make(chan error),
+		mixCoord:    syncutil.NewFuture[types.MixCoordClient](),
+	}
+	return s, nil
+}
+
+func (s *Server) Prepare() error {
+	listener, err := netutil.NewListener(
+		netutil.OptIP(paramtable.Get().QueryNodeGrpcServerCfg.IP),
+		netutil.OptHighPriorityToUsePort(paramtable.Get().QueryNodeGrpcServerCfg.Port.GetAsInt()),
+	)
+	if err != nil {
+		mlog.Warn(s.ctx, "QueryNode fail to create net listener", mlog.Err(err))
+		return err
+	}
+	s.listener = listener
+	mlog.Info(s.ctx, "QueryNode listen on", mlog.String("address", listener.Addr().String()), mlog.Int("port", listener.Port()))
+	paramtable.Get().Save(
+		paramtable.Get().QueryNodeGrpcServerCfg.Port.Key,
+		strconv.FormatInt(int64(listener.Port()), 10))
+	return nil
+}
+
+// init initializes QueryNode's grpc service.
+func (s *Server) init() error {
+	etcdConfig := &paramtable.Get().EtcdCfg
+
+	mlog.Debug(s.ctx, "QueryNode", mlog.Int("port", s.listener.Port()))
+
+	etcdCli, err := etcd.CreateEtcdClient(
+		etcdConfig.UseEmbedEtcd.GetAsBool(),
+		etcdConfig.EtcdEnableAuth.GetAsBool(),
+		etcdConfig.EtcdAuthUserName.GetValue(),
+		etcdConfig.EtcdAuthPassword.GetValue(),
+		etcdConfig.EtcdUseSSL.GetAsBool(),
+		etcdConfig.Endpoints.GetAsStrings(),
+		etcdConfig.EtcdTLSCert.GetValue(),
+		etcdConfig.EtcdTLSKey.GetValue(),
+		etcdConfig.EtcdTLSCACert.GetValue(),
+		etcdConfig.EtcdTLSMinVersion.GetValue(),
+		etcdConfig.ClientOptions()...)
+	if err != nil {
+		mlog.Debug(s.ctx, "QueryNode connect to etcd failed", mlog.Err(err))
+		return err
+	}
+	s.etcdCli = etcdCli
+	s.SetEtcdClient(etcdCli)
+	s.querynode.SetAddress(s.listener.Address())
+	mlog.Debug(s.ctx, "QueryNode connect to etcd successfully")
+	s.grpcWG.Add(1)
+	go s.startGrpcLoop()
+	// wait for grpc server loop start
+	err = <-s.grpcErrChan
+	if err != nil {
+		return err
+	}
+
+	s.querynode.UpdateStateCode(commonpb.StateCode_Initializing)
+	mlog.Debug(s.ctx, "QueryNode", mlog.Any("State", commonpb.StateCode_Initializing))
+	if err := s.querynode.Init(); err != nil {
+		mlog.Error(s.ctx, "QueryNode init error: ", mlog.Err(err))
+		return err
+	}
+	s.serverID.Store(s.querynode.GetNodeID())
+
+	// initialize MixCoord client asynchronously and set BinlogSaver on QueryNode
+	s.initMixCoord()
+	if qnImpl, ok := s.querynode.(*qn.QueryNode); ok {
+		qnImpl.SetBinlogSaver(&lazyBinlogSaver{mixCoordFuture: s.mixCoord})
+	}
+
+	return nil
+}
+
+func (s *Server) initMixCoord() {
+	go func() {
+		retry.Do(s.ctx, func() error {
+			mlog.Info(s.ctx, "QueryNode connect to mixCoord...")
+			mixCoord, err := mix.NewClient(s.ctx)
+			if err != nil {
+				return errors.Wrap(err, "QueryNode try to new mixCoord client failed")
+			}
+
+			mlog.Info(s.ctx, "QueryNode try to wait for mixCoord ready")
+			err = componentutil.WaitForComponentHealthy(s.ctx, mixCoord, "mixCoord", 1000000, time.Millisecond*200)
+			if err != nil {
+				return errors.Wrap(err, "QueryNode wait for mixCoord ready failed")
+			}
+			mlog.Info(s.ctx, "QueryNode mixCoord client ready")
+			s.mixCoord.Set(mixCoord)
+			return nil
+		}, retry.AttemptAlways())
+	}()
+}
+
+// start starts QueryNode's grpc service.
+func (s *Server) start() error {
+	if err := s.querynode.Start(); err != nil {
+		mlog.Error(s.ctx, "QueryNode start failed", mlog.Err(err))
+		return err
+	}
+	if err := s.querynode.Register(); err != nil {
+		mlog.Error(s.ctx, "QueryNode register service failed", mlog.Err(err))
+		return err
+	}
+	return nil
+}
+
+// startGrpcLoop starts the grpc loop of QueryNode component.
+func (s *Server) startGrpcLoop() {
+	defer s.grpcWG.Done()
+	Params := &paramtable.Get().QueryNodeGrpcServerCfg
+	kaep := keepalive.EnforcementPolicy{
+		MinTime:             5 * time.Second, // If a client pings more than once every 5 seconds, terminate the connection
+		PermitWithoutStream: true,            // Allow pings even when there are no active streams
+	}
+
+	kasp := keepalive.ServerParameters{
+		Time:    60 * time.Second, // Ping the client if it is idle for 60 seconds to ensure the connection is still active
+		Timeout: 10 * time.Second, // Wait 10 second for the ping ack before assuming the connection is dead
+	}
+
+	grpcOpts := []grpc.ServerOption{
+		grpc.KeepaliveEnforcementPolicy(kaep),
+		grpc.KeepaliveParams(kasp),
+		grpc.MaxRecvMsgSize(Params.ServerMaxRecvSize.GetAsInt()),
+		grpc.MaxSendMsgSize(Params.ServerMaxSendSize.GetAsInt()),
+		grpc.UnaryInterceptor(grpc_middleware.ChainUnaryServer(
+			// otelgrpc.UnaryServerInterceptor(opts...),
+			mlog.UnaryServerInterceptor(typeutil.QueryNodeRole),
+			interceptor.ClusterValidationUnaryServerInterceptor(),
+			interceptor.ServerIDValidationUnaryServerInterceptor(func() int64 {
+				if s.serverID.Load() == 0 {
+					s.serverID.Store(paramtable.GetNodeID())
+				}
+				return s.serverID.Load()
+			}),
+		)),
+		grpc.StreamInterceptor(grpc_middleware.ChainStreamServer(
+			// otelgrpc.StreamServerInterceptor(opts...),
+			mlog.StreamServerInterceptor(typeutil.QueryNodeRole),
+			interceptor.ClusterValidationStreamServerInterceptor(),
+			interceptor.ServerIDValidationStreamServerInterceptor(func() int64 {
+				if s.serverID.Load() == 0 {
+					s.serverID.Store(paramtable.GetNodeID())
+				}
+				return s.serverID.Load()
+			}),
+		)),
+		grpc.StatsHandler(tracer.GetDynamicOtelGrpcServerStatsHandler()),
+	}
+
+	grpcOpts = append(grpcOpts, utils.EnableInternalTLS("QueryNode"))
+	s.grpcServer = grpc.NewServer(grpcOpts...)
+	querypb.RegisterQueryNodeServer(s.grpcServer, s)
+
+	ctx, cancel := context.WithCancel(s.ctx)
+	defer cancel()
+
+	go funcutil.CheckGrpcReady(ctx, s.grpcErrChan)
+	if err := s.grpcServer.Serve(s.listener); err != nil {
+		mlog.Debug(s.ctx, "QueryNode Start Grpc Failed!!!!")
+		s.grpcErrChan <- err
+	}
+}
+
+// Run initializes and starts QueryNode's grpc service.
+func (s *Server) Run() error {
+	if err := s.init(); err != nil {
+		return err
+	}
+	mlog.Debug(s.ctx, "QueryNode init done ...")
+
+	if err := s.start(); err != nil {
+		return err
+	}
+	mlog.Debug(s.ctx, "QueryNode start done ...")
+	return nil
+}
+
+// Stop stops QueryNode's grpc service.
+func (s *Server) Stop() (err error) {
+	logger := mlog.With()
+	if s.listener != nil {
+		logger = logger.With(mlog.String("address", s.listener.Address()))
+	}
+	logger.Info(s.ctx, "QueryNode stopping")
+	defer func() {
+		logger.Info(s.ctx, "QueryNode stopped", mlog.Err(err))
+	}()
+
+	logger.Info(s.ctx, "internal server[querynode] start to stop")
+	err = s.querynode.Stop()
+	if err != nil {
+		logger.Error(s.ctx, "failed to close querynode", mlog.Err(err))
+		return err
+	}
+	if s.etcdCli != nil {
+		defer s.etcdCli.Close()
+	}
+
+	if s.grpcServer != nil {
+		utils.GracefulStopGRPCServer(s.grpcServer)
+	}
+	s.grpcWG.Wait()
+
+	s.cancel()
+	if s.listener != nil {
+		s.listener.Close()
+	}
+	return nil
+}
+
+// SetEtcdClient sets the etcd client for QueryNode component.
+func (s *Server) SetEtcdClient(etcdCli *clientv3.Client) {
+	s.querynode.SetEtcdClient(etcdCli)
+}
+
+// GetTimeTickChannel gets the time tick channel of QueryNode.
+func (s *Server) GetTimeTickChannel(ctx context.Context, req *internalpb.GetTimeTickChannelRequest) (*milvuspb.StringResponse, error) {
+	return s.querynode.GetTimeTickChannel(ctx, req)
+}
+
+// GetStatisticsChannel gets the statistics channel of QueryNode.
+func (s *Server) GetStatisticsChannel(ctx context.Context, req *internalpb.GetStatisticsChannelRequest) (*milvuspb.StringResponse, error) {
+	return s.querynode.GetStatisticsChannel(ctx, req)
+}
+
+// GetComponentStates gets the component states of QueryNode.
+func (s *Server) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
+	// ignore ctx and in
+	return s.querynode.GetComponentStates(ctx, req)
+}
+
+// WatchDmChannels watches the channels about data manipulation.
+func (s *Server) WatchDmChannels(ctx context.Context, req *querypb.WatchDmChannelsRequest) (*commonpb.Status, error) {
+	// ignore ctx
+	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
+	return s.querynode.LoadSegments(ctx, req)
+}
+
+// ReleaseCollection releases the data of the specified collection in QueryNode.
+func (s *Server) ReleaseCollection(ctx context.Context, req *querypb.ReleaseCollectionRequest) (*commonpb.Status, error) {
+	// ignore ctx
+	return s.querynode.ReleaseCollection(ctx, req)
+}
+
+// LoadPartitions updates partitions meta info in QueryNode.
+func (s *Server) LoadPartitions(ctx context.Context, req *querypb.LoadPartitionsRequest) (*commonpb.Status, error) {
+	return s.querynode.LoadPartitions(ctx, req)
+}
+
+// ReleasePartitions releases the data of the specified partitions in QueryNode.
+func (s *Server) ReleasePartitions(ctx context.Context, req *querypb.ReleasePartitionsRequest) (*commonpb.Status, error) {
+	// ignore ctx
+	return s.querynode.ReleasePartitions(ctx, req)
+}
+
+// ReleaseSegments releases the data of the specified segments in QueryNode.
+func (s *Server) ReleaseSegments(ctx context.Context, req *querypb.ReleaseSegmentsRequest) (*commonpb.Status, error) {
+	// ignore ctx
+	return s.querynode.ReleaseSegments(ctx, req)
+}
+
+// GetSegmentInfo gets the information of the specified segments in QueryNode.
+func (s *Server) GetSegmentInfo(ctx context.Context, req *querypb.GetSegmentInfoRequest) (*querypb.GetSegmentInfoResponse, error) {
+	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) (*internalpb.SearchResults, error) {
+	return s.querynode.Search(ctx, req)
+}
+
+func (s *Server) SearchSegments(ctx context.Context, req *querypb.SearchRequest) (*internalpb.SearchResults, error) {
+	return s.querynode.SearchSegments(ctx, req)
+}
+
+// Query performs query of streaming/historical replica on QueryNode.
+func (s *Server) Query(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
+	return s.querynode.Query(ctx, req)
+}
+
+func (s *Server) QueryStream(req *querypb.QueryRequest, srv querypb.QueryNode_QueryStreamServer) error {
+	return s.querynode.QueryStream(req, srv)
+}
+
+func (s *Server) QueryStreamSegments(req *querypb.QueryRequest, srv querypb.QueryNode_QueryStreamSegmentsServer) error {
+	return s.querynode.QueryStreamSegments(req, srv)
+}
+
+func (s *Server) QuerySegments(ctx context.Context, req *querypb.QueryRequest) (*internalpb.RetrieveResults, error) {
+	return s.querynode.QuerySegments(ctx, req)
+}
+
+// SyncReplicaSegments syncs replica segment information to shard leader
+func (s *Server) SyncReplicaSegments(ctx context.Context, req *querypb.SyncReplicaSegmentsRequest) (*commonpb.Status, error) {
+	return s.querynode.SyncReplicaSegments(ctx, req)
+}
+
+// ShowConfigurations gets specified configurations para of QueryNode
+func (s *Server) ShowConfigurations(ctx context.Context, req *internalpb.ShowConfigurationsRequest) (*internalpb.ShowConfigurationsResponse, error) {
+	return s.querynode.ShowConfigurations(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)
+}
+
+// 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)
+}
+
+// Delete is used to forward delete message between delegator and workers.
+func (s *Server) Delete(ctx context.Context, req *querypb.DeleteRequest) (*commonpb.Status, error) {
+	return s.querynode.Delete(ctx, req)
+}
+
+// DeleteBatch is the API to apply same delete data into multiple segments.
+// it's basically same as `Delete` but cost less memory pressure.
+func (s *Server) DeleteBatch(ctx context.Context, req *querypb.DeleteBatchRequest) (*querypb.DeleteBatchResponse, error) {
+	return s.querynode.DeleteBatch(ctx, req)
+}
+
+func (s *Server) UpdateSchema(ctx context.Context, req *querypb.UpdateSchemaRequest) (*commonpb.Status, error) {
+	return s.querynode.UpdateSchema(ctx, req)
+}
+
+func (s *Server) RunAnalyzer(ctx context.Context, req *querypb.RunAnalyzerRequest) (*milvuspb.RunAnalyzerResponse, error) {
+	return s.querynode.RunAnalyzer(ctx, req)
+}
+
+func (s *Server) DropIndex(ctx context.Context, req *querypb.DropIndexRequest) (*commonpb.Status, error) {
+	return s.querynode.DropIndex(ctx, req)
+}
+
+func (s *Server) UpdateIndex(ctx context.Context, req *querypb.UpdateIndexRequest) (*commonpb.Status, error) {
+	return s.querynode.UpdateIndex(ctx, req)
+}
+
+func (s *Server) ValidateAnalyzer(ctx context.Context, req *querypb.ValidateAnalyzerRequest) (*querypb.ValidateAnalyzerResponse, error) {
+	return s.querynode.ValidateAnalyzer(ctx, req)
+}
+
+func (s *Server) GetHighlight(ctx context.Context, req *querypb.GetHighlightRequest) (*querypb.GetHighlightResponse, error) {
+	return s.querynode.GetHighlight(ctx, req)
+}
+
+func (s *Server) SyncFileResource(ctx context.Context, req *internalpb.SyncFileResourceRequest) (*commonpb.Status, error) {
+	return s.querynode.SyncFileResource(ctx, req)
+}
+
+func (s *Server) ClearReadTaskQueue(ctx context.Context, req *internalpb.ClearReadTaskQueueRequest) (*internalpb.ClearReadTaskQueueResponse, error) {
+	return s.querynode.ClearReadTaskQueue(ctx, req)
+}
+
+func (s *Server) ComputePhraseMatchSlop(ctx context.Context, req *querypb.ComputePhraseMatchSlopRequest) (*querypb.ComputePhraseMatchSlopResponse, error) {
+	return s.querynode.ComputePhraseMatchSlop(ctx, req)
+}

internal/distributed/querynode/service_test.go

@@ -0,0 +1,343 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package grpcquerynode
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/cockroachdb/errors"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	clientv3 "go.etcd.io/etcd/client/v3"
+
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
+	"github.com/milvus-io/milvus/internal/mocks"
+	"github.com/milvus-io/milvus/internal/types"
+	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
+	"github.com/milvus-io/milvus/pkg/v3/util/merr"
+	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
+)
+
+type MockRootCoord struct {
+	types.RootCoord
+	initErr  error
+	startErr error
+	regErr   error
+	stopErr  error
+	stateErr commonpb.ErrorCode
+}
+
+func (m *MockRootCoord) Init() error {
+	return m.initErr
+}
+
+func (m *MockRootCoord) Start() error {
+	return m.startErr
+}
+
+func (m *MockRootCoord) Stop() error {
+	return m.stopErr
+}
+
+func (m *MockRootCoord) Register() error {
+	return m.regErr
+}
+
+func (m *MockRootCoord) SetEtcdClient(client *clientv3.Client) {
+}
+
+func (m *MockRootCoord) GetComponentStates(ctx context.Context, req *milvuspb.GetComponentStatesRequest) (*milvuspb.ComponentStates, error) {
+	return &milvuspb.ComponentStates{
+		State:  &milvuspb.ComponentInfo{StateCode: commonpb.StateCode_Healthy},
+		Status: &commonpb.Status{ErrorCode: m.stateErr},
+	}, nil
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+func TestMain(m *testing.M) {
+	paramtable.Init()
+	os.Exit(m.Run())
+}
+
+func Test_NewServer(t *testing.T) {
+	ctx := context.Background()
+	server, err := NewServer(ctx, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, server)
+
+	mockQN := mocks.NewMockQueryNode(t)
+	mockQN.EXPECT().Start().Return(nil).Maybe()
+	mockQN.EXPECT().Stop().Return(nil).Maybe()
+	mockQN.EXPECT().Register().Return(nil).Maybe()
+	mockQN.EXPECT().SetEtcdClient(mock.Anything).Maybe()
+	mockQN.EXPECT().SetAddress(mock.Anything).Maybe()
+	mockQN.EXPECT().UpdateStateCode(mock.Anything).Maybe()
+	mockQN.EXPECT().Init().Return(nil).Maybe()
+	mockQN.EXPECT().GetNodeID().Return(2).Maybe()
+	server.querynode = mockQN
+
+	t.Run("Run", func(t *testing.T) {
+		err = server.Prepare()
+		assert.NoError(t, err)
+		err = server.Run()
+		assert.NoError(t, err)
+	})
+
+	t.Run("GetComponentStates", func(t *testing.T) {
+		mockQN.EXPECT().GetComponentStates(mock.Anything, mock.Anything).Return(&milvuspb.ComponentStates{
+			State: &milvuspb.ComponentInfo{
+				StateCode: commonpb.StateCode_Healthy,
+			},
+		}, nil)
+		req := &milvuspb.GetComponentStatesRequest{}
+		states, err := server.GetComponentStates(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.StateCode_Healthy, states.State.StateCode)
+	})
+
+	t.Run("GetStatisticsChannel", func(t *testing.T) {
+		mockQN.EXPECT().GetStatisticsChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}}, nil)
+		req := &internalpb.GetStatisticsChannelRequest{}
+		resp, err := server.GetStatisticsChannel(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("GetTimeTickChannel", func(t *testing.T) {
+		mockQN.EXPECT().GetTimeTickChannel(mock.Anything, mock.Anything).Return(&milvuspb.StringResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}}, nil)
+		req := &internalpb.GetTimeTickChannelRequest{}
+		resp, err := server.GetTimeTickChannel(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("WatchDmChannels", func(t *testing.T) {
+		mockQN.EXPECT().WatchDmChannels(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.WatchDmChannelsRequest{}
+		resp, err := server.WatchDmChannels(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("LoadSegments", func(t *testing.T) {
+		mockQN.EXPECT().LoadSegments(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.LoadSegmentsRequest{}
+		resp, err := server.LoadSegments(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("ReleaseCollection", func(t *testing.T) {
+		mockQN.EXPECT().ReleaseCollection(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.ReleaseCollectionRequest{}
+		resp, err := server.ReleaseCollection(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("LoadPartitions", func(t *testing.T) {
+		mockQN.EXPECT().LoadPartitions(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.LoadPartitionsRequest{}
+		resp, err := server.LoadPartitions(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("ReleasePartitions", func(t *testing.T) {
+		mockQN.EXPECT().ReleasePartitions(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.ReleasePartitionsRequest{}
+		resp, err := server.ReleasePartitions(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("ReleaseSegments", func(t *testing.T) {
+		mockQN.EXPECT().ReleaseSegments(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.ReleaseSegmentsRequest{}
+		resp, err := server.ReleaseSegments(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.ErrorCode)
+	})
+
+	t.Run("GetSegmentInfo", func(t *testing.T) {
+		mockQN.EXPECT().GetSegmentInfo(mock.Anything, mock.Anything).Return(&querypb.GetSegmentInfoResponse{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.GetSegmentInfoRequest{}
+		resp, err := server.GetSegmentInfo(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("GetMetrics", func(t *testing.T) {
+		mockQN.EXPECT().GetMetrics(mock.Anything, mock.Anything).Return(
+			&milvuspb.GetMetricsResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}}, nil)
+		req := &milvuspb.GetMetricsRequest{
+			Request: "",
+		}
+		resp, err := server.GetMetrics(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("Search", func(t *testing.T) {
+		mockQN.EXPECT().Search(mock.Anything, mock.Anything).Return(&internalpb.SearchResults{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.SearchRequest{}
+		resp, err := server.Search(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("SearchSegments", func(t *testing.T) {
+		mockQN.EXPECT().SearchSegments(mock.Anything, mock.Anything).Return(&internalpb.SearchResults{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.SearchRequest{}
+		resp, err := server.SearchSegments(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("Query", func(t *testing.T) {
+		mockQN.EXPECT().Query(mock.Anything, mock.Anything).Return(&internalpb.RetrieveResults{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.QueryRequest{}
+		resp, err := server.Query(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("QueryStream", func(t *testing.T) {
+		mockQN.EXPECT().QueryStream(mock.Anything, mock.Anything).Return(nil)
+		ret := server.QueryStream(nil, nil)
+		assert.Nil(t, ret)
+	})
+
+	t.Run("QuerySegments", func(t *testing.T) {
+		mockQN.EXPECT().QuerySegments(mock.Anything, mock.Anything).Return(&internalpb.RetrieveResults{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.QueryRequest{}
+		resp, err := server.QuerySegments(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("QueryStreamSegments", func(t *testing.T) {
+		mockQN.EXPECT().QueryStreamSegments(mock.Anything, mock.Anything).Return(nil)
+		ret := server.QueryStreamSegments(nil, nil)
+		assert.Nil(t, ret)
+	})
+
+	t.Run("SyncReplicaSegments", func(t *testing.T) {
+		mockQN.EXPECT().SyncReplicaSegments(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		req := &querypb.SyncReplicaSegmentsRequest{}
+		resp, err := server.SyncReplicaSegments(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetErrorCode())
+	})
+
+	t.Run("ShowConfigurtaions", func(t *testing.T) {
+		mockQN.EXPECT().ShowConfigurations(mock.Anything, mock.Anything).Return(&internalpb.ShowConfigurationsResponse{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &internalpb.ShowConfigurationsRequest{
+			Pattern: "Cache",
+		}
+		resp, err := server.ShowConfigurations(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("DeleteBatch", func(t *testing.T) {
+		mockQN.EXPECT().DeleteBatch(mock.Anything, mock.Anything).Return(&querypb.DeleteBatchResponse{
+			Status: merr.Success(),
+		}, nil)
+
+		resp, err := server.DeleteBatch(ctx, &querypb.DeleteBatchRequest{})
+		assert.NoError(t, merr.CheckRPCCall(resp, err))
+	})
+
+	t.Run("RunAnalyzer", func(t *testing.T) {
+		mockQN.EXPECT().RunAnalyzer(mock.Anything, mock.Anything).Return(&milvuspb.RunAnalyzerResponse{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
+		}, nil)
+		req := &querypb.RunAnalyzerRequest{}
+		resp, err := server.RunAnalyzer(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("GetHighlight", func(t *testing.T) {
+		mockQN.EXPECT().GetHighlight(mock.Anything, mock.Anything).Return(&querypb.GetHighlightResponse{
+			Status: merr.Success(),
+		}, nil)
+
+		resp, err := server.GetHighlight(ctx, &querypb.GetHighlightRequest{
+			Channel: "test-channel",
+		})
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	t.Run("ValidateAnalyzer", func(t *testing.T) {
+		mockQN.EXPECT().ValidateAnalyzer(mock.Anything, mock.Anything).Return(&querypb.ValidateAnalyzerResponse{Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}}, nil)
+		req := &querypb.ValidateAnalyzerRequest{}
+		resp, err := server.ValidateAnalyzer(ctx, req)
+		assert.NoError(t, err)
+		assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
+	})
+
+	err = server.Stop()
+	assert.NoError(t, err)
+}
+
+func Test_Run(t *testing.T) {
+	ctx := context.Background()
+	server, err := NewServer(ctx, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, server)
+
+	mockQN := mocks.NewMockQueryNode(t)
+	mockQN.EXPECT().Start().Return(errors.New("Failed")).Maybe()
+	mockQN.EXPECT().Stop().Return(errors.New("Failed")).Maybe()
+	mockQN.EXPECT().Register().Return(errors.New("Failed")).Maybe()
+	mockQN.EXPECT().SetEtcdClient(mock.Anything).Maybe()
+	mockQN.EXPECT().SetAddress(mock.Anything).Maybe()
+	mockQN.EXPECT().UpdateStateCode(mock.Anything).Maybe()
+	mockQN.EXPECT().Init().Return(nil).Maybe()
+	mockQN.EXPECT().GetNodeID().Return(2).Maybe()
+	server.querynode = mockQN
+	err = server.Prepare()
+	assert.NoError(t, err)
+	err = server.Run()
+	assert.Error(t, err)
+
+	err = server.Run()
+	assert.Error(t, err)
+
+	err = server.Stop()
+	assert.Error(t, err)
+}