fix(promql): admit the label a label change writes before building it

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

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();

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<Option<EvaluatedChunk>> {
     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)?,
                         }
                     }

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);
+}