Trying from fix lock issue

Author: sujata-m

src/python_ms_core/core/topic/azure_topic.py

@@ -1,17 +1,16 @@
 import json
-
 import gc
 import logging
+import multiprocessing as mp
 import os
 import time
+import traceback
 from ..config.config import TopicConfig
-from ..resource_errors import ExceptionHandler
 from concurrent.futures import ThreadPoolExecutor
 from .abstract.topic_abstract import TopicAbstract
 from ..queue.models.queue_message import QueueMessage
 from azure.servicebus import ServiceBusClient, ServiceBusMessage
 from azure.servicebus import AutoLockRenewer
-import concurrent.futures as cf
 import threading
 
 try:
@@ -56,6 +55,9 @@ def __init__(self, config: TopicConfig=None, topic_name=None, max_concurrent_mes
         self.topic_name = topic_name
         self.publisher = self.client.get_topic_sender(topic_name=topic_name)
         self.executor = ThreadPoolExecutor(max_workers=max_concurrent_messages)
+        self.callback_execution_mode = self._get_callback_execution_mode()
+        self.callback_process_start_method = self._get_process_start_method()
+        self.process_context = self._get_process_context()
         self.internal_count = 0
         self.max_renewal_duration = self._get_positive_int_from_env(
             'TOPIC_MAX_LOCK_RENEWAL_DURATION_SECONDS',
@@ -81,9 +83,11 @@ def __init__(self, config: TopicConfig=None, topic_name=None, max_concurrent_mes
         self._process = None
         self._prime_runtime_samplers()
         logger.info(
-            'Configured lock renewal for topic %s: max_lock_renewal_duration=%s seconds, '
-            'renew_margin=%s seconds, renewer_max_workers=%s',
+            'Configured AzureTopic %s: execution_mode=%s, process_start_method=%s, '
+            'max_lock_renewal_duration=%s seconds, renew_margin=%s seconds, renewer_max_workers=%s',
             self.topic_name,
+            self.callback_execution_mode,
+            self.callback_process_start_method,
             self.max_renewal_duration,
             self.lock_renewal_margin,
             renewer_max_workers,
@@ -129,13 +133,13 @@ def subscribe(self, subscription: str, callback, max_receivable_messages=-1):
                         with self.thread_lock:
                             self.internal_count += len(messages)
                         for message in messages: 
+                            execution_task = self._submit_processing_task(message, callback)
                             self.lock_renewal.register(
                                 self.lock_renew_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)
                             self.pending_tasks.append((execution_task, message))
                     else:
                         if len(self.pending_tasks) > 0:
@@ -146,27 +150,64 @@ def subscribe(self, subscription: str, callback, max_receivable_messages=-1):
                     logger.error(f'Error in receiving messages: {e}')
 
 
-    def internal_callback(self, message, callbackfn):
+    def internal_callback(self, message_payload, callbackfn):
         """
         Internal callback function that processes a message and invokes the callback function.
         Args:
-            message (ServiceBusMessage): The message to process.
+            message_payload (str): The message payload to process.
             callbackfn (function): The callback function to invoke.
         Returns:
-            ServiceBusMessage: The processed message.
+            dict: The callback status payload.
         """
         try:
-            queue_message = QueueMessage.data_from(str(message))
+            queue_message = QueueMessage.data_from(message_payload)
             callbackfn(queue_message)
-            return [True,message]
+            return {'success': True, 'error': None}
         except Exception as e:
-            logger.error(f'Error in processing message: {e}')
-            return [False,message]
+            return {
+                'success': False,
+                'error': ''.join(traceback.format_exception(type(e), e, e.__traceback__)).strip(),
+            }
 
 
-    def settle_message(self, x: cf.Future):
+    def settle_message(self, x):
         return self._settle_task(x)
 
+    def _submit_processing_task(self, message, callback):
+        message_payload = str(message)
+        if self.callback_execution_mode == 'process':
+            try:
+                return self._submit_process_task(message_payload, callback)
+            except Exception as exc:
+                logger.warning(
+                    'Falling back to thread execution for message %s because process start failed: %s',
+                    self._get_message_id(message),
+                    exc,
+                )
+        return self._submit_thread_task(message_payload, callback)
+
+    def _submit_thread_task(self, message_payload, callback):
+        future = self.executor.submit(self.internal_callback, message_payload, callback)
+        return _FutureExecutionTask(future)
+
+    def _submit_process_task(self, message_payload, callback):
+        if self.process_context is None:
+            raise RuntimeError('Process execution mode is not available for this environment.')
+
+        parent_connection, child_connection = self.process_context.Pipe(duplex=False)
+        callback_process = self.process_context.Process(
+            target=_run_callback_in_subprocess,
+            args=(message_payload, callback, child_connection),
+        )
+        try:
+            callback_process.start()
+        except Exception:
+            parent_connection.close()
+            child_connection.close()
+            raise
+        child_connection.close()
+        return _ProcessExecutionTask(callback_process, parent_connection)
+
     def _get_receivable_count(self, max_receivable_messages=-1):
         with self.thread_lock:
             available_slots = self.max_concurrent_messages - self.internal_count
@@ -178,8 +219,14 @@ def _get_receivable_count(self, max_receivable_messages=-1):
     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)
+        if timeout <= 0:
+            self._settle_completed_tasks()
+            return
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            if any(task.done() for task, _ in self.pending_tasks):
+                break
+            time.sleep(min(0.1, max(deadline - time.time(), 0)))
         self._settle_completed_tasks()
 
     def _settle_completed_tasks(self):
@@ -191,16 +238,21 @@ def _settle_completed_tasks(self):
                 remaining_tasks.append((future, incoming_message))
         self.pending_tasks = remaining_tasks
 
-    def _settle_task(self, x: cf.Future, incoming_message=None):
+    def _settle_task(self, x, incoming_message=None):
         """
         Sets the message as completed and updates the internal count.
         Args:
-            x (cf.Future): The future object representing the message processing.
+            x: The task object representing the message processing.
         """
         try:
-            [is_success, settled_message] = x.result()
-            if settled_message is not None:
-                incoming_message = settled_message
+            task_result = x.result()
+        except Exception as e:
+            task_result = {
+                'success': False,
+                'error': f'Callback worker exited before returning a result: {e}',
+            }
+
+        try:
             if incoming_message is None:
                 return
             if getattr(incoming_message, '_lock_expired', False):
@@ -210,11 +262,15 @@ def _settle_task(self, x: cf.Future, incoming_message=None):
                     f'auto_renew_error={getattr(incoming_message, "auto_renew_error", None)}'
                 )
                 return
-            if is_success:
+            if task_result.get('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
+                logger.error(
+                    'Processing failed for message %s: %s',
+                    self._get_message_id(incoming_message),
+                    task_result.get('error', 'unknown processing failure'),
+                )
+                self.receiver.abandon_message(incoming_message)
         except Exception as e:
             logger.error(f'Error in settling message: {e}')
         finally:
@@ -309,6 +365,46 @@ def _get_positive_int_from_env(name, default):
         logger.warning(f'Invalid value for {name}: {value}. Using default {default}.')
         return default
 
+    @staticmethod
+    def _get_callback_execution_mode():
+        value = os.environ.get('TOPIC_CALLBACK_EXECUTION_MODE', 'process')
+        normalized = str(value).strip().lower()
+        if normalized in ('process', 'thread'):
+            return normalized
+        logger.warning(
+            'Invalid value for TOPIC_CALLBACK_EXECUTION_MODE: %s. Using default process.',
+            value,
+        )
+        return 'process'
+
+    @staticmethod
+    def _get_process_start_method():
+        available_methods = mp.get_all_start_methods()
+        default_method = 'fork' if 'fork' in available_methods else mp.get_start_method() or available_methods[0]
+        configured_method = os.environ.get('TOPIC_CALLBACK_PROCESS_START_METHOD', default_method)
+        normalized_method = str(configured_method).strip().lower()
+        if normalized_method in available_methods:
+            return normalized_method
+        logger.warning(
+            'Invalid value for TOPIC_CALLBACK_PROCESS_START_METHOD: %s. Using default %s.',
+            configured_method,
+            default_method,
+        )
+        return default_method
+
+    def _get_process_context(self):
+        if self.callback_execution_mode != 'process':
+            return None
+        try:
+            return mp.get_context(self.callback_process_start_method)
+        except ValueError:
+            logger.warning(
+                'Process start method %s is unavailable. Falling back to thread execution.',
+                self.callback_process_start_method,
+            )
+            self.callback_execution_mode = 'thread'
+            return None
+
 class _LockRenewLoggingReceiver:
     def __init__(self, topic):
         self._topic = topic
@@ -321,3 +417,56 @@ def renew_message_lock(self, renewable):
             self._topic._get_runtime_snapshot(),
         )
         return self._topic.receiver.renew_message_lock(renewable)
+
+
+def _run_callback_in_subprocess(message_payload, callbackfn, result_connection):
+    try:
+        queue_message = QueueMessage.data_from(message_payload)
+        callbackfn(queue_message)
+        result_connection.send({'success': True, 'error': None})
+    except BaseException as exc:  # pragma: no cover - exercised through the parent process wrapper
+        result_connection.send({
+            'success': False,
+            'error': ''.join(traceback.format_exception(type(exc), exc, exc.__traceback__)).strip(),
+        })
+    finally:
+        result_connection.close()
+
+
+class _FutureExecutionTask:
+    def __init__(self, future):
+        self._future = future
+
+    def done(self):
+        return self._future.done()
+
+    def result(self):
+        return self._future.result()
+
+
+class _ProcessExecutionTask:
+    def __init__(self, process, result_connection):
+        self._process = process
+        self._result_connection = result_connection
+        self._result = None
+
+    def done(self):
+        return not self._process.is_alive()
+
+    def result(self):
+        if self._result is not None:
+            return self._result
+
+        self._process.join()
+        try:
+            if self._result_connection.poll():
+                self._result = self._result_connection.recv()
+            else:
+                self._result = {
+                    'success': False,
+                    'error': f'Callback worker exited with code {self._process.exitcode} without returning a result.',
+                }
+        finally:
+            self._result_connection.close()
+
+        return self._result

src/python_ms_core/version.py

@@ -1 +1 @@
-__version__ = '0.2.5.1'
+__version__ = '0.2.5.2'