Skip to content

Refactor _executable_task_instances_to_queued to make logic more readable - #66878

Merged
ashb merged 2 commits into
apache:mainfrom
astronomer:scheduler-enqueue-readability
Sep 23, 2026
Merged

ashb merged 2 commits into
apache:mainfrom
astronomer:scheduler-enqueue-readability

Conversation

@ashb

@ashb ashb commented May 13, 2026

Copy link
Copy Markdown
Member

I was looking at the core scheduler logic, and the _executable_task_instances_to_queued was rather large, and hard for even me to understand, and past me wrote a good chunk of it!

This splits it into two (plus a helper) focused methods to aid readability:

  • _acquire_pool_capacity: takes the advisory lock (lifetime of the session, so longer than just this fn) and reads pool utilisation via SELECT FOR UPDATE. Returns (pools, max_tis, starved_pools) so caller can short-circuit when all pools are full before doing any TI selection work.

  • _select_task_instances_to_queue: given pre-computed pool capacity, selects eligible SCHEDULED TIs and moves them to QUEUED. Accepts the pools dict and starved_pools set as parameters, making it directly testable without needing a real lock or DB pool read. This uses the new _build_schedulable_tis_query helper fn to build the complex query.

_critical_section_enqueue_task_instances now calls these two methods in sequence, making the two-phase structure (acquire capacity, then select and queue) more explicit.

All test call sites updated to call _select_task_instances_to_queue directly with a make_pool_stats() helper, removing the dependency on pool row locking in unit tests.

No behaviour changes, just refactoring.

@ashb
ashb requested a review from XD-DENG as a code owner May 13, 2026 16:06
@ashb ashb changed the title Extract _acquire_pool_capacity from _critical_section_enqueue_task_instances Refactor _executable_task_instances_to_queued to make logic more readable May 13, 2026
@ashb ashb added the changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) label May 13, 2026
@ashb ashb added this to the Airflow 3.3.0 milestone May 13, 2026
@ashb

This comment was marked as resolved.

@ashb
ashb force-pushed the scheduler-enqueue-readability branch from 007fd8b to 0e16ac2 Compare May 14, 2026 11:21

@jscheffl jscheffl left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the break-up of the code and logic. That makes totally sense and due to the good pydoc makes it much more readable.

two nit from my side, then LGTM, hope we have this in 3.3.0? Would be cool as a clean start on a new minor.

@vatsrahul1001

Copy link
Copy Markdown
Contributor

@ashb could you help resolving conflicts?

…stances

Split the monolithic `_executable_task_instances_to_queued` into two
focused methods:

- `_acquire_pool_capacity`: takes the advisory lock and reads pool
  utilisation via SELECT FOR UPDATE. Returns (pools, max_tis,
  starved_pools) so callers can short-circuit when all pools are full
  before doing any TI selection work.

- `_select_task_instances_to_queue`: given pre-computed pool capacity,
  selects eligible SCHEDULED TIs and moves them to QUEUED. Accepts the
  pools dict and starved_pools set as parameters, making it directly
  testable without needing a real lock or DB pool read.

`_critical_section_enqueue_task_instances` now calls these two methods
in sequence, making the two-phase structure (acquire capacity, then
select and queue) visible at the orchestration level.

All test call sites updated to call `_select_task_instances_to_queue`
directly with a `make_pool_stats()` helper, removing the dependency on
pool row locking in unit tests.
@ashb
ashb force-pushed the scheduler-enqueue-readability branch from 0e16ac2 to dd58747 Compare September 23, 2026 09:53
@ashb
ashb merged commit 724c022 into apache:main Sep 23, 2026
78 checks passed
@ashb
ashb deleted the scheduler-enqueue-readability branch September 23, 2026 14:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:scheduler changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants