sdks/python/apache_beam/io/unbounded_source.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Experimental ``UnboundedSource`` support for the Python SDK.

``UnboundedSource`` support is currently experimental in the Python SDK; the
API may change in backwards-incompatible ways.

An unbounded source reads an effectively infinite stream of records (message
queues, change-data-capture feeds, and similar) with checkpoint-based
pause/resume, watermark reporting, and bundle finalization.

To define a source, implement :class:`UnboundedSource`, an
:class:`UnboundedReader`, and (when the reader has a resumable position) a
:class:`CheckpointMark`::

    import apache_beam as beam
    from apache_beam.io.unbounded_source import (
        CheckpointMark, UnboundedReader, UnboundedSource)
    from apache_beam.utils.timestamp import MAX_TIMESTAMP

    class MyCheckpointMark(CheckpointMark):
      def __init__(self, position):
        self.position = position

      def finalize_checkpoint(self):
        # Commit/acknowledge records up to ``position`` upstream, e.g. ack the
        # consumed messages on a queue.
        ...

    class MyReader(UnboundedReader):
      def start(self):
        # Position at the first record; return whether one is available.
        ...

      def advance(self):
        # Move to the next record; ``False`` means no data is available now.
        ...

      def get_current(self):
        ...

      def get_current_timestamp(self):
        ...  # event time of the current record

      def get_watermark(self):
        # Lower bound on the timestamps of future records. Return
        # ``MAX_TIMESTAMP`` to signal the reader has permanently finished.
        ...

      def get_checkpoint_mark(self):
        return MyCheckpointMark(...)

    class MySource(UnboundedSource):
      def split(self, desired_num_splits, options=None):
        # Return independent sub-sources, or ``[self]`` when not splittable.
        return [self]

      def create_reader(self, options, checkpoint_mark):
        # Build a reader; resume after ``checkpoint_mark`` when it is not None.
        return MyReader(...)

      def get_checkpoint_mark_coder(self):
        return ...  # a Coder for MyCheckpointMark

Read the source in a pipeline with :class:`apache_beam.io.Read`::

    with beam.Pipeline() as p:
      p | beam.io.Read(MySource()) | beam.Map(print)
"""

import collections
import dataclasses
import logging
import threading
import time
from typing import Any
from typing import Callable
from typing import Iterable
from typing import Optional

from typing_extensions import override

from apache_beam import coders
from apache_beam.coders.coders import BooleanCoder
from apache_beam.coders.coders import Coder
from apache_beam.coders.coders import NullableCoder
from apache_beam.coders.coders import TimestampCoder
from apache_beam.coders.coders import TupleCoder
from apache_beam.coders.coders import _MemoizingPickleCoder
from apache_beam.io import iobase
from apache_beam.io.watermark_estimators import ManualWatermarkEstimator
from apache_beam.options.pipeline_options import StreamingOptions
from apache_beam.options.pipeline_options_context import get_pipeline_options
from apache_beam.runners import sdf_utils
from apache_beam.transforms import PTransform
from apache_beam.transforms import core
from apache_beam.transforms.window import TimestampedValue
from apache_beam.utils.timestamp import MAX_TIMESTAMP
from apache_beam.utils.timestamp import MIN_TIMESTAMP
from apache_beam.utils.timestamp import Duration
from apache_beam.utils.timestamp import Timestamp

__all__ = [
    'CheckpointMark',
    'UnboundedReader',
    'UnboundedSource',
    'ReadFromUnboundedSource',
]

_LOGGER = logging.getLogger(__name__)

# Sentinel used when a reader has no data available right now.
# This is distinct from end-of-stream.
_NO_DATA = object()

_DEFAULT_POLL_INTERVAL_SECONDS = 1.0
_DEFAULT_DESIRED_NUM_SPLITS = 20
_DEFAULT_MAX_RECORDS_PER_BUNDLE = 10000
_DEFAULT_MAX_READ_TIME_SECONDS = 10.0
# A reader parked by a residual that never resumes is closed once idle this
# long; the cache also caps its entry count as a memory backstop.
_DEFAULT_READER_CACHE_IDLE_SECONDS = 60.0
_DEFAULT_READER_CACHE_MAX_SIZE = 100

# Encodes a source to a structural cache key. Internally consistent across park
# and acquire; need not match the restriction wire coder.
_SOURCE_KEY_CODER = _MemoizingPickleCoder()

# ------------------------------------------------------------------------------
# Public abstract base classes.
# ------------------------------------------------------------------------------


class CheckpointMark(object):
  """A durable, serializable position in an :class:`UnboundedSource`.

  Produced by :meth:`UnboundedReader.get_checkpoint_mark`, encoded with
  :meth:`UnboundedSource.get_checkpoint_mark_coder`, and used to resume a reader
  (see :meth:`UnboundedSource.create_reader`).
  """
  def finalize_checkpoint(self) -> None:
    """Called once the runner has durably committed work up to this mark.

    Override to acknowledge/commit upstream (for example, ack the consumed
    messages on a queue). The default is a no-op.

    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. 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


