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

"""Implementation of ``DataChannel``s to communicate across the data plane."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import logging
import queue
import sys
import threading
from builtins import object
from builtins import range
from typing import TYPE_CHECKING
from typing import Callable
from typing import DefaultDict
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional

import grpc
from future.utils import raise_
from future.utils import with_metaclass

from apache_beam.coders import coder_impl
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.portability.api import beam_fn_api_pb2_grpc
from apache_beam.runners.worker.channel_factory import GRPCChannelFactory
from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor

if TYPE_CHECKING:
  # TODO: remove from TYPE_CHECKING scope when we drop support for python < 3.6
  from typing import Collection

# This module is experimental. No backwards-compatibility guarantees.

_LOGGER = logging.getLogger(__name__)

_DEFAULT_SIZE_FLUSH_THRESHOLD = 10 << 20  # 10MB
_DEFAULT_TIME_FLUSH_THRESHOLD_MS = 0  # disable time-based flush by default


if TYPE_CHECKING:
  import apache_beam.coders.slow_stream
  OutputStream = apache_beam.coders.slow_stream.OutputStream
else:
  OutputStream = type(coder_impl.create_OutputStream())


class ClosableOutputStream(OutputStream):
  """A Outputstream for use with CoderImpls that has a close() method."""

  def __init__(self, close_callback=None):
    super(ClosableOutputStream, self).__init__()
    self._close_callback = close_callback

  def close(self):
    if self._close_callback:
      self._close_callback(self.get())

  @staticmethod
  def create(close_callback,
             flush_callback,
             data_buffer_time_limit_ms):
    if data_buffer_time_limit_ms > 0:
      return TimeBasedBufferingClosableOutputStream(
          close_callback,
          flush_callback=flush_callback,
          time_flush_threshold_ms=data_buffer_time_limit_ms)
    else:
      return SizeBasedBufferingClosableOutputStream(
          close_callback, flush_callback=flush_callback)


class SizeBasedBufferingClosableOutputStream(ClosableOutputStream):
  """A size-based buffering OutputStream."""

  def __init__(self,
               close_callback=None,  # type: Optional[Callable[[bytes], None]]
               flush_callback=None,  # type: Optional[Callable[[bytes], None]]
               size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD):
    super(SizeBasedBufferingClosableOutputStream, self).__init__(close_callback)
    self._flush_callback = flush_callback
    self._size_flush_threshold = size_flush_threshold

  # This must be called explicitly to avoid flushing partial elements.
  def maybe_flush(self):
    if self.size() > self._size_flush_threshold:
      self.flush()

  def flush(self):
    if self._flush_callback:
      self._flush_callback(self.get())
      self._clear()


class TimeBasedBufferingClosableOutputStream(
    SizeBasedBufferingClosableOutputStream):
  """A buffering OutputStream with both time-based and size-based."""

  def __init__(self,
               close_callback=None,
               flush_callback=None,
               size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD,
               time_flush_threshold_ms=_DEFAULT_TIME_FLUSH_THRESHOLD_MS):
    super(TimeBasedBufferingClosableOutputStream, self).__init__(
        close_callback, flush_callback, size_flush_threshold)
    assert time_flush_threshold_ms > 0
    self._time_flush_threshold_ms = time_flush_threshold_ms
    self._flush_lock = threading.Lock()
    self._schedule_lock = threading.Lock()
    self._closed = False
    self._schedule_periodic_flush()

  def flush(self):
    with self._flush_lock:
      super(TimeBasedBufferingClosableOutputStream, self).flush()

  def close(self):
    with self._schedule_lock:
      self._closed = True
      if self._flush_timer:
        self._flush_timer.cancel()
        self._flush_timer = None
    super(TimeBasedBufferingClosableOutputStream, self).close()

  def _schedule_periodic_flush(self):
    def _periodic_flush():
      with self._schedule_lock:
        if not self._closed:
          self.flush()
          self._schedule_periodic_flush()

    self._flush_timer = threading.Timer(
        self._time_flush_threshold_ms / 1000.0, _periodic_flush)
    self._flush_timer.start()


class DataChannel(with_metaclass(abc.ABCMeta, object)):  # type: ignore[misc]
  """Represents a channel for reading and writing data over the data plane.

  Read from this channel with the input_elements method::

    for elements_data in data_channel.input_elements(
        instruction_id, transform_ids):
      [process elements_data]

  Write to this channel using the output_stream method::

    out1 = data_channel.output_stream(instruction_id, transform_id)
    out1.write(...)
    out1.close()

  When all data for all instructions is written, close the channel::

    data_channel.close()
  """

  @abc.abstractmethod
  def input_elements(self,
                     instruction_id,  # type: str
                     expected_transforms,  # type: Collection[str]<commit_msg>[BEAM-7949] Introduce PeriodicThread for time-based cache threshold support.<commit_after>#
# 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.
#

"""Implementation of ``DataChannel``s to communicate across the data plane."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import abc
import collections
import logging
import queue
import sys
import threading
import time
from builtins import object
from builtins import range
from typing import TYPE_CHECKING
from typing import Callable
from typing import DefaultDict
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional

import grpc
from future.utils import raise_
from future.utils import with_metaclass

from apache_beam.coders import coder_impl
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.portability.api import beam_fn_api_pb2_grpc
from apache_beam.runners.worker.channel_factory import GRPCChannelFactory
from apache_beam.runners.worker.worker_id_interceptor import WorkerIdInterceptor

