diff --git a/src/ray/common/client_connection.cc b/src/ray/common/client_connection.cc index c1b237775..867199de8 100644 --- a/src/ray/common/client_connection.cc +++ b/src/ray/common/client_connection.cc @@ -122,7 +122,6 @@ void ClientConnection::ProcessMessageHeader(const boost::system::error_code & template void ClientConnection::ProcessMessage(const boost::system::error_code &error) { if (error) { - // TODO(hme): Disconnect differently & remove dependency on node_manager_generated.h read_type_ = protocol::MessageType_DisconnectClient; } message_handler_(this->shared_from_this(), read_type_, read_message_.data()); diff --git a/src/ray/object_manager/connection_pool.cc b/src/ray/object_manager/connection_pool.cc index 0fe0f2161..d897bcdfe 100644 --- a/src/ray/object_manager/connection_pool.cc +++ b/src/ray/object_manager/connection_pool.cc @@ -17,18 +17,15 @@ void ConnectionPool::RegisterReceiver(ConnectionType type, const ClientID &clien } } -void ConnectionPool::RemoveReceiver(ConnectionType type, const ClientID &client_id, - std::shared_ptr &conn) { +void ConnectionPool::RemoveReceiver(std::shared_ptr conn){ std::unique_lock guard(connection_mutex); - switch (type) { - case ConnectionType::MESSAGE: { + ClientID client_id = conn->GetClientID(); + if (message_receive_connections_.count(client_id) != 0){ Remove(message_receive_connections_, client_id, conn); - } break; - case ConnectionType::TRANSFER: { - Remove(transfer_receive_connections_, client_id, conn); - } break; } - // TODO(hme): appropriately dispose of client connection. + if (transfer_receive_connections_.count(client_id) != 0){ + Remove(transfer_receive_connections_, client_id, conn); + } } void ConnectionPool::RegisterSender(ConnectionType type, const ClientID &client_id, diff --git a/src/ray/object_manager/connection_pool.h b/src/ray/object_manager/connection_pool.h index 0fba1b50c..700d4b884 100644 --- a/src/ray/object_manager/connection_pool.h +++ b/src/ray/object_manager/connection_pool.h @@ -46,11 +46,8 @@ class ConnectionPool { /// Remove a receiver connection. /// - /// \param type The type of connection. - /// \param client_id The ClientID of the remote object manager. /// \param conn The actual connection. - void RemoveReceiver(ConnectionType type, const ClientID &client_id, - std::shared_ptr &conn); + void RemoveReceiver(std::shared_ptr conn); /// Register a receiver connection. /// diff --git a/src/ray/object_manager/format/object_manager.fbs b/src/ray/object_manager/format/object_manager.fbs index d1583e6c1..7d833d728 100644 --- a/src/ray/object_manager/format/object_manager.fbs +++ b/src/ray/object_manager/format/object_manager.fbs @@ -3,7 +3,6 @@ namespace ray.object_manager.protocol; enum MessageType:int { ConnectClient = 1, - DisconnectClient, PushRequest, PullRequest } @@ -28,10 +27,3 @@ table ConnectClientMessage { // Whether this is a transfer connection. is_transfer: bool; } - -table DisconnectClientMessage { - // ID of the connecting client. - client_id: string; - // Whether this is a transfer connection. - is_transfer: bool; -} diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 60672f788..f5294524a 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -425,7 +425,8 @@ void ObjectManager::ProcessClientMessage(std::shared_ptr co ConnectClient(conn, message); break; } - case object_manager_protocol::MessageType_DisconnectClient: { + case protocol::MessageType_DisconnectClient: { + // TODO(hme): Disconnect without depending on the node manager protocol. DisconnectClient(conn, message); break; } @@ -453,17 +454,7 @@ void ObjectManager::ConnectClient(std::shared_ptr &conn, void ObjectManager::DisconnectClient(std::shared_ptr &conn, const uint8_t *message) { - auto info = - flatbuffers::GetRoot(message); - ClientID client_id = ObjectID::from_binary(info->client_id()->str()); - bool is_transfer = info->is_transfer(); - if (is_transfer) { - connection_pool_.RemoveReceiver(ConnectionPool::ConnectionType::TRANSFER, client_id, - conn); - } else { - connection_pool_.RemoveReceiver(ConnectionPool::ConnectionType::MESSAGE, client_id, - conn); - } + connection_pool_.RemoveReceiver(conn); } void ObjectManager::ReceivePullRequest(std::shared_ptr &conn,