Fix service bus client lifecycle

CHANGELOG.md

@@ -8,6 +8,9 @@ ENHANCEMENTS:
 * Pin all GitHub Actions workflow steps to full commit SHAs to prevent supply chain attacks plus update to latest releases ([#4886](https://github.com/microsoft/AzureTRE/pull/4886))
 * Add Windows Server 2025 image support to Guacamole. ([#4890](https://github.com/microsoft/AzureTRE/issues/4890))
 
+BUG FIXES:
+* Fix inconsistent ServiceBusClient lifecycle management in deployment_status_updater.py, airlock_request_status_update.py, and runner.py to prevent connection socket and AMQP channel leaks
+
 ## (0.28.0) (March 2, 2026)
 **BREAKING CHANGES**
 * Sonatype Nexus shared service now requires explicit EULA acceptance (`accept_nexus_eula: true`) when deploying. This ensures compliance with Sonatype Nexus Community Edition licensing. ([#4842](https://github.com/microsoft/AzureTRE/issues/4842))

api_app/service_bus/airlock_request_status_update.py

@@ -34,43 +34,46 @@ async def receive_messages(self):
 
             while True:
                 try:
-                    current_time = time.time()
-                    polling_count += 1
-                    # Log a heartbeat message every 60 seconds to show the service is still working
-                    if current_time - last_heartbeat_time >= 60:
-                        logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue {polling_count} times in the last minute")
-                        last_heartbeat_time = current_time
-                        polling_count = 0
-
                     async with credentials.get_credential_async_context() as credential:
-                        service_bus_client = ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential)
-                        receiver = service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_STEP_RESULT_QUEUE)
-                        logger.debug(f"Looking for new messages on {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue...")
-                        async with receiver:
-                            received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=1)
-                            for msg in received_msgs:
-                                async with AutoLockRenewer() as renewer:
-                                    renewer.register(receiver, msg, max_lock_renewal_duration=60)
-                                    complete_message = await self.process_message(msg)
-                                    if complete_message:
-                                        await receiver.complete_message(msg)
-                                    else:
-                                        # could have been any kind of transient issue, we'll abandon back to the queue, and retry
-                                        await receiver.abandon_message(msg)
-
-                        await asyncio.sleep(10)
+                        async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client:
+                            while True:
+                                current_time = time.time()
+                                polling_count += 1
+                                # Log a heartbeat message every 60 seconds to show the service is still working
+                                if current_time - last_heartbeat_time >= 60:
+                                    logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue {polling_count} times in the last minute")
+                                    last_heartbeat_time = current_time
+                                    polling_count = 0
+
+                                receiver = service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_STEP_RESULT_QUEUE)
+                                logger.debug(f"Looking for new messages on {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue...")
+                                async with receiver:
+                                    received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=1)
+                                    for msg in received_msgs:
+                                        async with AutoLockRenewer() as renewer:
+                                            renewer.register(receiver, msg, max_lock_renewal_duration=60)
+                                            complete_message = await self.process_message(msg)
+                                            if complete_message:
+                                                await receiver.complete_message(msg)
+                                            else:
+                                                # could have been any kind of transient issue, we'll abandon back to the queue, and retry
+                                                await receiver.abandon_message(msg)
+
+                                await asyncio.sleep(10)
 
                 except OperationTimeoutError:
-                    # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available
+                    # Timeout occurred whilst connecting - this is expected and indicates no non-empty sessions are available
                     logger.debug("No sessions for this process. Will look again...")
 
                 except ServiceBusConnectionError:
                     # Occasionally there will be a transient / network-level error in connecting to SB.
                     logger.info("Unknown Service Bus connection error. Will retry...")
+                    await asyncio.sleep(10)
 
                 except Exception as e:
                     # Catch all other exceptions, log them via .exception to get the stack trace, and reconnect
                     logger.exception(f"Unknown exception. Will retry - {e}")
+                    await asyncio.sleep(10)
 
     async def process_message(self, msg):
         with tracer.start_as_current_span("process_message") as current_span:

api_app/service_bus/deployment_status_updater.py

@@ -43,43 +43,46 @@ async def receive_messages(self):
 
             while True:
                 try:
-                    current_time = time.time()
-                    polling_count += 1
-                    # Log a heartbeat message every 60 seconds to show the service is still working
-                    if current_time - last_heartbeat_time >= 60:
-                        logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue {polling_count} times in the last minute")
-                        last_heartbeat_time = current_time
-                        polling_count = 0
-
                     async with credentials.get_credential_async_context() as credential:
-                        service_bus_client = ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential)
-
-                        logger.debug(f"Looking for new messages on {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue...")
-                        # max_wait_time=1 -> don't hold the session open after processing of the message has finished
-                        async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=1, session_id=NEXT_AVAILABLE_SESSION) as receiver:
-                            logger.info(f"Got a session containing messages: {receiver.session.session_id}")
-                            async with AutoLockRenewer() as renewer:
-                                renewer.register(receiver, receiver.session, max_lock_renewal_duration=60)
-                                async for msg in receiver:
-                                    complete_message = await self.process_message(msg)
-                                    if complete_message:
-                                        await receiver.complete_message(msg)
-                                    else:
-                                        # could have been any kind of transient issue, we'll abandon back to the queue, and retry
-                                        await receiver.abandon_message(msg)
-                            logger.info(f"Closing session: {receiver.session.session_id}")
-
-                except OperationTimeoutError:
-                    # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available
-                    logger.debug("No sessions for this process. Will look again...")
+                        async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client:
+                            while True:
+                                try:
+                                    current_time = time.time()
+                                    polling_count += 1
+                                    # Log a heartbeat message every 60 seconds to show the service is still working
+                                    if current_time - last_heartbeat_time >= 60:
+                                        logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue {polling_count} times in the last minute")
+                                        last_heartbeat_time = current_time
+                                        polling_count = 0
+
+                                    logger.debug(f"Looking for new messages on {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue...")
+                                    # max_wait_time=1 -> don't hold the session open after processing of the message has finished
+                                    async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=1, session_id=NEXT_AVAILABLE_SESSION) as receiver:
+                                        logger.info(f"Got a session containing messages: {receiver.session.session_id}")
+                                        async with AutoLockRenewer() as renewer:
+                                            renewer.register(receiver, receiver.session, max_lock_renewal_duration=60)
+                                            async for msg in receiver:
+                                                complete_message = await self.process_message(msg)
+                                                if complete_message:
+                                                    await receiver.complete_message(msg)
+                                                else:
+                                                    # could have been any kind of transient issue, we'll abandon back to the queue, and retry
+                                                    await receiver.abandon_message(msg)
+                                        logger.info(f"Closing session: {receiver.session.session_id}")
+
+                                except OperationTimeoutError:
+                                    # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available
+                                    logger.debug("No sessions for this process. Will look again...")
 
                 except ServiceBusConnectionError:
                     # Occasionally there will be a transient / network-level error in connecting to SB.
                     logger.info("Unknown Service Bus connection error. Will retry...")
