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

"""Code for communicating with the Workers."""

# mypy: disallow-untyped-defs

import collections
import contextlib
import copy
import logging
import os
import queue
import subprocess
import sys
import threading
import time
from typing import TYPE_CHECKING
from typing import Any
from typing import BinaryIO  # pylint: disable=unused-import
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 Mapping
from typing import Optional
from typing import Tuple
from typing import Type
from typing import TypeVar
from typing import Union
from typing import cast
from typing import overload

import grpc

from apache_beam.io import filesystems
from apache_beam.io.filesystems import CompressionTypes
from apache_beam.portability import common_urns
from apache_beam.portability import python_urns
from apache_beam.portability.api import beam_artifact_api_pb2_grpc
from apache_beam.portability.api import beam_fn_api_pb2
from apache_beam.portability.api import beam_fn_api_pb2_grpc
from apache_beam.portability.api import beam_provision_api_pb2
from apache_beam.portability.api import beam_provision_api_pb2_grpc
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import endpoints_pb2
from apache_beam.runners.portability import artifact_service
from apache_beam.runners.portability.fn_api_runner.execution import Buffer
from apache_beam.runners.worker import data_plane
from apache_beam.runners.worker import sdk_worker
from apache_beam.runners.worker.channel_factory import GRPCChannelFactory
from apache_beam.runners.worker.log_handler import LOGENTRY_TO_LOG_LEVEL_MAP
from apache_beam.runners.worker.sdk_worker import _Future
from apache_beam.runners.worker.statecache import StateCache
from apache_beam.utils import proto_utils
from apache_beam.utils import thread_pool_executor
from apache_beam.utils.interactive_utils import is_in_notebook
from apache_beam.utils.sentinel import Sentinel

if TYPE_CHECKING:
  from grpc import ServicerContext
  from google.protobuf import message
  from apache_beam.runners.portability.fn_api_runner.fn_runner import ExtendedProvisionInfo  # pylint: disable=ungrouped-imports

# State caching is enabled in the fn_api_runner for testing, except for one
# test which runs without state caching (FnApiRunnerTestWithDisabledCaching).
# The cache is disabled in production for other runners.
STATE_CACHE_SIZE_MB = 100
MB_TO_BYTES = 1 << 20

# Time-based flush is enabled in the fn_api_runner by default.
DATA_BUFFER_TIME_LIMIT_MS = 1000

FNAPI_RUNNER_CAPABILITIES = frozenset([
    common_urns.runner_protocols.MULTIMAP_KEYS_VALUES_SIDE_INPUT.urn,
])

_LOGGER = logging.getLogger(__name__)

T = TypeVar('T')
ConstructorFn = Callable[[
    Union['message.Message', bytes],
    'sdk_worker.StateHandler',
    'ExtendedProvisionInfo',
    'GrpcServer'
],
                         'WorkerHandler']


class ControlConnection(object):

  _uid_counter = 0
  _lock = threading.Lock()

  def __init__(self):
    # type: () -> None
    self._push_queue = queue.Queue(
    )  # type: queue.Queue[Union[beam_fn_api_pb2.InstructionRequest, Sentinel]]
    self._input = None  # type: Optional[Iterable[beam_fn_api_pb2.InstructionResponse]]
    self._futures_by_id = {}  # type: Dict[str, ControlFuture]
    self._read_thread = threading.Thread(
        name='beam_control_read', target=self._read)
    self._state = BeamFnControlServicer.UNSTARTED_STATE

  def _read(self):
    # type: () -> None
    assert self._input is not None
    for data in self._input:
      self._futures_by_id.pop(data.instruction_id).set(data)

  @overload
  def push(self, req):
    # type: (Sentinel) -> None
    pass

  @overload
  def push(self, req):
    # type: (beam_fn_api_pb2.InstructionRequest) -> ControlFuture
    pass

  def push(self, req):
    # type: (Union[Sentinel, beam_fn_api_pb2.InstructionRequest]) -> Optional[ControlFuture]
    if req is BeamFnControlServicer._DONE_MARKER:
      self._push_queue.put(req)
      return None
    if not req.instruction_id:
      with ControlConnection._lock:
        ControlConnection._uid_counter += 1
        req.instruction_id = 'control_%s' % ControlConnection._uid_counter
    future = ControlFuture(req.instruction_id)
    self._futures_by_id[req.instruction_id] = future
    self._push_queue.put(req)
    return future

  def get_req(self):
    # type: () -> Union[Sentinel, beam_fn_api_pb2.InstructionRequest]
    return self._push_queue.get()

  def set_input(self, input):
    # type: (Iterable[beam_fn_api_pb2.InstructionResponse]) -> None
    with ControlConnection._lock:
      if self._input:
        raise RuntimeError('input is already set.')
      self._input = input
      self._read_thread.start()
      self._state = BeamFnControlServicer.STARTED_STATE

  def close(self):
    # type: () -> None
    with ControlConnection._lock:
      if self._state == BeamFnControlServicer.STARTED_STATE:
        self.push(BeamFnControlServicer._DONE_MARKER)
        self._read_thread.join()
      self._state = BeamFnControlServicer.DONE_STATE

  def abort(self, exn):
    # type: (Exception) -> None
    for future in self._futures_by_id.values():
      future.abort(exn)


