kaxil opened a new pull request, #75472:
URL: https://github.com/apache/airflow/pull/75472
The scheduler stops with `UniqueViolation` on `task_instance_current_key`
when it fails a task instance that still has retries left, from either of two
paths:
- **Heartbeat timeout:** any task with retries whose worker stops
heartbeating. The purge always sends a task callback request, so no callback
needs to be configured.
- **Executor event for a task still in `running`** (for example a task
process killed externally): only when the task has `on_failure_callback`,
`on_retry_callback` or email alerts.
## Cause
Both paths call `executor.send_callback(request)` before
`ti.handle_failure(...)`. `DatabaseCallbackSink.send` was `@provide_session`
and received no session. `create_session()` therefore handed it the scheduler
loop's own thread-local scoped session, and the sink committed and closed that
session. Closing expunges everything, so the task instance about to be retried
is detached. `prepare_db_for_next_try` then calls `archive()`, which sets
`working_set = None` and flushes, but nothing is written for a detached
instance. The insert of the next attempt collides with the row that is still
current.
The sink has committed the caller's scoped session since #21731 (Airflow
2.3), and that was harmless until #74222. That PR made retries archive the
attempt and insert a new row, which relies on `archive()` updating the row
through the ORM. 3.4.0b1 and 3.4.0b2 are affected; 3.3.x is not. This should be
backported to `v3-4-test`.
## Fix
`DatabaseCallbackSink.send` now adds the request to the caller's session
(the thread's scoped session when none is passed) and leaves the commit to the
caller. Every scheduler path that sends callbacks commits afterwards anyway:
| Path | Where the commit comes from |
|---|---|
| Executor-event processing, heartbeat purge | Run inside `create_session()`
|
| Stuck-in-queued handling | Commits per task instance |
| Paused-Dag state updates | Commit through their own `@provide_session` |
| `_do_scheduling` | Now commits after each Dag callback |
`_do_scheduling` was the one caller that relied on the sink's commit. After
sending Dag-level callbacks it runs the critical section, which expunges the
session and can roll back.
The Dag-lookup-failure branch of `_process_executor_events` called
`ti.set_state(state)` without a session, which committed and closed the same
shared session and detached the rest of the batch. It now passes
`session=session`.
**Why not pass the session through `executor.send_callback`?** Executors
don't take database sessions in Airflow 3, and provider executors override
`send_callback(request)` with the released signature. The sink is the
scheduler-side piece that writes to the database, so that is where the commit
comes out.
**Why not call `handle_failure` before `send_callback`?** The close would
still detach every task instance processed later in the same batch.
The new tests cover the callback, email, heartbeat-timeout and
Dag-lookup-failure cases. Each fails on `main` with the same `IntegrityError`.
## Reproduction
Postgres, LocalExecutor, `[scheduler] task_instance_heartbeat_timeout = 20`.
The task freezes its supervisor so heartbeats stop, the way a hung worker would:
```python
@dag(tags=["repro"])
def retry_after_heartbeat_timeout():
@task(retries=1)
def stall():
os.kill(os.getppid(), signal.SIGSTOP)
time.sleep(600)
stall()
```
On `main` (ccbd215ff29), the scheduler loop dies on the first timeout:
```
[error] Detected a task instance without a heartbeat: <TaskInstance:
retry_after_heartbeat_timeout.stall manual__2026-10-23T13:11:24.914003+00:00
[running] ...>
[error] Exception when executing SchedulerJob._run_scheduler_loop
File ".../scheduler_job_runner.py", line 4142, in
_purge_task_instances_without_heartbeats
ti.handle_failure(error=msg, session=session)
File ".../taskinstance.py", line 2734, in fetch_handle_failure_context
ti = ti.prepare_db_for_next_try(session)
File ".../taskinstance.py", line 1746, in prepare_db_for_next_try
session.execute(insert(TaskInstance.__table__).values(values))
sqlalchemy.exc.IntegrityError: (psycopg.errors.UniqueViolation) duplicate
key value violates unique constraint "task_instance_current_key"
DETAIL: Key (dag_id, task_id, run_id, region_id, region_index,
working_set)=(retry_after_heartbeat_timeout, stall,
manual__2026-10-23T13:11:24.914003+00:00, 00000000-0000-0000-0000-000000000000,
-1, t) already exists.
[info] Shutting down LocalExecutor; waiting for running tasks to finish.
```
With this change, the task goes to retry and the scheduler keeps running:
```
[error] Detected a task instance without a heartbeat: <TaskInstance:
retry_after_heartbeat_timeout.stall manual__2026-10-23T13:13:27.566460+00:00
[running] ...>
[info] Marking task as UP_FOR_RETRY. dag_id=retry_after_heartbeat_timeout,
task_id=stall, ...
```
```
try_number | state | working_set | archived_reason
------------+--------------+-------------+-----------------
1 | failed | | retry
2 | up_for_retry | t |
```
## Known issues
`TaskInstance.save_to_db`, which `handle_failure` calls, still commits the
session it is given. So the scan's row locks are released after the first
failed task instance in a batch, as they were before this change. This PR only
removes the close that detached the batch.
--
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]
