diff --git a/src/promql/src/extension_plan/fused_series.rs b/src/promql/src/extension_plan/fused_series.rs index 7ab19c1d163..b864dfe9460 100644 --- a/src/promql/src/extension_plan/fused_series.rs +++ b/src/promql/src/extension_plan/fused_series.rs @@ -54,7 +54,9 @@ use datafusion::common::DFSchemaRef; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; -use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation, UnboundedMemoryPool}; +use datafusion::execution::memory_pool::{ + MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, +}; use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; @@ -729,9 +731,10 @@ pub struct FusedSeriesExec { input: Arc, /// The columns this node produces. output_schema: SchemaRef, - /// The grid positions every label set in this plan's output already occupies. + /// The grid positions every label set in this plan's output already occupies, and what they + /// keep charged on the memory pool. /// - /// Plan-level rather than per-partition: see [`SharedLabelSets`]. + /// Plan-level rather than per-partition: see [`LabelSets`]. labels: SharedLabelSets, /// Per-node metrics, registered by each execution of the tree. metric: ExecutionPlanMetricsSet, @@ -793,9 +796,9 @@ impl FusedSeriesExec { input_properties.boundedness, )); // Shared by every partition's stream: a label set colliding is a property of the whole - // output vector, not of one partition. Creating it needs no pool; each stream reserves the - // bytes it adds in its own reservation. - let labels: SharedLabelSets = Arc::new(Mutex::new(HashMap::new())); + // output vector, not of one partition. Building it needs no pool; the owner's reservation is + // attached in `execute`, which is the first point that has one. + let labels: SharedLabelSets = Arc::new(Mutex::new(LabelSets::default())); Ok(Self { promql, tree, @@ -877,6 +880,14 @@ impl ExecutionPlan for FusedSeriesExec { partition: usize, context: Arc, ) -> DataFusionResult { + // The label-set state is the plan's, so the reservation that pays for it is attached exactly + // once, to the pool of the first execution. Every partition calls this and only the first + // takes effect: re-registering per stream would move a charge the other partitions keep + // growing, and would give it back when that stream ends while the entries stay. + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .attach(&context.runtime_env().memory_pool); let input = self.input.execute(partition, Arc::clone(&context))?; let mapping = input_mapping(&self.binding, &self.input.schema())?; // Batches become raw series chunks through the shared frame converter, which is what groups @@ -888,10 +899,8 @@ impl ExecutionPlan for FusedSeriesExec { &self.metric, partition, )?; - // What this stream adds to the shared label-set map is its own to account for; the map - // itself outlives it, so the entries are not given back with the reservation. - let reservation = - MemoryConsumer::new("FusedSeriesStream").register(&context.runtime_env().memory_pool); + // What the label sets hold is charged to the reservation attached above, so nothing is + // reserved per stream: the entries outlive every stream and are released with the plan. Ok(Box::pin(FusedSeriesStream { input: Some(input), executor: Some(NativeExecutor::new(state)), @@ -900,7 +909,6 @@ impl ExecutionPlan for FusedSeriesExec { binding: Arc::clone(&self.binding), output_schema: self.output_schema.clone(), labels: Arc::clone(&self.labels), - reservation, ready: VecDeque::new(), done: false, })) @@ -1062,8 +1070,75 @@ fn render_labels(metric: &MetricName) -> String { /// miss exactly those collisions. /// /// The map belongs to the plan, not to one stream: a partition that ends leaves what it claimed -/// for the partitions still running, so the entries are only released with the plan itself. -type SharedLabelSets = Arc>>>; +/// for the partitions still running, so the entries are only released with the plan itself. The +/// charge has the same lifetime, which is what makes bytes the map holds — one entry per label set +/// of the output, kept until the plan dies — visible to the memory pool. Charging each stream for +/// what it inserted instead returned those bytes when that stream ended, while the entries the +/// later partitions are checked against stayed: everything the map held after the first partition +/// was unaccounted for. +#[derive(Default)] +struct LabelSets { + /// The grid positions each final label set already occupies, as merged, sorted runs. + seen: HashMap>, + /// What `seen` keeps charged on the task's memory pool. `None` until the first execution, + /// since construction has no `TaskContext` and therefore no pool to register against. + reservation: Option, +} + +impl LabelSets { + /// Attaches the reservation to the pool of the first execution. + /// + /// Later executions keep the first reservation: it already pays for what `seen` holds, and + /// moving it to another pool would drop the bytes other partitions are still growing. + fn attach(&mut self, pool: &Arc) { + if self.reservation.is_none() { + self.reservation = + Some(MemoryConsumer::new("FusedSeriesExec label sets").register(pool)); + } + } + + /// Claims the occupied grid positions of one chunk of a final series for its label set. + /// + /// A collision is an error only when both series hold a sample at the same evaluation + /// timestamp, which is Prometheus' range-query rule. Series with disjoint points are legal: + /// Prometheus merges them into one output series, while this node emits them as separate + /// batches and defers that merge to matrix assembly, so every sample survives. + fn claim(&mut self, series: &Series) -> DataFusionResult<()> { + let runs = occupied_runs(series)?; + let key = LabelsKey::new(&series.metric); + let Self { seen, reservation } = self; + // Nothing reaches a claim without an `execute` first, and an execution that ran attached the + // reservation. Claiming without one would leave the map growing uncharged, so it is refused + // rather than silently skipped. + let Some(reservation) = reservation.as_mut() else { + return Err(DataFusionError::Execution( + "FusedSeriesExec label-set state is not attached to a memory pool".to_string(), + )); + }; + match seen.entry(key) { + Entry::Occupied(mut entry) => { + let stored = entry.get_mut(); + if runs_overlap(stored, &runs) { + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset: {}", + render_labels(&series.metric) + ))); + } + merge_runs(stored, &runs); + } + Entry::Vacant(entry) => { + // Only the label-set state is this node's own cross-poll state that grows with + // the output, so it is the only thing reserved here. + let bytes = entry.key().len() + runs.len() * 2 * size_of::(); + reservation.try_grow(bytes)?; + entry.insert(runs); + } + } + Ok(()) + } +} + +type SharedLabelSets = Arc>; /// The grid positions the rows of one chunk hold a value at, as sorted, disjoint runs. /// @@ -1190,9 +1265,9 @@ pub struct FusedSeriesStream { /// The schema of every batch this stream yields. output_schema: SchemaRef, /// The label-set occupancy shared with the other partitions' streams. + /// + /// What it holds is charged to the plan's own reservation, attached by the first execution. labels: SharedLabelSets, - /// Accounts for what this stream adds to `labels`; freed when the stream ends. - reservation: MemoryReservation, /// The batches the executor settled since the last poll of the input. ready: VecDeque, /// Whether the input ended and everything it settled has been handed out. @@ -1234,52 +1309,22 @@ impl FusedSeriesStream { self.done = true; self.executor = None; self.input = None; - // The label-set map is shared, so a stream stopping early must not clear entries another - // partition is still checking against: its own reservation is what this stream owns, and - // that is what it gives back. - self.reservation.free(); + // The label-set map and its charge are the plan's, so a stream stopping early gives back + // nothing: the entries another partition is still checking against have to stay charged, and + // are released when the plan itself is dropped. // A batch queued before the error was settled by a run that has now stopped. self.ready.clear(); } - /// Claims the occupied grid positions of one chunk of a final series for its label set. - /// - /// A collision is an error only when both series hold a sample at the same evaluation - /// timestamp, which is Prometheus' range-query rule. Series with disjoint points are legal: - /// Prometheus merges them into one output series, while this node emits them as separate - /// batches and defers that merge to matrix assembly, so every sample survives. - fn claim_label_set(&mut self, series: &Series) -> DataFusionResult<()> { - let runs = occupied_runs(series)?; - let key = LabelsKey::new(&series.metric); - let mut seen = self.labels.lock().unwrap_or_else(|e| e.into_inner()); - match seen.entry(key) { - Entry::Occupied(mut entry) => { - let stored = entry.get_mut(); - if runs_overlap(stored, &runs) { - return Err(DataFusionError::Execution(format!( - "vector cannot contain metrics with the same labelset: {}", - render_labels(&series.metric) - ))); - } - merge_runs(stored, &runs); - } - Entry::Vacant(entry) => { - // Only the label-set state is this node's own cross-poll state that grows with - // the output, so it is the only thing reserved here. - let bytes = entry.key().len() + runs.len() * 2 * size_of::(); - self.reservation.try_grow(bytes)?; - entry.insert(runs); - } - } - Ok(()) - } - /// Converts the chunks the executor settled into batches, in the order it produced them. fn convert(&mut self, chunks: Vec) -> DataFusionResult<()> { for chunk in &chunks { // The label set is claimed before its batch is queued, so a collision surfaces at the // chunk that reached it and nothing this plan would refuse is handed out after it. - self.claim_label_set(&chunk.series)?; + self.labels + .lock() + .unwrap_or_else(|e| e.into_inner()) + .claim(&chunk.series)?; let batch = series_to_record_batch( &self.classification, &self.binding.tags, diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs index c4335420059..f0ea868a3ee 100644 --- a/src/promql/src/extension_plan/fused_series/tests.rs +++ b/src/promql/src/extension_plan/fused_series/tests.rs @@ -42,7 +42,8 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result as DataFusionResult}; use datafusion::execution::context::TaskContext; -use datafusion::execution::memory_pool::UnboundedMemoryPool; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore}; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -248,6 +249,15 @@ fn memory_exec_partitioned(partitions: Vec>) -> Arc) -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(pool) as Arc) + .build() + .expect("the runtime environment builds"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) +} + /// A physical input whose one partition's stream `make` builds. fn staged( schema: SchemaRef, @@ -1465,6 +1475,84 @@ async fn a_duplicate_label_set_on_disjoint_grid_points_is_emitted() { } } +/// What a label set holds stays charged after the stream that first claimed it ends. +/// +/// The map belongs to the plan: it outlives every partition and is released with the plan. Charging +/// each stream for what it inserted returned the bytes of a label set when *that* stream ended, +/// while the entry the next partition is checked against stayed in the map, so everything the map +/// held after the first partition was invisible to the pool. What is charged here is therefore read +/// between the partitions: it has to survive the first stream and grow no further when the second +/// partition lands on the label set the first one claimed. +#[tokio::test] +async fn a_label_set_stays_charged_after_its_first_stream_ends() { + // The first series holds a value only at the first grid point and the second only at the last, + // so the two share a final label set without sharing an evaluation timestamp: the second + // partition is legal and has to be answered, and it has to find the first one's entry there. + let first = vec![series_batch( + "web-0", + &[(60_000, 60.0), (GRID_START, GRID_START as f64 / 1_000.0)], + )]; + let last = GRID_START + (GRID_LEN as i64 - 1) * GRID_STEP; + let second = vec![series_batch( + "web-0", + &[ + (GRID_START, GRID_START as f64 / 1_000.0), + (last, last as f64 / 1_000.0), + ], + )]; + let exec = fused(rate_tree()) + .to_execution_plan(memory_exec_partitioned(vec![first, second])) + .expect("the runtime executes a rate tree"); + assert_eq!( + 2, + exec.properties().output_partitioning().partition_count(), + "this test needs one series in each of two partitions" + ); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + + // The first partition claims the label set, answers it and ends. + let mut stream = exec + .execute(0, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the first partition runs")); + } + assert_eq!(1, answered.len(), "one series over one grid point"); + drop(stream); + let charged = pool.reserved(); + assert!( + charged > 0, + "the first partition's label set is still in the plan's map, but nothing is charged for it" + ); + + // The second partition lands on that same label set on points the first one does not hold: it + // is answered, and what it adds to the entry is already paid for. + let mut stream = exec + .execute(1, Arc::clone(&context)) + .expect("the node executes"); + let mut answered = Vec::new(); + while let Some(batch) = stream.next().await { + answered.push(batch.expect("the second partition runs")); + } + assert_eq!(1, answered.len(), "disjoint points are legal"); + drop(stream); + assert_eq!( + charged, + pool.reserved(), + "a label set the map already holds is not charged a second time" + ); + + // The entries live as long as the plan, so dropping it is what gives their bytes back. + drop(exec); + assert_eq!( + 0, + pool.reserved(), + "the plan's drop releases what its label sets held" + ); +} + /// Each partition runs its own series and answers it, with no label set claimed twice. #[tokio::test] async fn every_partition_evaluates_its_own_series() {