Fixed DirectRunner PubSub subscriber client lifecycle (#39079)

* Fix DirectRunner PubSub subscriber client cleanup

* Avoid global lock during PubSub subscription creation

Author: lalitium

sdks/python/apache_beam/io/gcp/pubsub_test.py

@@ -502,6 +502,70 @@ def finish_bundle(self):
 @unittest.skipIf(pubsub is None, 'GCP dependencies are not installed')
 @mock.patch('google.cloud.pubsub.SubscriberClient')
 class TestReadFromPubSub(unittest.TestCase):
+  def setUp(self):
+    _PubSubReadEvaluator._subscription_cache.clear()
+    _PubSubReadEvaluator._subscriber_client_cache.clear()
+
+  def test_subscriber_client_is_reused_for_transform(self, mock_pubsub):
+    class Transform(object):
+      pass
+
+    transform = Transform()
+    first_client = _PubSubReadEvaluator._get_subscriber_client(transform)
+    second_client = _PubSubReadEvaluator._get_subscriber_client(transform)
+
+    self.assertIs(first_client, second_client)
+    mock_pubsub.assert_called_once_with()
+    first_client.close.assert_not_called()
+
+  def test_subscription_creation_does_not_hold_global_cache_lock(
+      self, mock_pubsub):
+    class Transform(object):
+      pass
+
+    def subscription_path(project, subscription):
+      return 'projects/%s/subscriptions/%s' % (project, subscription)
+
+    transform = Transform()
+    client = mock_pubsub.return_value
+    client.subscription_path.side_effect = subscription_path
+    client.topic_path.return_value = 'projects/topic_project/topics/topic'
+    global_lock_available = []
+
+    def create_subscription(name, topic):
+      self.assertTrue(name.startswith('projects/sub_project/subscriptions/'))
+      self.assertEqual('projects/topic_project/topics/topic', topic)
+      acquired = _PubSubReadEvaluator._subscriber_client_cache_lock.acquire(
+          blocking=False)
+      global_lock_available.append(acquired)
+      if acquired:
+        _PubSubReadEvaluator._subscriber_client_cache_lock.release()
+
+    client.create_subscription.side_effect = create_subscription
+
+    sub_name = _PubSubReadEvaluator.get_subscription(
+        transform, 'topic_project', 'topic', 'sub_project', None)
+
+    self.assertTrue(global_lock_available)
+    self.assertTrue(global_lock_available[0])
+    self.assertTrue(sub_name.startswith('projects/sub_project/subscriptions/'))
+
+  def test_subscriber_client_cleanup_is_idempotent(self, unused_mock_pubsub):
+    client = mock.Mock()
+    subscriber_client = transform_evaluator._PubSubSubscriberClient(client)
+    subscriber_client.set_temporary_subscription('subscription')
+
+    subscriber_client.close()
+    subscriber_client.close()
+
+    client.assert_has_calls([
+        mock.call.delete_subscription(subscription='subscription'),
+        mock.call.close()
+    ])
+    client.delete_subscription.assert_called_once_with(
+        subscription='subscription')
+    client.close.assert_called_once_with()
+
   def test_read_messages_success(self, mock_pubsub):
     data = b'data'
     publish_time_secs = 1520861821
@@ -533,7 +597,8 @@ def test_read_messages_success(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.assert_called_once_with()
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_strings_success(self, mock_pubsub):
     data = '🤷 ¯\\_(ツ)_/¯'
@@ -555,7 +620,7 @@ def test_read_strings_success(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_data_success(self, mock_pubsub):
     data_encoded = '🤷 ¯\\_(ツ)_/¯'.encode('utf-8')
@@ -575,7 +640,7 @@ def test_read_data_success(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub):
     data = b'data'
@@ -610,7 +675,7 @@ def test_read_messages_timestamp_attribute_milli_success(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub):
     data = b'data'
@@ -645,7 +710,7 @@ def test_read_messages_timestamp_attribute_rfc3339_success(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_messages_timestamp_attribute_missing(self, mock_pubsub):
     data = b'data'
@@ -681,7 +746,7 @@ def test_read_messages_timestamp_attribute_missing(self, mock_pubsub):
     mock_pubsub.return_value.acknowledge.assert_has_calls(
         [mock.call(subscription=mock.ANY, ack_ids=[ack_id])])
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub):
     data = b'data'
@@ -710,7 +775,7 @@ def test_read_messages_timestamp_attribute_fail_parse(self, mock_pubsub):
       p.run()
     mock_pubsub.return_value.acknowledge.assert_not_called()
 
-    mock_pubsub.return_value.close.assert_has_calls([mock.call()])
+    mock_pubsub.return_value.close.assert_not_called()
 
   def test_read_message_id_label_unsupported(self, unused_mock_pubsub):
     # id_label is unsupported in DirectRunner.

sdks/python/apache_beam/runners/direct/transform_evaluator.py

@@ -19,11 +19,12 @@
 
 # pytype: skip-file
 
-import atexit
 import collections
 import logging
 import random
+import threading
 import time
+import weakref
 from collections import abc
 from typing import TYPE_CHECKING
 from typing import Any
@@ -577,13 +578,46 @@ def finish_bundle(self):
         self, self.bundles, unprocessed_bundles, None, {None: self.watermark})
 
 
+class _PubSubSubscriberClient(object):
+  """SubscriberClient state cached for one DirectRunner Pub/Sub read."""
+  def __init__(self, client):
+    self.client = client
+    self._temporary_subscription = None
+    self._closed = False
+    self._lock = threading.Lock()
+
+  def set_temporary_subscription(self, subscription):
+    self._temporary_subscription = subscription
+
+  def close(self):
+    if self._closed:
+      return
+    self._closed = True
+
+    try:
+      if self._temporary_subscription:
+        self.client.delete_subscription(
+            subscription=self._temporary_subscription)
+    except Exception:
+      _LOGGER.warning(
+          'Failed to delete temporary Pub/Sub subscription %s',
+          self._temporary_subscription,
+          exc_info=True)
+
+    try:
+      self.client.close()
+    except Exception:
+      _LOGGER.warning(
+          'Failed to close Pub/Sub subscriber client', exc_info=True)
+
+
 class _PubSubReadEvaluator(_TransformEvaluator):
   """TransformEvaluator for PubSub read."""
 
-  # A mapping of transform to _PubSubSubscriptionWrapper.
-  # TODO(https://github.com/apache/beam/issues/19751): Prevents garbage
-  # collection of pipeline instances.
-  _subscription_cache: dict[AppliedPTransform, str] = {}
+  # Weak-keyed per-transform caches avoid keeping completed pipelines alive.
+  _subscription_cache = weakref.WeakKeyDictionary()
+  _subscriber_client_cache = weakref.WeakKeyDictionary()
+  _subscriber_client_cache_lock = threading.Lock()
 
   def __init__(
       self,
@@ -627,18 +661,46 @@ def get_subscription(
     if short_sub_name:
       return pubsub.SubscriberClient.subscription_path(project, short_sub_name)
 
-    if transform in cls._subscription_cache:
-      return cls._subscription_cache[transform]
+    with cls._subscriber_client_cache_lock:
+      sub_name = cls._subscription_cache.get(transform)
+      if sub_name:
+        return sub_name
 
-    sub_client = pubsub.SubscriberClient()
-    sub_name = sub_client.subscription_path(
-        sub_project,
-        'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32)))
-    topic_name = sub_client.topic_path(project, short_topic_name)
-    sub_client.create_subscription(name=sub_name, topic=topic_name)
-    atexit.register(sub_client.delete_subscription, subscription=sub_name)
-    cls._subscription_cache[transform] = sub_name
-    return cls._subscription_cache[transform]
+      subscriber_client = cls._get_subscriber_client_state_unlocked(transform)
+
+    with subscriber_client._lock:
+      with cls._subscriber_client_cache_lock:
+        sub_name = cls._subscription_cache.get(transform)
+        if sub_name:
+          return sub_name
+
+      sub_client = subscriber_client.client
+      sub_name = sub_client.subscription_path(
+          sub_project,
+          'beam_%d_%x' % (int(time.time()), random.randrange(1 << 32)))
+      topic_name = sub_client.topic_path(project, short_topic_name)
+      sub_client.create_subscription(name=sub_name, topic=topic_name)
+      subscriber_client.set_temporary_subscription(sub_name)
+
+      with cls._subscriber_client_cache_lock:
+        cls._subscription_cache[transform] = sub_name
+
+      return sub_name
+
+  @classmethod
+  def _get_subscriber_client(cls, transform):
+    with cls._subscriber_client_cache_lock:
+      return cls._get_subscriber_client_state_unlocked(transform).client
+
+  @classmethod
+  def _get_subscriber_client_state_unlocked(cls, transform):
+    subscriber_client = cls._subscriber_client_cache.get(transform)
+    if subscriber_client is None:
+      from google.cloud import pubsub
+      subscriber_client = _PubSubSubscriberClient(pubsub.SubscriberClient())
+      cls._subscriber_client_cache[transform] = subscriber_client
+      weakref.finalize(transform, subscriber_client.close)
+    return subscriber_client
 
   def start_bundle(self):
     pass
@@ -648,8 +710,6 @@ def process_element(self, element):
 
   def _read_from_pubsub(
       self, timestamp_attribute) -> list[tuple[Timestamp, 'PubsubMessage']]:
-    from google.cloud import pubsub
-
     from apache_beam.io.gcp.pubsub import PubsubMessage
 
     def _get_element(message):
@@ -678,16 +738,13 @@ def _get_element(message):
     # evaluator fails with an exception before emitting a bundle. However,
     # the DirectRunner currently doesn't retry work items anyway, so the
     # pipeline would enter an inconsistent state on any error.
-    sub_client = pubsub.SubscriberClient()
-    try:
-      response = sub_client.pull(
-          subscription=self._sub_name, max_messages=10, timeout=30)
-      results = [_get_element(rm.message) for rm in response.received_messages]
-      ack_ids = [rm.ack_id for rm in response.received_messages]
-      if ack_ids:
-        sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids)
-    finally:
-      sub_client.close()
+    sub_client = self._get_subscriber_client(self._applied_ptransform)
+    response = sub_client.pull(
+        subscription=self._sub_name, max_messages=10, timeout=30)
+    results = [_get_element(rm.message) for rm in response.received_messages]
+    ack_ids = [rm.ack_id for rm in response.received_messages]
+    if ack_ids:
+      sub_client.acknowledge(subscription=self._sub_name, ack_ids=ack_ids)
 
     return results