disconnect bug fix. (#1837)

src/ray/common/client_connection.cc

@@ -122,7 +122,6 @@ void ClientConnection<T>::ProcessMessageHeader(const boost::system::error_code &
 template <class T>
 void ClientConnection<T>::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());

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<TcpClientConnection> &conn) {
+void ConnectionPool::RemoveReceiver(std::shared_ptr<TcpClientConnection> conn){
   std::unique_lock<std::mutex> 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,

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<TcpClientConnection> &conn);
+  void RemoveReceiver(std::shared_ptr<TcpClientConnection> conn);
 
   /// Register a receiver connection.
   ///

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

src/ray/object_manager/object_manager.cc

@@ -425,7 +425,8 @@ void ObjectManager::ProcessClientMessage(std::shared_ptr<TcpClientConnection> 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<TcpClientConnection> &conn,
 
 void ObjectManager::DisconnectClient(std::shared_ptr<TcpClientConnection> &conn,
                                      const uint8_t *message) {
-  auto info =
-      flatbuffers::GetRoot<object_manager_protocol::DisconnectClientMessage>(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<TcpClientConnection> &conn,