kaxil commented on code in PR #65590:
URL: https://github.com/apache/airflow/pull/65590#discussion_r3117170449
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1331,14 +1335,27 @@ async def run_trigger(
self.events.append((trigger_id, event))
span.set_status(Status(StatusCode.OK))
except asyncio.CancelledError as e:
- # We get cancelled by the scheduler changing the task state.
But if we do lets give a nice error
- # message about it
+ # A trigger can be cancelled for two reasons:
+ # - The user acted on the task (mark failed / clear / mark
succeeded).
+ # - The triggerer is shutting down, here cancel_triggers()
is not
+ # involved — the shutdown path cancels tasks directly
without a message.
+ # Only first case should invoke on_cancel().
+ #
+ # For timeout, raise immediately without calling on_cancel().
if timeout := timeout_after:
timeout = timeout.replace(tzinfo=timezone.utc) if not
timeout.tzinfo else timeout
if timeout < timezone.utcnow():
await self.log.aerror("Trigger cancelled due to
timeout")
span.set_status(Status(StatusCode.ERROR),
description=str(e))
raise
+ if e.args and e.args[0] == USER_ACTION_INDICATOR:
+ # Suppress exceptions from on_cancel() so they do not
digest the
+ # CancelledError errors.
+ with suppress(BaseException):
Review Comment:
`suppress(BaseException)` is too broad. It swallows `SystemExit`,
`KeyboardInterrupt`, and any real bug in `on_cancel()` implementations.
`on_cancel()` is the whole point of the feature -- if every provider
implementation silently fails, no one will know it's broken. Suggest
`suppress(Exception)` and log a warning on failure (similar to how `cleanup()`
is documented at line 190: "Exceptions raised by the cleanup method are
ignored"). Also worth a comment on why `CancelledError` specifically needs to
propagate, since that's the subtle bit.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -382,6 +382,62 @@ def test_run_inline_trigger_timeout(self, session,
cap_structlog) -> None:
)
assert {"event": "Trigger cancelled due to timeout", "log_level":
"error"} in cap_structlog
+ def test_run_trigger_calls_on_cancel_for_user_action(self, session,
cap_structlog) -> None:
+ """on_cancel() is called when CancelledError carries the "user-action"
message."""
+ trigger_runner = TriggerRunner()
+ trigger_runner.triggers = {
+ 1: {"task": MagicMock(spec=asyncio.Task), "is_watcher": False,
"name": "mock_name", "events": 0}
+ }
+ mock_trigger = MagicMock(spec=BaseTrigger)
+ mock_trigger.run.side_effect = asyncio.CancelledError("user-action")
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_cancel = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.on_cancel.assert_awaited_once()
+ assert any("on_cancel" in str(entry) for entry in cap_structlog), (
Review Comment:
`any("on_cancel" in str(entry) for entry in cap_structlog)` is a brittle
substring check. The source log message literally contains `"on_cancel"` so
this passes even if the event structure changes. Prefer asserting against the
structured event (matching the style on line 383):
```python
assert {"event": "Trigger cancelled by user action, invoking on_cancel",
"log_level": "info", "name": "mock_name"} in cap_structlog
```
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -1331,14 +1335,27 @@ async def run_trigger(
self.events.append((trigger_id, event))
span.set_status(Status(StatusCode.OK))
except asyncio.CancelledError as e:
- # We get cancelled by the scheduler changing the task state.
But if we do lets give a nice error
- # message about it
+ # A trigger can be cancelled for two reasons:
+ # - The user acted on the task (mark failed / clear / mark
succeeded).
+ # - The triggerer is shutting down, here cancel_triggers()
is not
+ # involved — the shutdown path cancels tasks directly
without a message.
+ # Only first case should invoke on_cancel().
+ #
+ # For timeout, raise immediately without calling on_cancel().
if timeout := timeout_after:
timeout = timeout.replace(tzinfo=timezone.utc) if not
timeout.tzinfo else timeout
if timeout < timezone.utcnow():
await self.log.aerror("Trigger cancelled due to
timeout")
span.set_status(Status(StatusCode.ERROR),
description=str(e))
raise
+ if e.args and e.args[0] == USER_ACTION_INDICATOR:
+ # Suppress exceptions from on_cancel() so they do not
digest the
+ # CancelledError errors.
+ with suppress(BaseException):
+ await self.log.ainfo(
+ "Trigger cancelled by user action, invoking
on_cancel", name=name
+ )
+ await trigger.on_cancel()
Review Comment:
`await trigger.on_cancel()` has no timeout. If a provider's `on_cancel()`
hangs (stuck HTTP to Databricks/BigQuery, slow external API), the triggerer
blocks on shutdown until SIGKILL. I verified this locally: even with
`suppress(BaseException)`, a hanging `on_cancel()` doesn't get cancelled
because we're already inside the except handler. Wrap in
`asyncio.wait_for(trigger.on_cancel(), timeout=...)` with a sensible bound
(e.g. 30s) and log a warning on timeout.
Also typo in the comment above at line 1352: "do not digest" should be "do
not swallow" (or "do not mask").
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -105,6 +105,7 @@
logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)
+USER_ACTION_INDICATOR = "user-action"
Review Comment:
Using a bare string as a cross-module sentinel is fragile. Anyone calling
`task.cancel("user-action")` from outside this module also triggers
`on_cancel()`. Worth a module-private constant like `_USER_ACTION_CANCEL_MSG =
"__airflow_user_action__"` so it reads as private and is unlikely to collide
with a typo elsewhere.
##########
airflow-core/src/airflow/triggers/base.py:
##########
@@ -192,6 +192,9 @@ async def cleanup(self) -> None:
and handle it appropriately (in async-compatible way).
"""
+ async def on_cancel(self) -> None:
+ """Override to cancel external jobs when a user cancels the deferred
task."""
Review Comment:
This docstring is too terse for a public override hook that provider authors
will implement. Worth documenting: (1) when exactly it fires (user acts on
task: clear/mark failed/mark succeeded), (2) when it does NOT fire (triggerer
shutdown, redistribution, trigger timeout, trigger completion), (3) that it
runs in the triggerer event loop so must be async-safe, (4) exception semantics
(currently swallowed), (5) that timeouts aren't enforced yet so impls should
use their own timeouts. This is the whole public API for the feature.
Also: this PR is a user-facing addition to `BaseTrigger` -- please add a
newsfragment (`airflow-core/newsfragments/65590.feature.rst`) so it shows up in
release notes.
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -382,6 +382,62 @@ def test_run_inline_trigger_timeout(self, session,
cap_structlog) -> None:
)
assert {"event": "Trigger cancelled due to timeout", "log_level":
"error"} in cap_structlog
+ def test_run_trigger_calls_on_cancel_for_user_action(self, session,
cap_structlog) -> None:
+ """on_cancel() is called when CancelledError carries the "user-action"
message."""
+ trigger_runner = TriggerRunner()
+ trigger_runner.triggers = {
+ 1: {"task": MagicMock(spec=asyncio.Task), "is_watcher": False,
"name": "mock_name", "events": 0}
+ }
+ mock_trigger = MagicMock(spec=BaseTrigger)
+ mock_trigger.run.side_effect = asyncio.CancelledError("user-action")
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_cancel = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.on_cancel.assert_awaited_once()
+ assert any("on_cancel" in str(entry) for entry in cap_structlog), (
+ "Expected 'on_cancel' in triggerer log"
+ )
+
+ def test_run_trigger_skips_on_cancel_without_user_action_message(self,
session) -> None:
+ """on_cancel() is not called when CancelledError has no "user-action"
message."""
+ trigger_runner = TriggerRunner()
+ trigger_runner.triggers = {
+ 1: {"task": MagicMock(spec=asyncio.Task), "is_watcher": False,
"name": "mock_name", "events": 0}
+ }
+ mock_trigger = MagicMock(spec=BaseTrigger)
+ mock_trigger.run.side_effect = asyncio.CancelledError()
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_cancel = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.on_cancel.assert_not_called()
+
+ def test_run_trigger_suppresses_on_cancel_exception(self, session) -> None:
+ """CancelledError propagates and cleanup() still runs even if
on_cancel() raises."""
+ trigger_runner = TriggerRunner()
+ trigger_runner.triggers = {
+ 1: {"task": MagicMock(spec=asyncio.Task), "is_watcher": False,
"name": "mock_name", "events": 0}
+ }
+ mock_trigger = MagicMock(spec=BaseTrigger)
+ mock_trigger.run.side_effect = asyncio.CancelledError("user-action")
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_cancel = AsyncMock(side_effect=RuntimeError("external
API down"))
+ mock_trigger.cleanup = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ # CancelledError propagated throuygh, but cleanup still ran
Review Comment:
Typo: "throuygh" -> "through". Also the test name says it tests that cleanup
still runs when `on_cancel()` raises, but cleanup runs in `finally` so it'd
always run regardless of `on_cancel()` behavior. The actual thing being
validated here is that the `CancelledError` still propagates (the first
`pytest.raises`). Worth clarifying the comment to match.
--
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]
