This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1812de86d6b Add on_kill() hook to BaseTrigger to handle user actions
on triggers (#65590)
1812de86d6b is described below
commit 1812de86d6b2969273156725c9f494374ea64a8b
Author: Amogh Desai <[email protected]>
AuthorDate: Wed Apr 22 18:25:30 2026 +0530
Add on_kill() hook to BaseTrigger to handle user actions on triggers
(#65590)
Deferred tasks don’t currently have a standard way to cancel external jobs
when users take actions (fail, clear, mark success), which can leave orphaned
jobs running.
Add `on_kill()` to BaseTrigger so trigger authors can handle cleanup and
cancel external work when a task is terminated.
---
airflow-core/newsfragments/65590.feature.rst | 1 +
.../src/airflow/config_templates/config.yml | 9 +++
.../src/airflow/jobs/triggerer_job_runner.py | 33 +++++++--
airflow-core/src/airflow/triggers/base.py | 36 ++++++++++
airflow-core/tests/unit/jobs/test_triggerer_job.py | 82 ++++++++++++++++++++++
5 files changed, 157 insertions(+), 4 deletions(-)
diff --git a/airflow-core/newsfragments/65590.feature.rst
b/airflow-core/newsfragments/65590.feature.rst
new file mode 100644
index 00000000000..840a15c6f04
--- /dev/null
+++ b/airflow-core/newsfragments/65590.feature.rst
@@ -0,0 +1 @@
+Add ``on_kill()`` hook to ``BaseTrigger`` so trigger authors can cancel
external jobs (BigQuery, Databricks, etc.) when a user marks a deferred task as
failed, clears it, or marks it succeeded. The hook is a no-op by default and is
only invoked for user-initiated cancellations — triggerer restarts and
redistributions do not trigger it.
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index d9ec3fe9a82..8afce430005 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -2690,6 +2690,15 @@ triggerer:
type: integer
example: ~
default: "50"
+ on_kill_timeout:
+ description: |
+ Maximum number of seconds the triggerer will wait for
``BaseTrigger.on_kill()`` to complete
+ before giving up and logging a warning. Prevents a slow or hung
external API call from
+ blocking the triggerer indefinitely when a deferred task is killed by
a user.
+ version_added: 3.2.0
+ type: integer
+ example: ~
+ default: "30"
queues_enabled:
description: |
When set to True, deferred tasks will register triggers with the task
queue they originated from,
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 74b83a3851a..5d03094032c 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -105,6 +105,10 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)
+# Private sentinel passed as the cancel message when a trigger is cancelled by
user action
+_USER_ACTION_CANCEL_MSG = "__airflow_user_action__"
+
+_ON_CANCEL_TIMEOUT: int = conf.getint("triggerer", "on_kill_timeout",
fallback=30)
def _make_trigger_span(
@@ -1137,12 +1141,15 @@ class TriggerRunner:
Drain the to_cancel queue and ensure all triggers that are not in the
DB are cancelled.
This allows the cleanup job to delete them.
+ Passes "user-action" as the cancel message so that run_trigger() knows
to invoke
+ on_kill(). Triggers in this queue are always removed because this is
the path in which
+ the user performed some action on the task. Trigger redistribution
goes through a separate
+ path.
"""
while self.to_cancel:
trigger_id = self.to_cancel.popleft()
if trigger_id in self.triggers:
- # We only delete if it did not exit already
- self.triggers[trigger_id]["task"].cancel()
+
self.triggers[trigger_id]["task"].cancel(_USER_ACTION_CANCEL_MSG)
await asyncio.sleep(0)
async def cleanup_finished_triggers(self) -> list[int]:
@@ -1331,14 +1338,32 @@ class TriggerRunner:
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_kill().
+ #
+ # For timeout, raise immediately without calling on_kill().
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_CANCEL_MSG:
+ await self.log.ainfo("Trigger cancelled by user action,
invoking on_kill", name=name)
+ try:
+ await asyncio.wait_for(trigger.on_kill(),
timeout=_ON_CANCEL_TIMEOUT)
+ except TimeoutError:
+ await self.log.awarning("on_kill() timed out",
timeout=_ON_CANCEL_TIMEOUT, name=name)
+ except asyncio.CancelledError:
+ # on_kill() was itself cancelled (e.g. triggerer
shutting down
+ # while on_kill() was running). Swallow it so the outer
+ # CancelledError (variable `e`) can be re-raised below.
+ pass
+ except Exception:
+ await self.log.awarning("on_kill() raised an
exception", name=name, exc_info=True)
span.set_status(Status(StatusCode.OK), description=str(e))
raise
except Exception as e:
diff --git a/airflow-core/src/airflow/triggers/base.py
b/airflow-core/src/airflow/triggers/base.py
index 205e830c365..b37448edfea 100644
--- a/airflow-core/src/airflow/triggers/base.py
+++ b/airflow-core/src/airflow/triggers/base.py
@@ -192,6 +192,42 @@ class BaseTrigger(abc.ABC, Templater, LoggingMixin):
and handle it appropriately (in async-compatible way).
"""
+ async def on_kill(self) -> None:
+ """
+ Kill the external job managed by this trigger when the task is killed
by a user.
+
+ Symmetric with ``BaseOperator.on_kill()`` on the worker side: override
this method
+ to stop external work (e.g. cancel a BigQuery job, terminate a
Databricks run) when
+ a user explicitly acts on the deferred task via mark-failed, clear, or
mark-succeeded.
+
+ **Distinction from** ``cleanup()``:
+
+ - ``cleanup()`` runs on every trigger exit — success, timeout,
shutdown, and user
+ kill. It is meant for releasing local resources held by this trigger
instance.
+ Putting external job cancellation in ``cleanup()`` would cancel
in-flight work
+ on every triggerer restart or rolling deploy.
+ - ``on_kill()`` runs only when a user explicitly kills the task. It is
the right
+ place to cancel external work you do not want to keep running after
the user performs an action.
+
+ This only fires when a user acts on the task. It does not fire on:
+
+ - Triggerer shutdown or restart — the trigger is redistributed, not
cancelled.
+ - Triggerer redistribution to another triggerer process.
+ - Trigger timeout — the trigger is killed, not cancelled by user.
+ - Normal trigger completion (the trigger fired an event).
+
+ This method runs in the triggerer's asyncio event loop, so
+ it must be async-safe. Use ``await`` for any I/O; do not block the
event loop.
+
+ Exceptions raised here are logged as warnings and do not
+ propagate — they will not affect the task state or the triggerer.
Implement your
+ own retry or error handling inside this method if needed.
+
+ ``on_kill()`` is given a bounded time to complete. Implementations
+ that call slow external APIs should apply their own timeouts rather
than relying on
+ the framework bound.
+ """
+
@staticmethod
def repr(classpath: str, kwargs: dict[str, Any]):
kwargs_str = ", ".join(f"{k}={v}" for k, v in kwargs.items())
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 503a3f4834c..08c960fd33f 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -46,6 +46,7 @@ from airflow.executors import workloads
from airflow.executors.workloads.task import TaskInstanceDTO
from airflow.jobs.job import Job
from airflow.jobs.triggerer_job_runner import (
+ _USER_ACTION_CANCEL_MSG,
ToTriggerRunner,
ToTriggerSupervisor,
TriggerCommsDecoder,
@@ -382,6 +383,87 @@ class TestTriggerRunner:
)
assert {"event": "Trigger cancelled due to timeout", "log_level":
"error"} in cap_structlog
+ def test_run_trigger_calls_on_kill_for_user_action(self, session,
cap_structlog) -> None:
+ """on_kill() is called when CancelledError carries the user-action
sentinel."""
+ 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_CANCEL_MSG)
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_kill = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.on_kill.assert_awaited_once()
+ assert {
+ "event": "Trigger cancelled by user action, invoking on_kill",
+ "log_level": "info",
+ "name": "mock_name",
+ } in cap_structlog
+
+ def test_run_trigger_skips_on_kill_without_user_action_message(self,
session) -> None:
+ """on_kill() is not called when CancelledError has no user-action
sentinel (shutdown/EOF)."""
+ 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_kill = AsyncMock()
+
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.on_kill.assert_not_called()
+
+ def
test_run_trigger_on_kill_exception_does_not_swallow_cancelled_error(self,
session) -> None:
+ """CancelledError propagates even if on_kill() 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_CANCEL_MSG)
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+ mock_trigger.on_kill = 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))
+
+ mock_trigger.on_kill.assert_awaited_once()
+ mock_trigger.cleanup.assert_awaited_once()
+
+ def test_run_trigger_on_kill_timeout_does_not_block_cleanup(self, session)
-> None:
+ """A hanging on_kill() is interrupted after the timeout and cleanup
still runs."""
+ 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_CANCEL_MSG)
+ mock_trigger.task_instance = MagicMock()
+ mock_trigger.task_instance.map_index = -1
+
+ async def hanging_on_kill():
+ await asyncio.sleep(9999)
+
+ mock_trigger.on_kill = hanging_on_kill
+ mock_trigger.cleanup = AsyncMock()
+
+ with patch("airflow.jobs.triggerer_job_runner._ON_CANCEL_TIMEOUT",
0.01):
+ with pytest.raises(asyncio.CancelledError):
+ asyncio.run(trigger_runner.run_trigger(1, mock_trigger))
+
+ mock_trigger.cleanup.assert_awaited_once()
+
@patch("airflow.jobs.triggerer_job_runner.Trigger._decrypt_kwargs")
@patch(
"airflow.jobs.triggerer_job_runner.TriggerRunner.get_trigger_by_classpath",
