<commit_before>from __future__ import absolute_import

import atexit
import logging
import time

try:
    from queue import Empty, Queue
except ImportError:
    from Queue import Empty, Queue
from collections import defaultdict

from threading import Thread, Event

import six

from kafka.common import (
    ProduceRequest, TopicAndPartition,
    UnsupportedCodecError, FailedPayloadsError
)
from kafka.protocol import CODEC_NONE, ALL_CODECS, create_message_set
from kafka.util import kafka_bytestring

log = logging.getLogger("kafka")

BATCH_SEND_DEFAULT_INTERVAL = 20
BATCH_SEND_MSG_COUNT = 20
BATCH_RETRY_BACKOFF_MS = 300
BATCH_RETRIES_LIMIT = 5

STOP_ASYNC_PRODUCER = -1


def _send_upstream(queue, client, codec, batch_time, batch_size,
                   req_acks, ack_timeout, retry_backoff, retries_limit, stop_event):
    """
    Listen on the queue for a specified number of messages or till
    a specified timeout and send them upstream to the brokers in one
    request
    """
    stop = False
    reqs = []
    client.reinit()

    while not stop_event.is_set():
        timeout = batch_time

        # it's a simplification: we're comparing message sets and
        # messages: each set can contain [1..batch_size] messages
        count = batch_size - len(reqs)
        send_at = time.time() + timeout
        msgset = defaultdict(list)

        # Keep fetching till we gather enough messages or a
        # timeout is reached
        while count > 0 and timeout >= 0:
            try:
                topic_partition, msg, key = queue.get(timeout=timeout)
            except Empty:
                break

            # Check if the controller has requested us to stop
            if topic_partition == STOP_ASYNC_PRODUCER:
                stop_event.set()
                break

            # Adjust the timeout to match the remaining period
            count -= 1
            timeout = send_at - time.time()
            msgset[topic_partition].append((msg, key))

        # Send collected requests upstream
        for topic_partition, msg in msgset.items():
            messages = create_message_set(msg, codec, key)
            req = ProduceRequest(topic_partition.topic,
                                 topic_partition.partition,
                                 messages)
            reqs.append(req)

        if not reqs:
            continue

        reqs_to_retry = []
        try:
            client.send_produce_request(reqs,
                                        acks=req_acks,
                                        timeout=ack_timeout)
        except FailedPayloadsError as ex:
            failed_reqs = ex.args[0]
            log.exception("Failed payloads count %s" % len(failed_reqs))

            if retries_limit is None:
                # retry all failed messages until success
                reqs_to_retry = failed_reqs
            elif not retries_limit < 0:
                #
                for req in failed_reqs:
                    if retries_limit and req.retries < retries_limit:
                        updated_req = req._replace(retries=req.retries+1)
                        reqs_to_retry.append(updated_req)
        except Exception as ex:
            log.exception("Unable to send message: %s" % type(ex))
        finally:
            reqs = []

        if reqs_to_retry and retry_backoff:
            reqs = reqs_to_retry
            log.warning("%s requests will be retried next call." % len(reqs))
            time.sleep(float(retry_backoff) / 1000)


class Producer(object):
    """
    Base class to be used by producers

    Arguments:<commit_msg>Fix: check failed reqs to retry only for positive limit<commit_after>from __future__ import absolute_import

import atexit
import logging
import time

try:
    from queue import Empty, Queue
except ImportError:
    from Queue import Empty, Queue
from collections import defaultdict

from threading import Thread, Event

import six

from kafka.common import (
    ProduceRequest, TopicAndPartition,
    UnsupportedCodecError, FailedPayloadsError
)
from kafka.protocol import CODEC_NONE, ALL_CODECS, create_message_set
from kafka.util import kafka_bytestring

log = logging.getLogger("kafka")

BATCH_SEND_DEFAULT_INTERVAL = 20
BATCH_SEND_MSG_COUNT = 20
BATCH_RETRY_BACKOFF_MS = 300
BATCH_RETRIES_LIMIT = 5

STOP_ASYNC_PRODUCER = -1


def _send_upstream(queue, client, codec, batch_time, batch_size,
                   req_acks, ack_timeout, retry_backoff, retries_limit, stop_event):
    """
    Listen on the queue for a specified number of messages or till
    a specified timeout and send them upstream to the brokers in one
    request
    """
    stop = False
    reqs = []
    client.reinit()

    while not stop_event.is_set():
        timeout = batch_time

        # it's a simplification: we're comparing message sets and
        # messages: each set can contain [1..batch_size] messages
        count = batch_size - len(reqs)
        send_at = time.time() + timeout
        msgset = defaultdict(list)

        # Keep fetching till we gather enough messages or a
        # timeout is reached
        while count > 0 and timeout >= 0:
            try:
                topic_partition, msg, key = queue.get(timeout=timeout)
            except Empty:
                break

            # Check if the controller has requested us to stop
            if topic_partition == STOP_ASYNC_PRODUCER:
                stop_event.set()
                break

            # Adjust the timeout to match the remaining period
            count -= 1
            timeout = send_at - time.time()
            msgset[topic_partition].append((msg, key))

        # Send collected requests upstream
        for topic_partition, msg in msgset.items():
            messages = create_message_set(msg, codec, key)
            req = ProduceRequest(topic_partition.topic,
                                 topic_partition.partition,
                                 messages)
            reqs.append(req)

        if not reqs:
            continue

        reqs_to_retry = []
        try:
            client.send_produce_request(reqs,
                                        acks=req_acks,
                                        timeout=ack_timeout)
        except FailedPayloadsError as ex:
            failed_reqs = ex.args[0]
            log.exception("Failed payloads count %s" % len(failed_reqs))

            # if no limit, retry all failed messages until success
            if retries_limit is None:
                reqs_to_retry = failed_reqs
            # makes sense to check failed reqs only if we have a limit > 0
            elif retries_limit > 0:
                for req in failed_reqs:
                    if retries_limit and req.retries < retries_limit:
                        updated_req = req._replace(retries=req.retries+1)
                        reqs_to_retry.append(updated_req)
        except Exception as ex:
            log.exception("Unable to send message: %s" % type(ex))
        finally:
            reqs = []

        if reqs_to_retry and retry_backoff:
            reqs = reqs_to_retry
            log.warning("%s requests will be retried next call." % len(reqs))
            time.sleep(float(retry_backoff) / 1000)


class Producer(object):
    """
    Base class to be used by producers

    Arguments:
{'fraction_non_alphanumeric': 0.06238056238056238, 'fraction_numerical': 0.0055965055965055965, 'mean_word_length': 1.9568200161420501, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3356958', 'n_tokens_mistral': 2053, 'n_tokens_neox': 1963, 'n_words': 603}