text: # Copyright 2020 The gRPC authors.
  #
  # Licensed 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.
  
  import argparse
  import collections
  from concurrent import futures
  import datetime
  import logging
  import signal
  import sys
  import threading
  import time
  from typing import DefaultDict, Dict, List, Mapping, Sequence, Set, Tuple
  
  import grpc
  import grpc_admin
  from grpc_channelz.v1 import channelz
  
  from src.proto.grpc.testing import empty_pb2
  from src.proto.grpc.testing import messages_pb2
  from src.proto.grpc.testing import test_pb2
  from src.proto.grpc.testing import test_pb2_grpc
  
  logger = logging.getLogger()
  console_handler = logging.StreamHandler()
  formatter = logging.Formatter(fmt='%(asctime)s: %(levelname)-8s %(message)s')
  console_handler.setFormatter(formatter)
  logger.addHandler(console_handler)
  
  _SUPPORTED_METHODS = (
      "UnaryCall",
      "EmptyCall",
  )
  
  _METHOD_CAMEL_TO_CAPS_SNAKE = {
      "UnaryCall": "UNARY_CALL",
      "EmptyCall": "EMPTY_CALL",
  }
  
  _METHOD_STR_TO_ENUM = {
      "UnaryCall": messages_pb2.ClientConfigureRequest.UNARY_CALL,
      "EmptyCall": messages_pb2.ClientConfigureRequest.EMPTY_CALL,
  }
  
  _METHOD_ENUM_TO_STR = {v: k for k, v in _METHOD_STR_TO_ENUM.items()}
  
  PerMethodMetadataType = Mapping[str, Sequence[Tuple[str, str]]]
  
  _CONFIG_CHANGE_TIMEOUT = datetime.timedelta(milliseconds=500)
  
  
  class _StatsWatcher:
      _start: int
      _end: int
      _rpcs_needed: int
      _rpcs_by_peer: DefaultDict[str, int]
      _rpcs_by_method: DefaultDict[str, DefaultDict[str, int]]
      _no_remote_peer: int
      _lock: threading.Lock
      _condition: threading.Condition
  
      def __init__(self, start: int, end: int):
          self._start = start
          self._end = end
          self._rpcs_needed = end - start
          self._rpcs_by_peer = collections.defaultdict(int)
          self._rpcs_by_method = collections.defaultdict(
              lambda: collections.defaultdict(int))
          self._condition = threading.Condition()
          self._no_remote_peer = 0
  
      def on_rpc_complete(self, request_id: int, peer: str, method: str) -> None:
          """Records statistics for a single RPC."""
          if self._start <= request_id < self._end:
              with self._condition:
                  if not peer:
                      self._no_remote_peer += 1
                  else:
                      self._rpcs_by_peer[peer] += 1
                      self._rpcs_by_method[method][peer] += 1
                  self._rpcs_needed -= 1
                  self._condition.notify()
  
      def await_rpc_stats_response(
              self, timeout_sec: int) -> messages_pb2.LoadBalancerStatsResponse:
          """Blocks until a full response has been collected."""
          with self._condition:
              self._condition.wait_for(lambda: not self._rpcs_needed,
                                       timeout=float(timeout_sec))
              response = messages_pb2.LoadBalancerStatsResponse()
              for peer, count in self._rpcs_by_peer.items():
                  response.rpcs_by_peer[peer] = count
              for method, count_by_peer in self._rpcs_by_method.items():
                  for peer, count in count_by_peer.items():
                      response.rpcs_by_method[method].rpcs_by_peer[peer] = count
              response.num_failures = self._no_remote_peer + self._rpcs_needed
          return response
  
  
  _global_lock = threading.Lock()
  _stop_event = threading.Event()
  _global_rpc_id: int = 0
  _watchers: Set[_StatsWatcher] = set()
  _global_server = None
  _global_rpcs_started: Mapping[str, int] = collections.defaultdict(int)
  _global_rpcs_succeeded: Mapping[str, int] = collections.defaultdict(int)
  _global_rpcs_failed: Mapping[str, int] = collections.defaultdict(int)
  
  # Mapping[method, Mapping[status_code, count]]
  _global_rpc_statuses: Mapping[str, Mapping[int, int]] = collections.defaultdict(
      lambda: collections.defaultdict(int))
  
  
  def _handle_sigint(sig, frame) -> None:
      _stop_event.set()
      _global_server.stop(None)
  
  
  class _LoadBalancerStatsServicer(test_pb2_grpc.LoadBalancerStatsServiceServicer
                                  ):
  
      def __init__(self):
          super(_LoadBalancerStatsServicer).__init__()
  
      def GetClientStats(
          self, request: messages_pb2.LoadBalancerStatsRequest,
          context: grpc.ServicerContext
      ) -> messages_pb2.LoadBalancerStatsResponse:
          logger.info("Received stats request.")
          start = None
          end = None
          watcher = None
          with _global_lock:
              start = _global_rpc_id + 1
              end = start + request.num_rpcs
              watcher = _StatsWatcher(start, end)
              _watchers.add(watcher)
          response = watcher.await_rpc_stats_response(request.timeout_sec)
          with _global_lock:
              _watchers.remove(watcher)
          logger.info("Returning stats response: %s", response)
          return response
  
      def GetClientAccumulatedStats(
          self, request: messages_pb2.LoadBalancerAccumulatedStatsRequest,
          context: grpc.ServicerContext
      ) -> messages_pb2.LoadBalancerAccumulatedStatsResponse:
          logger.info("Received cumulative stats request.")
          response = messages_pb2.LoadBalancerAccumulatedStatsResponse()
          with _global_lock:
              for method in _SUPPORTED_METHODS:
                  caps_method = _METHOD_CAMEL_TO_CAPS_SNAKE[method]
                  response.num_rpcs_started_by_method[
                      caps_method] = _global_rpcs_started[method]
                  response.num_rpcs_succeeded_by_method[
                      caps_method] = _global_rpcs_succeeded[method]
                  response.num_rpcs_failed_by_method[
                      caps_method] = _global_rpcs_failed[method]
                  response.stats_per_method[
                      caps_method].rpcs_started = _global_rpcs_started[method]
                  for code, count in _global_rpc_statuses[method].items():
                      response.stats_per_method[caps_method].result[code] = count
          logger.info("Returning cumulative stats response.")
          return response
  
  
  def _start_rpc(method: str, metadata: Sequence[Tuple[str, str]],
                 request_id: int, stub: test_pb2_grpc.TestServiceStub,
                 timeout: float, futures: Mapping[int, Tuple[grpc.Future,
                                                             str]]) -> None:
      logger.debug(f"Sending {method} request to backend: {request_id}")
      if method == "UnaryCall":
          future = stub.UnaryCall.future(messages_pb2.SimpleRequest(),
                                         metadata=metadata,
                                         timeout=timeout)
      elif method == "EmptyCall":
          future = stub.EmptyCall.future(empty_pb2.Empty(),
                                         metadata=metadata,
                                         timeout=timeout)
      else:
          raise ValueError(f"Unrecognized method '{method}'.")
      futures[request_id] = (future, method)
  
  
  def _on_rpc_done(rpc_id: int, future: grpc.Future, method: str,
                   print_response: bool) -> None:
      exception = future.exception()
      hostname = ""
      _global_rpc_statuses[method][future.code().value[0]] += 1
      if exception is not None:
          with _global_lock:
              _global_rpcs_failed[method] += 1
          if exception.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
              logger.error(f"RPC {rpc_id} timed out")
          else:
              logger.error(exception)
      else:
          response = future.result()
          hostname = None
          for metadatum in future.initial_metadata():
              if metadatum[0] == "hostname":
                  hostname = metadatum[1]
                  break
          else:
              hostname = response.hostname
          if future.code() == grpc.StatusCode.OK:
              with _global_lock:
                  _global_rpcs_succeeded[method] += 1
          else:
              with _global_lock:
                  _global_rpcs_failed[method] += 1
          if print_response:
              if future.code() == grpc.StatusCode.OK:
                  logger.debug("Successful response.")
              else:
                  logger.debug(f"RPC failed: {call}")
      with _global_lock:
          for watcher in _watchers:
              watcher.on_rpc_complete(rpc_id, hostname, method)
  
  
  def _remove_completed_rpcs(futures: Mapping[int, grpc.Future],
                             print_response: bool) -> None:
      logger.debug("Removing completed RPCs")
      done = []
      for future_id, (future, method) in futures.items():
          if future.done():
              _on_rpc_done(future_id, future, method, args.print_response)
              done.append(future_id)
      for rpc_id in done:
          del futures[rpc_id]
  
  
  def _cancel_all_rpcs(futures: Mapping[int, Tuple[grpc.Future, str]]) -> None:
      logger.info("Cancelling all remaining RPCs")
      for future, _ in futures.values():
          future.cancel()
  
  
  class _ChannelConfiguration:
      """Configuration for a single client channel.
  
      Instances of this class are meant to be dealt with as PODs. That is,
      data member should be accessed directly. This class is not thread-safe.
      When accessing any of its members, the lock member should be held.
      """
  
      def __init__(self, method: str, metadata: Sequence[Tuple[str, str]],
                   qps: int, server: str, rpc_timeout_sec: int,
                   print_response: bool, secure_mode: bool):
          # condition is signalled when a change is made to the config.
          self.condition = threading.Condition()
  
          self.method = method
          self.metadata = metadata
          self.qps = qps
          self.server = server
          self.rpc_timeout_sec = rpc_timeout_sec
          self.print_response = print_response
          self.secure_mode = secure_mode
  
  
  def _run_single_channel(config: _ChannelConfiguration) -> None:
      global _global_rpc_id  # pylint: disable=global-statement
      with config.condition:
          server = config.server
      channel = None
      if config.secure_mode:
          fallback_creds = grpc.experimental.insecure_channel_credentials()
          channel_creds = grpc.xds_channel_credentials(fallback_creds)
          channel = grpc.secure_channel(server, channel_creds)
      else:
          channel = grpc.insecure_channel(server)
      with channel:
          stub = test_pb2_grpc.TestServiceStub(channel)
          futures: Dict[int, Tuple[grpc.Future, str]] = {}
          while not _stop_event.is_set():
              with config.condition:
                  if config.qps == 0:
                      config.condition.wait(
                          timeout=_CONFIG_CHANGE_TIMEOUT.total_seconds())
                      continue
                  else:
                      duration_per_query = 1.0 / float(config.qps)
              request_id = None
              with _global_lock:
                  request_id = _global_rpc_id
                  _global_rpc_id += 1
                  _global_rpcs_started[config.method] += 1
              start = time.time()
              end = start + duration_per_query
              with config.condition:
                  _start_rpc(config.method, config.metadata, request_id, stub,
                             float(config.rpc_timeout_sec), futures)
              with config.condition:
                  _remove_completed_rpcs(futures, config.print_response)
              logger.debug(f"Currently {len(futures)} in-flight RPCs")
              now = time.time()
              while now < end:
                  time.sleep(end - now)
                  now = time.time()
          _cancel_all_rpcs(futures)
  
  
  class _XdsUpdateClientConfigureServicer(
          test_pb2_grpc.XdsUpdateClientConfigureServiceServicer):
  
      def __init__(self, per_method_configs: Mapping[str, _ChannelConfiguration],
                   qps: int):
          super(_XdsUpdateClientConfigureServicer).__init__()
          self._per_method_configs = per_method_configs
          self._qps = qps
  
      def Configure(
              self, request: messages_pb2.ClientConfigureRequest,
              context: grpc.ServicerContext
      ) -> messages_pb2.ClientConfigureResponse:
          logger.info("Received Configure RPC: %s", request)
          method_strs = (_METHOD_ENUM_TO_STR[t] for t in request.types)
          for method in _SUPPORTED_METHODS:
              method_enum = _METHOD_STR_TO_ENUM[method]
              channel_config = self._per_method_configs[method]
              if method in method_strs:
                  qps = self._qps
                  metadata = ((md.key, md.value)
                              for md in request.metadata
                              if md.type == method_enum)
                  # For backward compatibility, do not change timeout when we
                  # receive a default value timeout.
                  if request.timeout_sec == 0:
                      timeout_sec = channel_config.rpc_timeout_sec
                  else:
                      timeout_sec = request.timeout_sec
              else:
                  qps = 0
                  metadata = ()
                  # Leave timeout unchanged for backward compatibility.
                  timeout_sec = channel_config.rpc_timeout_sec
              with channel_config.condition:
                  channel_config.qps = qps
                  channel_config.metadata = list(metadata)
                  channel_config.rpc_timeout_sec = timeout_sec
                  channel_config.condition.notify_all()
          return messages_pb2.ClientConfigureResponse()
  
  
  class _MethodHandle:
      """An object grouping together threads driving RPCs for a method."""
  
      _channel_threads: List[threading.Thread]
  
      def __init__(self, num_channels: int,
                   channel_config: _ChannelConfiguration):
          """Creates and starts a group of threads running the indicated method."""
          self._channel_threads = []
          for i in range(num_channels):
              thread = threading.Thread(target=_run_single_channel,
                                        args=(channel_config,))
              thread.start()
              self._channel_threads.append(thread)
  
      def stop(self) -> None:
          """Joins all threads referenced by the handle."""
          for channel_thread in self._channel_threads:
              channel_thread.join()
  
  
  def _run(args: argparse.Namespace, methods: Sequence[str],
           per_method_metadata: PerMethodMetadataType) -> None:
      logger.info("Starting python xDS Interop Client.")
      global _global_server  # pylint: disable=global-statement
      method_handles = []
      channel_configs = {}
      for method in _SUPPORTED_METHODS:
          if method in methods:
              qps = args.qps
          else:
              qps = 0
          channel_config = _ChannelConfiguration(
              method, per_method_metadata.get(method, []), qps, args.server,
              args.rpc_timeout_sec, args.print_response, args.secure_mode)
          channel_configs[method] = channel_config
          method_handles.append(_MethodHandle(args.num_channels, channel_config))
      _global_server = grpc.server(futures.ThreadPoolExecutor())
      _global_server.add_insecure_port(f"0.0.0.0:{args.stats_port}")
      test_pb2_grpc.add_LoadBalancerStatsServiceServicer_to_server(
          _LoadBalancerStatsServicer(), _global_server)
      test_pb2_grpc.add_XdsUpdateClientConfigureServiceServicer_to_server(
          _XdsUpdateClientConfigureServicer(channel_configs, args.qps),
          _global_server)
      channelz.add_channelz_servicer(_global_server)
      grpc_admin.add_admin_servicers(_global_server)
      _global_server.start()
      _global_server.wait_for_termination()
      for method_handle in method_handles:
          method_handle.stop()
  
  
  def parse_metadata_arg(metadata_arg: str) -> PerMethodMetadataType:
      metadata = metadata_arg.split(",") if args.metadata else []
      per_method_metadata = collections.defaultdict(list)
      for metadatum in metadata:
          elems = metadatum.split(":")
          if len(elems) != 3:
              raise ValueError(
                  f"'{metadatum}' was not in the form 'METHOD:KEY:VALUE'")
          if elems[0] not in _SUPPORTED_METHODS:
              raise ValueError(f"Unrecognized method '{elems[0]}'")
          per_method_metadata[elems[0]].append((elems[1], elems[2]))
      return per_method_metadata
  
  
  def parse_rpc_arg(rpc_arg: str) -> Sequence[str]:
      methods = rpc_arg.split(",")
      if set(methods) - set(_SUPPORTED_METHODS):
          raise ValueError("--rpc supported methods: {}".format(
              ", ".join(_SUPPORTED_METHODS)))
      return methods
  
  
  def bool_arg(arg: str) -> bool:
      if arg.lower() in ("true", "yes", "y"):
          return True
      elif arg.lower() in ("false", "no", "n"):
          return False
      else:
          raise argparse.ArgumentTypeError(f"Could not parse '{arg}' as a bool.")
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser(
          description='Run Python XDS interop client.')
      parser.add_argument(
          "--num_channels",
          default=1,
          type=int,
          help="The number of channels from which to send requests.")
      parser.add_argument("--print_response",
                          default="False",
                          type=bool_arg,
                          help="Write RPC response to STDOUT.")
      parser.add_argument(
          "--qps",
          default=1,
          type=int,
          help="The number of queries to send from each channel per second.")
      parser.add_argument("--rpc_timeout_sec",
                          default=30,
                          type=int,
                          help="The per-RPC timeout in seconds.")
      parser.add_argument("--server",
                          default="localhost:50051",
                          help="The address of the server.")
      parser.add_argument(
          "--stats_port",
          default=50052,
          type=int,
          help="The port on which to expose the peer distribution stats service.")
      parser.add_argument(
          "--secure_mode",
          default="False",
          type=bool_arg,
          help="If specified, uses xDS credentials to connect to the server.")
      parser.add_argument('--verbose',
                          help='verbose log output',
                          default=False,
                          action='store_true')
      parser.add_argument("--log_file",
                          default=None,
                          type=str,
                          help="A file to log to.")
      rpc_help = "A comma-delimited list of RPC methods to run. Must be one of "
      rpc_help += ", ".join(_SUPPORTED_METHODS)
      rpc_help += "."
      parser.add_argument("--rpc", default="UnaryCall", type=str, help=rpc_help)
      metadata_help = (
          "A comma-delimited list of 3-tuples of the form " +
          "METHOD:KEY:VALUE, e.g. " +
          "EmptyCall:key1:value1,UnaryCall:key2:value2,EmptyCall:k3:v3")
      parser.add_argument("--metadata", default="", type=str, help=metadata_help)
      args = parser.parse_args()
      signal.signal(signal.SIGINT, _handle_sigint)
      if args.verbose:
          logger.setLevel(logging.DEBUG)
      if args.log_file:
          file_handler = logging.FileHandler(args.log_file, mode='a')
          file_handler.setFormatter(formatter)
          logger.addHandler(file_handler)
      _run(args, parse_rpc_arg(args.rpc), parse_metadata_arg(args.metadata))
  