class UnboundedReader(object):
  """Reads records from an :class:`UnboundedSource`.

  Lifecycle: exactly one :meth:`start`, then any number of :meth:`advance`
  calls; whenever one returns ``True`` the current record is available via
  :meth:`get_current` / :meth:`get_current_timestamp`. A ``False`` return means
  "no data available right now", which is distinct from end-of-stream: a reader
  signals a permanent end by reporting a watermark of ``MAX_TIMESTAMP``.
  """
  def start(self) -> bool:
    """Positions at the first record; returns whether one is available."""
    raise NotImplementedError

  def advance(self) -> bool:
    """Advances to the next record. ``False`` means no data is available now.

    Should not block. The wrapper enforces the per-bundle record and time caps
    only between records, so a blocking ``start``/``advance`` can overrun the
    time cap and stall the bundle. Return ``False`` when no data is currently
    available instead of waiting.
    """
    raise NotImplementedError

  def get_current(self) -> Any:
    """Returns the record claimed by the last successful start/advance."""
    raise NotImplementedError

  def get_current_timestamp(self) -> Timestamp:
    """Returns the event-time timestamp of the current record."""
    raise NotImplementedError

  def get_watermark(self) -> Timestamp:
    """An approximate lower bound on timestamps of future records.

    Treated as monotonic by the wrapper. Return ``MAX_TIMESTAMP`` to signal that
    this reader has permanently finished.
    """
    raise NotImplementedError

  def get_checkpoint_mark(self) -> CheckpointMark:
    """Returns a durable mark to resume from. Call only at a bundle boundary."""
    raise NotImplementedError

  def close(self) -> None:
    """Releases reader resources. Default no-op."""
    pass


class UnboundedSource(iobase.SourceBase):
  """A source producing an unbounded stream of records with checkpointing.

  Read it in a pipeline with :class:`apache_beam.io.Read`::

      p | beam.io.Read(MyUnboundedSource())
  """
  def split(self,
            desired_num_splits: int,
            options: Optional[Any] = None) -> Iterable['UnboundedSource']:
    """Splits into at most ``desired_num_splits`` independent sub-sources.

    Each returned sub-source must be independent and must not share mutable
    state with siblings (the runner may execute them concurrently across
    workers). Return ``[self]`` if the source cannot be split. Splitting is
    performed once, before any checkpoint exists; once a reader has
    checkpointed, the restriction is kept intact.
    """
    raise NotImplementedError

  def create_reader(
      self, options: Optional[Any],
      checkpoint_mark: Optional[CheckpointMark]) -> UnboundedReader:
    """Creates a reader, optionally resuming from ``checkpoint_mark``.

    Contract:
      * When ``checkpoint_mark`` is ``None``, the returned reader's ``start()``
        produces the very first record of the source (or returns ``False`` if
        none yet).
      * When ``checkpoint_mark`` is not ``None``, the returned reader's
        ``start()`` produces the first record strictly after the position
        encoded by ``checkpoint_mark``. The reader must not re-deliver records
        already covered by the prior bundle.
    """
    raise NotImplementedError

  def get_checkpoint_mark_coder(self) -> Coder:
    """Returns the coder for this source's :class:`CheckpointMark` instances.

    The SDK may call this while encoding or decoding source restrictions.
    Implementations should be deterministic, side-effect free, and should not
    perform I/O.
    """
    raise NotImplementedError(
        '%s must override get_checkpoint_mark_coder() to return a Coder for '
        'its CheckpointMark subclass.' % type(self).__name__)

  def is_bounded(self) -> bool:
    # SourceBase.is_bounded raises; an unbounded source is, by definition, not.
    return False

  def default_output_coder(self) -> Coder:
    # Permissive default; override for a tighter coder.
    return coders.registry.get_coder(object)


# ------------------------------------------------------------------------------
# SDF wrapper internals: a private implementation detail of
# ReadFromUnboundedSource.
# ------------------------------------------------------------------------------


