codeant-ai-for-open-source[bot] commented on code in PR #42305:
URL: https://github.com/apache/superset/pull/42305#discussion_r3658165602
##########
superset/async_events/async_query_manager.py:
##########
@@ -385,3 +417,79 @@ 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.
+
+ Called from the worker's exception handler, so any cache failure is
+ swallowed and treated as "not cancelled" — a Redis blip must never mask
+ the original error (e.g. a genuine timeout) with a connection error.
+ """
+ if not self._cache:
+ return False
+ try:
+ raw = self._cache.get(self._job_registry_key(job_id))
+ if raw is None:
+ return False
+ return bool(json.loads(raw).get("cancelled"))
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Failed to read cancellation flag for job %s", job_id,
exc_info=True
+ )
+ 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. Write only if the
+ # key still exists (``xx``): if the job finished and cleared its record
+ # between the read above and here, don't recreate a stale record or
+ # revoke a task that is already gone — report it as not found instead.
+ flagged = self._cache.set(
+ key,
+ json.dumps({**record, "cancelled": True}),
+ ex=self._jwt_expiration_seconds or None,
+ xx=True,
+ )
+ if not flagged:
+ raise AsyncQueryJobException("Job not found or already completed")
Review Comment:
**Suggestion:** The read, conditional flag, and terminal-state deletion are
separate Redis operations. If the worker publishes and deletes the registry
after the initial read but before the conditional `set`, cancellation can still
race with terminal processing; conversely, if cancellation flags the record
just before the worker publishes `done`, the worker can emit `done` while the
revoke handler emits `cancelled`, producing multiple contradictory terminal
events. Use an atomic state transition or verify the terminal state before
publishing the cancellation event. [race condition]
<details>
<summary><b>Severity Level:</b> Major ⚠️</summary>
```mdx
- ⚠️ Async query streams can contain conflicting terminal states.
- ⚠️ Explore may render completion after the user stopped it.
- ⚠️ Cancellation responses can race with result delivery.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Submit a chart or Explore query through `submit_chart_data_job()` at
`superset/async_events/async_query_manager.py:343-369` or
`submit_explore_json_job()` at
lines 315-341; `init_job()` registers the job at lines 286-312.
2. Let the worker finish concurrently while `cancel_job()` at lines 449-495
is processing.
The worker publishes a terminal event and deletes the registry in
`update_job()` at lines
418-426.
3. If `cancel_job()` reads the registry at lines 466-473 before that
deletion, it can
subsequently set the cancellation flag at lines 480-485 while the worker is
completing.
4. The worker may then publish `STATUS_DONE` while the revoke path causes
soft-time-limit
handling to publish `STATUS_CANCELLED`, because the terminal publication and
registry
deletion are not one atomic state transition.
5. Poll the job's channel stream through `read_events()` at lines 371-395
and observe
multiple terminal events or a terminal status inconsistent with the user's
cancellation
action.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=be4bd5e30a564ff1995fb6bc750bafd8&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=be4bd5e30a564ff1995fb6bc750bafd8&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:** 480:487
**Comment:**
*Race Condition: The read, conditional flag, and terminal-state
deletion are separate Redis operations. If the worker publishes and deletes the
registry after the initial read but before the conditional `set`, cancellation
can still race with terminal processing; conversely, if cancellation flags the
record just before the worker publishes `done`, the worker can emit `done`
while the revoke handler emits `cancelled`, producing multiple contradictory
terminal events. Use an atomic state transition or verify the terminal state
before publishing the cancellation event.
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=25235e63681545f1a277a675190437e6279753da76ce54230f301aa5eca2a06b&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42305&comment_hash=25235e63681545f1a277a675190437e6279753da76ce54230f301aa5eca2a06b&reaction=dislike'>👎</a>
##########
superset/async_events/async_query_manager.py:
##########
@@ -385,3 +417,79 @@ 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.
+
+ Called from the worker's exception handler, so any cache failure is
+ swallowed and treated as "not cancelled" — a Redis blip must never mask
+ the original error (e.g. a genuine timeout) with a connection error.
+ """
+ if not self._cache:
+ return False
+ try:
+ raw = self._cache.get(self._job_registry_key(job_id))
+ if raw is None:
+ return False
+ return bool(json.loads(raw).get("cancelled"))
+ except Exception: # pylint: disable=broad-except
+ logger.warning(
+ "Failed to read cancellation flag for job %s", job_id,
exc_info=True
+ )
+ 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. Write only if the
+ # key still exists (``xx``): if the job finished and cleared its record
+ # between the read above and here, don't recreate a stale record or
+ # revoke a task that is already gone — report it as not found instead.
+ flagged = self._cache.set(
+ key,
+ json.dumps({**record, "cancelled": True}),
+ ex=self._jwt_expiration_seconds or None,
+ xx=True,
+ )
+ if not flagged:
+ raise AsyncQueryJobException("Job not found or already completed")
+
+ # pylint: disable=import-outside-toplevel
+ from superset.extensions import celery_app
+
+ # SIGUSR1 raises SoftTimeLimitExceeded inside the running task rather
+ # than hard-killing the process, so the worker can emit its terminal
+ # event before exiting.
+ celery_app.control.revoke(job_id, terminate=True, signal="SIGUSR1")
Review Comment:
**Suggestion:** Revoking a task before a worker starts it prevents the task
body from running, so `_handle_soft_time_limit` is never invoked and no
`cancelled` event is emitted. The client can therefore wait indefinitely for a
terminal event for jobs that were still queued when cancellation occurred.
Cancellation needs an explicit terminal-event path for revoked/unstarted tasks,
or a worker-side mechanism that guarantees the event is published. [logic error]
<details>
<summary><b>Severity Level:</b> Critical 🚨</summary>
```mdx
- ❌ Queued Explore queries receive no cancellation event.
- ❌ Polling can wait indefinitely for terminal job state.
- ⚠️ Stop behavior remains stuck during worker saturation.
```
</details>
<details>
<summary><b>Steps of Reproduction ✅ </b></summary>
```mdx
1. Start an Explore async query through
`AsyncQueryManager.submit_chart_data_job()` at
`superset/async_events/async_query_manager.py:343-369`; this registers the
job and submits
the Celery task with the client-visible `job_id` as its task ID.
2. Keep the Celery task queued, for example while all workers are busy, so
the task has
not started executing its worker body.
3. Invoke `AsyncQueryManager.cancel_job()` at
`superset/async_events/async_query_manager.py:449-495`; it sets the Redis
cancellation
flag and calls `celery_app.control.revoke()` at line 495.
4. Celery revocation prevents an unstarted task from being executed, so the
worker-side
soft-time-limit handling described by `cancel_job()` at lines 455-458 is
never reached and
no `STATUS_CANCELLED` event is written by `update_job()` at lines 397-426.
5. The frontend continues polling for a terminal event and does not receive
cancellation
for this queued job, leaving the query-stop flow without the expected
terminal state.
```
</details>
[](https://app.codeant.ai/fix-in-ide?tool=cursor&prompt_id=47681534fd6047f288e526d45980129a&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=47681534fd6047f288e526d45980129a&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:** 495:495
**Comment:**
*Logic Error: Revoking a task before a worker starts it prevents the
task body from running, so `_handle_soft_time_limit` is never invoked and no
`cancelled` event is emitted. The client can therefore wait indefinitely for a
terminal event for jobs that were still queued when cancellation occurred.
Cancellation needs an explicit terminal-event path for revoked/unstarted tasks,
or a worker-side mechanism that guarantees the event is published.
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=61e730afe17df37b5c5e1cc7b3c34fa9863519c38088aae800f5ff52e86e033a&reaction=like'>👍</a>
| <a
href='https://app.codeant.ai/feedback?pr_url=https%3A%2F%2Fgithub.com%2Fapache%2Fsuperset%2Fpull%2F42305&comment_hash=61e730afe17df37b5c5e1cc7b3c34fa9863519c38088aae800f5ff52e86e033a&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]
