codeant-ai-for-open-source[bot] commented on code in PR #42305:
URL: https://github.com/apache/superset/pull/42305#discussion_r3628983970
##########
superset/async_events/async_query_manager.py:
##########
@@ -380,3 +412,64 @@ def update_job(
self._cache.xadd(scoped_stream_name, event_data, "*",
self._stream_limit)
self._cache.xadd(full_stream_name, event_data, "*",
self._stream_limit_firehose)
+
+ # Once a job reaches a terminal state its cancel record is dead weight
+ # (and a stale record would let a caller "cancel" a finished job), so
+ # drop it. TTL is the backstop if this is never reached.
+ if status in (self.STATUS_DONE, self.STATUS_ERROR,
self.STATUS_CANCELLED):
+ if job_id := job_metadata.get("job_id"):
+ self._cache.delete(self._job_registry_key(job_id))
+
+ def is_job_cancelled(self, job_id: str) -> bool:
+ """Whether ``cancel_job`` has flagged this job for cancellation."""
+ if not self._cache:
+ return False
+ raw = self._cache.get(self._job_registry_key(job_id))
+ if raw is None:
+ return False
+ try:
+ return bool(json.loads(raw).get("cancelled"))
+ except (ValueError, AttributeError):
+ return False
+
+ def cancel_job(self, job_id: str, channel_id: str, user_id: Optional[int])
-> None:
+ """
+ Authorize and cancel a running async job.
+
+ The caller's ``channel_id`` and ``user_id`` (resolved server-side from
+ the request, never taken from the client) must match the job's original
+ owner. On success the running Celery task is revoked; the SIGUSR1 it
+ receives surfaces as ``SoftTimeLimitExceeded`` in the worker, which
+ reads the cancelled flag set here and emits the terminal
+ ``STATUS_CANCELLED`` event — so exactly one terminal event is written.
+
+ :raises AsyncQueryJobException: the job is unknown or already terminal
+ :raises AsyncQueryTokenException: the caller does not own the job
+ """
+ if not self._cache:
+ raise CacheBackendNotInitialized("Cache backend not initialized")
+
+ key = self._job_registry_key(job_id)
+ raw = self._cache.get(key)
+ if raw is None:
+ raise AsyncQueryJobException("Job not found or already completed")
+
+ record = json.loads(raw)
+ if record.get("channel_id") != channel_id or record.get("user_id") !=
user_id:
+ raise AsyncQueryTokenException("Not authorized to cancel this job")
+
+ # Flag before revoking so the worker's timeout handler, which may fire
+ # almost immediately, reliably sees the cancellation.
+ self._cache.set(
+ key,
+ json.dumps({**record, "cancelled": True}),
+ ex=self._jwt_expiration_seconds or None,
+ )
Review Comment:
**Suggestion:** There is a race condition between reading the job record and
writing the cancellation flag: if the worker finishes and deletes the registry
key in between, this code recreates the key and still returns success, so
callers can “cancel” already-finished jobs and revoke is sent for stale state.
Make this update atomic (for example, conditional write only if the key still
exists and fail otherwise, ideally in a Redis transaction/Lua script). [race
condition]
<details>
<summary><b>Severity Level:</b> Major ⚠️</summary>
```mdx
- ❌ Cancel endpoint reports success cancelling already-completed async jobs.
- ⚠️ Workers may see stale cancellation flag for finished jobs.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. In `superset/async_events/async_query_manager.py:338-363`, a client
starts an async
chart query via `AsyncQueryManager.submit_chart_data_job(...)`, which calls
`init_job(...)` at `281-286`; `init_job` registers the cancellable job in
Redis through
`_register_cancellable_job(...)` at `292-307`.
2. The Celery task created by
`_load_chart_data_into_cache_job.apply_async(...)` at
`352-362` runs and, on successful completion, calls
`AsyncQueryManager.update_job(...)` at
`392-421` with `status=self.STATUS_DONE`, which emits the terminal event and
deletes the
job registry key via `self._cache.delete(self._job_registry_key(job_id))` at
`419-421`.
3. Around the same time the client presses Stop in Explore, hitting the new
cancel
endpoint described in the PR, which calls
`AsyncQueryManager.cancel_job(...)` at
`435-475`; inside `cancel_job`, `key = self._job_registry_key(job_id)` (line
452) and `raw
= self._cache.get(key)` (line 453) read the current record before the
`update_job` call in
step 2 performs its delete.
4. After `update_job` deletes the key, `cancel_job` still proceeds to
`self._cache.set(...)` at `463-467`, recreating the key with `{"cancelled":
True}` and
returning success while the job is already completed;
`celery_app.control.revoke(job_id,
...)` at `475` is then issued for a finished task, so from the caller’s
perspective a
“cancel” appears to succeed based on stale state instead of failing with
“Job not found or
already completed”.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=ff548a405f794cf8bd8e1a035c1a8339&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
[](https://app.codeant.ai/fix-in-ide?tool=vscode-claude&prompt_id=ff548a405f794cf8bd8e1a035c1a8339&service=github&base_url=https%3A%2F%2Fgithub.com&org=apache&repo=apache%2Fsuperset)
*(Use Cmd/Ctrl + Click for best experience)*
<details>
<summary><b>Prompt for AI Agent 🤖 </b></summary>
```mdx
This is a comment left during a code review.
**Path:** superset/async_events/async_query_manager.py
**Line:** 453:467
**Comment:**
*Race Condition: There is a race condition between reading the job
record and writing the cancellation flag: if the worker finishes and deletes
the registry key in between, this code recreates the key and still returns
success, so callers can “cancel” already-finished jobs and revoke is sent for
stale state. Make this update atomic (for example, conditional write only if
the key still exists and fail otherwise, ideally in a Redis transaction/Lua
script).
Validate the correctness of the flagged issue. If correct, How can I resolve
this? If you propose a fix, implement it and please make it concise.
Once fix is implemented, also check other comments on the same PR, and ask
user if the user wants to fix the rest of the comments as well. if said yes,
then fetch all the comments validate the correctness and implement a minimal fix
```
</details>
<a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42305&comment_hash=a200c306dcf18ecf7b5e79441b1d392a8c8b579568ee8b3d4f947217d226d8d4&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42305&comment_hash=a200c306dcf18ecf7b5e79441b1d392a8c8b579568ee8b3d4f947217d226d8d4&reaction=dislike'>👎</a>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
