[Python] Bound UnboundedSource read per bundle by records and time

A continuously busy reader stayed in process() forever: the loop only
broke at EOF or when data ran out, so a hot source never checkpointed,
finalized, or committed its watermark on the Fn API path. Cap the
has-data path at max_records_per_bundle (10000) or max_read_time_seconds
(10), then self-checkpoint via defer_remainder so the runner commits
progress and resumes immediately. Both caps are configurable on
ReadFromUnboundedSource.

Author: Eliaaazzz

sdks/python/apache_beam/io/unbounded_source.py

@@ -86,6 +86,7 @@ def get_checkpoint_mark_coder(self):
 import dataclasses
 import logging
 import threading
+import time
 from typing import Any
 from typing import Iterable
 from typing import Optional
@@ -123,6 +124,8 @@ def get_checkpoint_mark_coder(self):
 
 _DEFAULT_POLL_INTERVAL_SECONDS = 1.0
 _DEFAULT_DESIRED_NUM_SPLITS = 20
+_DEFAULT_MAX_RECORDS_PER_BUNDLE = 10000
+_DEFAULT_MAX_READ_TIME_SECONDS = 10.0
 
 # ------------------------------------------------------------------------------
 # Public abstract base classes.
@@ -613,8 +616,14 @@ class _ReadFromUnboundedSourceDoFn(core.DoFn):
   Module-level so stdlib pickle and cloudpickle can serialise the DoFn. The
   restriction provider is the module-level :data:`_PROVIDER` singleton.
   """
-  def __init__(self, poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS):
+  def __init__(
+      self,
+      poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS,
+      max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE,
+      max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS):
     self._poll_interval = poll_interval
+    self._max_records_per_bundle = max_records_per_bundle
+    self._max_read_time_seconds = max_read_time_seconds
 
   @core.DoFn.unbounded_per_element()
   def process(
@@ -628,6 +637,8 @@ def process(
     # kwarg-injected ones (tracker, watermark estimator).
     assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
     initial = tracker.current_restriction()
+    records_emitted = 0
+    read_deadline = time.monotonic() + self._max_read_time_seconds
     try:
       while True:
         holder = [None]
@@ -639,10 +650,8 @@ def process(
           break
         record = holder[0]
         if record is _NO_DATA:
-          # No data now: advance the watermark and self-checkpoint with an
-          # explicit delay. defer_remainder(None) resumes ASAP and would
-          # busy-loop an idle source, so (unlike Java's runner-applied backoff)
-          # we set the resume delay here.
+          # No data now: advance the watermark and self-checkpoint with the
+          # poll delay so an idle source backs off before resuming.
           _set_watermark_if_greater(
               watermark_estimator, tracker.current_restriction().watermark)
           tracker.defer_remainder(Duration(seconds=self._poll_interval))
@@ -652,6 +661,15 @@ def process(
         value, record_timestamp, source_watermark = record
         _set_watermark_if_greater(watermark_estimator, source_watermark)
         yield TimestampedValue(value, record_timestamp)
+        records_emitted += 1
+        # A busy reader never reaches the EOF or no-data branch, so bound the
+        # bundle by record count and elapsed time. The checkpoint lets the
+        # runner commit progress and run finalization. Resume with no delay
+        # because data is still available.
+        if (records_emitted >= self._max_records_per_bundle or
+            time.monotonic() >= read_deadline):
+          tracker.defer_remainder()
+          break
     finally:
       current = tracker.current_restriction()
       try:
@@ -695,17 +713,23 @@ class ReadFromUnboundedSource(PTransform):
       p | beam.io.Read(MyUnboundedSource())
 
   ``poll_interval`` is the resume delay used when the reader has no data, which
-  bounds how often an idle source is polled.
+  bounds how often an idle source is polled. ``max_records_per_bundle`` and
+  ``max_read_time_seconds`` bound how much a busy reader produces before the
+  bundle self-checkpoints.
   """
   def __init__(
       self,
       source: UnboundedSource,
-      poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS):
+      poll_interval: float = _DEFAULT_POLL_INTERVAL_SECONDS,
+      max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE,
+      max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS):
     if not isinstance(source, UnboundedSource):
       raise TypeError('source must be an UnboundedSource, got %r' % (source, ))
     super().__init__()
     self._source = source
     self._poll_interval = poll_interval
+    self._max_records_per_bundle = max_records_per_bundle
+    self._max_read_time_seconds = max_read_time_seconds
 
   def expand(self, pbegin):
     source = self._source
