Fix service bus client lifecycle - #4930
Conversation
…status_updater.py, airlock_request_status_update.py, and runner.py to prevent connection socket and AMQP channel leaks
…and deployment_status_updater.py for improved heartbeat logging and error handling; update test_runner.py to mock ServiceBusClient correctly; enhance runner.py with consistent exception handling and retry logic.
Unit Test Results761 tests 761 ✅ 10s ⏱️ Results for commit 8430c9c. ♻️ This comment has been updated with latest results. |
There was a problem hiding this comment.
Pull request overview
This PR fixes inconsistent Azure Service Bus client lifecycle handling across the API service-bus listeners and the resource processor runner to reduce the risk of leaking sockets/AMQP channels during long-running polling and reconnect loops.
Changes:
- Wrap
ServiceBusClientusage in async context managers so clients are deterministically closed. - Add backoff (
asyncio.sleep(10)) on connection/unknown exceptions in the runner and service-bus listeners. - Update runner unit tests to account for the async context manager usage; add a changelog entry.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
api_app/service_bus/deployment_status_updater.py |
Uses async with ServiceBusClient(...) to ensure the client is closed across reconnect cycles. |
api_app/service_bus/airlock_request_status_update.py |
Uses async with ServiceBusClient(...) to ensure the client is closed; adjusts retry sleeps. |
resource_processor/vmss_porter/runner.py |
Uses async with ServiceBusClient(...) for proper client teardown; adds retry sleeps on exceptions. |
resource_processor/tests_rp/test_runner.py |
Updates mocks to handle the ServiceBusClient async context manager behavior. |
CHANGELOG.md |
Adds an unreleased BUG FIXES entry for the lifecycle fix. |
…0.25.17 and 0.13.4; refine debug logging in airlock_request_status_update.py
|
/test-extended |
|
🤖 pr-bot 🤖 🏃 Running extended tests: https://github.com/microsoft/AzureTRE/actions/runs/28157965154 (with refid (in response to this comment from Jack Morris (@rudolphjacksonm)) |
|
/test-extended |
|
🤖 pr-bot 🤖 🏃 Running extended tests: https://github.com/microsoft/AzureTRE/actions/runs/28224500075 (with refid (in response to this comment from maxmartin-cgi) |
|
/test-extended |
|
🤖 pr-bot 🤖 🏃 Running extended tests: https://github.com/microsoft/AzureTRE/actions/runs/28243761740 (with refid (in response to this comment from maxmartin-cgi) |
|
/test-extended |
|
🤖 pr-bot 🤖 (in response to this comment from maxmartin-cgi) |
Marcus Robinson (marrobi)
left a comment
Marcus Robinson (marrobi)
left a comment
There was a problem hiding this comment.
From Opus 4.8
I found one correctness issue and one behavior risk:
-
In
api_app/service_bus/airlock_request_status_update.py,ServiceBusClientis now created once outside the inner polling loop and reused forever. That means if the connection drops or the client becomes unhealthy after a transient failure, the code can keep retrying with the same stale client instead of recreating it. The previous pattern re-established the client each iteration; this change may reduce leaks, but it also makes recovery less robust. Consider moving theasync with ServiceBusClient(...)back inside the retry loop, or explicitly recreating the client after connection-related failures. -
In
deployment_status_updater.py, the logic still keeps a singleServiceBusClientalive across the inner loop as well. If the intent is just to avoid leaks, that’s fine, but the current structure can hold the client open indefinitely while also suppressing reconnect churn. Make sure this is intentional and that there’s a clean exit/recreate path when the receiver or client encounters a fatal state. -
resource_processor/vmss_porter/runner.pynow wrapsServiceBusClientinasync with, but the exception handlers only sleep and continue. Ifreceive_message()ever exits due to a non-fatal error, the outerasync withwill close the client and then immediately reopen it on the nextrunner()invocation, which is fine. Just verify the tests cover the__aenter__/__aexit__path for the new context-manager usage. -
The changelog entry mentions all three files, but the behavioral change in
airlock_request_status_update.pyis materially different from the others: it changes reconnection semantics, not just lifecycle management. That should be called out more explicitly in the PR description or handled with a narrower refactor.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (1)
CHANGELOG.md:18
- This changelog entry links to PR #4930, but the PR description/issue context indicates this fix resolves #4977. Using the wrong reference makes the entry hard to trace and contradicts the stated “resolves #4977” intent.
* 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))
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.
Suppressed comments (1)
CHANGELOG.md:20
- This changelog entry references PR #4930, but the PR description/issue context for this change says it resolves #4977. Using the wrong link makes the changelog hard to trace back to the actual work item.
* 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))
Chris Chapman (ChrisChapman-gh)
left a comment
Chris Chapman (ChrisChapman-gh)
left a comment
There was a problem hiding this comment.
Assuming the conflicts are resolved and the tests pass - looks ok to me.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.
Suppressed comments (2)
api_app/service_bus/deployment_status_updater.py:86
- On OperationTimeoutError, the loop immediately retries without any backoff. If the SDK raises this quickly (or if max_wait_time isn't effectively enforcing a wait in some cases), this can create a hot loop and high log volume. Consider adding a small sleep/backoff here (even a short delay) or otherwise ensuring the retry cadence is bounded.
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...")
resource_processor/tests_rp/test_runner.py:75
- These are async context manager methods; asserting they were 'called' is weaker than asserting they were 'awaited'. To make the intent precise (and avoid false positives with AsyncMock), prefer using assert_awaited_once / assert_awaited_once_with for aenter and aexit (and apply the same change to the similar assertions in the other tests in this file).
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()
…/JC-wk/AzureTRE into fix-service-bus-client-lifecycle
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.
Suppressed comments (2)
Previously missed (1) — in code that hasn't changed since the last review.
resource_processor/vmss_porter/runner.py:283
- This changes the Resource Processor runtime, but
resource_processor/_version.py:1remains at0.13.6. Component changes in this repository require their own semantic-version bump; otherwise the Resource Processor artifact does not expose a new version for this fix.
async with ServiceBusClient(config["service_bus_namespace"], credential) as service_bus_client:
await receive_message(service_bus_client, config)
api_app/_version.py:1
- This lowers the API version from
0.26.5to0.26.1, so builds from this commit would advertise an older version and may be treated as a downgrade. Since this is a backward-compatible bug fix, increment the existing patch version instead.
__version__ = "0.26.6"
| async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client: | ||
| client_created_time = time.time() | ||
| while True: |
| 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: |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.
Suppressed comments (5)
Previously missed (3) — in code that hasn't changed since the last review.
resource_processor/tests_rp/test_runner.py:117
- This mocked
__aexit__assertion does not verify that the real credential is closed on an exception.default_credentialscurrently awaitscredential.close()only afteryield, without afinally; whenreceive_messageraises, the exception is thrown into the generator and that close call is skipped. Wrap the yield intry/finallyand test the real helper'sclose()call so the new “even on exception” claim is valid.
# 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()
resource_processor/vmss_porter/runner.py:283
- This changes Resource Processor runtime behavior, but
resource_processor/_version.pyremains at 0.13.6 while only the API version was incremented. The repository's component versioning policy requires each edited component to be bumped; increment the Resource Processor patch version so deployments can identify and package this fix.
async with ServiceBusClient(config["service_bus_namespace"], credential) as service_bus_client:
await receive_message(service_bus_client, config)
api_app/service_bus/airlock_request_status_update.py:73
- This reduces the successful idle-poll delay from 10 seconds to 1 second. Combined with
max_wait_time=1, an empty queue now creates and closes a receiver roughly every 2 seconds instead of every 11 seconds, substantially increasing Service Bus receiver/link churn. Preserve the previous polling cadence unless a separate latency change is intended and assessed.
await asyncio.sleep(1)
api_app/service_bus/airlock_request_status_update.py:42
- The existing airlock-status tests invoke only
process_message, so none of this new client reuse/cleanup loop is exercised. Add a boundedreceive_messagestest that verifies one client is reused across polls and closed when receiver errors or the hourly rollover exits the inner loop.
async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client:
client_created_time = time.time()
api_app/service_bus/deployment_status_updater.py:52
- The primary lifecycle change in
receive_messagesis untested: the existing deployment-status tests invoke onlyprocess_message. Add a bounded-loop test that verifies the client is reused and that__aexit__runs on a receiver failure and on the hourly rollover; otherwise a regression in the leak fix would not be detected.
async with ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) as service_bus_client:
client_created_time = time.time()
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
resource_processor/vmss_porter/runner.py:285
- This changes the packaged Resource Processor, but
resource_processor/_version.pyremains at0.13.6. That version is used to build and push the Resource Processor image (Makefile targets at lines 101 and 130), so without a patch bump the lifecycle fix may not be published as a new component version. Please increment the Resource Processor version to0.13.7as well.
async with ServiceBusClient(config["service_bus_namespace"], credential) as service_bus_client:
await receive_message(service_bus_client, config)
| 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 |
| # 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() |
| 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) |
| 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) |
| 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) |
| 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) |
| import time | ||
|
|
||
| from mock import AsyncMock, patch | ||
| from mock import AsyncMock, MagicMock, patch |
resolves #4977
What is being addressed
inconsistent ServiceBusClient lifecycle management in deployment_status_updater.py, airlock_request_status_update.py, and runner.py.
key message listeners repeatedly instantiate ServiceBusClient inside infinite loops but discard the instances without calling close() or wrapping them in context managers. This can leak connection sockets and AMQP channels under transient reconnect loops.
How is this addressed
Prevents AMQP channel leaks: Ensures channels are cleaned up even under transient reconnect scenarios
Consistent pattern: All three files now follow the same best practice of using