@dataclasses.dataclass(frozen=True)
class _UnboundedSourceRestriction(object):
  """Durable SDF restriction describing where a reader should (re)start.

  Holds only serializable state -- never a live reader. ``is_done`` marks the
  terminal (MAX-watermark) transition. ``finalization_checkpoint_mark`` is kept
  separate from ``checkpoint_mark`` so a done primary can carry a commit hook
  without polluting the resume state.

  Field roles:
    * ``checkpoint_mark`` -- RESUME state. A reader rebuilt from this mark
      must produce the first record strictly after it.
    * ``finalization_checkpoint_mark`` -- COMMIT hook. Only set on a done
      primary that was just cut this bundle. Registered with the runner's
      bundle finalizer to acknowledge upstream. Independent of
      ``checkpoint_mark`` so a residual's resume state can be ``None`` while
      still recording what should be committed.
  """
  source: UnboundedSource
  checkpoint_mark: Optional[CheckpointMark] = None
  watermark: Timestamp = MIN_TIMESTAMP
  is_done: bool = False
  finalization_checkpoint_mark: Optional[CheckpointMark] = None


class _UnboundedSourceRestrictionCoder(Coder):
  """Encodes :class:`_UnboundedSourceRestriction` as a fixed 5-tuple.

  Stateless: at encode time the source's own
  :meth:`UnboundedSource.get_checkpoint_mark_coder` is looked up from the
  restriction; at decode time the source is decoded first and its coder
  drives the checkpoint-mark decoding. This avoids passing source-specific
  coder state into the coder's constructor, which in turn lets
  :class:`_UnboundedSourceRestrictionProvider` and
  :class:`_ReadFromUnboundedSourceDoFn` be module-level classes.

  Wire shape: source_bytes / checkpoint_bytes / watermark / done /
  finalization_checkpoint_bytes -- the checkpoint and finalization bytes
  are independently encoded with the (source-declared) checkpoint coder
  wrapped in :class:`NullableCoder`.
  """
  def __init__(self):
    self._source_coder = _MemoizingPickleCoder()
    self._bytes_coder = coders.BytesCoder()
    self._tuple_coder = TupleCoder((
        self._bytes_coder,  # source (pickled bytes)
        self._bytes_coder,  # checkpoint_mark (nullable-encoded bytes)
        TimestampCoder(),  # watermark
        BooleanCoder(),  # is_done
        self._bytes_coder))  # finalization_checkpoint_mark (nullable-encoded)

  def _checkpoint_coder(self, source: UnboundedSource) -> Coder:
    return NullableCoder(source.get_checkpoint_mark_coder())

  def encode(self, restriction: '_UnboundedSourceRestriction') -> bytes:
    source_bytes = self._source_coder.encode(restriction.source)
    cp_coder = self._checkpoint_coder(restriction.source)
    return self._tuple_coder.encode((
        source_bytes,
        cp_coder.encode(restriction.checkpoint_mark),
        restriction.watermark,
        restriction.is_done,
        cp_coder.encode(restriction.finalization_checkpoint_mark)))

  def decode(self, encoded: bytes) -> '_UnboundedSourceRestriction':
    (source_bytes, checkpoint_bytes, watermark, is_done,
     finalization_bytes) = self._tuple_coder.decode(encoded)
    source = self._source_coder.decode(source_bytes)
    cp_coder = self._checkpoint_coder(source)
    return _UnboundedSourceRestriction(
        source=source,
        checkpoint_mark=cp_coder.decode(checkpoint_bytes),
        watermark=watermark,
        is_done=is_done,
        finalization_checkpoint_mark=cp_coder.decode(finalization_bytes))

  def is_deterministic(self) -> bool:
    # Pickled source and checkpoint are not guaranteed deterministic.
    return False


class _ReaderCache(object):
  """Holds live readers between an SDF self-checkpoint and its resume.

  A fresh tracker is built for every bundle, so a reader parked at a
  self-checkpoint would otherwise be closed and rebuilt from its checkpoint
  mark on the next bundle. Parking it here lets the resuming bundle reclaim the
  same started reader, keyed by the residual's structural ``(source, checkpoint
  mark)``. ``acquire`` removes the entry, so two trackers never drive one
  reader.

  A residual may be reassigned to another worker and never resume here; such a
  reader is released once it falls idle past ``idle_seconds`` or when the entry
  count exceeds ``max_size``, and the owning DoFn's teardown releases the rest.
  One DoFn instance drives several trackers across threads, so access is locked.
  """
  def __init__(
      self,
      idle_seconds: float = _DEFAULT_READER_CACHE_IDLE_SECONDS,
      max_size: int = _DEFAULT_READER_CACHE_MAX_SIZE,
      now: Optional[Callable[[], float]] = None):
    self._idle_seconds = idle_seconds
    self._max_size = max_size
    self._now = now or time.monotonic
    self._lock = threading.Lock()
    # key -> (reader, started, parked_at); ordered oldest-first for eviction.
    self._entries = collections.OrderedDict(
    )  # type: collections.OrderedDict[Any, tuple[UnboundedReader, bool, float]]

  def acquire(self, key: Any) -> Optional[tuple['UnboundedReader', bool]]:
    """Removes and returns ``(reader, started)`` for ``key``, or None."""
    with self._lock:
      entry = self._entries.pop(key, None)
      stale = self._evict_idle()
    self._close_readers(stale)
    if entry is None:
      return None
    return entry[0], entry[1]

  def park(self, key: Any, reader: 'UnboundedReader', started: bool) -> None:
    """Stores ``reader`` under ``key`` for a later bundle to reclaim. A reader
    already parked under ``key`` is closed."""
    with self._lock:
      replaced = self._entries.pop(key, None)
      self._entries[key] = (reader, started, self._now())
      stale = self._evict_idle()
      while len(self._entries) > self._max_size:
        _, oldest = self._entries.popitem(last=False)
        stale.append(oldest[0])
    if replaced is not None and replaced[0] is not reader:
      stale.append(replaced[0])
    self._close_readers(stale)

  def close_all(self) -> None:
    """Closes every parked reader. Called from the owning DoFn's teardown."""
    with self._lock:
      entries = list(self._entries.values())
      self._entries.clear()
    self._close_readers(entry[0] for entry in entries)

  def _evict_idle(self) -> list:
    # Caller holds the lock. Pops entries idle past the window (oldest first)
    # and returns their readers for the caller to close after unlocking.
    deadline = self._now() - self._idle_seconds
    stale = []
    while self._entries:
      key, entry = next(iter(self._entries.items()))
      if entry[2] > deadline:
        break
      del self._entries[key]
      stale.append(entry[0])
    return stale

  def _close_readers(self, readers) -> None:
    for reader in readers:
      try:
        reader.close()
      except Exception:  # pylint: disable=broad-except
        _LOGGER.warning('Error closing UnboundedReader', exc_info=True)


class _UnboundedSourceRestrictionTracker(iobase.RestrictionTracker):
  """Drives an :class:`UnboundedReader` for one SDF restriction.

  Owns the live reader (lazily created, never serialized): both runner-initiated
  ``defer_remainder`` self-checkpoints with ``try_split(0)``, which must
  checkpoint the live reader.

  A self-checkpoint parks the reader in the DoFn's :class:`_ReaderCache` for the
  next bundle to reclaim, keeping one started reader alive across bundles. The
  DoFn injects ``_reader_cache`` at the start of ``process()``; with no cache
  the tracker builds a fresh reader each bundle.

  ``process()`` only sees a ``RestrictionTrackerView``, which hides custom
  methods and whose ``try_claim`` returns just a bool, so the freshly-read
  record is handed back through a one-element holder list passed as the
  ``try_claim`` *position* argument.
  """
  def __init__(
      self,
      restriction: _UnboundedSourceRestriction,
      options: Optional[Any] = None):
    self._restriction = restriction
    self._options = options
    self._reader = None  # type: Optional[UnboundedReader]
    self._started = False
    # True once a checkpoint has been cut this bundle (EOF or self-checkpoint).
    self._checkpoint_taken = False
    # Cross-bundle reader cache, injected by the DoFn; None disables caching.
    self._reader_cache = None  # type: Optional[_ReaderCache]

  def _ensure_reader(self) -> None:
    if self._reader is not None:
      return
    cached = self._acquire_cached_reader()
    if cached is not None:
      # A parked reader is already started and positioned past its checkpoint.
      self._reader, self._started = cached
      return
    self._reader = self._restriction.source.create_reader(
        self._options, self._restriction.checkpoint_mark)

  def _cache_key(self,
                 restriction: _UnboundedSourceRestriction) -> Optional[Any]:
    """Structural ``(source, checkpoint)`` key, or None when uncacheable.

    Built from the source pickle and the source's own checkpoint coder so a
    parked reader and its resuming restriction map to the same entry. A None
    key disables caching for that restriction; the resume then rebuilds from
    the checkpoint mark under the source's ``create_reader`` contract.
    """
    try:
      source_bytes = _SOURCE_KEY_CODER.encode(restriction.source)
      cp_coder = NullableCoder(restriction.source.get_checkpoint_mark_coder())
      return source_bytes, cp_coder.encode(restriction.checkpoint_mark)
    except Exception:  # pylint: disable=broad-except
      return None

  def _acquire_cached_reader(self) -> Optional[tuple['UnboundedReader', bool]]:
    if self._reader_cache is None:
      return None
    key = self._cache_key(self._restriction)
    if key is None:
      return None
    return self._reader_cache.acquire(key)

  def _park_or_close_reader(
      self, residual: _UnboundedSourceRestriction) -> None:
    """Hands the live reader to the cache for ``residual`` to reclaim, or
    closes it when no cache is available or the restriction is uncacheable."""
    if self._reader is None:
      return
    key = (
        self._cache_key(residual) if self._reader_cache is not None else None)
    if key is None:
      self._close_reader_if_open()
      return
    reader, self._reader = self._reader, None
    self._reader_cache.park(key, reader, self._started)

  def _clone_checkpoint(
      self, checkpoint: Optional[CheckpointMark]) -> Optional[CheckpointMark]:
    """Returns an independent copy of a mark via the source's checkpoint coder.

    Used to keep the primary's finalize hook and the residual's resume state
    from sharing one object, since a user ``finalize_checkpoint()`` may mutate
    the mark.
    """
    if checkpoint is None:
      return None
    coder = self._restriction.source.get_checkpoint_mark_coder()
    return coder.decode(coder.encode(checkpoint))

  def _close_reader_if_open(self) -> None:
    """Idempotent reader release. Called by the EOF and split paths, and by
    the DoFn's ``finally`` so an exception inside ``process()`` does not leak
    sockets / file descriptors held by the live :class:`UnboundedReader`.
    """
    if self._reader is None:
      return
    try:
      self._reader.close()
    except Exception:  # pylint: disable=broad-except
      _LOGGER.warning('Error closing UnboundedReader', exc_info=True)
    finally:
      self._reader = None

  def current_restriction(self) -> _UnboundedSourceRestriction:
    return self._restriction

  def try_claim(self, out: list[Any]) -> bool:
    """Advances the reader by one record.

    ``out[0]`` receives ``(value, record_timestamp, source_watermark)`` on the
    has-data path, or the :data:`_NO_DATA` sentinel otherwise. The watermark is
    the source's reported watermark, not the record's event time: the DoFn
    advances the output watermark with the former and timestamps the record
    with the latter. The argument doubles as the output holder (the
    ``RestrictionTracker`` ABC treats it as a claim position), which the
    threadsafe-tracker chain forwards opaquely.
    """
    try:
      return self._try_claim_inner(out)
    except Exception:
      # Reader state is now indeterminate; release it before re-raising.
      self._close_reader_if_open()
      raise

  def _try_claim_inner(self, out: list[Any]) -> bool:
    if self._restriction.is_done:
      out[0] = _NO_DATA
      return False
    self._ensure_reader()
    if not self._started:
      has_data = self._reader.start()
    else:
      has_data = self._reader.advance()
    self._started = True
    if has_data:
      # Emit an available record before checking the watermark: a reader may
      # report its last record and a MAX_TIMESTAMP watermark on the same call,
      # and EOF is realized on the next data-less claim.
      out[0] = (
          self._reader.get_current(),
          self._reader.get_current_timestamp(),
          self._reader.get_watermark())
      return True
    watermark = self._reader.get_watermark()
    if watermark >= MAX_TIMESTAMP:
      # No data and watermark at MAX: cut a final checkpoint, close, mark done.
      checkpoint = self._reader.get_checkpoint_mark()
      self._close_reader_if_open()
      self._restriction = dataclasses.replace(
          self._restriction,
          checkpoint_mark=None,  # nothing left to resume from
          watermark=MAX_TIMESTAMP,
          is_done=True,
          finalization_checkpoint_mark=checkpoint)
      self._checkpoint_taken = True
      out[0] = _NO_DATA
      return False
    # No data is available now. Refresh the watermark before deferring.
    self._restriction = dataclasses.replace(
        self._restriction, watermark=watermark)
    out[0] = _NO_DATA
    return True

  def try_split(
      self, fraction_of_remainder
  ) -> Optional[tuple[_UnboundedSourceRestriction,
                      _UnboundedSourceRestriction]]:
    """Cuts a checkpoint, returning (primary, residual) or None.

    The cut checkpoint goes into ``primary.finalization_checkpoint_mark`` so
    the DoFn can register a bundle-finalize callback for it. The same
    checkpoint object also goes into ``residual.checkpoint_mark`` so the
    resumed reader rebuilds at the correct position. The two fields are
    independent on purpose (see :class:`_UnboundedSourceRestriction`
    docstring): a runner that re-processes the primary alone must not see
    a stale resume state, and a residual must not register finalize again
    until ITS checkpoint is cut in a future bundle.
    """
    try:
      return self._try_split_inner(fraction_of_remainder)
    except Exception:
      self._close_reader_if_open()
      raise

  def _try_split_inner(self, fraction_of_remainder):
    # Only self-checkpoint (fraction 0) is supported; decline runner-initiated
    # dynamic splits.
    if fraction_of_remainder != 0:
      return None
    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. The residual gets its own clone
    # so a finalize_checkpoint() that mutates the primary's mark cannot corrupt
    # the residual's resume position before the runner encodes it.
    primary = dataclasses.replace(
        self._restriction,
        checkpoint_mark=None,
        is_done=True,
        finalization_checkpoint_mark=checkpoint)
    residual = _UnboundedSourceRestriction(
        source=self._restriction.source,
        checkpoint_mark=self._clone_checkpoint(checkpoint),
        watermark=watermark,
        is_done=False,
        finalization_checkpoint_mark=None)
    self._restriction = primary
    self._checkpoint_taken = True
    # Park the reader so the resuming bundle reclaims it; on a cache miss the
    # residual rebuilds one from its checkpoint mark.
    self._park_or_close_reader(residual)
    return primary, residual

  def check_done(self) -> bool:
    # Called after every process(); must raise if work is left unaccounted for.
    if self._restriction.is_done or self._checkpoint_taken:
      return True
    raise ValueError(
        'UnboundedSource restriction was neither finished nor checkpointed; '
        'process() must self-checkpoint via defer_remainder() or run to EOF: '
        '%r' % (self._restriction, ))

  def current_progress(self) -> 'iobase.RestrictionProgress':
    # Backlog-based progress is not implemented; report a coarse done/not-done
    # signal via ``completed`` / ``remaining``.
    if self._restriction.is_done:
      return iobase.RestrictionProgress(completed=1.0, remaining=0.0)
    return iobase.RestrictionProgress(completed=0.0, remaining=1.0)

  def is_bounded(self) -> bool:
    return False