@@ -717,7 +741,10 @@ def expand(self, pbegin):
         pbegin
         | 'Create' >> core.Create([source])
         | 'ReadUnbounded' >> core.ParDo(
-            _ReadFromUnboundedSourceDoFn(self._poll_interval)))
+            _ReadFromUnboundedSourceDoFn(
+                self._poll_interval,
+                self._max_records_per_bundle,
+                self._max_read_time_seconds)))
     # Surface an element type only when the global registry already maps it to
     # an equivalent coder. Avoid mutating ``coders.registry`` for a
     # parameterized coder whose instance state would be lost.

sdks/python/apache_beam/io/unbounded_source_test.py

@@ -35,6 +35,7 @@
 from apache_beam.io.unbounded_source import UnboundedReader
 from apache_beam.io.unbounded_source import UnboundedSource
 from apache_beam.io.unbounded_source import _FinalizeCheckpointOnce
+from apache_beam.io.unbounded_source import _ReadFromUnboundedSourceDoFn
 from apache_beam.io.unbounded_source import _set_watermark_if_greater
 from apache_beam.io.unbounded_source import _UnboundedSourceRestriction
 from apache_beam.io.unbounded_source import _UnboundedSourceRestrictionCoder
@@ -560,6 +561,66 @@ def test_is_bounded_false(self):
     self.assertFalse(_new_tracker(UnboundedCountingSource(3)).is_bounded())
 
 
+class _RecordingBundleFinalizer:
+  def __init__(self):
+    self.registered = []
+
+  def register(self, callback):
+    self.registered.append(callback)
+
+
+class BundleCapTest(unittest.TestCase):
+  """A busy reader self-checkpoints once the per-bundle record or time cap is
+  reached, so the runner can commit progress and run finalization."""
+  def _run_process(self, dofn, source):
+    tracker = _UnboundedSourceRestrictionTracker(
+        _UnboundedSourceRestriction(source=source))
+    threadsafe = sdf_utils.ThreadsafeRestrictionTracker(tracker)
+    view = sdf_utils.RestrictionTrackerView(threadsafe)
+    outputs = list(
+        dofn.process(
+            None,
+            bundle_finalizer=_RecordingBundleFinalizer(),
+            tracker=view,
+            watermark_estimator=ManualWatermarkEstimator(None)))
+    return outputs, tracker, threadsafe
+
+  def test_record_cap_checkpoints_busy_source(self):
+    dofn = _ReadFromUnboundedSourceDoFn(
+        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+    # 1000 records is effectively unbounded against a cap of 5.
+    outputs, tracker, threadsafe = self._run_process(
+        dofn, UnboundedCountingSource(1000))
+
+    self.assertEqual([tv.value for tv in outputs], [0, 1, 2, 3, 4])
+    self.assertTrue(tracker.current_restriction().is_done)
+    self.assertTrue(tracker.check_done())
+    residual, _delay = threadsafe.deferred_status()
+    self.assertEqual(residual.checkpoint_mark.last_index, 4)
+
+  def test_time_cap_checkpoints_busy_source(self):
+    dofn = _ReadFromUnboundedSourceDoFn(
+        poll_interval=0, max_records_per_bundle=10**9, max_read_time_seconds=0)
+    outputs, tracker, threadsafe = self._run_process(
+        dofn, UnboundedCountingSource(1000))
+
+    # A zero time budget trips the deadline right after the first record.
+    self.assertEqual([tv.value for tv in outputs], [0])
+    self.assertTrue(tracker.current_restriction().is_done)
+    residual, _delay = threadsafe.deferred_status()
+    self.assertEqual(residual.checkpoint_mark.last_index, 0)
+
+  def test_eof_before_cap_finishes_without_residual(self):
+    dofn = _ReadFromUnboundedSourceDoFn(
+        poll_interval=0, max_records_per_bundle=100, max_read_time_seconds=1e9)
+    outputs, tracker, threadsafe = self._run_process(
+        dofn, UnboundedCountingSource(3))
+
+    self.assertEqual([tv.value for tv in outputs], [0, 1, 2])
+    self.assertTrue(tracker.current_restriction().is_done)
+    self.assertIsNone(threadsafe.deferred_status())
+
+
 class WatermarkTest(unittest.TestCase):
   def test_set_watermark_is_monotonic(self):
     estimator = ManualWatermarkEstimator(None)