Feature fix lock issue

Author: sujata-m

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Change log
 
+# Version 0.0.25
+### Fixes
+- **Azure Topic Settlement Stability:** Moved Azure Service Bus message settlement back onto the receiver-owning loop instead of settling from worker callback threads. This keeps receive and complete/abandon operations on the same receiver flow for long-running jobs.
+- **Receiver Slot Tracking:** Reserved and released in-flight message slots on the receive loop so concurrency limits remain accurate while messages are still processing.
+- **Lock Renewal Diagnostics:** Added logging for Service Bus lock-renew failures to make long-running lock-loss issues visible before final settlement.
+
 # Version 0.0.23
 - Updated unit test cases pipeline
 - Added support to upload test cases results on Azure blob
@@ -109,4 +115,4 @@ Reference task:
 - Added classes and methods
     - Topic
         - publish
-        - subscribe
+        - subscribe

MANIFEST.in

@@ -0,0 +1 @@
+include requirements.txt

src/python_ms_core/core/topic/azure_topic.py

@@ -54,6 +54,7 @@ def __init__(self, config: TopicConfig=None, topic_name=None, max_concurrent_mes
         self.max_renewal_duration = 86400 # Renew the message upto 1 day
         self.wait_time_for_message = 5
         self.thread_lock = threading.Lock()
+        self.pending_tasks = []
 
 
     def publish(self, data: QueueMessage):
@@ -78,23 +79,36 @@ def subscribe(self, subscription: str, callback, max_receivable_messages=-1):
         self.receiver.local_received_messages = 0
         while True:
                 try:
-                    to_receive = (self.max_concurrent_messages - self.internal_count)
-                    total_messages_to_recieve_more = max_receivable_messages - self.receiver.local_received_messages
-                    if max_receivable_messages > 0:
-                        to_receive = min(to_receive, total_messages_to_recieve_more)
+                    self._settle_completed_tasks()
+                    to_receive = self._get_receivable_count(max_receivable_messages=max_receivable_messages)
+                    if max_receivable_messages > 0 and self.receiver.local_received_messages >= max_receivable_messages:
+                        if len(self.pending_tasks) == 0:
+                            break
+                        self._wait_for_pending_tasks(timeout=0.5)
+                        continue
                     if to_receive > 0:
                         messages = self.receiver.receive_messages(max_message_count=to_receive, max_wait_time=self.wait_time_for_message)
                         if not messages or len(messages) == 0:
+                            if len(self.pending_tasks) > 0:
+                                self._wait_for_pending_tasks(timeout=0.5)
                             continue
                         self.receiver.local_received_messages += len(messages)
+                        with self.thread_lock:
+                            self.internal_count += len(messages)
                         for message in messages: 
-                            self.lock_renewal.register(self.receiver, message, max_lock_renewal_duration=self.max_renewal_duration)
+                            self.lock_renewal.register(
+                                self.receiver,
+                                message,
+                                max_lock_renewal_duration=self.max_renewal_duration,
+                                on_lock_renew_failure=self._handle_lock_renew_failure,
+                            )
                             execution_task = self.executor.submit(self.internal_callback, message, callback)
-                            execution_task.add_done_callback(lambda x: self.settle_message(x))
-                        if self.receiver.local_received_messages >= max_receivable_messages and max_receivable_messages > 0: # Break if the messages are more than max_receivable_messages
-                            break
+                            self.pending_tasks.append((execution_task, message))
                     else:
-                        time.sleep(self.wait_time_for_message)
+                        if len(self.pending_tasks) > 0:
+                            self._wait_for_pending_tasks(timeout=0.5)
+                        else:
+                            time.sleep(self.wait_time_for_message)
                 except Exception as e:
                     logger.error(f'Error in receiving messages: {e}')
 
@@ -109,8 +123,6 @@ def internal_callback(self, message, callbackfn):
             ServiceBusMessage: The processed message.
         """
         try:
-            with self.thread_lock:
-                self.internal_count += 1 # thread safe.
             queue_message = QueueMessage.data_from(str(message))
             callbackfn(queue_message)
             return [True,message]
@@ -120,21 +132,58 @@ def internal_callback(self, message, callbackfn):
 
 
     def settle_message(self, x: cf.Future):
+        return self._settle_task(x)
+
+    def _get_receivable_count(self, max_receivable_messages=-1):
+        with self.thread_lock:
+            available_slots = self.max_concurrent_messages - self.internal_count
+        if max_receivable_messages > 0:
+            remaining_messages = max_receivable_messages - self.receiver.local_received_messages
+            available_slots = min(available_slots, remaining_messages)
+        return max(available_slots, 0)
+
+    def _wait_for_pending_tasks(self, timeout=0.5):
+        if len(self.pending_tasks) == 0:
+            return
+        futures = [future for future, _ in self.pending_tasks]
+        cf.wait(futures, timeout=timeout, return_when=cf.FIRST_COMPLETED)
+        self._settle_completed_tasks()
+
+    def _settle_completed_tasks(self):
+        remaining_tasks = []
+        for future, incoming_message in self.pending_tasks:
+            if future.done():
+                self._settle_task(future, incoming_message=incoming_message)
+            else:
+                remaining_tasks.append((future, incoming_message))
+        self.pending_tasks = remaining_tasks
+
+    def _settle_task(self, x: cf.Future, incoming_message=None):
         """
         Sets the message as completed and updates the internal count.
         Args:
             x (cf.Future): The future object representing the message processing.
         """
-        # Lock the internal count
-        with self.thread_lock:
-            self.internal_count -= 1
-        # Check if the future has an exception
-        [is_success,incoming_message] = x.result()
-        if is_success:
-            self.receiver.complete_message(incoming_message)
-        else:
-            print(f'Abandoning message: {incoming_message}')
-            self.receiver.abandon_message(incoming_message) # send back to the topic
-        return  
+        try:
+            [is_success, settled_message] = x.result()
+            if settled_message is not None:
+                incoming_message = settled_message
+            if incoming_message is None:
+                return
+            if is_success:
+                self.receiver.complete_message(incoming_message)
+            else:
+                logger.info(f'Abandoning message: {incoming_message}')
+                self.receiver.abandon_message(incoming_message) # send back to the topic
+        except Exception as e:
+            logger.error(f'Error in settling message: {e}')
+        finally:
+            with self.thread_lock:
+                self.internal_count = max(self.internal_count - 1, 0)
+        return
+
+    def _handle_lock_renew_failure(self, renewable, error):
+        message_id = getattr(renewable, 'message_id', None) or getattr(renewable, 'messageId', 'unknown')
+        logger.error(f'Error renewing lock for message {message_id}: {error}')
 
-    
+    

src/python_ms_core/version.py

@@ -1 +1 @@
-__version__ = '0.0.24'
+__version__ = '0.0.25'

tests/unit_tests/test_topic/test_azure_topic.py

@@ -0,0 +1,70 @@
+import unittest
+import concurrent.futures as cf
+from unittest.mock import MagicMock, patch
+
+from src.python_ms_core.core.topic.azure_topic import AzureTopic
+from src.python_ms_core.core.queue.models.queue_message import QueueMessage
+
+
+class TestAzureTopic(unittest.TestCase):
+
+    @patch('src.python_ms_core.core.topic.azure_topic.AutoLockRenewer')
+    @patch('src.python_ms_core.core.topic.azure_topic.ServiceBusClient')
+    def test_subscribe_settles_completed_tasks_on_receiver_loop(self, mock_service_bus_client, mock_auto_lock_renewer):
+        mock_client = MagicMock()
+        mock_receiver = MagicMock()
+        mock_message = MagicMock()
+        mock_future = MagicMock()
+        mock_config = MagicMock(connection_string='Endpoint=sb://test/')
+
+        mock_service_bus_client.from_connection_string.return_value = mock_client
+        mock_client.get_topic_sender.return_value = MagicMock()
+        mock_client.get_subscription_receiver.return_value = mock_receiver
+        mock_receiver.receive_messages.side_effect = [[mock_message]]
+
+        topic = AzureTopic(config=mock_config, topic_name='mock-topic', max_concurrent_messages=1)
+        callback = MagicMock()
+        
+        def submit_side_effect(fn, *args, **kwargs):
+            mock_future = cf.Future()
+            mock_future.set_result(fn(*args, **kwargs))
+            return mock_future
+
+        topic.executor.submit = MagicMock(side_effect=submit_side_effect)
+        with patch.object(QueueMessage, 'data_from', return_value=QueueMessage()):
+            topic.subscribe(subscription='mock-subscription', callback=callback, max_receivable_messages=1)
+
+        callback.assert_called_once()
+        mock_auto_lock_renewer.return_value.register.assert_called_once()
+        mock_receiver.complete_message.assert_called_once_with(mock_message)
+        mock_receiver.abandon_message.assert_not_called()
+
+    @patch('src.python_ms_core.core.topic.azure_topic.logger')
+    @patch('src.python_ms_core.core.topic.azure_topic.AutoLockRenewer')
+    @patch('src.python_ms_core.core.topic.azure_topic.ServiceBusClient')
+    def test_settle_task_logs_error_and_releases_slot(self, mock_service_bus_client, mock_auto_lock_renewer, mock_logger):
+        mock_client = MagicMock()
+        mock_receiver = MagicMock()
+        mock_message = MagicMock()
+        mock_future = MagicMock()
+        mock_config = MagicMock(connection_string='Endpoint=sb://test/')
+
+        mock_service_bus_client.from_connection_string.return_value = mock_client
+        mock_client.get_topic_sender.return_value = MagicMock()
+        mock_client.get_subscription_receiver.return_value = mock_receiver
+        mock_future.result.return_value = [True, mock_message]
+        mock_receiver.complete_message.side_effect = Exception('Mocked settlement failure')
+
+        topic = AzureTopic(config=mock_config, topic_name='mock-topic', max_concurrent_messages=1)
+        topic.receiver = mock_receiver
+        topic.internal_count = 1
+
+        topic._settle_task(mock_future, incoming_message=mock_message)
+
+        mock_receiver.complete_message.assert_called_once_with(mock_message)
+        mock_logger.error.assert_called_once_with('Error in settling message: Mocked settlement failure')
+        self.assertEqual(topic.internal_count, 0)
+
+
+if __name__ == '__main__':
+    unittest.main()

version.py

@@ -0,0 +1 @@
+version = '0.0.25'