class _UnboundedSourceRestrictionProvider(core.RestrictionProvider):
  """Wraps an :class:`UnboundedSource` element as an SDF restriction.

  Stateless module-level singleton (see :data:`_PROVIDER`): all
  source-specific state (e.g. the source's checkpoint coder) is derived
  per-call from the restriction's ``source`` field, which lets
  :class:`_ReadFromUnboundedSourceDoFn` live at module level too. The provider
  currently passes ``None`` for the ``options`` forwarded to
  :meth:`UnboundedSource.split`.
  """
  def __init__(self):
    self._restriction_coder = _UnboundedSourceRestrictionCoder()

  def initial_restriction(
      self, element: UnboundedSource) -> _UnboundedSourceRestriction:
    if not isinstance(element, UnboundedSource):
      raise TypeError(
          'ReadFromUnboundedSource expected an UnboundedSource element, got %r'
          % (element, ))
    return _UnboundedSourceRestriction(source=element)

  def create_tracker(
      self, restriction: _UnboundedSourceRestriction
  ) -> _UnboundedSourceRestrictionTracker:
    return _UnboundedSourceRestrictionTracker(restriction)

  def split(self, element,
            restriction) -> Iterable[_UnboundedSourceRestriction]:
    if restriction.is_done or restriction.checkpoint_mark is not None:
      yield restriction
      return

    desired_num_splits = _DEFAULT_DESIRED_NUM_SPLITS
    options = get_pipeline_options()
    if options is not None:
      configured_splits = options.view_as(
          StreamingOptions).desired_num_unbounded_source_splits
      if configured_splits > 0:
        desired_num_splits = configured_splits

    # ``source.split`` is user code and may refuse to split; fall back to a
    # single restriction on error.
    try:
      split_sources = list(restriction.source.split(desired_num_splits, None))
    except Exception:  # pylint: disable=broad-except
      _LOGGER.warning(
          'Exception while splitting UnboundedSource. Source not split.',
          exc_info=True)
      yield restriction
      return

    if not split_sources:
      yield restriction
      return

    # A non-UnboundedSource split result is a contract violation, not a
    # refusal, so fail loudly (outside the try/except above).
    for split_source in split_sources:
      if not isinstance(split_source, UnboundedSource):
        raise TypeError(
            'UnboundedSource.split() produced %r, expected UnboundedSource' %
            (split_source, ))

    for split_source in split_sources:
      yield dataclasses.replace(
          restriction,
          source=split_source,
          checkpoint_mark=None,
          is_done=False,
          finalization_checkpoint_mark=None)

  def restriction_size(self, element, restriction) -> int:
    # TODO(https://github.com/apache/beam/issues/19137): implement backlog
    # estimation.
    return 1

  def restriction_coder(self) -> Coder:
    return self._restriction_coder

  def truncate(self, element, restriction):
    # On drain, stop emitting new records.
    return None