+                    await asyncio.sleep(10)
 
                 except Exception as e:
                     # Catch all other exceptions, log them via .exception to get the stack trace, and reconnect
                     logger.exception(f"Unknown exception. Will retry - {e}")
+                    await asyncio.sleep(10)
 
     async def process_message(self, msg):
         complete_message = False

resource_processor/tests_rp/test_runner.py

@@ -50,6 +50,7 @@ async def setup_service_bus_client_and_credential(mock_service_bus_client, mock_
     mock_credential = AsyncMock()
     mock_default_credential.return_value.__aenter__.return_value = mock_credential
     mock_service_bus_client_instance = mock_service_bus_client.return_value
+    mock_service_bus_client.return_value.__aenter__.return_value = mock_service_bus_client_instance
     return mock_service_bus_client_instance, mock_credential
 
 

resource_processor/vmss_porter/runner.py

@@ -100,11 +100,12 @@ async def receive_message(service_bus_client, config: dict, keep_running=lambda:
         except ServiceBusConnectionError:
             # Occasionally there will be a transient / network-level error in connecting to SB.
             logger.info("Unknown Service Bus connection error. Will retry...")
+            await asyncio.sleep(10)
 
         except Exception:
             # Catch all other exceptions, log them via .exception to get the stack trace, sleep, and reconnect
-
             logger.exception("Unknown exception. Will retry...")
+            await asyncio.sleep(10)
 
 
 async def run_porter(command_parts_list: list, config: dict):
@@ -278,8 +279,8 @@ async def get_porter_outputs(msg_body: dict, config: dict):
 async def runner(process_number: int, config: dict):
     with tracer.start_as_current_span(process_number):
         async with default_credentials(config["vmss_msi_id"]) as credential:
-            service_bus_client = ServiceBusClient(config["service_bus_namespace"], credential)
-            await receive_message(service_bus_client, config)
+            async with ServiceBusClient(config["service_bus_namespace"], credential) as service_bus_client:
+                await receive_message(service_bus_client, config)
 
 
 async def check_runners(processes: list, httpserver: Process, keep_running=lambda: True):