Improve handling of high message volumes:

Block when trying to produce a message and the librdkafka internal queue is full.
Allow specifying a timeout for producer.close.
Provide an interface to check the queue size.

Author: cnweaver

adc/producer.py

@@ -47,11 +47,26 @@ def write(self,
                                 "Either configure a topic when constructing the Producer, "
                                 "or specify the topic argument to write()")
         self.logger.debug("writing message to %s", topic)
+        produce_kwargs = {"headers": headers, "key": key}
         if delivery_callback is not None:
-            self._producer.produce(topic, msg, headers=headers, key=key,
-                                   on_delivery=delivery_callback)
-        else:
-            self._producer.produce(topic, msg, headers=headers, key=key,)
+            produce_kwargs["on_delivery"] = delivery_callback
+        while True:
+            try:
+                self._producer.produce(topic, msg, **produce_kwargs)
+                break
+            except BufferError:
+                # It's hard to know what the size limit on the buffer is, so we will try to
+                # wait until the number of items in it decreases (or it is empty, in case all
+                # messages were successfully sent in the time it took to handle the error).
+                buffer_len = len(self._producer)
+                self.logger.debug(f"Blocking due to BufferError, buffer size: {buffer_len}")
+                while buffer_len > 0 and len(self._producer) == buffer_len:
+                    self._producer.poll(0.01)
+
+    def queued_message_count(self):
+        """Get the number of messages waiting to be sent to the broker.
+        """
+        return len(self._producer)
 
     def flush(self, timeout: timedelta = timedelta(seconds=10)) -> int:
         """Attempt to flush enqueued messages. Return the number of messages still
@@ -65,9 +80,9 @@ def flush(self, timeout: timedelta = timedelta(seconds=10)) -> int:
             self.logger.debug("flushed all messages")
         return n
 
-    def close(self) -> int:
+    def close(self, timeout: timedelta = timedelta(seconds=10)) -> int:
         self.logger.debug("shutting down producer")
-        return self.flush()
+        return self.flush(timeout)
 
     def __enter__(self) -> 'Producer':
         return self

tests/test_kafka_integration.py

@@ -84,6 +84,44 @@ def test_message_with_key(self):
         self.assertEqual(msg.value(), b"can you hear me?")
         self.assertEqual(msg.key(), b"test_msg")
 
+    def test_message_with_callback(self):
+        """Try writing a message into the Kafka broker, with a delivery callback.
+        """
+        topic = "test_message_with_callback"
+        callback_called = False
+        callback_error = False
+
+        def callback(err, msg):
+            nonlocal callback_called
+            nonlocal callback_error
+            callback_called = True
+            callback_error = err is not None
+
+        ap = adc.producer
+        with ap.Producer(ap.ProducerConfig(broker_urls=[self.kafka.address],
+                                           topic=topic, auth=self.kafka.auth)) as producer:
+            producer.write("message data", delivery_callback=callback)
+            producer.flush()
+            consumer = adc.consumer.Consumer(adc.consumer.ConsumerConfig(
+                broker_urls=[self.kafka.address],
+                group_id="test_consumer",
+                auth=self.kafka.auth,
+            ))
+            consumer.subscribe(topic)
+            stream = consumer.stream()
+
+            msg = next(stream)
+            if msg.error() is not None:
+                raise Exception(msg.error())
+            # give the producer's background thread time to notice the ack from the broker and
+            # fire the callback
+            producer._producer.poll(0.1)
+
+        self.assertEqual(msg.topic(), topic)
+        self.assertEqual(msg.value(), b"message data")
+        self.assertEqual(callback_called, True)
+        self.assertEqual(callback_error, False)
+
     def test_reset_to_end(self):
         # Write a few messages.
         topic = "test_reset_to_end"
@@ -486,7 +524,7 @@ def poll_for_kafka_broker_address(self, maxiter=20, sleep=timedelta(milliseconds
         """Block until the Docker daemon tells us the IP and Port of the Kafa broker.
 
         Returns the ip and port as a string in the form "ip:port."
-"""
+        """
         i = 0
         while (not self.query_kafka_broker_address()) and i < maxiter:
             logger.info("polling to wait for container to acquire port...")
@@ -508,7 +546,7 @@ def query_kafka_broker_address(self):
         port = addrs[0]['HostPort']
         return f"{ip}:{port}"
 
-    def poll_for_kafka_active(self, maxiter=20, sleep=timedelta(milliseconds=500)):
+    def poll_for_kafka_active(self, maxiter=40, sleep=timedelta(milliseconds=500)):
         """Block until Kafka's network listener is accepting connections."""
         i = 0
         while (not self.query_kafka_active()) and i < maxiter: