Release sender connections as soon as WriteMessageAsync completes (#3313)

src/ray/object_manager/connection_pool.cc

@@ -38,6 +38,17 @@ void ConnectionPool::RegisterSender(ConnectionType type, const ClientID &client_
   // Don't add to available connections. It will become available once it is released.
 }
 
+void ConnectionPool::RemoveSender(const std::shared_ptr<SenderConnection> &conn) {
+  std::unique_lock<std::mutex> guard(connection_mutex);
+  ClientID client_id = conn->GetClientID();
+  if (message_send_connections_.count(client_id) != 0) {
+    Remove(message_send_connections_, client_id, conn);
+  }
+  if (transfer_send_connections_.count(client_id) != 0) {
+    Remove(transfer_send_connections_, client_id, conn);
+  }
+}
+
 void ConnectionPool::GetSender(ConnectionType type, const ClientID &client_id,
                                std::shared_ptr<SenderConnection> *conn) {
   std::unique_lock<std::mutex> guard(connection_mutex);
@@ -85,6 +96,21 @@ void ConnectionPool::Remove(ReceiverMapType &conn_map, const ClientID &client_id
   connections.erase(connections.begin() + pos);
 }
 
+void ConnectionPool::Remove(SenderMapType &conn_map, const ClientID &client_id,
+                            const std::shared_ptr<SenderConnection> &conn) {
+  auto it = conn_map.find(client_id);
+  if (it == conn_map.end()) {
+    return;
+  }
+  auto &connections = it->second;
+  int64_t pos =
+      std::find(connections.begin(), connections.end(), conn) - connections.begin();
+  if (pos >= static_cast<int64_t>(connections.size())) {
+    return;
+  }
+  connections.erase(connections.begin() + pos);
+}
+
 uint64_t ConnectionPool::Count(SenderMapType &conn_map, const ClientID &client_id) {
   auto it = conn_map.find(client_id);
   if (it == conn_map.end()) {

src/ray/object_manager/connection_pool.h

@@ -57,6 +57,11 @@ class ConnectionPool {
   void RegisterSender(ConnectionType type, const ClientID &client_id,
                       std::shared_ptr<SenderConnection> &conn);
 
+  /// Remove a sender connection.
+  ///
+  /// \param conn The actual connection.
+  void RemoveSender(const std::shared_ptr<SenderConnection> &conn);
+
   /// Get a sender connection from the connection pool.
   /// The connection must be released or removed when the operation for which the
   /// connection was obtained is completed. If the connection pool is empty, the
@@ -108,6 +113,10 @@ class ConnectionPool {
   void Remove(ReceiverMapType &conn_map, const ClientID &client_id,
               std::shared_ptr<TcpClientConnection> &conn);
 
+  /// Removes the given sender for ClientID from the given map.
+  void Remove(SenderMapType &conn_map, const ClientID &client_id,
+              const std::shared_ptr<SenderConnection> &conn);
+
   /// Returns the count of sender connections to ClientID.
   uint64_t Count(SenderMapType &conn_map, const ClientID &client_id);
 

src/ray/object_manager/object_manager.cc

@@ -6,15 +6,6 @@ namespace asio = boost::asio;
 
 namespace object_manager_protocol = ray::object_manager::protocol;
 
-namespace {
-
-void CheckIOError(ray::Status &status, const std::string &operation) {
-  RAY_CHECK(status.IsIOError());
-  RAY_LOG(ERROR) << "Failed to contact remote object manager during " << operation;
-}
-
-}  // namespace
-
 namespace ray {
 
 ObjectManager::ObjectManager(asio::io_service &main_service,
@@ -278,6 +269,7 @@ void ObjectManager::PullEstablishConnection(const ObjectID &object_id,
 
   if (conn != nullptr) {
     PullSendRequest(object_id, conn);
+    connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
   }
 }
 
@@ -292,11 +284,11 @@ void ObjectManager::PullSendRequest(const ObjectID &object_id,
   fbb.Finish(message);
   conn->WriteMessageAsync(
       static_cast<int64_t>(object_manager_protocol::MessageType::PullRequest),
-      fbb.GetSize(), fbb.GetBufferPointer(), [this, conn](ray::Status status) mutable {
-        if (status.ok()) {
-          connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
-        } else {
-          CheckIOError(status, "Pull");
+      fbb.GetSize(), fbb.GetBufferPointer(), [this, conn](ray::Status status) {
+        if (!status.ok()) {
+          RAY_CHECK(status.IsIOError())
+              << "Failed to contact remote object manager during Pull";
+          connection_pool_.RemoveSender(conn);
         }
       });
 }
@@ -439,7 +431,9 @@ ray::Status ObjectManager::ExecuteSendObject(
   if (conn != nullptr) {
     status = SendObjectHeaders(object_id, data_size, metadata_size, chunk_index, conn);
     if (!status.ok()) {
-      CheckIOError(status, "Push");
+      RAY_CHECK(status.IsIOError())
+          << "Failed to contact remote object manager during Push";
+      connection_pool_.RemoveSender(conn);
     }
   }
   return status;
@@ -683,8 +677,6 @@ std::shared_ptr<SenderConnection> ObjectManager::CreateSenderConnection(
     RAY_LOG(ERROR) << "Failed to connect to remote object manager.";
   } else {
     // Register the new connection.
-    // TODO(Yuhong): Implement ConnectionPool::RemoveSender and call it if the client
-    // disconnects.
     connection_pool_.RegisterSender(type, info.client_id, conn);
     // Prepare client connection info buffer
     flatbuffers::FlatBufferBuilder fbb;
@@ -880,13 +872,16 @@ void ObjectManager::SpreadFreeObjectRequest(const std::vector<ObjectID> &object_
     }
 
     if (conn != nullptr) {
-      // TODO(swang): Make this a WriteMessageAsync.
-      ray::Status status = conn->WriteMessage(
+      conn->WriteMessageAsync(
           static_cast<int64_t>(object_manager_protocol::MessageType::FreeRequest),
-          fbb.GetSize(), fbb.GetBufferPointer());
-      if (status.ok()) {
-        connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
-      }
+          fbb.GetSize(), fbb.GetBufferPointer(), [this, conn](ray::Status status) {
+            if (!status.ok()) {
+              RAY_CHECK(status.IsIOError())
+                  << "Failed to contact remote object manager during Free";
+              connection_pool_.RemoveSender(conn);
+            }
+          });
+      connection_pool_.ReleaseSender(ConnectionPool::ConnectionType::MESSAGE, conn);
     }
   }
 }