From 577c1dda748f53adfe927e7448313b4b739fb3e4 Mon Sep 17 00:00:00 2001 From: Stephanie Wang Date: Tue, 13 Nov 2018 18:32:24 -0800 Subject: [PATCH] Release sender connections as soon as WriteMessageAsync completes (#3313) --- src/ray/object_manager/connection_pool.cc | 26 ++++++++++++++ src/ray/object_manager/connection_pool.h | 9 +++++ src/ray/object_manager/object_manager.cc | 41 ++++++++++------------- 3 files changed, 53 insertions(+), 23 deletions(-) diff --git a/src/ray/object_manager/connection_pool.cc b/src/ray/object_manager/connection_pool.cc index 2104eaa4a..10ed2b8f3 100644 --- a/src/ray/object_manager/connection_pool.cc +++ b/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 &conn) { + std::unique_lock 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 *conn) { std::unique_lock 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 &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(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()) { diff --git a/src/ray/object_manager/connection_pool.h b/src/ray/object_manager/connection_pool.h index 6f96293dc..d2e92ee2a 100644 --- a/src/ray/object_manager/connection_pool.h +++ b/src/ray/object_manager/connection_pool.h @@ -57,6 +57,11 @@ class ConnectionPool { void RegisterSender(ConnectionType type, const ClientID &client_id, std::shared_ptr &conn); + /// Remove a sender connection. + /// + /// \param conn The actual connection. + void RemoveSender(const std::shared_ptr &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 &conn); + /// Removes the given sender for ClientID from the given map. + void Remove(SenderMapType &conn_map, const ClientID &client_id, + const std::shared_ptr &conn); + /// Returns the count of sender connections to ClientID. uint64_t Count(SenderMapType &conn_map, const ClientID &client_id); diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 033a40b01..1a6d48a68 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/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(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 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 &object_ } if (conn != nullptr) { - // TODO(swang): Make this a WriteMessageAsync. - ray::Status status = conn->WriteMessage( + conn->WriteMessageAsync( static_cast(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); } } }