#

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

#

# pytype: skip-file

"""Common utility class to help SDK harness to execute an SDF. """

import logging

import threading

from typing import TYPE_CHECKING

from typing import Any

from typing import NamedTuple

from typing import Optional

from typing import Tuple

from typing import Union

from apache_beam.transforms.core import WatermarkEstimatorProvider

from apache_beam.utils.timestamp import Duration

from apache_beam.utils.timestamp import Timestamp

from apache_beam.utils.windowed_value import WindowedValue

if TYPE_CHECKING:

from apache_beam.io.iobase import RestrictionProgress

from apache_beam.io.iobase import RestrictionTracker

from apache_beam.io.iobase import WatermarkEstimator

_LOGGER = logging.getLogger(__name__)

SplitResultPrimary = NamedTuple(

'SplitResultPrimary', [('primary_value', WindowedValue)])

SplitResultResidual = NamedTuple(

'SplitResultResidual',

[('residual_value', WindowedValue), ('current_watermark', Timestamp),

('deferred_timestamp', Optional[Duration])])

class ThreadsafeRestrictionTracker(object):

"""A thread-safe wrapper which wraps a `RestrictionTracker`.

This wrapper guarantees synchronization of modifying restrictions across

multi-thread.

"""

def __init__(self, restriction_tracker: 'RestrictionTracker') -> None:

from apache_beam.io.iobase import RestrictionTracker

if not isinstance(restriction_tracker, RestrictionTracker):

raise ValueError(

'Initialize ThreadsafeRestrictionTracker requires'

'RestrictionTracker.')

self._restriction_tracker = restriction_tracker

# Records an absolute timestamp when defer_remainder is called.

self._timestamp = None

self._lock = threading.RLock()

self._deferred_residual = None

self._deferred_timestamp: Optional[Union[Timestamp, Duration]] = None

def current_restriction(self):

with self._lock:

return self._restriction_tracker.current_restriction()

def try_claim(self, position):

with self._lock:

return self._restriction_tracker.try_claim(position)

def defer_remainder(self, deferred_time=None):

"""Performs self-checkpoint on current processing restriction with an

expected resuming time.

Self-checkpoint could happen during processing elements. When executing an

DoFn.process(), you may want to stop processing an element and resuming

later if current element has been processed quit a long time or you also

want to have some outputs from other elements. ``defer_remainder()`` can be

called on per element if needed.

Args:

deferred_time: A relative ``Duration`` that indicates the ideal time gap

between now and resuming, or an absolute ``Timestamp`` for resuming

execution time. If the time_delay is None, the deferred work will be

executed as soon as possible.

"""

# Record current time for calculating deferred_time later.

with self._lock:

self._timestamp = Timestamp.now()

if deferred_time and not isinstance(deferred_time, (Duration, Timestamp)):

raise ValueError(

'The timestamp of deter_remainder() should be a '

'Duration or a Timestamp, or None.')

self._deferred_timestamp = deferred_time

checkpoint = self.try_split(0)

if checkpoint:

_, self._deferred_residual = checkpoint

def check_done(self):

with self._lock:

return self._restriction_tracker.check_done()

def current_progress(self) -> 'RestrictionProgress':

with self._lock:

return self._restriction_tracker.current_progress()

def try_split(self, fraction_of_remainder):

with self._lock:

return self._restriction_tracker.try_split(fraction_of_remainder)

def deferred_status(self) -> Optional[Tuple[Any, Duration]]:

"""Returns deferred work which is produced by ``defer_remainder()``.

When there is a self-checkpoint performed, the system needs to fulfill the

DelayedBundleApplication with deferred_work for a ProcessBundleResponse.

The system calls this API to get deferred_residual with watermark together

to help the runner to schedule a future work.

Returns: (deferred_residual, time_delay) if having any residual, else None.

"""

if self._deferred_residual:

# If _deferred_timestamp is None, create Duration(0).

if not self._deferred_timestamp:

self._deferred_timestamp = Duration()

# If an absolute timestamp is provided, calculate the delta between

# the absoluted time and the time deferred_status() is called.

elif isinstance(self._deferred_timestamp, Timestamp):

self._deferred_timestamp = (self._deferred_timestamp - Timestamp.now())

# If a Duration is provided, the deferred time should be:

# provided duration - the spent time since the defer_remainder() is

# called.

elif isinstance(self._deferred_timestamp, Duration):

self._deferred_timestamp -= (Timestamp.now() - self._timestamp)

return self._deferred_residual, self._deferred_timestamp

return None

def is_bounded(self):

return self._restriction_tracker.is_bounded()

class RestrictionTrackerView(object):

"""A DoFn view of thread-safe RestrictionTracker.

The RestrictionTrackerView wraps a ThreadsafeRestrictionTracker and only

exposes APIs that will be called by a ``DoFn.process()``. During execution

time, the RestrictionTrackerView will be fed into the ``DoFn.process`` as a

restriction_tracker.

"""

def __init__(

self,

threadsafe_restriction_tracker: ThreadsafeRestrictionTracker) -> None:

if not isinstance(threadsafe_restriction_tracker,

ThreadsafeRestrictionTracker):

raise ValueError(

'Initialize RestrictionTrackerView requires '

'ThreadsafeRestrictionTracker.')

self._threadsafe_restriction_tracker = threadsafe_restriction_tracker

def current_restriction(self):

return self._threadsafe_restriction_tracker.current_restriction()

def try_claim(self, position):

return self._threadsafe_restriction_tracker.try_claim(position)

def defer_remainder(self, deferred_time=None):

self._threadsafe_restriction_tracker.defer_remainder(deferred_time)

def is_bounded(self):

self._threadsafe_restriction_tracker.is_bounded()

class ThreadsafeWatermarkEstimator(object):

"""A threadsafe wrapper which wraps a WatermarkEstimator with locking

mechanism to guarantee multi-thread safety.

"""

def __init__(self, watermark_estimator: 'WatermarkEstimator') -> None:

from apache_beam.io.iobase import WatermarkEstimator

if not isinstance(watermark_estimator, WatermarkEstimator):

raise ValueError('Initializing Threadsafe requires a WatermarkEstimator')

self._watermark_estimator = watermark_estimator

self._lock = threading.Lock()

def __getattr__(self, attr):

if hasattr(self._watermark_estimator, attr):

def method_wrapper(*args, **kw):

with self._lock:

return getattr(self._watermark_estimator, attr)(*args, **kw)

return method_wrapper

raise AttributeError(attr)

def get_estimator_state(self):

with self._lock:

return self._watermark_estimator.get_estimator_state()

def current_watermark(self) -> Timestamp:

with self._lock:

return self._watermark_estimator.current_watermark()

def observe_timestamp(self, timestamp: Timestamp) -> None:

if not isinstance(timestamp, Timestamp):

raise ValueError(

'Input of observe_timestamp should be a Timestamp '

'object')

with self._lock:

self._watermark_estimator.observe_timestamp(timestamp)

class NoOpWatermarkEstimatorProvider(WatermarkEstimatorProvider):

"""A WatermarkEstimatorProvider which creates NoOpWatermarkEstimator for the

framework.

"""

def initial_estimator_state(self, element, restriction):

return None

def create_watermark_estimator(self, estimator_state):

from apache_beam.io.iobase import WatermarkEstimator

class _NoOpWatermarkEstimator(WatermarkEstimator):

"""A No-op WatermarkEstimator which is provided for the framework if there

is no custom one.

"""

def observe_timestamp(self, timestamp):

pass

def current_watermark(self):

return None

def get_estimator_state(self):

return None

return _NoOpWatermarkEstimator()