class BeamFnControlServicer(beam_fn_api_pb2_grpc.BeamFnControlServicer):
  """Implementation of BeamFnControlServicer for clients."""

  UNSTARTED_STATE = 'unstarted'
  STARTED_STATE = 'started'
  DONE_STATE = 'done'

  _DONE_MARKER = Sentinel.sentinel

  def __init__(
      self,
      worker_manager,  # type: WorkerHandlerManager
  ):
    # type: (...) -> None
    self._worker_manager = worker_manager
    self._lock = threading.Lock()
    self._uid_counter = 0
    self._state = self.UNSTARTED_STATE
    # following self._req_* variables are used for debugging purpose, data is
    # added only when self._log_req is True.
    self._req_sent = collections.defaultdict(int)  # type: DefaultDict[str, int]
    self._log_req = logging.getLogger().getEffectiveLevel() <= logging.DEBUG
    self._connections_by_worker_id = collections.defaultdict(
        ControlConnection)  # type: DefaultDict[str, ControlConnection]

  def get_conn_by_worker_id(self, worker_id):
    # type: (str) -> ControlConnection
    with self._lock:
      return self._connections_by_worker_id[worker_id]

  def Control(self,
              iterator,  # type: Iterable[beam_fn_api_pb2.InstructionResponse]
              context  # type: ServicerContext
             ):
    # type: (...) -> Iterator[beam_fn_api_pb2.InstructionRequest]
    with self._lock:
      if self._state == self.DONE_STATE:
        return
      else:
        self._state = self.STARTED_STATE

    worker_id = dict(context.invocation_metadata()).get('worker_id')
    if not worker_id:
      raise RuntimeError(
          'All workers communicate through gRPC should have '
          'worker_id. Received None.')

    control_conn = self.get_conn_by_worker_id(worker_id)
    control_conn.set_input(iterator)

    while True:
      to_push = control_conn.get_req()
      if to_push is self._DONE_MARKER:
        return
      yield to_push
      if self._log_req:
        self._req_sent[to_push.instruction_id] += 1

  def done(self):
    # type: () -> None
    self._state = self.DONE_STATE
    _LOGGER.debug(
        'Runner: Requests sent by runner: %s',
        [(str(req), cnt) for req, cnt in self._req_sent.items()])

  def GetProcessBundleDescriptor(self, id, context=None):
    # type: (beam_fn_api_pb2.GetProcessBundleDescriptorRequest, Any) -> beam_fn_api_pb2.ProcessBundleDescriptor
    return self._worker_manager.get_process_bundle_descriptor(id)


class WorkerHandler(object):
  """worker_handler for a worker.

  It provides utilities to start / stop the worker, provision any resources for
  it, as well as provide descriptors for the data, state and logging APIs for
  it.
  """

  _registered_environments = {}  # type: Dict[str, Tuple[ConstructorFn, type]]
  _worker_id_counter = -1
  _lock = threading.Lock()

  control_conn = None  # type: ControlConnection
  data_conn = None  # type: data_plane._GrpcDataChannel

  def __init__(self,
               control_handler,  # type: Any
               data_plane_handler,  # type: Any
               state,  # type: sdk_worker.StateHandler
               provision_info  # type: ExtendedProvisionInfo
              ):
    # type: (...) -> None

    """Initialize a WorkerHandler.

    Args:
      control_handler:
      data_plane_handler (data_plane.DataChannel):
      state:
      provision_info:
    """
    self.control_handler = control_handler
    self.data_plane_handler = data_plane_handler
    self.state = state
    self.provision_info = provision_info

    with WorkerHandler._lock:
      WorkerHandler._worker_id_counter += 1
      self.worker_id = 'worker_%s' % WorkerHandler._worker_id_counter

  def close(self):
    # type: () -> None
    self.stop_worker()

  def start_worker(self):
    # type: () -> None
    raise NotImplementedError

  def stop_worker(self):
    # type: () -> None
    raise NotImplementedError

  def control_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    raise NotImplementedError

  def artifact_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    raise NotImplementedError

  def data_api_service_descriptor(self):
    # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor]
    raise NotImplementedError

  def state_api_service_descriptor(self):
    # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor]
    raise NotImplementedError

  def logging_api_service_descriptor(self):
    # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor]
    raise NotImplementedError

  @classmethod
  def register_environment(
      cls,
      urn,  # type: str
      payload_type  # type: Optional[Type[T]]
  ):
    # type: (...) -> Callable[[Type[WorkerHandler]], Callable[[T, sdk_worker.StateHandler, ExtendedProvisionInfo, GrpcServer], WorkerHandler]]
    def wrapper(constructor):
      # type: (Callable) -> Callable
      cls._registered_environments[urn] = constructor, payload_type  # type: ignore[assignment]
      return constructor

    return wrapper

  @classmethod
  def create(cls,
             environment,  # type: beam_runner_api_pb2.Environment
             state,  # type: sdk_worker.StateHandler
             provision_info,  # type: ExtendedProvisionInfo
             grpc_server  # type: GrpcServer
            ):
    # type: (...) -> WorkerHandler
    constructor, payload_type = cls._registered_environments[environment.urn]
    return constructor(
        proto_utils.parse_Bytes(environment.payload, payload_type),
        state,
        provision_info,
        grpc_server)


