Disconnect object manager clients if receiving an object fails (#4141)

* Disconnect object manager clients if ReadBuffer fails

* unused

* put back EINTR handling

src/ray/common/client_connection.cc

@@ -74,25 +74,26 @@ Status ServerConnection<T>::WriteBuffer(
 }
 
 template <class T>
-void ServerConnection<T>::ReadBuffer(
-    const std::vector<boost::asio::mutable_buffer> &buffer,
-    boost::system::error_code &ec) {
+Status ServerConnection<T>::ReadBuffer(
+    const std::vector<boost::asio::mutable_buffer> &buffer) {
+  boost::system::error_code error;
   // Loop until all bytes are read while handling interrupts.
   for (const auto &b : buffer) {
     uint64_t bytes_remaining = boost::asio::buffer_size(b);
     uint64_t position = 0;
     while (bytes_remaining != 0) {
       size_t bytes_read =
-          socket_.read_some(boost::asio::buffer(b + position, bytes_remaining), ec);
+          socket_.read_some(boost::asio::buffer(b + position, bytes_remaining), error);
       position += bytes_read;
       bytes_remaining -= bytes_read;
-      if (ec.value() == EINTR) {
+      if (error.value() == EINTR) {
         continue;
-      } else if (ec.value() != boost::system::errc::errc_t::success) {
-        return;
+      } else if (error.value() != boost::system::errc::errc_t::success) {
+        return boost_to_ray_status(error);
       }
     }
   }
+  return Status::OK();
 }
 
 template <class T>

src/ray/common/client_connection.h

@@ -59,15 +59,14 @@ class ServerConnection : public std::enable_shared_from_this<ServerConnection<T>
   /// Write a buffer to this connection.
   ///
   /// \param buffer The buffer.
-  /// \param ec The error code object in which to store error codes.
+  /// \return Status.
   Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer);
 
   /// Read a buffer from this connection.
   ///
   /// \param buffer The buffer.
-  /// \param ec The error code object in which to store error codes.
-  void ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer,
-                  boost::system::error_code &ec);
+  /// \return Status.
+  Status ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer);
 
   /// Shuts down socket for this connection.
   void Close() {

src/ray/object_manager/object_manager.cc

@@ -839,16 +839,17 @@ ray::Status ObjectManager::ExecuteReceiveObject(
 
   std::pair<const ObjectBufferPool::ChunkInfo &, ray::Status> chunk_status =
       buffer_pool_.CreateChunk(object_id, data_size, metadata_size, chunk_index);
+  ray::Status status;
   ObjectBufferPool::ChunkInfo chunk_info = chunk_status.first;
   if (chunk_status.second.ok()) {
     // Avoid handling this chunk if it's already being handled by another process.
     std::vector<boost::asio::mutable_buffer> buffer;
     buffer.push_back(asio::buffer(chunk_info.data, chunk_info.buffer_length));
-    boost::system::error_code ec;
-    conn.ReadBuffer(buffer, ec);
-    if (ec.value() == boost::system::errc::success) {
+    status = conn.ReadBuffer(buffer);
+    if (status.ok()) {
       buffer_pool_.SealChunk(object_id, chunk_index);
     } else {
+      // We may have not have read out the correct data, so abort this chunk.
       buffer_pool_.AbortCreateChunk(object_id, chunk_index);
       // TODO(hme): This chunk failed, so create a pull request for this chunk.
     }
@@ -860,19 +861,24 @@ ray::Status ObjectManager::ExecuteReceiveObject(
     mutable_vec.resize(buffer_length);
     std::vector<boost::asio::mutable_buffer> buffer;
     buffer.push_back(asio::buffer(mutable_vec, buffer_length));
-    boost::system::error_code ec;
-    conn.ReadBuffer(buffer, ec);
-    if (ec.value() != boost::system::errc::success) {
-      RAY_LOG(ERROR) << boost_to_ray_status(ec).ToString();
-    }
+    status = conn.ReadBuffer(buffer);
     // TODO(hme): If the object isn't local, create a pull request for this chunk.
   }
-  conn.ProcessMessages();
+
   RAY_LOG(DEBUG) << "ExecuteReceiveObject completed on " << client_id_ << " from "
                  << client_id << " of object " << object_id << " chunk " << chunk_index
                  << " at " << current_sys_time_ms();
+  if (status.ok()) {
+    // We successfully read the buffer, so we are ready to receive the next
+    // message.
+    conn.ProcessMessages();
+  } else {
+    // Close the connection by skipping the call to ProcessMessages.
+    RAY_LOG(ERROR) << "Failed to ExecuteReceiveObject from remote object manager, error: "
+                   << status;
+  }
 
-  return chunk_status.second;
+  return status;
 }
 
 void ObjectManager::ReceiveFreeRequest(std::shared_ptr<TcpClientConnection> &conn,

src/ray/object_manager/object_manager_client_connection.h

@@ -59,7 +59,7 @@ class SenderConnection : public boost::enable_shared_from_this<SenderConnection>
   /// Write a buffer to this connection.
   ///
   /// \param buffer The buffer.
-  /// \param ec The error code object in which to store error codes.
+  /// \return Status.
   Status WriteBuffer(const std::vector<boost::asio::const_buffer> &buffer) {
     return conn_->WriteBuffer(buffer);
   }
@@ -67,10 +67,9 @@ class SenderConnection : public boost::enable_shared_from_this<SenderConnection>
   /// Read a buffer from this connection.
   ///
   /// \param buffer The buffer.
-  /// \param ec The error code object in which to store error codes.
-  void ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer,
-                  boost::system::error_code &ec) {
-    return conn_->ReadBuffer(buffer, ec);
+  /// \return Status.
+  Status ReadBuffer(const std::vector<boost::asio::mutable_buffer> &buffer) {
+    return conn_->ReadBuffer(buffer);
   }
 
   /// \return The ClientID of this connection.

src/ray/raylet/raylet.cc

@@ -99,7 +99,10 @@ ray::Status Raylet::RegisterGcs(const std::string &node_ip_address,
 
   RAY_LOG(DEBUG) << "Node manager " << gcs_client_->client_table().GetLocalClientId()
                  << " started on " << client_info.node_manager_address << ":"
-                 << client_info.node_manager_port;
+                 << client_info.node_manager_port << " object manager at "
+                 << client_info.node_manager_address << ":"
+                 << client_info.object_manager_port;
+  ;
   RAY_RETURN_NOT_OK(gcs_client_->client_table().Connect(client_info));
 
   RAY_RETURN_NOT_OK(node_manager_.RegisterGcs());