Mark instructions as cleaned up in the GRPC data channel if processing an instruction fails.

Author: tvalentyn

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

@@ -502,7 +502,7 @@ def _clean_receiving_queue(self, instruction_id):
     instruction_id cannot be reused for new queue.
     """
     with self._receive_lock:
-      self._received.pop(instruction_id)
+      self._received.pop(instruction_id, None)
       self._cleaned_instruction_ids[instruction_id] = True
       while len(self._cleaned_instruction_ids) > _MAX_CLEANED_INSTRUCTIONS:
         self._cleaned_instruction_ids.popitem(last=False)
@@ -787,6 +787,12 @@ def close(self):
     """Close all channels that this factory owns."""
     raise NotImplementedError(type(self))
 
+  def cleanup(self, instruction_id):
+    # type: (str) -> None
+
+    """Clean up resources for a given instruction."""
+    pass
+
 
 class GrpcClientDataChannelFactory(DataChannelFactory):
   """A factory for ``GrpcClientDataChannel``.
@@ -854,6 +860,11 @@ def close(self):
       channel.close()
     self._data_channel_cache.clear()
 
+  def cleanup(self, instruction_id):
+    # type: (str) -> None
+    for channel in self._data_channel_cache.values():
+      channel._clean_receiving_queue(instruction_id)
+
 
 class InMemoryDataChannelFactory(DataChannelFactory):
   """A singleton factory for ``InMemoryDataChannel``."""

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

@@ -568,6 +568,7 @@ def discard(self, instruction_id, exception):
 
     # Perform the shutdown while not holding the lock.
     processor.shutdown()
+    self.data_channel_factory.cleanup(instruction_id)
 
   def release(self, instruction_id):
     # type: (str) -> None