Add success and fail callback of grpc sending reply (#5141)

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

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.

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<ObjectID> 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<ObjectID> &object_ids,

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
   ///

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(

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_;

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<ServerCall *>(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;

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`.

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`.

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<void(Status)>;
+/// \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<void(Status status, std::function<void()> success,
+                                             std::function<void()> 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 <class ServiceHandler, class Request, class Reply>
 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<void()> success,
+               std::function<void()> 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<void()> send_reply_success_callback_ = nullptr;
+
+  /// The callback when sending reply fails.
+  std::function<void()> send_reply_failure_callback_ = nullptr;
+
   template <class T1, class T2, class T3, class T4>
   friend class ServerCallFactoryImpl;
 };

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`.