# Module-level stateless singleton, captured via ``RestrictionParam`` at the
# DoFn's class-definition time.
_PROVIDER = _UnboundedSourceRestrictionProvider()


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

  def __call__(self) -> None:
    with self._lock:
      if self._finalized:
        return
      self._finalized = True
    # Finalization is best effort: log and swallow so a failing user override
    # does not fail the bundle (matches CheckpointMark.finalize_checkpoint).
    try:
      self._checkpoint_mark.finalize_checkpoint()
    except Exception:  # pylint: disable=broad-except
      _LOGGER.warning(
          'Error finalizing UnboundedSource checkpoint mark.', exc_info=True)


class _ReadFromUnboundedSourceDoFn(core.DoFn):
  """SDF wrapper driving an :class:`UnboundedReader` for one restriction.

  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,
      max_records_per_bundle: int = _DEFAULT_MAX_RECORDS_PER_BUNDLE,
      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
    # Per-worker reader cache; created in setup(), never serialized.
    self._reader_cache = None  # type: Optional[_ReaderCache]

  def setup(self):
    self._reader_cache = _ReaderCache()

  def teardown(self):
    if self._reader_cache is not None:
      self._reader_cache.close_all()
      self._reader_cache = None

  @core.DoFn.unbounded_per_element()
  def process(
      self,
      unused_element,
      bundle_finalizer=core.DoFn.BundleFinalizerParam,
      tracker=core.DoFn.RestrictionParam(_PROVIDER),
      watermark_estimator=core.DoFn.WatermarkEstimatorParam(
          ManualWatermarkEstimator.default_provider())):
    # Positional params (element, bundle finalizer) must precede the
    # kwarg-injected ones (tracker, watermark estimator).
    assert isinstance(tracker, sdf_utils.RestrictionTrackerView)
    inner_tracker = _unwrap_tracker(tracker)
    if inner_tracker is not None and self._reader_cache is not None:
      # Let this bundle reclaim a reader parked by the prior bundle and re-park
      # it on self-checkpoint. No cache means setup() was skipped.
      inner_tracker._reader_cache = self._reader_cache
    initial = tracker.current_restriction()
    now = self._now or time.monotonic
    records_emitted = 0
    # Armed on the first emitted record so reader startup is excluded.
    read_deadline = None  # type: Optional[float]
    try:
      while True:
        holder = [None]
        if not tracker.try_claim(holder):
          # EOF: advance the estimator to the tracker's MAX watermark so
          # downstream event-time windows can close.
          _set_watermark_if_greater(
              watermark_estimator, tracker.current_restriction().watermark)
          break
        record = holder[0]
        if record is _NO_DATA:
          # 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))
          break
        # The third tuple field is the source watermark. The record timestamp
        # remains the output event time. Emit the element before advancing the
        # estimator so a reader that reports MAX on the same claim as its final
        # record cannot push the output watermark past that record first.
        value, record_timestamp, source_watermark = record
        yield TimestampedValue(value, record_timestamp)
        _set_watermark_if_greater(watermark_estimator, source_watermark)
        records_emitted += 1
        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))
      finally:
        # The EOF and self-checkpoint paths already closed or parked the
        # reader, so this is a no-op there. It closes a reader still held when
        # process() exits early, e.g. a downstream yield raised before any
        # checkpoint.
        if inner_tracker is not None:
          inner_tracker._close_reader_if_open()
        else:
          _LOGGER.warning(
              'UnboundedSource DoFn could not close a reader because the SDF '
              'tracker wrapper did not expose '
              '_UnboundedSourceRestrictionTracker (got %s). Reader resources '
              'may remain open until garbage collection.',
              type(tracker).__name__)


def _unwrap_tracker(
    tracker: Any) -> Optional['_UnboundedSourceRestrictionTracker']:
  """Returns the :class:`_UnboundedSourceRestrictionTracker` behind the SDF
  view and threadsafe wrappers, or None when the chain is unexpected."""
  inner = tracker
  if hasattr(inner, '_threadsafe_restriction_tracker'):
    inner = inner._threadsafe_restriction_tracker
  if hasattr(inner, '_restriction_tracker'):
    inner = inner._restriction_tracker
  if isinstance(inner, _UnboundedSourceRestrictionTracker):
    return inner
  return None


def _set_watermark_if_greater(
    watermark_estimator, new_watermark: Timestamp) -> None:
  # ManualWatermarkEstimator.set_watermark raises on regression, so only ever
  # advance it (a regressing reader watermark is absorbed here).
  current = watermark_estimator.current_watermark()
  if current is None or new_watermark > current:
    watermark_estimator.set_watermark(new_watermark)


class ReadFromUnboundedSource(PTransform):
  """Reads an :class:`UnboundedSource`.

  Most users should prefer :class:`apache_beam.io.Read`, which dispatches an
  ``UnboundedSource`` here automatically::

      p | beam.io.Read(MyUnboundedSource())

  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,
      source: UnboundedSource,
      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, ))
    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
    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
    output_coder = source.default_output_coder()
    # The source is the SDF element used to derive the initial restriction.
    # process() reads from the restriction, so it does not use the element
    # directly.
    output = (
        pbegin
        | 'Create' >> core.Create([source])
        | 'ReadUnbounded' >> core.ParDo(
            _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.
    try:
      type_hint = output_coder.to_type_hint()
    except NotImplementedError:
      type_hint = None
    if type_hint is not None:
      try:
        registered_coder = coders.registry.get_coder(type_hint)
      except Exception:  # pylint: disable=broad-except
        _LOGGER.warning(
            'Could not look up the registered coder for element type %s.',
            type_hint,
            exc_info=True)
      else:
        if registered_coder == output_coder:
          output.element_type = type_hint
    return output

  def _infer_output_coder(self, input_type=None, input_coder=None):
    return self._source.default_output_coder()


# ------------------------------------------------------------------------------
# Internal in-memory counting source, shared by Beam's own tests and examples.
# ------------------------------------------------------------------------------

# Realistic event-time base away from the Unix epoch.
_EVENT_TIME_BASE = Timestamp(1729987200)  # 2024-10-27T00:00:00Z


class _CountingCheckpointMark(CheckpointMark):
  def __init__(self, last_index, finalize_log=None):
    self.last_index = last_index
    self._finalize_log = finalize_log

  @override
  def finalize_checkpoint(self):
    if self._finalize_log is not None:
      self._finalize_log.append(self.last_index)

  def __eq__(self, other):
    return (
        isinstance(other, _CountingCheckpointMark) and
        other.last_index == self.last_index)

  def __hash__(self):
    return hash(self.last_index)

  def __repr__(self):
    return '_CountingCheckpointMark(last_index=%r)' % (self.last_index, )


class _CountingReader(UnboundedReader):
  def __init__(
      self, count, start_index, finalize_log=None, modulus=1, residue=0):
    self._count = count
    self._next = start_index
    self._modulus = modulus
    self._residue = residue
    self._current = None
    self._exhausted = False
    self._finalize_log = finalize_log
    self.closed = False

  def _read_next(self):
    while self._next < self._count:
      index = self._next
      self._next += 1
      if index % self._modulus == self._residue:
        self._current = index
        return True
    self._exhausted = True
    return False

  @override
  def start(self):
    return self._read_next()

  @override
  def advance(self):
    return self._read_next()

  @override
  def get_current(self):
    return self._current

  @override
  def get_current_timestamp(self):
    return _EVENT_TIME_BASE + self._current

  @override
  def get_watermark(self):
    if self._exhausted:
      return MAX_TIMESTAMP
    if self._current is None:
      return MIN_TIMESTAMP
    return _EVENT_TIME_BASE + self._current

  @override
  def get_checkpoint_mark(self):
    last = self._current if self._current is not None else self._next - 1
    return _CountingCheckpointMark(last, finalize_log=self._finalize_log)

  @override
  def close(self):
    self.closed = True


class UnboundedCountingSource(UnboundedSource):
  """In-memory counting source for Beam's own tests and examples.

  Internal: not part of the public ``UnboundedSource`` API. Emits integers
  ``0..count-1`` with event time ``_EVENT_TIME_BASE + index``, self-terminates
  at EOF, resumes from ``last_index + 1``, and optionally splits into
  independent even/odd sub-sources.
  """
  def __init__(
      self,
      count,
      finalize_log=None,
      is_splittable=False,
      modulus=1,
      residue=0):
    self._count = count
    self._finalize_log = finalize_log
    self._is_splittable = is_splittable
    self._modulus = modulus
    self._residue = residue
    self.last_reader = None

  @override
  def split(self, desired_num_splits, options=None):
    if not self._is_splittable or desired_num_splits < 2:
      return [self]
    # Split into independent even/odd sub-sources (each non-splittable).
    return [
        UnboundedCountingSource(
            self._count,
            finalize_log=self._finalize_log,
            modulus=2,
            residue=residue) for residue in (0, 1)
    ]

  @override
  def create_reader(self, options, checkpoint_mark):
    start_index = (
        0 if checkpoint_mark is None else checkpoint_mark.last_index + 1)
    self.last_reader = _CountingReader(
        self._count,
        start_index,
        finalize_log=self._finalize_log,
        modulus=self._modulus,
        residue=self._residue)
    return self.last_reader

  @override
  def get_checkpoint_mark_coder(self):
    return coders.PickleCoder()