From 5733690aa640ba39131782637cf1c5f00be869e1 Mon Sep 17 00:00:00 2001 From: Joey Jiang <452084368@qq.com> Date: Tue, 9 Jul 2019 17:03:57 +0800 Subject: [PATCH] Add success and fail callback of grpc sending reply (#5141) --- .../core_worker/transport/raylet_transport.cc | 4 +- .../core_worker/transport/raylet_transport.h | 6 +- src/ray/object_manager/object_manager.cc | 12 ++-- src/ray/object_manager/object_manager.h | 12 ++-- src/ray/raylet/node_manager.cc | 4 +- src/ray/raylet/node_manager.h | 2 +- src/ray/rpc/grpc_server.cc | 18 ++++-- .../rpc/node_manager/node_manager_server.h | 6 +- .../object_manager/object_manager_server.h | 10 +-- src/ray/rpc/server_call.h | 63 ++++++++++++++----- src/ray/rpc/worker/worker_server.h | 6 +- 11 files changed, 92 insertions(+), 51 deletions(-) diff --git a/src/ray/core_worker/transport/raylet_transport.cc b/src/ray/core_worker/transport/raylet_transport.cc index df9d2fbe9..11f66dc04 100644 --- a/src/ray/core_worker/transport/raylet_transport.cc +++ b/src/ray/core_worker/transport/raylet_transport.cc @@ -37,11 +37,11 @@ CoreWorkerRayletTaskReceiver::CoreWorkerRayletTaskReceiver( void CoreWorkerRayletTaskReceiver::HandleAssignTask( const rpc::AssignTaskRequest &request, rpc::AssignTaskReply *reply, - rpc::RequestDoneCallback done_callback) { + rpc::SendReplyCallback send_reply_callback) { const raylet::Task task(request.task()); const auto &spec = task.GetTaskSpecification(); auto status = task_handler_(spec); - done_callback(status); + send_reply_callback(status, nullptr, nullptr); } } // namespace ray diff --git a/src/ray/core_worker/transport/raylet_transport.h b/src/ray/core_worker/transport/raylet_transport.h index cdf8deebb..7e8fea1ee 100644 --- a/src/ray/core_worker/transport/raylet_transport.h +++ b/src/ray/core_worker/transport/raylet_transport.h @@ -44,14 +44,14 @@ class CoreWorkerRayletTaskReceiver : public CoreWorkerTaskReceiver, /// /// Handle a `AssignTask` request. /// The implementation can handle this request asynchronously. When hanling is done, the - /// `done_callback` should be called. + /// `send_reply_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. - /// \param[in] done_callback The callback to be called when the request is done. + /// \param[in] send_reply_callback The callback to be called when the request is done. void HandleAssignTask(const rpc::AssignTaskRequest &request, rpc::AssignTaskReply *reply, - rpc::RequestDoneCallback done_callback) override; + rpc::SendReplyCallback send_reply_callback) override; private: /// Raylet client. diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 67a830a32..08c05a0f4 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -664,7 +664,7 @@ void ObjectManager::WaitComplete(const UniqueID &wait_id) { /// Implementation of ObjectManagerServiceHandler void ObjectManager::HandlePushRequest(const rpc::PushRequest &request, rpc::PushReply *reply, - rpc::RequestDoneCallback done_callback) { + rpc::SendReplyCallback send_reply_callback) { ObjectID object_id = ObjectID::FromBinary(request.object_id()); ClientID client_id = ClientID::FromBinary(request.client_id()); @@ -680,7 +680,7 @@ void ObjectManager::HandlePushRequest(const rpc::PushRequest &request, double end_time = current_sys_time_seconds(); HandleReceiveFinished(object_id, client_id, chunk_index, start_time, end_time, status); - done_callback(status); + send_reply_callback(status, nullptr, nullptr); } ray::Status ObjectManager::ReceiveObjectChunk(const ClientID &client_id, @@ -711,7 +711,7 @@ ray::Status ObjectManager::ReceiveObjectChunk(const ClientID &client_id, void ObjectManager::HandlePullRequest(const rpc::PullRequest &request, rpc::PullReply *reply, - rpc::RequestDoneCallback done_callback) { + rpc::SendReplyCallback send_reply_callback) { ObjectID object_id = ObjectID::FromBinary(request.object_id()); ClientID client_id = ClientID::FromBinary(request.client_id()); RAY_LOG(DEBUG) << "Received pull request from client " << client_id << " for object [" @@ -729,18 +729,18 @@ void ObjectManager::HandlePullRequest(const rpc::PullRequest &request, } main_service_->post([this, object_id, client_id]() { Push(object_id, client_id); }); - done_callback(Status::OK()); + send_reply_callback(Status::OK(), nullptr, nullptr); } void ObjectManager::HandleFreeObjectsRequest(const rpc::FreeObjectsRequest &request, rpc::FreeObjectsReply *reply, - rpc::RequestDoneCallback done_callback) { + rpc::SendReplyCallback send_reply_callback) { std::vector object_ids; for (const auto &e : request.object_ids()) { object_ids.emplace_back(ObjectID::FromBinary(e)); } FreeObjects(object_ids, /* local_only */ true); - done_callback(Status::OK()); + send_reply_callback(Status::OK(), nullptr, nullptr); } void ObjectManager::FreeObjects(const std::vector &object_ids, diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index 019b92bb9..00d5d9eef 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -79,26 +79,26 @@ class ObjectManager : public ObjectManagerInterface, /// /// \param request Push request including the object chunk data /// \param reply Reply to the sender - /// \param done_callback Callback of the request + /// \param send_reply_callback Callback of the request void HandlePushRequest(const rpc::PushRequest &request, rpc::PushReply *reply, - rpc::RequestDoneCallback done_callback) override; + rpc::SendReplyCallback send_reply_callback) override; /// Handle pull request from remote object manager /// /// \param request Pull request /// \param reply Reply - /// \param done_callback Callback of request + /// \param send_reply_callback Callback of request void HandlePullRequest(const rpc::PullRequest &request, rpc::PullReply *reply, - rpc::RequestDoneCallback done_callback) override; + rpc::SendReplyCallback send_reply_callback) override; /// Handle free objects request /// /// \param request Free objects request /// \param reply Reply - /// \param done_callback + /// \param send_reply_callback void HandleFreeObjectsRequest(const rpc::FreeObjectsRequest &request, rpc::FreeObjectsReply *reply, - rpc::RequestDoneCallback done_callback) override; + rpc::SendReplyCallback send_reply_callback) override; /// Send object to remote object manager /// diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index f9c53d99a..f9b5864ae 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -1215,7 +1215,7 @@ void NodeManager::ProcessNewNodeManager(TcpClientConnection &node_manager_client void NodeManager::HandleForwardTask(const rpc::ForwardTaskRequest &request, rpc::ForwardTaskReply *reply, - rpc::RequestDoneCallback done_callback) { + rpc::SendReplyCallback send_reply_callback) { // Get the forwarded task and its uncommitted lineage from the request. TaskID task_id = TaskID::FromBinary(request.task_id()); Lineage uncommitted_lineage; @@ -1228,7 +1228,7 @@ void NodeManager::HandleForwardTask(const rpc::ForwardTaskRequest &request, << " on node " << gcs_client_->client_table().GetLocalClientId() << " spillback=" << task.GetTaskExecutionSpec().NumForwards(); SubmitTask(task, uncommitted_lineage, /* forwarded = */ true); - done_callback(Status::OK()); + send_reply_callback(Status::OK(), nullptr, nullptr); } void NodeManager::ProcessSetResourceRequest( diff --git a/src/ray/raylet/node_manager.h b/src/ray/raylet/node_manager.h index 582f6bd77..3c30c20c4 100644 --- a/src/ray/raylet/node_manager.h +++ b/src/ray/raylet/node_manager.h @@ -459,7 +459,7 @@ class NodeManager : public rpc::NodeManagerServiceHandler { /// Handle a `ForwardTask` request. void HandleForwardTask(const rpc::ForwardTaskRequest &request, rpc::ForwardTaskReply *reply, - rpc::RequestDoneCallback done_callback) override; + rpc::SendReplyCallback send_reply_callback) override; // GCS client ID for this node. ClientID client_id_; diff --git a/src/ray/rpc/grpc_server.cc b/src/ray/rpc/grpc_server.cc index 80ad81d98..e8988bd2d 100644 --- a/src/ray/rpc/grpc_server.cc +++ b/src/ray/rpc/grpc_server.cc @@ -49,9 +49,7 @@ void GrpcServer::PollEventsFromCompletionQueue() { // Keep reading events from the `CompletionQueue` until it's shutdown. while (cq_->Next(&tag, &ok)) { auto *server_call = static_cast(tag); - // `ok == false` indicates that the server has been shut down. - // We should delete the call object in this case. - bool delete_call = !ok; + bool delete_call = false; if (ok) { switch (server_call->GetState()) { case ServerCallState::PENDING: @@ -63,14 +61,24 @@ void GrpcServer::PollEventsFromCompletionQueue() { server_call->HandleRequest(); break; case ServerCallState::SENDING_REPLY: - // The reply has been sent, this call can be deleted now. - // This event is triggered by `ServerCallImpl::Finish`. + // GRPC has sent reply successfully, invoking the callback. + server_call->OnReplySent(); + // The rpc call has finished and can be deleted now. delete_call = true; break; default: RAY_LOG(FATAL) << "Shouldn't reach here."; break; } + } else { + // `ok == false` will occur in two situations: + // First, the server has been shut down, the server call's status is PENDING + // Second, server has sent reply to client and failed, the server call's status is + // SENDING_REPLY + if (server_call->GetState() == ServerCallState::SENDING_REPLY) { + server_call->OnReplyFailed(); + } + delete_call = true; } if (delete_call) { delete server_call; diff --git a/src/ray/rpc/node_manager/node_manager_server.h b/src/ray/rpc/node_manager/node_manager_server.h index 4505d1017..76531f10e 100644 --- a/src/ray/rpc/node_manager/node_manager_server.h +++ b/src/ray/rpc/node_manager/node_manager_server.h @@ -15,14 +15,14 @@ class NodeManagerServiceHandler { public: /// Handle a `ForwardTask` request. /// The implementation can handle this request asynchronously. When handling is done, - /// the `done_callback` should be called. + /// the `send_reply_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. - /// \param[in] done_callback The callback to be called when the request is done. + /// \param[in] send_reply_callback The callback to be called when the request is done. virtual void HandleForwardTask(const ForwardTaskRequest &request, ForwardTaskReply *reply, - RequestDoneCallback done_callback) = 0; + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcService` for `NodeManagerService`. diff --git a/src/ray/rpc/object_manager/object_manager_server.h b/src/ray/rpc/object_manager/object_manager_server.h index dbabc11a9..c0af15ffb 100644 --- a/src/ray/rpc/object_manager/object_manager_server.h +++ b/src/ray/rpc/object_manager/object_manager_server.h @@ -16,20 +16,20 @@ class ObjectManagerServiceHandler { public: /// Handle a `Push` request. /// The implementation can handle this request asynchronously. When handling is done, - /// the `done_callback` should be called. + /// the `send_reply_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. - /// \param[in] done_callback The callback to be called when the request is done. + /// \param[in] send_reply_callback The callback to be called when the request is done. virtual void HandlePushRequest(const PushRequest &request, PushReply *reply, - RequestDoneCallback done_callback) = 0; + SendReplyCallback send_reply_callback) = 0; /// Handle a `Pull` request virtual void HandlePullRequest(const PullRequest &request, PullReply *reply, - RequestDoneCallback done_callback) = 0; + SendReplyCallback send_reply_callback) = 0; /// Handle a `FreeObjects` request virtual void HandleFreeObjectsRequest(const FreeObjectsRequest &request, FreeObjectsReply *reply, - RequestDoneCallback done_callback) = 0; + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcService` for `ObjectManagerGrpcService`. diff --git a/src/ray/rpc/server_call.h b/src/ray/rpc/server_call.h index 263250e73..19964c5bd 100644 --- a/src/ray/rpc/server_call.h +++ b/src/ray/rpc/server_call.h @@ -11,7 +11,13 @@ namespace rpc { /// Represents the callback function to be called when a `ServiceHandler` finishes /// handling a request. -using RequestDoneCallback = std::function; +/// \param status The status would be returned to client. +/// \param success Success callback which will be invoked when the reply is successfully +/// sent to the client. +/// \param failure Failure callback which will be invoked when the reply fails to be +/// sent to the client. +using SendReplyCallback = std::function success, + std::function failure)>; /// Represents state of a `ServerCall`. enum class ServerCallState { @@ -59,8 +65,11 @@ class ServerCall { /// Get the factory that created this `ServerCall`. virtual const ServerCallFactory &GetFactory() const = 0; - /// Finish the `ServerCall`. - virtual void Finish(Status status) = 0; + /// Invoked when sending reply successes. + virtual void OnReplySent() = 0; + + // Invoked when sending reply fails. + virtual void OnReplyFailed() = 0; /// Virtual destruct function to make sure subclass would destruct properly. virtual ~ServerCall() = default; @@ -84,7 +93,7 @@ class ServerCallFactory { /// \tparam Reply Type of the reply message. template using HandleRequestFunction = void (ServiceHandler::*)(const Request &, Reply *, - RequestDoneCallback); + SendReplyCallback); /// Implementation of `ServerCall`. It represents `ServerCall` for a particular /// RPC method. @@ -123,30 +132,48 @@ class ServerCallImpl : public ServerCall { // Handle service for rpc call has stopped, we must handle the call here // to send reply and remove it from cq RAY_LOG(DEBUG) << "Handle service has been closed."; - Finish(Status::Invalid("HandleServiceClosed")); + SendReply(Status::Invalid("HandleServiceClosed")); } } void HandleRequestImpl() { state_ = ServerCallState::PROCESSING; - (service_handler_.*handle_request_function_)(request_, &reply_, - [this](Status status) { - // When the handler is done with the - // request, tell gRPC to finish this - // request. - Finish(status); - }); + (service_handler_.*handle_request_function_)( + request_, &reply_, + [this](Status status, std::function success, + std::function failure) { + // When the handler is done with the + // request, tell gRPC to finish this + // request. + SendReply(status); + send_reply_success_callback_ = std::move(success); + send_reply_failure_callback_ = std::move(failure); + }); } const ServerCallFactory &GetFactory() const override { return factory_; } - /// Tell gRPC to finish this request. - void Finish(Status status) override { + void OnReplySent() { + if (send_reply_success_callback_ && !io_service_.stopped()) { + auto callback = std::move(send_reply_success_callback_); + io_service_.post([callback]() { callback(); }); + } + } + + void OnReplyFailed() { + if (send_reply_failure_callback_ && !io_service_.stopped()) { + auto callback = std::move(send_reply_failure_callback_); + io_service_.post([callback]() { callback(); }); + } + } + + private: + /// Tell gRPC to finish this request and send reply asynchronously. + void SendReply(Status status) { state_ = ServerCallState::SENDING_REPLY; response_writer_.Finish(reply_, RayStatusToGrpcStatus(status), this); } - private: /// State of this call. ServerCallState state_; @@ -175,6 +202,12 @@ class ServerCallImpl : public ServerCall { /// The reply message. Reply reply_; + /// The callback when sending reply successes. + std::function send_reply_success_callback_ = nullptr; + + /// The callback when sending reply fails. + std::function send_reply_failure_callback_ = nullptr; + template friend class ServerCallFactoryImpl; }; diff --git a/src/ray/rpc/worker/worker_server.h b/src/ray/rpc/worker/worker_server.h index adf9bb982..fe25f9c94 100644 --- a/src/ray/rpc/worker/worker_server.h +++ b/src/ray/rpc/worker/worker_server.h @@ -15,13 +15,13 @@ class WorkerTaskHandler { public: /// Handle a `AssignTask` request. /// The implementation can handle this request asynchronously. When handling is done, - /// the `done_callback` should be called. + /// the `send_reply_callback` should be called. /// /// \param[in] request The request message. /// \param[out] reply The reply message. - /// \param[in] done_callback The callback to be called when the request is done. + /// \param[in] send_reply_callback The callback to be called when the request is done. virtual void HandleAssignTask(const AssignTaskRequest &request, AssignTaskReply *reply, - RequestDoneCallback done_callback) = 0; + SendReplyCallback send_reply_callback) = 0; }; /// The `GrpcServer` for `WorkerService`.