diff options
| -rw-r--r-- | kafka/producer/base.py | 9 | ||||
| -rw-r--r-- | test/test_producer.py | 18 |
2 files changed, 23 insertions, 4 deletions
diff --git a/kafka/producer/base.py b/kafka/producer/base.py index a0d9ac1..0e005c5 100644 --- a/kafka/producer/base.py +++ b/kafka/producer/base.py @@ -31,6 +31,8 @@ BATCH_SEND_MSG_COUNT = 20 BATCH_RETRY_OPTIONS = RetryOptions( limit=0, backoff_ms=300, retry_on_timeouts=False) +# unlimited +ASYNC_QUEUE_MAXSIZE = 0 STOP_ASYNC_PRODUCER = -1 @@ -159,12 +161,14 @@ class Producer(object): batch_send=False, batch_send_every_n=BATCH_SEND_MSG_COUNT, batch_send_every_t=BATCH_SEND_DEFAULT_INTERVAL, - batch_retry_options=BATCH_RETRY_OPTIONS): + batch_retry_options=BATCH_RETRY_OPTIONS, + async_queue_maxsize=ASYNC_QUEUE_MAXSIZE): if batch_send: async = True assert batch_send_every_n > 0 assert batch_send_every_t > 0 + assert async_queue_maxsize >= 0 else: batch_send_every_n = 1 batch_send_every_t = 3600 @@ -186,7 +190,8 @@ class Producer(object): log.warning("async producer does not guarantee message delivery!") log.warning("Current implementation does not retry Failed messages") log.warning("Use at your own risk! (or help improve with a PR!)") - self.queue = Queue() # Messages are sent through this queue + # Messages are sent through this queue + self.queue = Queue(async_queue_maxsize) self.thread_stop_event = Event() self.thread = Thread(target=_send_upstream, args=(self.queue, diff --git a/test/test_producer.py b/test/test_producer.py index b57dfd8..627178d 100644 --- a/test/test_producer.py +++ b/test/test_producer.py @@ -53,15 +53,29 @@ class TestKafkaProducer(unittest.TestCase): assert client.send_produce_request.called @patch('kafka.producer.base.Process') - def test_producer_batch_send_queue_overfilled(self, process_mock): + def test_producer_async_queue_overfilled_batch_send(self, process_mock): queue_size = 2 producer = Producer(MagicMock(), batch_send=True, - batch_send_queue_maxsize=queue_size) + async_queue_maxsize=queue_size) topic = b'test-topic' partition = 0 + message = b'test-message' + + with self.assertRaises(BatchQueueOverfilledError): + message_list = [message] * (queue_size + 1) + producer.send_messages(topic, partition, *message_list) + @patch('kafka.producer.base.Process') + def test_producer_async_queue_overfilled(self, process_mock): + queue_size = 2 + producer = Producer(MagicMock(), async=True, + async_queue_maxsize=queue_size) + + topic = b'test-topic' + partition = 0 message = b'test-message' + with self.assertRaises(BatchQueueOverfilledError): message_list = [message] * (queue_size + 1) producer.send_messages(topic, partition, *message_list) |