# This takes a WorkerHandlerManager instead of GrpcServer, so it is not
# compatible with WorkerHandler.register_environment.  There is a special case
# in WorkerHandlerManager.get_worker_handlers() that allows it to work.
@WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON, None)
class EmbeddedWorkerHandler(WorkerHandler):
  """An in-memory worker_handler for fn API control, state and data planes."""

  def __init__(self,
               unused_payload,  # type: None
               state,  # type: sdk_worker.StateHandler
               provision_info,  # type: ExtendedProvisionInfo
               worker_manager,  # type: WorkerHandlerManager
              ):
    # type: (...) -> None
    super().__init__(
        self, data_plane.InMemoryDataChannel(), state, provision_info)
    self.control_conn = self  # type: ignore  # need Protocol to describe this
    self.data_conn = self.data_plane_handler
    state_cache = StateCache(STATE_CACHE_SIZE_MB * MB_TO_BYTES)
    self.bundle_processor_cache = sdk_worker.BundleProcessorCache(
        FNAPI_RUNNER_CAPABILITIES,
        SingletonStateHandlerFactory(
            sdk_worker.GlobalCachingStateHandler(state_cache, state)),
        data_plane.InMemoryDataChannelFactory(
            self.data_plane_handler.inverse()),
        worker_manager._process_bundle_descriptors)
    self.worker = sdk_worker.SdkWorker(self.bundle_processor_cache)
    self._uid_counter = 0

  def push(self, request):
    # type: (beam_fn_api_pb2.InstructionRequest) -> ControlFuture
    if not request.instruction_id:
      self._uid_counter += 1
      request.instruction_id = 'control_%s' % self._uid_counter
    response = self.worker.do_instruction(request)
    return ControlFuture(request.instruction_id, response)

  def start_worker(self):
    # type: () -> None
    pass

  def stop_worker(self):
    # type: () -> None
    self.bundle_processor_cache.shutdown()

  def done(self):
    # type: () -> None
    pass

  def data_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    # A fake endpoint is needed for properly constructing timer info map in
    # bundle_processor for fnapi_runner.
    return endpoints_pb2.ApiServiceDescriptor(url='fake')

  def state_api_service_descriptor(self):
    # type: () -> None
    return None

  def logging_api_service_descriptor(self):
    # type: () -> None
    return None


class BasicLoggingService(beam_fn_api_pb2_grpc.BeamFnLoggingServicer):
  def Logging(self, log_messages, context=None):
    # type: (Iterable[beam_fn_api_pb2.LogEntry.List], Any) -> Iterator[beam_fn_api_pb2.LogControl]
    yield beam_fn_api_pb2.LogControl()
    for log_message in log_messages:
      for log in log_message.log_entries:
        logging.log(LOGENTRY_TO_LOG_LEVEL_MAP[log.severity], str(log))


class BasicProvisionService(beam_provision_api_pb2_grpc.ProvisionServiceServicer
                            ):
  def __init__(self, base_info, worker_manager):
    # type: (beam_provision_api_pb2.ProvisionInfo, WorkerHandlerManager) -> None
    self._base_info = base_info
    self._worker_manager = worker_manager

  def GetProvisionInfo(self, request, context=None):
    # type: (Any, Optional[ServicerContext]) -> beam_provision_api_pb2.GetProvisionInfoResponse
    if context:
      worker_id = dict(context.invocation_metadata())['worker_id']
      worker = self._worker_manager.get_worker(worker_id)
      info = copy.copy(worker.provision_info.provision_info)
      info.logging_endpoint.CopyFrom(worker.logging_api_service_descriptor())
      info.artifact_endpoint.CopyFrom(worker.artifact_api_service_descriptor())
      info.control_endpoint.CopyFrom(worker.control_api_service_descriptor())
    else:
      info = self._base_info
    info.runner_capabilities[:] = FNAPI_RUNNER_CAPABILITIES
    return beam_provision_api_pb2.GetProvisionInfoResponse(info=info)


