UN-3721 [FIX] PG finalization strand — terminal guard at the model layer (backend stale-object clobber)#2178
Conversation
Load-test root cause: ~4% PG strands persisted despite the #2172 endpoint guard. Log trace of 4 stranded IDs showed the callback wrote COMPLETED (200), then ~12ms later the row reverted to EXECUTING+NULL with NO status-transition log anywhere — a backend-internal stale-object save (an execution created as EXECUTING with NULL counters, re-.save()d after the callback committed COMPLETED; Django's full-row save reverts status + counters). The #2172 guard is only in the update_status HTTP endpoint, so it can't see backend-internal update_execution/.save() writes. - WorkflowExecution.update_execution + WorkflowExecutionServiceHelper.update_execution: terminal-one-way guard (PG-scoped via queue_message_id) — re-read committed status and skip the ENTIRE save when it would revert a protected COMPLETED/STOPPED (also prevents the stale NULL counters from clobbering). ERROR stays correctable; Celery rows unaffected. - WorkflowHelper._set_result_acknowledge: save(update_fields=['result_acknowledged']) so acknowledging never rewrites status/counters (the confirmed clobber site). 6 backend tests. Full backend finalization suite 22 green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| Filename | Overview |
|---|---|
| backend/workflow_manager/workflow_v2/models/execution.py | Adds persisted PG-path routing, row locking, terminal-state protection, and field-scoped updates. |
| backend/workflow_manager/workflow_v2/execution.py | Routes regular and late-error status writes through the guarded model method. |
| backend/workflow_manager/workflow_v2/workflow_helper.py | Updates result acknowledgement directly without running stale model cache hooks. |
| unstract/core/src/unstract/core/data_models.py | Defines protected terminal statuses while keeping ERROR correctable. |
| backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py | Adds tests for stale PG writes, terminal guards, counters, direct error updates, and result acknowledgement. |
Reviews (2): Last reviewed commit: "Merge remote-tracking branch 'origin/fea..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
muhammad-ali-e
left a comment
muhammad-ali-e
left a comment
There was a problem hiding this comment.
PR Review Toolkit — consolidated findings (Code Reviewer · Silent-Failure Hunter · Type-Design · Test Analyzer · Comment Analyzer · Simplifier).
The terminal-one-way guard is a sound fix for the reported clobber. Below are findings that are new relative to the existing greptile P1 comments (same-status counter clobber @397, check-save race @388-392, stale-cache publish @Helper:443, update_execution_err bypass @execution:200) — I've deliberately not repeated those. Highest-value items first. Inline comments follow.
| # update_status guard can't see. ERROR is deliberately NOT protected (a | ||
| # premature ERROR stays correctable to COMPLETED). Celery rows | ||
| # (queue_message_id NULL) are unaffected. | ||
| if self.queue_message_id is not None: |
There was a problem hiding this comment.
[High] Guard gates on the stale self.queue_message_id — can fail open and NULL out the PG marker.
The guard re-reads the committed status fresh (lines 387-391, correct), but decides whether to run at all from self.queue_message_id — a field on the same stale in-memory object the guard exists to defend against. If this object was snapshotted before _record_dispatch_handle recorded the handle (queue_message_id is written via a scoped save(update_fields=["queue_message_id"]) at execution.py:466), then self.queue_message_id is None, the gate is skipped entirely, and the unconditional self.save() at line 419 not only reverts the protected status but writes NULL back into the queue_message_id column — permanently making the row look like a Celery row and disabling this guard plus the sibling guards (all gate on queue_message_id is not None) for every later write.
Fix: read the marker from the same query as the status and gate on the persisted value:
committed_qmid, committed = (
WorkflowExecution.objects.filter(pk=self.pk)
.values_list("queue_message_id", "status")
.first()
) or (None, None)
if committed_qmid is not None and committed in protected and status.value != committed:
...
returnNo current test exercises this (every _stale() helper reads the row after queue_message_id is set).
There was a problem hiding this comment.
Fixed (20e1c8ec0). The write now ROUTES on the persisted queue_message_id (fresh values_list read), never stale self. A PG row snapshotted before dispatch → routed to the guarded path; and PG writes are field-scoped (update_fields) so queue_message_id is never written → can't be nulled. New test test_pg_stale_null_marker_still_guarded_and_not_nulled.
| status.value, | ||
| self.id, | ||
| ) | ||
| return |
There was a problem hiding this comment.
[High] Early return silently drops error and increment_attempt, not just the status.
The guard returns from inside the if status is not None: block, above the if error: (line ~430) and if increment_attempt: (line ~432) handling. So a call on a protected-terminal row like:
update_execution(status=ExecutionStatus.ERROR, error="post-finalization cleanup failed", increment_attempt=True)correctly refuses the status revert but also silently discards error_message (a real post-terminal failure diagnostic vanishes) and skips attempts += 1 (retry accounting corrupted) — with only a WARNING that mentions status. This re-introduces a silent failure on a different field.
Fix: protect only the status field — on a protected row skip the status assignment / execution_time / rate-limit logic, but fall through to apply error/increment_attempt and save(update_fields=[...]) only the fields you touched. At minimum, include the dropped error/increment_attempt payload in the WARNING so it isn't lost.
There was a problem hiding this comment.
Fixed. The guard now refuses ONLY the status field — it sets self.status = committed and falls through, so error / increment_attempt still apply (attempts incremented from the locked value). Warning names both. New test test_pg_refused_status_still_applies_error_and_attempt.
| status.value, | ||
| execution.id, | ||
| ) | ||
| return |
There was a problem hiding this comment.
[High] HTTP-layer guard: (a) early return drops error/increment_attempt, and (b) unlocked read-modify-write (TOCTOU).
(a) Same issue as the model layer — this return sits inside if status is not None:, above the if error: / if increment_attempt: blocks (lines ~215-218), so a refused call also silently drops the error message and attempt increment.
(b) This method reads execution at line 177, evaluates the guard, then execution.save() at line 220 — with no transaction.atomic() / select_for_update(), unlike the internal_views guard it says it "mirrors" (which locks the row precisely to close this window). A callback can commit COMPLETED between the read and the save; the guard saw a pre-terminal status, passes, and the full-row save() clobbers the finalized row. Wrap the read-check-save in with transaction.atomic(): and re-fetch via select_for_update() when queue_message_id is not None.
(This is the same class as the greptile model-layer race @388-392, but a distinct, unflagged location.)
There was a problem hiding this comment.
Both fixed. The service-helper method now just delegates to the model update_execution (single guard, atomic + field-scoped), eliminating the hand-copied guard, its early-return drop, and its unlocked TOCTOU. New test test_service_helper_update_execution_is_guarded.
| # above, so its status is the committed one. ERROR stays correctable to | ||
| # COMPLETED; Celery rows (queue_message_id NULL) are unaffected. | ||
| if ( | ||
| execution.queue_message_id is not None |
There was a problem hiding this comment.
[Med] Design/simplification: the protected-terminal rule is duplicated as inline literals and has already diverged.
The protected set {COMPLETED, STOPPED} is written two different ways — a tuple here (execution.py:187-190) and a set literal in models/execution.py:392-395 — with duplicated warning text and duplicated 4-/10-line rationale comments. ExecutionStatus already owns every other status subset as a single-sourced classmethod (terminal_values(), failure_statuses()), each with a "one definition so checks stay in sync" comment. This third subset is the only one left inline, so adding a future terminal status (e.g. CANCELLED) silently misses these two spots.
Suggest: add ExecutionStatus.protected_terminal_values() (derive as terminal_values() - {ERROR.value} so new terminals are protected-by-default) and a shared refuses_overwrite(committed, target) predicate; name the PG discriminator as a WorkflowExecution.is_pg_execution property (queue_message_id is not None is repeated at execution.py:185/458, models/execution.py:386, internal_views.py:530). Note the model layer's committed-status re-read is load-bearing (must stay) — only the set/predicate/naming should be shared, so a "make these consistent" refactor mustn't collapse the re-read away.
There was a problem hiding this comment.
Done. Added ExecutionStatus.protected_terminal_values() (= terminal_values() - {ERROR}, so a new terminal is protected-by-default) in shared core; both guard sites use it. The service-helper's duplicated guard is gone (now delegates).
| ex.refresh_from_db() | ||
| assert ex.status == ExecutionStatus.COMPLETED.value | ||
|
|
||
| def test_pg_allows_idempotent_completed_rewrite(self): |
There was a problem hiding this comment.
[Med] Test coverage gaps.
- Idempotent same-status counter clobber is not exercised. This test builds the writer with a fresh
WorkflowExecution.objects.get(pk=ex.pk)(correct counters) and asserts onlystatus— so it can't catch the greptile-reported counter clobber, where the guard is a no-op (status.value == committed) and the trailing fullsave()writes stale/NULL counters. Add a stale-instance variant that nulls the counters and assertsex.successful_files == 1survives. - The HTTP-layer guard (
WorkflowExecutionServiceHelper.update_execution, execution.py:171-220) has no direct test — only the model method and theinternal_viewsview are covered. It carries its own hand-copied guard that can drift; mirror the model cases against the service helper. - Silent side-effect drop (error/increment_attempt on a protected row) is untested at both layers.
- STOPPED counter preservation isn't asserted (
test_pg_refuses_stale_revert_of_stoppedcreates no counters and checks only status) — make it symmetric with the COMPLETED case.
There was a problem hiding this comment.
All added (20e1c8ec0): same-status counter clobber, STOPPED counter preservation, dropped error/increment on a protected row, the service-helper guard, and update_execution_err. 11 stale-writer tests; full suite 29 green.
… field/marker/cache clobber - Guard now atomic (select_for_update) + routes on the persisted queue_message_id (not stale self); PG writes are update_fields-scoped so counters/marker/cache can't be clobbered; refused status still applies error/increment_attempt. - _set_result_acknowledge → queryset .update() (no save/cache republish). - update_execution_err + service-helper update_execution route through the one guard. - ExecutionStatus.protected_terminal_values() single-sources the set. - 11 stale-writer tests; backend finalization suite 27 green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…n' into fix/UN-3445-stale-writer-model-guard # Conflicts: # backend/workflow_manager/execution/tests/test_pg_finalization_fixes.py
Review addressed (
|
| Finding | Fix |
|---|---|
Gated on stale self.queue_message_id → could NULL the marker & disable all guards |
Route on the persisted marker (fresh read); PG writes are update_fields-scoped so queue_message_id is never written |
| Check-save race (not atomic) | PG path reads committed status + writes inside transaction.atomic() + select_for_update() |
| Same-status full save clobbers counters | update_fields only — never touches counters |
Early return drops error/increment_attempt |
Refuse ONLY the status field; error/increment still apply (attempts from the locked value) |
_set_result_acknowledge republishes stale cache |
DB-side .update(result_acknowledged=True) — no save(), no cache write |
update_execution_err bypasses guard |
Routes through update_execution |
| Service-helper hand-copied guard (TOCTOU + dropped fields) | Now delegates to the one model method |
| Duplicated protected set | ExecutionStatus.protected_terminal_values() (= terminal_values() − {ERROR}) in shared core |
Also merged the latest feat (resolved the test-file conflict — kept #2177's RetrieveNotFoundTests alongside the new ModelStaleWriterGuardTests).
11 new stale-writer tests; full backend finalization suite 29 green. SonarCloud was already clean.
|
1 participant
What
Move the terminal-one-way guard from the HTTP endpoint (#2172) down to the model layer, so backend-internal writers can't revert a finalized PG execution.
Why
An 800-user load test still stranded ~4% of PG executions (
EXECUTING+ NULL counters after the callback finalizedCOMPLETED). A log trace of 4 stranded IDs proved it is not a worker/guard-path race:COMPLETED(HTTP 200), then ~12ms later the row reverted toEXECUTING+NULL, with no status-transition log in any container.WorkflowExecutioncreated asEXECUTING(NULL counters) is re-save()d after the callback committedCOMPLETED; Django's full-row.save()reverts status + counters from the stale in-memory copy.The #2172 guard lives only in the
update_statusHTTP endpoint (worker→backend), so it cannot see backend-internalupdate_execution/.save()writes — and those don't log the transition (why the trace looked "clean").How
WorkflowExecution.update_execution(models/execution.py) andWorkflowExecutionServiceHelper.update_execution(execution.py): terminal-one-way guard, PG-scoped viaqueue_message_id. Re-read the committed status; if it isCOMPLETED/STOPPEDand the write would change it, skip the entire save — which also stops the stale NULL counters from clobbering the real ones.ERRORstays correctable toCOMPLETED; Celery rows (NULLqueue_message_id) are untouched.WorkflowHelper._set_result_acknowledge(workflow_helper.py):save(update_fields=["result_acknowledged"])— acknowledging a result must never rewrite status/counters (the confirmed full-save-after-completion site).Tests
6 new backend tests: refuses stale revert of
COMPLETED/STOPPED(status and counters preserved), allowsERROR→COMPLETED, allows idempotent rewrite, Celery row unaffected, acknowledge doesn't touch status/counters. Full backend finalization suite: 22 green.Safety
PG-scoped (
queue_message_id), so the Celery path is unchanged. The guard skips only a reverting write of a protected-terminal status — legitimate finalization (EXECUTING→COMPLETED) and the premature-ERROR→COMPLETEDcorrection still work.UN-3721 — follow-up to #2172 (endpoint guard was necessary but at the wrong layer for this clobber).
🤖 Generated with Claude Code