def ensure_connection(service_name, method_name, method):
    """Ensure that client is connected before executing method.

    .. note:: Class to which this decorator is applied **must** have
       ``get_connection_exception_class()`` method that would return the
       class of the exception to be thrown when retryable connection
       exception is repacked.

    This decorator can only be applied to class methods,
    not a non-class function.

    Args:
        method_name: A string, the name of the method to be ensured.
        method: A string, the actual method to be ensured.

    Returns:
        Whatever the executed method returns.

    """
    @functools.wraps(method)
    def method_wrapper(self, *args, **kwargs):

        req_timeout_ms = kwargs.pop('rpc_timeout_ms', self.timeout)
        conn_timeout_ms = kwargs.pop('rpc_timeout_ms',
                                     self.socket_connection_timeout)
        if conn_timeout_ms is None:
            conn_timeout_ms = req_timeout_ms
        retries_left = self.retry_count
        while retries_left:
            start_time = datetime.datetime.now()
            try:
                # Ensure connection.
                try:
                    self.connect(conn_timeout_ms, req_timeout_ms)
                except socket.timeout:
                    raise ThriftConnectionTimeoutError()
                result = method(self._client, *args, **kwargs)
                time_taken = datetime.datetime.now() - start_time
                # compute time taken into milliseconds
                time_taken_ms = time_taken.total_seconds() * 1000
                self.statsd_client.timing(
                    "client.requests.{0}.{1}".format(service_name, method_name),
                    time_taken_ms, sample_rate=0.001)
                self.refresh_connection_if_needed()
                return result
            except TApplicationException as e:
                handler_args = args_to_str(*args, **kwargs)
                time_taken = datetime.datetime.now() - start_time
                # compute time taken into milliseconds
                time_taken_ms = time_taken.total_seconds() * 1000
                log.info(
                    "Thrift call failed TApplicationException : %s(%s) : "
                    "%s:%d : time_taken_ms : %s : %s" % (
                        method_name, handler_args, self.host,
                        self.port, time_taken_ms, e))
                raise
            except Exception as e:
                t, v, tb = sys.exc_info()
                retries_left -= 1
                handler_args = args_to_str(*args, **kwargs)
                time_taken = datetime.datetime.now() - start_time
                # compute time taken into milliseconds
                time_taken_ms = time_taken.total_seconds() * 1000

                # application exception, if it is retriable as determined by
                # RetryPolicy then we simply raise the exception, no connection
                # teardown is needed, because the exception was thrown by the
                # server and transported back to the client.
                if _is_application_exception(e):
                    retry_policy_to_apply = self.retry_policy
                    if not retry_policy_to_apply:
                        retry_policy_to_apply = DEFAULT_RETRY_POLICY
                    if not retry_policy_to_apply.should_retry(e):
                        if random.random() < self.failed_retry_policy_log_sample_rate:
                            # Sample logging in case logging is too overwhelming.
                            log.info(
                                "Thrift call failed retry policy : %s(%s) :"
                                "%s:%d : time_taken_ms : %s : %s" % (
                                    method_name, handler_args, self.host,
                                    self.port, time_taken_ms, e))
                            # raise exception to stop it from being retried
                        raise t, v, tb
                elif _is_rpc_timeout(e):
                    # rpc socket timeout, not connection socket timeout
                    log.info(
                        "Thrift call failed rpc timeout : %s(%s) :"
                        "%s:%d : time_taken_ms : %s : %s" % (
                            method_name, handler_args, self.host,
                            self.port, time_taken_ms, e))
                    self.statsd_client.increment(
                        "errors.thriftclient.RpcTimeoutError",
                        sample_rate=0.01,
                        tags={'client': self.host})
                    # socket timeout, only reliable way to recover is to tear
                    # down the connection, it is probably good to select a
                    # new host, regardless whether we should retry this request
                    # or not.
                    self.teardown_connection(select_new_host=True)
                    # TODO(Yongsheng): temporarily disable this feature, we need
                    # a way to gauge the server healthiness before we can bring
                    # this feature back.
                    # raise exception to keep it from being retried.
                    # raise self.get_connection_exception_class()(e)
                else:
                    # at this point, we assume it is connectivity issue,
                    # socket read/write errors, or failing to establish
                    # connection, we will need to tear down the connection
                    # and re-establish it for subsequent calls
                    log.info(
                        "Thrift client connection fail : %s(%s) : %s:%d : "
                        "retries_left=%d : time_taken_ms : %s  %r",
                        method_name, handler_args, self.host,
                        self.port, retries_left, time_taken_ms, e)
                    self.statsd_client.increment(
                        "errors.thriftclient.ConnectionError",
                        sample_rate=0.01,
                        tags={'client': self.host})
                    # By default, for the first two retries, try the same host
                    # to rule out intermittent connectivity issue. For the last
                    # retry select a new host randomly.
                    # If ``always_retry_on_new_host`` is set True, always retry
                    # on a new host.
                    if self.always_retry_on_new_host or retries_left == 1:
                        # turn this on when we are ready to penalize bad hosts
                        # self._host_selector.invalidate()
                        self.teardown_connection(select_new_host=True)
                    else:
                        self.teardown_connection(select_new_host=False)

                # Retriable errors, but no retries left, bail.
                if not retries_left:
                    log.info(
                        "Thrift call failed all retries : %s(%s) : "
                        "%s:%d : time_taken_ms: %s %s" % (
                            method_name, handler_args, self.host, self.port,
                            time_taken_ms, e))
                    self.statsd_client.increment(
                        "errors.thriftclient.AllConnectionError",
                        sample_rate=0.01,
                        tags={'client': self.host})
                    # Repack the message and raise as a different exception.
                    raise self.get_connection_exception_class()(e), None, tb

    return method_wrapper
dataset/ETHPy150Open pinterest/kingpin/kingpin/thrift_utils/thrift_client_mixin.py/ensure_connection/original