mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-06 08:09:24 +00:00
416a6cce32
## Summary Design at: https://github.com/openobserve/designs/blob/main/infrastructure/work-group/work_group_v2_impl.md This PR adds slot-based distributed admission control for search queries and fixes file dispatch consistency when a subset of querier nodes is selected per query. To enable this feature we need to set: ``` O2_WORK_GROUP_MAX_NODES_PER_QUERY > 0 O2_WORK_GROUP_NODE_SELECTION_STRATEGY = org | stream ``` ### Slot-based admission - Add `TryAcquire / StartQuery / ReleaseQuery` gRPC RPCs — Follow nodes reserve slots before execution starts and self-release via `FlightEncoderStream::drop` - Wire slot admission in both Arrow Flight search and PromQL paths, co-located with the existing org/user concurrency queue check - Delegate DataFusion `target_partitions` / `memory_size` to enterprise `get_resource_limit()`, which applies per-query slot fraction with optional dynamic expansion - TTL sweep interval derived from `slot_reserved_ttl_ms / 4` (min 100 ms) instead of a hardcoded 500 ms - Move `QUERY_PENDING_NUMS` / `QUERY_RUNNING_NUMS` metric updates to after slot admission completes (a blocked request is no longer counted as "running") ### Consistent-hash file dispatch fix (`max_nodes_per_query > 0`) - Add `get_node_from_consistent_hash_within(key, role, group, allowed)` to `infra::cluster`: walks the global ring clockwise from `hash(key)` but skips nodes outside `allowed`. Preserves all consistent-hash properties (even distribution via vnodes, minimal remapping on membership change) while confining dispatch to the selected subset - `partition_file_by_hash` now passes only the selected querier nodes as the `allowed` set, so every file maps to one of the N selected nodes — no more fallback-to-index-0 imbalance when `strategy=org/stream` - File-list broadcast (enterprise) pre-computes per-item target nodes using the same filtered ring walk via `broadcast_item_targets()` in `o2-enterprise`, so broadcast and dispatch always agree on which node owns each file ### Proto fix - Reserve field 10 in `QueryStatus` (was `WorkGroup` enum / varint wire type) and move `work_group` string to field 13 — prevents wire-type mismatch during rolling upgrades ### Other fixes from code review - `get_first_stream_key()` now returns `"{stream_type}/"` sentinel instead of `""` for multi-stream / cross-index queries, so `select_nodes` always receives a non-empty deterministic key - PromQL node selection uses `"metrics/"` stream key (consistent with the sentinel format) instead of bare `"metrics"` - Move `_orig_trace_id` binding inside `#[cfg(feature = "enterprise")]` block in `FlightEncoderStream::drop` ## Key config (all default to safe/off values) | Env var | Default | Description | |---|---|---| | `O2_WORK_GROUP_MAX_NODES_PER_QUERY` | `0` | Nodes per query; 0 = all nodes (disables slot admission) | | `O2_WORK_GROUP_NODE_SELECTION_STRATEGY` | `all` | `all` / `org` / `stream` | | `O2_WORK_GROUP_SHORT_SLOT_RATIO` | `0.8` | Max fraction of node slots for short queries | | `O2_WORK_GROUP_DYNAMIC_RESOURCE` | `true` | Expand DataFusion allocation on idle nodes | ## Test plan - [x] Default (`max_nodes_per_query=0`): behavior identical to before, zero overhead - [x] `strategy=org`, `max_nodes_per_query=3`, 10-node cluster: each org consistently uses the same 3 nodes; files distributed evenly across those 3 nodes (no node-0 pile-up) - [x] Broadcast routes each file to the same node that `partition_file_by_hash` selects at query time → cache hits on re-queries - [x] Slot admission: 3rd request waits when budget=2 and 2 queries running; slots released on completion - [x] PromQL queries respect slot budget - [x] Rolling upgrade: old node sends `QueryStatus` without field 10 conflict ## Related Requires companion PR in `o2-enterprise`: https://github.com/openobserve/o2-enterprise/pull/1378 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>