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.6"

api_app/service_bus/airlock_request_status_update.py

@@ -34,43 +34,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_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:
+                        # We keep a single ServiceBusClient alive across the inner loop to avoid excessive connection
+                        # and reconnection churn. Any fatal connection-related errors or other exceptions
+                        # 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:
-                            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
-                    logger.debug("No sessions for this process. Will look again...")
+                            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_STEP_RESULT_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_STEP_RESULT_QUEUE} queue...")
+                                    receiver = service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_STEP_RESULT_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 - 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

api_app/tests_ma/test_service_bus/test_airlock_request_status_update.py

@@ -3,7 +3,7 @@
 import pytest
 import time
 
-from mock import AsyncMock, patch
+from mock import AsyncMock, MagicMock, patch
 from service_bus.airlock_request_status_update import AirlockStatusUpdater
 from models.domain.events import AirlockNotificationUserData, AirlockFile
 from models.domain.airlock_request import AirlockRequest, AirlockRequestStatus, AirlockRequestType
@@ -104,6 +104,112 @@ def __str__(self):
         return self.message
 
 
+class StopReceiveMessages(BaseException):
+    pass
+
+
+def service_bus_client_context():
+    client = MagicMock()
+    client.__aenter__ = AsyncMock(return_value=client)
+    client.__aexit__ = AsyncMock(return_value=False)
+    return client
+
+
+def credential_context():
+    context = MagicMock()
+    context.__aenter__ = AsyncMock(return_value=MagicMock())
+    context.__aexit__ = AsyncMock(return_value=False)
+    return context
+
+
+def queue_receiver_context():
+    receiver = MagicMock()
+    receiver.__aenter__ = AsyncMock(return_value=receiver)
+    receiver.__aexit__ = AsyncMock(return_value=False)
+    receiver.receive_messages = AsyncMock(return_value=[])
+    return receiver
+
+
+async def run_receive_messages_with_mocks(service_bus_client, time_values, client_side_effect=None):
+    updater = AirlockStatusUpdater()
+    credential = credential_context()
+    receiver = queue_receiver_context()
+    service_bus_client.get_queue_receiver.return_value = receiver
+
+    with patch("service_bus.airlock_request_status_update.credentials.get_credential_async_context", return_value=credential), \
+            patch("service_bus.airlock_request_status_update.ServiceBusClient", return_value=service_bus_client, side_effect=client_side_effect), \
+            patch("service_bus.airlock_request_status_update.time.time", side_effect=time_values), \
+            patch("service_bus.airlock_request_status_update.asyncio.sleep", new_callable=AsyncMock):
+        await updater.receive_messages()
+
+
+async def test_receive_messages_reuses_client_for_multiple_polls():
+    service_bus_client = service_bus_client_context()
+    time_call_count = 0
+    client_call_count = 0
+
+    def time_after_two_polls():
+        nonlocal time_call_count
+        time_call_count += 1
+        return 0 if time_call_count <= 5 else 3601
+
+    def create_client(*args, **kwargs):
+        nonlocal client_call_count
+        client_call_count += 1
+        if client_call_count == 1:
+            return service_bus_client
+        raise StopReceiveMessages()
+
+    with pytest.raises(StopReceiveMessages):
+        await run_receive_messages_with_mocks(service_bus_client, time_after_two_polls, create_client)
+
+    assert service_bus_client.get_queue_receiver.call_count == 2
+    service_bus_client.__aenter__.assert_awaited_once()
+    service_bus_client.__aexit__.assert_awaited_once()
+
+
+async def test_receive_messages_closes_client_before_hourly_recreation():
+    first_client = service_bus_client_context()
+
+    def create_client(*args, **kwargs):
+        if create_client.called:
+            raise StopReceiveMessages()
+        create_client.called = True
+        return first_client
+
+    create_client.called = False
+
+    with pytest.raises(StopReceiveMessages):
+        with patch("service_bus.airlock_request_status_update.credentials.get_credential_async_context", return_value=credential_context()), \
+                patch("service_bus.airlock_request_status_update.ServiceBusClient", side_effect=create_client), \
+                patch("service_bus.airlock_request_status_update.time.time", side_effect=[0, 0, 0, 3601]), \
+                patch("service_bus.airlock_request_status_update.asyncio.sleep", new_callable=AsyncMock):
+            first_client.get_queue_receiver.return_value = queue_receiver_context()
+            await AirlockStatusUpdater().receive_messages()
+
+    first_client.__aexit__.assert_awaited_once()
+    assert first_client.get_queue_receiver.call_count == 1
+
+
+async def test_receive_messages_closes_client_after_receiver_failure():
+    service_bus_client = service_bus_client_context()
+    service_bus_client.get_queue_receiver.side_effect = RuntimeError("receiver failed")
+    client_call_count = 0
+
+    def create_client(*args, **kwargs):
+        nonlocal client_call_count
+        client_call_count += 1
+        if client_call_count == 1:
+            return service_bus_client
+        raise StopReceiveMessages()
+
+    with pytest.raises(StopReceiveMessages):
+        await run_receive_messages_with_mocks(service_bus_client, lambda: 0, create_client)
+
+    service_bus_client.__aenter__.assert_awaited_once()
+    service_bus_client.__aexit__.assert_awaited_once()
+
+
 @patch("event_grid.helpers.EventGridPublisherClient")
 @patch('service_bus.airlock_request_status_update.AirlockRequestRepository.create')
 @patch('service_bus.airlock_request_status_update.WorkspaceRepository.create')