[Python] Refine UnboundedSource read-bound caps after review

Validate the caps at the transform boundary (max_records_per_bundle >= 1,
max_read_time_seconds > 0, poll_interval >= 0). Arm the read deadline on the
first emitted record so reader startup is excluded, and add an injectable
clock seam for deterministic time-cap tests. Document the between-records
time-cap limitation, the finalize idempotency-on-retry contract, and the
advisory residual watermark. Expand BundleCapTest with cross-bundle resume,
EOF-exactly-at-cap, finalize registration, watermark, drain-truncate, and
cap-validation coverage.

Author: Eliaaazzz

sdks/python/apache_beam/io/unbounded_source.py

@@ -88,6 +88,7 @@ def get_checkpoint_mark_coder(self):
 import threading
 import time
 from typing import Any
+from typing import Callable
 from typing import Iterable
 from typing import Optional
 
@@ -147,7 +148,9 @@ def finalize_checkpoint(self) -> None:
 
     The runner calls this at most once for a committed checkpoint mark.
     Finalization is best effort; a mark may never be finalized. An exception
-    raised here is logged.
+    raised here is logged. On bundle retry an uncommitted mark may be re-cut
+    over an overlapping span, so this method must be idempotent (acknowledge by
+    absolute position).
     """
     pass
 
@@ -468,6 +471,8 @@ def _try_split_inner(self, fraction_of_remainder):
     if self._reader is None or not self._started or self._restriction.is_done:
       return None
     checkpoint = self._reader.get_checkpoint_mark()
+    # The residual watermark is advisory; the SDF watermark estimator state is
+    # the authoritative cross-bundle watermark.
     watermark = self._reader.get_watermark()
     # Keep the two channels independent: the primary carries only the finalize
     # hook, the residual only the resume state.
@@ -593,6 +598,8 @@ def truncate(self, element, restriction):
 class _FinalizeCheckpointOnce(object):
   def __init__(self, checkpoint_mark: CheckpointMark):
     self._checkpoint_mark = checkpoint_mark
+    # The lock keeps finalization idempotent if a runner ever invokes the
+    # callback more than once.
     self._lock = threading.Lock()
     self._finalized = False
 
@@ -620,10 +627,13 @@ 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):
+      max_read_time_seconds: float = _DEFAULT_MAX_READ_TIME_SECONDS,
+      _now: Optional[Callable[[], float]] = None):
     self._poll_interval = poll_interval
     self._max_records_per_bundle = max_records_per_bundle
     self._max_read_time_seconds = max_read_time_seconds
+    # Monotonic clock seam; tests inject a deterministic clock.
+    self._now = _now
 
   @core.DoFn.unbounded_per_element()
   def process(
@@ -637,8 +647,10 @@ def process(
     # kwarg-injected ones (tracker, watermark estimator).
     assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
     initial = tracker.current_restriction()
+    now = self._now or time.monotonic
     records_emitted = 0
-    read_deadline = time.monotonic() + self._max_read_time_seconds
+    # Armed on the first emitted record so reader startup is excluded.
+    read_deadline = None  # type: Optional[float]
     try:
       while True:
         holder = [None]
@@ -662,18 +674,23 @@ def process(
         _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):
+        if read_deadline is None:
+          read_deadline = now() + self._max_read_time_seconds
+        # A busy reader never hits the EOF or no-data branch. Bound the bundle
+        # by record count and elapsed time so the runner commits the checkpoint
+        # and runs finalization, then resume with no delay. The deadline is
+        # checked between records; a reader that blocks inside advance() can
+        # overrun it, so the record cap is the hard backstop.
+        reached_record_cap = records_emitted >= self._max_records_per_bundle
+        if reached_record_cap or now() >= read_deadline:
           tracker.defer_remainder()
           break
     finally:
       current = tracker.current_restriction()
       try:
         # Register finalization only when a checkpoint was cut this bundle.
+        # The SDK bundle finalizer applies no deadline, so finalization is
+        # unbounded best effort.
         finalize_mark = current.finalization_checkpoint_mark
         if current is not initial and finalize_mark is not None:
           bundle_finalizer.register(_FinalizeCheckpointOnce(finalize_mark))
@@ -712,10 +729,18 @@ 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. ``max_records_per_bundle`` and
-  ``max_read_time_seconds`` bound how much a busy reader produces before the
-  bundle self-checkpoints.
+  Args:
+    source: the :class:`UnboundedSource` to read.
+    poll_interval: resume delay in seconds applied when the reader has no data,
+      which bounds how often an idle source is polled. Must be >= 0.
+    max_records_per_bundle: a busy reader self-checkpoints after emitting this
+      many records in one bundle. Must be >= 1. Defaults to 10000.
+    max_read_time_seconds: a busy reader self-checkpoints after this many
+      seconds in one bundle. Must be > 0. Defaults to 10.0. The deadline is
+      checked between records, so a reader that blocks inside ``advance()`` may
+      overrun it; ``max_records_per_bundle`` is the hard backstop.
+
+  The bundle self-checkpoints as soon as either cap is reached.
   """
   def __init__(
       self,
@@ -725,6 +750,17 @@ def __init__(
       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, ))
+    if max_records_per_bundle < 1:
+      raise ValueError(
+          'max_records_per_bundle must be >= 1, got %r' %
+          (max_records_per_bundle, ))
+    if max_read_time_seconds <= 0:
+      raise ValueError(
+          'max_read_time_seconds must be > 0, got %r' %
+          (max_read_time_seconds, ))
+    if poll_interval < 0:
+      raise ValueError(
+          'poll_interval must be >= 0, got %r' % (poll_interval, ))
     super().__init__()
     self._source = source
     self._poll_interval = poll_interval

sdks/python/apache_beam/io/unbounded_source_test.py

@@ -419,6 +419,13 @@ def split(self, desired_num_splits, options=None):
 
     self.assertEqual(list(provider.split(source, restriction)), [restriction])
 
+  def test_truncate_returns_none_for_drain(self):
+    # On drain the SDF stops emitting; truncate yields no residual.
+    provider = _UnboundedSourceRestrictionProvider()
+    source = UnboundedCountingSource(5)
+    restriction = _UnboundedSourceRestriction(source=source)
+    self.assertIsNone(provider.truncate(source, restriction))
+
   def test_splittable_source_partitions_into_independent_subsources(self):
     # A splittable source fans out into two sub-sources; reading each in
     # isolation yields the even and the odd integers, and their union is the
@@ -569,54 +576,122 @@ def register(self, callback):
     self.registered.append(callback)
 
 
+class _ManualClock:
+  """A deterministic monotonic clock for the time-cap tests."""
+  def __init__(self, now=0.0):
+    self.now = now
+
+  def __call__(self):
+    return self.now
+
+
 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):
+  def _bundle(self, dofn, source, checkpoint=None, estimator=None):
+    """Builds the SDF tracker chain and returns the process() generator plus the
+    tracker, threadsafe tracker, finalizer, and watermark estimator."""
     tracker = _UnboundedSourceRestrictionTracker(
-        _UnboundedSourceRestriction(source=source))
+        _UnboundedSourceRestriction(source=source, checkpoint_mark=checkpoint))
     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
+    finalizer = _RecordingBundleFinalizer()
+    estimator = estimator or ManualWatermarkEstimator(None)
+    gen = dofn.process(
+        None,
+        bundle_finalizer=finalizer,
+        tracker=view,
+        watermark_estimator=estimator)
+    return gen, tracker, threadsafe, finalizer, estimator
 
   def test_record_cap_checkpoints_busy_source(self):
+    finalize_log = []
     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))
+    gen, tracker, threadsafe, finalizer, estimator = self._bundle(
+        dofn, UnboundedCountingSource(1000, finalize_log=finalize_log))
+    outputs = list(gen)
 
     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()
+    # The estimator holds the last emitted record's source watermark.
+    self.assertEqual(estimator.current_watermark(), _EVENT_TIME_BASE + 4)
+    # Residual resumes after the cut and carries no finalize hook.
+    residual, _ = threadsafe.deferred_status()
     self.assertEqual(residual.checkpoint_mark.last_index, 4)
+    self.assertIsNone(residual.finalization_checkpoint_mark)
+    # Exactly one finalizer is registered; firing it commits the cut index once.
+    self.assertEqual(len(finalizer.registered), 1)
+    finalizer.registered[0]()
+    finalizer.registered[0]()
+    self.assertEqual(finalize_log, [4])
 
   def test_time_cap_checkpoints_busy_source(self):
