Fix service bus client lifecycle

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ENHANCEMENTS:
 
+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 ([#4930](https://github.com/microsoft/AzureTRE/pull/4930))
+
 ## (0.29.0) (August 14, 2026)
 **BREAKING CHANGES**
 * Remove Windows 10 and dsvm image support from Guacamole. ([#4890](https://github.com/microsoft/AzureTRE/issues/4890))

api_app/_version.py

@@ -1 +1 @@
-__version__ = "0.26.5"
+__version__ = "0.26.1"

api_app/service_bus/airlock_request_status_update.py

@@ -58,19 +58,21 @@ async def receive_messages(self):
                                             # 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)
+                    await asyncio.sleep(10)
 
                 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...")
+                    # Timeout occurred whilst connecting - this is expected and indicates no messages are available
+                    logger.debug("No messages 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,57 @@ 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...")
+                        # We keep a single ServiceBusClient alive across the inner loop to avoid excessive connection
+                        # and reconnection churn, as get_queue_receiver with NEXT_AVAILABLE_SESSION is polled frequently.
+                        # Any fatal connection-related errors or other exceptions (other than OperationTimeoutError)
+                        # will propagate out of the inner loop, closing this context manager and recreating the client.
+                        async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client:
+                            client_created_time = time.time()
+                            while True:
+                                try:
+                                    # Recreate the client periodically (every hour) to ensure connection freshness
+                                    # and avoid holding a potentially stale client open indefinitely.
+                                    if time.time() - client_created_time > 3600:
+                                        logger.info("ServiceBusClient has been active for 1 hour. Recreating for freshness...")
+                                        break
+
+                                    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
 
 
@@ -67,6 +68,12 @@ async def test_runner(mock_receive_message, mock_service_bus_client, mock_defaul
     mock_service_bus_client.assert_called_once_with("test_namespace", mock_credential)
     mock_receive_message.assert_called_once_with(mock_service_bus_client_instance, config)
 
+    # Verify context manager entered and exited cleanly
+    mock_default_credential.return_value.__aenter__.assert_called_once()
+    mock_default_credential.return_value.__aexit__.assert_called_once()
+    mock_service_bus_client.return_value.__aenter__.assert_called_once()
+    mock_service_bus_client.return_value.__aexit__.assert_called_once()
+
 
 @pytest.mark.asyncio
 @patch("vmss_porter.runner.receive_message")
@@ -82,6 +89,12 @@ async def test_runner_no_msi_id(mock_receive_message, mock_service_bus_client, m
     mock_service_bus_client.assert_called_once_with("test_namespace", mock_credential)
     mock_receive_message.assert_called_once_with(mock_service_bus_client_instance, config)
 
+    # Verify context manager entered and exited cleanly
+    mock_default_credential.return_value.__aenter__.assert_called_once()
+    mock_default_credential.return_value.__aexit__.assert_called_once()
+    mock_service_bus_client.return_value.__aenter__.assert_called_once()
+    mock_service_bus_client.return_value.__aexit__.assert_called_once()
+
 
 @pytest.mark.asyncio
 @patch("vmss_porter.runner.receive_message")
@@ -99,6 +112,12 @@ async def test_runner_exception(mock_receive_message, mock_service_bus_client, m
     mock_service_bus_client.assert_called_once_with("test_namespace", mock_credential)
     mock_receive_message.assert_called_once_with(mock_service_bus_client_instance, config)
 
+    # Verify context manager entered and exited cleanly, even on exception
+    mock_default_credential.return_value.__aenter__.assert_called_once()
+    mock_default_credential.return_value.__aexit__.assert_called_once()
+    mock_service_bus_client.return_value.__aenter__.assert_called_once()
+    mock_service_bus_client.return_value.__aexit__.assert_called_once()
+
 
 @pytest.mark.asyncio
 @patch("vmss_porter.runner.invoke_porter_action", return_value=True)

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):