class GrpcServer(object):

  _DEFAULT_SHUTDOWN_TIMEOUT_SECS = 5

  def __init__(self,
               state,  # type: StateServicer
               provision_info,  # type: Optional[ExtendedProvisionInfo]
               worker_manager,  # type: WorkerHandlerManager
              ):
    # type: (...) -> None

    # Options to have no limits (-1) on the size of the messages
    # received or sent over the data plane. The actual buffer size
    # is controlled in a layer above. Also, options to keep the server alive
    # when too many pings are received.
    options = [("grpc.max_receive_message_length", -1),
               ("grpc.max_send_message_length", -1),
               ("grpc.http2.max_pings_without_data", 0),
               ("grpc.http2.max_ping_strikes", 0)]

    self.state = state
    self.provision_info = provision_info
    self.control_server = grpc.server(
        thread_pool_executor.shared_unbounded_instance(), options=options)
    self.control_port = self.control_server.add_insecure_port('[::]:0')
    self.control_address = 'localhost:%s' % self.control_port

    self.data_server = grpc.server(
        thread_pool_executor.shared_unbounded_instance(), options=options)
    self.data_port = self.data_server.add_insecure_port('[::]:0')

    self.state_server = grpc.server(
        thread_pool_executor.shared_unbounded_instance(), options=options)
    self.state_port = self.state_server.add_insecure_port('[::]:0')

    self.control_handler = BeamFnControlServicer(worker_manager)
    beam_fn_api_pb2_grpc.add_BeamFnControlServicer_to_server(
        self.control_handler, self.control_server)

    # If we have provision info, serve these off the control port as well.
    if self.provision_info:
      if self.provision_info.provision_info:
        beam_provision_api_pb2_grpc.add_ProvisionServiceServicer_to_server(
            BasicProvisionService(
                self.provision_info.provision_info, worker_manager),
            self.control_server)

      def open_uncompressed(f):
        # type: (str) -> BinaryIO
        return filesystems.FileSystems.open(
            f, compression_type=CompressionTypes.UNCOMPRESSED)

      beam_artifact_api_pb2_grpc.add_ArtifactRetrievalServiceServicer_to_server(
          artifact_service.ArtifactRetrievalService(
              file_reader=open_uncompressed),
          self.control_server)

    self.data_plane_handler = data_plane.BeamFnDataServicer(
        DATA_BUFFER_TIME_LIMIT_MS)
    beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(
        self.data_plane_handler, self.data_server)

    beam_fn_api_pb2_grpc.add_BeamFnStateServicer_to_server(
        GrpcStateServicer(state), self.state_server)

    self.logging_server = grpc.server(
        thread_pool_executor.shared_unbounded_instance(), options=options)
    self.logging_port = self.logging_server.add_insecure_port('[::]:0')
    beam_fn_api_pb2_grpc.add_BeamFnLoggingServicer_to_server(
        BasicLoggingService(), self.logging_server)

    _LOGGER.info('starting control server on port %s', self.control_port)
    _LOGGER.info('starting data server on port %s', self.data_port)
    _LOGGER.info('starting state server on port %s', self.state_port)
    _LOGGER.info('starting logging server on port %s', self.logging_port)
    self.logging_server.start()
    self.state_server.start()
    self.data_server.start()
    self.control_server.start()

  def close(self):
    # type: () -> None
    self.control_handler.done()
    to_wait = [
        self.control_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS),
        self.data_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS),
        self.state_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS),
        self.logging_server.stop(self._DEFAULT_SHUTDOWN_TIMEOUT_SECS)
    ]
    for w in to_wait:
      w.wait()


class GrpcWorkerHandler(WorkerHandler):
  """An grpc based worker_handler for fn API control, state and data planes."""

  def __init__(self,
               state,  # type: StateServicer
               provision_info,  # type: ExtendedProvisionInfo
               grpc_server  # type: GrpcServer
              ):
    # type: (...) -> None
    self._grpc_server = grpc_server
    super().__init__(
        self._grpc_server.control_handler,
        self._grpc_server.data_plane_handler,
        state,
        provision_info)
    self.state = state

    self.control_address = self.port_from_worker(self._grpc_server.control_port)
    self.control_conn = self._grpc_server.control_handler.get_conn_by_worker_id(
        self.worker_id)

    self.data_conn = self._grpc_server.data_plane_handler.get_conn_by_worker_id(
        self.worker_id)

  def control_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    return endpoints_pb2.ApiServiceDescriptor(
        url=self.port_from_worker(self._grpc_server.control_port))

  def artifact_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    return endpoints_pb2.ApiServiceDescriptor(
        url=self.port_from_worker(self._grpc_server.control_port))

  def data_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    return endpoints_pb2.ApiServiceDescriptor(
        url=self.port_from_worker(self._grpc_server.data_port))

  def state_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    return endpoints_pb2.ApiServiceDescriptor(
        url=self.port_from_worker(self._grpc_server.state_port))

  def logging_api_service_descriptor(self):
    # type: () -> endpoints_pb2.ApiServiceDescriptor
    return endpoints_pb2.ApiServiceDescriptor(
        url=self.port_from_worker(self._grpc_server.logging_port))

  def close(self):
    # type: () -> None
    self.control_conn.close()
    self.data_conn.close()
    super().close()

  def port_from_worker(self, port):
    # type: (int) -> str
    return '%s:%s' % (self.host_from_worker(), port)

  def host_from_worker(self):
    # type: () -> str
    return 'localhost'


@WorkerHandler.register_environment(
    common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload)
