feat(events): improve events kafka pool & client retry

* Check consumers on access from kafka consumer pool
* Allow infinite retries from python actions consumer
* Locking fixes, race condition fixes with consumer pool

Author: david-leifker

datahub-actions/src/datahub_actions/plugin/source/acryl/datahub_cloud_event_source.py

@@ -62,6 +62,7 @@ class DataHubEventsSourceConfig(ConfigModel):
     consumer_id: Optional[str] = None  # Used to store offset for the consumer.
     lookback_days: Optional[int] = None
     reset_offsets: Optional[bool] = False
+    infinite_retry: Optional[bool] = False
 
     # Time and Exit Conditions.
     kill_after_idle_timeout: bool = False
@@ -106,6 +107,7 @@ def __init__(self, config: DataHubEventsSourceConfig, ctx: PipelineContext):
             graph=self.ctx.graph.graph,
             lookback_days=self.source_config.lookback_days,
             reset_offsets=self.source_config.reset_offsets,
+            infinite_retry=self.source_config.infinite_retry,
         )
 
         self.ack_manager = AckManager()
@@ -120,6 +122,7 @@ def _initialize_topic_consumers(
         graph: DataHubGraph,
         lookback_days: Optional[int],
         reset_offsets: Optional[bool],
+        infinite_retry: Optional[bool],
     ) -> Dict[str, DataHubEventsConsumer]:
         """
         Initialize DataHub consumers for each topic with appropriate consumer IDs.
@@ -156,6 +159,7 @@ def _initialize_topic_consumers(
                 consumer_id=topic_consumer_id,
                 lookback_days=lookback_days,
                 reset_offsets=reset_offsets,
+                infinite_retry=infinite_retry,
             )
 
         return topic_consumers

datahub-actions/src/datahub_actions/plugin/source/acryl/datahub_cloud_events_consumer.py

@@ -11,9 +11,10 @@
     Timeout,
 )
 from tenacity import (
-    retry,
+    Retrying,
     retry_if_exception_type,
     stop_after_attempt,
+    stop_never,
     wait_exponential,
 )
 
@@ -49,6 +50,7 @@ def __init__(
         offset_id: Optional[str] = None,
         lookback_days: Optional[int] = None,
         reset_offsets: Optional[bool] = False,
+        infinite_retry: Optional[bool] = False,
     ):
         # 1) Always set self.consumer_id, even if None, so tests can assert it safely.
         self.consumer_id: Optional[str] = consumer_id
@@ -57,6 +59,9 @@ def __init__(
         self.graph: DataHubGraph = graph
         self.offset_id: Optional[str] = offset_id
         self.default_lookback_days: Optional[int] = lookback_days
+        self.infinite_retry: bool = (
+            infinite_retry if infinite_retry is not None else False
+        )
         self.offsets_store: Optional[
             DataHubEventsConsumerPlatformResourceOffsetsStore
         ] = None
@@ -83,14 +88,6 @@ def __init__(
         else:
             logger.debug("Starting DataHub Events Consumer with no consumer ID.")
 
-    @retry(
-        retry=retry_if_exception_type(
-            (HTTPError, ConnectionError, ChunkedEncodingError, Timeout)
-        ),
-        wait=wait_exponential(multiplier=1, min=2, max=30),
-        stop=stop_after_attempt(3),
-        reraise=True,
-    )
     def poll_events(
         self,
         topic: str,
@@ -101,6 +98,36 @@ def poll_events(
         """
         Fetch events for a specific topic.
         """
+        stop_condition = stop_never if self.infinite_retry else stop_after_attempt(15)
+
+        retry_strategy = Retrying(
+            retry=retry_if_exception_type(
+                (HTTPError, ConnectionError, ChunkedEncodingError, Timeout)
+            ),
+            wait=wait_exponential(multiplier=1, min=2, max=60),
+            stop=stop_condition,
+            reraise=True,
+        )
+
+        for attempt in retry_strategy:
+            with attempt:
+                return self._poll_events_impl(
+                    topic, offset_id, limit, poll_timeout_seconds
+                )
+
+        # This should never be reached due to reraise=True, but mypy needs it
+        raise RuntimeError("Retry strategy exhausted without returning or raising")
+
+    def _poll_events_impl(
+        self,
+        topic: str,
+        offset_id: Optional[str] = None,
+        limit: Optional[int] = None,
+        poll_timeout_seconds: Optional[int] = None,
+    ) -> ExternalEventsResponse:
+        """
+        Internal implementation of poll_events.
+        """
         endpoint = f"{self.base_url}/v1/events/poll"
 
         # If the caller provided an offset_id, use it; otherwise fall back to self.offset_id.

datahub-actions/tests/unit/plugin/source/acryl/test_datahub_cloud_events_consumer.py

@@ -200,12 +200,12 @@ def test_poll_events_http_error(mock_graph: DataHubGraph) -> None:
     with patch(
         "requests.get", side_effect=HTTPError(response=dummy_response)
     ) as mock_get:
-        # The default Tenacity stop_after_attempt=3
+        # The default Tenacity stop_after_attempt=15
         with pytest.raises(HTTPError):
             consumer.poll_events(topic="TestTopic")
 
         # requests.get should be called multiple times due to retry
-        assert mock_get.call_count == 3
+        assert mock_get.call_count == 15
 
 
 def test_poll_events_connection_error(mock_graph: DataHubGraph) -> None:
@@ -225,7 +225,7 @@ def test_poll_events_connection_error(mock_graph: DataHubGraph) -> None:
             consumer.poll_events(topic="TestTopic")
 
         # requests.get should be called multiple times due to retry
-        assert mock_get.call_count == 3
+        assert mock_get.call_count == 15
 
 
 def test_poll_events_chunked_encoding_error(mock_graph: DataHubGraph) -> None:
@@ -245,7 +245,7 @@ def test_poll_events_chunked_encoding_error(mock_graph: DataHubGraph) -> None:
             consumer.poll_events(topic="TestTopic")
 
         # requests.get should be called multiple times due to retry
-        assert mock_get.call_count == 3
+        assert mock_get.call_count == 15
 
 
 def test_poll_events_timeout(mock_graph: DataHubGraph) -> None:
@@ -263,7 +263,7 @@ def test_poll_events_timeout(mock_graph: DataHubGraph) -> None:
             consumer.poll_events(topic="TestTopic")
 
         # requests.get should be called multiple times due to retry
-        assert mock_get.call_count == 3
+        assert mock_get.call_count == 15
 
 
 def test_get_events(
@@ -324,3 +324,143 @@ def test_commit_offsets_no_store(mock_graph: DataHubGraph) -> None:
     # No error should occur, and no calls to the store.
     # We just assert that the code doesn't raise an exception.
     # (No additional assert needed here.)
+
+
+def test_poll_events_infinite_retry_retries_more_than_default(
+    mock_graph: DataHubGraph,
+) -> None:
+    """
+    Test that when infinite_retry=True, poll_events retries more than the default 15 times.
+    """
+    consumer = DataHubEventsConsumer(
+        graph=mock_graph,
+        consumer_id="test-consumer",
+        offset_id="initial-offset",
+        infinite_retry=True,
+    )
+
+    dummy_response = Response()
+    # Fail 5 times, then succeed on the 6th attempt
+    call_count = 0
+
+    def side_effect(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count < 6:
+            raise HTTPError(response=dummy_response)
+        # Return a successful response on the 6th attempt
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "offsetId": "success-offset",
+            "count": 0,
+            "events": [],
+        }
+        mock_response.raise_for_status.return_value = None
+        return mock_response
+
+    with patch("requests.get", side_effect=side_effect) as mock_get:
+        result = consumer.poll_events(topic="TestTopic")
+
+        # Should have been called 6 times (5 failures + 1 success)
+        assert mock_get.call_count == 6
+        assert result.offsetId == "success-offset"
+
+
+def test_poll_events_infinite_retry_false_uses_default_retries(
+    mock_graph: DataHubGraph,
+) -> None:
+    """
+    Test that when infinite_retry=False (default), poll_events only retries 15 times.
+    """
+    consumer = DataHubEventsConsumer(
+        graph=mock_graph,
+        consumer_id="test-consumer",
+        offset_id="initial-offset",
+        infinite_retry=False,
+    )
+
+    dummy_response = Response()
+    with patch(
+        "requests.get", side_effect=HTTPError(response=dummy_response)
+    ) as mock_get:
+        with pytest.raises(HTTPError):
+            consumer.poll_events(topic="TestTopic")
+
+        # Should only retry 15 times (default behavior)
+        assert mock_get.call_count == 15
+
+
+def test_poll_events_infinite_retry_connection_error(
+    mock_graph: DataHubGraph,
+) -> None:
+    """
+    Test that infinite_retry works with ConnectionError and retries more than default (15).
+    """
+    consumer = DataHubEventsConsumer(
+        graph=mock_graph,
+        consumer_id="test-consumer",
+        offset_id="initial-offset",
+        infinite_retry=True,
+    )
+
+    # Fail 4 times, then succeed
+    call_count = 0
+
+    def side_effect(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count < 5:
+            raise ConnectionError("Connection Error")
+        # Return a successful response on the 5th attempt
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "offsetId": "success-offset",
+            "count": 0,
+            "events": [],
+        }
+        mock_response.raise_for_status.return_value = None
+        return mock_response
+
+    with patch("requests.get", side_effect=side_effect) as mock_get:
+        result = consumer.poll_events(topic="TestTopic")
+
+        # Should have been called 5 times (4 failures + 1 success)
+        assert mock_get.call_count == 5
+        assert result.offsetId == "success-offset"
+
+
+def test_poll_events_infinite_retry_timeout(mock_graph: DataHubGraph) -> None:
+    """
+    Test that infinite_retry works with Timeout and retries more than default.
+    """
+    consumer = DataHubEventsConsumer(
+        graph=mock_graph,
+        consumer_id="test-consumer",
+        offset_id="initial-offset",
+        infinite_retry=True,
+    )
+
+    # Fail 7 times, then succeed
+    call_count = 0
+
+    def side_effect(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if call_count < 8:
+            raise Timeout("Request Timeout")
+        # Return a successful response on the 8th attempt
+        mock_response = MagicMock()
+        mock_response.json.return_value = {
+            "offsetId": "success-offset",
+            "count": 0,
+            "events": [],
+        }
+        mock_response.raise_for_status.return_value = None
+        return mock_response
+
+    with patch("requests.get", side_effect=side_effect) as mock_get:
+        result = consumer.poll_events(topic="TestTopic")
+
+        # Should have been called 8 times (7 failures + 1 success)
+        assert mock_get.call_count == 8
+        assert result.offsetId == "success-offset"

docs/actions/sources/datahub-cloud-event-source.md

@@ -77,6 +77,7 @@ source:
     topics: ["PlatformEvent_v1"] # Add MetadataChangeLog_Versioned_v1 and / or MetadataChangeLog_Timeseries_v1 to generate raw MCL events.
     lookback_days: 7 # Look back 7 days for events
     reset_offsets: true # Ignore stored offsets and start fresh
+    infinite_retry: true # Enable infinite retry for connection failures (default: false)
     kill_after_idle_timeout: true # Enable shutdown after idle period
     idle_timeout_duration_seconds: 60 # Idle timeout set to 60 seconds
     event_processing_time_max_duration_seconds: 45 # Max processing time of 45 seconds per batch
@@ -94,6 +95,7 @@ Note that the `datahub` configuration block is **required** to connect to your D
 | `topics`                                     |    ❌    | `PlatformEvent_v1` | The name of the topic from which events will be consumed. By default only produces `EntityChangeEvent_v1` events. To include `MetadataChangeLog_v1` events, set this value to include ["MetadataChangeLog_Versioned_v1", "MetadataChangeLog_Timeseries_v1"] |
 | `lookback_days`                              |    ❌    |        None        | Optional number of days to look back when polling for events.                                                                                                                                                                                               |
 | `reset_offsets`                              |    ❌    |      `False`       | When set to `True`, the consumer will ignore any stored offsets and start fresh.                                                                                                                                                                            |
+| `infinite_retry`                             |    ❌    |      `False`       | When set to `True`, the consumer will retry indefinitely on connection failures (HTTPError, ConnectionError, ChunkedEncodingError, Timeout) with exponential backoff (2s to 60s). When `False` (default), it retries up to 15 times before failing.         |
 | `kill_after_idle_timeout`                    |    ❌    |      `False`       | If `True`, stops the consumer after being idle for the specified timeout duration.                                                                                                                                                                          |
 | `idle_timeout_duration_seconds`              |    ❌    |        `30`        | Duration in seconds after which, if no events are received, the consumer is considered idle.                                                                                                                                                                |
 | `event_processing_time_max_duration_seconds` |    ❌    |        `30`        | Maximum allowed time in seconds for processing events before timing out.                                                                                                                                                                                    |
@@ -110,3 +112,7 @@ Yes, simply set `reset_offsets` to True for a single run of the action. Remember
 
 Today, there is undefined behavior deploying multiple actions with the same name using the DataHub Cloud Events Source.
 All events must be processed by a single running action
+
+3. How do I handle transient connection failures?
+
+By default, the consumer will retry up to 15 times on connection failures (HTTPError, ConnectionError, ChunkedEncodingError, Timeout) with exponential backoff (2s to 60s). If you expect longer outages or want to ensure the consumer continues retrying indefinitely, set `infinite_retry: true` in the source configuration. When enabled, the consumer will retry indefinitely with exponential backoff (starting at 2 seconds, increasing up to 60 seconds, then continuing at 60 seconds) until the connection is restored.

metadata-service/configuration/src/main/resources/application.yaml

@@ -583,6 +583,8 @@ kafka:
   consumerPool:
     initialSize: ${KAFKA_CONSUMER_POOL_INITIAL_SIZE:1}
     maxSize: ${KAFKA_CONSUMER_POOL_MAX_SIZE:5}
+    validationTimeoutSeconds: ${KAFKA_CONSUMER_POOL_VALIDATION_TIMEOUT_SECONDS:5}
+    validationCacheIntervalMinutes: ${KAFKA_CONSUMER_POOL_VALIDATION_CACHE_INTERVAL_MINUTES:5}
 
   schemaRegistry:
     type: ${SCHEMA_REGISTRY_TYPE:KAFKA} # INTERNAL or KAFKA or AWS_GLUE

metadata-service/events-service/build.gradle

@@ -9,6 +9,7 @@ dependencies {
 
     implementation externalDependency.kafkaClients
     implementation externalDependency.kafkaAvroSerde
+    implementation project(':metadata-utils')
 
     testImplementation externalDependency.mockito
     testImplementation externalDependency.testng