Move lock from _next_id to register.

Author: shunping

sdks/python/apache_beam/utils/subprocess_server.py

@@ -76,63 +76,38 @@ def __init__(self, constructor, destructor):
     self._cache = {}
     self._lock = threading.Lock()
     self._counter = 0
-    self._log_buffer = []
+
+  def _next_id(self):
+    # Caller must hold self._lock.
+    self._counter += 1
+    return self._counter
 
   def register(self):
-    import traceback
-    stack = "".join(traceback.format_stack())
     with self._lock:
-      self._counter += 1
-      owner = self._counter
+      owner = self._next_id()
       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:
-        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()
-
+        _LOGGER.warning(
+            "Subprocess owner %s already purged. If this occurs during atexit "
+            "shutdown, the subprocess was already cleaned up earlier.",
+            owner)
+        return
+      self._live_owners.remove(owner)
+      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]
     # 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.")