class ExternalWorkerHandler(GrpcWorkerHandler):
  def __init__(self,
               external_payload,  # type: beam_runner_api_pb2.ExternalPayload
               state,  # type: StateServicer
               provision_info,  # type: ExtendedProvisionInfo
               grpc_server  # type: GrpcServer
              ):
    # type: (...) -> None
    super().__init__(state, provision_info, grpc_server)
    self._external_payload = external_payload

  def start_worker(self):
    # type: () -> None
    _LOGGER.info("Requesting worker at %s", self._external_payload.endpoint.url)
    stub = beam_fn_api_pb2_grpc.BeamFnExternalWorkerPoolStub(
        GRPCChannelFactory.insecure_channel(
            self._external_payload.endpoint.url))
    control_descriptor = endpoints_pb2.ApiServiceDescriptor(
        url=self.control_address)
    response = stub.StartWorker(
        beam_fn_api_pb2.StartWorkerRequest(
            worker_id=self.worker_id,
            control_endpoint=control_descriptor,
            artifact_endpoint=control_descriptor,
            provision_endpoint=control_descriptor,
            logging_endpoint=self.logging_api_service_descriptor(),
            params=self._external_payload.params))
    if response.error:
      raise RuntimeError("Error starting worker: %s" % response.error)

  def stop_worker(self):
    # type: () -> None
    pass

  def host_from_worker(self):
    # type: () -> str
    # TODO(https://github.com/apache/beam/issues/19947): Reconcile across
    # platforms.
    if sys.platform in ['win32', 'darwin']:
      return 'localhost'
    import socket
    return socket.getfqdn()


@WorkerHandler.register_environment(python_urns.EMBEDDED_PYTHON_GRPC, bytes)
class EmbeddedGrpcWorkerHandler(GrpcWorkerHandler):
  def __init__(self,
               payload,  # type: bytes
               state,  # type: StateServicer
               provision_info,  # type: ExtendedProvisionInfo
               grpc_server  # type: GrpcServer
              ):
    # type: (...) -> None
    super().__init__(state, provision_info, grpc_server)

    from apache_beam.transforms.environments import EmbeddedPythonGrpcEnvironment
    config = EmbeddedPythonGrpcEnvironment.parse_config(payload.decode('utf-8'))
    self._state_cache_size = (
        config.get('state_cache_size') or STATE_CACHE_SIZE_MB) << 20
    self._data_buffer_time_limit_ms = \
        config.get('data_buffer_time_limit_ms') or DATA_BUFFER_TIME_LIMIT_MS

  def start_worker(self):
    # type: () -> None
    self.worker = sdk_worker.SdkHarness(
        self.control_address,
        state_cache_size=self._state_cache_size,
        data_buffer_time_limit_ms=self._data_buffer_time_limit_ms,
        worker_id=self.worker_id,
        runner_capabilities=FNAPI_RUNNER_CAPABILITIES)
    self.worker_thread = threading.Thread(
        name='run_worker', target=self.worker.run)
    self.worker_thread.daemon = True
    self.worker_thread.start()

  def stop_worker(self):
    # type: () -> None
    self.worker_thread.join()


# The subprocesses module is not threadsafe on Python 2.7. Use this lock to
# prevent concurrent calls to Popen().
SUBPROCESS_LOCK = threading.Lock()


@WorkerHandler.register_environment(python_urns.SUBPROCESS_SDK, bytes)
class SubprocessSdkWorkerHandler(GrpcWorkerHandler):
  def __init__(self,
               worker_command_line,  # type: bytes
               state,  # type: StateServicer
               provision_info,  # type: ExtendedProvisionInfo
               grpc_server  # type: GrpcServer
              ):
    # type: (...) -> None
    super().__init__(state, provision_info, grpc_server)
    self._worker_command_line = worker_command_line

  def start_worker(self):
    # type: () -> None
    from apache_beam.runners.portability import local_job_service
    self.worker = local_job_service.SubprocessSdkWorker(
        self._worker_command_line,
        self.control_address,
        self.provision_info,
        self.worker_id)
    self.worker_thread = threading.Thread(
        name='run_worker', target=self.worker.run)
    self.worker_thread.start()

  def stop_worker(self):
    # type: () -> None
    self.worker_thread.join()


@WorkerHandler.register_environment(
    common_urns.environments.DOCKER.urn, beam_runner_api_pb2.DockerPayload)
