Address feedback.

Author: tvalentyn

sdks/python/apache_beam/runners/worker/data_plane.py

@@ -502,6 +502,10 @@ def _clean_receiving_queue(self, instruction_id):
     instruction_id cannot be reused for new queue.
     """
     with self._receive_lock:
+      # Per-instruction read queue may or may not be created yet when
+      # we mark an instruction as 'cleaned up' when creating
+      # a bundle processor failed, e.g. due to a flake in DoFn.setup().
+      # We want to mark an instruction as cleaned up regardless.
       self._received.pop(instruction_id, None)
       self._cleaned_instruction_ids[instruction_id] = True
       while len(self._cleaned_instruction_ids) > _MAX_CLEANED_INSTRUCTIONS:
@@ -857,13 +861,13 @@ def create_data_channel(self, remote_grpc_port):
   def close(self):
     # type: () -> None
     _LOGGER.info('Closing all cached grpc data channels.')
-    for _, channel in self._data_channel_cache.items():
+    for channel in list(self._data_channel_cache.values()):
       channel.close()
     self._data_channel_cache.clear()
 
   def cleanup(self, instruction_id):
     # type: (str) -> None
-    for channel in self._data_channel_cache.values():
+    for channel in list(self._data_channel_cache.values()):
       channel._clean_receiving_queue(instruction_id)