Support closures and clean hard-deletion in MultiProcessShared spawn_process (#40673)

sdks/python/apache_beam/utils/multi_process_shared.py

@@ -26,6 +26,7 @@
 import logging
 import multiprocessing.managers
 import os
+import sys
 import tempfile
 import threading
 import time
@@ -38,6 +39,7 @@
 
 import fasteners
 
+from apache_beam.internal import pickler
 from apache_beam.utils import retry
 
 # In some python versions, there is a bug where AutoProxy doesn't handle
@@ -150,8 +152,22 @@ def release(self, proxy):
   def unsafe_hard_delete(self):
     with self.lock:
       if self.initialied:
+        obj = self.obj
         del self.obj
         self.initialied = False
+        # When hard-deleting in a spawned server process that exits via
+        # os._exit(0), an in-flight RPC thread (e.g. a timed-out call) may
+        # still hold a stack reference to `obj`, preventing CPython from
+        # invoking `obj.__del__()` before `os._exit(0)`.
+        if (self._hard_delete_callback is not None and
+            sys.getrefcount(obj) > 2 and hasattr(type(obj), '__del__')):
+          try:
+            type(obj).__del__(obj)
+          except Exception:
+            logging.warning(
+                'Exception while finalizing object during unsafe_hard_delete',
+                exc_info=True)
+        del obj
       if self._hard_delete_callback:
         self._hard_delete_callback()
 
@@ -232,7 +248,8 @@ def unsafe_hard_delete(self):
     self._proxyObject.unsafe_hard_delete()
 
 
-def _run_server_process(address_file, tag, constructor, authkey, life_line):
+def _run_server_process(
+    address_file, tag, serialized_constructor, authkey, life_line):
   """
     Runs in a separate process.
     Includes a 'Suicide Pact' monitor: If parent dies, I die.
@@ -276,6 +293,7 @@ def _monitor_parent():
     logging.getLogger().setLevel(logging.INFO)
     multiprocessing.current_process().authkey = authkey
 
+    constructor = pickler.loads(serialized_constructor)
     serving_manager = _SingletonRegistrar(
         address=('localhost', 0), authkey=authkey)
     _process_level_singleton_manager.set_hard_delete_callback(
@@ -439,7 +457,19 @@ def unsafe_hard_delete(self):
       to this object exist, or (b) you are ok with all existing references to
       this object throwing strange errors when derefrenced.
     """
-    self._get_manager().unsafe_hard_delete_singleton(self._tag)
+    address_file = os.path.join(self._path, self._tag) + ".address"
+    if (self._manager is None and
+        not _process_level_singleton_manager.has_singleton(self._tag) and
+        not os.path.exists(address_file)):
+      return
+    try:
+      self._get_manager().unsafe_hard_delete_singleton(self._tag)
+    except (EOFError, OSError, BrokenPipeError):
+      # Expected when a spawned server process exits via os._exit(0) before
+      # sending an RPC reply.
+      pass
+    finally:
+      self._manager = None
 
   def _create_server(self, address_file):
     if self._spawn_process:
@@ -459,7 +489,12 @@ def _create_server(self, address_file):
       ctx = multiprocessing.get_context('spawn')
       p = ctx.Process(
           target=_run_server_process,
-          args=(address_file, self._tag, self._constructor, AUTH_KEY, reader),
+          args=(
+              address_file,
+              self._tag,
+              pickler.dumps(self._constructor),
+              AUTH_KEY,
+              reader),
           daemon=False  # Must be False for nested proxies
       )
       p.start()

sdks/python/apache_beam/utils/multi_process_shared_test.py

@@ -16,11 +16,14 @@
 #
 # pytype: skip-file
 
+import gc
 import logging
 import multiprocessing
+import multiprocessing.connection
 import os
 import tempfile
 import threading
+import time
 import unittest
 from typing import Any
 from unittest.mock import patch
@@ -106,6 +109,39 @@ def __getattribute__(self, __name: str) -> Any:
       return object.__getattribute__(self, __name)
 
 
+class PidCounter(object):
+  def __init__(self, start=0):
+    self.running = start
+    self.lock = threading.Lock()
+
+  def get_pid(self):
+    return os.getpid()
+
+  def get(self):
+    return self.running
+
+  def increment(self, value=1):
+    with self.lock:
+      self.running += value
+      return self.running
+
+
+class SlowFinalizingObject(object):
+  def __init__(self, entered_path, finalized_path):
+    self._entered_path = entered_path
+    self._finalized_path = finalized_path
+
+  def block_call(self):
+    with open(self._entered_path, 'w') as f:
+      f.write('entered')
+    time.sleep(30)
+    return 'done'
+
+  def __del__(self):
+    with open(self._finalized_path, 'w') as f:
+      f.write('finalized')
+
+
 class MultiProcessSharedTest(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
@@ -492,6 +528,153 @@ def side_effect_connect(self_mgr, *args, **kwargs):
     self.assertEqual(counter1.increment(), 1)
     self.assertEqual(counter2.increment(), 2)
 
+  def test_spawn_process_runs_in_dedicated_pid(self):
+    parent_pid = os.getpid()
+    proxies = [
+        multi_process_shared.MultiProcessShared(
+            PidCounter,
+            tag=f'dedicated_pid_{i}',
+            always_proxy=True,
+            spawn_process=True).acquire() for i in range(3)
+    ]
+    pids = [proxy.get_pid() for proxy in proxies]
+    self.assertEqual(len(set(pids)), 3)
+    for pid in pids:
+      self.assertNotEqual(pid, parent_pid)
+
+    same_tag_proxy = multi_process_shared.MultiProcessShared(
+        PidCounter,
+        tag='dedicated_pid_0',
+        always_proxy=True,
+        spawn_process=True).acquire()
+    self.assertEqual(same_tag_proxy.get_pid(), pids[0])
+    self.assertEqual(proxies[0].increment(5), 5)
+    self.assertEqual(same_tag_proxy.get(), 5)
+
+  def test_no_cyclic_gc_self_deadlock_issue_40673(self):
+    """Regression test for https://github.com/apache/beam/issues/40673.
+
+    When always_proxy=True and spawn_process=False, cyclic GC collecting an
+    unreachable BaseProxy on the in-process Server.accepter thread deadlocks
+    the server against itself. With spawn_process=True, the server runs in a
+    dedicated child process with no BaseProxy references, preventing deadlock.
+    """
+    orig_accept = multiprocessing.connection.Listener.accept
+
+    def gc_on_accepter_accept(listener_self):
+      if 'accepter' in threading.current_thread().name:
+        gc.collect()
+      return orig_accept(listener_self)
+
+    stop_bg_gc = threading.Event()
+
+    def background_gc():
+      while not stop_bg_gc.is_set():
+        gc.collect()
+        time.sleep(0.005)
+
+    errors = []
+
+    def workload():
+      try:
+        for i in range(15):
+          handle = multi_process_shared.MultiProcessShared(
+              Counter,
+              tag='issue_40673_tag',
+              always_proxy=True,
+              spawn_process=True)
+          proxy = handle.acquire()
+          self.assertEqual(proxy.increment(), i + 1)
+          cycle = []
+          cycle.append(cycle)
+          cycle.append(proxy)
+          cycle.append(handle)
+          del proxy
+          del handle
+          del cycle
+          if i % 3 == 0:
+            gc.collect()
+      except BaseException as e:  # pylint: disable=broad-exception-caught
+        errors.append(e)
+
+    with patch.object(multiprocessing.connection.Listener,
+                      'accept',
+                      gc_on_accepter_accept):
+      gc_thread = threading.Thread(target=background_gc, daemon=True)
+      worker_thread = threading.Thread(target=workload, daemon=True)
+      gc_thread.start()
+      worker_thread.start()
+      worker_thread.join(timeout=30)
+      stop_bg_gc.set()
+      gc_thread.join(timeout=5)
+
+    self.assertFalse(
+        worker_thread.is_alive(),
+        'Deadlock detected while acquiring/collecting cyclic BaseProxy '
+        'instances with spawn_process=True')
+    self.assertEqual(errors, [])
+
+    final_proxy = multi_process_shared.MultiProcessShared(
+        Counter, tag='issue_40673_tag', always_proxy=True,
+        spawn_process=True).acquire()
+    self.assertEqual(final_proxy.get(), 15)
+
+  def test_spawn_process_supports_lambda_and_closure_constructors(self):
+    captured_offset = 42
+
+    def local_factory(extra):
+      return PidCounter(start=captured_offset + extra)
+
+    proxy = multi_process_shared.MultiProcessShared(
+        lambda: local_factory(8),
+        tag='closure_constructor',
+        always_proxy=True,
+        spawn_process=True).acquire()
+    self.assertNotEqual(proxy.get_pid(), os.getpid())
+    self.assertEqual(proxy.get(), 50)
+    self.assertEqual(proxy.increment(5), 55)
+
+  def test_unsafe_hard_delete_missing_tag_is_noop(self):
+    handle = multi_process_shared.MultiProcessShared(
+        Counter, tag='nonexistent_tag', always_proxy=True, spawn_process=True)
+    handle.unsafe_hard_delete()
+    self.assertEqual(multiprocessing.active_children(), [])
+
+  def test_unsafe_hard_delete_runs_del_during_inflight_blocked_call(self):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+      entered_path = os.path.join(tmp_dir, 'entered')
+      finalized_path = os.path.join(tmp_dir, 'finalized')
+      handle = multi_process_shared.MultiProcessShared(
+          lambda: SlowFinalizingObject(entered_path, finalized_path),
+          tag='inflight_hard_delete',
+          always_proxy=True,
+          spawn_process=True)
+      proxy = handle.acquire()
+      server_proc = next(
+          p for p in multiprocessing.active_children() if p.is_alive())
+
+      def call_ignoring_disconnect():
+        try:
+          proxy.block_call()
+        except (EOFError, OSError):
+          pass
+
+      call_thread = threading.Thread(
+          target=call_ignoring_disconnect, daemon=True)
+      call_thread.start()
+
+      deadline = time.time() + 5
+      while not os.path.exists(entered_path) and time.time() < deadline:
+        time.sleep(0.02)
+      self.assertTrue(os.path.exists(entered_path))
+
+      handle.unsafe_hard_delete()
+      server_proc.join(timeout=5)
+      self.assertFalse(server_proc.is_alive())
+      self.assertTrue(os.path.exists(finalized_path))
+      with open(finalized_path) as f:
+        self.assertEqual(f.read(), 'finalized')
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)