Add an optional parameter naming 'client_tag' on client constructor in c++ sdk (#4377)

* Add an optional parameter naming  on client constructor in c++ sdk

Signed-off-by: fishpenguin <kun.yu@zilliz.com>

* Fix wrong indent

Signed-off-by: fishpenguin <kun.yu@zilliz.com>

CHANGELOG.md

@@ -11,6 +11,7 @@ Please mark all changes in change log and use the issue from GitHub
 -   \#3132 Add new indexes: RHNSW_PQ RHNSW_SQ
 -   \#3920 Allow an optional parameter `nbits` when creating IVF_PQ index
 -   \#4363 Put starting process of prometheus and grafana into docker-compose
+-   \#4376 Add an optional parameter naming `client_tag` on client constructor in c++ sdk
 
 ## Improvement
 

sdk/CMakeLists.txt

@@ -24,6 +24,7 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 include_directories(include)
 include_directories(grpc-gen/gen-milvus)
 include_directories(grpc-gen/gen-status)
+include_directories(grpc/interceptor)
 
 # set build type
 if (CMAKE_BUILD_TYPE STREQUAL "Release")
@@ -75,6 +76,10 @@ set(grpc_service_files
         grpc-gen/gen-status/status.pb.cc
         )
 
+set(grpc_interceptor_files
+        grpc/interceptor/GrpcClientInterceptor.cpp
+        grpc/interceptor/GrpcInterceptorHookHandler.cpp)
+
 set(grpc_lib
         grpcpp_channelz
         grpc++
@@ -87,6 +92,7 @@ add_library(milvus_sdk SHARED
         ${interface_files}
         ${grpc_client_files}
         ${grpc_service_files}
+        ${grpc_interceptor_files}
         )
 
 target_link_libraries(milvus_sdk

sdk/grpc/ClientProxy.cpp

@@ -10,6 +10,7 @@
 // or implied. See the License for the specific language governing permissions and limitations under the License.
 
 #include "grpc/ClientProxy.h"
+#include "interceptor/GrpcClientInterceptor.h"
 #include "thirdparty/nlohmann/json.hpp"
 
 #include <memory>
@@ -388,6 +389,17 @@ CopyEntityToJson(::milvus::grpc::Entities& grpc_entities, JSON& json_entity) {
     }
 }
 
+void
+ClientProxy::OnPreSendMessage(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                              ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) {
+    client_rpc_info->client_context()->AddMetadata("client_tag", client_tag_);
+}
+
+void
+ClientProxy::OnPostRecvInitialMetaData(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                                       ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) {
+}
+
 Status
 ClientProxy::Connect(const ConnectParam& param) {
     std::string uri = param.ip_address + ":" + param.port;
@@ -395,7 +407,14 @@ ClientProxy::Connect(const ConnectParam& param) {
     ::grpc::ChannelArguments args;
     args.SetMaxSendMessageSize(-1);
     args.SetMaxReceiveMessageSize(-1);
-    channel_ = ::grpc::CreateCustomChannel(uri, ::grpc::InsecureChannelCredentials(), args);
+    client_tag_ = param.client_tag_;
+
+    std::vector<std::unique_ptr<::grpc::experimental::ClientInterceptorFactoryInterface>> interceptor_creators;
+    interceptor_creators.emplace_back(std::unique_ptr<::grpc::experimental::ClientInterceptorFactoryInterface>(
+        new GrpcClientInterceptorFactory(this)));
+    channel_ = ::grpc::experimental::CreateCustomChannelWithInterceptors(uri, ::grpc::InsecureChannelCredentials(),
+                                                                         args, std::move(interceptor_creators));
+    //    channel_ = ::grpc::CreateCustomChannel(uri, ::grpc::InsecureChannelCredentials(), args);
     if (channel_ != nullptr) {
         connected_ = true;
         client_ptr_ = std::make_shared<GrpcClient>(channel_);

sdk/grpc/ClientProxy.h

@@ -13,6 +13,7 @@
 
 #include "GrpcClient.h"
 #include "MilvusApi.h"
+#include "interceptor/GrpcInterceptorHookHandler.h"
 
 #include <memory>
 #include <string>
@@ -20,9 +21,17 @@
 
 namespace milvus {
 
-class ClientProxy : public Connection {
+class ClientProxy : public Connection, public GrpcInterceptorHookHandler {
  public:
     // Implementations of the Connection interface
+    void
+    OnPostRecvInitialMetaData(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                              ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) override;
+
+    void
+    OnPreSendMessage(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                     ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) override;
+
     Status
     Connect(const ConnectParam& connect_param) override;
 
@@ -107,6 +116,7 @@ class ClientProxy : public Connection {
     std::shared_ptr<::grpc::Channel> channel_;
     std::shared_ptr<GrpcClient> client_ptr_;
     bool connected_ = false;
+    std::string client_tag_ = "";
 };
 
 }  // namespace milvus

sdk/grpc/interceptor/GrpcClientInterceptor.cpp

@@ -0,0 +1,37 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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.
+
+#include <grpc/grpc.h>
+#include <grpcpp/client_context.h>
+
+#include "GrpcClientInterceptor.h"
+
+namespace milvus {
+
+GrpcClientInterceptor::GrpcClientInterceptor(::grpc::experimental::ClientRpcInfo* info,
+                                             GrpcInterceptorHookHandler* hook_handler)
+    : info_(info), hook_handler_(hook_handler) {
+}
+
+void
+GrpcClientInterceptor::Intercept(::grpc::experimental::InterceptorBatchMethods* methods) {
+    if (methods->QueryInterceptionHookPoint(::grpc::experimental::InterceptionHookPoints::PRE_SEND_INITIAL_METADATA)) {
+        hook_handler_->OnPreSendMessage(info_, methods);
+    }
+    methods->Proceed();
+}
+
+::grpc::experimental::Interceptor*
+GrpcClientInterceptorFactory::CreateClientInterceptor(::grpc::experimental::ClientRpcInfo* info) {
+    return new GrpcClientInterceptor(info, hook_handler_);
+}
+
+}  // namespace milvus

sdk/grpc/interceptor/GrpcClientInterceptor.h

@@ -0,0 +1,41 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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.
+
+#pragma once
+
+#include "grpc/interceptor/GrpcInterceptorHookHandler.h"
+
+namespace milvus {
+
+class GrpcClientInterceptor : public ::grpc::experimental::Interceptor {
+ public:
+    explicit GrpcClientInterceptor(::grpc::experimental::ClientRpcInfo* info, GrpcInterceptorHookHandler* hook_handler);
+
+    void
+    Intercept(::grpc::experimental::InterceptorBatchMethods* methods) override;
+
+ private:
+    ::grpc::experimental::ClientRpcInfo* info_ = nullptr;
+    GrpcInterceptorHookHandler* hook_handler_;
+};
+
+class GrpcClientInterceptorFactory : public ::grpc::experimental::ClientInterceptorFactoryInterface {
+ public:
+    explicit GrpcClientInterceptorFactory(GrpcInterceptorHookHandler* hook_handler) : hook_handler_(hook_handler){};
+
+    ::grpc::experimental::Interceptor*
+    CreateClientInterceptor(::grpc::experimental::ClientRpcInfo* info) override;
+
+ private:
+    GrpcInterceptorHookHandler* hook_handler_;
+};
+
+}  // namespace milvus

sdk/grpc/interceptor/GrpcInterceptorHookHandler.cpp

@@ -0,0 +1,27 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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.
+
+#include "GrpcInterceptorHookHandler.h"
+
+namespace milvus {
+
+void
+GrpcInterceptorHookHandler::OnPostRecvInitialMetaData(
+    ::grpc::experimental::ClientRpcInfo* client_rpc_info,
+    ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) {
+}
+
+void
+GrpcInterceptorHookHandler::OnPreSendMessage(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                                             ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods) {
+}
+
+}  // namespace milvus

sdk/grpc/interceptor/GrpcInterceptorHookHandler.h

@@ -0,0 +1,30 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed 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.
+
+#pragma once
+
+#include <grpcpp/impl/codegen/interceptor.h>
+#include <grpcpp/impl/codegen/server_interceptor.h>
+
+namespace milvus {
+
+class GrpcInterceptorHookHandler {
+ public:
+    virtual void
+    OnPostRecvInitialMetaData(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                              ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods);
+
+    virtual void
+    OnPreSendMessage(::grpc::experimental::ClientRpcInfo* client_rpc_info,
+                     ::grpc::experimental::InterceptorBatchMethods* interceptor_batch_methods);
+};
+
+}  // namespace milvus

sdk/include/MilvusApi.h

@@ -62,8 +62,9 @@ enum class MetricType {
  * @brief Connect API parameter
  */
 struct ConnectParam {
-    std::string ip_address;  ///< Server IP address
-    std::string port;        ///< Server PORT
+    std::string ip_address;          ///< Server IP address
+    std::string port;                ///< Server PORT
+    std::string client_tag_ = "";    ///< Client Tag
 };
 
 /**