Zonally aware ingest controller - #6795
nadav-govari wants to merge 3 commits into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f82b498415
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let candidates = if !same_zone_minima.is_empty() { | ||
| same_zone_minima | ||
| } else { | ||
| minima |
There was a problem hiding this comment.
Keep balanced allocations within the requested zone
When zones have unequal node counts or existing loads, the requested zone may contain no globally minimum-loaded ingester; this fallback then assigns the shard to another zone. Since ShardPlacement::Balanced first divides every source's new shards evenly by zone but uses this same picker, a deployment with multiple nodes in one zone can place most net-new shards there instead, defeating the zonal balancing this change is intended to provide. Restrict balanced allocations to the requested zone, reserving cross-zone fallback for the rebalance path that explicitly requires global-load priority.
Useful? React with 👍 / 👎.
| let opened = self | ||
| .try_open_shards_by_zone( | ||
| num_shards_by_source, | ||
| requested_zone.clone(), | ||
| eligible_ingesters, | ||
| model, | ||
| progress, | ||
| ) | ||
| .await?; |
There was a problem hiding this comment.
Preserve completed zone batches when a later batch fails
When a request spans multiple zones and a later metastore call fails, this ? discards the counts accumulated for earlier zone batches even though those batches have already been committed to the metastore and inserted into the model. The scale-up caller consequently releases all permits and may retry past its target, while the rebalance caller returns before closing any corresponding originals, leaving surplus open shards. The operation needs to retain or reconcile earlier successes rather than representing the whole multi-transaction operation as failed.
Useful? React with 👍 / 👎.
| enum ShardPlacement { | ||
| Balanced(SourceShardCount), | ||
| Zoned(HashMap<Option<Zone>, SourceShardCount>), |
There was a problem hiding this comment.
Document the zonal placement architecture
This adds a new availability-zone-aware placement mode and changes how shard creation and rebalancing interact with cluster topology, but the commit updates only this Rust file and leaves the repository's architecture documentation without the activation rule, balancing guarantees, or cross-zone fallback behavior. Document these contracts so operators and future scheduler changes do not rely on assumptions that differ from the controller's behavior.
AGENTS.md reference: AGENTS.md:L23-L24
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 04a67d41c0
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| let mut num_shards_by_source_by_zone: HashMap<Option<Zone>, SourceShardCount> = HashMap::new(); | ||
| for (source_uid, num_shards) in source_shard_counts { | ||
| // Number of shards to open for this source in each zone. | ||
| for (zone, count) in distribute_shards_across_zones(num_shards, &zones) { |
There was a problem hiding this comment.
Include existing shards when selecting zones
When a source scales up by fewer shards than there are zones—commonly one shard at a time—this redistributes only the new request and randomly chooses its starting zone without considering that source's existing shards. Repeated scale-ups can therefore keep selecting the same zone (for example, a balanced 1/1/1 source can become 4/1/1), so the resulting source is not zonally balanced even if every allocation stays within its requested zone. Seed the distribution with the source's current per-zone counts from the model and allocate to the least-populated zones.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
https://github.com/quickwit-oss/quickwit/blob/0910f03e9c9f532d9b9b67ecffbf0f475d16a5fc/quickwit-control-plane/src/ingest/ingest_controller.rs#L265-L268
Reject missing ingesters when matching the unzoned AZ
When AZ awareness is disabled, original_zone is None, but this and_then also returns None when the original ingester disappeared from the pool during the async open. That departed shard is therefore treated as a valid match; close_shards subsequently skips it because its ingester is unavailable, while the replacement is still counted as successfully rebalanced. This silently leaves both shards open in the model and can select a departed shard instead of another live unzoned predecessor in the same bucket, so require the pool lookup itself to succeed before comparing zones.
AGENTS.md reference: AGENTS.md:L19-L22
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Description
Refactors the ingest controller to favor zonal balance by spreading shard assignments evenly among all known zones, if availability zone awareness is enabled. If AZ awareness is disabled, all of the zonal logic reduces to one global zonal group (keyed by
None) and falls back to the logic as it exists today.It does this by turning try_open_shards into a zonal function called for each zone.
try_open_shardshas three callsites:get_or_create_open_shards and scale_up_shards create net-new shards, which we ideally want to spread evenly across zones.
rebalance_shards, however, favors keeping shards within the same zone; otherwise, each zone will net-lose shards on a rebalance (since we have to close the original ones). It still favors global equilibrium first, but breaks ties within the same zone.
This allows us to then have an even, strongly local and dense scheduling plan, which is introduced in #6694.
How was this PR tested?
Unit testing and extensive live cluster testing.