# Source code for kafka.client

    from __future__ import absolute_import

    import collections
    import copy
    import functools
    import logging
    import random
    import time
    import select

    from kafka.vendor import six

    import kafka.errors
    from kafka.errors import (UnknownError, KafkaConnectionError, FailedPayloadsError,
                              KafkaTimeoutError, KafkaUnavailableError,
                              LeaderNotAvailableError, UnknownTopicOrPartitionError,
                              NotLeaderForPartitionError, ReplicaNotAvailableError)
    from kafka.structs import TopicPartition, BrokerMetadata

    from kafka.conn import (
        collect_hosts, BrokerConnection,
        ConnectionStates, get_ip_port_afi)
    from kafka.protocol import KafkaProtocol

    # New KafkaClient
    # this is not exposed in top-level imports yet,
    # due to conflicts with legacy SimpleConsumer / SimpleProducer usage
    from kafka.client_async import KafkaClient

    log = logging.getLogger(__name__)

    # Legacy KafkaClient interface -- will be deprecated soon
    class SimpleClient(object):

        CLIENT_ID = b'kafka-python'
        DEFAULT_SOCKET_TIMEOUT_SECONDS = 120

        # NOTE: The timeout given to the client should always be greater than the
        # one passed to SimpleConsumer.get_message(), otherwise you can get a
        # socket timeout.
        def __init__(self, hosts, client_id=CLIENT_ID,
                     timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS,
                     correlation_id=0):
            # We need one connection to bootstrap
            self.client_id = client_id
            self.timeout = timeout
            self.hosts = collect_hosts(hosts)
            self.correlation_id = correlation_id

            self._conns = {}
            self.brokers = {}            # broker_id -> BrokerMetadata
            self.topics_to_brokers = {}  # TopicPartition -> BrokerMetadata
            self.topic_partitions = {}   # topic -> partition -> leader

            self.load_metadata_for_topics()  # bootstrap with all metadata

        ##################
        #   Private API  #
        ##################

        def _get_conn(self, host, port, afi):
            """Get or create a connection to a broker using host and port"""
            host_key = (host, port)
            if host_key not in self._conns:
                self._conns[host_key] = BrokerConnection(
                    host, port, afi,
                    request_timeout_ms=self.timeout * 1000,
                    client_id=self.client_id
                )

            conn = self._conns[host_key]
            if not conn.connect_blocking(self.timeout):
                conn.close()
                raise KafkaConnectionError("%s:%s (%s)" % (host, port, afi))
            return conn

        def _get_leader_for_partition(self, topic, partition):
            """
            Returns the leader for a partition or None if the partition exists
            but has no leader.

            Raises:
                UnknownTopicOrPartitionError: If the topic or partition is not part
                    of the metadata.
                LeaderNotAvailableError: If the server has metadata, but there is no
            current leader.
            """

            key = TopicPartition(topic, partition)

            # Use cached metadata if it is there
            if self.topics_to_brokers.get(key) is not None:
                return self.topics_to_brokers[key]

            # Otherwise refresh metadata

            # If topic does not already exist, this will raise
            # UnknownTopicOrPartitionError if not auto-creating
            # LeaderNotAvailableError otherwise until partitions are created
            self.load_metadata_for_topics(topic)

            # If the partition doesn't actually exist, raise
            if partition not in self.topic_partitions.get(topic, []):
                raise UnknownTopicOrPartitionError(key)

            # If there's no leader for the partition, raise
            leader = self.topic_partitions[topic][partition]
            if leader == -1:
                raise LeaderNotAvailableError((topic, partition))

            # Otherwise return the BrokerMetadata
            return self.brokers[leader]

        def _get_coordinator_for_group(self, group):
            """
            Returns the coordinator broker for a consumer group.

            GroupCoordinatorNotAvailableError will be raised if the coordinator
            does not currently exist for the group.

            GroupLoadInProgressError is raised if the coordinator is available
            but is still loading offsets from the internal topic
            """

            resp = self.send_consumer_metadata_request(group)

            # If there's a problem with finding the coordinator, raise the
            # provided error
            kafka.errors.check_error(resp)

            # Otherwise return the BrokerMetadata
            return BrokerMetadata(resp.nodeId, resp.host, resp.port, None)

        def _next_id(self):
            """Generate a new correlation id"""
            # modulo to keep w/i int32
            self.correlation_id = (self.correlation_id + 1) % 2**31
            return self.correlation_id

        def _send_broker_unaware_request(self, payloads, encoder_fn, decoder_fn):
            """
            Attempt to send a broker-agnostic request to one of the available
            brokers. Keep trying until you succeed.
            """
            hosts = set()
            for broker in self.brokers.values():
                host, port, afi = get_ip_port_afi(broker.host)
                hosts.add((host, broker.port, afi))

            hosts.update(self.hosts)
            hosts = list(hosts)
            random.shuffle(hosts)

            for (host, port, afi) in hosts:
                try:
                    conn = self._get_conn(host, port, afi)
                except KafkaConnectionError:
                    log.warning("Skipping unconnected connection: %s:%s (AFI %s)",
                                host, port, afi)
                    continue
                request = encoder_fn(payloads=payloads)
                future = conn.send(request)

                # Block
                while not future.is_done:
                    for r, f in conn.recv():
                        f.success(r)

                if future.failed():
                    log.error("Request failed: %s", future.exception)
                    continue

                return decoder_fn(future.value)

            raise KafkaUnavailableError('All servers failed to process request: %s' % (hosts,))

        def _payloads_by_broker(self, payloads):
            payloads_by_broker = collections.defaultdict(list)
            for payload in payloads:
                try:
                    leader = self._get_leader_for_partition(payload.topic, payload.partition)
                except (KafkaUnavailableError, LeaderNotAvailableError,
                        UnknownTopicOrPartitionError):
                    leader = None
                payloads_by_broker[leader].append(payload)
            return dict(payloads_by_broker)

        def _send_broker_aware_request(self, payloads, encoder_fn, decoder_fn):
            """
            Group a list of request payloads by topic+partition and send them to
            the leader broker for that partition using the supplied encode/decode
            functions

            Arguments:

            payloads: list of object-like entities with a topic (str) and
                partition (int) attribute; payloads with duplicate topic-partitions
                are not supported.

            encode_fn: a method to encode the list of payloads to a request body,
                must accept client_id, correlation_id, and payloads as
                keyword arguments

            decode_fn: a method to decode a response body into response objects.
                The response objects must be object-like and have topic
                and partition attributes

            Returns:

            List of response objects in the same order as the supplied payloads
            """
            # encoders / decoders do not maintain ordering currently
            # so we need to keep this so we can rebuild order before returning
            original_ordering = [(p.topic, p.partition) for p in payloads]

            # Connection errors generally mean stale metadata
            # although sometimes it means incorrect api request
            # Unfortunately there is no good way to tell the difference
            # so we'll just reset metadata on all errors to be safe
            refresh_metadata = False

            # For each broker, send the list of request payloads
            # and collect the responses and errors
            payloads_by_broker = self._payloads_by_broker(payloads)
            responses = {}

            def failed_payloads(payloads):
                for payload in payloads:
                    topic_partition = (str(payload.topic), payload.partition)
                    responses[(topic_partition)] = FailedPayloadsError(payload)

            # For each BrokerConnection keep the real socket so that we can use
            # a select to perform unblocking I/O
            connections_by_future = {}
            for broker, broker_payloads in six.iteritems(payloads_by_broker):
                if broker is None:
                    failed_payloads(broker_payloads)
                    continue

                host, port, afi = get_ip_port_afi(broker.host)
                try:
                    conn = self._get_conn(host, broker.port, afi)
                except KafkaConnectionError:
                    refresh_metadata = True
                    failed_payloads(broker_payloads)
                    continue

                request = encoder_fn(payloads=broker_payloads)
                future = conn.send(request)

                if future.failed():
                    refresh_metadata = True
                    failed_payloads(broker_payloads)
                    continue

                if not request.expect_response():
                    for payload in broker_payloads:
                        topic_partition = (str(payload.topic), payload.partition)
                        responses[topic_partition] = None
                    continue

                connections_by_future[future] = (conn, broker)

            conn = None
            while connections_by_future:
                futures = list(connections_by_future.keys())

                # block until a socket is ready to be read
                sockets = [
                    conn._sock
                    for future, (conn, _) in six.iteritems(connections_by_future)
                    if not future.is_done and conn._sock is not None]
                if sockets:
                    read_socks, _, _ = select.select(sockets, [], [])

                for future in futures:

                    if not future.is_done:
                        conn, _ = connections_by_future[future]
                        for r, f in conn.recv():
                            f.success(r)
                        continue

                    _, broker = connections_by_future.pop(future)
                    if future.failed():
                        refresh_metadata = True
                        failed_payloads(payloads_by_broker[broker])

                    else:
                        for payload_response in decoder_fn(future.value):
                            topic_partition = (str(payload_response.topic),
                                               payload_response.partition)
                            responses[topic_partition] = payload_response

            if refresh_metadata:
                self.reset_all_metadata()

            # Return responses in the same order as provided
            return [responses[tp] for tp in original_ordering]

        def _send_consumer_aware_request(self, group, payloads, encoder_fn, decoder_fn):
            """
            Send a list of requests to the consumer coordinator for the group
            specified using the supplied encode/decode functions. As the payloads
            that use consumer-aware requests do not contain the group (e.g.
            OffsetFetchRequest), all payloads must be for a single group.

            Arguments:

            group: the name of the consumer group (str) the payloads are for
            payloads: list of object-like entities with topic (str) and
                partition (int) attributes; payloads with duplicate
                topic+partition are not supported.

            encode_fn: a method to encode the list of payloads to a request body,
                must accept client_id, correlation_id, and payloads as
                keyword arguments

            decode_fn: a method to decode a response body into response objects.
                The response objects must be object-like and have topic
                and partition attributes

            Returns:

            List of response objects in the same order as the supplied payloads
            """
            # encoders / decoders do not maintain ordering currently
            # so we need to keep this so we can rebuild order before returning
            original_ordering = [(p.topic, p.partition) for p in payloads]

            broker = self._get_coordinator_for_group(group)

            # Send the list of request payloads and collect the responses and
            # errors
            responses = {}
            request_id = self._next_id()
            log.debug('Request %s to %s: %s', request_id, broker, payloads)
            request = encoder_fn(client_id=self.client_id,
                                 correlation_id=request_id, payloads=payloads)

            # Send the request, recv the response
            try:
                host, port, afi = get_ip_port_afi(broker.host)
                conn = self._get_conn(host, broker.port, afi)
            except KafkaConnectionError as e:
                log.warning('KafkaConnectionError attempting to send request %s '
                            'to server %s: %s', request_id, broker, e)

                for payload in payloads:
                    topic_partition = (payload.topic, payload.partition)
                    responses[topic_partition] = FailedPayloadsError(payload)

            # No exception, try to get response
            else:

                future = conn.send(request_id, request)
                while not future.is_done:
                    for r, f in conn.recv():
                        f.success(r)

                # decoder_fn=None signal that the server is expected to not
                # send a response.  This probably only applies to
                # ProduceRequest w/ acks = 0
                if decoder_fn is None:
                    log.debug('Request %s does not expect a response '
                              '(skipping conn.recv)', request_id)
                    for payload in payloads:
                        topic_partition = (payload.topic, payload.partition)
                        responses[topic_partition] = None
                    return []

                if future.failed():
                    log.warning('Error attempting to receive a '
                                'response to request %s from server %s: %s',
                                request_id, broker, future.exception)

                    for payload in payloads:
                        topic_partition = (payload.topic, payload.partition)
                        responses[topic_partition] = FailedPayloadsError(payload)

                else:
                    response = future.value
                    _resps = []
                    for payload_response in decoder_fn(response):
                        topic_partition = (payload_response.topic,
                                           payload_response.partition)
                        responses[topic_partition] = payload_response
                        _resps.append(payload_response)
                    log.debug('Response %s: %s', request_id, _resps)

            # Return responses in the same order as provided
            return [responses[tp] for tp in original_ordering]

        def __repr__(self):
            return '<SimpleClient client_id=%s>' % (self.client_id)

        def _raise_on_response_error(self, resp):

            # Response can be an unraised exception object (FailedPayloadsError)
            if isinstance(resp, Exception):
                raise resp

            # Or a server api error response
            try:
                kafka.errors.check_error(resp)
            except (UnknownTopicOrPartitionError, NotLeaderForPartitionError):
                self.reset_topic_metadata(resp.topic)
                raise

            # Return False if no error to enable list comprehensions
            return False

        #################
        #   Public API  #
        #################
        def close(self):
            for conn in self._conns.values():
                conn.close()

        def copy(self):
            """
            Create an inactive copy of the client object, suitable for passing
            to a separate thread.

            Note that the copied connections are not initialized, so :meth:`.reinit`
            must be called on the returned copy.
            """
            _conns = self._conns
            self._conns = {}
            c = copy.deepcopy(self)
            self._conns = _conns
            return c

        def reinit(self):
            timeout = time.time() + self.timeout
            conns = set(self._conns.values())
            for conn in conns:
                conn.close()
                conn.connect()

            while time.time() < timeout:
                for conn in list(conns):
                    conn.connect()
                    if conn.connected():
                        conns.remove(conn)
                if not conns:
                    break

        def reset_topic_metadata(self, *topics):
            for topic in topics:
                for topic_partition in list(self.topics_to_brokers.keys()):
                    if topic_partition.topic == topic:
                        del self.topics_to_brokers[topic_partition]
                if topic in self.topic_partitions:
                    del self.topic_partitions[topic]

        def reset_all_metadata(self):
            self.topics_to_brokers.clear()
            self.topic_partitions.clear()

        def has_metadata_for_topic(self, topic):
            return (
                topic in self.topic_partitions
                and len(self.topic_partitions[topic]) > 0
            )

        def get_partition_ids_for_topic(self, topic):
            if topic not in self.topic_partitions:
                return []

            return sorted(list(self.topic_partitions[topic]))

        @property
        def topics(self):
            return list(self.topic_partitions.keys())

        def ensure_topic_exists(self, topic, timeout=30):
            start_time = time.time()

            while not self.has_metadata_for_topic(topic):
                if time.time() > start_time + timeout:
                    raise KafkaTimeoutError('Unable to create topic {0}'.format(topic))
                self.load_metadata_for_topics(topic, ignore_leadernotavailable=True)
                time.sleep(.5)

        def load_metadata_for_topics(self, *topics, **kwargs):
            """Fetch broker and topic-partition metadata from the server.

            Updates internal data: broker list, topic/partition list, and
            topic/partition -> broker map. This method should be called after
            receiving any error.

            Note: Exceptions *will not* be raised in a full refresh (i.e. no topic
            list). In this case, error codes will be logged as errors.
            Partition-level errors will also not be raised here (a single partition
            w/o a leader, for example).

            Arguments:
                *topics (optional): If a list of topics is provided,
                    the metadata refresh will be limited to the specified topics
                    only.
                ignore_leadernotavailable (bool): suppress LeaderNotAvailableError
                    so that metadata is loaded correctly during auto-create.
                    Default: False.

            Raises:
                UnknownTopicOrPartitionError: Raised for topics that do not exist,
                    unless the broker is configured to auto-create topics.
                LeaderNotAvailableError: Raised for topics that do not exist yet,
                    when the broker is configured to auto-create topics. Retry
                    after a short backoff (topics/partitions are initializing).
            """
            if 'ignore_leadernotavailable' in kwargs:
                ignore_leadernotavailable = kwargs['ignore_leadernotavailable']
            else:
                ignore_leadernotavailable = False

            if topics:
                self.reset_topic_metadata(*topics)
            else:
                self.reset_all_metadata()

            resp = self.send_metadata_request(topics)

            log.debug('Updating broker metadata: %s', resp.brokers)
            log.debug('Updating topic metadata: %s', [topic for _, topic, _ in resp.topics])

            self.brokers = dict([(nodeId, BrokerMetadata(nodeId, host, port, None))
                                 for nodeId, host, port in resp.brokers])

            for error, topic, partitions in resp.topics:
                # Errors expected for new topics
                if error:
                    error_type = kafka.errors.kafka_errors.get(error, UnknownError)
                    if error_type in (UnknownTopicOrPartitionError, LeaderNotAvailableError):
                        log.error('Error loading topic metadata for %s: %s (%s)',
                                  topic, error_type, error)
                        if topic not in topics:
                            continue
                        elif (error_type is LeaderNotAvailableError and
                              ignore_leadernotavailable):
                            continue
                    raise error_type(topic)

                self.topic_partitions[topic] = {}
                for error, partition, leader, _, _ in partitions:

                    self.topic_partitions[topic][partition] = leader

                    # Populate topics_to_brokers dict
                    topic_part = TopicPartition(topic, partition)

                    # Check for partition errors
                    if error:
                        error_type = kafka.errors.kafka_errors.get(error, UnknownError)

                        # If No Leader, topics_to_brokers topic_partition -> None
                        if error_type is LeaderNotAvailableError:
                            log.error('No leader for topic %s partition %d', topic, partition)
                            self.topics_to_brokers[topic_part] = None
                            continue

                        # If one of the replicas is unavailable -- ignore
                        # this error code is provided for admin purposes only
                        # we never talk to replicas, only the leader
                        elif error_type is ReplicaNotAvailableError:
                            log.debug('Some (non-leader) replicas not available for topic %s partition %d', topic, partition)

                        else:
                            raise error_type(topic_part)

                    # If Known Broker, topic_partition -> BrokerMetadata
                    if leader in self.brokers:
                        self.topics_to_brokers[topic_part] = self.brokers[leader]

                    # If Unknown Broker, fake BrokerMetadata so we don't lose the id
                    # (not sure how this could happen. server could be in bad state)
                    else:
                        self.topics_to_brokers[topic_part] = BrokerMetadata(
                            leader, None, None, None
                        )

        def send_metadata_request(self, payloads=(), fail_on_error=True,
                                  callback=None):
            encoder = KafkaProtocol.encode_metadata_request
            decoder = KafkaProtocol.decode_metadata_response

            return self._send_broker_unaware_request(payloads, encoder, decoder)

        def send_consumer_metadata_request(self, payloads=(), fail_on_error=True,
                                           callback=None):
            encoder = KafkaProtocol.encode_consumer_metadata_request
            decoder = KafkaProtocol.decode_consumer_metadata_response

            return self._send_broker_unaware_request(payloads, encoder, decoder)

        def send_produce_request(self, payloads=(), acks=1, timeout=1000,
                                 fail_on_error=True, callback=None):
            """
            Encode and send some ProduceRequests

            ProduceRequests will be grouped by (topic, partition) and then
            sent to a specific broker. Output is a list of responses in the
            same order as the list of payloads specified

            Arguments:
                payloads (list of ProduceRequest): produce requests to send to kafka
                    ProduceRequest payloads must not contain duplicates for any
                    topic-partition.
                acks (int, optional): how many acks the servers should receive from replica
                    brokers before responding to the request. If it is 0, the server
                    will not send any response. If it is 1, the server will wait
                    until the data is written to the local log before sending a
                    response.  If it is -1, the server will wait until the message
                    is committed by all in-sync replicas before sending a response.
                    For any value > 1, the server will wait for this number of acks to
                    occur (but the server will never wait for more acknowledgements than
                    there are in-sync replicas). defaults to 1.
                timeout (int, optional): maximum time in milliseconds the server can
                    await the receipt of the number of acks, defaults to 1000.
                fail_on_error (bool, optional): raise exceptions on connection and
                    server response errors, defaults to True.
                callback (function, optional): instead of returning the ProduceResponse,
                    first pass it through this function, defaults to None.

            Returns:
                list of ProduceResponses, or callback results if supplied, in the
                order of input payloads
            """

            encoder = functools.partial(
                KafkaProtocol.encode_produce_request,
                acks=acks,
                timeout=timeout)

            if acks == 0:
                decoder = None
            else:
                decoder = KafkaProtocol.decode_produce_response

            resps = self._send_broker_aware_request(payloads, encoder, decoder)

            return [resp if not callback else callback(resp) for resp in resps
                    if resp is not None and
                    (not fail_on_error or not self._raise_on_response_error(resp))]

        def send_fetch_request(self, payloads=(), fail_on_error=True,
                               callback=None, max_wait_time=100, min_bytes=4096):
            """
            Encode and send a FetchRequest

            Payloads are grouped by topic and partition so they can be pipelined
            to the same brokers.
            """

            encoder = functools.partial(KafkaProtocol.encode_fetch_request,
                              max_wait_time=max_wait_time,
                              min_bytes=min_bytes)

            resps = self._send_broker_aware_request(
                payloads, encoder,
                KafkaProtocol.decode_fetch_response)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]

        def send_offset_request(self, payloads=(), fail_on_error=True,
                                callback=None):
            resps = self._send_broker_aware_request(
                payloads,
                KafkaProtocol.encode_offset_request,
                KafkaProtocol.decode_offset_response)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]

        def send_list_offset_request(self, payloads=(), fail_on_error=True,
                                callback=None):
            resps = self._send_broker_aware_request(
                payloads,
                KafkaProtocol.encode_list_offset_request,
                KafkaProtocol.decode_list_offset_response)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]

        def send_offset_commit_request(self, group, payloads=(),
                                       fail_on_error=True, callback=None):
            encoder = functools.partial(KafkaProtocol.encode_offset_commit_request,
                              group=group)
            decoder = KafkaProtocol.decode_offset_commit_response
            resps = self._send_broker_aware_request(payloads, encoder, decoder)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]

        def send_offset_fetch_request(self, group, payloads=(),
                                      fail_on_error=True, callback=None):

            encoder = functools.partial(KafkaProtocol.encode_offset_fetch_request,
                              group=group)
            decoder = KafkaProtocol.decode_offset_fetch_response
            resps = self._send_broker_aware_request(payloads, encoder, decoder)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]

        def send_offset_fetch_request_kafka(self, group, payloads=(),
                                      fail_on_error=True, callback=None):

            encoder = functools.partial(KafkaProtocol.encode_offset_fetch_request,
                              group=group, from_kafka=True)
            decoder = KafkaProtocol.decode_offset_fetch_response
            resps = self._send_consumer_aware_request(group, payloads, encoder, decoder)

            return [resp if not callback else callback(resp) for resp in resps
                    if not fail_on_error or not self._raise_on_response_error(resp)]