Implements address space cleanup when a workspace service is uninstalled - #4744
Implements address space cleanup when a workspace service is uninstalled#4744James Chapman (JC-wk) wants to merge 112 commits into
Conversation
Unit Test Results907 tests 907 ✅ 13s ⏱️ Results for commit 672526b. ♻️ This comment has been updated with latest results. |
|
I have been testing this for a few days, I am not sure if unit tests are needed and how best to write them if anyone wants to assist. |
There was a problem hiding this comment.
Pull request overview
This PR addresses IP range exhaustion risk by ensuring workspace address spaces allocated by workspace services are freed on successful uninstall, and by triggering a workspace upgrade so downstream infra reflects the removal.
Changes:
- Add post-uninstall cleanup in the service bus deployment status handler to remove a workspace-service
address_spacefrom the parent workspace’saddress_spaces. - Update AzureML and Databricks workspace-service templates to run a workspace
upgradestep after uninstall. - Bump API + template versions and add a changelog entry.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| templates/workspace_services/databricks/template_schema.json | Adds a workspace upgrade step after uninstall (and JSON formatting changes). |
| templates/workspace_services/databricks/porter.yaml | Patch version bump. |
| templates/workspace_services/azureml/template_schema.json | Adds a workspace upgrade step after uninstall. |
| templates/workspace_services/azureml/porter.yaml | Patch version bump. |
| api_app/service_bus/deployment_status_updater.py | Implements address space cleanup after successful uninstall main step. |
| api_app/_version.py | API patch version bump. |
| CHANGELOG.md | Adds an Unreleased entry describing the change. |
There was a problem hiding this comment.
🟡 Changes recommended
Two critical and three moderate workflow, reconciliation, and lease-handling issues remain unresolved.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (5)
api_app/api/routes/workspaces.py:302
- The workspace lease is already held here, but this refresh is outside the cleanup
try. A transient read or deployment-validation failure leaves an orphan lease with no operation document, blocking all workspace mutations for up to two hours. Ensure this failure path releases the newly acquired lease.
workspace = await workspace_repo.get_deployed_workspace_by_id(workspace.id, operations_repo)
api_app/api/routes/workspaces.py:376
- This post-lease service refresh occurs before the exception handlers that release the lease. If it fails, the lease has no corresponding operation and prevents further workspace changes until stale recovery. Add release-on-failure around this refresh.
workspace_service = await workspace_service_repo.get_workspace_service_by_id(
workspace_service.workspaceId, workspace_service.id)
api_app/api/routes/workspaces.py:530
- Both resource refreshes happen after lease acquisition but outside the guarded block. If either read fails, the orphan lease blocks the workspace for the full expiry period because no operation document was created. Release the lease on refresh failure.
workspace = await workspace_repo.get_deployed_workspace_by_id(workspace.id, operations_repo)
workspace_service = await workspace_service_repo.get_deployed_workspace_service_by_id(
workspace.id, workspace_service.id, operations_repo)
api_app/db/repositories/operations.py:525
- Marking the operation
reconciledbefore resource reconciliation makes a failed reconciliation non-retryable. The terminal operation drops out of the next active-operation query, and lease acquisition treats a terminal reconciled operation as safe to replace, so a new mutation can start while resource statuses remain stale. Persist the terminal operation withreconciled=False, reconcile resources, then set and persistreconciled=True, matching the recovery path above.
op.reconciled = True
api_app/services/airlock.py:550
- With workspace leasing enabled, this non-waiting path cannot dispatch uninstall immediately:
disable_user_resourcehas just started an active upgrade holding the workspace lease, sosend_uninstall_messageacquires a different operation ID and deterministically receives 409. Consequently unhealthy review-VM replacement fails before the redeploy workflow is queued. Queue the workflow after the disable operation and let it wait for disable before starting uninstall, rather than starting both operations concurrently.
if not wait_for_completion:
logger.info(f"Starting deletion of user resource {user_resource.id} after disable operation {disable_op.id}")
return await send_uninstall_message(
- Files reviewed: 37/38 changed files
- Comments generated: 3
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Airlock cleanup sequencing and redeployment recovery can fail permanently under the new lease and retry behavior.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
api_app/services/airlock.py:557
- When
wait_for_completionis false, the disable upgrade has just acquired the workspace lease, so immediately starting an uninstall under a new operation ID is rejected with 409. This makes the unhealthy-review-VM path fail before it can enqueue the redeploy workflow. Return the disable operation instead and enqueue the redeploy withuninstall_started=False, allowing the worker to wait for disable completion before uninstalling.
if wait_for_completion and disable_op and hasattr(disable_op, "id"):
await wait_for_successful_operation(operations_repo, disable_op.id)
if not wait_for_completion:
logger.info(f"Starting deletion of user resource {user_resource.id} after disable operation {disable_op.id}")
return await send_uninstall_message(
- Files reviewed: 39/40 changed files
- Comments generated: 2
- Review effort level: Balanced
There was a problem hiding this comment.
🔵 Needs a closer look
The asynchronous unhealthy-VM path starts uninstall while the disable operation still owns the workspace lease, causing a 409 before redeployment is queued.
Review details
Suppressed comments (1)
api_app/services/airlock.py:557
wait_for_completion=Falsestill dispatches the uninstall immediately after dispatching the disable operation. The disable operation retains the new workspace lease until it reaches a terminal status, sosend_uninstall_messageattempts to acquire the same workspace under a different operation ID and deterministically returns 409; consequentlycreate_review_vmnever reaches the code that enqueues the redeploy workflow for an unhealthy VM. In the asynchronous path, return/enqueue the disable operation and let the workflow wait for it before starting uninstall (withuninstall_started=False), rather than starting both operations here.
if not wait_for_completion:
logger.info(f"Starting deletion of user resource {user_resource.id} after disable operation {disable_op.id}")
return await send_uninstall_message(
- Files reviewed: 39/40 changed files
- Comments generated: 0 new
- Review effort level: Balanced
…ureTRE into 4744-deallocate-ip-addresses
There was a problem hiding this comment.
🟡 Changes recommended
Airlock redeployment sequencing conflicts with its own lease, and stale reconciliation can be marked complete before resource recovery succeeds.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
api_app/services/airlock.py:557
- When
wait_for_completionis false, the disable operation has only been dispatched and still owns the workspace lease. Starting uninstall immediately uses a new operation ID, sosend_uninstall_messageattempts to acquire the same workspace lease and deterministically returns 409. The unhealthy-review-VM path therefore never reaches the code that enqueues the redeploy workflow. Queue the continuation after the disable operation (or have the worker wait for that operation and then start uninstall) instead of dispatching both operations here.
if not wait_for_completion:
logger.info(f"Starting deletion of user resource {user_resource.id} after disable operation {disable_op.id}")
return await send_uninstall_message(
- Files reviewed: 41/42 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Airlock retry recovery can strand redeployments, and workspace-scoped locking does not protect the global CIDR allocator.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
api_app/service_bus/airlock_workflow.py:136
- A failed dispatch leaves a terminal operation document with this persisted workflow operation ID while
save_and_deploy_resourcedeletes the replacement resource. On redelivery no resource is recovered, but_deploy_vmreuses the same ID, socreate_operation_itemhits the existing Cosmos item on every retry and the workflow eventually dead-letters. Allocate and persist a fresh operation ID whenever there is no replacement resource to resume.
if replacement_resource is None:
workflow_state = AirlockRedeployWorkflow(
workflowId=workflow_id,
phase="deploying",
operationId=operation_id)
api_app/service_bus/airlock_workflow.py:128
- Recovering solely from the resource document can skip deployment entirely.
_deploy_vmsaves that document before creating the operation and dispatching its message, so a process termination in that window leaves this lookup successful; the retry then bypasses_deploy_vmand marks the workflow completed even though no install operation exists. Verifyworkflow_state.operationIdexists before accepting the recovered resource, and remove/restart an orphan with a fresh operation ID otherwise.
replacement_resource = await self.user_resource_repo.get_user_resource_by_workflow_id(
workspace_id=payload["review_workspace_id"],
service_id=payload["review_workspace_service_id"],
workflow_id=workflow_id)
- Files reviewed: 41/42 changed files
- Comments generated: 2
- Review effort level: Balanced
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Lease validation races, incomplete reconciliation, and Airlock retry handling can leave resources or workflows in inconsistent states.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (5)
Previously missed (1) — in code that hasn't changed since the last review.
api_app/service_bus/airlock_workflow.py:50
- This renewal window only equals the two sequential operation waits below: cleanup/redeploy can wait up to
WORKSPACE_LEASE_EXPIRY_SECONDSfor disable and then the same duration for uninstall, plus processing overhead. Near those limits the lock expires before settlement, causing redelivery and repeated workflow work. Allow an additional margin (for example, three expiry windows) or split the workflow into independently settled messages.
api_app/db/repositories/operations.py:530
- The operation is persisted with
reconciled=Truebefore_reconcile_operation_resourcessucceeds. If a resource write then fails, subsequent active-operation queries exclude the now-terminal operation, and lease takeover treats it as already reconciled, so later mutations can proceed while resources remain in an active status. Persistreconciled=False, reconcile resources, then persistreconciled=True, matching the two-phase recovery used above in this repository.
op.reconciled = True
api_app/service_bus/airlock_workflow.py:116
- A failed replacement dispatch can permanently wedge this workflow. The workflow state has already persisted this operation ID, while
save_and_deploy_resourcedeletes the replacement resource but leaves the terminal operation record; on redelivery no replacement is found, this same ID is reused, andcreate_operation_itemrepeatedly conflicts when creating the existing operation. Recovery needs to advance the workflow to a fresh operation ID (and safely release/reconcile the prior lease) after a failed dispatch.
replacement_resource = None
operation_id = workflow_state.operationId if workflow_state is not None else str(uuid.uuid4())
api_app/api/routes/workspaces.py:445
- The child-resource validation occurs before the workspace lease is acquired by
send_uninstall_message. A concurrent user-resource creation can complete in this gap; the cascade snapshot taken after lease acquisition then includes and deletes that newly enabled resource without applying the disabled-resource precondition. Acquire the lease first, reload and validate dependencies under it, then dispatch with the same operation ID.
if await operations_repo.resource_has_active_operation(workspace.id):
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=strings.WORKSPACE_HAS_ACTIVE_OPERATION)
if await delete_validation(workspace_service, workspace_service_repo):
api_app/api/routes/workspaces.py:635
- The enabled-state check is made on the dependency-injected snapshot before
send_uninstall_messageacquires the workspace lease. A concurrent patch can complete in between, after which deletion proceeds with the stale disabled object even if the resource is now enabled. Acquire the lease first and reload/revalidate the resource before dispatching the uninstall with that same operation ID.
if await operations_repo.resource_has_active_operation(user_resource.workspaceId):
raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=strings.WORKSPACE_HAS_ACTIVE_OPERATION)
if user_resource.isEnabled:
- Files reviewed: 41/42 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Untrusted workflow authorization, incomplete mutation serialization, and reconciliation ordering can permit unauthorized actions or inconsistent workspace state.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
api_app/db/repositories/operations.py:537
- This marks the stale operation reconciled before
_reconcile_operation_resourcessucceeds. If the operation write succeeds but any resource reconciliation fails, the operation remains terminal withreconciled=True; subsequent active-operation queries exclude it, and lease acquisition skips the terminal-unreconciled recovery branch, so the lease can be replaced while resources remain stuck in active statuses. Persistreconciled=False, reconcile all resources, and only then persistreconciled=True, matching the two-phase recovery used inacquire_workspace_lease.
- Files reviewed: 41/42 changed files
- Comments generated: 3
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Airlock identity handling and lease-retention paths can prevent cleanup or allow concurrent workspace mutations.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (3)
api_app/service_bus/airlock_workflow.py:97
- These assignments replace the persisted review resource location with the workspace's current Airlock configuration before the old resource is fetched. If that configuration changed after the VM was created, cleanup queries the new location, treats
EntityDoesNotExistas success, and leaks the VM at its recorded location. Keep thereview_resourceworkspace/service IDs for lookup and deletion, and use separate current-config IDs only as the redeploy destination.
if airlock_request.type == AirlockRequestType.Import:
review_config = workspace.properties["airlock_review_config"]["import"]
review_workspace_id = review_config["import_vm_workspace_id"]
review_workspace_service_id = review_config["import_vm_workspace_service_id"]
user_resource_template_name = review_config["import_vm_user_resource_template_name"]
api_app/service_bus/airlock_workflow.py:124
- On a retry after the workflow has already persisted a
deployingstate, the original message still hasuninstall_started=False, so this block submits another uninstall for the old (usually already deleted) resource before attempting deployment recovery. Use the persisted workflow state to recognize that uninstall already completed; otherwise transient redeploy failures create duplicate operations and can stall behind the replacement deployment's lease.
if user_resource is not None and not payload.get("uninstall_started", False):
api_app/service_bus/airlock_workflow.py:177
_deploy_vmdeliberately propagateslease_retained=Truewhen the operation may have advanced concurrently, but this handler unconditionally releases that lease and rotates the operation ID. That permits another workspace mutation to overlap the possibly active deployment and defeats the rollback guard. Only release and replace the operation ID when the exception does not retain the lease.
except Exception:
try:
await self.operations_repo.release_workspace_lease(
review_workspace_id, operation_id)
- Files reviewed: 41/42 changed files
- Comments generated: 2
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Lease handling, legacy Airlock recovery, and local queue configuration contain unresolved reliability issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 41/42 changed files
- Comments generated: 3
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Normal lease contention can exhaust the Airlock queue’s delivery attempts and permanently dead-letter cleanup or redeployment work.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 41/42 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Stale active operations can permanently lock workspaces, and Airlock redeployments are marked complete before deployment succeeds.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 43/44 changed files
- Comments generated: 2
- Review effort level: Balanced
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
🟡 Changes recommended
Stale active operations can indefinitely lock workspaces, while the Airlock consumer can block unrelated workflows and discard cleanup when configuration is unavailable.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (2)
api_app/db/repositories/operations.py:217
- The configured expiry is never applied when the lease's operation document still has an active status: this branch always returns 409, regardless of the lease/operation age. If status delivery is lost or dead-lettered, that operation remains active and the workspace is locked indefinitely, so the advertised stale-lease recovery cannot recover this common failure mode. Please add a safe heartbeat/expiry reconciliation path (or another explicit recovery mechanism) for stale active operations.
await self._reconcile_operation(existing_op, timestamp)
await self.release_workspace_lease(workspace_id, existing_op.id)
api_app/service_bus/airlock_workflow.py:103
- Cleanup workflows do not use the source workspace's current Airlock configuration, but this lookup happens for both workflow types. If the source workspace/config was removed or changed while the message was queued, the resulting
EntityDoesNotExist/KeyErroris treated as an invalid message and completed, leaving the review VM undeleted. Resolve the workspace and redeploy configuration only forredeploy; cleanup should rely on the IDs already stored inreviewUserResources.
workspace = await self.workspace_repo.get_workspace_by_id(airlock_request.workspaceId)
if airlock_request.type == AirlockRequestType.Import:
review_config = workspace.properties["airlock_review_config"]["import"]
redeploy_workspace_id = review_config["import_vm_workspace_id"]
redeploy_workspace_service_id = review_config["import_vm_workspace_service_id"]
- Files reviewed: 43/44 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Redeployment recovery can select a failed resource, stale reconciliation leaves active step states, and the workflow consumer can block unrelated work for hours.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
api_app/service_bus/airlock_workflow.py:57
- This receiver awaits each entire workflow before accepting the next message. A workflow can execute several
wait_for_successful_operationcalls, each with a two-hour timeout, so one slow operation can block all Airlock cleanup/redeploy work handled by an API instance for hours. Use short-lived state-machine messages that reschedule while an operation is pending, or process bounded concurrent tasks while retaining workspace-level serialization.
if await self.process_message(message):
- Files reviewed: 43/44 changed files
- Comments generated: 2
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
Lease fencing, stale reconciliation, and Airlock redeployment recovery contain correctness issues that can duplicate resources or persist inconsistent state.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
api_app/db/repositories/operations.py:125
- Reconciliation applies statuses sequentially, but one pipeline can contain multiple steps for the same resource (for example, OHDSI has an earlier workspace update and the final cleanup workspace update). The first entry changes an active resource to a terminal status, so
_reconcile_resource_statusignores the later entry; a timed-out cleanup can therefore leave the workspace shown asUpdatedinstead ofUpdatingFailed. Collapse entries by resource ID, retaining the last step's status, before writing them.
- Files reviewed: 43/44 changed files
- Comments generated: 2
- Review effort level: Balanced
Resolves #4727
PR
What is being addressed
address_spacewhen a workspace service is uninstalled, preventing IP range exhaustion.How is this addressed
CHANGELOG.md, and incremented template versions.