diff --git a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc index 655003b9a9b..cc86a360441 100644 --- a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc +++ b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.cc @@ -70,8 +70,7 @@ bool MetaServerNode::InitTCPServer() { return true; } -std::shared_ptr MetaServerNode::HandleMessage(const std::shared_ptr &message) { - MS_EXCEPTION_IF_NULL(message); +MessageBase *const MetaServerNode::HandleMessage(MessageBase *const message) { const auto &name = message->Name(); // Handle system messages. @@ -83,6 +82,7 @@ std::shared_ptr MetaServerNode::HandleMessage(const std::shared_ptr return rpc::NULL_MSG; } system_msg_handlers_[message_name](message); + return rpc::NULL_MSG; // Handle user defined messages. } else { @@ -92,13 +92,11 @@ std::shared_ptr MetaServerNode::HandleMessage(const std::shared_ptr return rpc::NULL_MSG; } (*message_handlers_[name])(message->Body()); + return rpc::NULL_MSG; } - return rpc::NULL_MSG; } -std::shared_ptr MetaServerNode::ProcessRegister(const std::shared_ptr &message) { - MS_EXCEPTION_IF_NULL(message); - +MessageBase *const MetaServerNode::ProcessRegister(MessageBase *const message) { RegistrationMessage registration; const std::string &body = message->Body(); registration.ParseFromArray(body.c_str(), body.length()); @@ -116,9 +114,7 @@ std::shared_ptr MetaServerNode::ProcessRegister(const std::shared_p return rpc::NULL_MSG; } -std::shared_ptr MetaServerNode::ProcessUnregister(const std::shared_ptr &message) { - MS_EXCEPTION_IF_NULL(message); - +MessageBase *const MetaServerNode::ProcessUnregister(MessageBase *const message) { UnregistrationMessage unregistration; const std::string &body = message->Body(); unregistration.ParseFromArray(body.c_str(), body.length()); @@ -133,9 +129,7 @@ std::shared_ptr MetaServerNode::ProcessUnregister(const std::shared return rpc::NULL_MSG; } -std::shared_ptr MetaServerNode::ProcessHeartbeat(const std::shared_ptr &message) { - MS_EXCEPTION_IF_NULL(message); - +MessageBase *const MetaServerNode::ProcessHeartbeat(MessageBase *const message) { HeartbeatMessage heartbeat; const std::string &body = message->Body(); heartbeat.ParseFromArray(body.c_str(), body.length()); @@ -185,7 +179,7 @@ size_t MetaServerNode::GetAliveNodeNum() { } bool MetaServerNode::RegisterMessageHandler(const std::string &name, - std::shared_ptr> handler) { + std::shared_ptr> handler) { if (message_handlers_.find(name) != message_handlers_.end()) { MS_LOG(ERROR) << "The message name: " << name << " have already been registered"; return false; diff --git a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h index 0dd0264c0b6..bbf03b61cff 100644 --- a/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h +++ b/mindspore/ccsrc/distributed/cluster/topology/meta_server_node.h @@ -73,23 +73,23 @@ class MetaServerNode : public NodeBase { // Register the message handler for the user defined message which is specified by the `name` parameter. bool RegisterMessageHandler(const std::string &name, - std::shared_ptr> handler); + std::shared_ptr> handler); private: // Create and init the tcp server. bool InitTCPServer(); // Handle the message received by the tcp server. - std::shared_ptr HandleMessage(const std::shared_ptr &message); + MessageBase *const HandleMessage(MessageBase *const message); // Process the received register message sent from compute graph nodes. - std::shared_ptr ProcessRegister(const std::shared_ptr &message); + MessageBase *const ProcessRegister(MessageBase *const message); // Process the received unregister message sent from compute graph nodes. - std::shared_ptr ProcessUnregister(const std::shared_ptr &message); + MessageBase *const ProcessUnregister(MessageBase *const message); // Process the received heartbeat message sent from compute graph nodes. - std::shared_ptr ProcessHeartbeat(const std::shared_ptr &message); + MessageBase *const ProcessHeartbeat(MessageBase *const message); // Maintain the state which is type of `TopoState` of this cluster topology. void UpdateTopoState(); @@ -106,7 +106,7 @@ class MetaServerNode : public NodeBase { // All the handlers for compute graph node's user-defined messages processing. // The `user-defined` means that this kind of message is user defined and has customized message handler. - std::map>> message_handlers_; + std::map>> message_handlers_; // Stores the registered compute graph nodes. std::map> nodes_; diff --git a/mindspore/ccsrc/distributed/rpc/tcp/connection.cc b/mindspore/ccsrc/distributed/rpc/tcp/connection.cc index 7f1a43e39b2..187148bf961 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/connection.cc +++ b/mindspore/ccsrc/distributed/rpc/tcp/connection.cc @@ -254,12 +254,9 @@ int Connection::ReceiveMessage() { return 0; } - std::shared_ptr msg(recv_message); - recv_message = nullptr; - // Call msg handler if set if (message_handler) { - message_handler(msg); + message_handler(recv_message); } else { MS_LOG(INFO) << "Message handler was not found"; } @@ -354,7 +351,6 @@ void Connection::FillSendMessage(MessageBase *msg, const std::string &advertiseU // update metrics send_metrics->UpdateMax(msg->body.size()); send_metrics->last_send_msg_name = msg->name; - return; } else { if (advertise_addr_.empty()) { diff --git a/mindspore/ccsrc/distributed/rpc/tcp/constants.h b/mindspore/ccsrc/distributed/rpc/tcp/constants.h index 175892683e8..4ff91de8767 100644 --- a/mindspore/ccsrc/distributed/rpc/tcp/constants.h +++ b/mindspore/ccsrc/distributed/rpc/tcp/constants.h @@ -30,7 +30,7 @@ namespace mindspore { namespace distributed { namespace rpc { -using MessageHandler = std::function(const std::shared_ptr &)>; +using MessageHandler = std::function; using DeleteCallBack = void (*)(const std::string &from, const std::string &to); using ConnectionCallBack = void (*)(void *conn); @@ -57,7 +57,7 @@ static const int g_httpKmsgEnable = -1; using IntTypeMetrics = std::queue; using StringTypeMetrics = std::queue; -static const std::shared_ptr NULL_MSG = nullptr; +static MessageBase *const NULL_MSG = nullptr; // Server socket listen backlog. static const int SOCKET_LISTEN_BACKLOG = 2048; diff --git a/mindspore/ccsrc/plugin/device/cpu/kernel/rpc/rpc_kernel.h b/mindspore/ccsrc/plugin/device/cpu/kernel/rpc/rpc_kernel.h index 128256bf87f..5eadbc05229 100644 --- a/mindspore/ccsrc/plugin/device/cpu/kernel/rpc/rpc_kernel.h +++ b/mindspore/ccsrc/plugin/device/cpu/kernel/rpc/rpc_kernel.h @@ -37,10 +37,10 @@ class RpcKernelMod : public NativeCpuKernelMod { void InitKernel(const CNodePtr &kernel_node) override { return; } // Set remote data as input. - void SetRemoteInput(const std::shared_ptr &msg) { remote_input_ = msg; } + void SetRemoteInput(MessageBase *const msg) { remote_input_ = msg; } protected: - std::shared_ptr remote_input_; + MessageBase *remote_input_; }; } // namespace kernel } // namespace mindspore diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc index 6d96c79e8c5..64be19598fa 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.cc @@ -73,7 +73,7 @@ bool RecvActor::StartServer() { return true; } -void RecvActor::RunOpInterProcessData(const std::shared_ptr &msg, OpContext *const context) { +void RecvActor::RunOpInterProcessData(MessageBase *const msg, OpContext *const context) { // Once recv actor is launched, reset the op_context so that the next step's recv will not be launched in advance. ResetOpcontext(); @@ -132,7 +132,7 @@ void RecvActor::EraseInput(const OpContext *context) { } } -std::shared_ptr RecvActor::HandleMessage(const std::shared_ptr &msg) { +MessageBase *RecvActor::HandleMessage(MessageBase *const msg) { // Block the message handler if the context is invalid. std::unique_lock lock(context_mtx_); context_cv_.wait(lock, [this] { return is_context_valid_; }); diff --git a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h index 76bda913972..99493ba16ca 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h +++ b/mindspore/ccsrc/runtime/graph_scheduler/actor/rpc/recv_actor.h @@ -59,14 +59,14 @@ class RecvActor : public RpcActor { bool CheckRunningCondition(const OpContext *context) const override; // When an inter-process data received, this method is called. - void RunOpInterProcessData(const std::shared_ptr &msg, OpContext *const context); + void RunOpInterProcessData(MessageBase *const msg, OpContext *const context); // Besides erasing input data and input controls when finish actor running, inter-process inputs should be erased. void EraseInput(const OpContext *context) override; private: // The message callback of the tcp server. - std::shared_ptr HandleMessage(const std::shared_ptr &msg); + MessageBase *HandleMessage(MessageBase *const msg); // The network address of this recv actor. It's generated automatically by rpc module. std::string ip_; diff --git a/tests/ut/cpp/distributed/cluster/topology/test_dynamic_networking.cc b/tests/ut/cpp/distributed/cluster/topology/test_dynamic_networking.cc index 3795f76a88a..bcb2a461f64 100644 --- a/tests/ut/cpp/distributed/cluster/topology/test_dynamic_networking.cc +++ b/tests/ut/cpp/distributed/cluster/topology/test_dynamic_networking.cc @@ -83,8 +83,11 @@ TEST_F(TestDynamicNetworking, AddMessageHandler) { std::string message_name = "route"; static std::string received_message; - auto func = std::make_shared>( - [](const std::string &message) { received_message = message; }); + auto func = + std::make_shared>([](const std::string &message) -> std::string { + received_message = message; + return message; + }); msn.RegisterMessageHandler(message_name, func); ComputeGraphNode cgn("compute_graph_node"); diff --git a/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc b/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc index 9f800cacaad..c8f216968b7 100644 --- a/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc +++ b/tests/ut/cpp/distributed/rpc/tcp/tcp_test.cc @@ -154,7 +154,7 @@ TEST_F(TCPTest, SendOneMessage) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const { IncrDataMsgNum(1); return NULL_MSG; }); @@ -196,7 +196,7 @@ TEST_F(TCPTest, SendTwoMessages) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const { IncrDataMsgNum(1); return NULL_MSG; }); @@ -251,7 +251,7 @@ TEST_F(TCPTest, SendSyncMessage) { bool ret = server->Initialize(server_url); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const { IncrDataMsgNum(1); return NULL_MSG; }); @@ -292,7 +292,7 @@ TEST_F(TCPTest, SendLargeMessages) { bool ret = server->Initialize(); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const { IncrDataMsgNum(1); return NULL_MSG; }); @@ -348,7 +348,7 @@ TEST_F(TCPTest, CreateManyConnectionPairs) { auto port = server->GetPort(); ASSERT_TRUE(ret); - server->SetMessageHandler([](const std::shared_ptr &message) -> std::shared_ptr { + server->SetMessageHandler([](MessageBase *const message) -> MessageBase *const { IncrDataMsgNum(1); return NULL_MSG; });