This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8f60e69177b Fixed DirectRunner PubSub subscriber client lifecycle
(#39079)
8f60e69177b is described below
commit 8f60e69177beef5dbcdc038a98795e4c464fa138
Author: Lalit Yadav <[email protected]>
AuthorDate: Tue Jun 30 10:14:53 2026 -0500
Fixed DirectRunner PubSub subscriber client lifecycle (#39079)
* Fix DirectRunner PubSub subscriber client cleanup
* Avoid global lock during PubSub subscription creation
---
sdks/python/apache_beam/io/gcp/pubsub_test.py | 79 ++++++++++++--
.../runners/direct/transform_evaluator.py | 113 ++++++++++++++++-----
2 files changed, 157 insertions(+), 35 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/pubsub_test.py
b/sdks/python/apache_beam/io/gcp/pubsub_test.py
index c35de62fca1..050a69aff6c 100644
--- a/sdks/python/apache_beam/io/gcp/pubsub_test.py
+++ b/sdks/python/apache_beam/io/gcp/pubsub_test.py
@@ -502,6 +502,70 @@
transform_evaluator.TransformEvaluatorRegistry._test_evaluators_overrides = {
@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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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 @@ class TestReadFromPubSub(unittest.TestCase):
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.
diff --git a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
b/sdks/python/apache_beam/runners/direct/transform_evaluator.py
index 2349a0881e4..6702ec3362b 100644
--- a/sdks/python/apache_beam/runners/direct/transform_evaluator.py
+++ b/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 @@ class _TestStreamEvaluator(_TransformEvaluator):
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 @@ class _PubSubReadEvaluator(_TransformEvaluator):
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 @@ class _PubSubReadEvaluator(_TransformEvaluator):
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 @@ class _PubSubReadEvaluator(_TransformEvaluator):
# 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
