diff --git a/src/promql/src/map_label.rs b/src/promql/src/map_label.rs index beb75081127..f84e2e7cc77 100644 --- a/src/promql/src/map_label.rs +++ b/src/promql/src/map_label.rs @@ -139,6 +139,39 @@ impl LabelRewrite { })) } + /// An upper bound on the bytes the label this change writes to `metric` takes, read before it + /// is built. + /// + /// The written value can be far larger than the input: `label_join` concatenates its sources, + /// listed any number of times, and every `$` reference of a `label_replace` template can expand + /// to the whole source value. The bound counts each source once per listing and each `$` as + /// one whole source value. + pub fn written_bytes(&self, metric: &MetricName) -> usize { + let value = |name: &str| metric.get(name).map_or(0, str::len); + match &self.0 { + Change::Replace { + dst, + replacement, + src, + .. + } => { + let references = replacement.matches('$').count(); + dst.len() + .saturating_add(replacement.len()) + .saturating_add(references.saturating_mul(value(src))) + } + Change::Join { + dst, + separator, + src, + } => src.iter().fold( + dst.len() + .saturating_add(separator.len().saturating_mul(src.len())), + |total, name| total.saturating_add(value(name)), + ), + } + } + /// Changes `metric` the way the function changes one series. pub fn apply(&self, metric: &mut MetricName) -> Result<()> { let mut value = String::new(); diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index e0f6d85f2c7..e04ed6ad94e 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -2308,7 +2308,11 @@ impl NodeRuntime { // the input again. let _timer = metrics.timer(); match self.map.as_ref() { - Some(map) => map_chunk(&map.schema, &chunk, &map.ops), + // A label change can write a label far larger than the chunk, so each one is + // admitted before it is built. + Some(map) => map_chunk(&map.schema, &chunk, &map.ops, &mut |bytes| { + self.buffer.grow(bytes) + }), None => Err(DataFusionError::Execution(format!( "the node at {path} holds no map" ))), @@ -4017,15 +4021,22 @@ impl NativeExecutor { /// the operation emptied every position, which yields no chunk at all /// unless this chunk closed its series: a closing chunk is still handed over, holding no position, /// because a consumer tells "nothing more here" from "no more input" by that chunk alone. +/// +/// `admit` is handed the bytes a label change can write before the change builds them; an error +/// from it stops the map. fn map_chunk( schema: &SeriesSchema, chunk: &EvaluatedChunk, ops: &[MapOp], + admit: &mut dyn FnMut(usize) -> Result<()>, ) -> Result> { let mut series = chunk.series.clone(); for (index, op) in ops.iter().enumerate() { match op { - MapOp::Label(rewrite) => rewrite.apply(&mut series.metric)?, + MapOp::Label(rewrite) => { + admit(rewrite.written_bytes(&series.metric))?; + rewrite.apply(&mut series.metric)?; + } MapOp::Math(math) => match map_math(schema, &series, *math)? { Some(mapped) => series = mapped, // `clamp` with `max < min` is the empty vector, and math skips native histograms, @@ -4040,7 +4051,10 @@ fn map_chunk( let mut metric = series.metric.clone(); for op in &ops[index..] { match op { - MapOp::Label(rewrite) => rewrite.apply(&mut metric)?, + MapOp::Label(rewrite) => { + admit(rewrite.written_bytes(&metric))?; + rewrite.apply(&mut metric)?; + } MapOp::Math(_) => metric = drop_metric_name(&metric)?, } } diff --git a/src/promql/src/native/execute/tests.rs b/src/promql/src/native/execute/tests.rs index bff148bead5..00cfa1c1c7d 100644 --- a/src/promql/src/native/execute/tests.rs +++ b/src/promql/src/native/execute/tests.rs @@ -947,9 +947,14 @@ fn a_window_map_window_chain_keeps_each_layer_on_its_own_grid() { let inner_series = inner_evaluator.push(&chunk).unwrap().unwrap(); inner_evaluator.finish().unwrap(); let map_schema = evaluated_schema(Arc::clone(&inner_grid)); - let mapped_chunk = map_chunk(&map_schema, &inner_series, &[MapOp::Math(MathOp::Abs)]) - .unwrap() - .expect("`abs` maps every position, so the chunk keeps its rows"); + let mapped_chunk = map_chunk( + &map_schema, + &inner_series, + &[MapOp::Math(MathOp::Abs)], + &mut |_| Ok(()), + ) + .unwrap() + .expect("`abs` maps every position, so the chunk keeps its rows"); assert!( Arc::ptr_eq( mapped_chunk.series.time.grid().expect("a grid chunk"), @@ -7520,9 +7525,14 @@ fn a_label_only_map_keeps_histogram_samples() { true, ); let rewrite = LabelRewrite::replace("job", "y", "job", "x").unwrap(); - let mapped = map_chunk(&schema, &chunk, &[MapOp::Label(Arc::new(rewrite))]) - .unwrap() - .unwrap(); + let mapped = map_chunk( + &schema, + &chunk, + &[MapOp::Label(Arc::new(rewrite))], + &mut |_| Ok(()), + ) + .unwrap() + .unwrap(); assert_eq!(mapped.series.metric.get("job"), Some("y")); assert_eq!(mapped.series.time.len(), 3); let samples = mapped.series.fields.as_slice()[0] @@ -7537,3 +7547,47 @@ fn a_label_only_map_keeps_histogram_samples() { ] )); } + +/// `label_join` over one source listed a thousand times writes a label a thousand times the +/// input's, so the map admits that before it builds the label: the bound covers the label it +/// writes, and a pool that cannot hold it refuses the chunk and gets every byte back. +#[test] +fn a_label_change_admits_the_label_it_writes_before_building_it() { + let grid = grid(0, 60_000, 1); + let schema = evaluated_schema(Arc::clone(&grid)); + let job = "j".repeat(1024); + let sources = vec!["job"; 1024]; + let join = Arc::new(LabelRewrite::join("dst", "", &sources).unwrap()); + let chunk = evaluated_chunk_of(&schema, named_metric("m", &job), 0, vec![Some(1.0)], true); + + let mut admitted = 0; + map_chunk( + &schema, + &chunk, + &[MapOp::Label(Arc::clone(&join))], + &mut |bytes| { + admitted += bytes; + Ok(()) + }, + ) + .unwrap(); + assert!(admitted >= 1024 * 1024, "{admitted}"); + // Each capture reference of a `label_replace` template can repeat the whole source value. + let replace = LabelRewrite::replace("dst", &"$1".repeat(8), "job", "(.*)").unwrap(); + let mut metric = named_metric("m", &job); + let bound = replace.written_bytes(&metric); + replace.apply(&mut metric).unwrap(); + assert!(bound >= metric.get("dst").unwrap().len() + "dst".len()); + + let pool = bounded_pool(256 * 1024); + let tree = NativeTree::try_new(map_ops( + evaluated_input(Arc::clone(&grid)), + vec![MapOp::Label(join)], + )) + .unwrap(); + let mut execution = executor(&tree, Arc::clone(&pool), &ExecutionPlanMetricsSet::new()); + execution.push_evaluated(chunk).unwrap(); + let error = execution.next_chunk().unwrap_err(); + assert!(error.to_string().contains("Resources exhausted"), "{error}"); + assert_eq!(pool.reserved(), 0); +}