Support to manage connections (milvus-io#24224)

Signed-off-by: longjiquan <jiquan.long@zilliz.com>
yah01 · May 19, 2023 · bd34355 · bd34355
1 parent 473c65c
commit bd34355
Show file tree

Hide file tree

Showing 26 changed files with 4,486 additions and 1,879 deletions.
diff --git a/go.mod b/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/golang/protobuf v1.5.3
 	github.com/klauspost/compress v1.14.4
 	github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
-	github.com/milvus-io/milvus-proto/go-api v0.0.0-20230515081521-d963c95b041f
+	github.com/milvus-io/milvus-proto/go-api v0.0.0-20230517025117-8ba62a3f3a63
 	github.com/milvus-io/milvus/pkg v0.0.0-00010101000000-000000000000
 	github.com/minio/minio-go/v7 v7.0.17
 	github.com/panjf2000/ants/v2 v2.7.2

diff --git a/go.sum b/go.sum
@@ -583,6 +583,8 @@ github.com/milvus-io/milvus-proto/go-api v0.0.0-20230421091228-eaa38c831a61 h1:E
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20230421091228-eaa38c831a61/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20230515081521-d963c95b041f h1:uZzVaSbUtxMdEix9By6z+M/H/XNkXRQJdZQ9HP/wHtc=
 github.com/milvus-io/milvus-proto/go-api v0.0.0-20230515081521-d963c95b041f/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230517025117-8ba62a3f3a63 h1:V962mSHjOFUbuMgAXziBdbYPOCVZmN1MkqEeKpME+MA=
+github.com/milvus-io/milvus-proto/go-api v0.0.0-20230517025117-8ba62a3f3a63/go.mod h1:148qnlmZ0Fdm1Fq+Mj/OW2uDoEP25g3mjh0vMGtkgmk=
 github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
 github.com/milvus-io/pulsar-client-go v0.6.10/go.mod h1:lQqCkgwDF8YFYjKA+zOheTk1tev2B+bKj5j7+nm8M1w=
 github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=

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/proxy/client/client.go b/internal/distributed/proxy/client/client.go
@@ -236,3 +236,21 @@ func (c *Client) SetRates(ctx context.Context, req *proxypb.SetRatesRequest) (*c
 	}
 	return ret.(*commonpb.Status), err
 }
+
+func (c *Client) ListClientInfos(ctx context.Context, req *proxypb.ListClientInfosRequest) (*proxypb.ListClientInfosResponse, error) {
+	req = typeutil.Clone(req)
+	commonpbutil.UpdateMsgBase(
+		req.GetBase(),
+		commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.grpcClient.GetNodeID())),
+	)
+	ret, err := c.grpcClient.ReCall(ctx, func(client proxypb.ProxyClient) (any, error) {
+		if !funcutil.CheckCtxValid(ctx) {
+			return nil, ctx.Err()
+		}
+		return client.ListClientInfos(ctx, req)
+	})
+	if err != nil {
+		return nil, err
+	}
+	return ret.(*proxypb.ListClientInfosResponse), nil
+}
diff --git a/internal/distributed/proxy/service.go b/internal/distributed/proxy/service.go
@@ -176,6 +176,7 @@ func (s *Server) startExternalGrpc(grpcPort int, errChan chan error) {
 			logutil.UnaryTraceLoggerInterceptor,
 			proxy.RateLimitInterceptor(limiter),
 			accesslog.UnaryAccessLoggerInterceptor,
+			proxy.KeepActiveInterceptor,
 		)),
 	}
 
@@ -923,3 +924,11 @@ func (s *Server) ListIndexedSegment(ctx context.Context, req *federpb.ListIndexe
 func (s *Server) DescribeSegmentIndexData(ctx context.Context, req *federpb.DescribeSegmentIndexDataRequest) (*federpb.DescribeSegmentIndexDataResponse, error) {
 	panic("TODO: implement me")
 }
+
+func (s *Server) Connect(ctx context.Context, req *milvuspb.ConnectRequest) (*milvuspb.ConnectResponse, error) {
+	return s.proxy.Connect(ctx, req)
+}
+
+func (s *Server) ListClientInfos(ctx context.Context, req *proxypb.ListClientInfosRequest) (*proxypb.ListClientInfosResponse, error) {
+	return s.proxy.ListClientInfos(ctx, req)
+}
diff --git a/internal/distributed/proxy/service_test.go b/internal/distributed/proxy/service_test.go
@@ -514,6 +514,10 @@ func (m *MockProxy) Register() error {
 	return m.regErr
 }
 
+func (m *MockProxy) ListClientInfos(ctx context.Context, request *proxypb.ListClientInfosRequest) (*proxypb.ListClientInfosResponse, error) {
+	return nil, nil
+}
+
 func (m *MockProxy) InvalidateCollectionMetaCache(ctx context.Context, request *proxypb.InvalidateCollMetaCacheRequest) (*commonpb.Status, error) {
 	return nil, nil
 }
@@ -852,6 +856,10 @@ func (m *MockProxy) RenameCollection(ctx context.Context, req *milvuspb.RenameCo
 	return nil, nil
 }
 
+func (m *MockProxy) Connect(ctx context.Context, req *milvuspb.ConnectRequest) (*milvuspb.ConnectResponse, error) {
+	return nil, nil
+}
+
 ///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 type WaitOption struct {

diff --git a/internal/mocks/mock_proxy.go b/internal/mocks/mock_proxy.go
diff --git a/internal/proto/proxy.proto b/internal/proto/proxy.proto
@@ -20,6 +20,8 @@ service Proxy {
   rpc RefreshPolicyInfoCache(RefreshPolicyInfoCacheRequest) returns (common.Status) {}
   rpc GetProxyMetrics(milvus.GetMetricsRequest) returns (milvus.GetMetricsResponse) {}
   rpc SetRates(SetRatesRequest) returns (common.Status) {}
+
+  rpc ListClientInfos(ListClientInfosRequest) returns (ListClientInfosResponse) {}
 }
 
 message InvalidateCollMetaCacheRequest {
@@ -61,3 +63,12 @@ message SetRatesRequest {
   common.MsgBase base = 1;
   repeated CollectionRate rates = 2;
 }
+
+message ListClientInfosRequest {
+  common.MsgBase base = 1;
+}
+
+message ListClientInfosResponse {
+  common.Status status = 1;
+  repeated common.ClientInfo client_infos = 2;
+}