class DockerSdkWorkerHandler(GrpcWorkerHandler):
  def __init__(self,
               payload,  # type: beam_runner_api_pb2.DockerPayload
               state,  # type: StateServicer
               provision_info,  # type: ExtendedProvisionInfo
               grpc_server  # type: GrpcServer
              ):
    # type: (...) -> None
    super().__init__(state, provision_info, grpc_server)
    self._container_image = payload.container_image
    self._container_id = None  # type: Optional[bytes]

  def host_from_worker(self):
    # type: () -> str
    if sys.platform == 'darwin':
      # See https://docs.docker.com/docker-for-mac/networking/
      return 'host.docker.internal'
    if sys.platform == 'linux' and is_in_notebook():
      import socket
      # Gets ipv4 address of current host. Note the host is not guaranteed to
      # be localhost because the python SDK could be running within a container.
      return socket.gethostbyname(socket.getfqdn())
    return super().host_from_worker()

  def start_worker(self):
    # type: () -> None
    credential_options = []
    try:
      # This is the public facing API, skip if it is not available.
      # (If this succeeds but the imports below fail, better to actually raise
      # an error below rather than silently fail.)
      # pylint: disable=unused-import
      import google.auth
    except ImportError:
      pass
    else:
      from google.auth import environment_vars
      from google.auth import _cloud_sdk
      gcloud_cred_file = os.environ.get(
          environment_vars.CREDENTIALS,
          _cloud_sdk.get_application_default_credentials_path())
      if os.path.exists(gcloud_cred_file):
        docker_cred_file = '/docker_cred_file.json'
        credential_options.extend([
            '--mount',
            f'type=bind,source={gcloud_cred_file},target={docker_cred_file}',
            '--env',
            f'{environment_vars.CREDENTIALS}={docker_cred_file}'
        ])
    with SUBPROCESS_LOCK:
      try:
        _LOGGER.info('Attempting to pull image %s', self._container_image)
        subprocess.check_call(['docker', 'pull', self._container_image])
      except Exception:
        _LOGGER.info(
            'Unable to pull image %s, defaulting to local image if it exists' %
            self._container_image)
      self._container_id = subprocess.check_output([
          'docker',
          'run',
          '-d',
          '--network=host',
      ] + credential_options + [
          self._container_image,
          '--id=%s' % self.worker_id,
          '--logging_endpoint=%s' % self.logging_api_service_descriptor().url,
          '--control_endpoint=%s' % self.control_address,
          '--artifact_endpoint=%s' % self.control_address,
          '--provision_endpoint=%s' % self.control_address,
      ]).strip()
      assert self._container_id is not None
      while True:
        status = subprocess.check_output([
            'docker', 'inspect', '-f', '{{.State.Status}}', self._container_id
        ]).strip()
        _LOGGER.info(
            'Waiting for docker to start up. Current status is %s' %
            status.decode('utf-8'))
        if status == b'running':
          _LOGGER.info(
              'Docker container is running. container_id = %s, '
              'worker_id = %s',
              self._container_id,
              self.worker_id)
          break
        elif status in (b'dead', b'exited'):
          subprocess.call(['docker', 'container', 'logs', self._container_id])
          raise RuntimeError(
              'SDK failed to start. Final status is %s' %
              status.decode('utf-8'))
      time.sleep(1)
    self._done = False
    t = threading.Thread(target=self.watch_container)
    t.daemon = True
    t.start()

  def watch_container(self):
    # type: () -> None
    while not self._done:
      assert self._container_id is not None
      status = subprocess.check_output(
          ['docker', 'inspect', '-f', '{{.State.Status}}',
           self._container_id]).strip()
      if status != b'running':
        if not self._done:
          logs = subprocess.check_output([
              'docker', 'container', 'logs', '--tail', '10', self._container_id
          ],
                                         stderr=subprocess.STDOUT)
          _LOGGER.info(logs)
          self.control_conn.abort(
              RuntimeError(
                  'SDK exited unexpectedly. '
                  'Final status is %s. Final log line is %s' % (
                      status.decode('utf-8'),
                      logs.decode('utf-8').strip().rsplit('\n',
                                                          maxsplit=1)[-1])))
      time.sleep(5)

  def stop_worker(self):
    # type: () -> None
    self._done = True
    if self._container_id:
      with SUBPROCESS_LOCK:
        subprocess.call(['docker', 'kill', self._container_id])