+    clock = _ManualClock(1000.0)
     dofn = _ReadFromUnboundedSourceDoFn(
-        poll_interval=0, max_records_per_bundle=10**9, max_read_time_seconds=0)
-    outputs, tracker, threadsafe = self._run_process(
+        poll_interval=0,
+        max_records_per_bundle=10**9,
+        max_read_time_seconds=5.0,
+        _now=clock)
+    gen, tracker, threadsafe, _, _ = self._bundle(
         dofn, UnboundedCountingSource(1000))
 
-    # A zero time budget trips the deadline right after the first record.
-    self.assertEqual([tv.value for tv in outputs], [0])
+    # The deadline arms at 1000 + 5 after the first record and is checked
+    # between records, so records keep flowing until the clock passes it.
+    self.assertEqual(next(gen).value, 0)
+    self.assertEqual(next(gen).value, 1)
+    self.assertEqual(next(gen).value, 2)
+    clock.now = 1006.0
+    with self.assertRaises(StopIteration):
+      next(gen)
+
     self.assertTrue(tracker.current_restriction().is_done)
-    residual, _delay = threadsafe.deferred_status()
-    self.assertEqual(residual.checkpoint_mark.last_index, 0)
+    residual, _ = threadsafe.deferred_status()
+    self.assertEqual(residual.checkpoint_mark.last_index, 2)
+
+  def test_cap_residual_resumes_in_next_bundle(self):
+    dofn = _ReadFromUnboundedSourceDoFn(
+        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+    source = UnboundedCountingSource(1000)
+    # Bundle 1 emits 0-4 and cuts a residual at index 4.
+    gen1, _, threadsafe1, _, _ = self._bundle(dofn, source)
+    self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
+    residual1, _ = threadsafe1.deferred_status()
+
+    # Bundle 2 rebuilds the reader from the residual and emits 5-9.
+    gen2, _, threadsafe2, _, _ = self._bundle(
+        dofn, source, checkpoint=residual1.checkpoint_mark)
+    self.assertEqual([tv.value for tv in gen2], [5, 6, 7, 8, 9])
+    residual2, _ = threadsafe2.deferred_status()
+    self.assertEqual(residual2.checkpoint_mark.last_index, 9)
+
+  def test_eof_exactly_at_cap_resumes_then_finishes(self):
+    dofn = _ReadFromUnboundedSourceDoFn(
+        poll_interval=0, max_records_per_bundle=5, max_read_time_seconds=1e9)
+    source = UnboundedCountingSource(5)  # exactly cap records
+    # Bundle 1 hits the cap on the last record before observing EOF.
+    gen1, t1, threadsafe1, _, _ = self._bundle(dofn, source)
+    self.assertEqual([tv.value for tv in gen1], [0, 1, 2, 3, 4])
+    self.assertTrue(t1.current_restriction().is_done)
+    residual1, _ = threadsafe1.deferred_status()
+    self.assertEqual(residual1.checkpoint_mark.last_index, 4)
+
+    # Bundle 2 resumes at index 5, finds EOF, and finishes with no output.
+    gen2, t2, threadsafe2, _, _ = self._bundle(
+        dofn, source, checkpoint=residual1.checkpoint_mark)
+    self.assertEqual(list(gen2), [])
+    self.assertTrue(t2.current_restriction().is_done)
+    self.assertIsNone(threadsafe2.deferred_status())
 
   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(
+    gen, tracker, threadsafe, _, _ = self._bundle(
         dofn, UnboundedCountingSource(3))
 
-    self.assertEqual([tv.value for tv in outputs], [0, 1, 2])
+    self.assertEqual([tv.value for tv in gen], [0, 1, 2])
     self.assertTrue(tracker.current_restriction().is_done)
     self.assertIsNone(threadsafe.deferred_status())
 
@@ -1022,6 +1097,15 @@ def test_non_source_argument_raises(self):
     with self.assertRaises(TypeError):
       ReadFromUnboundedSource('not-a-source')  # type: ignore[arg-type]
 
+  def test_invalid_caps_raise(self):
+    source = UnboundedCountingSource(1)
+    with self.assertRaises(ValueError):
+      ReadFromUnboundedSource(source, max_records_per_bundle=0)
+    with self.assertRaises(ValueError):
+      ReadFromUnboundedSource(source, max_read_time_seconds=0)
+    with self.assertRaises(ValueError):
+      ReadFromUnboundedSource(source, poll_interval=-1)
+
 
 if __name__ == '__main__':
   logging.getLogger().setLevel(logging.INFO)