Fix consume() breaking out of chunked loop on first result

src/confluent_kafka/src/Consumer.c

@@ -1090,15 +1090,14 @@
 }
 
 /**
- * @brief Consume a batch of messages from the subscribed topics.
+ * @brief Consume a batch of up to num_messages messages from the subscribed
+ * topics.
  *
- * Instead of a single blocking call to rd_kafka_consume_batch_queue() with the
- * full timeout, this function:
- * 1. Splits the timeout into 200ms chunks
- * 2. Calls rd_kafka_consume_batch_queue() with chunk timeout
- * 3. Between chunks, re-acquires GIL and calls PyErr_CheckSignals()
- * 4. If signal detected, returns NULL (raises KeyboardInterrupt)
- * 5. Continues until messages received, timeout expired, or signal detected.
+ * This makes a single blocking call to rd_kafka_consume_batch_queue(), which
+ * gathers up to num_messages messages, blocking for up to the full timeout.
+ * The call is not interruptible while it is blocked: a pending signal such as
+ * Ctrl+C is only observed once the call returns, so with the default infinite
+ * timeout an idle consume() can block until a message arrives.
  *
  * @param self Consumer handle
  * @param args Positional arguments (unused)
@@ -1107,8 +1106,9 @@
  *                consume per call. Default: 1. Maximum: 1000000.
  *              - timeout (float, optional): Timeout in seconds.
  *                Default: -1.0 (infinite timeout)
- * @return PyObject* List of Message objects, empty list if timeout, or NULL on
- * error (raises KeyboardInterrupt if signal detected)
+ * @return PyObject* List of up to num_messages Message objects, empty list if
+ * the timeout elapses with none available, or NULL on error or when a pending
+ * signal (e.g. KeyboardInterrupt) is raised after the call returns.
  */
 static PyObject *
 Consumer_consume(Handle *self, PyObject *args, PyObject *kwargs) {
@@ -1119,11 +1119,7 @@
         PyObject *msglist;
         rd_kafka_queue_t *rkqu = self->u.Consumer.rkqu;
         CallState cs;
-        Py_ssize_t i, n = 0;
-        const int CHUNK_TIMEOUT_MS = 200; /* 200ms chunks for signal checking */
-        int total_timeout_ms;
-        int chunk_timeout_ms;
-        int chunk_count = 0;
+        Py_ssize_t i, n;
 
         if (!self->rk) {
                 PyErr_SetString(PyExc_RuntimeError, ERR_MSG_CONSUMER_CLOSED);
@@ -1141,8 +1137,6 @@
                 return NULL;
         }
 
-        total_timeout_ms = cfl_timeout_ms(tmout);
-
         rkmessages = malloc(num_messages * sizeof(rd_kafka_message_t *));
         if (!rkmessages) {
                 PyErr_NoMemory();
@@ -1151,64 +1145,9 @@
 
         CallState_begin(self, &cs);
 
-        /* Skip wakeable poll pattern for non-blocking or very short timeouts.
-         * This avoids unnecessary GIL re-acquisition that can interfere with
-         * ThreadPool. Only use wakeable poll for
-         * blocking calls that need to be interruptible. */
-        if (total_timeout_ms >= 0 && total_timeout_ms < CHUNK_TIMEOUT_MS) {
-                n = (Py_ssize_t)rd_kafka_consume_batch_queue(
-                    rkqu, total_timeout_ms, rkmessages, num_messages);
+        n = (Py_ssize_t)rd_kafka_consume_batch_queue(
+            rkqu, cfl_timeout_ms(tmout), rkmessages, num_messages);
 
-                if (n < 0) {
-                        /* Error - need to restore GIL before setting error */
-                        PyEval_RestoreThread(cs.thread_state);
-                        free(rkmessages);
-                        cfl_PyErr_Format(
-                            rd_kafka_last_error(), "%s",
-                            rd_kafka_err2str(rd_kafka_last_error()));
-                        return NULL;
-                }
-        } else {
-                while (1) {
-                        /* Calculate timeout for this chunk */
-                        chunk_timeout_ms = calculate_chunk_timeout(
-                            total_timeout_ms, chunk_count, CHUNK_TIMEOUT_MS);
-                        if (chunk_timeout_ms == 0) {
-                                /* Timeout expired */
-                                break;
-                        }
-
-                        /* Consume with chunk timeout */
-                        n = (Py_ssize_t)rd_kafka_consume_batch_queue(
-                            rkqu, chunk_timeout_ms, rkmessages, num_messages);
-
-                        if (n < 0) {
-                                /* Error - need to restore GIL before setting
-                                 * error */
-                                PyEval_RestoreThread(cs.thread_state);
-                                free(rkmessages);
-                                cfl_PyErr_Format(
-                                    rd_kafka_last_error(), "%s",
-                                    rd_kafka_err2str(rd_kafka_last_error()));
-                                return NULL;
-                        }
-
-                        /* If we got messages, exit the loop */
-                        if (n > 0) {
-                                break;
-                        }
-
-                        chunk_count++;
-
-                        /* Check for signals between chunks */
-                        if (check_signals_between_chunks(self, &cs)) {
-                                free(rkmessages);
-                                return NULL;
-                        }
-                }
-        }
-
-        /* Final GIL restore and signal check */
         if (!CallState_end(self, &cs)) {
                 for (i = 0; i < n; i++) {
                         rd_kafka_message_destroy(rkmessages[i]);
@@ -1217,7 +1156,13 @@
                 return NULL;
         }
 
-        /* Create Python list from messages  */
+        if (n < 0) {
+                free(rkmessages);
+                cfl_PyErr_Format(rd_kafka_last_error(), "%s",
+                                 rd_kafka_err2str(rd_kafka_last_error()));
+                return NULL;
+        }
+
         msglist = PyList_New(n);
 
         for (i = 0; i < n; i++) {
@@ -1393,6 +1338,13 @@
      "  .. note: Callbacks may be called from this method, "
      "such as ``on_assign``, ``on_revoke``, et.al.\n"
      "\n"
+     "  .. note:: This is a blocking call and does not respond to signals "
+     "(e.g. Ctrl+C / SIGINT) while it is waiting; a pending signal is only "
+     "delivered once the call returns. With the default infinite timeout an "
+     "idle ``consume()`` can block indefinitely. Applications that need to "
+     "stay responsive to interrupts should pass a short ``timeout`` and call "
+     "``consume()`` in a loop, handling signals between calls.\n"
+     "\n"
      "  :param int num_messages: The maximum number of messages to return "
      "(default: 1).\n"
      "  :param float timeout: The maximum time to block waiting for message, "

src/confluent_kafka/src/Producer.c

@@ -353,7 +353,8 @@
  * 2. Calls rd_kafka_poll() with chunk timeout
  * 3. Between chunks, re-acquires GIL and calls PyErr_CheckSignals()
  * 4. If signal detected, returns -1 (raises KeyboardInterrupt)
- * 5. Continues until events processed, timeout expired, or signal detected
+ * 5. Returns as soon as a chunk serves any events, the timeout expires, or a
+ *    signal is detected
  *
  * @param self Producer handle
  * @param tmout Timeout in milliseconds (-1 for infinite)
@@ -377,7 +378,7 @@
         if (total_timeout_ms >= 0 && total_timeout_ms < CHUNK_TIMEOUT_MS) {
                 r = rd_kafka_poll(self->rk, total_timeout_ms);
         } else {
                 while (1) {
                         /* Calculate timeout for this chunk */
                         chunk_timeout_ms = calculate_chunk_timeout(
                             total_timeout_ms, chunk_count, CHUNK_TIMEOUT_MS);
@@ -394,7 +395,10 @@
                                 r = chunk_result;
                                 break;
                         }
-                        r += chunk_result; /* Accumulate events processed */
+                        r = chunk_result;
+
+                        if (chunk_result > 0)
+                                break;
 
                         chunk_count++;
 

tests/integration/consumer/test_consumer_wakeable_poll_consume.py

@@ -142,3 +142,139 @@ def test_consume_message_delivery_with_wakeable_pattern(kafka_cluster):
         assert msg.value() == expected_value, expected_msg
 
     consumer.close()
+
+
+def test_consume_accumulates_messages_up_to_num_messages(kafka_cluster):
+    """Test that consume() gathers up to num_messages within the timeout,
+    rather than returning as soon as the first messages become available.
+    """
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-accumulate-chunks')
+
+    # Produce 10 messages
+    producer = kafka_cluster.cimpl_producer()
+    num_produced = 10
+    for i in range(num_produced):
+        producer.produce(topic, value=f'msg-{i}'.encode())
+    producer.flush(timeout=5.0)
+
+    # Create consumer
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-accumulate',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    # Wait for subscription and partition assignment
+    time.sleep(2.0)
+
+    # Consume num_messages=10 with a generous timeout; the batch call gathers
+    # all 10 within the timeout rather than returning on the first few.
+    msglist = consumer.consume(num_messages=num_produced, timeout=10.0)
+
+    assert len(msglist) == num_produced, (
+        f"Expected {num_produced} messages but got {len(msglist)}. "
+        f"consume() did not gather the full batch within the timeout."
+    )
+
+    for i, msg in enumerate(msglist):
+        assert not msg.error(), f"Message {i} has error: {msg.error()}"
+
+    consumer.close()
+
+
+def test_consume_returns_partial_on_timeout(kafka_cluster):
+    """Test that consume() returns partial results when timeout expires
+    before num_messages is reached."""
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-partial-timeout')
+
+    # Produce only 3 messages, but request 100
+    producer = kafka_cluster.cimpl_producer()
+    num_produced = 3
+    for i in range(num_produced):
+        producer.produce(topic, value=f'partial-{i}'.encode())
+    producer.flush(timeout=5.0)
+
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-partial',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    time.sleep(2.0)
+
+    # Request 100 messages but only 3 exist — should return 3 after timeout
+    start = time.time()
+    msglist = consumer.consume(num_messages=100, timeout=3.0)
+    elapsed = time.time() - start
+
+    assert len(msglist) == num_produced, f"Expected {num_produced} messages (partial), got {len(msglist)}"
+    # Should have waited close to the full timeout since num_messages wasn't reached
+    assert elapsed >= 2.0, f"Should wait near full timeout for more messages, but returned in {elapsed:.2f}s"
+
+    for i, msg in enumerate(msglist):
+        assert not msg.error(), f"Message {i} has error: {msg.error()}"
+        assert msg.value() == f'partial-{i}'.encode()
+
+    consumer.close()
+
+
+def test_consume_accumulates_messages_produced_in_waves(kafka_cluster):
+    """Test that consume() accumulates messages that arrive in multiple waves,
+    rather than returning as soon as the first wave becomes available.
+    """
+    import threading
+
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-consume-waves')
+
+    producer = kafka_cluster.cimpl_producer()
+
+    def produce_in_waves():
+        """Produce 3 waves of 4 messages each, with 1s gaps."""
+        for wave in range(3):
+            time.sleep(1.0)
+            for i in range(4):
+                msg_num = wave * 4 + i
+                producer.produce(topic, value=f'wave-{msg_num}'.encode())
+            producer.flush(timeout=5.0)
+
+    consumer_conf = kafka_cluster.client_conf(
+        {
+            'group.id': 'test-consume-waves',
+            'socket.timeout.ms': 100,
+            'session.timeout.ms': 6000,
+            'auto.offset.reset': 'earliest',
+        }
+    )
+    consumer = TestConsumer(consumer_conf)
+    consumer.subscribe([topic])
+
+    time.sleep(2.0)
+
+    # Start producing in background
+    producer_thread = threading.Thread(target=produce_in_waves, daemon=True)
+    producer_thread.start()
+
+    # Request 10 messages with a long timeout (waves take ~3s to complete)
+    msglist = consumer.consume(num_messages=10, timeout=10.0)
+
+    # Should have accumulated messages across multiple waves
+    assert len(msglist) == 10, (
+        f"Expected exactly 10 messages accumulated across waves, got {len(msglist)}. "
+        f"consume() may be returning early on the first wave."
+    )
+
+    for msg in msglist:
+        assert not msg.error(), f"Message has error: {msg.error()}"
+
+    producer_thread.join(timeout=5.0)
+    consumer.close()

tests/integration/producer/test_producer_wakeable_poll_flush.py

@@ -119,6 +119,41 @@ def delivery_callback(err, msg):
     consumer.close()
 
 
+def test_poll_returns_early_after_delivery_callback(kafka_cluster):
+    """Test that poll() returns early after delivery callback fires."""
+    topic = kafka_cluster.create_topic_and_wait_propogation('test-poll-early-return')
+
+    delivery_called = []
+
+    def delivery_callback(err, msg):
+        delivery_called.append(time.time())
+
+    producer_conf = kafka_cluster.client_conf(
+        {
+            'socket.timeout.ms': 100,
+            'message.timeout.ms': 10000,
+        }
+    )
+    producer = kafka_cluster.cimpl_producer(producer_conf)
+
+    producer.produce(topic, value=b'early-return-test', on_delivery=delivery_callback)
+
+    # Poll with a long timeout — should return early once callback fires
+    poll_timeout = 5.0
+    start = time.time()
+    events = producer.poll(timeout=poll_timeout)
+    elapsed = time.time() - start
+
+    assert len(delivery_called) == 1, "Expected delivery callback to fire"
+    assert events > 0, "Expected at least 1 event served"
+    assert elapsed < poll_timeout - 1.0, (
+        f"poll({poll_timeout}) took {elapsed:.2f}s — should have returned "
+        f"early after delivery callback, not blocked for full timeout"
+    )
+
+    producer.close()
+
+
 def test_flush_message_delivery_with_wakeable_pattern(kafka_cluster):
     """Test that flush() correctly delivers messages when using wakeable pattern.