class WorkerHandlerManager(object):
  """
  Manages creation of ``WorkerHandler``s.

  Caches ``WorkerHandler``s based on environment id.
  """
  def __init__(self,
               environments,  # type: Mapping[str, beam_runner_api_pb2.Environment]
               job_provision_info  # type: ExtendedProvisionInfo
              ):
    # type: (...) -> None
    self._environments = environments
    self._job_provision_info = job_provision_info
    self._cached_handlers = collections.defaultdict(
        list)  # type: DefaultDict[str, List[WorkerHandler]]
    self._workers_by_id = {}  # type: Dict[str, WorkerHandler]
    self.state_servicer = StateServicer()
    self._grpc_server = None  # type: Optional[GrpcServer]
    self._process_bundle_descriptors = {
    }  # type: Dict[str, beam_fn_api_pb2.ProcessBundleDescriptor]

  def register_process_bundle_descriptor(self, process_bundle_descriptor):
    # type: (beam_fn_api_pb2.ProcessBundleDescriptor) -> None
    self._process_bundle_descriptors[
        process_bundle_descriptor.id] = process_bundle_descriptor

  def get_process_bundle_descriptor(self, request):
    # type: (beam_fn_api_pb2.GetProcessBundleDescriptorRequest) -> beam_fn_api_pb2.ProcessBundleDescriptor
    return self._process_bundle_descriptors[
        request.process_bundle_descriptor_id]

  def get_worker_handlers(
      self,
      environment_id,  # type: Optional[str]
      num_workers  # type: int
  ):
    # type: (...) -> List[WorkerHandler]
    if environment_id is None:
      # Any environment will do, pick one arbitrarily.
      environment_id = next(iter(self._environments.keys()))
    environment = self._environments[environment_id]

    if environment.urn == common_urns.environments.ANYOF.urn:
      payload = beam_runner_api_pb2.AnyOfEnvironmentPayload.FromString(
          environment.payload)
      env_rankings = {
          python_urns.EMBEDDED_PYTHON: 10,
          common_urns.environments.EXTERNAL.urn: 5,
          common_urns.environments.DOCKER.urn: 1,
      }
      environment = sorted(
          payload.environments,
          key=lambda env: env_rankings.get(env.urn, -1))[-1]

    # assume all environments except EMBEDDED_PYTHON use gRPC.
    if environment.urn == python_urns.EMBEDDED_PYTHON:
      # special case for EmbeddedWorkerHandler: there's no need for a gRPC
      # server, but we need to pass self instead.  Cast to make the type check
      # on WorkerHandler.create() think we have a GrpcServer instance.
      grpc_server = cast(GrpcServer, self)
    elif self._grpc_server is None:
      self._grpc_server = GrpcServer(
          self.state_servicer, self._job_provision_info, self)
      grpc_server = self._grpc_server
    else:
      grpc_server = self._grpc_server

    worker_handler_list = self._cached_handlers[environment_id]
    if len(worker_handler_list) < num_workers:
      for _ in range(len(worker_handler_list), num_workers):
        worker_handler = WorkerHandler.create(
            environment,
            self.state_servicer,
            self._job_provision_info.for_environment(environment),
            grpc_server)
        _LOGGER.info(
            "Created Worker handler %s for environment %s (%s, %r)",
            worker_handler,
            environment_id,
            environment.urn,
            environment.payload)
        self._cached_handlers[environment_id].append(worker_handler)
        self._workers_by_id[worker_handler.worker_id] = worker_handler
        worker_handler.start_worker()
    return self._cached_handlers[environment_id][:num_workers]

  def close_all(self):
    # type: () -> None
    for worker_handler_list in self._cached_handlers.values():
      for worker_handler in set(worker_handler_list):
        try:
          worker_handler.close()
        except Exception:
          _LOGGER.error(
              "Error closing worker_handler %s" % worker_handler, exc_info=True)
    self._cached_handlers = {}  # type: ignore[assignment]
    self._workers_by_id = {}
    if self._grpc_server is not None:
      self._grpc_server.close()
      self._grpc_server = None

  def get_worker(self, worker_id):
    # type: (str) -> WorkerHandler
    return self._workers_by_id[worker_id]


class StateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer,
                    sdk_worker.StateHandler):
  _SUPPORTED_STATE_TYPES = frozenset([
      'runner',
      'multimap_side_input',
      'multimap_keys_side_input',
      'multimap_keys_values_side_input',
      'iterable_side_input',
      'bag_user_state',
      'multimap_user_state'
  ])

  class CopyOnWriteState(object):
    def __init__(self, underlying):
      # type: (DefaultDict[bytes, Buffer]) -> None
      self._underlying = underlying
      self._overlay = {}  # type: Dict[bytes, Buffer]

    def __getitem__(self, key):
      # type: (bytes) -> Buffer
      if key in self._overlay:
        return self._overlay[key]
      else:
        return StateServicer.CopyOnWriteList(
            self._underlying, self._overlay, key)

    def __delitem__(self, key):
      # type: (bytes) -> None
      self._overlay[key] = []

    def commit(self):
      # type: () -> DefaultDict[bytes, Buffer]
      self._underlying.update(self._overlay)
      return self._underlying

  class CopyOnWriteList(object):
    def __init__(self,
        underlying,  # type: DefaultDict[bytes, Buffer]
        overlay,  # type: Dict[bytes, Buffer]
        key  # type: bytes
    ):
      # type: (...) -> None
      self._underlying = underlying
      self._overlay = overlay
      self._key = key

    def __iter__(self):
      # type: () -> Iterator[bytes]
      if self._key in self._overlay:
        return iter(self._overlay[self._key])
      else:
        return iter(self._underlying[self._key])

    def append(self, item):
      # type: (bytes) -> None
      if self._key not in self._overlay:
        self._overlay[self._key] = list(self._underlying[self._key])
      self._overlay[self._key].append(item)

    def extend(self, other: Buffer) -> None:
      raise NotImplementedError()

  StateType = Union[CopyOnWriteState, DefaultDict[bytes, Buffer]]

  def __init__(self):
    # type: () -> None
    self._lock = threading.Lock()
    self._state = collections.defaultdict(list)  # type: StateServicer.StateType
    self._checkpoint = None  # type: Optional[StateServicer.StateType]
    self._use_continuation_tokens = False
    self._continuations = {}  # type: Dict[bytes, Tuple[bytes, ...]]

  def checkpoint(self):
    # type: () -> None
    assert self._checkpoint is None and not \
      isinstance(self._state, StateServicer.CopyOnWriteState)
    self._checkpoint = self._state
    self._state = StateServicer.CopyOnWriteState(self._state)

  def commit(self):
    # type: () -> None
    assert isinstance(self._state,
                      StateServicer.CopyOnWriteState) and \
           isinstance(self._checkpoint,
                      StateServicer.CopyOnWriteState)
    self._state.commit()
    self._state = self._checkpoint.commit()
    self._checkpoint = None

  def restore(self):
    # type: () -> None
    assert self._checkpoint is not None
    self._state = self._checkpoint
    self._checkpoint = None

  @contextlib.contextmanager
  def process_instruction_id(self, unused_instruction_id):
    # type: (Any) -> Iterator
    yield

  def get_raw(self,
      state_key,  # type: beam_fn_api_pb2.StateKey
      continuation_token=None  # type: Optional[bytes]
              ):
    # type: (...) -> Tuple[bytes, Optional[bytes]]

    if state_key.WhichOneof('type') not in self._SUPPORTED_STATE_TYPES:
      raise NotImplementedError(
          'Unknown state type: ' + state_key.WhichOneof('type'))

    with self._lock:
      full_state = self._state[self._to_key(state_key)]
      if self._use_continuation_tokens:
        # The token is "nonce:index".
        if not continuation_token:
          token_base = b'token_%x' % len(self._continuations)
          self._continuations[token_base] = tuple(full_state)
          return b'', b'%s:0' % token_base
        else:
          token_base, index = continuation_token.split(b':')
          ix = int(index)
          full_state_cont = self._continuations[token_base]
          if ix == len(full_state_cont):
            return b'', None
          else:
            return full_state_cont[ix], b'%s:%d' % (token_base, ix + 1)
      else:
        assert not continuation_token
        return b''.join(full_state), None

  def append_raw(
      self,
      state_key,  # type: beam_fn_api_pb2.StateKey
      data  # type: bytes
  ):
    # type: (...) -> _Future
    with self._lock:
      self._state[self._to_key(state_key)].append(data)
    return _Future.done()

  def clear(self, state_key):
    # type: (beam_fn_api_pb2.StateKey) -> _Future
    with self._lock:
      try:
        del self._state[self._to_key(state_key)]
      except KeyError:
        # This may happen with the caching layer across bundles. Caching may
        # skip this storage layer for a blocking_get(key) request. Without
        # the caching, the state for a key would be initialized via the
        # defaultdict that _state uses.
        pass
    return _Future.done()

  def done(self):
    # type: () -> None
    pass

  @staticmethod
  def _to_key(state_key):
    # type: (beam_fn_api_pb2.StateKey) -> bytes
    return state_key.SerializeToString()


class GrpcStateServicer(beam_fn_api_pb2_grpc.BeamFnStateServicer):
  def __init__(self, state):
    # type: (StateServicer) -> None
    self._state = state

  def State(self,
      request_stream,  # type: Iterable[beam_fn_api_pb2.StateRequest]
      context=None  # type: Any
            ):
    # type: (...) -> Iterator[beam_fn_api_pb2.StateResponse]
    # Note that this eagerly mutates state, assuming any failures are fatal.
    # Thus it is safe to ignore instruction_id.
    for request in request_stream:
      try:
        request_type = request.WhichOneof('request')
        if request_type == 'get':
          data, continuation_token = self._state.get_raw(
              request.state_key, request.get.continuation_token)
          yield beam_fn_api_pb2.StateResponse(
              id=request.id,
              get=beam_fn_api_pb2.StateGetResponse(
                  data=data, continuation_token=continuation_token))
        elif request_type == 'append':
          self._state.append_raw(request.state_key, request.append.data)
          yield beam_fn_api_pb2.StateResponse(
              id=request.id, append=beam_fn_api_pb2.StateAppendResponse())
        elif request_type == 'clear':
          self._state.clear(request.state_key)
          yield beam_fn_api_pb2.StateResponse(
              id=request.id, clear=beam_fn_api_pb2.StateClearResponse())
        else:
          raise NotImplementedError('Unknown state request: %s' % request_type)
      except Exception as exn:
        yield beam_fn_api_pb2.StateResponse(id=request.id, error=str(exn))


class SingletonStateHandlerFactory(sdk_worker.StateHandlerFactory):
  """A singleton cache for a StateServicer."""
  def __init__(self, state_handler):
    # type: (sdk_worker.CachingStateHandler) -> None
    self._state_handler = state_handler

  def create_state_handler(self, api_service_descriptor):
    # type: (endpoints_pb2.ApiServiceDescriptor) -> sdk_worker.CachingStateHandler

    """Returns the singleton state handler."""
    return self._state_handler

  def close(self):
    # type: () -> None

    """Does nothing."""
    pass


class ControlFuture(object):
  def __init__(self,
               instruction_id,  # type: str
               response=None  # type: Optional[beam_fn_api_pb2.InstructionResponse]
              ):
    # type: (...) -> None
    self.instruction_id = instruction_id
    self._response = response
    if response is None:
      self._condition = threading.Condition()
    self._exception = None  # type: Optional[Exception]

  def is_done(self):
    # type: () -> bool
    return self._response is not None

  def set(self, response):
    # type: (beam_fn_api_pb2.InstructionResponse) -> None
    with self._condition:
      self._response = response
      self._condition.notify_all()

  def get(self, timeout=None):
    # type: (Optional[float]) -> beam_fn_api_pb2.InstructionResponse
    if not self._response and not self._exception:
      with self._condition:
        if not self._response and not self._exception:
          self._condition.wait(timeout)
    if self._exception:
      raise self._exception
    else:
      assert self._response is not None
      return self._response

  def abort(self, exception):
    # type: (Exception) -> None
    with self._condition:
      self._exception = exception
      self._condition.notify_all()
