Merge pull request #308 from sylwiaszunejko/keyerror

Remove endpoint to host_id mapping when removing host by host_id

cassandra/cluster.py

@@ -4007,7 +4007,7 @@ def _refresh_node_list_and_token_map(self, connection, preloaded_results=None,
             if old_host_id not in found_host_ids:
                 should_rebuild_token_map = True
                 log.debug("[control connection] Removing host not found in peers metadata: %r", old_host)
-                self._cluster.metadata.remove_host_by_host_id(old_host_id)
+                self._cluster.metadata.remove_host_by_host_id(old_host_id, old_host.endpoint)
 
         log.debug("[control connection] Finished fetching ring info")
         if partitioner and should_rebuild_token_map:

cassandra/metadata.py

@@ -344,8 +344,10 @@ def remove_host(self, host):
             self._host_id_by_endpoint.pop(host.endpoint, False)
             return bool(self._hosts.pop(host.host_id, False))
 
-    def remove_host_by_host_id(self, host_id):
+    def remove_host_by_host_id(self, host_id, endpoint=None):
         with self._hosts_lock:
+            if endpoint and self._host_id_by_endpoint[endpoint] == host_id:
+                self._host_id_by_endpoint.pop(endpoint, False)
             return bool(self._hosts.pop(host_id, False))
 
     def update_host(self, host, old_endpoint):

tests/unit/test_control_connection.py

@@ -88,7 +88,9 @@ def update_host(self, host, old_endpoint):
     def all_hosts_items(self):
         return list(self.hosts.items())
 
-    def remove_host_by_host_id(self, host_id):
+    def remove_host_by_host_id(self, host_id, endpoint=None):
+        if endpoint and self._host_id_by_endpoint[endpoint] == host_id:
+            self._host_id_by_endpoint.pop(endpoint, False)
         self.removed_hosts.append(self.hosts.pop(host_id, False))
         return bool(self.hosts.pop(host_id, False))