AZ and decommissioning-aware indexing planning

[nadav-govari]

Description

This change adds support for availability zone aware indexing plans, and for allowing decommissioning ingesters to index their own shards to help speed up decommissioning.

The change sits behind a feature flag, default off, that falls back to the existing flow if it's off. The existing unit tests were not touched (other than struct stuff for compilation). Extensive new tests were added, including tests at huge scale, and to ensure that there isn't too much cluster churn between runs.

How was this PR tested?

Extensive unit testing. Cluster tests to follow.

AI Summary of changes (worth reading, I worked hard on making it readable)

Solver level — fixtures where the numbers force the answer, called directly so capacity inflation can't blur the branch.

Test	Pins
test_flag_off_ignores_locality_and_eligibility	AZ and eligibility data cannot reach the legacy solver
test_self_hosted_only_indexer_removes_foreign_work	A draining indexer sheds foreign sharded and non-sharded work
test_draining_indexer_reclaims_own_shards_from_peers	It reclaims its own shards from peers, bounded by its room, and settles
test_reproduce_non_monotonic_inflation	The capacity ladder counterexample that broke plan stability
test_reproduce_6524 (existing)	Idempotency regression, still green after removing its old workaround

Physical plan level — shard identity, not just counts.

Test	Pins
test_locality_aware_topology_matrix_is_stable	Exact local/nearby/remote across no AZ, one, two and three zones; stable on replan and under reordered indexers
test_draining_indexer_keeps_only_hosted_shard_ids	A swapped plan is corrected to exact hosted ownership
test_small_clusters_are_not_unbalanced	2/4/5 indexers stay within fair share, matching the locality-disabled path
test_maybe_build_plan_from_scratch	Below threshold adopts a from-scratch plan; cooldown blocks; above threshold no attempt

Admission and reconciliation.

Test	Pins
test_select_ready_and_draining_indexers	Which statuses are admitted, per variant
test_build_indexer_infos_assigns_draining_eligibility	Eligibility per status, and blanked with the flag off
test_indexing_plans_diff	A status-only transition triggers a rebuild, both argument orders

Property-based. test_proptest_locality_aware_idempotence — random zones, eligibility and affinities; asserts solve is idempotent and no draining indexer exceeds its hosted count. Passes 25k cases. The pre-existing flag-off proptest is untouched.

Scale, 500 indexers, 3 zones, 1000 sources, ≤20k shards, fixed seed: all ready, 250 draining spread across zones, and a whole zone draining. Each asserts every shard placed exactly once, pipeline limits, load balance, no idle indexer, and that a draining indexer only ever indexes shards it hosts.

Successive-rebuild churn — 200 consecutive rebuilds under ingest-driven scale up and down plus rolling decommission, replayed identically with locality disabled.

Headline results

• Locality: 24% → 96% of shards indexed on their host; cross-zone traffic 51% → 0%
• Churn over 200 rebuilds: 8.2% of shards move, against 11.3% with locality disabled — we churn less, not more
• Mass decommission: draining indexers retain 96% of their own shards; zero foreign shards on any draining node
• Plan build cost, release: ~5ms at 100 indexes, ~100ms at 1000

Not covered

No test asserts flag-specific behaviour at the rebuild_plan layer, since the flag is a process-global env read; the suite passes with it on and off. Nothing has run against a real cluster. fmt, clippy, machete and typos not run.

[chatgpt-codex-connector]

💡 Codex Review

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-control-plane/src/indexing_scheduler/scheduling/scheduling_logic.rs#L477-L478
Reclaim draining shards even when the source is fully assigned

When every shard of a source is already assigned to ready peers, unassigned_sources contains no entry for that source, so this loop never calls reclaim_self_hosted_shards_from_peers. For example, a draining indexer with free capacity that hosts a shard currently indexed by a ready indexer in the same AZ will never reclaim it; the physical-plan pass removes and then reassigns it to the ready indexer's unchanged quota, while the locality threshold sees it as nearby and does not rebuild from scratch. This defeats the decommissioning optimization for settled plans, so reclaim candidates need to include fully assigned sources too.

---

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs#L603-L604
Consider every replica when determining a shard's AZ

For a shard replicated across AZs, selecting only the first hosting node makes is_shard_nearby ignore the AZs of all other replicas. Consequently, an assignment to a non-hosting indexer colocated with the second replica is misclassified as remote, and find_nearby_indexer may choose an actually cross-AZ indexer instead. The nearby check should succeed when the candidate shares an availability zone with any shard location, matching is_shard_local's treatment of replicas.

---

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-control-plane/src/indexing_scheduler/mod.rs#L493-L495
Apply the rebuild cooldown after unsuccessful attempts

When locality remains below the threshold but a scratch plan is no better, this return leaves next_plan_from_scratch_timestamp unset. Every subsequent scheduling event therefore solves the complete placement problem twice again, despite the documented 30-minute cooldown; on large clusters with persistently unavoidable low locality, ordinary model changes can repeatedly incur the expensive scratch solve. Record the next-attempt timestamp whenever a scratch rebuild is attempted, not only when its plan is selected.

---

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-control-plane/src/indexing_scheduler/mod.rs#L319-L323
Ignore zero-capacity ready nodes when granting eligibility

If the pool contains a Ready indexer with zero indexing capacity plus positive-capacity draining indexers, this check marks every drainer SelfHostedOnly, after which build_indexer_infos drops the only ready indexer. Any shard not hosted by a drainer then has no eligible destination: capacity inflation cannot change eligibility, so minimal_feasible_inflation_attempt exhausts its attempts and panics. Treat a ready indexer as available here only when its indexing capacity is nonzero, consistent with the subsequent filtering.

---

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-control-plane/src/indexing_scheduler/scheduling/mod.rs#L698-L706
Prevent physical assignment of foreign shards to drainers

The physical allocator does not filter remaining_num_shards_per_node by eligibility, so a SelfHostedOnly indexer can receive a foreign shard when its logical quota cannot be matched to distinct local shard IDs. This occurs with replicated shards, for example when two draining indexers both host the same shard: the count-based solver can assign each a local quota from the overlapping affinities, but after one gets that shard, this fallback assigns another, non-hosted shard to the other drainer. Enforce self-hosting while resolving shard IDs, or make the logical quotas account for overlapping replicas.

---

https://github.com/quickwit-oss/quickwit/blob/43f7a4624f2893725b0a749d0908a631fc15eddd/quickwit-serve/src/lib.rs#L1518
Refresh the pool when an indexer's AZ changes

Although the pool entry now caches availability_zone, indexer_node_changed still filters updates using only status, capacity, and tasks. If a ClusterChange::Update changes or supplies the AZ without changing those other fields, the update is discarded and the scheduler continues using the stale AZ for placement and locality metrics until the node is removed and re-added. Include availability-zone changes in the pool's meaningful-update predicate.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".