uranusjr commented on code in PR #65920:
URL: https://github.com/apache/airflow/pull/65920#discussion_r3257306941
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -2878,25 +2881,45 @@ def check_trigger_timeouts(
self, max_retries: int = MAX_DB_RETRIES, session: Session = NEW_SESSION
) -> None:
"""Mark any "deferred" task as failed if the trigger or execution
timeout has passed."""
- for attempt in run_with_db_retries(max_retries, logger=self.log):
- with attempt:
- result = session.execute(
- update(TI)
- .where(
- TI.state == TaskInstanceState.DEFERRED,
- TI.trigger_timeout < timezone.utcnow(),
+ while True:
+ task_instance_ids = []
+ for attempt in run_with_db_retries(max_retries, logger=self.log):
+ with attempt:
+ now = timezone.utcnow()
+ candidates = (
+ select(TI.id)
+ .where(
+ TI.state == TaskInstanceState.DEFERRED,
+ TI.trigger_timeout < now,
+ )
+ .order_by(TI.id)
+ .limit(_TRIGGER_TIMEOUT_BATCH_SIZE)
)
- .values(
- state=TaskInstanceState.SCHEDULED,
- next_method=TRIGGER_FAIL_REPR,
- next_kwargs={"error":
TriggerFailureReason.TRIGGER_TIMEOUT},
- scheduled_dttm=timezone.utcnow(),
- trigger_id=None,
+ task_instance_ids = list(
+ session.scalars(
+ with_row_locks(candidates, of=TI, session=session,
skip_locked=True)
+ ).all()
)
- )
- num_timed_out_tasks = getattr(result, "rowcount", 0)
- if num_timed_out_tasks:
- self.log.info("Timed out %i deferred tasks without fired
triggers", num_timed_out_tasks)
+ if task_instance_ids:
+ result = session.execute(
+ update(TI)
+ .where(TI.id.in_(task_instance_ids))
+ .values(
+ state=TaskInstanceState.SCHEDULED,
+ next_method=TRIGGER_FAIL_REPR,
+ next_kwargs={"error":
TriggerFailureReason.TRIGGER_TIMEOUT},
+ scheduled_dttm=now,
+ trigger_id=None,
+ )
+ .execution_options(synchronize_session=False)
+ )
+ num_timed_out_tasks = getattr(result, "rowcount", 0)
+ if num_timed_out_tasks:
+ self.log.info(
+ "Timed out %i deferred tasks without fired
triggers", num_timed_out_tasks
+ )
+ if len(task_instance_ids) < _TRIGGER_TIMEOUT_BATCH_SIZE:
+ break
Review Comment:
Instead of checking `< _TRIGGER_TIMEOUT_BATCH_SIZE`, maybe better to just do
```python
if not task_instance_ids:
continue
... do update ...
```
This may cause one more empty query, but may be more robust since the last
query might not contain all remaining rows due to them being locked in the
previous iteration? The code would also be much cleaner.
--
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]
