Fix a memory leak in the failed instruction id cache. (#39405)

* Sanitize the exception message stored for failed instructions to avoid inadvertently capturing the content of stack frames in the heap and limit RAM growth due to cache size.

* Use Py3.10 syntax

Author: tvalentyn

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

@@ -80,7 +80,7 @@
 MAX_KNOWN_NOT_RUNNING_INSTRUCTIONS = 1000
 # The number of ProcessBundleRequest instruction ids that BundleProcessorCache
 # will remember for failed instructions.
-MAX_FAILED_INSTRUCTIONS = 10000
+MAX_FAILED_INSTRUCTIONS = 1000
 
 # retry on transient UNAVAILABLE grpc error from state channels.
 _GRPC_SERVICE_CONFIG = json.dumps({
@@ -559,7 +559,15 @@ def discard(self, instruction_id, exception):
     """
     processor = None
     with self._lock:
-      self.failed_instruction_ids[instruction_id] = exception
+      tb_str = "".join(traceback.format_exception(exception))
+      if len(tb_str) > 10240:
+        tb_str = (
+            tb_str[:5000] + "\n... [traceback truncated] ...\n" +
+            tb_str[-5000:])
+      clean_exception = RuntimeError(
+          f"Original Exception: {type(exception).__name__}: {str(exception)[:2000]}\n{tb_str}"
+      )
+      self.failed_instruction_ids[instruction_id] = clean_exception
       while len(self.failed_instruction_ids) > MAX_FAILED_INSTRUCTIONS:
         self.failed_instruction_ids.popitem(last=False)
       if instruction_id in self.active_bundle_processors:

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

@@ -296,6 +296,44 @@ def test_failed_bundle_processor_returns_failed_split_response(self):
         worker.do_instruction(split_request).error,
         hc.contains_string('test message'))
 
+  def test_failed_instruction_id_cache_size_is_capped(self):
+    data_channel_factory = mock.create_autospec(
+        data_plane.GrpcClientDataChannelFactory)
+    bundle_processor_cache = BundleProcessorCache(
+        None, None, data_channel_factory, {})
+    if bundle_processor_cache.periodic_shutdown:
+      bundle_processor_cache.periodic_shutdown.cancel()
+
+    with mock.patch(
+        'apache_beam.runners.worker.sdk_worker.MAX_FAILED_INSTRUCTIONS', 10):
+      for i in range(15):
+        bundle_processor_cache.discard(f'inst_{i}', RuntimeError(f'error {i}'))
+
+      for i in range(5):
+        self.assertNotIn(
+            f'inst_{i}', bundle_processor_cache.failed_instruction_ids)
+      for i in range(5, 15):
+        self.assertIn(
+            f'inst_{i}', bundle_processor_cache.failed_instruction_ids)
+
+  def test_failed_instruction_tracebacks_are_truncated_when_too_long(self):
+    data_channel_factory = mock.create_autospec(
+        data_plane.GrpcClientDataChannelFactory)
+    bundle_processor_cache = BundleProcessorCache(
+        None, None, data_channel_factory, {})
+    if bundle_processor_cache.periodic_shutdown:
+      bundle_processor_cache.periodic_shutdown.cancel()
+
+    long_message = "x" * 15000
+    bundle_processor_cache.discard('instruction_id', RuntimeError(long_message))
+
+    stored_exception = bundle_processor_cache.failed_instruction_ids[
+        'instruction_id']
+    tb_str = str(stored_exception)
+
+    self.assertLessEqual(len(tb_str), 13000)
+    self.assertIn('[traceback truncated]', tb_str)
+
   def test_data_sampling_response(self):
     # Create a data sampler with some fake sampled data. This data will be seen
     # in the sample response.