if TYPE_CHECKING:
  # TODO: remove from TYPE_CHECKING scope when we drop support for python < 3.6
  from typing import Collection

# This module is experimental. No backwards-compatibility guarantees.

_LOGGER = logging.getLogger(__name__)

_DEFAULT_SIZE_FLUSH_THRESHOLD = 10 << 20  # 10MB
_DEFAULT_TIME_FLUSH_THRESHOLD_MS = 0  # disable time-based flush by default


if TYPE_CHECKING:
  import apache_beam.coders.slow_stream
  OutputStream = apache_beam.coders.slow_stream.OutputStream
else:
  OutputStream = type(coder_impl.create_OutputStream())


class ClosableOutputStream(OutputStream):
  """A Outputstream for use with CoderImpls that has a close() method."""

  def __init__(self, close_callback=None):
    super(ClosableOutputStream, self).__init__()
    self._close_callback = close_callback

  def close(self):
    if self._close_callback:
      self._close_callback(self.get())

  @staticmethod
  def create(close_callback,
             flush_callback,
             data_buffer_time_limit_ms):
    if data_buffer_time_limit_ms > 0:
      return TimeBasedBufferingClosableOutputStream(
          close_callback,
          flush_callback=flush_callback,
          time_flush_threshold_ms=data_buffer_time_limit_ms)
    else:
      return SizeBasedBufferingClosableOutputStream(
          close_callback, flush_callback=flush_callback)


class SizeBasedBufferingClosableOutputStream(ClosableOutputStream):
  """A size-based buffering OutputStream."""

  def __init__(self,
               close_callback=None,  # type: Optional[Callable[[bytes], None]]
               flush_callback=None,  # type: Optional[Callable[[bytes], None]]
               size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD):
    super(SizeBasedBufferingClosableOutputStream, self).__init__(close_callback)
    self._flush_callback = flush_callback
    self._size_flush_threshold = size_flush_threshold

  # This must be called explicitly to avoid flushing partial elements.
  def maybe_flush(self):
    if self.size() > self._size_flush_threshold:
      self.flush()

  def flush(self):
    if self._flush_callback:
      self._flush_callback(self.get())
      self._clear()


class TimeBasedBufferingClosableOutputStream(
    SizeBasedBufferingClosableOutputStream):
  """A buffering OutputStream with both time-based and size-based."""

  def __init__(self,
               close_callback=None,
               flush_callback=None,
               size_flush_threshold=_DEFAULT_SIZE_FLUSH_THRESHOLD,
               time_flush_threshold_ms=_DEFAULT_TIME_FLUSH_THRESHOLD_MS):
    super(TimeBasedBufferingClosableOutputStream, self).__init__(
        close_callback, flush_callback, size_flush_threshold)
    assert time_flush_threshold_ms > 0
    self._time_flush_threshold_ms = time_flush_threshold_ms
    self._flush_lock = threading.Lock()
    self._schedule_lock = threading.Lock()
    self._closed = False
    self._schedule_periodic_flush()

  def flush(self):
    with self._flush_lock:
      super(TimeBasedBufferingClosableOutputStream, self).flush()

  def close(self):
    with self._schedule_lock:
      self._closed = True
      if self._periodic_flusher:
        self._periodic_flusher.cancel()
        self._periodic_flusher = None
    super(TimeBasedBufferingClosableOutputStream, self).close()

  def _schedule_periodic_flush(self):
    def _flush():
      with self._schedule_lock:
        if not self._closed:
          self.flush()

    self._periodic_flusher = PeriodicThread(
        self._time_flush_threshold_ms / 1000.0, _flush)
    self._periodic_flusher.daemon = True
    self._periodic_flusher.start()


class PeriodicThread(threading.Thread):
  """Call a function periodically with the specified number of seconds"""

  def __init__(self,
               interval,
               function,
               args=None,
               kwargs=None):
    threading.Thread.__init__(self)
    self._interval = interval
    self._function = function
    self._args = args if args is not None else []
    self._kwargs = kwargs if kwargs is not None else {}
    self._finished = threading.Event()

  def run(self):
    next_call = time.time() + self._interval
    while not self._finished.wait(next_call - time.time()):
      next_call = next_call + self._interval
      self._function(*self._args, **self._kwargs)

  def cancel(self):
    """Stop the thread if it hasn't finished yet."""
    self._finished.set()


class DataChannel(with_metaclass(abc.ABCMeta, object)):  # type: ignore[misc]
  """Represents a channel for reading and writing data over the data plane.

  Read from this channel with the input_elements method::

    for elements_data in data_channel.input_elements(
        instruction_id, transform_ids):
      [process elements_data]

  Write to this channel using the output_stream method::

    out1 = data_channel.output_stream(instruction_id, transform_id)
    out1.write(...)
    out1.close()

  When all data for all instructions is written, close the channel::

    data_channel.close()
  """

  @abc.abstractmethod
  def input_elements(self,
                     instruction_id,  # type: str
                     expected_transforms,  # type: Collection[str]
{'fraction_non_alphanumeric': 0.06855971452433376, 'fraction_numerical': 0.00409991648318275, 'mean_word_length': 3.9931766489764975, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1081457', 'n_tokens_mistral': 3726, 'n_tokens_neox': 3600, 'n_words': 1051}