Enhanced the BufferTimeoutManager to flush the librdkafka queue (#2103)

Author: k-raina

src/confluent_kafka/experimental/aio/producer/_buffer_timeout_manager.py

@@ -106,8 +106,7 @@ async def _monitor_timeout(self):
 
             # Check if buffer has been inactive for too long
             time_since_activity = time.time() - manager._last_activity
-            if (time_since_activity >= manager._timeout and
-                    not manager._batch_processor.is_buffer_empty()):
+            if (time_since_activity >= manager._timeout):
 
                 try:
                     # Flush the buffer due to timeout
@@ -125,6 +124,11 @@ async def _flush_buffer_due_to_timeout(self):
         This method handles the complete timeout flush workflow:
         1. Create batches from the batch processor
         2. Execute batches from the batch processor
+        3. Flush librdkafka queue to ensure messages are delivered
         """
-        # Create batches from current buffer
+        # Create batches from current buffer and send to librdkafka queue
         await self._batch_processor.flush_buffer()
+
+        # Flush librdkafka queue to ensure messages are delivered to broker
+        # 0 timeout means non-blocking flush
+        await self._kafka_executor.flush_librdkafka_queue(0)

src/confluent_kafka/experimental/aio/producer/_kafka_batch_executor.py

@@ -15,6 +15,8 @@
 import asyncio
 import logging
 
+from .. import _common
+
 logger = logging.getLogger(__name__)
 
 
@@ -88,6 +90,23 @@ def _produce_batch_and_poll():
         loop = asyncio.get_running_loop()
         return await loop.run_in_executor(self._executor, _produce_batch_and_poll)
 
+    async def flush_librdkafka_queue(self, timeout=-1):
+        """Flush the librdkafka queue and wait for all messages to be delivered
+
+        This method awaits until all outstanding produce requests are completed
+        or the timeout is reached, unless the timeout is set to 0 (non-blocking).
+
+        Args:
+            timeout: Maximum time to wait in seconds:
+                    - -1 = wait indefinitely (default)
+                    - 0 = non-blocking, return immediately
+                    - >0 = wait up to timeout seconds
+
+        Returns:
+            Number of messages still in queue after flush attempt
+        """
+        return await _common.async_call(self._executor, self._producer.flush, timeout)
+
     def _handle_partial_failures(self, batch_messages):
         """Handle messages that failed during produce_batch
 

tests/test_buffer_timeout_manager.py

@@ -0,0 +1,294 @@
+#!/usr/bin/env python3
+"""
+Unit tests for the BufferTimeoutManager class (_buffer_timeout_manager.py)
+
+This module tests the BufferTimeoutManager class to ensure proper
+buffer timeout monitoring and automatic flush handling.
+"""
+
+import asyncio
+import unittest
+import time
+from unittest.mock import Mock, AsyncMock
+
+from confluent_kafka.experimental.aio.producer._buffer_timeout_manager import BufferTimeoutManager
+
+
+class TestBufferTimeoutManager(unittest.TestCase):
+    """Test cases for BufferTimeoutManager class"""
+
+    def setUp(self):
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+
+        # Create mock dependencies
+        self.mock_batch_processor = Mock()
+        self.mock_kafka_executor = Mock()
+
+        # Configure mock methods as async
+        self.mock_batch_processor.flush_buffer = AsyncMock()
+        self.mock_batch_processor.is_buffer_empty = Mock(return_value=False)
+        self.mock_kafka_executor.flush_librdkafka_queue = AsyncMock(return_value=0)
+
+        # Create timeout manager with 1 second timeout
+        self.timeout_manager = BufferTimeoutManager(
+            self.mock_batch_processor,
+            self.mock_kafka_executor,
+            timeout=1.0
+        )
+
+    def tearDown(self):
+        # Stop any running timeout monitoring tasks
+        if hasattr(self.timeout_manager, '_timeout_task') and self.timeout_manager._timeout_task:
+            self.timeout_manager.stop_timeout_monitoring()
+        self.loop.close()
+
+    def test_initialization(self):
+        """Test that BufferTimeoutManager initializes correctly"""
+        self.assertEqual(self.timeout_manager._batch_processor, self.mock_batch_processor)
+        self.assertEqual(self.timeout_manager._kafka_executor, self.mock_kafka_executor)
+        self.assertEqual(self.timeout_manager._timeout, 1.0)
+        self.assertFalse(self.timeout_manager._running)
+        self.assertIsNone(self.timeout_manager._timeout_task)
+        self.assertIsInstance(self.timeout_manager._last_activity, float)
+        self.assertGreater(self.timeout_manager._last_activity, 0)
+
+    def test_mark_activity(self):
+        """Test that mark_activity updates the last activity timestamp"""
+        initial_time = self.timeout_manager._last_activity
+        time.sleep(0.01)  # Ensure time difference
+
+        self.timeout_manager.mark_activity()
+
+        self.assertGreater(self.timeout_manager._last_activity, initial_time)
+
+    def test_start_timeout_monitoring(self):
+        """Test that timeout monitoring task starts correctly"""
+        async def async_test():
+            self.timeout_manager.start_timeout_monitoring()
+
+            self.assertTrue(self.timeout_manager._running)
+            self.assertIsNotNone(self.timeout_manager._timeout_task)
+            self.assertFalse(self.timeout_manager._timeout_task.done())
+
+            # Clean up
+            self.timeout_manager.stop_timeout_monitoring()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_stop_timeout_monitoring(self):
+        """Test that timeout monitoring task stops correctly"""
+        async def async_test():
+            self.timeout_manager.start_timeout_monitoring()
+            self.assertTrue(self.timeout_manager._running)
+
+            self.timeout_manager.stop_timeout_monitoring()
+
+            self.assertFalse(self.timeout_manager._running)
+            # Task should be cancelled or None
+            self.assertTrue(
+                self.timeout_manager._timeout_task is None or
+                self.timeout_manager._timeout_task.done()
+            )
+
+        self.loop.run_until_complete(async_test())
+
+    def test_start_timeout_monitoring_disabled(self):
+        """Test that timeout monitoring doesn't start when timeout is 0"""
+        manager = BufferTimeoutManager(
+            self.mock_batch_processor,
+            self.mock_kafka_executor,
+            timeout=0  # Disabled
+        )
+
+        manager.start_timeout_monitoring()
+
+        self.assertFalse(manager._running)
+        self.assertIsNone(manager._timeout_task)
+
+    def test_flush_buffer_due_to_timeout(self):
+        """Test that _flush_buffer_due_to_timeout calls both flush methods"""
+        async def async_test():
+            # Call the flush method
+            await self.timeout_manager._flush_buffer_due_to_timeout()
+
+            # Verify both flush operations were called
+            self.mock_batch_processor.flush_buffer.assert_called_once()
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_called_once()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_flush_buffer_due_to_timeout_order(self):
+        """Test that flush operations are called in the correct order"""
+        call_order = []
+
+        async def track_batch_flush():
+            call_order.append('batch_processor')
+
+        async def track_kafka_flush(timeout=0):
+            call_order.append('kafka_executor')
+
+        self.mock_batch_processor.flush_buffer.side_effect = track_batch_flush
+        self.mock_kafka_executor.flush_librdkafka_queue.side_effect = track_kafka_flush
+
+        async def async_test():
+            await self.timeout_manager._flush_buffer_due_to_timeout()
+
+            # Verify order: batch processor first, then kafka executor
+            self.assertEqual(call_order, ['batch_processor', 'kafka_executor'])
+
+        self.loop.run_until_complete(async_test())
+
+    def test_flush_buffer_due_to_timeout_batch_processor_error(self):
+        """Test that errors from batch processor are propagated"""
+        self.mock_batch_processor.flush_buffer.side_effect = Exception("Batch flush failed")
+
+        async def async_test():
+            with self.assertRaises(Exception) as context:
+                await self.timeout_manager._flush_buffer_due_to_timeout()
+
+            self.assertEqual(str(context.exception), "Batch flush failed")
+            # Batch processor flush was attempted
+            self.mock_batch_processor.flush_buffer.assert_called_once()
+            # Kafka executor flush should not be called if batch processor fails
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_not_called()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_flush_buffer_due_to_timeout_kafka_executor_error(self):
+        """Test that errors from kafka executor are propagated"""
+        self.mock_kafka_executor.flush_librdkafka_queue.side_effect = Exception("Kafka flush failed")
+
+        async def async_test():
+            with self.assertRaises(Exception) as context:
+                await self.timeout_manager._flush_buffer_due_to_timeout()
+
+            self.assertEqual(str(context.exception), "Kafka flush failed")
+            # Both flush operations were attempted
+            self.mock_batch_processor.flush_buffer.assert_called_once()
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_called_once()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_monitor_timeout_triggers_flush(self):
+        """Test that timeout monitoring triggers flush when buffer is inactive"""
+        async def async_test():
+            # Set last activity to past time (simulating inactivity)
+            self.timeout_manager._last_activity = time.time() - 2.0  # 2 seconds ago
+
+            # Start monitoring
+            self.timeout_manager.start_timeout_monitoring()
+
+            # Wait for timeout check to trigger
+            # Check interval for 1s timeout should be 0.5s, add buffer time
+            await asyncio.sleep(0.7)
+
+            # Stop monitoring
+            self.timeout_manager.stop_timeout_monitoring()
+
+            # Verify flush was called
+            self.mock_batch_processor.flush_buffer.assert_called()
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_called()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_monitor_timeout_flushes_librdkafka_even_when_buffer_empty(self):
+        """Test that timeout monitoring flushes librdkafka queue even when local buffer is empty
+        This ensures that messages sitting in librdkafka's internal queue get delivered,
+        even if there are no messages in our local buffer.
+        """
+        # Configure buffer as empty
+        self.mock_batch_processor.is_buffer_empty.return_value = True
+
+        async def async_test():
+            # Set last activity to past time
+            self.timeout_manager._last_activity = time.time() - 2.0
+
+            # Start monitoring
+            self.timeout_manager.start_timeout_monitoring()
+
+            # Wait for potential timeout check
+            await asyncio.sleep(0.7)
+
+            # Stop monitoring
+            self.timeout_manager.stop_timeout_monitoring()
+
+            # flush_buffer() returns early since buffer is empty
+            self.mock_batch_processor.flush_buffer.assert_called_once()
+            # BUT librdkafka queue flush should still happen
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_called_once()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_monitor_timeout_does_not_flush_recent_activity(self):
+        """Test that timeout monitoring doesn't flush if buffer has recent activity"""
+        async def async_test():
+            # Set last activity to recent time (within timeout)
+            self.timeout_manager._last_activity = time.time() - 0.3  # 0.3 seconds ago
+
+            # Start monitoring
+            self.timeout_manager.start_timeout_monitoring()
+
+            # Wait for potential timeout check (less than timeout duration)
+            # 0.3s (initial) + 0.4s (sleep) = 0.7s < 1.0s timeout
+            await asyncio.sleep(0.4)
+
+            # Stop monitoring
+            self.timeout_manager.stop_timeout_monitoring()
+
+            # Verify flush was NOT called since activity is recent
+            self.mock_batch_processor.flush_buffer.assert_not_called()
+            self.mock_kafka_executor.flush_librdkafka_queue.assert_not_called()
+
+        self.loop.run_until_complete(async_test())
+
+    def test_monitor_timeout_updates_activity_after_flush(self):
+        """Test that timeout monitoring updates activity timestamp after successful flush"""
+        async def async_test():
+            # Set last activity to past time
+            self.timeout_manager._last_activity = time.time() - 2.0
+            initial_time = self.timeout_manager._last_activity
+
+            # Start monitoring
+            self.timeout_manager.start_timeout_monitoring()
+
+            # Wait for timeout check to trigger
+            await asyncio.sleep(0.7)
+
+            # Stop monitoring
+            self.timeout_manager.stop_timeout_monitoring()
+
+            # Verify activity was updated after flush
+            self.assertGreater(self.timeout_manager._last_activity, initial_time)
+
+        self.loop.run_until_complete(async_test())
+
+    def test_weak_reference_cleanup(self):
+        """Test that weak references allow proper garbage collection"""
+        async def async_test():
+            # Create a manager with timeout monitoring
+            manager = BufferTimeoutManager(
+                self.mock_batch_processor,
+                self.mock_kafka_executor,
+                timeout=1.0
+            )
+            manager.start_timeout_monitoring()
+
+            # Wait a bit for task to start
+            await asyncio.sleep(0.1)
+
+            # Get reference to the task before deleting manager
+            task = manager._timeout_task
+            self.assertIsNotNone(task)
+            self.assertFalse(task.done())
+
+            # Stop the manager first to ensure clean shutdown
+            manager.stop_timeout_monitoring()
+
+            # Wait a bit for the cancellation to take effect
+            await asyncio.sleep(0.8)
+
+            # Verify the task stopped when we explicitly stopped monitoring
+            self.assertTrue(task.done())
+
+        self.loop.run_until_complete(async_test())