Store the registration and purge in memory and output at the end.

Author: shunping

sdks/python/apache_beam/utils/subprocess_server.py

@@ -76,41 +76,63 @@ def __init__(self, constructor, destructor):
     self._cache = {}
     self._lock = threading.Lock()
     self._counter = 0
+    self._log_buffer = []
 
-  def _next_id(self):
+  def register(self):
+    import traceback
+    stack = "".join(traceback.format_stack())
     with self._lock:
       self._counter += 1
-      return self._counter
-
-  def register(self):
-    owner = self._next_id()
-    self._live_owners.add(owner)
-    _LOGGER.warning("Registered subprocess owner %s (PID %s)", owner, os.getpid(), stack_info=True)
+      owner = self._counter
+      self._live_owners.add(owner)
+      msg = f"Registered subprocess owner {owner} (PID {os.getpid()})"
+      self._log_buffer.append((msg, stack))
     return owner
 
   def purge(self, owner):
     to_delete = []
+    should_flush = False
+    import traceback
+    stack = "".join(traceback.format_stack())
+    
     with self._lock:
       if owner not in self._live_owners:
-        _LOGGER.warning(
-            "Subprocess owner %s (PID %s) already purged. If this occurs during atexit "
-            "shutdown, the subprocess was already cleaned up earlier.",
-            owner,
-            os.getpid(),
-            stack_info=True)
-        return
-      self._live_owners.remove(owner)
-      _LOGGER.warning("Purging subprocess owner %s (PID %s)", owner, os.getpid(), stack_info=True)
-      for key, entry in list(self._cache.items()):
-        if owner in entry.owners:
-          entry.owners.remove(owner)
-        if not entry.owners:
-          to_delete.append(entry.obj)
-          del self._cache[key]
+        msg = (
+            f"Subprocess owner {owner} (PID {os.getpid()}) already purged. "
+            "If this occurs during atexit shutdown, the subprocess was already cleaned up earlier."
+        )
+        self._log_buffer.append((msg, stack))
+        should_flush = True
+      else:
+        self._live_owners.remove(owner)
+        msg = f"Purging subprocess owner {owner} (PID {os.getpid()})"
+        self._log_buffer.append((msg, stack))
+        
+        for key, entry in list(self._cache.items()):
+          if owner in entry.owners:
+            entry.owners.remove(owner)
+          if not entry.owners:
+            to_delete.append(entry.obj)
+            del self._cache[key]
+        
+        if to_delete:
+          should_flush = True
+
+    if should_flush:
+      self._flush_logs()
+
     # Actually call the destructors outside of the lock.
     for value in to_delete:
       self._destructor(value)
 
+  def _flush_logs(self):
+    with self._lock:
+      logs_to_print = list(self._log_buffer)
+      self._log_buffer.clear()
+    
+    for msg, stack in logs_to_print:
+      _LOGGER.warning("%s\nStack Trace:\n%s", msg, stack)
+
   def get(self, *key):
     if not self._live_owners:
       raise RuntimeError("At least one owner must be registered.")