From bb20057985de2933a9d84d3e8e3e7f41cfbe875d Mon Sep 17 00:00:00 2001 From: discord9 Date: Thu, 29 Oct 2026 17:38:44 +0800 Subject: [PATCH] feat(promql): route evaluated fragments to their canonical owner (#10038) * feat(promql): route evaluated chains independently by canonical key Signed-off-by: discord9 * test(promql): verify canonical routing key semantics Signed-off-by: discord9 * fix(promql): release closed-output chains and yield routing producers Signed-off-by: discord9 * test(promql): check closed-chain budget before releasing the child Signed-off-by: discord9 * fix(promql): account histogram slices and verify routing completion Signed-off-by: discord9 * fix(promql): finalize contiguous series at repartition owners Wait for successful producer EOF before publishing exact complete metric-name/tag identities. Reuse the existing sample merge for position ordering, complementary fields, and conflicting-cell errors; recompute one final close per owner identity. Keep N producer tasks and M outputs. Admit nonwaiting queue payloads through the query pool, reserve merge workspace, and retain cancellation, failure, memory-release and unread-owner coverage. Signed-off-by: discord9 * feat(promql): publish completed repartition identities incrementally Signed-off-by: discord9 * test(promql): expose histogram amplification in public repartition Signed-off-by: discord9 * fix(datatypes): encode only referenced histogram table rows Signed-off-by: discord9 * test(promql): wait for queued retention before draining a failed read Signed-off-by: discord9 * perf(promql): bound identity rescans and share repeated histogram rows Signed-off-by: discord9 * docs(query): scope the evaluated validator to one chain Signed-off-by: discord9 --------- Signed-off-by: discord9 --- src/common/query/src/series_frame.rs | 7 + src/common/query/src/series_frame/tests.rs | 25 +- .../query/src/series_frame/tests/evaluated.rs | 6 +- .../query/src/series_frame/tests/lifecycle.rs | 53 +- src/datatypes/src/series/frame.rs | 62 +- src/datatypes/src/series/frame/tests.rs | 414 +- src/promql/src/extension_plan.rs | 2 + .../src/extension_plan/series_repartition.rs | 1666 +++++ .../series_repartition/tests.rs | 5765 +++++++++++++++++ src/promql/src/native/execute.rs | 4 +- 10 files changed, 7726 insertions(+), 278 deletions(-) create mode 100644 src/promql/src/extension_plan/series_repartition.rs create mode 100644 src/promql/src/extension_plan/series_repartition/tests.rs diff --git a/src/common/query/src/series_frame.rs b/src/common/query/src/series_frame.rs index 94022e4eb3a..b32e7ff529b 100644 --- a/src/common/query/src/series_frame.rs +++ b/src/common/query/src/series_frame.rs @@ -571,6 +571,13 @@ impl SequenceValidator { /// Cross-chunk validation state for evaluated chains. /// +/// This is a **chain-local** validator. It checks exactly three things: that a chunk does not change +/// identity while its chain is open, that a chain closes exactly once, and that the positions of an +/// open chain keep increasing across adjacent chunks. It keeps at most the one open chain, so it +/// never certifies that a complete identity does not recur later in the stream; that global +/// non-recurrence is a producer contract enforced where identities are completed (the repartition +/// owner), not here. +/// /// `is_last` closes one chunk chain and nothing more. It is not the end of a label set: a range /// function or a label rewrite can make two input series show the same complete labels, and each of /// them closes itself, so a closed chain may be followed by another chain whose labels happen to be diff --git a/src/common/query/src/series_frame/tests.rs b/src/common/query/src/series_frame/tests.rs index 1ed826e3ce4..659a83566b8 100644 --- a/src/common/query/src/series_frame/tests.rs +++ b/src/common/query/src/series_frame/tests.rs @@ -2185,10 +2185,10 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { struct Expected { time: i64, float_bits: Option, - hist_index: Option, + hist_sum: Option, } - // Table rows: 0 => sum 10.5, 1 => sum -9.5, 2 => sum 7.25. Row-local indices stay - // unchanged, so equal indices in different rows keep the same sum. + // Table rows: 0 => sum 10.5, 1 => sum -9.5, 2 => sum 7.25. Each frame row compacts + // referenced histogram rows densely, so the same local index can carry different sums. let nan_bits = f64::from_bits(0x7ff8_0000_0000_0042).to_bits(); let rows: [(Vec, bool); 4] = [ ( @@ -2196,12 +2196,12 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { Expected { time: 10, float_bits: Some(1.0_f64.to_bits()), - hist_index: Some(1), + hist_sum: Some(-9.5), }, Expected { time: 20, float_bits: None, - hist_index: None, + hist_sum: None, }, ], false, @@ -2210,7 +2210,7 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { vec![Expected { time: 30, float_bits: Some(nan_bits), - hist_index: Some(0), + hist_sum: Some(10.5), }], false, ), @@ -2218,7 +2218,7 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { vec![Expected { time: 40, float_bits: Some((-0.0_f64).to_bits()), - hist_index: Some(0), + hist_sum: Some(10.5), }], false, ), @@ -2226,12 +2226,11 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { vec![Expected { time: 50, float_bits: Some(crate::prometheus::PROMETHEUS_STALE_NAN_BITS), - hist_index: Some(2), + hist_sum: Some(7.25), }], true, ), ]; - let sums = [10.5_f64, -9.5, 7.25]; for (expected, last) in rows { let chunk = decoded.next().await.unwrap().unwrap(); assert_eq!(chunk.primary_key.as_slice(), b"series-key"); @@ -2250,12 +2249,12 @@ async fn multirow_ipc_roundtrip_preserves_cross_chunk_variable_segments() { SampleRef::Missing => assert_eq!(expected.float_bits, None), SampleRef::Histogram(_) => panic!("expected a float field"), } - match expected.hist_index { - Some(index_value) => { + match expected.hist_sum { + Some(sum) => { assert!(hist.indices().is_valid(index)); - assert_eq!(hist.indices().value(index), index_value); + assert_eq!(hist.indices().value(index), 0); let view = hist.view_at(index).expect("histogram sample present"); - assert_eq!(view.sum.to_bits(), sums[index_value as usize].to_bits()); + assert_eq!(view.sum.to_bits(), sum.to_bits()); assert_eq!(view.schema, -7); assert_eq!(view.reset_hint, CounterResetHint::Unrecognized(37)); assert_eq!(view.custom_values, [1.5, 4.5]); diff --git a/src/common/query/src/series_frame/tests/evaluated.rs b/src/common/query/src/series_frame/tests/evaluated.rs index b936247be64..26dfdc871a8 100644 --- a/src/common/query/src/series_frame/tests/evaluated.rs +++ b/src/common/query/src/series_frame/tests/evaluated.rs @@ -1286,8 +1286,8 @@ async fn mixed_payload_roundtrips_without_dropping_either_kind() { assert_eq!(mixed.histograms().len(), 4); assert_eq!( mixed.histograms().batch().len(), - 2, - "the referenced table survives whole" + 1, + "only the referenced histogram row survives compaction" ); match field.sample(0) { SampleRef::Float(value) => assert_eq!(value.to_bits(), 1.5_f64.to_bits()), @@ -1297,10 +1297,12 @@ async fn mixed_payload_roundtrips_without_dropping_either_kind() { SampleRef::Histogram(histogram) => { assert_eq!(histogram.count.to_bits(), 17.25_f64.to_bits()); assert_eq!(histogram.sum.to_bits(), (-9.5_f64).to_bits()); + assert_eq!(histogram.start_timestamp, Some(-123)); } _ => panic!("position 1 was a histogram"), } assert!(matches!(field.sample(2), SampleRef::Missing)); + assert!(mixed.histograms().view_at(2).is_none()); match field.sample(3) { // A NaN float stays a float value, and its histogram slot stays empty. SampleRef::Float(value) => assert!(value.is_nan()), diff --git a/src/common/query/src/series_frame/tests/lifecycle.rs b/src/common/query/src/series_frame/tests/lifecycle.rs index 70fe6ede10c..a7b3c36af62 100644 --- a/src/common/query/src/series_frame/tests/lifecycle.rs +++ b/src/common/query/src/series_frame/tests/lifecycle.rs @@ -346,9 +346,9 @@ async fn streaming_round_trip_preserves_sliced_multifield_payload_and_owners() { assert_eq!(mixed_seen, [Some(1.25_f64.to_bits()), None, None]); } -/// One row of one sample, with every array a nonzero-offset slice of a larger owner: the decoded -/// chunk keeps the source buffers by pointer, and the values stay readable and bit-exact after -/// every source owner is dropped. +/// One row of one sample, with payload arrays as nonzero-offset slices of larger owners: the +/// decoded chunk keeps the float and time buffers by pointer, and values stay readable and bit-exact +/// after every source owner is dropped. #[tokio::test] async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() { let schema = Arc::new(SeriesSchema { @@ -388,9 +388,9 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() let index_owner = UInt32Array::from(vec![Some(3), Some(0)]); let indices = index_owner.slice(1, 1); let hist = Samples::try_from_histogram(indices.clone(), sliced_table.clone()).unwrap(); - let mixed_float_owner = Float64Array::from(vec![Some(0.0), Some(-0.0)]); + let mixed_float_owner = Float64Array::from(vec![Some(0.0), None]); let mixed_floats = mixed_float_owner.slice(1, 1); - let mixed_index_owner = UInt32Array::from(vec![Some(0), None]); + let mixed_index_owner = UInt32Array::from(vec![None, Some(0)]); let mixed_indices = mixed_index_owner.slice(1, 1); let mixed = Samples::try_from_mixed( mixed_floats.clone(), @@ -409,9 +409,7 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() let input_times_ptr = chunk.time.as_buffer().as_ptr(); let input_float_ptr = float.values().as_ptr(); - let input_hist_indices_ptr = indices.values().as_ptr(); let input_mixed_float_ptr = mixed_floats.values().as_ptr(); - let input_mixed_indices_ptr = mixed_indices.values().as_ptr(); let input: SendableRawSeriesStream = Box::pin(RawFixture { schema: schema.clone(), @@ -438,18 +436,9 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() else { panic!("expected histogram") }; - assert_eq!( - decoded_hist.indices().values().as_ptr(), - input_hist_indices_ptr - ); - assert!( - decoded_hist - .batch() - .as_arrow() - .to_data() - .ptr_eq(&sliced_table.as_arrow().to_data()), - "a single row keeps the sliced histogram table by pointer" - ); + assert!(decoded_hist.indices().is_valid(0)); + assert_eq!(decoded_hist.indices().value(0), 0); + assert_histogram(decoded_hist.batch().view(0).unwrap(), -9.5, Some(456)); let DenseSamplesRef::Mixed(decoded_mixed) = output[0].fields.get(2).unwrap().dense().unwrap() else { panic!("expected mixed") @@ -458,17 +447,13 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() decoded_mixed.floats().values().as_ptr(), input_mixed_float_ptr ); - assert_eq!( - decoded_mixed.histograms().indices().values().as_ptr(), - input_mixed_indices_ptr - ); - assert!( - decoded_mixed - .histograms() - .batch() - .as_arrow() - .to_data() - .ptr_eq(&sliced_table.as_arrow().to_data()) + assert!(decoded_mixed.floats().is_null(0)); + assert!(decoded_mixed.histograms().indices().is_valid(0)); + assert_eq!(decoded_mixed.histograms().indices().value(0), 0); + assert_histogram( + decoded_mixed.histograms().batch().view(0).unwrap(), + -9.5, + Some(456), ); assert_eq!(output[0].time.as_buffer().as_ptr(), input_times_ptr); @@ -506,8 +491,8 @@ async fn single_row_sliced_payload_keeps_zero_copy_buffers_and_releases_owners() panic!("expected histogram") }; assert_histogram(value, -9.5, Some(456)); - match chunk.fields.get(2).unwrap().sample(0) { - SampleRef::Float(value) => assert_eq!(value.to_bits(), (-0.0_f64).to_bits()), - SampleRef::Missing | SampleRef::Histogram(_) => panic!("expected mixed float"), - } + let SampleRef::Histogram(value) = chunk.fields.get(2).unwrap().sample(0) else { + panic!("expected mixed histogram") + }; + assert_histogram(value, -9.5, Some(456)); } diff --git a/src/datatypes/src/series/frame.rs b/src/datatypes/src/series/frame.rs index 3bdcf581c65..e8934d5f4f4 100644 --- a/src/datatypes/src/series/frame.rs +++ b/src/datatypes/src/series/frame.rs @@ -73,7 +73,7 @@ #[cfg(test)] mod tests; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow::array::{ @@ -81,7 +81,7 @@ use arrow::array::{ StructArray, UInt32Array, UInt64Array, }; use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; -use arrow::compute::concat; +use arrow::compute::{concat, take}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use serde::{Deserialize, Serialize}; @@ -1457,13 +1457,43 @@ fn histogram_charge(values: &HistogramSamples) -> Result { ) } +fn compact_histogram(values: &HistogramSamples) -> Result<(ArrayRef, UInt32Array)> { + let indices = values.indices(); + // One encoded entry per *distinct* referenced row: a segment that references the same row many + // times would otherwise copy that row once per reference, so a small table plus a long run of + // equal indices would still grow the payload with every position. + let mut carried: HashMap = HashMap::new(); + let mut selected = Vec::new(); + let mut remapped = Vec::with_capacity(indices.len()); + for index in 0..indices.len() { + if indices.is_valid(index) { + let original = indices.value(index); + let local_index = match carried.get(&original) { + Some(local_index) => *local_index, + None => { + let local_index = u32::try_from(selected.len()) + .map_err(|_| err("too many referenced histograms in one frame segment"))?; + selected.push(original); + carried.insert(original, local_index); + local_index + } + }; + remapped.push(Some(local_index)); + } else { + remapped.push(None); + } + } + let selected = UInt32Array::from(selected); + let table = arrow(take(values.batch().as_arrow(), &selected, None))?; + Ok((table, UInt32Array::from(remapped))) +} + /// Encodes buffered rows as one batch with direct N-row Arrow array construction. /// -/// Row-local histogram indices stay exactly as encoded: physical tables are concatenated under -/// cumulative row list offsets, so no index is remapped. Each histogram-bearing row carries its -/// complete table, including unreferenced entries; many segments referencing one large table can -/// therefore repeat that table in the encoded batch. A single constituent child array is reused -/// by clone (sharing its buffers); multiple constituents are concatenated, which copies. +/// Each histogram segment carries each *distinct* table row it references exactly once, and its +/// local indices are remapped into that compact table, so the payload grows with the rows a segment +/// actually names rather than with its positions. A single constituent child array is reused by +/// clone (sharing its buffers); multiple constituents are concatenated, which copies. fn encode_rows(schema: &SchemaRef, shape: FrameShape, rows: &[FrameRow]) -> Result { let count = rows.len(); let DataType::Struct(time_fields) = schema.field(2).data_type() else { @@ -1590,20 +1620,22 @@ fn encode_rows(schema: &SchemaRef, shape: FrameShape, rows: &[FrameRow]) -> Resu DenseSamplesRef::Histogram(values) => { float_lengths.push(0); float_valid.push(false); - index_lengths.push(values.len()); - index_parts.push(Arc::new(values.indices().clone())); - table_lengths.push(values.batch().len()); - table_parts.push(Arc::new(values.batch().as_arrow().clone())); + let (table, indices) = compact_histogram(values)?; + index_lengths.push(indices.len()); + index_parts.push(Arc::new(indices)); + table_lengths.push(table.len()); + table_parts.push(table); histogram_valid.push(true); } DenseSamplesRef::Mixed(values) => { float_lengths.push(values.floats().len()); float_parts.push(Arc::new(values.floats().clone())); float_valid.push(true); - index_lengths.push(values.histograms().len()); - index_parts.push(Arc::new(values.histograms().indices().clone())); - table_lengths.push(values.histograms().batch().len()); - table_parts.push(Arc::new(values.histograms().batch().as_arrow().clone())); + let (table, indices) = compact_histogram(values.histograms())?; + index_lengths.push(indices.len()); + index_parts.push(Arc::new(indices)); + table_lengths.push(table.len()); + table_parts.push(table); histogram_valid.push(true); } } diff --git a/src/datatypes/src/series/frame/tests.rs b/src/datatypes/src/series/frame/tests.rs index bf6e5e35d4c..6854ad6b006 100644 --- a/src/datatypes/src/series/frame/tests.rs +++ b/src/datatypes/src/series/frame/tests.rs @@ -231,11 +231,6 @@ fn assert_histogram_with_start( ); } -fn assert_unreferenced_histogram(histogram: crate::series::histogram::HistogramView<'_>) { - assert_eq!(histogram.schema, -7); - assert_histogram_with_start(histogram, 90.5, None); -} - fn assert_float_bits(actual: &Samples, expected: &[Option]) { assert_eq!(actual.len(), expected.len()); for (index, expected) in expected.iter().enumerate() { @@ -347,8 +342,7 @@ fn frame_histogram_lists(batch: &RecordBatch, field: usize) -> (ListArray, ListA (indices, tables) } -/// Every row is expected to resolve index 0 inside its own table only, with the whole table -/// retained and the other index entry null. +/// Each row's referenced histogram is preserved in its compact row-local table. fn assert_row_local_histogram_rows(rows: &[RawSeriesChunk]) { assert_eq!(rows.len(), 2); assert_eq!(rows[0].time.as_buffer().as_ref(), &[1_i64, 2]); @@ -357,23 +351,20 @@ fn assert_row_local_histogram_rows(rows: &[RawSeriesChunk]) { let DenseSamplesRef::Histogram(first) = rows[0].fields.get(0).unwrap().dense().unwrap() else { panic!("expected histogram samples"); }; - assert_eq!(first.batch().len(), 2); + assert_eq!(first.batch().len(), 1); assert_eq!(first.indices().len(), 2); assert_eq!(first.indices().value(0), 0); assert!(!first.indices().is_valid(1)); assert_histogram(first.view_at(0).unwrap(), -9.5); - assert_unreferenced_histogram(first.batch().view(1).unwrap()); let DenseSamplesRef::Histogram(second) = rows[1].fields.get(0).unwrap().dense().unwrap() else { panic!("expected histogram samples"); }; - assert_eq!(second.batch().len(), 2); + assert_eq!(second.batch().len(), 1); assert_eq!(second.indices().len(), 2); assert!(!second.indices().is_valid(0)); assert_eq!(second.indices().value(1), 0); - // A global index remap would resolve index 0 into the first row's table instead. assert_histogram(second.view_at(1).unwrap(), 42.25); - assert_histogram_with_start(second.batch().view(1).unwrap(), 142.25, None); } #[test] @@ -431,23 +422,10 @@ fn sliced_histogram_frames_densify_without_duplicating_shared_table() { DenseSamplesRef::Histogram(histograms) => histograms.batch(), _ => panic!("expected histogram samples"), }; - assert!(!Arc::ptr_eq(&first_table, second_table)); - let first_values = first_table - .as_arrow() - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - let second_values = second_table - .as_arrow() - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - first_values.values().as_ptr(), - second_values.values().as_ptr() - ); + assert_eq!(first_table.len(), 1); + assert_eq!(second_table.len(), 1); + assert_histogram(first_table.view(0).unwrap(), -9.5); + assert_histogram(second_table.view(0).unwrap(), -9.5); // The default budget groups the same 64 rows into one multirow batch; every row still decodes // to its own slice of the shared table with bit-exact values. @@ -490,21 +468,9 @@ fn sliced_histogram_frames_densify_without_duplicating_shared_table() { panic!("expected dense histogram samples"); }; assert_eq!(histograms.len(), 64); - assert_eq!(histograms.batch().len(), first_table.len()); for position in 0..64 { assert_histogram(histograms.view_at(position).unwrap(), -9.5); } - let dense_values = histograms - .batch() - .as_arrow() - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - dense_values.values().as_ptr(), - first_values.values().as_ptr() - ); } #[test] @@ -2730,7 +2696,7 @@ fn rejects_malformed_histogram_and_payload_encodings() { DataType::List(item) => item.clone(), _ => unreachable!(), }, - OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1, 1])), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 1])), Arc::new(arrow::array::Int32Array::from(vec![1])), None, ) @@ -2740,8 +2706,8 @@ fn rejects_malformed_histogram_and_payload_encodings() { DataType::List(item) => item.clone(), _ => unreachable!(), }, - OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 5])), - Arc::new(Float64Array::from(vec![3.25, 7.5, 3.0, 7.0, 8.0])), + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2])), + Arc::new(Float64Array::from(vec![3.25, 7.5])), None, ) .unwrap(); @@ -3165,67 +3131,19 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { .as_any() .downcast_ref::() .unwrap(); - for (packed_indices, source_indices) in [ - (packed_mixed_indices, source_indices), - (packed_hist_indices, &histogram_indices), - ] { - assert_eq!( - packed_indices.values().as_ptr(), - source_indices.values().as_ptr() - ); - assert_eq!( - packed_indices.nulls().unwrap().buffer().as_ptr(), - source_indices.nulls().unwrap().buffer().as_ptr() - ); - assert_eq!( - packed_indices.nulls().unwrap().offset(), - source_indices.nulls().unwrap().offset() - ); - } + assert_eq!( + packed_mixed_indices, + &UInt32Array::from(vec![Some(0), None]) + ); + assert_eq!( + packed_hist_indices, + &UInt32Array::from(vec![Some(0), Some(0)]), + "both positions name the same distinct row" + ); let packed_mixed_table = frame_histogram_table(packed, 0); let packed_hist_table = frame_histogram_table(packed, 1); - for packed_table in [&packed_mixed_table, &packed_hist_table] { - let packed_sum = packed_table - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(packed_sum.values().as_ptr(), sliced_sum.values().as_ptr()); - let packed_positive = packed_table - .column(10) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(packed_positive.offset(), sliced_positive.offset()); - assert_eq!( - packed_positive.value_offsets().as_ptr(), - sliced_positive.value_offsets().as_ptr() - ); - assert_eq!( - packed_positive.value_offsets()[sliced_table_offset], - sliced_positive.value_offsets()[sliced_table_offset] - ); - assert_eq!( - packed_positive.value_offsets()[sliced_table_offset + 1], - sliced_positive.value_offsets()[sliced_table_offset + 1] - ); - assert_eq!( - packed_positive - .values() - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_ptr(), - sliced_positive - .values() - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_ptr() - ); - } + assert_eq!(packed_mixed_table.len(), 1); + assert_eq!(packed_hist_table.len(), 1); let decoded = unpack_raw_series_frame( &schema, FrameLabelMode::Key, @@ -3245,30 +3163,14 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { let decoded_time_values = decoded[0].time.as_buffer(); assert_eq!(decoded_time_values.as_ptr(), packed_times.values().as_ptr()); assert_eq!(decoded_time_values, source_times); - for (decoded_indices, source_indices) in [ - (decoded_mixed.histograms().indices(), source_indices), - (decoded_histogram.indices(), &histogram_indices), - ] { - assert_eq!( - decoded_indices.values().as_ptr(), - source_indices.values().as_ptr() - ); - assert_eq!( - decoded_indices.nulls().unwrap().buffer().as_ptr(), - source_indices.nulls().unwrap().buffer().as_ptr() - ); - assert_eq!( - decoded_indices.nulls().unwrap().offset(), - source_indices.nulls().unwrap().offset() - ); - } assert_eq!( - decoded_mixed.histograms().indices().values().as_ptr(), - indices_owner.values().as_ptr().wrapping_add(1) + decoded_mixed.histograms().indices(), + &UInt32Array::from(vec![Some(0), None]) ); assert_eq!( - decoded_histogram.indices().values().as_ptr(), - histogram_indices_owner.values().as_ptr().wrapping_add(1) + decoded_histogram.indices(), + &UInt32Array::from(vec![Some(0), Some(0)]), + "both positions name the one distinct row" ); assert_eq!(decoded_mixed.histograms().indices().value(0), 0); assert!(!decoded_mixed.histograms().indices().is_valid(1)); @@ -3281,84 +3183,16 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { float_owner.nulls().unwrap().buffer().as_ptr() ); assert_eq!(decoded_mixed.floats().nulls().unwrap().offset(), 1); - assert_eq!(decoded_histogram.indices().value(0), 1); - assert_eq!(decoded_histogram.indices().value(1), 1); - assert_eq!(decoded_histogram.batch().len(), 2); - let decoded_table = decoded_histogram.batch().as_arrow(); - let decoded_sum = decoded_table - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(decoded_sum.values().as_ptr(), sliced_sum.values().as_ptr()); - let decoded_mixed_table = decoded_mixed.histograms().batch().as_arrow(); - for decoded_table_view in [decoded_mixed_table, decoded_table] { - let primitive = decoded_table_view - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - primitive.values().as_ptr(), - source_sum.values().as_ptr().wrapping_add(1) - ); - let positive = decoded_table_view - .column(10) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(positive.offset(), source_positive.offset()); - assert_eq!( - positive.value_offsets()[sliced_table_offset], - source_positive.value_offsets()[1] - ); - assert_eq!( - positive.value_offsets().as_ptr(), - sliced_positive.value_offsets().as_ptr() - ); - assert_eq!( - positive.value_offsets()[sliced_table_offset + 1], - source_positive.value_offsets()[2] - ); - assert_eq!( - positive - .values() - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_ptr(), - source_positive - .values() - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_ptr() - ); - } - let decoded_positive_buckets = decoded_table - .column(10) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_any() - .downcast_ref::() - .unwrap(); - let source_positive_buckets = source_table - .column(10) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!( - decoded_positive_buckets.values().as_ptr(), - source_positive_buckets.values().as_ptr() + assert_eq!(decoded_histogram.indices().value(0), 0); + assert_eq!(decoded_histogram.indices().value(1), 0); + assert_eq!(decoded_histogram.batch().len(), 1); + assert_histogram_with_start( + decoded_mixed.histograms().batch().view(0).unwrap(), + -9.5, + Some(-123), ); + // Both positions name the one distinct row of the compact table. + assert_histogram_with_start(decoded_histogram.batch().view(0).unwrap(), 90.5, None); drop(packed_mixed_table); drop(packed_hist_table); @@ -3372,11 +3206,13 @@ fn sliced_canonical_mixed_histogram_round_trip_keeps_buffers_alive() { drop(chunk); drop(frames); - assert_eq!(decoded_table.len(), 2); - assert_histogram(decoded_mixed.histograms().batch().view(0).unwrap(), -9.5); - assert_unreferenced_histogram(decoded_mixed.histograms().batch().view(1).unwrap()); - assert_histogram(decoded_histogram.batch().view(0).unwrap(), -9.5); - assert_unreferenced_histogram(decoded_histogram.batch().view(1).unwrap()); + assert_eq!(decoded_histogram.batch().len(), 1); + assert_histogram_with_start( + decoded_mixed.histograms().batch().view(0).unwrap(), + -9.5, + Some(-123), + ); + assert_histogram_with_start(decoded_histogram.batch().view(0).unwrap(), 90.5, None); assert!(matches!( decoded[0].fields.get(0).unwrap().sample(0), SampleRef::Histogram(_) @@ -3884,6 +3720,160 @@ fn byte_budget_flushes_before_exceeding_and_retries_the_pending_row() { assert!(offset_buffer(&[i32::MAX as usize, 1]).is_err()); } +#[test] +fn histogram_tables_encode_only_referenced_rows_and_remap_indices() { + let schema = test_schema(&["histogram", "mixed", "all_null"]); + let table = histogram_batch(); + let indices = UInt32Array::from(vec![Some(1), None, Some(0), Some(1)]); + let mixed = Samples::try_from_mixed( + Float64Array::from(vec![None, Some(12.5), None, None]), + crate::series::samples::HistogramSamples::try_new(indices.clone(), table.clone()).unwrap(), + ) + .unwrap(); + let fields = SeriesFields::try_from_vec(vec![ + Samples::try_from_histogram(indices.clone(), table.clone()).unwrap(), + mixed, + Samples::try_from_histogram(UInt32Array::new_null(4), table.clone()).unwrap(), + ]) + .unwrap(); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x71_u8]), + None, + RawTimestamps::try_new(vec![1_i64, 2, 3, 4].into()).unwrap(), + fields, + true, + ) + .unwrap(); + let packed = pack_raw_series_chunk( + &schema, + FrameLabelMode::Key, + Some(&dense_description()), + &chunk, + ) + .unwrap(); + let (encoded_indices, _) = frame_histogram_lists(&packed[0], 0); + let encoded_indices = encoded_indices + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + // Source indices are [1, null, 0, 1]: row 1 is named first and then row 0, and the second + // reference to row 1 reuses the first encoded row. + assert_eq!( + encoded_indices, + UInt32Array::from(vec![Some(0), None, Some(1), Some(0)]) + ); + assert_eq!( + frame_histogram_table(&packed[0], 0).len(), + 2, + "one encoded row per distinct referenced row, not per reference" + ); + assert_eq!(frame_histogram_table(&packed[0], 1).len(), 2); + assert_eq!(frame_histogram_table(&packed[0], 2).len(), 0); + + let decoded = decode_batches(&schema, FrameLabelMode::Key, &packed); + let DenseSamplesRef::Histogram(histograms) = decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("histogram representation was normalized"); + }; + assert_eq!( + histograms.indices(), + &UInt32Array::from(vec![Some(0), None, Some(1), Some(0)]) + ); + assert_histogram_with_start(histograms.view_at(0).unwrap(), 90.5, None); + assert_histogram_with_start(histograms.view_at(2).unwrap(), -9.5, Some(-123)); + assert_histogram_with_start(histograms.view_at(3).unwrap(), 90.5, None); + + let DenseSamplesRef::Mixed(mixed) = decoded[0].fields.get(1).unwrap().dense().unwrap() else { + panic!("mixed representation was normalized"); + }; + assert_eq!(mixed.histograms().indices(), histograms.indices()); + for (index, expected) in [None, Some(12.5_f64), None, None].into_iter().enumerate() { + match ( + mixed.floats().is_valid(index), + mixed.floats().value(index), + expected, + ) { + (false, _, None) => {} + (true, actual, Some(expected)) => { + assert_eq!(actual.to_bits(), expected.to_bits(), "position {index}"); + } + _ => panic!("unexpected mixed float at position {index}"), + } + } + assert!(mixed.floats().is_null(0)); + assert!(mixed.histograms().indices().is_valid(0)); + assert!(mixed.floats().is_valid(1)); + assert!(!mixed.histograms().indices().is_valid(1)); + assert_histogram_with_start(mixed.histograms().view_at(0).unwrap(), 90.5, None); + assert_histogram_with_start(mixed.histograms().view_at(2).unwrap(), -9.5, Some(-123)); + assert_histogram_with_start(mixed.histograms().view_at(3).unwrap(), 90.5, None); + + let DenseSamplesRef::Histogram(all_null) = decoded[0].fields.get(2).unwrap().dense().unwrap() + else { + panic!("all-null histogram representation was normalized"); + }; + assert_eq!(all_null.indices().len(), 4); + assert!((0..4).all(|index| !all_null.indices().is_valid(index))); + assert_eq!(all_null.batch().len(), 0); +} + +#[test] +fn repeated_references_to_one_histogram_row_encode_it_once() { + let schema = test_schema(&["histogram"]); + let table = histogram_batch(); + // Every position names the same row: the encoded payload must grow with distinct named rows, + // not with positions. + let positions = 10_000; + let indices = UInt32Array::from(vec![Some(1); positions]); + let chunk = RawSeriesChunk::try_new( + &schema, + Buffer::from(vec![0x33_u8]), + None, + RawTimestamps::try_new((0..positions as i64).collect::>().into()).unwrap(), + SeriesFields::one(Samples::try_from_histogram(indices, table).unwrap()), + true, + ) + .unwrap(); + let packed = pack_raw_series_chunk( + &schema, + FrameLabelMode::Key, + Some(&dense_description()), + &chunk, + ) + .unwrap(); + let (encoded_indices, _) = frame_histogram_lists(&packed[0], 0); + let encoded = encoded_indices + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(encoded.len(), positions); + assert!( + (0..positions).all(|position| encoded.value(position) == 0), + "every position reuses the one encoded row" + ); + assert_eq!( + frame_histogram_table(&packed[0], 0).len(), + 1, + "one distinct referenced row, regardless of how often it is referenced" + ); + + let decoded = decode_batches(&schema, FrameLabelMode::Key, &packed); + let DenseSamplesRef::Histogram(histograms) = decoded[0].fields.get(0).unwrap().dense().unwrap() + else { + panic!("histogram representation was normalized"); + }; + assert_eq!(histograms.batch().len(), 1); + // Row 1 of the fixture is the one with sum 90.5 and no start timestamp. + for position in [0, positions / 2, positions - 1] { + assert_histogram_with_start(histograms.view_at(position).unwrap(), 90.5, None); + } +} + #[test] fn multirow_histogram_rows_keep_row_local_indices_and_tables() { let schema = test_schema(&["histogram"]); @@ -3932,16 +3922,16 @@ fn multirow_histogram_rows_keep_row_local_indices_and_tables() { assert_eq!( encoded_second.value(1), 0, - "row-local indices must not be remapped" + "indices must be remapped within their row's compact table" ); - // The second row references its own table through a nonzero slice of the concatenated child. - assert_eq!(tables.value_offsets(), [0, 2, 4]); + // Each row's sole referenced table entry is encoded independently. + assert_eq!(tables.value_offsets(), [0, 1, 2]); let concatenated = tables .values() .as_any() .downcast_ref::() .unwrap(); - assert_eq!(concatenated.len(), 4); + assert_eq!(concatenated.len(), 2); let decoded = decode_batches( &schema, @@ -3950,8 +3940,8 @@ fn multirow_histogram_rows_keep_row_local_indices_and_tables() { ); assert_row_local_histogram_rows(&decoded); - // A real Arrow IPC round trip keeps the schema, the whole tables, the null index entries and - // the bit-exact payloads. + // A real Arrow IPC round trip keeps the schema, compact row tables, null index entries and + // the bit-exact referenced payloads. let mut ipc = Vec::new(); { let mut writer = StreamWriter::try_new(&mut ipc, &packed[0].schema()).unwrap(); diff --git a/src/promql/src/extension_plan.rs b/src/promql/src/extension_plan.rs index 102dd918c9e..ba4cf65e0d5 100644 --- a/src/promql/src/extension_plan.rs +++ b/src/promql/src/extension_plan.rs @@ -26,6 +26,7 @@ mod raw_series_source; mod scalar_calculate; mod series_divide; mod series_gather; +mod series_repartition; #[cfg(test)] mod test_util; mod union_distinct_on; @@ -56,6 +57,7 @@ pub use range_manipulate::{RangeManipulate, RangeManipulateExec, RangeManipulate pub use raw_series_source::{RawSeriesSource, RawSeriesSourceExec}; pub use scalar_calculate::ScalarCalculate; pub use series_divide::{SeriesDivide, SeriesDivideExec, SeriesDivideStream}; +pub use series_repartition::SeriesRepartitionExec; pub use union_distinct_on::{UnionDistinctOn, UnionDistinctOnExec, UnionDistinctOnStream}; pub type Millisecond = ::Native; diff --git a/src/promql/src/extension_plan/series_repartition.rs b/src/promql/src/extension_plan/series_repartition.rs new file mode 100644 index 00000000000..af53ef796d9 --- /dev/null +++ b/src/promql/src/extension_plan/series_repartition.rs @@ -0,0 +1,1666 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Routes evaluated series chains by canonical label set across DataFusion input partitions. +//! +//! This node routes chains by canonical label-set hash, then merges exact metric-name/tag +//! identities at each owner after their close/EOF barriers. This is owner-local grouping, not a claim +//! of global canonical completeness. It operates on DataFusion partitions only; arbitrary `MatchKey` +//! and `GroupKey` routing is unsupported. +//! +//! Chains are transported intact to owners with their input-partition ID, followed by one charged +//! normal-EOF message per live output from each completed input. Owners group exact identities and +//! may publish an identity as soon as every input has contributed its final chain or normal EOF. The +//! frame is the child's own: evaluated chunks are decoded only to read routing labels, and outputs +//! publish the child's declared layout. +//! +//! # Routing +//! +//! The key is [`label_set`] of the chunk's metric, hashed whole: coarsening it (a subset of the +//! labels, a metric name alone) would send two label sets to one output for sharing a prefix, which +//! is not the relation a downstream merge groups by. One [`Routing`] is built per execution and +//! every chunk is reduced with it, so the destination of a label set cannot depend on which chunk of +//! it was routed first. +//! +//! The node declares [`Partitioning::UnknownPartitioning`], never [`Partitioning::Hash`]: the key is +//! decoded from the rows, and a hash declaration carries an Arrow expression DataFusion could use to +//! prove two plans agree on a partitioning. There is no such expression here, so the claim would be +//! one no optimizer could check and a hash exchange inserted on top of it could silently disagree +//! with. No ordering is declared either: the node reorders, and it says so by claiming nothing. +//! +//! # Chain boundaries +//! +//! A consumer of these frames requires a series' identity not to change and its positions to increase +//! strictly while open, with `is_last` closing it. An input task therefore assembles each child chain +//! before routing it as one atomic unit. This prevents chains from interleaving during transport; it +//! does not establish the output order or final series boundaries. For a full exact identity (metric +//! name and every tag name/value), the input contract is that each identity closes at most once per +//! input. Each owner therefore waits only for each distinct input's final chain or normal EOF before +//! merging and publishing that identity. +//! +//! What that costs is stated rather than hidden. Nothing of a chain is handed over before the chain +//! closes, so the first chunk of a series reaches its output later than a chunk-at-a-time design would +//! deliver it, and a long chain occupies memory for as long as it is being read. Those bytes are +//! charged to a [`MemoryConsumer`] as the chain grows, so a pool under pressure ends the read with an +//! error instead of letting a chain grow without bound: the node refuses, it does not spin and it does +//! not buffer past the pool's limit. +//! +//! Queues carry whole chains rather than individual chunks, so two sources' partial chains cannot +//! interleave on one shared output queue; each output needs only one queue. +//! +//! # Where the reads run +//! +//! One task per child partition, and no task above them: each input's task reads its own child, +//! assembles that input's chains and sends them to their owners itself. No task makes a routing +//! decision for another input, and nothing rotates over the inputs. +//! +//! That is not an optimisation of the routing: an evaluated partition is one stream whose computation +//! happens inside its `poll_next`, so a task polling every partition in turn runs the child's native +//! kernels for all of them itself, one chunk at a time, and the partition count the child declared +//! buys no parallelism at all. It is also what keeps routing local, which is what this node needs +//! when two of them are composed: +//! +//! - A partial chain never occupies a shared output queue, so no output waits for a continuation: a +//! chain is either fully assembled in the input task that read it, or not sent at all. +//! - A task never waits for another task's routing decision, and the tasks share no lock while +//! routing, so where one input's chain goes is not something another input learns or needs. +//! +//! Producer queues are unbounded: unread owners do not stall routing, but their retained chains +//! continue to consume pool-accounted memory. Owners publish each exact identity after its per-input +//! close/EOF barrier, with finality recomputed for that output series. +//! +//! # Memory, backpressure and end of stream +//! +//! Output queues are unbounded; each input task holds at most the one chain it is assembling. Every +//! queued chain is admitted through the task context's pool and carries its charge. The public packed +//! stream releases a completed cursor's charges after its batch is flushed, while retaining the +//! current cursor's charges until the cursor is exhausted. The pool controls +//! admission, but queue length is not a constant-byte bound. Owners retain each pending exact identity +//! until its per-input close/EOF barrier is satisfied, then +//! reserve merge workspace for that identity alone. This bounds merge workspace to one pending +//! identity, while its inputs remain charged until handoff. No task waits for memory admission. +//! +//! An output ends only when every input partition finished and its queue was drained; a chain that +//! closed locally is not the end of anything. Each input sends its own EOF to each live output before +//! recording normal completion. With one task per input the tasks count themselves down: the last +//! records the read outcome, then sender closure lets outputs detect missing EOF or report failure. If the child errors, the failing task records it before cancelling its peers; sender closure makes +//! every output report that recorded error instead of a clean end. +//! +//! A queue closing is not by itself that end: it also closes when the read is stopped - aborted, or +//! panicked - and a stopped read is missing exactly what it did not route. The read therefore records +//! how it ended, and an output whose queue has closed reads that record, so only a read that ended on +//! its own terms ends an output and anything else is an error saying the producer stopped +//! unexpectedly. +//! +//! An output stream owns the read, not the plan it was executed from: DataFusion's +//! `execute_stream_partitioned` takes the plan by value and keeps only the streams, so a read tied to +//! the plan would be stopped while its streams are still being consumed. The read therefore lives as +//! long as any of its outputs does, and stops only when there is nothing left it could deliver to: +//! every output was handed out and every one of them is gone. An output that ends - normally, or by +//! being dropped - while another partition was never executed stops nothing by itself, because that +//! partition can still claim its queue; the execution going away as a whole, which no later `execute` +//! can reach, is the other point at which everything is released. + +use std::collections::{BTreeMap, VecDeque}; +use std::fmt::{self, Formatter}; +use std::hash::{Hash, Hasher}; +use std::pin::Pin; +use std::sync::{Arc, Mutex, Weak}; +use std::task::{Context, Poll}; + +use common_error::ext::BoxedError; +use common_query::series_frame::{ + SendableEvaluatedSeriesStream, pack_evaluated_series, unpack_evaluated_series, +}; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::error::{DataFusionError, Result as DataFusionResult}; +use datafusion::execution::context::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; +use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{EmissionType, EvaluationType, SchedulingType}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, +}; +use datatypes::series::SeriesLabels; +use datatypes::series::frame::{FrameContent, SeriesFrameDeclaration}; +use datatypes::series::samples::{Series, SeriesSchema}; +use futures::{Stream, StreamExt}; +use tokio::sync::mpsc; +use twox_hash::XxHash64; + +use crate::grid_series::{label_set, merge_same_labelset}; +use crate::native::execute::merge_workspace_bytes; +use crate::native::execute::operand_chunk_bytes; +use datatypes::series::samples::DenseSamplesRef; + +/// The seed behind every routing decision. +/// +/// Fixed rather than drawn per execution, so which output a label set belongs on is a function of +/// the label set and the output count alone: an execution's own decisions have to agree with each +/// other, and a seeded router can be reasoned about - and tested - without one. +const ROUTING_SEED: u64 = 0x517c_c1b7_2722_0a95; + +/// The Repartition stage: one child's partitions, redistributed by canonical label set. +/// +/// The node holds one execution's shared state: its output partitions are separate `execute` calls +/// that must agree on one read of the child, so the first call starts the read and every call takes +/// its own output from it. +pub struct SeriesRepartitionExec { + input: Arc, + /// The frame the child declares, re-declared by every output batch. + series_schema: Arc, + /// The one hash configuration of this node's execution. + routing: Routing, + /// Started by the first `execute` call, since the output partitions share one read of the child. + shared: Mutex>>, + properties: Arc, +} + +impl SeriesRepartitionExec { + /// The name this node is shown under. + pub const fn name() -> &'static str { + "SeriesRepartitionExec" + } + + /// Creates the node over `input`, routing into `outputs` partitions. + /// + /// `outputs` is fixed here, by the planner, which knows the target partition count; the node + /// reads no session configuration of its own. A child that does not declare an evaluated frame + /// this node can re-declare is refused before the plan can run: the output schema is the child's, + /// and a frame whose declaration does not survive the round trip would be published under a + /// schema its batches do not match. + pub fn try_new(input: Arc, outputs: usize) -> DataFusionResult { + if outputs == 0 { + return Err(DataFusionError::Plan( + "SeriesRepartition requires at least one output partition".to_string(), + )); + } + if input.properties().boundedness.is_unbounded() { + return Err(DataFusionError::Plan( + "SeriesRepartition requires a finite input to finalize series identities".into(), + )); + } + let declaration = + SeriesFrameDeclaration::parse(input.schema().as_ref()).map_err(|error| { + DataFusionError::Plan(format!( + "SeriesRepartition input is not a series frame: {error}" + )) + })?; + if declaration.content() != FrameContent::Evaluated { + return Err(DataFusionError::Plan( + "SeriesRepartition routes evaluated series fragments, not raw samples".to_string(), + )); + } + let series_schema = Arc::new(declaration.series().clone()); + let declared = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + )? + .schema(); + if declared != input.schema() { + return Err(DataFusionError::Plan( + "SeriesRepartition input frame carries a schema this node cannot re-declare" + .to_string(), + )); + } + let properties = PlanProperties::new( + EquivalenceProperties::new(input.schema()), + Partitioning::UnknownPartitioning(outputs), + // Each exact identity becomes final independently at its per-input close/EOF barrier. + EmissionType::Incremental, + // The node ends when its child does, and spends its life waiting until then. + input.properties().boundedness, + ) + .with_evaluation_type(EvaluationType::Eager) + .with_scheduling_type(SchedulingType::Cooperative); + let properties = Arc::new(properties); + Ok(Self { + input, + series_schema, + routing: Routing::new(outputs), + shared: Mutex::new(None), + properties, + }) + } + + /// How many output partitions this node has. + fn outputs(&self) -> usize { + self.properties.partitioning.partition_count() + } +} + +impl fmt::Debug for SeriesRepartitionExec { + /// The node and what it routes: the child, the schema and the running execution are noise in a + /// plan's debug form. + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("SeriesRepartitionExec") + .field("outputs", &self.outputs()) + .finish_non_exhaustive() + } +} + +impl DisplayAs for SeriesRepartitionExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!( + f, + "SeriesRepartitionExec: outputs={}, input_partitions={}", + self.outputs(), + self.input.properties().partitioning.partition_count() + ) + } +} + +impl ExecutionPlan for SeriesRepartitionExec { + fn name(&self) -> &str { + Self::name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + /// Nothing is asked of the input's distribution, and nothing may be inserted above it. + /// + /// The child's own partitioning is what this node reads: it reads every partition as it is, and + /// an exchange inserted between the two would reorder chunks across the chain boundaries the node + /// relies on, which no requirement it states could survive. A hash requirement would be worse + /// still, since the routing key is decoded from the rows rather than computed by an expression. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(vec![Distribution::UnspecifiedDistribution]) + } + + /// Splitting the child further would not help: this node reads every input partition itself, so + /// a repartition below it only adds a hop that reorders what it must read in order. + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + // The routing key is decoded from the rows, so this node binds no expression of its own. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> DataFusionResult> { + if children.len() != 1 { + return Err(DataFusionError::Plan(format!( + "SeriesRepartitionExec requires one child, got {}", + children.len() + ))); + } + // Rebuilt rather than patched, so a child whose schema changed is refused here instead of + // routed under a frame declaration the new child does not make. + Ok(Arc::new(Self::try_new( + children.swap_remove(0), + self.outputs(), + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + if partition >= self.outputs() { + return Err(DataFusionError::Execution(format!( + "SeriesRepartition output partition {partition} is out of range" + ))); + } + let shared = { + let mut state = self + .shared + .lock() + .unwrap_or_else(|error| error.into_inner()); + if state.is_none() { + // The whole execution is started here, once: the output partitions are separate + // calls, and reading the child once for all of them is what keeps a chunk from being + // emitted twice or read by a scan that was never meant to run twice. + *state = Some(SharedExecution::start( + &self.input, + self.routing, + Arc::clone(&self.series_schema), + &context, + )?); + } + Arc::clone(state.as_ref().expect("the execution was just started")) + }; + let (receiver, lease) = SharedExecution::claim(&shared, partition).ok_or_else(|| { + DataFusionError::Execution(format!( + "SeriesRepartition output partition {partition} was executed twice" + )) + })?; + // Take the queue and its claim on the read together. + let inner = RepartitionStream::new( + receiver, + lease, + Arc::clone(&self.series_schema), + &context.runtime_env().memory_pool, + self.input.properties().partitioning.partition_count(), + ); + let handoff = Arc::clone(&inner.handoff); + let output = pack_evaluated_series(Box::pin(inner), Arc::clone(&self.series_schema))?; + let schema = output.schema(); + let output = Box::pin(PackedHandoffStream { + inner: Some(output), + schema, + handoff, + }); + Ok(datafusion::physical_plan::coop::make_cooperative(output)) + } +} + +/// One execution's state: the child read once, and the outputs it is being distributed into. +/// +/// The plan that started it holds this, and so does every output stream it hands out. The read it +/// owns lives while any of them does - which is what makes it survive the plan - and what stops it is +/// [`Running::claimed`] and [`Running::remaining`] both reaching zero, or the state itself going away. +struct SharedExecution { + /// The queues, the tasks filling them, and how many of the execution's outputs the read still has. + running: Mutex, + /// How the read ended, and the failure it owes its outputs if it failed. + ending: Arc>, +} + +/// What one execution's state holds that has to be stopped together. +struct Running { + /// One receiver per output partition, taken by the `execute` call of that partition. A receiver + /// that is never taken stays here with its pool-accounted queued chains. + receivers: Vec>>, + /// The tasks reading the child's partitions into those queues, and how many of them are left. + inputs: Arc, + /// How many output streams were handed out and are not yet released. + /// + /// Counted rather than derived from a reference count, because the plan holds this state too and + /// is not an output: what has to reach zero for the read to be over is the number of streams, not + /// the number of owners. + claimed: usize, + /// How many outputs were never handed out, and can still be. + /// + /// Kept apart from [`claimed`](Self::claimed) because the two say different things: an output that + /// is claimed is one somebody is reading, and one that is only `remaining` is one nobody has asked + /// for yet. A claimed count that reaches zero with outputs left here has no active consumers + /// yet, but those queues survive so a later `execute` can claim them. + remaining: usize, +} + +impl Running { + /// Stops the read: every queue nobody took is dropped and every input task is aborted. + /// + /// Idempotent - the queues are already gone on a second call, and aborting a finished or already + /// aborted task is a no-op. Dropping the queues of outputs that were never handed out is part of + /// it: what was queued for them is read-ahead nobody will take, and holding it until the process + /// ends would keep both the memory and the child's scans alive for nothing. + fn stop(&mut self) { + self.receivers.clear(); + // Aborting a read that has not recorded an outcome is what the outputs still open have to + // see: nothing was recorded, so the end they read is "the producer stopped unexpectedly". + self.inputs.stop(); + } +} + +impl SharedExecution { + /// Starts reading every partition of `input` into `routing`'s outputs on `context`'s pool. + /// + /// One task per child partition, started here: the child's computation happens inside its own + /// `poll_next`, so a partition read on a task of another partition would run there, and the + /// partition count the child declared would buy no parallelism at all. + fn start( + input: &Arc, + routing: Routing, + series_schema: Arc, + context: &Arc, + ) -> DataFusionResult> { + let partitions = input.properties().partitioning.partition_count(); + let mut streams = Vec::with_capacity(partitions); + for child_partition in 0..partitions { + let frames = input.execute(child_partition, Arc::clone(context))?; + // The expected schema is a check and never an input: the decoder reads the frame's own + // declaration, and a partition declaring another frame is refused here rather than + // routed under labels it does not carry. + streams.push(unpack_evaluated_series( + frames, + Some(Arc::clone(&series_schema)), + )?); + } + let mut senders = Vec::with_capacity(routing.outputs); + let mut receivers = Vec::with_capacity(routing.outputs); + for _ in 0..routing.outputs { + let (sender, receiver) = mpsc::unbounded_channel(); + senders.push(sender); + receivers.push(Some(receiver)); + } + let ending = Arc::new(Mutex::new(Ending::default())); + let shared = Arc::new_cyclic(|weak| SharedExecution { + running: Mutex::new(Running { + receivers, + inputs: Inputs::new(Arc::clone(&ending), partitions, weak.clone()), + claimed: 0, + remaining: routing.outputs, + }), + ending: Arc::clone(&ending), + }); + let inputs = { + let running = shared + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + Arc::clone(&running.inputs) + }; + // A child with no partitions has nothing to read, so the read finished before it started: the + // count that the last input task would bring to zero is already there. + if partitions == 0 { + inputs.record(Ending::Finished); + } + let pool = Arc::clone(&context.runtime_env().memory_pool); + for (partition, stream) in streams.into_iter().enumerate() { + // Each task is handed its own copy of the senders rather than a lock over one copy: a + // send is then a send of its own, and the queues close when the last task returns instead + // of when some shared lock says so. + let report = ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition, + reported: false, + }; + let task = common_runtime::spawn_query(route_partition( + stream, + partition, + routing, + senders.clone(), + // What this partition reads ahead is charged to its own consumer, so the read-ahead + // of one partition is attributable to the task that did it. + MemoryConsumer::new("SeriesRepartitionInput").register(&pool), + Arc::clone(&inputs), + report, + )); + inputs.attach(partition, task); + } + // This copy is dropped before any output exists, so it is no part of when the queues close. + drop(senders); + Ok(shared) + } + + /// The queue of `partition` and the lease on this read that reading it takes. + /// + /// The queue and the claim are taken under one lock, and that is what keeps the counts honest: an + /// output stops being one nobody asked for and becomes one the read is running for in the same + /// step, so a release on another task cannot see "nothing left" in between and stop a read that + /// is being handed out. `None` means that output was claimed before - its queue is gone and this + /// call does not count as one the read is running for. + fn claim( + shared: &Arc, + partition: usize, + ) -> Option<(mpsc::UnboundedReceiver, OutputLease)> { + let mut running = shared + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + let receiver = running + .receivers + .get_mut(partition) + .and_then(Option::take)?; + running.remaining -= 1; + running.claimed += 1; + Some(( + receiver, + OutputLease { + shared: Arc::clone(shared), + }, + )) + } + + /// Gives up one claim, and stops the read once there is nothing left it could deliver to. + /// + /// An output that is released is one that finished being read; outputs nobody has claimed yet can + /// still be claimed later. Stop only when no output is being read and none remains to hand out. + fn release_output(&self) { + let mut running = self + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + running.claimed -= 1; + if running.claimed == 0 && running.remaining == 0 { + running.stop(); + } + } + + /// How the read ended, and the failure it owes its outputs if it failed. + /// + /// Both are read under the one lock that records them, so an output never sees a failure without + /// the error that belongs to it. + fn ending(&self) -> Ending { + self.ending + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + + /// Stops the read and releases everything it holds. + /// + /// Called when the execution itself goes away - the plan and every output are gone, so no later + /// `execute` can claim an output that was never handed out - and by + /// [`release_output`](Self::release_output) when the last claim is given up on an execution that + /// has nothing left to hand out. Idempotent, because both can happen and because the state's own + /// drop runs it once more. + fn stop(&self) { + let mut running = self + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + running.stop(); + } +} + +impl Drop for SharedExecution { + fn drop(&mut self) { + // Nothing holds the read any more: every output is gone and so is the plan. No input task + // may outlive both of them. + self.stop(); + } +} + +/// The input tasks of one read: what is left of them, and what the last of them records. +/// +/// One task reads each child partition. The tasks count themselves down; the last records completion +/// before returning and dropping its senders, so outputs read the recorded outcome when queues close. +struct Inputs { + /// The tasks themselves, and the count of those that have not stopped. + state: Mutex, + /// How the read ended, and the failure it owes its outputs if it failed. + ending: Arc>, + shared: Weak, +} + +/// What [`Inputs`] guards. +struct InputState { + /// The task of each input partition, by partition, `None` once it stopped. + /// + /// Kept rather than dropped so that the read can be stopped as a whole, and so that the one task + /// that records a failure can abort the rest: a partition whose peer already failed has + /// nothing left to route that an output would still take. + tasks: Vec>>, + /// How many of those tasks have not stopped yet. + /// + /// The last task to stop brings this to zero and records that the read finished. Cancelled peers + /// are aborted and their report guards count them out; a task attached after cancellation is + /// aborted by [`Inputs::attach`] and its guard does the same. + live: usize, + /// Whether the read has been cancelled: a failure was elected, or [`Inputs::stop`] was called. + /// + /// Recorded under the same lock that guards [`tasks`](Self::tasks), so a task attached after the + /// cancellation is aborted by [`Inputs::attach`] rather than stored. The sweep that cancels the + /// read can only abort the handles that exist when it runs, and a task spawned afterwards would + /// otherwise continue polling a child for a read that is over while holding a sender to every + /// output, keeping its queues open. + cancelled: bool, +} + +impl Inputs { + /// The bookkeeping of a read over `partitions` input partitions. + fn new( + ending: Arc>, + partitions: usize, + shared: Weak, + ) -> Arc { + Arc::new(Self { + state: Mutex::new(InputState { + tasks: (0..partitions).map(|_| None).collect(), + live: partitions, + cancelled: false, + }), + ending, + shared, + }) + } + + /// Gives `task` the slot of `partition`, so that the read can still abort it. + /// + /// Stored even if the task has already stopped - a partition that answers at once can return + /// before its handle is stored - because the slot is only consulted to stop what is still running, + /// and aborting a finished task is a no-op. + /// + /// A task attached after cancellation is aborted rather than stored, so it cannot keep output + /// queues open after the read is over. Its report guard counts it out if it is polled. + fn attach(&self, partition: usize, task: common_runtime::JoinHandle<()>) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if state.cancelled { + state.tasks[partition] = None; + task.abort(); + return; + } + state.tasks[partition] = Some(task); + } + + /// Records that `partition` ended, and that the read finished if it was the last one running. + /// + /// A stopped or cancelled read cannot be reported as finished. + fn ended(&self, partition: usize) { + if self.count_out(partition) { + self.record(Ending::Finished); + } + } + + /// Records that `partition`'s task stopped without reporting an outcome, and stops the read with + /// it. + /// + /// This runs where no arm of [`route_partition`]'s match can: in a `Drop`. A panic inside + /// [`read_chains`] unwinds the task and drops the future that was polling it, and an abort drops + /// that same future where it stands, so both land here - which is the point, because neither + /// records anything otherwise. What that costs when nothing stops it is the whole read: the peers + /// go on reading a child that may never answer, their sender clones stay alive, no queue closes, + /// and every output waits for an ending that is not coming. So a task that died is a failure of + /// the read, and it is recorded before peer tasks are aborted. + /// + /// A read that is already cancelled - stopped from outside or by a failure - only counts this task out. The outcome that cancelled it is the one + /// the outputs report, and the aborts it performs are what killed this task: a task dying there is + /// a consequence of that outcome, not a second one for the outputs to choose between. + /// + /// Nothing is sent to the outputs: a `Drop` cannot send asynchronously, and closing the senders + /// lets every output observe the recorded failure. + /// Delivery here is structural: a queue closes when the last of its senders is dropped, every input + /// task holds a clone of every sender, and the peers cancelled here drop theirs as their futures + /// are dropped - so the queues close with nobody sending anything. What an output whose queue has + /// closed reads is [`RepartitionStream::poll_next`] at `Ready(None)`, which consults + /// [`OutputLease::ending`] before calling the read over, and that turns the [`Ending::Failed`] + /// recorded here into the error built from its terminal error. Every output therefore reports this + /// failure out of the read's own record after all input senders close. + fn task_died(&self, partition: usize) { + if self.is_cancelled() { + // The read has an outcome already, or another task's failure is the one being handed + // over. Counting out only keeps the bookkeeping honest: the read is cancelled, so this + // cannot make it look finished. + self.count_out(partition); + return; + } + // Record the error before cancellation drops this task's senders. + self.fail(Arc::new(DataFusionError::Execution(format!( + "series repartition input partition {partition} stopped without reporting an outcome" + )))); + self.cancel(partition); + self.count_out(partition); + } + + /// Records the first failure before cancelling peers. Outputs read it after sender closure. + fn failed(&self, partition: usize, error: &Arc) { + self.fail(Arc::clone(error)); + self.cancel(partition); + self.count_out(partition); + } + + fn is_cancelled(&self) -> bool { + self.state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .cancelled + } + + /// Marks the read cancelled under the task lock, then aborts its peer tasks. + /// + /// Cancellation keeps peers from recording a successful completion after this read has failed. + fn cancel(&self, partition: usize) { + { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.cancelled = true; + } + self.abort_others(partition); + } + + /// Counts `partition` off the read, and says whether it was the last task still running. + fn count_out(&self, partition: usize) -> bool { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + // Taken and dropped: this task is the one stopping, and holding its own handle would leave the + // read waiting for a task that has already finished. + state.tasks[partition] = None; + state.live -= 1; + // A cancelled read never records that it finished: a failure is the outcome its outputs + // report, and a stop is not an end. + state.live == 0 && !state.cancelled + } + + /// Aborts every input task but `partition`'s own. + /// + /// Stops peer tasks after a failure or an unexpected task death. `partition`'s own task is spared + /// so it can finish its bookkeeping, or because it has already died. + fn abort_others(&self, partition: usize) { + let state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + for (index, task) in state.tasks.iter().enumerate() { + // Aborting a task that already stopped is a no-op, so the finished ones cost nothing + // here. + if index != partition + && let Some(task) = task + { + task.abort(); + } + } + } + + /// Stops every input task that has not stopped yet and marks the read cancelled, so a task + /// attached after this sweep is aborted by [`attach`](Self::attach) rather than keeping queues + /// open after the read is over. + fn stop(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.cancelled = true; + for task in state.tasks.iter_mut() { + if let Some(task) = task.take() { + task.abort(); + } + } + } + + /// Records `outcome` as how the read ended, unless an ending is already recorded. + /// + /// A normal ending is recorded once. A failure is recorded by [`fail`](Self::fail), which stores + /// its error in the same write. + fn record(&self, outcome: Ending) { + debug_assert!( + !matches!(outcome, Ending::Failed(_)), + "a failure is recorded by `fail`, with its error" + ); + let mut ending = self + .ending + .lock() + .unwrap_or_else(|error| error.into_inner()); + if matches!(*ending, Ending::Pending) { + *ending = outcome; + } + } + + /// Records that the read failed with `error`, keeping the error where its outputs can read it. + /// + /// The first failure and its error are recorded together so outputs all report the same error. + fn fail(&self, error: Arc) { + let mut ending = self + .ending + .lock() + .unwrap_or_else(|error| error.into_inner()); + let first_failure = !matches!(*ending, Ending::Failed(_)); + if first_failure { + *ending = Ending::Failed(error); + } + drop(ending); + if first_failure && let Some(shared) = self.shared.upgrade() { + let mut running = shared + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + for receiver in running.receivers.iter_mut().flatten() { + receiver.close(); + while receiver.try_recv().is_ok() {} + } + } + } + + /// How many input tasks have not stopped, which is what a test sees the read stop by. + #[cfg(test)] + fn live(&self) -> usize { + self.state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .live + } + + /// Aborts the task of one input partition, for a test that has to make exactly one producer stop. + /// + /// [`stop`](Self::stop) cannot state that: it stops every task and records the read as stopped, + /// so what an output reports afterwards is the stop rather than one producer dying. + #[cfg(test)] + fn abort(&self, partition: usize) { + let state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(task) = &state.tasks[partition] { + task.abort(); + } + } +} + +/// One input partition's task: reads its own child, assembles chains, and sends each complete chain to +/// its canonical owner. Output queues are unbounded, so sending does not wait for an owner to poll; +/// queued chains remain charged to the pool until consumed or dropped. Each task records its outcome +/// before returning and dropping its senders, so output streams observe the final read outcome. +async fn route_partition( + mut stream: SendableEvaluatedSeriesStream, + partition: usize, + routing: Routing, + senders: Vec>, + read_ahead: MemoryReservation, + inputs: Arc, + mut report: ReportIfUnreported, +) { + let stop = read_chains(&mut stream, partition, &routing, &senders, &read_ahead).await; + if matches!(&stop, Stop::Ended) { + for sender in &senders { + if sender.is_closed() { + continue; + } + let metadata = read_ahead.new_empty(); + if let Err(error) = metadata.try_grow(std::mem::size_of::()) { + inputs.failed(partition, &Arc::new(error)); + report.reported = true; + return; + } + if let Err(error) = sender.send(Queued::Eof { + input: partition, + _metadata: metadata, + }) { + drop(error.0); + } + } + } + match &stop { + Stop::Ended => inputs.ended(partition), + Stop::Failed(error) => inputs.failed(partition, error), + } + // This task said what became of its partition, so the guard has nothing left to report. + report.reported = true; + drop(report); +} + +/// Reports an input task that stopped without recording an outcome, on the paths that would +/// otherwise report nothing. +/// +/// [`route_partition`] records ordinary input completion and errors; this guard records when a task +/// stops without returning. A panic unwinds the task and takes the future polling it with it, and the task +/// being aborted drops that future where it stands; neither reaches an arm of that match, so nothing +/// is recorded, the count is untouched and the peers go on reading. A peer parked on a child that +/// never answers then holds a sender to every output for good, no queue ever closes, and every output +/// waits for an ending that is not coming - the `Pending` fallback is no help there, since it is +/// only read once a queue has closed. +/// +/// A guard is what covers both without [`std::panic::catch_unwind`], which covers neither an abort nor +/// a panic the runtime catches elsewhere: unwinding and cancellation both drop this, so both report. +struct ReportIfUnreported { + /// The bookkeeping of the read this task belongs to. + inputs: Arc, + /// The input partition whose task holds this guard. + partition: usize, + /// Whether that task recorded its own outcome, which is what makes this a no-op. + reported: bool, +} + +impl Drop for ReportIfUnreported { + fn drop(&mut self) { + // Also runs on the way out of a normal return, which is what the flag is for: a task that + // reached its own ending reported it there and has nothing to add here. + if !self.reported { + self.inputs.task_died(self.partition); + } + } +} + +/// Why one input task stopped reading its partition. +enum Stop { + /// The partition ended: every chunk it had was read and handed to the output it belongs on. + Ended, + /// The partition failed, or a chain could not be buffered: every output has to learn that. + Failed(Arc), +} + +/// Accounts for the portion of a decoded histogram table retained by one series slice. +pub(crate) fn histogram_table_slice_bytes( + batch: &datatypes::series::histogram::HistogramBatch, +) -> DataFusionResult { + use datafusion::arrow::array::{Array, ListArray}; + + let array = batch.as_arrow(); + let mut bytes = std::mem::size_of_val(batch); + if array.nulls().is_some() { + bytes = bytes.saturating_add(array.len().div_ceil(8)); + } + for (index, column) in array.columns().iter().enumerate() { + if (7..=13).contains(&index) { + let list = column.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Execution(format!( + "histogram list column {index} has invalid type" + )) + })?; + let offsets = list.value_offsets(); + let first = *offsets + .first() + .ok_or_else(|| DataFusionError::Execution("empty histogram offsets".into()))? + as usize; + let last = *offsets + .last() + .ok_or_else(|| DataFusionError::Execution("empty histogram offsets".into()))? + as usize; + let values = list.values().slice(first, last.saturating_sub(first)); + bytes = + bytes.saturating_add((list.len() + 1).saturating_mul(std::mem::size_of::())); + bytes = bytes.saturating_add(values.to_data().get_slice_memory_size()?); + if list.nulls().is_some() { + bytes = bytes.saturating_add(list.len().div_ceil(8)); + } + } else { + bytes = bytes.saturating_add(column.to_data().get_slice_memory_size()?); + } + } + Ok(bytes) +} + +/// Reads one input partition to its end, one whole chain at a time, routing each chain as it closes. +/// +/// A chain is assembled here and handed over atomically after it closes; its growing memory is charged +/// to the pool. Output queues are unbounded, so routing does not wait for output polling. Each queued +/// chain retains its pool charge until the owner consumes it or the queue is dropped. This function +/// otherwise waits only on its own child. +async fn read_chains( + stream: &mut SendableEvaluatedSeriesStream, + partition: usize, + routing: &Routing, + senders: &[mpsc::UnboundedSender], + read_ahead: &MemoryReservation, +) -> Stop { + loop { + // One whole chain, assembled here: the chunk that closes the series ends it, and so does a + // stream that ends between two chains. + let mut chain: Chain = Vec::new(); + let mut chain_output: Option = None; + let mut discarding = false; + let mut ended = false; + loop { + // Input streams may return a long run of ready chunks without yielding to Tokio. Consume + // the normal task budget per chunk so sibling producers can run and report failures. + tokio::task::coop::consume_budget().await; + let next = match (chain_output, discarding) { + (Some(output), false) if !chain.is_empty() => { + tokio::select! { + biased; + _ = senders[output].closed() => { + chain.clear(); + discarding = true; + continue; + } + next = stream.next() => next, + } + } + _ => stream.next().await, + }; + match next { + Some(Ok((series, is_last))) => { + let output = match chain_output { + Some(output) => output, + None => { + let output = routing.output(&series.metric); + chain_output = Some(output); + output + } + }; + // A dropped output is no longer a destination. Still read and validate its + // entire chain, but do not retain or reserve its chunks. + if senders[output].is_closed() { + chain.clear(); + discarding = true; + if is_last { + break; + } + continue; + } + // Charged as the chain grows, so a chain that does not fit the pool is refused + // here rather than after it has been held. + if discarding { + if is_last { + break; + } + continue; + } + let charge = read_ahead.new_empty(); + let mut tables = 0usize; + for segment in series.fields.iter().flat_map(|field| field.segments()) { + let batch = match segment.dense() { + Some(DenseSamplesRef::Histogram(histograms)) => { + Some(histograms.batch()) + } + Some(DenseSamplesRef::Mixed(mixed)) => Some(mixed.histograms().batch()), + _ => None, + }; + if let Some(batch) = batch { + match histogram_table_slice_bytes(batch) { + Ok(bytes) => tables = tables.saturating_add(bytes), + Err(error) => return Stop::Failed(Arc::new(error)), + } + } + } + let bytes = operand_chunk_bytes(&series).saturating_add(tables); + if let Err(error) = charge.try_grow(bytes) { + if senders[output].is_closed() { + chain.clear(); + discarding = true; + if is_last { + break; + } + continue; + } + return Stop::Failed(Arc::new(error)); + } + chain.push((series, is_last, charge)); + if is_last { + break; + } + } + Some(Err(error)) => { + // The error is the last thing this partition says: it is handed to every output + // rather than ending this one task quietly. + return Stop::Failed(Arc::new(DataFusionError::External(Box::new(error)))); + } + None => { + ended = true; + break; + } + } + } + if ended { + // A partition that ends is a partition with nothing left to route - unless it ended + // inside a chain, which no later chunk can close. The decoders refuse that, so this is a + // producer that bypassed them, and its output would hold a series no chunk can close. + return if chain_output.is_none() { + Stop::Ended + } else { + Stop::Failed(Arc::new(DataFusionError::Execution( + "series repartition input ended with its series open".to_string(), + ))) + }; + } + // One hash per chain, over the labels of its first chunk: the chunks of a chain are one + // series, so the first of them names the output for all of them. + let Some(output) = chain_output else { + return Stop::Failed(Arc::new(DataFusionError::Internal( + "series repartition completed a chain without a routing destination".into(), + ))); + }; + if !discarding && { + let metadata = read_ahead.new_empty(); + if let Err(error) = metadata.try_grow(std::mem::size_of::()) { + return Stop::Failed(Arc::new(error)); + } + senders[output] + .send(Queued::Chain { + input: partition, + chain, + _metadata: metadata, + }) + .is_err() + } { + // A closed output is no longer a destination. Drop this chain and keep consuming. + continue; + } + } +} + +fn exact_identity(series: &Series) -> ExactIdentity { + ( + series.metric.name().map(Arc::::from), + series + .metric + .tags() + .iter() + .map(|tag| (Arc::clone(&tag.name), Arc::clone(&tag.value))) + .collect(), + ) +} + +/// Merges exact metric-name and complete-tag identities, preserving first-seen group order. +/// +/// The Fused series stack calls this helper for evaluated frame finalization; this branch does not +/// yet contain that caller, so the shared helper is not dead code in the integrated feature. +#[allow(dead_code)] +pub(crate) fn merge_exact_series( + schema: &Arc, + series: Vec, +) -> DataFusionResult> { + let mut indices = BTreeMap::new(); + let mut groups: Vec> = Vec::new(); + for series in series { + let key = ( + series.metric.name().map(Arc::::from), + series + .metric + .tags() + .iter() + .map(|tag| (Arc::clone(&tag.name), Arc::clone(&tag.value))) + .collect::>(), + ); + let group = match indices.get(&key) { + Some(index) => *index, + None => { + let index = groups.len(); + indices.insert(key, index); + groups.push(Vec::new()); + index + } + }; + groups[group].push(series); + } + let mut output = Vec::with_capacity(groups.len()); + for group in groups { + output.extend( + merge_same_labelset(schema, group) + .map_err(|error| DataFusionError::External(Box::new(error)))?, + ); + } + Ok(output) +} + +/// What one output stream holds so that the read belongs to the streams rather than to the plan. +/// +/// DataFusion's `execute_stream_partitioned` takes the plan by value and keeps only the streams, so +/// the state that started the read is dropped before the first chunk is read: a read stopped there +/// would end every output early, and an end is not something a consumer can tell from a complete +/// answer. The lease is what keeps it running instead; and because only streams hold one, a lease +/// going away is the point at which that one output can no longer take a chunk - which +/// [`SharedExecution::release_output`] turns into the end of the read only when no output is left that +/// could. +struct OutputLease { + shared: Arc, +} + +impl OutputLease { + /// Records an owner-side failure and cancels producer tasks without replacing an earlier error. + fn fail(&self, error: Arc) { + let inputs = { + let running = self + .shared + .running + .lock() + .unwrap_or_else(|error| error.into_inner()); + Arc::clone(&running.inputs) + }; + inputs.fail(Arc::clone(&error)); + inputs.stop(); + } + + /// Reports successful EOF only when every producer ended successfully. + fn ending(&self) -> Option> { + match self.shared.ending() { + Ending::Finished => None, + Ending::Failed(error) => Some(Err(boxed(DataFusionError::from(&error)))), + Ending::Pending => Some(Err(boxed(DataFusionError::Execution( + "series repartition producer stopped unexpectedly before the read finished" + .to_string(), + )))), + } + } +} + +impl Drop for OutputLease { + fn drop(&mut self) { + // An output that is done is not by itself the end of the read: the outputs that were never + // claimed are still to come, and what is queued for them is what they will read. Giving up + // the claim therefore only says that this one output will take nothing more - whether the + // read can go on is decided from every output, claimed and unclaimed alike. + self.shared.release_output(); + } +} + +/// How a read ended, as an output whose queue has closed has to report it. +#[derive(Clone, Debug, Default)] +enum Ending { + #[default] + Pending, + Finished, + Failed(Arc), +} + +/// One whole chain of one series: its chunks from the one that opened it to the one that closed it, +/// each with the bytes it is charged. +/// +/// The transport unit is a chain so queued source contributions remain intact. Owners collect all +/// chains before locally merging exact identities and producing their final output. +type Chain = Vec<(Series, bool, MemoryReservation)>; + +/// What one output queue carries: a whole chain and its memory charge. +enum Queued { + Chain { + input: usize, + chain: Chain, + _metadata: MemoryReservation, + }, + Eof { + input: usize, + _metadata: MemoryReservation, + }, +} + +type ExactIdentity = (Option>, Vec<(Arc, Arc)>); + +struct PendingGroup { + fragments: Vec, + charges: Vec, + received: Vec, + metadata: MemoryReservation, +} + +fn group_ready(group: &PendingGroup, eof: &[bool]) -> bool { + group + .received + .iter() + .zip(eof) + .all(|(received, eof)| *received || *eof) +} + +struct ReadySeries { + series: Series, + charges: ReadyCharges, +} + +struct ReadyCharges { + _charges: Vec, + _metadata: MemoryReservation, + _workspace: Option, +} + +/// One owner processes identities as soon as every input has closed or contributed its final chain. +struct RepartitionStream { + receiver: mpsc::UnboundedReceiver, + lease: OutputLease, + schema: Arc, + pool: Arc, + groups: BTreeMap, + eof: Vec, + eof_metadata: Option, + inputs: usize, + output: VecDeque, + handoff: Arc>, + /// Whether a new pass over the pending identities is owed after an input EOF. + /// + /// A chain can only complete the identity it carries, so the chain arm checks that group + /// directly. An EOF can complete any waiting identity, so it opens a pass instead; the pass + /// resumes after [`resume_after`](Self::resume_after) so the same prefix is not rescanned for + /// every identity it releases. + rescan: bool, + /// The last identity a pass examined, exclusive. Retains at most one identity descriptor. + resume_after: Option, + reported: bool, +} + +/// Charges move out of the current cursor only after it is exhausted; the public batch stream +/// releases completed cursors after the packer has flushed the frames they back. +#[derive(Default)] +struct Handoff { + current: Option, + completed: Vec, +} + +struct PackedHandoffStream { + inner: Option, + schema: SchemaRef, + handoff: Arc>, +} + +impl RecordBatchStream for PackedHandoffStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl Stream for PackedHandoffStream { + type Item = DataFusionResult; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let result = match self.inner.as_mut() { + Some(inner) => inner.as_mut().poll_next(cx), + None => return Poll::Ready(None), + }; + match &result { + Poll::Ready(Some(Ok(_))) | Poll::Pending => { + let mut handoff = self.handoff.lock().unwrap_or_else(|e| e.into_inner()); + // A pending packer has exhausted its cursor and its frame builder is empty, so + // completed charges no longer retain any payload; `current` may still own the + // whole merged series behind a live cursor. + handoff.completed.clear(); + } + Poll::Ready(None) | Poll::Ready(Some(Err(_))) => { + drop(self.inner.take()); + let mut handoff = self.handoff.lock().unwrap_or_else(|e| e.into_inner()); + handoff.current = None; + handoff.completed.clear(); + } + } + result + } +} + +impl RepartitionStream { + fn new( + receiver: mpsc::UnboundedReceiver, + lease: OutputLease, + schema: Arc, + pool: &Arc, + inputs: usize, + ) -> Self { + Self { + receiver, + lease, + schema, + pool: Arc::clone(pool), + groups: BTreeMap::new(), + eof: Vec::new(), + eof_metadata: None, + inputs, + output: VecDeque::new(), + handoff: Arc::new(Mutex::new(Handoff::default())), + rescan: false, + resume_after: None, + reported: false, + } + } + + fn finish_group(&mut self, key: &ExactIdentity) -> DataFusionResult<()> { + let Some(group) = self.groups.remove(key) else { + return Ok(()); + }; + let workspace = if group.fragments.len() > 1 { + let rows = group.fragments.iter().fold(0usize, |rows, series| { + rows.saturating_add(series.time.len()) + }); + let reservation = MemoryConsumer::new("SeriesRepartitionOwner").register(&self.pool); + reservation.try_grow(merge_workspace_bytes( + rows, + group.fragments.len(), + self.schema.fields.len(), + ))?; + Some(reservation) + } else { + None + }; + let merged = merge_same_labelset(&self.schema, group.fragments) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + if merged.len() != 1 { + return Err(DataFusionError::Internal( + "exact identity produced multiple merged series".into(), + )); + } + let Some(series) = merged.into_iter().next() else { + return Err(DataFusionError::Internal( + "exact identity produced no merged series".into(), + )); + }; + self.output.push_back(ReadySeries { + series, + charges: ReadyCharges { + _charges: group.charges, + _metadata: group.metadata, + _workspace: workspace, + }, + }); + Ok(()) + } +} + +impl Stream for RepartitionStream { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + { + let mut handoff = this.handoff.lock().unwrap_or_else(|e| e.into_inner()); + if let Some(current) = handoff.current.take() { + handoff.completed.push(current); + } + } + let mut drained = 0; + loop { + if let Ending::Failed(error) = this.lease.shared.ending() { + if this.reported { + return Poll::Ready(None); + } + this.reported = true; + this.receiver.close(); + while this.receiver.try_recv().is_ok() {} + this.groups.clear(); + this.output.clear(); + this.eof.clear(); + this.eof_metadata = None; + return Poll::Ready(Some(Err(boxed(DataFusionError::from(&error))))); + } + if this.eof_metadata.is_none() { + let metadata = MemoryConsumer::new("SeriesRepartitionEof").register(&this.pool); + let bytes = this.inputs.saturating_mul(std::mem::size_of::()); + if let Err(error) = metadata.try_grow(bytes) { + this.lease.fail(Arc::new(error)); + continue; + } + this.eof = vec![false; this.inputs]; + this.eof_metadata = Some(metadata); + } + if let Ending::Failed(error) = this.lease.shared.ending() { + if this.reported { + return Poll::Ready(None); + } + this.reported = true; + this.receiver.close(); + while this.receiver.try_recv().is_ok() {} + this.groups.clear(); + this.output.clear(); + this.eof.clear(); + this.eof_metadata = None; + return Poll::Ready(Some(Err(boxed(DataFusionError::from(&error))))); + } + if this.rescan { + let next = match &this.resume_after { + Some(after) => this + .groups + .range(( + std::ops::Bound::Excluded(after.clone()), + std::ops::Bound::Unbounded, + )) + .find_map(|(key, group)| { + group_ready(group, &this.eof).then(|| key.clone()) + }), + None => this.groups.iter().find_map(|(key, group)| { + group_ready(group, &this.eof).then(|| key.clone()) + }), + }; + match next { + Some(key) => { + this.resume_after = Some(key.clone()); + if let Err(error) = this.finish_group(&key) { + this.rescan = false; + this.resume_after = None; + this.lease.fail(Arc::new(error)); + continue; + } + } + None => { + // Every pending identity after the resume point was examined and none is + // ready; only a new chain or EOF can change that. + this.rescan = false; + this.resume_after = None; + } + } + } + if let Some(ready) = this.output.pop_front() { + this.handoff + .lock() + .unwrap_or_else(|e| e.into_inner()) + .current = Some(ready.charges); + return Poll::Ready(Some(Ok((ready.series, true)))); + } + if this.reported { + return Poll::Ready(None); + } + if drained == 128 { + cx.waker().wake_by_ref(); + return Poll::Pending; + } + match this.receiver.poll_recv(cx) { + Poll::Ready(Some(Queued::Chain { + input, + chain, + _metadata, + })) => { + drained += 1; + let Some(first) = chain.first() else { + continue; + }; + let key = exact_identity(&first.0); + if chain + .iter() + .any(|(series, _, _)| exact_identity(series) != key) + { + this.lease.fail(Arc::new(DataFusionError::Execution( + "one series chain contains multiple exact identities".into(), + ))); + continue; + } + if !this.groups.contains_key(&key) { + let metadata = + MemoryConsumer::new("SeriesRepartitionIdentity").register(&this.pool); + let bytes = std::mem::size_of::() + .saturating_add(std::mem::size_of::()) + .saturating_add(key.0.as_ref().map_or(0, |name| name.len())) + .saturating_add(key.1.capacity().saturating_mul(std::mem::size_of::<( + Arc, + Arc, + )>( + ))) + .saturating_add( + key.1 + .iter() + .map(|(n, v)| n.len().saturating_add(v.len())) + .sum::(), + ) + .saturating_add(this.eof.len().saturating_mul(2)); + if let Err(error) = metadata.try_grow(bytes) { + this.lease.fail(Arc::new(error)); + continue; + } + this.groups.insert( + key.clone(), + PendingGroup { + fragments: Vec::new(), + charges: Vec::new(), + received: vec![false; this.eof.len()], + metadata, + }, + ); + } + let Some(group) = this.groups.get_mut(&key) else { + this.lease.fail(Arc::new(DataFusionError::Internal( + "series identity group was not retained".into(), + ))); + continue; + }; + if input >= group.received.len() || group.received[input] { + this.lease.fail(Arc::new(DataFusionError::Execution( + "input repeated an exact series identity".into(), + ))); + continue; + } + group.received[input] = true; + for (series, _, charge) in chain { + group.fragments.push(series); + group.charges.push(charge); + } + if group_ready(group, &this.eof) + && let Err(error) = this.finish_group(&key) + { + this.lease.fail(Arc::new(error)); + } + } + Poll::Ready(Some(Queued::Eof { input, _metadata })) => { + drained += 1; + if input >= this.eof.len() || this.eof[input] { + this.lease.fail(Arc::new(DataFusionError::Execution( + "duplicate series repartition input EOF".into(), + ))); + continue; + } + this.eof[input] = true; + // An EOF can complete any identity waiting only for this input. + this.rescan = true; + this.resume_after = None; + } + Poll::Ready(None) => { + this.reported = true; + if let Some(Err(error)) = this.lease.ending() { + this.groups.clear(); + this.output.clear(); + return Poll::Ready(Some(Err(error))); + } + if this.eof.iter().any(|seen| !seen) || !this.groups.is_empty() { + let error = boxed(DataFusionError::Execution( + "series repartition queue closed before every input EOF".into(), + )); + this.lease.fail(Arc::new(DataFusionError::Execution( + "series repartition queue closed before every input EOF".into(), + ))); + this.groups.clear(); + return Poll::Ready(Some(Err(error))); + } + if let Some(ready) = this.output.pop_front() { + this.handoff + .lock() + .unwrap_or_else(|e| e.into_inner()) + .current = Some(ready.charges); + return Poll::Ready(Some(Ok((ready.series, true)))); + } + return Poll::Ready(None); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// The boxed error a frame stream reports. +fn boxed(error: DataFusionError) -> BoxedError { + BoxedError::new(common_query::error::Error::from(error)) +} + +/// The hash configuration one execution routes with. +/// +/// The seed is what an execution builds once; a [`XxHash64`] accumulates, so a fresh hasher starts +/// from this configuration for each chunk, and that is the only thing each chunk builds. +#[derive(Clone, Copy, Debug)] +struct Routing { + /// The number of output partitions a hash is reduced modulo. + outputs: usize, + /// The seed every chunk of this execution is hashed with. + seed: u64, +} + +impl Routing { + /// The routing of an execution with `outputs` output partitions. + fn new(outputs: usize) -> Self { + Self { + outputs, + seed: ROUTING_SEED, + } + } + + /// The output partition `metric`'s canonical label set belongs on. + /// + /// The whole label set is hashed, since that is what a merge groups by: pieces of two series + /// whose label sets differ have to be able to land apart, or the split this node exists to make + /// would not happen. [`label_set`] is what canonicalizes the labels - dropping empty values and + /// keeping the metric name - so two pieces that merge hash the same. + fn output(&self, metric: &SeriesLabels) -> usize { + let mut hasher = XxHash64::with_seed(self.seed); + label_set(metric).hash(&mut hasher); + // `outputs` is at least one, which [`SeriesRepartitionExec::try_new`] is what guarantees. + (hasher.finish() % self.outputs as u64) as usize + } +} + +#[cfg(test)] +mod tests; diff --git a/src/promql/src/extension_plan/series_repartition/tests.rs b/src/promql/src/extension_plan/series_repartition/tests.rs new file mode 100644 index 00000000000..6dcbff6ef61 --- /dev/null +++ b/src/promql/src/extension_plan/series_repartition/tests.rs @@ -0,0 +1,5765 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! What this node has to get right: where a chunk lands, that a chain crosses it in one piece, what +//! a stalled output does to the read, what an output learns when the child fails, and that one output +//! is the degenerate case rather than a special one. +//! +//! Some properties need an input that behaves in time rather than only in sequence, so they are tested +//! with children that hold: partitions that stay silent until the test opens a [`Switch`], which is +//! how the streams' ownership of the read is checked, by making the input produce nothing before the +//! plan is dropped. +//! +//! That a partition is read on a task of its own is not one of those properties, and is deliberately +//! not tested as one: it is a structural fact about the read, so it is observed structurally, by +//! recording the task that polls each partition (see [`Recording`]), rather than by requiring two +//! tasks to run on two threads at the same instant - which the test cannot require of the runtime. +//! +//! The child here is scripted rather than planned: each of its partitions answers a list of steps, +//! so a test says which series a chunk belongs to, where on the grid it sits and whether it closes +//! the series instead of relying on where a batch boundary happens to fall. The frames are packed and +//! unpacked by the same helpers the real path uses, so an assertion about an output stream is an +//! assertion about what a consumer of this node decodes. + +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, Weak}; +use std::task::{Poll, Waker}; +use std::time::Duration; + +use datafusion::arrow::array::{Array, Float64Array, ListArray, StructArray}; +use datafusion::common::config::ConfigOptions; +use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::ensure_coop::EnsureCooperative; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::execution_plan::Boundedness; +use datafusion::physical_plan::limit::LocalLimitExec; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::prelude::SessionContext; +use datatypes::series::Label; +use datatypes::series::histogram::{ + BucketSideView, CounterResetHint, HistogramBatchBuilder, HistogramView, +}; +use datatypes::series::samples::{FieldSpec, Samples, SeriesFields}; +use datatypes::series::time::{EvalGrid, GridSpan, GridSpec, TimeAxis}; +use futures::TryStreamExt; + +use super::*; +use crate::extension_plan::test_util::prepare_test_data; +use crate::grid_series::{render_label_set, series_bytes}; +use crate::native::execute::{merge_workspace_bytes, operand_chunk_bytes}; + +const TEST_FIXTURE_VOLUME: usize = 4; +const GRID_START: i64 = 0; +const GRID_STEP: i64 = 1_000; +const GRID_LEN: u64 = 256; + +/// The grid every frame of this module is addressed on: a value per second, 256 of them. +fn grid() -> Arc { + Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, GRID_LEN).unwrap()).unwrap(), + ) +} + +/// A frame schema of one field on [`grid`]. +fn schema() -> Arc { + Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: Some(grid()), + }) +} + +/// Complete labels of the `cpu` metric on one host. +fn labels(host: &str) -> SeriesLabels { + SeriesLabels::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(host), + }], + ) + .expect("the labels of these tests are legal") +} + +fn histogram_chunk(metric: SeriesLabels, position: usize) -> (Series, bool) { + histogram_rows_chunk(metric, position, 1) +} + +fn histogram_rows_chunk(metric: SeriesLabels, position: usize, rows: usize) -> (Series, bool) { + histogram_rows_chunk_with_sum(metric, position, rows, 3.0) +} + +fn histogram_rows_chunk_with_sum( + metric: SeriesLabels, + position: usize, + rows: usize, + sum: f64, +) -> (Series, bool) { + histogram_rows_chunk_with_schema(&schema(), metric, position, rows, sum) +} + +fn histogram_rows_chunk_with_schema( + schema: &Arc, + metric: SeriesLabels, + position: usize, + rows: usize, + sum: f64, +) -> (Series, bool) { + let mut builder = HistogramBatchBuilder::new(); + let mut indices = Vec::with_capacity(rows); + let offsets = [0]; + let lengths = [4]; + let buckets = [1.0, 2.0, 3.0, 4.0]; + for _ in 0..rows { + indices.push( + builder + .push(HistogramView { + schema: 0, + zero_threshold: 0.0, + sum, + count: 4.0, + zero_count: 0.0, + reset_hint: CounterResetHint::Unknown, + start_timestamp: None, + custom_values: &[0.5, 1.5, 2.5, 3.5], + positive: BucketSideView { + span_offsets: &offsets, + span_lengths: &lengths, + buckets: &buckets, + }, + negative: BucketSideView { + span_offsets: &offsets, + span_lengths: &lengths, + buckets: &buckets, + }, + }) + .expect("histogram row builds"), + ); + } + let table = builder.finish().expect("histogram table builds"); + let fields = SeriesFields::one( + Samples::try_from_histogram(datafusion::arrow::array::UInt32Array::from(indices), table) + .expect("histogram indices reference their table"), + ); + let time = TimeAxis::from_grid_spans( + schema.eval_grid.clone().unwrap(), + vec![GridSpan::new(position, rows)], + ) + .expect("histogram rows lie on the grid"); + (Series::try_new(schema, metric, time, fields).unwrap(), true) +} + +fn sparse_histogram_rows_chunk( + metric: SeriesLabels, + positions: impl IntoIterator, + rows: usize, + sum: f64, +) -> (Series, bool) { + let (mut series, is_last) = histogram_rows_chunk_with_sum(metric, 0, rows, sum); + series.time = TimeAxis::from_grid_spans( + schema().eval_grid.clone().unwrap(), + positions + .into_iter() + .map(|position| GridSpan::new(position, 1)) + .collect(), + ) + .expect("sparse histogram rows lie on the grid"); + (series, is_last) +} + +fn sliced_chunks(metric: SeriesLabels) -> [(Series, bool); 2] { + let schema = schema(); + let values = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])); + [0, 2].map(|start| { + let time = TimeAxis::from_grid_spans( + schema.eval_grid.clone().unwrap(), + vec![GridSpan::new(start, 2)], + ) + .expect("slice lies on the grid"); + let samples = Samples::from_float(values.slice(start, 2)); + let series = Series::try_new(&schema, metric.clone(), time, SeriesFields::one(samples)) + .expect("sliced series builds"); + (series, start == 2) + }) +} + +/// A single chain split into more output fragments than one frame can encode. +#[tokio::test] +async fn public_frame_retains_large_series_charge_until_remainder_is_dropped() { + const FRAGMENTS: usize = 257; + let metric = labels("frame-remainder"); + let eval_grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, FRAGMENTS as u64).unwrap()) + .unwrap(), + ); + let series_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: Some(eval_grid), + }); + let steps: Vec = (0..FRAGMENTS) + .map(|position| { + let time = TimeAxis::from_grid_spans( + series_schema.eval_grid.clone().unwrap(), + vec![GridSpan::new(position, 1)], + ) + .unwrap(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + position as f64, + ]))); + let series = Series::try_new(&series_schema, metric.clone(), time, fields).unwrap(); + Step::Chunk(series, position + 1 == FRAGMENTS) + }) + .collect(); + let probe_pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let probe_context = task_context(&probe_pool); + let probe = Scripted::plan_with_schema(vec![steps.clone()], Arc::clone(&series_schema)); + let mut probe = + unpack_evaluated_series(probe.execute(0, probe_context).unwrap(), None).unwrap(); + let (decoded, errors) = until_end(&mut probe).await; + assert!(errors.is_empty()); + assert_eq!(FRAGMENTS, decoded.len()); + let payload = decoded + .iter() + .map(|(series, _)| operand_chunk_bytes(series)) + .sum::(); + let expected_values = decoded + .iter() + .flat_map(|(series, _)| { + series + .fields + .iter() + .flat_map(|field| field.iter()) + .map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => panic!("expected float sample"), + }) + }) + .collect::>(); + assert_eq!( + (0..FRAGMENTS).map(|p| p as f64).collect::>(), + expected_values + ); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let plan = node( + Scripted::plan_with_schema(vec![steps.clone()], Arc::clone(&series_schema)), + 1, + ); + let mut frames = plan.execute(0, Arc::clone(&context)).unwrap(); + wait_until("the producer queues the completed series", || { + producer_stopped(&execution(&plan)) + }) + .await; + let first_frame = futures::future::poll_fn(|cx| frames.as_mut().poll_next(cx)) + .await + .expect("public frame arrives") + .expect("public frame succeeds"); + assert_eq!( + 256, + first_frame.num_rows(), + "first public RecordBatch is full" + ); + assert!( + pool.reserved() >= payload + merge_workspace_bytes(FRAGMENTS, 1, 1), + "the unread remainder retains its payload plus merge workspace" + ); + drop(frames); + assert_eq!( + 0, + pool.reserved(), + "dropping output releases remainder charges" + ); + drop(first_frame); + + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let plan = node(Scripted::plan_with_schema(vec![steps], series_schema), 1); + let mut frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let mut retained_frames = Vec::new(); + while let Some(frame) = futures::future::poll_fn(|cx| frames.as_mut().poll_next(cx)).await { + retained_frames.push(frame.expect("public frame succeeds")); + } + assert_eq!(2, retained_frames.len()); + assert_eq!(256, retained_frames[0].num_rows()); + assert_eq!(1, retained_frames[1].num_rows()); + // The public frames are one series split by the encoder's row budget, so they are decoded as + // one stream: only the last frame's row closes the chain. + let frame_schema = retained_frames[0].schema(); + let decoded = unpack_evaluated_series( + Box::pin(RecordBatchStreamAdapter::new( + frame_schema, + futures::stream::iter(std::mem::take(&mut retained_frames).into_iter().map(Ok)), + )), + None, + ) + .unwrap(); + futures::pin_mut!(decoded); + let mut values = Vec::new(); + let mut closures = 0usize; + while let Some(item) = decoded.as_mut().next().await { + let (series, is_last) = item.unwrap(); + closures += usize::from(is_last); + let field = series.fields.get(0).expect("one field"); + for (position, row) in series.time.positions() { + let value = match field.sample(row) { + datatypes::series::samples::SampleRef::Float(value) => value, + _ => panic!("expected float sample"), + }; + values.push((position, value)); + } + } + assert_eq!(1, closures, "the split series closes exactly once"); + let expected_values = (0..FRAGMENTS) + .map(|position| (position, position as f64)) + .collect::>(); + assert_eq!( + expected_values, values, + "all decoded positions stay in order" + ); + // The decoder still owns the frames it drained until it is dropped. + let _ = decoded; + drop(frames); + assert_eq!( + 0, + pool.reserved(), + "dropping delivered batches releases all charges" + ); +} + +/// One evaluated chunk: `len` rows from grid position `start`, and whether it closes its series. +/// +/// A chunk of no rows is the closing chunk a series sends when nothing is left of it, which is a +/// shape of its own: it carries the finality and no position. +fn chunk(metric: SeriesLabels, start: usize, len: usize, is_last: bool) -> (Series, bool) { + let schema = schema(); + let spans = if len == 0 { + Vec::new() + } else { + vec![GridSpan::new(start, len)] + }; + let time = TimeAxis::from_grid_spans(schema.eval_grid.clone().unwrap(), spans) + .expect("one run on the grid"); + let values = (0..len) + .map(|index| Some((start + index) as f64)) + .collect::>(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(values))); + ( + Series::try_new(&schema, metric, time, fields).expect("a chunk of this schema"), + is_last, + ) +} + +/// One step of a scripted child partition. +#[derive(Clone)] +enum Step { + /// A chunk and its finality. + Chunk(Series, bool), + /// The child fails here. + Error, + /// The child fails here with a message of its own, for a test that has to tell two failures apart. + ErrorSaying(&'static str), +} + +impl Step { + fn chunk((series, is_last): (Series, bool)) -> Self { + Self::Chunk(series, is_last) + } +} + +/// What one scripted step hands over. +fn step_answer(step: Step) -> Result<(Series, bool), BoxedError> { + match step { + Step::Chunk(series, is_last) => Ok((series, is_last)), + Step::Error => Err(boxed(DataFusionError::Execution( + "scripted child failure".to_string(), + ))), + Step::ErrorSaying(message) => Err(boxed(DataFusionError::Execution(message.to_string()))), + } +} + +/// The steps of one partition, as the frame stream that hands them over. +struct Script { + steps: std::vec::IntoIter, +} + +impl Stream for Script { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(self.steps.next().map(step_answer)) + } +} + +/// A child whose partitions answer scripted steps, each as one evaluated frame. +struct Scripted { + steps: Vec>, + series_schema: Arc, + properties: Arc, +} + +impl Scripted { + /// The child plan over `steps`, one partition per list. + fn plan(steps: Vec>) -> Arc { + Self::plan_with_schema(steps, schema()) + } + + /// The child plan over `steps` with an explicitly declared evaluated frame schema. + fn plan_with_schema( + steps: Vec>, + series_schema: Arc, + ) -> Arc { + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(steps.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Arc::new(Self { + steps, + series_schema, + properties, + }) + } +} + +impl fmt::Debug for Scripted { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Scripted") + .field("partitions", &self.steps.len()) + .finish() + } +} + +impl DisplayAs for Scripted { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Scripted: partitions={}", self.steps.len()) + } +} + +impl ExecutionPlan for Scripted { + fn name(&self) -> &str { + "Scripted" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a scripted child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let steps = self.steps.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("scripted child has no partition {partition}")) + })?; + pack_evaluated_series( + Box::pin(Script { + steps: steps.into_iter(), + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// A task context whose memory pool is `pool`, so what the node charges can be read afterwards. +fn task_context(pool: &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))) +} + +/// The node over `child` routing into `outputs` partitions. +fn node(child: Arc, outputs: usize) -> Arc { + Arc::new(SeriesRepartitionExec::try_new(child, outputs).expect("the node is built")) +} + +/// Every output partition of `plan`, decoded, and the error it ended with if it ended with one. +/// +/// All partitions are read at once, as in normal execution. Draining them concurrently makes this +/// helper work for outputs that can publish before global EOF as well as outputs that cannot. +async fn drain( + plan: &Arc, + context: &Arc, +) -> Vec, BoxedError>> { + let mut streams = Vec::new(); + for partition in 0..plan.properties().partitioning.partition_count() { + let stream = plan + .execute(partition, Arc::clone(context)) + .expect("the output executes"); + streams.push( + unpack_evaluated_series(stream, None).expect("the output declares an evaluated frame"), + ); + } + futures::future::join_all( + streams + .into_iter() + .map(|stream| stream.try_collect::>()), + ) + .await +} + +/// The output partition of `plan` at `partition`, as the frame stream a consumer reads. +fn output( + plan: &Arc, + partition: usize, + context: &Arc, +) -> SendableEvaluatedSeriesStream { + unpack_evaluated_series( + plan.execute(partition, Arc::clone(context)) + .expect("the output executes"), + None, + ) + .expect("the output declares an evaluated frame") +} + +/// The output partition of `plan` at `partition`, for a test that keeps the node itself. +/// +/// [`execution`] reaches the read through the node rather than through the plan a consumer would hold, +/// so a test that has to look at the tasks behind an output claims it here instead of coercing the +/// node to the plan it is: the two are the same call, one on a reference the test already has. +fn output_of( + plan: &Arc, + partition: usize, + context: &Arc, +) -> SendableEvaluatedSeriesStream { + unpack_evaluated_series( + plan.execute(partition, Arc::clone(context)) + .expect("the output executes"), + None, + ) + .expect("the output declares an evaluated frame") +} + +/// The labels of a chunk in the form a failure message can show. +fn named(metric: &SeriesLabels) -> String { + render_label_set(&label_set(metric)) +} + +/// What one chunk says, in the terms a test compares: its canonical label set, where its rows start, +/// how many there are, and whether it closed its series. +fn summary((series, is_last): &(Series, bool)) -> (String, Option, usize, bool) { + ( + named(&series.metric), + series.time.first_position(), + series.time.len(), + *is_last, + ) +} + +/// The summaries of a whole output stream. +fn summaries(chunks: &[(Series, bool)]) -> Vec<(String, Option, usize, bool)> { + chunks.iter().map(summary).collect() +} + +/// Asserts that every chain `chunks` hands over is whole, in the terms the frame validator reads a +/// chain in. +/// +/// A chunk that leaves a series open must be followed by the same series' next chunk, at a position +/// after the last one that chain carried. Another series' chunk in between is what a merge downstream +/// cannot repair, since it would have to hold two open series to read either; a position that does not +/// advance is a chunk the merge would have to reorder to read. Ending with a chain open fails too, +/// because that means the finality of a chunk was lost on the way, and a chunk of no rows fails unless +/// it is the closing chunk of the chain it belongs to - that being the shape a series ends in when +/// nothing of it is left. +/// +/// Identity here is the label set the frame carries and not its rendering: two label sets that print +/// the same are still two series to a merge - an empty label value is not printed - so comparing the +/// displayed labels would accept a stream that changed identity mid-chain. +fn assert_chains_are_whole(chunks: &[(Series, bool)]) { + // What the validator keeps of a chain it has not seen close: the identity it opened with, and the + // last position it carried. + let mut open: Option<(SeriesLabels, usize)> = None; + for (series, is_last) in chunks { + let (Some(first), Some(last)) = (series.time.first_position(), series.time.last_position()) + else { + // No position, so the chunk carries no rows: only a closing chunk is shaped that way, and + // it closes the chain it belongs to without advancing it. + assert!( + *is_last, + "{} carries no positions and does not close its series", + named(&series.metric) + ); + if let Some((open_labels, _)) = &open { + assert_eq!( + open_labels, &series.metric, + "another series' chunk interrupted an open chain" + ); + } + open = None; + continue; + }; + if let Some((open_labels, previous)) = &open { + assert_eq!( + open_labels, &series.metric, + "another series' chunk interrupted an open chain" + ); + assert!( + first > *previous, + "the positions of one chain did not advance across its chunks: {} starts at {first} \ + after the chain's previous chunk ended at {previous}", + named(&series.metric), + ); + } + open = (!*is_last).then(|| (series.metric.clone(), last)); + } + assert!(open.is_none(), "the output ended with a chain open"); +} + +/// Checks completed owner outputs against exact series identities, positions, and finality. +fn assert_output_contract(chunks: &[(Series, bool)], expected: &[(SeriesLabels, Vec)]) { + assert_chains_are_whole(chunks); + let mut actual: Vec<(SeriesLabels, Vec, usize)> = Vec::new(); + let mut closed_identities = Vec::::new(); + for (series, is_last) in chunks { + assert!( + !closed_identities.contains(&series.metric), + "a closed full identity must not reappear: {}", + named(&series.metric) + ); + assert_eq!(series.time.len(), series.time.positions().count()); + let (_, positions, closes) = match actual + .iter_mut() + .find(|(metric, _, _)| metric == &series.metric) + { + Some(entry) => entry, + None => { + actual.push((series.metric.clone(), Vec::new(), 0)); + actual.last_mut().unwrap() + } + }; + positions.extend(series.time.positions().map(|(position, _)| position)); + *closes += usize::from(*is_last); + if *is_last { + closed_identities.push(series.metric.clone()); + } + } + assert_eq!( + expected.len(), + actual.len(), + "one closed series per full identity" + ); + for (metric, positions) in expected { + let index = actual + .iter() + .position(|(identity, _, _)| identity == metric) + .expect("expected exact full label identity"); + let (_, actual_positions, closes) = actual.remove(index); + assert_eq!( + positions, + &actual_positions, + "positions for {}", + named(metric) + ); + assert_eq!(1, closes, "one terminal output close for {}", named(metric)); + } + assert!(actual.is_empty()); +} + +/// A label set whose output is not `other`'s, found rather than assumed. +fn labels_elsewhere(routing: &Routing, other: &SeriesLabels) -> SeriesLabels { + let taken = routing.output(other); + for index in 0..64 { + let candidate = labels(&format!("elsewhere-{index}")); + if routing.output(&candidate) != taken { + return candidate; + } + } + panic!("no label set of these tests routes to another output"); +} + +/// A node with no outputs, or over a child that does not declare a series frame, is refused before +/// the plan can run. +#[test] +fn construction_rejects_a_zero_output_count_and_a_plain_child() { + let none = SeriesRepartitionExec::try_new(Scripted::plan(vec![Vec::new()]), 0) + .expect_err("a node has to have an output"); + assert!( + none.to_string().contains("at least one output partition"), + "{none}" + ); + + let plain: Arc = Arc::new(prepare_test_data()); + let refused = SeriesRepartitionExec::try_new(plain, 2) + .expect_err("a plain record batch stream carries no series labels"); + assert!( + refused.to_string().contains("is not a series frame"), + "{refused}" + ); +} + +/// T1: two partitions each hold a piece of one canonical label set, and both pieces reach the output +/// that label set hashes to; another label set reaches the output its own hash names. +#[tokio::test] +async fn pieces_of_one_label_set_meet_on_one_output() { + let routing = Routing::new(3); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + assert_ne!( + routing.output(&shared), + routing.output(&other), + "the two label sets have to be able to land apart for this test to mean anything" + ); + // One label set split over the child's two partitions: the first piece opens and closes on its + // own, the second continues at a later position. Neither piece alone is the series. + let child = Scripted::plan(vec![ + vec![ + Step::chunk(chunk(shared.clone(), 0, 2, false)), + Step::chunk(chunk(shared.clone(), 0, 0, true)), + Step::chunk(chunk(other.clone(), 0, 1, true)), + ], + vec![Step::chunk(chunk(shared.clone(), 8, 2, true))], + ]); + let plan: Arc = node(child, 3); + assert_eq!(3, plan.properties().partitioning.partition_count()); + + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + assert_eq!(3, outputs.len()); + let outputs = outputs + .into_iter() + .map(|chunks| { + let chunks = chunks.expect("every output decodes"); + assert_chains_are_whole(&chunks); + summaries(&chunks) + }) + .collect::>(); + + let other_piece = named(&other); + let shared_output = routing.output(&shared); + let shared_chunks = outputs[shared_output] + .iter() + .map(|summary| summary.2) + .sum::(); + assert_eq!( + 4, shared_chunks, + "both producer partitions contribute positions" + ); + assert_eq!( + 1, + outputs[shared_output] + .iter() + .filter(|chunk| chunk.3) + .count() + ); + assert_eq!( + vec![(other_piece, Some(0), 1, true)], + outputs[routing.output(&other)], + "another full identity lands on its own output" + ); + let empty = outputs + .iter() + .enumerate() + .filter(|(partition, chunks)| *partition != shared_output && chunks.is_empty()) + .map(|(partition, _)| partition) + .collect::>(); + assert_eq!( + 1, + empty.len(), + "the third output receives neither label set" + ); +} + +/// T2: two chains of one label set cross the node whole, in the order the read reached them, and a +/// closing chunk that carries no row survives the round trip. +#[tokio::test] +async fn chains_cross_the_node_whole_and_a_closing_chunk_survives() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + // Both of these chains hash to the same output, so a node that routed chunk by chunk without + // reading one chain to its end would interleave them there - which is the failure this test is + // built to see. + let child = Scripted::plan(vec![ + vec![ + Step::chunk(chunk(shared.clone(), 0, 2, false)), + Step::chunk(chunk(shared.clone(), 2, 2, true)), + Step::chunk(chunk(other.clone(), 0, 1, true)), + ], + vec![ + Step::chunk(chunk(shared.clone(), 8, 2, true)), + Step::chunk(chunk(other.clone(), 0, 0, true)), + ], + ]); + let plan: Arc = node(child, 2); + + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + let shared_output = routing.output(&shared); + let chunks = outputs + .into_iter() + .nth(shared_output) + .expect("an output per partition") + .expect("the output decodes"); + assert_chains_are_whole(&chunks); + + assert_output_contract(&chunks, &[(shared.clone(), vec![0, 1, 2, 3, 8, 9])]); +} + +/// An unread output does not prevent routing to another owner, and a completed series is readable +/// while a later input chain is gated. +#[tokio::test] +async fn an_unread_owner_does_not_stall_routing_and_outputs_publish_incrementally() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + let stalled = routing.output(&shared); + let moving = routing.output(&other); + // More chains than one queue can hold, so the producer has to stop mid-read rather than finish. + let chains = TEST_FIXTURE_VOLUME + 2; + let completed = (0..chains) + .map(|index| labels(&format!("unread-{index}"))) + .filter(|metric| routing.output(metric) == stalled) + .collect::>(); + assert!( + completed.len() > 1, + "enough completed identities route to the stalled owner" + ); + let mut steps = completed + .iter() + .enumerate() + .map(|(index, metric)| Step::chunk(chunk(metric.clone(), index, 1, true))) + .collect::>(); + steps.push(Step::chunk(chunk(other.clone(), 0, 1, true))); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![GatedPartition { + before: steps, + after: Vec::new(), + switch: None, + }], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let concrete_plan = node(child, 2); + let plan: Arc = concrete_plan.clone(); + let context = SessionContext::default().task_ctx(); + + let mut moving_stream = output(&plan, moving, &context); + let stalled_stream = output(&plan, stalled, &context); + let execution = execution(&concrete_plan); + // The source holds terminal EOF behind the switch. Thus the producer cannot finish early merely + // because this finite input happened to reach its legitimate end. + assert!(!switch.is_open()); + let moving_chunk = tokio::time::timeout(Duration::from_secs(10), moving_stream.next()) + .await + .expect("completed moving identity arrives before gated EOF") + .expect("moving identity exists") + .expect("moving output succeeds"); + assert_eq!(other, moving_chunk.0.metric); + assert!(moving_chunk.1); + assert!( + futures::future::poll_fn(|cx| match moving_stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "no second identity is ready before the gate" + ); + assert!( + !producer_stopped(&execution), + "terminal EOF gate remains closed" + ); + switch.open(); + wait_until("unbounded producer finishes after terminal EOF", || { + producer_stopped(&execution) + }) + .await; + // The unread owner's completed identities remain queued while the moving output was drained. + let stalled_chunks = tokio::time::timeout( + Duration::from_secs(10), + stalled_stream.try_collect::>(), + ) + .await + .expect("stalled owner publishes after global EOF") + .expect("the stalled output decodes"); + let expected_stalled = completed + .iter() + .enumerate() + .map(|(index, metric)| (metric.clone(), vec![index])) + .collect::>(); + assert_output_contract(&stalled_chunks, &expected_stalled); + + let moving_rest = tokio::time::timeout( + Duration::from_secs(10), + moving_stream.try_collect::>(), + ) + .await + .expect("the moving output ends after global EOF") + .expect("the moving output decodes"); + assert!( + moving_rest.is_empty(), + "completed identity is emitted only once" + ); +} + +/// An exactly full output frame releases its completed identity before the next poll yields. +#[tokio::test] +async fn exact_full_frame_releases_completed_identity_while_other_chains_stay_open() { + let metric = labels("exact-full-frame"); + let a_bytes = (0..GRID_LEN as usize) + .map(|position| { + operand_chunk_bytes( + &chunk( + metric.clone(), + position, + 1, + position + 1 == GRID_LEN as usize, + ) + .0, + ) + }) + .sum::(); + let held_b = labels("exact-full-held-b"); + let held_c = labels("exact-full-held-c"); + let held_bytes = operand_chunk_bytes(&chunk(held_b.clone(), 0, 2, false).0) + + operand_chunk_bytes(&chunk(held_c.clone(), 0, 2, false).0); + let switch = Arc::new(Switch::default()); + let ranges = [(0, 86), (86, 171), (171, GRID_LEN as usize)]; + let partitions = ranges + .into_iter() + .enumerate() + .map(|(partition, (start, end))| { + let mut before = (start..end) + .map(|position| { + Step::chunk(chunk(metric.clone(), position, 1, position + 1 == end)) + }) + .collect::>(); + let held = match partition { + 1 => &held_b, + 2 => &held_c, + _ => return GatedPartition::open(before), + }; + before.push(Step::chunk(chunk(held.clone(), 0, 2, false))); + GatedPartition { + before, + after: vec![Step::chunk(chunk(held.clone(), 2, 1, true))], + switch: None, + } + }) + .collect(); + let child = Gated::plan( + partitions, + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan = node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let mut frames = plan.execute(0, Arc::clone(&context)).unwrap(); + wait_until("A and the open B/C prefixes are retained", || { + pool.reserved() >= a_bytes + held_bytes + }) + .await; + let reserved_before = pool.reserved(); + let first = futures::future::poll_fn(|cx| frames.as_mut().poll_next(cx)) + .await + .expect("the exact-full public frame arrives") + .expect("the frame succeeds"); + assert_eq!(256, first.num_rows()); + assert!( + futures::future::poll_fn(|cx| match frames.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "with B/C still open, the next poll waits rather than requiring a later frame" + ); + assert!( + pool.reserved() < reserved_before, + "completed A charge is released" + ); + assert!( + pool.reserved() >= held_bytes, + "open B/C payload charges remain retained" + ); + switch.open(); + drop(frames); + wait_until("dropping the stream releases remaining charges", || { + pool.reserved() == 0 + }) + .await; +} + +/// An absent identity contribution waits for that input's normal EOF, not another identity's close. +#[tokio::test] +async fn absent_identity_waits_for_normal_eof_before_publishing_other_input_chain() { + let metric = labels_routed_to(0, 2).remove(0); + let other = labels_routed_to(1, 2).remove(0); + let input_zero = Arc::new(Switch::default()); + let input_one = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + GatedPartition { + before: vec![Step::chunk(chunk(metric.clone(), 0, 2, true))], + after: vec![Step::chunk(chunk(other.clone(), 0, 1, true))], + switch: Some(Arc::clone(&input_zero)), + }, + GatedPartition { + before: vec![Step::chunk(chunk(other.clone(), 2, 1, true))], + after: Vec::new(), + switch: Some(Arc::clone(&input_one)), + }, + ], + Arc::new(Switch::default()), + Arc::new(AtomicUsize::new(0)), + ); + let plan = node(child, 2); + let context = SessionContext::default().task_ctx(); + let mut stream = output_of(&plan, 0, &context); + assert!( + futures::future::poll_fn(|cx| match stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "A is not ready while input 1 has contributed only another identity and remains open" + ); + input_one.open(); + let mut chunks = Vec::new(); + loop { + let next = tokio::time::timeout(Duration::from_secs(10), async { + futures::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await + }) + .await + .expect("input 1 EOF releases A while input 0 remains gated"); + match next { + Some(Ok(chunk)) => { + let done = chunk.1; + chunks.push(chunk); + if done { + break; + } + } + Some(Err(error)) => panic!("A fails: {error}"), + None => panic!("output ended while input 0 remains open"), + } + } + assert_output_contract(&chunks, &[(metric, vec![0, 1])]); + assert!(!input_zero.is_open(), "input 0 remains independently gated"); + input_zero.open(); + let (rest, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut stream)) + .await + .expect("the remaining input finishes after its independent release"); + assert!(errors.is_empty(), "remaining output succeeds: {errors:?}"); + assert_chains_are_whole(&rest); +} + +/// A complete identity can be read from the public output while every input remains live. +#[tokio::test] +async fn completed_identity_is_published_before_inputs_finish() { + let metric = labels("early-series"); + let routing = Routing::new(2); + let switch = Arc::new(Switch::default()); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let fragments = (0..3) + .map(|input| { + vec![ + Step::chunk(chunk(metric.clone(), input * 10, 1, false)), + Step::chunk(chunk(metric.clone(), input * 10 + 1, 1, true)), + ] + }) + .collect::>(); + let partitions = fragments + .into_iter() + .enumerate() + .map(|(input, before)| { + let held = labels(&format!("held-{input}")); + GatedPartition { + before: { + let mut before = before; + before.push(Step::chunk(chunk(held.clone(), 40, 2, false))); + before + }, + after: vec![Step::chunk(chunk(held, 42, 1, true))], + switch: None, + } + }) + .collect::>(); + let child = Gated::plan( + partitions, + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan = node(child, 2); + let owner = routing.output(&metric); + let mut stream = output_of(&plan, owner, &context); + let a_bytes = (0..3) + .flat_map(|input| [input * 10, input * 10 + 1]) + .map(|position| operand_chunk_bytes(&chunk(metric.clone(), position, 1, true).0)) + .sum::(); + let held_bytes = (0..3) + .map(|input| operand_chunk_bytes(&chunk(labels(&format!("held-{input}")), 40, 2, false).0)) + .sum::(); + assert!(held_bytes >= 3 * operand_chunk_bytes(&chunk(labels("held-0"), 40, 2, false).0)); + wait_until( + "completed A and open held fragments are retained before handoff", + || pool.reserved() >= a_bytes + held_bytes, + ) + .await; + let reserved_before_a = pool.reserved(); + + let mut received = Vec::new(); + let first = tokio::time::timeout(Duration::from_secs(10), async { + futures::future::poll_fn(|cx| match stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Some(Ok(chunk))) => Poll::Ready(Ok(chunk)), + Poll::Ready(Some(Err(error))) => Poll::Ready(Err(error)), + Poll::Ready(None) => Poll::Ready(Err(boxed(DataFusionError::Execution( + "output ended before the gated inputs".into(), + )))), + }) + .await + }) + .await + .expect("completed identity is delivered while inputs are gated") + .expect("output succeeds"); + received.push(first); + while !received + .last() + .is_some_and(|(series, is_last)| series.metric == metric && *is_last) + { + let next = tokio::time::timeout(Duration::from_secs(10), async { + futures::future::poll_fn(|cx| match stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Pending, + Poll::Ready(Some(Ok(chunk))) => Poll::Ready(Ok(chunk)), + Poll::Ready(Some(Err(error))) => Poll::Ready(Err(error)), + Poll::Ready(None) => Poll::Ready(Err(boxed(DataFusionError::Execution( + "output ended before the gated inputs".into(), + )))), + }) + .await + }) + .await + .expect("completed identity continues to arrive") + .expect("output succeeds"); + received.push(next); + } + assert_output_contract(&received, &[(metric.clone(), vec![0, 1, 10, 11, 20, 21])]); + assert!(!switch.is_open(), "later identities are still blocked"); + assert!( + pool.reserved() < reserved_before_a, + "handing off A releases its per-identity charge: before {}, after {}", + reserved_before_a, + pool.reserved() + ); + assert!( + pool.reserved() >= held_bytes, + "the held identities retain their independent payload charge: expected at least {held_bytes}, got {}", + pool.reserved() + ); + + switch.open(); + let (rest, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut stream)) + .await + .expect("gated inputs finish"); + assert!(errors.is_empty(), "later output succeeds: {errors:?}"); + received.extend(rest); + assert_chains_are_whole(&received); +} + +/// A one-fragment large identity needs no merge workspace and fits its own finite payload budget. +#[tokio::test] +async fn single_fragment_large_identity_fits_finite_payload_budget() { + let metric = labels("single-fragment-large"); + let long_grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, 1_000).unwrap()).unwrap(), + ); + let long_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: Some(Arc::clone(&long_grid)), + }); + let time = TimeAxis::from_grid_spans(long_grid, vec![GridSpan::new(0, 1_000)]) + .expect("one run on the long grid"); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from( + (0..1_000) + .map(|value| Some(value as f64)) + .collect::>(), + ))); + let series = Series::try_new(&long_schema, metric.clone(), time, fields) + .expect("one large float series chunk"); + let payload = operand_chunk_bytes(&series); + let budget = payload + 8 * 1024; + let workspace = merge_workspace_bytes(1_000, 1, 1); + assert!( + workspace > budget, + "estimated workspace must exceed payload budget: workspace={workspace}, budget={budget}" + ); + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let context = task_context(&pool); + let plan = node( + Scripted::plan_with_schema(vec![vec![Step::Chunk(series, true)]], long_schema), + 1, + ); + let mut output = output_of(&plan, 0, &context); + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .expect("single-fragment identity fits its finite pool"); + assert!( + errors.is_empty(), + "single fragment needs no merge workspace: {errors:?}" + ); + assert_eq!(1, chunks.len()); + let positions = chunks[0] + .0 + .time + .positions() + .map(|(position, _)| position) + .collect::>(); + assert_eq!((0..1_000).collect::>(), positions); + assert_eq!(metric, chunks[0].0.metric); + assert!(chunks[0].1); + assert!(payload < budget); + assert_eq!( + 0, + pool.reserved(), + "delivery releases the single-fragment charge" + ); +} + +/// An input EOF already consumed still counts when a different input first contributes an identity. +#[tokio::test] +async fn eof_before_first_identity_contribution_makes_group_ready() { + let metric = labels("late-first-chain"); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + GatedPartition::open(Vec::new()), + GatedPartition::held(vec![Step::chunk(chunk(metric.clone(), 7, 1, true))]), + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan = node(child, 1); + let context = SessionContext::default().task_ctx(); + let mut stream = output_of(&plan, 0, &context); + assert!( + futures::future::poll_fn(|cx| match stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "output waits for the not-yet-arrived identity" + ); + switch.open(); + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut stream)) + .await + .expect("late first chain is processed after the other input EOF"); + assert!(errors.is_empty(), "output succeeds: {errors:?}"); + assert_output_contract(&chunks, &[(metric, vec![7])]); +} + +/// T3, second half: what the producer queues is charged to the execution's pool, and dropping the +/// execution releases it. +#[tokio::test] +async fn dropping_the_execution_releases_what_the_producer_reserved() { + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let shared = labels("web-1"); + let child = Scripted::plan(vec![vec![ + Step::chunk(chunk(shared.clone(), 0, 1, false)), + Step::chunk(chunk(shared.clone(), 1, 1, false)), + Step::chunk(chunk(shared.clone(), 2, 1, true)), + ]]); + let plan: Arc = node(child, 2); + let stream = plan + .execute(0, Arc::clone(&context)) + .expect("the first output starts the read"); + + // The producer reads ahead into the queue, and what it holds is charged to the pool. + let held = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let reserved = pool.reserved(); + if reserved > 0 { + return reserved; + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("the producer charges what it queues"); + assert!( + held > 0, + "the queued chain has a nonzero reservation: {held}" + ); + + drop(stream); + drop(plan); + tokio::time::timeout(Duration::from_secs(10), async { + while pool.reserved() != 0 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap_or_else(|_| { + panic!( + "the dropped execution still holds {} bytes", + pool.reserved() + ) + }); +} + +/// T5: one output is the same node, not a special case - every chunk arrives, in one chain after +/// another, with nothing filtered out. +#[tokio::test] +async fn one_output_keeps_every_chunk_and_its_chains_whole() { + let shared = labels("web-1"); + let other = labels("elsewhere-1"); + let child = Scripted::plan(vec![ + vec![ + Step::chunk(chunk(shared.clone(), 0, 2, false)), + Step::chunk(chunk(shared.clone(), 2, 2, true)), + Step::chunk(chunk(other.clone(), 0, 1, false)), + Step::chunk(chunk(other.clone(), 0, 0, true)), + ], + vec![Step::chunk(chunk(shared.clone(), 8, 2, true))], + ]); + let plan: Arc = node(child, 1); + assert_eq!(1, plan.properties().partitioning.partition_count()); + + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + assert_eq!(1, outputs.len()); + let chunks = outputs + .into_iter() + .next() + .expect("one output") + .expect("the output decodes"); + assert_chains_are_whole(&chunks); + assert_output_contract( + &chunks, + &[ + (shared.clone(), vec![0, 1, 2, 3, 8, 9]), + (other.clone(), vec![0]), + ], + ); + assert_eq!(2, chunks.iter().filter(|(_, closed)| *closed).count()); +} + +/// A switch a test opens to let gated input streams produce. +/// +/// Opening it wakes the streams already parked, and a stream polled after the opening sees it +/// straight away: registering before the second read is what makes an opening that lands between the +/// two impossible to lose. +#[derive(Default)] +struct Switch { + open: AtomicBool, + wakers: Mutex>, +} + +impl Switch { + /// Lets every stream waiting on this switch produce, and wakes the ones already parked. + fn open(&self) { + self.open.store(true, Ordering::Release); + let wakers = std::mem::take(&mut *self.wakers.lock().unwrap()); + for waker in wakers { + waker.wake(); + } + } + + fn is_open(&self) -> bool { + self.open.load(Ordering::Acquire) + } + + /// Parks `cx` until the switch is opened, unless it already is. + fn wait(&self, cx: &Context<'_>) { + if self.is_open() { + return; + } + self.wakers.lock().unwrap().push(cx.waker().clone()); + if self.is_open() { + // The opening landed between the two reads, so its wake did not find this waker. + cx.waker().wake_by_ref(); + } + } +} + +/// One partition of a [`Gated`] child: what it hands over at once, and what it holds back. +/// `switch` overrides the child's shared gate when a test needs independent input progress. +#[derive(Clone, Default)] +struct GatedPartition { + before: Vec, + after: Vec, + switch: Option>, +} + +impl GatedPartition { + /// A partition that hands over `steps` only once the switch is opened. + fn held(steps: Vec) -> Self { + Self { + before: Vec::new(), + after: steps, + switch: None, + } + } + + /// A partition that hands over `steps` as soon as it is polled. + fn open(steps: Vec) -> Self { + Self { + before: steps, + after: Vec::new(), + switch: None, + } + } +} + +/// A child whose partitions hold their second half back until `switch` is opened. +/// +/// The gate is how a test states that the read starts only after the plan is gone: the input is a +/// real source that has not produced yet, not one that has already answered before the interesting +/// moment. `dropped` counts the streams a consumer let go of, which is how a test sees the read +/// cancelling the child it was reading. +struct Gated { + partitions: Vec, + switch: Arc, + dropped: Arc, + series_schema: Arc, + properties: Arc, +} + +impl Gated { + /// The child over `partitions`, gated on `switch` and counting drops in `dropped`. + fn plan( + partitions: Vec, + switch: Arc, + dropped: Arc, + ) -> Arc { + Self::plan_with_schema(partitions, switch, dropped, schema()) + } + + /// The same child over an explicit schema, for trees with more than one field. + fn plan_with_schema( + partitions: Vec, + switch: Arc, + dropped: Arc, + series_schema: Arc, + ) -> Arc { + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Arc::new(Self { + partitions, + switch, + dropped, + series_schema, + properties, + }) + } +} + +impl fmt::Debug for Gated { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Gated") + .field("partitions", &self.partitions.len()) + .finish() + } +} + +impl DisplayAs for Gated { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Gated: partitions={}", self.partitions.len()) + } +} + +impl ExecutionPlan for Gated { + fn name(&self) -> &str { + "Gated" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a gated child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let steps = self.partitions.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("gated child has no partition {partition}")) + })?; + let switch = steps.switch.unwrap_or_else(|| Arc::clone(&self.switch)); + pack_evaluated_series( + Box::pin(GatedScript { + before: steps.before.into_iter(), + after: steps.after.into_iter(), + switch, + dropped: Arc::clone(&self.dropped), + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// One partition of a [`Gated`] child. +struct GatedScript { + before: std::vec::IntoIter, + after: std::vec::IntoIter, + switch: Arc, + dropped: Arc, +} + +impl Stream for GatedScript { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if let Some(step) = this.before.next() { + return Poll::Ready(Some(step_answer(step))); + } + // Nothing of the second half is handed over before the test opens the switch, so the read + // can be started and then left waiting on a source that has not produced anything yet. + this.switch.wait(cx); + if !this.switch.is_open() { + return Poll::Pending; + } + Poll::Ready(this.after.next().map(step_answer)) + } +} + +impl Drop for GatedScript { + fn drop(&mut self) { + self.dropped.fetch_add(1, Ordering::SeqCst); + } +} + +/// The tokio task that first polled each partition of a [`Recording`] child. +type ReadTasks = Arc>>>; + +/// A child whose partitions answer scripted steps, recording the task that first polls each one. +/// +/// A partition count says how many parts a plan declares, not how many tasks compute them: one task +/// polling every partition in turn declares the same number. What is recorded here is the task +/// itself, read from [`tokio::task::try_id`] inside the poll - which is where a child partition is +/// really read, and hence what the partition count actually bought. +struct Recording { + steps: Vec>, + polled: ReadTasks, + series_schema: Arc, + properties: Arc, +} + +impl Recording { + /// The child over `steps`, one partition per list, and the record of which task polled each. + fn plan(steps: Vec>) -> (Arc, ReadTasks) { + let polled: ReadTasks = Arc::new(Mutex::new(vec![None; steps.len()])); + let series_schema = schema(); + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(steps.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + let child = Arc::new(Self { + steps, + polled: Arc::clone(&polled), + series_schema, + properties, + }); + (child, polled) + } +} + +impl fmt::Debug for Recording { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Recording") + .field("partitions", &self.steps.len()) + .finish() + } +} + +impl DisplayAs for Recording { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Recording: partitions={}", self.steps.len()) + } +} + +impl ExecutionPlan for Recording { + fn name(&self) -> &str { + "Recording" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a recording child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let steps = self.steps.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("recording child has no partition {partition}")) + })?; + pack_evaluated_series( + Box::pin(RecordingScript { + partition, + steps: steps.into_iter(), + polled: Arc::clone(&self.polled), + recorded: false, + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// One partition of a [`Recording`] child. +struct RecordingScript { + partition: usize, + steps: std::vec::IntoIter, + polled: ReadTasks, + /// Whether this partition's first poll has been recorded. + recorded: bool, +} + +impl Stream for RecordingScript { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + // Recorded on the first poll rather than in `execute`: `execute` may run on the consumer's + // task while the reading is what another task does. + if !this.recorded { + this.recorded = true; + let mut polled = this + .polled + .lock() + .unwrap_or_else(|error| error.into_inner()); + polled[this.partition] = tokio::task::try_id(); + } + Poll::Ready(this.steps.next().map(step_answer)) + } +} + +/// Waits until `done` holds, so a test states a condition rather than a duration. +async fn wait_until(what: &str, done: impl Fn() -> bool) { + tokio::time::timeout(Duration::from_secs(10), async { + while !done() { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting until {what}")); +} + +/// The read a node started, for a test that has to look at the task behind the outputs. +fn execution(plan: &Arc) -> Arc { + let state = plan.shared.lock().unwrap(); + Arc::clone(state.as_ref().expect("the read was started by execute")) +} + +/// Whether every input task of the read has stopped. +/// +/// There is one task per input partition and no single producer, so "the read stopped" is what has to +/// be read off the tasks themselves: none of them is still running. +fn producer_stopped(execution: &SharedExecution) -> bool { + execution.running.lock().unwrap().inputs.live() == 0 +} + +/// The error `stream` reports after the batches it delivers before it, or `None` if it ends. +async fn terminal_error(stream: &mut SendableRecordBatchStream) -> Option { + loop { + match std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await { + None => return None, + Some(Ok(_)) => {} + Some(Err(error)) => return Some(error), + } + } +} + +async fn evaluated_terminal_error( + stream: &mut SendableEvaluatedSeriesStream, +) -> Option { + loop { + match std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await { + None => return None, + Some(Ok(_)) => {} + Some(Err(error)) => return Some(error), + } + } +} + +/// Everything `stream` hands over until it ends: the chunks it delivered and the errors it reported. +/// +/// Reading past the first error is what a test of the failure path needs. An error item is terminal, +/// so a stream that reports one and then answers again is what these tests are looking for - and so +/// is one that ends without ever reporting one, which is the clean end a failed read must not produce. +/// It also assumes the stream does end: a failure that leaves an output waiting for a chunk that will +/// never come is a hang, and the caller's timeout is what turns that into a failure. +async fn until_end( + stream: &mut SendableEvaluatedSeriesStream, +) -> (Vec<(Series, bool)>, Vec) { + let mut chunks = Vec::new(); + let mut errors = Vec::new(); + loop { + match std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await { + None => return (chunks, errors), + Some(Ok(chunk)) => chunks.push(chunk), + Some(Err(error)) => errors.push(error.to_string()), + } + } +} + +/// the read drives each input partition on a task of its own. +/// +/// The child records the task that first polls each of its partitions, and what is asserted is that +/// the two partitions were polled by two distinct tasks, neither of them the consumer's. That is the +/// property the node provides - one task per child partition - and it is what a partition count on +/// its own cannot say: one task polling every partition in turn declares the same count. +/// +/// Two distinct tasks is deliberately not a claim about simultaneous wall-clock execution: the two +/// tasks may share a worker thread and be polled one after the other, and the read is correct either +/// way. Nor does the assertion need two threads to be free at the same instant, which is what ties it +/// to the scheduling of the runtime the read runs on rather than to the read. What a read that drove +/// the partitions from one task would record instead is one id for both, or the consumer's own. +#[tokio::test] +async fn every_input_partition_is_read_on_a_task_of_its_own() { + let shared = labels("web-1"); + let (child, polled) = Recording::plan(vec![ + vec![Step::chunk(chunk(shared.clone(), 0, 1, true))], + vec![Step::chunk(chunk(shared.clone(), 4, 1, true))], + ]); + let plan: Arc = node(child, 1); + // Taken before anything is read: the task that asks for an output is the one a partition must + // not be read on. + let consumer = tokio::task::try_id(); + let outputs = tokio::time::timeout( + Duration::from_secs(10), + drain(&plan, &SessionContext::default().task_ctx()), + ) + .await + .expect("both partitions of the child answer, so the output ends"); + + let recorded = polled + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone(); + assert_eq!( + 2, + recorded.len(), + "the child holds one partition per scripted list" + ); + let tasks = recorded + .iter() + .map(|task| task.unwrap_or_else(|| panic!("a child partition was never polled"))) + .collect::>(); + assert_ne!( + tasks[0], tasks[1], + "the two partitions were polled by two different tasks, so neither was read on the task \ + that read the other: {tasks:?}" + ); + assert!( + tasks.iter().all(|task| Some(*task) != consumer), + "neither partition ran on the consumer's task, so the read stayed on the tasks the node \ + started for the child: {tasks:?} against {consumer:?}" + ); + + let mut chunks = outputs + .into_iter() + .next() + .expect("one output") + .expect("the output decodes"); + chunks.sort_by_key(|(series, _)| series.time.first_position()); + assert_output_contract(&chunks, &[(shared.clone(), vec![0, 4])]); + assert_eq!(1, chunks.iter().filter(|(_, closed)| *closed).count()); +} + +/// the streams own the read, so dropping the plan does not truncate them. +/// +/// The input produces nothing at all until the switch is opened, and the switch is opened only after +/// every plan reference is gone. Before that, an output is polled once and delivers nothing, which is +/// what makes the read that follows a read of a plan that no longer exists: if dropping the plan had +/// aborted the producer, the outputs would end here - the answer would be missing every chunk rather +/// than late. +#[tokio::test] +async fn dropping_the_plan_leaves_the_read_to_its_streams() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + let switch = Arc::new(Switch::default()); + let dropped = Arc::new(AtomicUsize::new(0)); + let child = Gated::plan( + vec![ + GatedPartition::held(vec![ + Step::chunk(chunk(shared.clone(), 0, 2, false)), + Step::chunk(chunk(shared.clone(), 2, 1, true)), + ]), + GatedPartition::held(vec![Step::chunk(chunk(other.clone(), 0, 1, true))]), + ], + Arc::clone(&switch), + Arc::clone(&dropped), + ); + let plan: Arc = node(child, 2); + let context = SessionContext::default().task_ctx(); + let mut streams = (0..2) + .map(|partition| output(&plan, partition, &context)) + .collect::>(); + // Every plan reference goes away here: DataFusion's `execute_stream_partitioned` hands out the + // streams and drops the plan, and this node has to keep reading for them. + drop(plan); + + assert!( + tokio::time::timeout(Duration::from_millis(200), streams[0].next()) + .await + .is_err(), + "the gated input produced before the plan was dropped, so this test proves nothing" + ); + switch.open(); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all( + streams + .into_iter() + .map(|stream| stream.try_collect::>()), + ), + ) + .await + .expect("the read outlives the plan it was executed from") + .into_iter() + .map(|output| output.expect("every output decodes")) + .collect::>(); + for chunks in &outputs { + assert_chains_are_whole(chunks); + } + + let piece = named(&shared); + assert_eq!( + vec![ + (piece.clone(), Some(0), 2, false), + (piece, Some(2), 1, true), + ], + summaries(&outputs[routing.output(&shared)]), + "both chunks of the chain that was held back arrive, in order" + ); + assert_eq!( + vec![(named(&other), Some(0), 1, true)], + summaries(&outputs[routing.output(&other)]), + "the other partition's chunk arrives too" + ); +} + +/// EOF makes every queued identity ready, but each poll hands off only one merged identity. +#[tokio::test] +async fn eof_readies_two_groups_but_poll_merges_one_at_a_time() { + let first_metric = labels("eof-ready-first"); + let second_metric = labels("eof-ready-second"); + let child = Scripted::plan(vec![vec![ + Step::chunk(chunk(first_metric.clone(), 0, 1, true)), + Step::chunk(chunk(second_metric.clone(), 0, 1, true)), + ]]); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let shared = + SharedExecution::start(&child, Routing::new(1), Arc::clone(&schema()), &context).unwrap(); + let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); + let mut stream = RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + ); + wait_until("both identities and EOF are queued", || { + producer_stopped(&shared) + }) + .await; + + let first = futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) + .await + .expect("first ready identity") + .expect("first identity succeeds"); + assert_eq!(first_metric, first.0.metric); + assert!(first.1); + assert!( + stream.output.is_empty(), + "no eagerly merged second result is buffered" + ); + + let second = futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) + .await + .expect("second ready identity") + .expect("second identity succeeds"); + assert_eq!(second_metric, second.0.metric); + assert!(second.1); + assert!(stream.groups.is_empty()); + drop(first); + drop(second); + assert!( + futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) + .await + .is_none() + ); + drop(stream); + assert_eq!(0, pool.reserved(), "finished stream releases all charges"); +} + +/// when every output is dropped the read stops, even though the plan is still alive. +/// +/// The plan is kept for the whole test, as a metrics collector holding a plan would keep it: what +/// stops the read is the outputs going away, not the plan being freed. Asserted are the three things +/// that has to release - the reading task, the child it was reading, and the memory the read +/// reserved. +#[tokio::test] +async fn dropping_every_output_stops_the_read_and_releases_it() { + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let switch = Arc::new(Switch::default()); + let dropped = Arc::new(AtomicUsize::new(0)); + let shared = labels("web-1"); + // More closed chains than the two outputs can queue, so the read cannot finish on its own. + let steps = (0..TEST_FIXTURE_VOLUME * 3) + .map(|index| Step::chunk(chunk(shared.clone(), index, 1, true))) + .collect::>(); + let child = Gated::plan( + vec![ + GatedPartition::open(steps.clone()), + GatedPartition::open(steps), + ], + Arc::clone(&switch), + Arc::clone(&dropped), + ); + let plan = node(child, 2); + let streams = (0..2) + .map(|partition| { + unpack_evaluated_series( + plan.execute(partition, Arc::clone(&context)) + .expect("the output executes"), + None, + ) + .expect("the output declares an evaluated frame") + }) + .collect::>(); + let execution = execution(&plan); + wait_until("the read reserved what it queued", || pool.reserved() > 0).await; + + drop(streams); + + wait_until("the read's reservation is released", || { + pool.reserved() == 0 + }) + .await; + wait_until("the child's streams are dropped", || { + dropped.load(Ordering::SeqCst) == 2 + }) + .await; + wait_until("the reading task stops", || producer_stopped(&execution)).await; + assert!( + plan.shared.lock().unwrap().is_some(), + "the plan is still alive, so it is the outputs that stopped the read" + ); +} + +/// a child failure is a terminal item on every output, and the end comes after it. +/// +/// The second half matters as much as the first: an output that reports the failure and then answers +/// again would be a stream whose error was not terminal, and one that answers with a chunk before the +/// error is already covered by the contents. What must never happen is a clean end in place of the +/// failure, so the error is required on every output. +#[tokio::test] +async fn a_child_failure_is_terminal_on_every_output() { + let shared = labels("web-1"); + let child = Scripted::plan(vec![ + vec![Step::chunk(chunk(shared.clone(), 0, 2, true)), Step::Error], + vec![Step::chunk(chunk(shared.clone(), 8, 1, true)), Step::Error], + ]); + let plan: Arc = node(child, 3); + let context = SessionContext::default().task_ctx(); + let mut streams = (0..3) + .map(|partition| { + plan.execute(partition, Arc::clone(&context)) + .expect("the output executes") + }) + .collect::>(); + + for (partition, stream) in streams.iter_mut().enumerate() { + let error = tokio::time::timeout(Duration::from_secs(10), terminal_error(stream)) + .await + .unwrap_or_else(|_| panic!("output partition {partition} never ended")) + .unwrap_or_else(|| { + panic!("output partition {partition} reported a clean end, not the failure") + }); + assert!( + error.to_string().contains("scripted child failure"), + "output partition {partition} ended with {error}" + ); + let after = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await; + assert!( + after.is_none(), + "output partition {partition} answered again after its terminal error" + ); + } +} + +/// an output that ends takes nothing else down with it when another output was never executed. +/// +/// This is DataFusion's ordinary order rather than an unusual one: `CoalescePartitionsExec` executes +/// and consumes one partition per task, and nothing makes every stream exist before the first one is +/// read. The input here is empty, so the read finishes while partition 0 alone is handed out, and +/// reading that partition to its end releases its lease - which is exactly the moment the read must +/// not be taken as over: partition 1 has not been asked for yet, and its queue is what a later +/// `execute` comes for. A read that stopped there would answer `execute(1)` with "executed twice" for +/// a partition that was never executed. +#[tokio::test] +async fn an_output_that_ends_does_not_close_the_queue_of_an_unclaimed_output() { + let plan: Arc = node(Scripted::plan(vec![Vec::new()]), 2); + let context = SessionContext::default().task_ctx(); + + let first = output(&plan, 0, &context); + let first_chunks = tokio::time::timeout(Duration::from_secs(10), first.try_collect::>()) + .await + .expect("an empty input still ends the first output") + .expect("the first output ends without a failure"); + assert!( + first_chunks.is_empty(), + "an empty input routed nothing, so the first output has nothing to hand over" + ); + + let second = unpack_evaluated_series( + plan.execute(1, Arc::clone(&context)) + .unwrap_or_else(|error| panic!("the second output is still claimable: {error}")), + None, + ) + .expect("the second output declares an evaluated frame"); + let second_chunks = + tokio::time::timeout(Duration::from_secs(10), second.try_collect::>()) + .await + .expect("the second output is executed after the first one ended") + .expect("the second output ends without a failure"); + assert!( + second_chunks.is_empty(), + "an empty input routed nothing, so the second output has nothing to hand over" + ); +} + +/// the output claimed last still receives its whole share after a claimed sibling ended. +/// +/// Partition 0 is read to its end before partition 1 is executed. Its released lease must not discard +/// the unclaimed sibling queue, whose complete result is verified after the later claim. +#[tokio::test] +async fn an_unclaimed_output_delivers_its_whole_share_after_a_sibling_ended() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + let ended_first = routing.output(&shared); + let claimed_later = routing.output(&other); + assert_ne!( + ended_first, claimed_later, + "the two label sets have to land apart for this test to mean anything" + ); + // One distinct label set per identity and output; chunks belonging to an identity form one chain. + let chains = TEST_FIXTURE_VOLUME - 1; + let identities = labels_routed_to(claimed_later, chains); + let steps = (0..chains) + .map(|index| Step::chunk(chunk(shared.clone(), index, 1, index + 1 == chains))) + .chain( + identities + .iter() + .enumerate() + .map(|(index, metric)| Step::chunk(chunk(metric.clone(), index, 1, true))), + ) + .collect::>(); + let plan: Arc = node(Scripted::plan(vec![steps]), 2); + let context = SessionContext::default().task_ctx(); + + let first = output(&plan, ended_first, &context); + let first_chunks = tokio::time::timeout(Duration::from_secs(10), first.try_collect::>()) + .await + .expect("the first output reaches its end before the other is executed") + .expect("the first output ends without a failure"); + assert_output_contract(&first_chunks, &[(shared.clone(), (0..chains).collect())]); + + let second = unpack_evaluated_series( + plan.execute(claimed_later, Arc::clone(&context)) + .unwrap_or_else(|error| panic!("the other output is still claimable: {error}")), + None, + ) + .expect("the other output declares an evaluated frame"); + let second_chunks = + tokio::time::timeout(Duration::from_secs(10), second.try_collect::>()) + .await + .expect("the output executed last reaches its end") + .expect("the output executed last ends without a failure"); + let expected = identities + .iter() + .enumerate() + .map(|(index, metric)| (metric.clone(), vec![index])) + .collect::>(); + assert_output_contract(&second_chunks, &expected); +} + +/// A claimed owner completes while the other owner has received no chain yet. After opening the gate, +/// the sibling's result remains available for a later claim. +#[tokio::test] +async fn an_unclaimed_output_keeps_the_producer_of_a_claimed_one_running() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + let read_first = routing.output(&shared); + let claimed_later = routing.output(&other); + assert_ne!( + read_first, claimed_later, + "the two label sets have to land apart for this test to mean anything" + ); + let chains = TEST_FIXTURE_VOLUME + 2; + let queued_identities = labels_routed_to(claimed_later, 2 * chains); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + GatedPartition { + before: vec![Step::chunk(chunk(shared.clone(), 0, 1, true))], + after: queued_identities[..chains] + .iter() + .enumerate() + .map(|(index, metric)| Step::chunk(chunk(metric.clone(), index, 1, true))) + .collect(), + switch: None, + }, + GatedPartition { + before: vec![Step::chunk(chunk(shared.clone(), 1, 1, true))], + after: queued_identities[chains..] + .iter() + .enumerate() + .map(|(index, metric)| Step::chunk(chunk(metric.clone(), index, 1, true))) + .collect(), + switch: None, + }, + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let concrete_plan = node(child, 2); + let plan: Arc = concrete_plan.clone(); + let context = SessionContext::default().task_ctx(); + let mut first = unpack_evaluated_series( + plan.execute(read_first, Arc::clone(&context)) + .expect("claimed output executes"), + None, + ) + .expect("the output declares an evaluated frame"); + let execution = execution(&concrete_plan); + + let mut first_chunks = Vec::new(); + loop { + let item = tokio::time::timeout(Duration::from_secs(10), first.next()) + .await + .expect("completed identity is emitted while sibling chain is gated") + .expect("output remains open until all inputs finish") + .expect("output succeeds"); + let done = item.1; + first_chunks.push(item); + if done { + break; + } + } + assert_output_contract(&first_chunks, &[(shared.clone(), vec![0, 1])]); + assert!(!switch.is_open(), "sibling input remains gated"); + assert!( + !producer_stopped(&execution), + "another input remains live despite the completed output" + ); + switch.open(); + wait_until("the read finishes after terminal EOF", || { + producer_stopped(&execution) + }) + .await; + + // The sibling has ended, but the unclaimed owner's entire queue remains available to claim. + let second = output(&plan, claimed_later, &context); + let second_chunks = + tokio::time::timeout(Duration::from_secs(10), second.try_collect::>()) + .await + .expect("the retained unclaimed owner completes") + .expect("the retained owner succeeds"); + let expected = queued_identities + .iter() + .enumerate() + .map(|(index, metric)| (metric.clone(), vec![index % chains])) + .collect::>(); + assert_output_contract(&second_chunks, &expected); +} + +/// A task context whose memory pool never refuses a reservation. +/// +/// The composition below reads further ahead than a bounded pool would allow, and a read that failed +/// on memory would end with an error instead of showing whether it finishes at all - which is the one +/// thing this test is about. +fn unbounded_context() -> Arc { + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(UnboundedMemoryPool::default()) as Arc) + .build() + .expect("the runtime environment builds"); + Arc::new(TaskContext::default().with_runtime(Arc::new(runtime))) +} + +/// `count` label sets of distinct series that a two-output node routes to `output`. +/// +/// Found with the node's own [`Routing`] over candidate labels rather than assumed: which output a +/// label set belongs on is a hash of it, and a label set hard-coded from watching one run would stop +/// being a label set of `output` the moment the key or the seed changed, leaving the test about +/// something else. Reading it off the same function the node reduces with is what keeps the premise +/// - every chunk lands on `output` - true by construction. +fn labels_routed_to(output: usize, count: usize) -> Vec { + let routing = Routing::new(2); + let found = (0..) + .filter_map(|index| { + let candidate = labels(&format!("many-{index}")); + (routing.output(&candidate) == output).then_some(candidate) + }) + .take(count) + .collect::>(); + assert_eq!( + count, + found.len(), + "two outputs split the candidates between them, so {count} of each are found" + ); + found +} + +/// two of these nodes composed do not deadlock while every output of the outer one is drained. +/// +/// The child produces `SERIES` single-point chains, each with a label set of its own, every one of +/// them hashing to output 1 of a two-output node - so the inner node's output 0 is empty and output 1 +/// carries all of them. Composed as `outer(inner(child))`, the outer node's input 0 is the inner +/// node's empty output, and what the read does with that is what this test states. +/// +/// The input has many more chains than the bounded output queues can hold, so the test needs the +/// empty inner output to close through independent input-task progress rather than by buffering the +/// child whole. Both outer outputs are drained concurrently, as required by the node. +/// +/// Both outputs of the outer node are read at once and each is drained to its end, so nothing here +/// can be blamed on an output nobody drained - which is the one requirement this node states of its +/// caller. The count that comes out is asserted against the count that went in, so a passing run is a +/// run that lost nothing. +/// +/// Independent input tasks let the outer node route the inner node's non-empty output while its +/// other task waits for the inner empty output to close, allowing that inner execution to finish. +#[tokio::test] +async fn composing_two_nodes_does_not_deadlock_when_the_inner_output_is_empty() { + const SERIES: usize = 4096; + let routed = labels_routed_to(1, SERIES); + let source = || { + Scripted::plan(vec![ + routed + .iter() + .map(|metric| Step::chunk(chunk(metric.clone(), 0, 1, true))) + .collect::>(), + ]) + }; + let context = unbounded_context(); + + // The same child over one node finishes, and finishes with everything on output 1: what stalls + // below is the composition, not the child and not the volume. + let alone: Arc = node(source(), 2); + let alone_outputs = tokio::time::timeout(Duration::from_secs(10), drain(&alone, &context)) + .await + .expect("one node over this child finishes"); + assert!( + alone_outputs[0] + .as_ref() + .expect("the first output decodes") + .is_empty(), + "every label set of this test hashes to output 1, so output 0 has nothing to carry" + ); + assert_eq!( + SERIES, + alone_outputs[1] + .as_ref() + .expect("the second output decodes") + .len(), + "one node delivers every series it was given" + ); + + let inner: Arc = node(source(), 2); + let plan: Arc = node(inner, 2); + + // Both outputs at once, each to its end, and each chunk counted where it arrives: the count is + // what a failure has to say when the read never gets to the assertion below. + let delivered = Arc::new([AtomicUsize::new(0), AtomicUsize::new(0)]); + let reads = (0..2) + .map(|partition| { + let plan = Arc::clone(&plan); + let context = Arc::clone(&context); + let delivered = Arc::clone(&delivered); + async move { + let mut stream = output(&plan, partition, &context); + let mut chunks = Vec::new(); + while let Some(chunk) = stream.try_next().await? { + delivered[partition].fetch_add(1, Ordering::Relaxed); + chunks.push(chunk); + } + Ok::<_, BoxedError>(chunks) + } + }) + .collect::>(); + let outputs = tokio::time::timeout(Duration::from_secs(10), futures::future::join_all(reads)) + .await + .unwrap_or_else(|_| { + panic!( + "two composed nodes never finished although both outputs of the outer one were \ + drained: {} of {SERIES} series arrived on output 0 and {} on output 1", + delivered[0].load(Ordering::Relaxed), + delivered[1].load(Ordering::Relaxed), + ) + }) + .into_iter() + .map(|chunks| chunks.expect("both outputs decode")) + .collect::>(); + + assert!( + outputs[0].is_empty(), + "the same label sets hash to output 1 of the outer node as of the inner one" + ); + for chunks in &outputs { + assert_chains_are_whole(chunks); + } + let mut names = outputs + .iter() + .flatten() + .map(|(series, is_last)| { + assert!(*is_last, "every chain of this test closes on its own"); + named(&series.metric) + }) + .collect::>(); + names.sort(); + names.dedup(); + assert_eq!( + SERIES, + names.len(), + "every series that went in comes out exactly once, none twice and none truncated" + ); + assert_eq!( + SERIES, + outputs.iter().map(Vec::len).sum::(), + "the composition delivers as many chunks as the child handed over" + ); +} + +#[tokio::test] +async fn owner_merges_across_inputs_and_sorts_late_positions_until_eof() { + let metric = labels("owner-order"); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + GatedPartition::open(vec![Step::chunk(chunk(metric.clone(), 10, 1, true))]), + GatedPartition::held(vec![Step::chunk(chunk(metric.clone(), 0, 1, true))]), + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan: Arc = node(child, 1); + let context = SessionContext::default().task_ctx(); + let mut stream = output(&plan, 0, &context); + assert!( + futures::future::poll_fn(|cx| match stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "the identity waits while input 1 has not reached its final chain" + ); + switch.open(); + let chunks = tokio::time::timeout(Duration::from_secs(10), stream.try_collect::>()) + .await + .expect("gated input completes") + .expect("owner succeeds"); + assert_output_contract(&chunks, &[(metric.clone(), vec![0, 10])]); + assert_eq!(1, chunks.iter().filter(|(_, is_last)| *is_last).count()); +} + +#[tokio::test] +async fn multiple_producers_and_owners_merge_exact_full_label_identities() { + let routing = Routing::new(2); + let on_zero = labels_routed_to(0, 1).remove(0); + let on_one = labels_routed_to(1, 1).remove(0); + // These identities must not collapse merely because rendering omits empty values or a name. + let empty_value = SeriesLabels::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("zone"), + value: Arc::from(""), + }], + ) + .unwrap(); + let no_empty_value = SeriesLabels::try_new(Some(Arc::from("cpu")), Vec::new()).unwrap(); + let absent_name = SeriesLabels::try_new(None, Vec::new()).unwrap(); + let present_name = SeriesLabels::try_new(Some(Arc::from("")), Vec::new()).unwrap(); + assert_ne!( + empty_value, no_empty_value, + "empty label presence is identity" + ); + assert_ne!(absent_name, present_name, "name presence is identity"); + let identities = [ + on_zero, + on_one, + empty_value, + no_empty_value, + absent_name, + present_name, + ]; + let mut partitions = vec![Vec::new(), Vec::new(), Vec::new()]; + for metric in identities.iter() { + for (partition, position) in [0, 10, 20].into_iter().enumerate() { + partitions[partition].push(Step::chunk(chunk(metric.clone(), position, 1, true))); + } + } + let plan: Arc = node(Scripted::plan(partitions), 2); + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + for (owner, output) in outputs.iter().enumerate() { + let chunks = output.as_ref().expect("each owner succeeds"); + let expected = identities + .iter() + .filter(|metric| routing.output(metric) == owner) + .map(|metric| (metric.clone(), vec![0, 10, 20])) + .collect::>(); + assert!(!expected.is_empty(), "each output receives a real identity"); + assert_output_contract(chunks, &expected); + assert_eq!( + expected.len(), + chunks.iter().filter(|(_, closed)| *closed).count() + ); + } +} + +#[tokio::test] +async fn error_after_early_public_chain_is_reported_once_and_releases_pool() { + let metric = labels_routed_to(0, 2).remove(0); + let switch = Arc::new(Switch::default()); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + // Both inputs contribute A before their gates; only input 0 fails after the gate opens. + let child = Gated::plan( + vec![ + GatedPartition { + before: vec![Step::chunk(chunk(metric.clone(), 0, 1, true))], + after: vec![Step::ErrorSaying("late gated failure")], + switch: None, + }, + GatedPartition { + before: vec![Step::chunk(chunk(metric.clone(), 1, 1, true))], + after: Vec::new(), + switch: None, + }, + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan = node(child, 2); + let mut stream = output_of(&plan, 0, &context); + let mut emitted = Vec::new(); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + match futures::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await { + Some(Ok(chunk)) => { + let done = chunk.1; + emitted.push(chunk); + if done { + break; + } + } + Some(Err(error)) => panic!("A failed before publication: {error}"), + None => panic!("A ended before publication"), + } + } + }) + .await + .expect("A is published before the late failure"); + assert_output_contract(&emitted, &[(metric, vec![0, 1])]); + switch.open(); + let error = evaluated_terminal_error(&mut stream) + .await + .expect("late failure is reported"); + assert!(error.to_string().contains("late gated failure")); + assert!( + evaluated_terminal_error(&mut stream).await.is_none(), + "failure is reported once" + ); + assert_eq!(0, pool.reserved(), "failure releases all retained charges"); +} + +#[tokio::test] +async fn late_unclaimed_output_receives_first_merge_error_once_after_early_output_fails() { + let metric = labels_routed_to(0, 1).remove(0); + let other = labels_routed_to(1, 1).remove(0); + let conflicting = |value| { + let time = TimeAxis::from_grid_spans(grid(), vec![GridSpan::new(0, 1)]).unwrap(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![value]))); + Step::Chunk( + Series::try_new(&schema(), metric.clone(), time, fields).unwrap(), + true, + ) + }; + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let child = Scripted::plan(vec![ + vec![ + Step::chunk(chunk(other.clone(), 0, 1, true)), + conflicting(1.0), + ], + vec![conflicting(2.0)], + ]); + let plan = node(child, 2); + let mut early = output_of(&plan, 0, &context); + let mut emitted = Vec::new(); + let error = tokio::time::timeout(Duration::from_secs(10), async { + loop { + match futures::future::poll_fn(|cx| early.as_mut().poll_next(cx)).await { + None => panic!("conflicting merge ended without its error"), + Some(Ok(chunk)) => emitted.push(chunk), + Some(Err(error)) => break error.to_string(), + } + } + }) + .await + .expect("failure terminates the claimed output"); + assert!( + error.contains("cannot contain metrics with the same labelset"), + "actual merge error preserved: {error}" + ); + assert!( + emitted.is_empty(), + "the conflicting identity is never emitted" + ); + assert_eq!( + 0, + pool.reserved(), + "failure clears queued receiver data immediately" + ); + assert!( + futures::future::poll_fn(|cx| early.as_mut().poll_next(cx)) + .await + .is_none(), + "error is terminal" + ); + + let mut late = plan + .execute(1, Arc::clone(&context)) + .expect("late output executes"); + let late_error = terminal_error(&mut late) + .await + .expect("late owner gets failure"); + assert_eq!(error, late_error.to_string(), "first error is retained"); + assert!( + terminal_error(&mut late).await.is_none(), + "late output reports it once" + ); + assert_eq!(0, pool.reserved(), "late error leaves no retained charges"); +} + +#[tokio::test] +async fn owner_merges_complementary_fields_across_input_positions() { + let metric = labels("complementary-fields"); + let schema = Arc::new(SeriesSchema { + fields: Arc::from([ + FieldSpec { + name: Arc::from("left"), + }, + FieldSpec { + name: Arc::from("right"), + }, + ]), + eval_grid: Some(grid()), + }); + let make = |position, left, right| { + let time = TimeAxis::from_grid_spans( + schema.eval_grid.clone().unwrap(), + vec![GridSpan::new(position, 1)], + ) + .unwrap(); + let fields = SeriesFields::try_from_vec(vec![ + Samples::from_float(Float64Array::from(vec![left])), + Samples::from_float(Float64Array::from(vec![right])), + ]) + .unwrap(); + Step::Chunk( + Series::try_new(&schema, metric.clone(), time, fields).unwrap(), + true, + ) + }; + let switch = Arc::new(Switch::default()); + let child = Gated::plan_with_schema( + vec![ + GatedPartition::open(vec![make(0, Some(10.0), None)]), + GatedPartition { + before: Vec::new(), + after: vec![make(10, None, Some(20.0))], + switch: None, + }, + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + Arc::clone(&schema), + ); + let plan: Arc = node(child, 1); + let mut output = output(&plan, 0, &SessionContext::default().task_ctx()); + match tokio::time::timeout(Duration::from_millis(50), output.next()).await { + Err(_) => {} + Ok(Some(Ok((series, last)))) => panic!( + "the identity must wait while one input has not reached its final chain, got {} last={last}", + named(&series.metric) + ), + Ok(Some(Err(error))) => panic!("the output failed instead of waiting: {error}"), + Ok(None) => panic!("the output ended before the gated input contributed"), + } + switch.open(); + let chunks = output + .try_collect::>() + .await + .expect("merged output succeeds"); + assert_output_contract(&chunks, &[(metric.clone(), vec![0, 10])]); + assert_eq!(1, chunks.iter().filter(|(_, is_last)| *is_last).count()); + let series = chunks.iter().map(|(series, _)| series).collect::>(); + assert_eq!(2, series[0].fields.len()); + for (index, expected) in [10.0, 20.0].into_iter().enumerate() { + let values = series + .iter() + .flat_map(|series| series.fields.get(index).unwrap().iter()) + .filter_map(|sample| match sample { + datatypes::series::samples::SampleRef::Float(value) => Some(value), + _ => None, + }) + .collect::>(); + assert_eq!(1, values.len()); + assert_eq!(expected, values[0]); + } +} + +#[tokio::test] +async fn owner_rejects_conflicting_values_at_same_identity_and_position() { + let metric = labels("owner-conflict"); + let first = chunk(metric.clone(), 0, 1, true).0; + let schema = schema(); + let time = + TimeAxis::from_grid_spans(schema.eval_grid.clone().unwrap(), vec![GridSpan::new(0, 1)]) + .unwrap(); + let conflicting = Series::try_new( + &schema, + metric.clone(), + time, + SeriesFields::one(Samples::from_float(Float64Array::from(vec![99.0]))), + ) + .unwrap(); + let child = Scripted::plan(vec![ + vec![Step::Chunk(first, true)], + vec![Step::Chunk(conflicting, true)], + ]); + let plan: Arc = node(child, 1); + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + assert!( + outputs[0].is_err(), + "conflicting populated cell is an error" + ); +} + +/// Same-identity input chunks are collected and emitted as one closed series. +#[tokio::test] +async fn same_canonical_key_chains_merge_into_one_series() { + let metric = labels("same"); + let child = Scripted::plan(vec![ + vec![Step::chunk(chunk(metric.clone(), 0, 2, true))], + vec![Step::chunk(chunk(metric.clone(), 8, 2, true))], + ]); + let plan: Arc = node(child, 2); + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + let owner = Routing::new(2).output(&metric); + let chunks = outputs[owner].as_ref().expect("the owner output decodes"); + assert_output_contract(chunks, &[(metric.clone(), vec![0, 1, 8, 9])]); + assert_eq!(1, chunks.iter().filter(|(_, closed)| *closed).count()); + assert!(outputs[1 - owner].as_ref().unwrap().is_empty()); +} + +#[tokio::test] +async fn canonical_label_key_variants_route_consistently_without_merging_chains() { + let base = SeriesLabels::try_new(Some(Arc::from("cpu")), Vec::new()).unwrap(); + let metric_name_change = SeriesLabels::try_new(Some(Arc::from("memory")), Vec::new()).unwrap(); + let extra_tag = labels("host-a"); + let empty_tag = SeriesLabels::try_new( + Some(Arc::from("cpu")), + vec![Label { + name: Arc::from("host"), + value: Arc::from(""), + }], + ) + .unwrap(); + let absent_tag = base.clone(); + let no_name = SeriesLabels::try_new(None, Vec::new()).unwrap(); + let empty_name = SeriesLabels::try_new(Some(Arc::from("")), Vec::new()).unwrap(); + + let cases = [ + ( + "metric name is part of key", + base.clone(), + metric_name_change, + false, + ), + ( + "additional tag is part of key", + base.clone(), + extra_tag, + false, + ), + ("empty tag value is absent", empty_tag, absent_tag, true), + ( + "missing and empty metric names canonicalize alike", + no_name, + empty_name, + true, + ), + ]; + for (description, left, right, expected_same_key) in cases { + let left_key = label_set(&left); + let right_key = label_set(&right); + assert_eq!( + expected_same_key, + left_key == right_key, + "canonical keys follow the expected label semantics: {description}" + ); + if expected_same_key { + assert_eq!( + Routing::new(5).output(&left), + Routing::new(5).output(&right), + "equal canonical keys must share an owner: {description}" + ); + assert_eq!( + Routing::new(3).output(&left), + Routing::new(3).output(&right), + "canonical equivalence holds for any output count: {description}" + ); + } + + let child = Scripted::plan(vec![ + vec![Step::chunk(chunk(left.clone(), 0, 1, true))], + vec![Step::chunk(chunk(right.clone(), 4, 1, true))], + ]); + let plan: Arc = node(child, 5); + let outputs = drain(&plan, &SessionContext::default().task_ctx()).await; + let seen = outputs + .iter() + .enumerate() + .flat_map(|(partition, output)| { + output + .as_ref() + .expect("output decodes") + .iter() + .map(move |(series, is_last)| { + let owner = Routing::new(5).output(&series.metric); + assert_eq!( + partition, owner, + "each returned chain is on its computed owner: {description}" + ); + (label_set(&series.metric), *is_last) + }) + }) + .collect::>(); + assert_eq!( + 2, + seen.len(), + "both independent closed chains survive: {description}" + ); + assert!(seen.iter().all(|(_, is_last)| *is_last)); + let mut expected_keys = vec![left_key, right_key]; + expected_keys.sort(); + let mut observed_keys = seen.into_iter().map(|(key, _)| key).collect::>(); + observed_keys.sort(); + assert_eq!( + expected_keys, observed_keys, + "each chain's canonical key remains observable: {description}" + ); + } +} + +/// two sources whose chains belong on one output land there whole, never inside one another. +/// +/// Both label sets hash to the same output and both sources have a chain for it. One of the chains is +/// held open across a gate, so the other source's chain is the only thing that can be delivered; what +/// must not happen is that the held chain's first chunk sits on the output ahead of the other source's +/// chain, because two sources' chunks inside one open series is exactly what a downstream frame reader +/// refuses. What does happen is that the closed chain arrives first, whole, and the held one follows. +#[tokio::test] +async fn two_sources_keep_their_chains_whole_on_one_output() { + let on_one = labels_routed_to(0, 2); + let (held, closed) = (on_one[0].clone(), on_one[1].clone()); + assert_eq!( + Routing::new(2).output(&held), + Routing::new(2).output(&closed), + "two label sets of one output, found rather than assumed" + ); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + // The chain that stays open across the gate: its first chunk is handed over at once and + // the chunk that closes it only after the switch opens. + GatedPartition { + before: vec![ + Step::chunk(chunk(closed.clone(), 0, 2, true)), + Step::chunk(chunk(held.clone(), 0, 2, false)), + ], + after: vec![Step::chunk(chunk(held.clone(), 2, 2, true))], + switch: None, + }, + // The other source's final complementary chunk is handed over as soon as it is polled. + GatedPartition { + before: vec![Step::chunk(chunk(closed.clone(), 2, 2, true))], + after: Vec::new(), + switch: None, + }, + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan: Arc = node(child, 2); + let context = SessionContext::default().task_ctx(); + let mut routed = output(&plan, 0, &context); + let other = output(&plan, 1, &context); + + // The owner publishes the other identity even while the first input's chain remains open. + let mut chunks = Vec::new(); + loop { + let chunk = tokio::time::timeout(Duration::from_secs(10), routed.next()) + .await + .expect("closed identity emits before gated input completes") + .expect("closed identity is present") + .expect("owner succeeds"); + assert_eq!(closed, chunk.0.metric); + let done = chunk.1; + chunks.push(chunk); + if done { + break; + } + } + assert!( + futures::future::poll_fn(|cx| match routed.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "held identity cannot emit before its chain closes" + ); + switch.open(); + chunks.extend( + tokio::time::timeout(Duration::from_secs(10), routed.try_collect::>()) + .await + .expect("held chain emits after gate opens") + .expect("owner succeeds"), + ); + assert_chains_are_whole(&chunks); + assert_output_contract( + &chunks, + &[ + (closed.clone(), vec![0, 1, 2, 3]), + (held.clone(), vec![0, 1, 2, 3]), + ], + ); + let empty = tokio::time::timeout(Duration::from_secs(10), other.try_collect::>()) + .await + .expect("the other output ends") + .expect("the other output decodes"); + assert!( + empty.is_empty(), + "neither label set hashes to the other output" + ); +} + +/// a closing chunk that carries no rows is delivered, and is not mistaken for the end of the +/// stream. +/// +/// A series can close with nothing left of it, and that chunk carries the finality and no position. It +/// is a chunk of its own on the way out - it is how a consumer closes the series - and a stream that is +/// still going must not be ended by it: the chunk after it is what says the empty one was a chunk and +/// not the end. +#[tokio::test] +async fn an_empty_closing_chunk_is_a_chunk_and_not_the_end() { + let routing = Routing::new(2); + let shared = labels("web-1"); + let other = labels_elsewhere(&routing, &shared); + let child = Scripted::plan(vec![ + // One source closes a series with nothing left of it, and keeps going afterwards. + vec![ + Step::chunk(chunk(other.clone(), 0, 0, true)), + Step::chunk(chunk(labels("after-empty-close"), 0, 1, true)), + ], + // The other source hands over a chain with positions, and closes it with an empty chunk. + vec![ + Step::chunk(chunk(labels("empty-before-data"), 0, 0, true)), + Step::chunk(chunk(shared.clone(), 4, 2, false)), + Step::chunk(chunk(shared.clone(), 6, 0, true)), + ], + ]); + let plan: Arc = node(child, 2); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + drain(&plan, &SessionContext::default().task_ctx()), + ) + .await + .expect("an input that only closes empty chains still ends") + .into_iter() + .map(|output| output.expect("no output ends with an error")) + .collect::>(); + for chunks in &outputs { + assert_chains_are_whole(chunks); + } + + let fixture = [ + (shared.clone(), vec![4, 5]), + (other.clone(), vec![]), + (labels("empty-before-data"), vec![]), + (labels("after-empty-close"), vec![0]), + ]; + for (owner, chunks) in outputs.iter().enumerate() { + let expected = fixture + .iter() + .filter(|(metric, _)| routing.output(metric) == owner) + .cloned() + .collect::>(); + assert_output_contract(chunks, &expected); + } +} + +/// while one owner's chain is still open, the other owner's output still delivers. +/// +/// Two label sets on two outputs, one long chain each, and the first one is held open across a gate. +/// The other output has to deliver its whole chain in that time: the wait on the gated source belongs +/// to the task reading that source, and no other source's chunks wait with it. What this rules out is +/// the shape this node used to have - one reader for every input, parked on the gated source and +/// delivering nothing at all until the chain closed. +#[tokio::test] +async fn an_open_chain_on_one_output_does_not_hold_the_other_output_back() { + let routing = Routing::new(2); + let held = labels("web-1"); + let free = labels_elsewhere(&routing, &held); + let held_output = routing.output(&held); + let free_output = routing.output(&free); + assert_ne!( + held_output, free_output, + "the two label sets have to land apart for this test to mean anything" + ); + let switch = Arc::new(Switch::default()); + let child = Gated::plan( + vec![ + GatedPartition::open(vec![ + Step::chunk(chunk(free.clone(), 0, 2, false)), + Step::chunk(chunk(free.clone(), 2, 2, true)), + ]), + GatedPartition::open(vec![ + Step::chunk(chunk(free.clone(), 4, 2, false)), + Step::chunk(chunk(free.clone(), 6, 2, true)), + ]), + // A long chain, gated after its first chunk. + GatedPartition { + before: vec![ + Step::chunk(chunk(free.clone(), 8, 2, false)), + Step::chunk(chunk(free.clone(), 10, 2, true)), + Step::chunk(chunk(held.clone(), 0, 2, false)), + ], + after: vec![ + Step::chunk(chunk(held.clone(), 2, 2, false)), + Step::chunk(chunk(held.clone(), 4, 2, false)), + Step::chunk(chunk(held.clone(), 6, 2, true)), + ], + switch: None, + }, + ], + Arc::clone(&switch), + Arc::new(AtomicUsize::new(0)), + ); + let plan: Arc = node(child, 2); + let context = SessionContext::default().task_ctx(); + let mut held_stream = output(&plan, held_output, &context); + let mut free_stream = output(&plan, free_output, &context); + + // All contributors to the free identity close it before the separate held chain gates. + let mut free_chunks = Vec::new(); + loop { + let chunk = tokio::time::timeout(Duration::from_secs(10), free_stream.next()) + .await + .expect("completed free identity emits before gate opens") + .expect("free identity is present") + .expect("free owner succeeds"); + assert_eq!(free, chunk.0.metric); + let done = chunk.1; + free_chunks.push(chunk); + if done { + break; + } + } + assert_output_contract( + &free_chunks, + &[(free.clone(), vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])], + ); + assert!( + futures::future::poll_fn(|cx| match held_stream.as_mut().poll_next(cx) { + Poll::Pending => Poll::Ready(true), + Poll::Ready(_) => Poll::Ready(false), + }) + .await, + "the gated chain cannot complete before the switch opens" + ); + switch.open(); + let free_rest = tokio::time::timeout( + Duration::from_secs(10), + (&mut free_stream).try_collect::>(), + ) + .await + .expect("free output ends after remaining input closes") + .expect("free owner succeeds"); + assert!( + free_rest.is_empty(), + "completed identity is not emitted twice" + ); + + switch.open(); + let held_chunks = + tokio::time::timeout(Duration::from_secs(10), held_stream.try_collect::>()) + .await + .expect("the held chain arrives once the gate opens") + .expect("the held output decodes"); + assert_chains_are_whole(&held_chunks); + assert_eq!( + vec![ + (named(&held), Some(0), 2, false), + (named(&held), Some(2), 2, false), + (named(&held), Some(4), 2, false), + (named(&held), Some(6), 2, true), + ], + summaries(&held_chunks), + "the long chain arrives whole and in order once its source closes it" + ); + let rest = tokio::time::timeout( + Duration::from_secs(10), + (&mut free_stream).try_collect::>(), + ) + .await + .expect("the other output ends") + .expect("the other output decodes"); + assert!( + rest.is_empty(), + "the other output carried its whole chain already" + ); +} + +/// a partition that fails mid-stream ends every output with that error, none of them cleanly. +/// +/// One of the two input partitions fails after handing over a chunk and the other ends normally, so +/// one output has data and an end of its own that the failure has to replace. Three outputs make sure +/// the error is not merely the fate of the output the failed partition was routing to: an output that +/// ended cleanly would report the series that never arrived as series that do not exist. +#[tokio::test] +async fn a_child_partition_that_fails_mid_stream_ends_every_output_with_its_error() { + let shared = labels("web-1"); + let other = labels("elsewhere"); + let child = Scripted::plan(vec![ + // This partition fails with its series open. + vec![Step::chunk(chunk(shared.clone(), 0, 2, false)), Step::Error], + // This one ends normally. + vec![Step::chunk(chunk(other.clone(), 0, 1, true))], + ]); + let plan: Arc = node(child, 3); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + drain(&plan, &SessionContext::default().task_ctx()), + ) + .await + .expect("a failed read still ends every output"); + assert_eq!(3, outputs.len()); + for (partition, output) in outputs.iter().enumerate() { + let error = output + .as_ref() + .err() + .unwrap_or_else(|| panic!("output partition {partition} reported a clean end")); + assert!( + error.to_string().contains("scripted child failure"), + "output partition {partition} ended with {error}" + ); + } +} + +/// two inputs make their own routing decisions, and neither waits for the other to make one. +/// +/// The child will not hand over a second chunk of a partition until every partition has handed over a +/// first, so nothing gets past the first chunk of either source unless both are being read at the same +/// time: a read that took one input's chain to its end before it looked at the other waits forever +/// there for a chunk that only the other input's own progress can release. Each source's chain is two +/// chunks long, so that wait lands in the middle of an open chain - the point at which a +/// chain-continuous reader that drove the inputs itself would have to stop. +/// +/// The two label sets hash to two different outputs, so what each task does here is a routing decision +/// of its own: one hash over one chain, sent to one owner. Two outputs means each of them has a +/// different decision to make and no decision of the other's to wait for, which is the property this +/// test states; a node with a single output is [`one_output_keeps_every_chunk_and_its_chains_whole`], +/// where there is nothing to decide and every chunk still arrives. +/// +/// The gate is this child and not a stopwatch: nothing here depends on two cores happening to run at +/// the same instant, only on the two inputs being driven, and routed, independently of each other. +#[tokio::test] +async fn two_inputs_are_routed_without_waiting_for_each_other() { + let routing = Routing::new(2); + let first = labels("web-1"); + let second = labels_elsewhere(&routing, &first); + assert_ne!( + routing.output(&first), + routing.output(&second), + "the two label sets have to land apart for this test to mean anything" + ); + let child = Together::plan(vec![ + vec![ + Step::chunk(chunk(first.clone(), 0, 1, false)), + Step::chunk(chunk(first.clone(), 1, 1, true)), + ], + vec![ + Step::chunk(chunk(second.clone(), 0, 1, false)), + Step::chunk(chunk(second.clone(), 1, 1, true)), + ], + ]); + let plan: Arc = node(child, 2); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + drain(&plan, &SessionContext::default().task_ctx()), + ) + .await + .expect("both inputs are read at once, so the rendezvous between them opens") + .into_iter() + .map(|output| output.expect("no output ends with an error")) + .collect::>(); + assert_eq!(2, outputs.len()); + for chunks in &outputs { + assert_chains_are_whole(chunks); + } + assert_eq!( + vec![ + (named(&first), Some(0), 1, false), + (named(&first), Some(1), 1, true), + ], + summaries(&outputs[routing.output(&first)]), + "the first input's chain is routed, whole and in order, to the output its own labels name" + ); + assert_eq!( + vec![ + (named(&second), Some(0), 1, false), + (named(&second), Some(1), 1, true), + ], + summaries(&outputs[routing.output(&second)]), + "and the second input's chain to its own output, neither decision waiting on the other" + ); +} + +/// The chunks each partition of a [`Together`] child has handed over so far. +#[derive(Default)] +struct FirstChunks { + /// Whether each partition handed over its first chunk, by partition. + handed: Mutex>, + /// The streams parked until the partitions that have not handed theirs over do. + wakers: Mutex>, +} + +impl FirstChunks { + fn new(partitions: usize) -> Arc { + Arc::new(Self { + handed: Mutex::new(vec![false; partitions]), + wakers: Mutex::new(Vec::new()), + }) + } + + /// Records that `partition` handed over its first chunk, and wakes what waited for it. + fn handed(&self, partition: usize) { + self.handed + .lock() + .unwrap_or_else(|error| error.into_inner())[partition] = true; + for waker in std::mem::take( + &mut *self + .wakers + .lock() + .unwrap_or_else(|error| error.into_inner()), + ) { + waker.wake(); + } + } + + /// Whether every partition handed over its first chunk. + fn all(&self) -> bool { + self.handed + .lock() + .unwrap_or_else(|error| error.into_inner()) + .iter() + .all(|handed| *handed) + } + + /// Parks `cx` until every partition has, unless they already have. + fn wait(&self, cx: &Context<'_>) { + if self.all() { + return; + } + self.wakers + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(cx.waker().clone()); + if self.all() { + // The last partition handed its chunk over between the two reads, so its wake did not + // find this waker. + cx.waker().wake_by_ref(); + } + } +} + +/// A child each of whose partitions hands over a second chunk only once every partition handed over a +/// first. +/// +/// Progress through one partition's chain requires every partition to have progressed too. This is +/// a rendezvous between inputs, not a stopwatch: nothing depends on two threads running at once. +struct Together { + steps: Vec>, + chunks: Arc, + series_schema: Arc, + properties: Arc, +} + +impl Together { + /// The child over `steps`, one partition per list. + fn plan(steps: Vec>) -> Arc { + let series_schema = schema(); + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(steps.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Arc::new(Self { + chunks: FirstChunks::new(steps.len()), + steps, + series_schema, + properties, + }) + } +} + +impl fmt::Debug for Together { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Together") + .field("partitions", &self.steps.len()) + .finish() + } +} + +impl DisplayAs for Together { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Together: partitions={}", self.steps.len()) + } +} + +impl ExecutionPlan for Together { + fn name(&self) -> &str { + "Together" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a together child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let steps = self.steps.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("together child has no partition {partition}")) + })?; + pack_evaluated_series( + Box::pin(TogetherScript { + partition, + steps: steps.into_iter(), + chunks: Arc::clone(&self.chunks), + handed: false, + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// One partition of a [`Together`] child. +struct TogetherScript { + partition: usize, + steps: std::vec::IntoIter, + chunks: Arc, + /// Whether this partition handed over its first chunk already. + handed: bool, +} + +impl Stream for TogetherScript { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if this.handed { + // Everything after the first chunk waits for the other partitions to have handed theirs + // over, so this stream cannot be read past its first chunk on its own. + this.chunks.wait(cx); + if !this.chunks.all() { + return Poll::Pending; + } + } + match this.steps.next() { + Some(step) => { + if !this.handed { + this.handed = true; + this.chunks.handed(this.partition); + } + Poll::Ready(Some(step_answer(step))) + } + None => Poll::Ready(None), + } + } +} + +/// A child whose partitions hand over their `before` steps, wait until every partition got that far, +/// and only then hand over their `after` steps. +/// +/// The barrier sits between the two halves rather than before a partition's first chunk, which is what +/// [`Together`] gates: what a test of the failure path needs is two partitions that fail *at once*. +/// Every partition that gets to the gate is parked until the last one arrives, so the `after` steps of +/// two partitions are handed over in the same wake - which is the interleaving +/// [`two_failures_hand_every_output_one_error`] is built on. The gate itself is [`FirstChunks`], the +/// all-arrived barrier [`Together`] already uses, since "has every partition got here" is what both +/// need. +struct Simultaneous { + partitions: Vec, + arrived: Arc, + series_schema: Arc, + properties: Arc, +} + +impl Simultaneous { + /// The child over `partitions`, and the barrier every one of them waits at. + fn plan(partitions: Vec) -> (Arc, Arc) { + let series_schema = schema(); + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + let arrived = FirstChunks::new(partitions.len()); + let child = Arc::new(Self { + partitions, + arrived: Arc::clone(&arrived), + series_schema, + properties, + }); + (child, arrived) + } +} + +impl fmt::Debug for Simultaneous { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Simultaneous") + .field("partitions", &self.partitions.len()) + .finish() + } +} + +impl DisplayAs for Simultaneous { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Simultaneous: partitions={}", self.partitions.len()) + } +} + +impl ExecutionPlan for Simultaneous { + fn name(&self) -> &str { + "Simultaneous" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a simultaneous child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let steps = self.partitions.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("simultaneous child has no partition {partition}")) + })?; + pack_evaluated_series( + Box::pin(SimultaneousScript { + partition, + before: steps.before.into_iter(), + after: steps.after.into_iter(), + arrived: Arc::clone(&self.arrived), + arrived_here: false, + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// One partition of a [`Simultaneous`] child. +struct SimultaneousScript { + partition: usize, + before: std::vec::IntoIter, + after: std::vec::IntoIter, + arrived: Arc, + /// Whether this partition already reached the barrier. + arrived_here: bool, +} + +impl Stream for SimultaneousScript { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + if let Some(step) = this.before.next() { + return Poll::Ready(Some(step_answer(step))); + } + if !this.arrived_here { + // Everything after the first half waits for every partition to have handed its first half + // over, so the second halves - and the failures among them - leave two tasks together. + this.arrived_here = true; + this.arrived.handed(this.partition); + } + this.arrived.wait(cx); + if !this.arrived.all() { + return Poll::Pending; + } + Poll::Ready(this.after.next().map(step_answer)) + } +} + +/// Every output of `plan`, read to its end at once: the chunks each delivered and the errors each +/// reported. +/// +/// All partitions are drained concurrently, which is what an executor does and what a failed read +/// needs: an input suspends on an output nobody is draining, so an output drained to its end before +/// the next one is started would wait for a sender that is waiting for that next output. +async fn every_output_until_end( + plan: &Arc, + context: &Arc, +) -> Vec<(Vec<(Series, bool)>, Vec)> { + let mut streams = Vec::new(); + for partition in 0..plan.outputs() { + streams.push( + unpack_evaluated_series( + plan.execute(partition, Arc::clone(context)) + .expect("the output executes"), + None, + ) + .expect("the output declares an evaluated frame"), + ); + } + futures::future::join_all(streams.iter_mut().map(until_end)).await +} + +#[tokio::test] +async fn coalesce_local_limit_reads_repartitioned_outputs_without_sibling_error() { + let first = labels_routed_to(0, 1).remove(0); + let second = labels_routed_to(1, 1).remove(0); + let steps = labels_routed_to(0, 3 * TEST_FIXTURE_VOLUME) + .into_iter() + .map(|metric| Step::chunk(chunk(metric, 0, 1, true))) + .chain(std::iter::once(Step::chunk(chunk( + second.clone(), + 0, + 1, + true, + )))) + .collect(); + let routed: Arc = node( + Fragile::plan( + vec![Behaviour::OneAtATime(steps)], + Arc::new(Switch::default()), + ) + .plan, + 2, + ); + let limited: Arc = Arc::new(LocalLimitExec::new(routed, 1)); + let plan = CoalescePartitionsExec::new(limited); + let context = SessionContext::default().task_ctx(); + let stream = plan.execute(0, context).expect("coalesce executes"); + let mut output = unpack_evaluated_series(stream, None).expect("evaluated frame"); + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .expect("coalesce and local limit make progress"); + assert!(errors.is_empty(), "no sibling output failure: {errors:?}"); + assert_eq!( + 2, + chunks.len(), + "LocalLimit fetch is applied per output partition" + ); + let mut actual = chunks + .iter() + .map(|(series, is_last)| { + assert_eq!( + 1, + series.time.len(), + "fetch=1 returns one row per partition" + ); + assert!(*is_last, "each returned series is closed"); + series.metric.clone() + }) + .collect::>(); + actual.sort_by_key(|metric| metric.tags()[0].value.clone()); + let mut expected = vec![first, second]; + expected.sort_by_key(|metric| metric.tags()[0].value.clone()); + assert_eq!( + expected, actual, + "both routed outputs contribute their expected series" + ); +} + +#[test] +fn ensure_cooperative_preserves_direct_repartition_child() { + let repartition: Arc = node(Scripted::plan(vec![Vec::new()]), 2); + let optimized = EnsureCooperative::new() + .optimize(repartition, &ConfigOptions::default()) + .expect("EnsureCooperative optimizes the plan"); + assert!( + optimized.downcast_ref::().is_some(), + "EnsureCooperative preserves the bare repartition root" + ); +} +/// Regression: a paced producer fills output 0 with closed chains, then fails while output 1 is drained first. +#[tokio::test] +async fn failure_after_filling_another_output_does_not_hang_early_output_drain() { + const MESSAGE: &str = "failure after output 0 filled"; + let mut steps = labels_routed_to(0, TEST_FIXTURE_VOLUME - 1) + .into_iter() + .map(|metric| Step::chunk(chunk(metric, 0, 1, true))) + .collect::>(); + steps.extend( + labels_routed_to(1, 1) + .into_iter() + .map(|metric| Step::chunk(chunk(metric, 0, 1, true))), + ); + steps.push(Step::ErrorSaying(MESSAGE)); + let plan = node( + Fragile::plan( + vec![Behaviour::OneAtATime(steps)], + Arc::new(Switch::default()), + ) + .plan, + 2, + ); + let context = SessionContext::default().task_ctx(); + let mut output1 = output_of(&plan, 1, &context); + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output1)) + .await + .expect("output 1 sees the failure despite output 0 being full"); + let _ = chunks; + assert_eq!(1, errors.len()); + assert!(errors[0].contains(MESSAGE)); +} + +/// two partitions that fail at the same moment leave every output exactly one error. +/// +/// Between them the two partitions hand over [`TEST_FIXTURE_VOLUME`] chains per output before they fail, +/// so both queues are exactly full and then two failures land on a read whose outputs are not being +/// drained. A failing task cancels its peers while another input may be suspended on a full output; +/// every output must still learn the elected failure from the shared ending record when senders close. +/// +/// Both partitions are known to have reached the barrier before either failed, which is what makes +/// the two failures overlap rather than merely possible: neither can hand its failing step over before +/// the other has arrived. What is asserted is that the first failure and its error are retained, and +/// every output ends with exactly one error, never a clean end and never a duplicate. +#[tokio::test] +async fn two_failures_hand_every_output_one_error() { + const FIRST: &str = "the first partition failed"; + const SECOND: &str = "the second partition failed"; + assert_eq!( + 0, + TEST_FIXTURE_VOLUME % 2, + "the chains of one queue are split evenly between the two partitions" + ); + let half = TEST_FIXTURE_VOLUME / 2; + let to_first = labels_routed_to(0, TEST_FIXTURE_VOLUME); + let to_second = labels_routed_to(1, TEST_FIXTURE_VOLUME); + + /// The steps of one partition: its half of every output's queue, a place at the barrier, and then + /// its failure. + /// + /// The failure is what comes *after* the barrier, which is the whole point: a partition cannot hand + /// it over until every partition has handed over its chains, so the two failures are released into + /// the same wake instead of one of them landing while the other task is still reading. + fn failing( + first: &[SeriesLabels], + second: &[SeriesLabels], + message: &'static str, + ) -> GatedPartition { + GatedPartition { + before: first + .iter() + .chain(second) + .map(|labels| Step::chunk(chunk(labels.clone(), 0, 1, true))) + .collect::>(), + after: vec![Step::ErrorSaying(message)], + switch: None, + } + } + + let (child, arrived) = Simultaneous::plan(vec![ + failing(&to_first[..half], &to_second[..half], FIRST), + failing(&to_first[half..], &to_second[half..], SECOND), + ]); + let plan = node(child, 2); + let context = SessionContext::default().task_ctx(); + // Claim both owners before reading so either partition can observe its terminal failure. + let mut streams = Vec::new(); + for partition in 0..plan.outputs() { + streams.push( + unpack_evaluated_series( + plan.execute(partition, Arc::clone(&context)) + .expect("the output executes"), + None, + ) + .expect("the output declares an evaluated frame"), + ); + } + let execution = execution(&plan); + + wait_until("both partitions reach the failure path", || arrived.all()).await; + wait_until("the read records its failure", || { + matches!(execution.ending(), Ending::Failed(_)) + }) + .await; + let recorded = match execution.ending() { + Ending::Failed(error) => error, + other => panic!("the read records its failure: {other:?}"), + }; + let expected = recorded.to_string(); + let named = [FIRST, SECOND] + .into_iter() + .find(|message| expected.contains(message)) + .unwrap_or_else(|| panic!("the recorded error names one of the two failures: {expected}")); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(streams.iter_mut().map(until_end)), + ) + .await + .expect("a read with two failures still ends every output"); + + for (partition, (chunks, errors)) in outputs.iter().enumerate() { + assert_eq!( + 1, + errors.len(), + "output partition {partition} reports the failure once, not once per failing input: \ + {errors:?}" + ); + // The item and the recorded error are the same failure, but each is wrapped by the path it + // takes out of the read, so what they are compared on is the failure they name. + assert!( + errors[0].contains(named), + "output partition {partition} reports the failure the read recorded ({named}), not {}", + errors[0] + ); + assert_chains_are_whole(chunks); + } +} + +/// a task attached after the failure is swept is aborted instead of waited on. +/// +/// Partition 0 fails on its first poll, which can happen before the tasks of the partitions after it +/// are even spawned - and a handle attached after the failure's sweep is one the sweep cannot see. +/// Stored as live, it waits on a child that never answers while holding a sender to every output, so +/// the queues never close and both outputs hang here. The other partitions are gated on a switch +/// nothing opens, so the only way this test finishes is those tasks being aborted: their streams +/// record the drop an aborted task's future goes through. +#[tokio::test] +async fn a_partition_failing_before_its_peers_are_attached_still_ends_every_output() { + const MESSAGE: &str = "the first partition failed at once"; + let switch = Arc::new(Switch::default()); + let dropped = Arc::new(AtomicUsize::new(0)); + let shared = labels("web-1"); + // Ten partitions, so the failing one's sweep runs while the later ones are still being spawned and + // attached: the read must stop them anyway. `read_partitions` counts them with the failing one. + let mut partitions = vec![GatedPartition::open(vec![Step::ErrorSaying(MESSAGE)])]; + partitions.extend( + (0..10).map(|_| GatedPartition::held(vec![Step::chunk(chunk(shared.clone(), 0, 1, true))])), + ); + let read_partitions = partitions.len(); + let child = Gated::plan(partitions, switch, Arc::clone(&dropped)); + let plan = node(child, 2); + let context = SessionContext::default().task_ctx(); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + every_output_until_end(&plan, &context), + ) + .await + .expect( + "a read whose first partition fails before its peers were attached still ends: a task left \ + waiting on a child holds a sender to every output, so nothing would close", + ); + for (partition, (_, errors)) in outputs.iter().enumerate() { + assert_eq!( + 1, + errors.len(), + "output partition {partition} reports the failure: {errors:?}" + ); + assert!( + errors[0].contains(MESSAGE), + "output partition {partition} ended with {}", + errors[0] + ); + } + wait_until("every child partition's stream is dropped", || { + dropped.load(Ordering::SeqCst) == read_partitions + }) + .await; + // The read is over, which it could not be if a task attached after the sweep were still holding + // a sender to every output. + let execution = execution(&plan); + wait_until("the read stops", || producer_stopped(&execution)).await; +} + +/// A route task attached after failure is aborted by `attach`; its captured report guard counts the +/// partition out when the aborted future is dropped. +#[tokio::test] +async fn a_task_attached_after_the_failure_is_aborted_and_not_counted() { + let ending = Arc::new(Mutex::new(Ending::default())); + let inputs = Inputs::new(Arc::clone(&ending), 2, Weak::::new()); + let error = Arc::new(DataFusionError::Execution( + "the first partition failed".to_string(), + )); + inputs.failed(0, &error); + assert_eq!( + 1, + inputs.live(), + "the late partition is still counted until its guard drops" + ); + + let task = tokio::spawn(route_partition( + Box::pin(futures::stream::pending()), + 1, + Routing::new(1), + Vec::new(), + MemoryConsumer::new("late-attach-test") + .register(&(Arc::new(GreedyMemoryPool::new(1 << 20)) as Arc)), + Arc::clone(&inputs), + ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition: 1, + reported: false, + }, + )); + inputs.attach(1, task); + assert!( + inputs + .state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .tasks[1] + .is_none(), + "the task of a cancelled read is aborted rather than stored as one to wait for" + ); + tokio::time::timeout(Duration::from_secs(10), async { + while inputs.live() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the attached route task's ReportIfUnreported guard counts it out on abort"); +} + +/// a read that finished owes its outputs no error, whatever a failure elsewhere recorded. +/// +/// A read that finished records `Finished`, and every output ends cleanly once. +#[tokio::test] +async fn a_finished_read_ends_cleanly_and_owes_no_error() { + let shared = labels("web-1"); + let child = Scripted::plan(vec![ + vec![Step::chunk(chunk(shared.clone(), 0, 2, true))], + vec![Step::chunk(chunk(shared.clone(), 8, 2, true))], + ]); + let plan = node(child, 2); + let context = SessionContext::default().task_ctx(); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + every_output_until_end(&plan, &context), + ) + .await + .expect("both partitions end, so both outputs end"); + // Read after the outputs are done, so the read has recorded everything it is going to. + let execution = execution(&plan); + for (partition, (chunks, errors)) in outputs.iter().enumerate() { + assert!( + errors.is_empty(), + "output partition {partition} reported {errors:?} for a read that finished" + ); + assert_chains_are_whole(chunks); + } + assert!( + matches!(execution.ending(), Ending::Finished), + "a read that finished records no failure for its outputs to report" + ); +} + +/// a failure recorded before outputs are drained is still reported by every output. +/// +/// One input fills the output queues and then fails while the outputs are not being drained. The +/// failure cancels its peers; outputs observe closed queues and must report the recorded terminal +/// error exactly once rather than ending cleanly. +#[tokio::test] +async fn a_recorded_failure_reaches_every_output_once() { + const MESSAGE: &str = "the partition failed"; + let to_zero = labels_routed_to(0, TEST_FIXTURE_VOLUME); + let to_one = labels_routed_to(1, TEST_FIXTURE_VOLUME); + let mut steps = to_zero + .into_iter() + .zip(to_one) + .flat_map(|(zero, one)| [zero, one].map(|metric| Step::chunk(chunk(metric, 0, 1, true)))) + .collect::>(); + steps.push(Step::ErrorSaying(MESSAGE)); + let plan = node( + Fragile::plan( + vec![Behaviour::OneAtATime(steps)], + Arc::new(Switch::default()), + ) + .plan, + 2, + ); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + // Claim internal queues so assertions observe retained chains, not encoded frame batches. + let child = plan.children()[0].clone(); + let shared = + SharedExecution::start(&child, Routing::new(2), Arc::clone(&schema()), &context).unwrap(); + let mut streams = Vec::new(); + for partition in 0..2 { + let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); + streams.push(Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + )) as SendableEvaluatedSeriesStream); + } + let execution = shared; + + wait_until( + "the read records its failure while no output is drained", + || matches!(execution.ending(), Ending::Failed(_)), + ) + .await; + // The producer records the failure before its queued chains are all visible, so wait for the + // retention this assertion is about instead of reading the pool once. + wait_until("queued output chains retain memory before drain", || { + pool.reserved() > 0 + }) + .await; + let outputs = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(streams.iter_mut().map(until_end)), + ) + .await + .expect("a failed read still ends every output"); + for (partition, (chunks, errors)) in outputs.iter().enumerate() { + assert!( + chunks.is_empty(), + "no incomplete owner series escapes on failure" + ); + assert_eq!( + 1, + errors.len(), + "output partition {partition} reports the failure the read recorded exactly once: \ + {errors:?}" + ); + assert!( + errors[0].contains(MESSAGE), + "output partition {partition} ended with {} rather than the failure", + errors[0] + ); + } + assert_eq!(0, pool.reserved(), "draining releases all queued charges"); +} + +/// What one partition of a [`Fragile`] child does when it is polled. +#[derive(Clone)] +enum Behaviour { + /// Hands `steps` over only once the switch is opened, then ends. + Held(Vec), + /// Panics on the first poll, which is a task dying where no arm of the read's ending match is + /// reached: the unwind takes the whole future with it. + Panics, + /// Never answers, so the task reading it stays parked on it until it is aborted. + Parked, + /// Hands one step over per poll, staying pending in between, so each of them becomes a frame + /// batch of its own. + /// + /// The frame boundary coalesces what a producer hands over in one poll into one batch - the + /// builder only flushes when the source is pending or done - so a child that hands several chunks + /// over back to back reaches a consumer as ONE chunk. A chunk per poll is what keeps them apart, + /// and a test about how a chain accumulates is about chunks rather than about batches. + OneAtATime(Vec), + /// Delivers one chunk, waits until the test opens the gate, then delivers the rest. + Gated(Vec), +} + +/// A child whose partitions stop in ways of their own: one dies or is parked, another stays parked. +/// +/// A test of what a read does with a producer that dies has to be able to make exactly one producer +/// die, and to see that the others really were reading and really were stopped. [`Behaviour::Panics`] +/// is the death nothing in the read reaches, [`Behaviour::Parked`] is the peer that would wait for a +/// child that never answers, and the counters answer the two questions a timeout cannot: `polled` +/// says a partition was being read rather than skipped, and `dropped` says the read let the stream of +/// a parked partition go - which is what an aborted task does. +struct Fragile { + behaviours: Vec, + switch: Arc, + /// Whether each partition has been polled once, which every partition waits for before it does + /// what its behaviour says. + /// + /// [`Behaviour::Panics`] needs it, and needs it for the peer rather than for itself: a producer + /// that dies on its very first poll can die before the task reading the other partition has been + /// polled at all, and then the peer is aborted without ever having read anything - which is a read + /// that never had a parked input to hang on. Waiting here is what makes the shape the test is + /// about exist before the death happens. + reading: Arc, + polled: Vec>, + dropped: Vec>, + series_schema: Arc, + properties: Arc, +} + +/// A [`Fragile`] child, and the two counters of every one of its partitions. +struct FragileChild { + plan: Arc, + /// How many times each partition was polled. + polled: Vec>, + /// How many of each partition's streams were dropped. + dropped: Vec>, +} + +impl Fragile { + /// The child over `behaviours`, one partition per behaviour, gated on `switch`. + fn plan(behaviours: Vec, switch: Arc) -> FragileChild { + let series_schema = schema(); + let frame = pack_evaluated_series( + Box::pin(futures::stream::empty()), + Arc::clone(&series_schema), + ) + .expect("an evaluated schema declares a frame") + .schema(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(frame), + Partitioning::UnknownPartitioning(behaviours.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + let polled = (0..behaviours.len()) + .map(|_| Arc::new(AtomicUsize::new(0))) + .collect::>(); + let dropped = (0..behaviours.len()) + .map(|_| Arc::new(AtomicUsize::new(0))) + .collect::>(); + FragileChild { + plan: Arc::new(Self { + behaviours: behaviours.clone(), + switch, + reading: FirstChunks::new(behaviours.len()), + polled: polled.clone(), + dropped: dropped.clone(), + series_schema, + properties, + }), + polled, + dropped, + } + } +} + +impl fmt::Debug for Fragile { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("Fragile") + .field("partitions", &self.behaviours.len()) + .finish() + } +} + +impl DisplayAs for Fragile { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "Fragile: partitions={}", self.behaviours.len()) + } +} + +impl ExecutionPlan for Fragile { + fn name(&self) -> &str { + "Fragile" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> DataFusionResult, + ) -> DataFusionResult { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> DataFusionResult> { + assert!(children.is_empty(), "a fragile child has no children"); + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> DataFusionResult { + let behaviour = self.behaviours.get(partition).cloned().ok_or_else(|| { + DataFusionError::Execution(format!("fragile child has no partition {partition}")) + })?; + let steps = match &behaviour { + Behaviour::Held(steps) | Behaviour::OneAtATime(steps) | Behaviour::Gated(steps) => { + steps.clone() + } + Behaviour::Panics | Behaviour::Parked => Vec::new(), + }; + pack_evaluated_series( + Box::pin(FragileScript { + partition, + behaviour, + steps: steps.into_iter(), + switch: Arc::clone(&self.switch), + reading: Arc::clone(&self.reading), + polled: Arc::clone(&self.polled[partition]), + dropped: Arc::clone(&self.dropped[partition]), + stayed_pending: false, + }), + Arc::clone(&self.series_schema), + ) + } +} + +/// One partition of a [`Fragile`] child. +struct FragileScript { + partition: usize, + behaviour: Behaviour, + steps: std::vec::IntoIter, + switch: Arc, + reading: Arc, + polled: Arc, + dropped: Arc, + /// Whether a [`Behaviour::OneAtATime`] partition has already stayed pending since its last chunk. + stayed_pending: bool, +} + +impl Stream for FragileScript { + type Item = Result<(Series, bool), BoxedError>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = &mut *self; + this.polled.fetch_add(1, Ordering::SeqCst); + // Every partition is read once before any of them does what it is here for, so a partition + // that dies cannot do so before its peer was read. + this.reading.handed(this.partition); + this.reading.wait(cx); + if !this.reading.all() { + return Poll::Pending; + } + match &this.behaviour { + Behaviour::Held(_) => { + this.switch.wait(cx); + if !this.switch.is_open() { + return Poll::Pending; + } + Poll::Ready(this.steps.next().map(step_answer)) + } + Behaviour::Gated(_) => { + if !this.stayed_pending { + this.stayed_pending = true; + return Poll::Ready(this.steps.next().map(step_answer)); + } + this.switch.wait(cx); + if !this.switch.is_open() { + return Poll::Pending; + } + Poll::Ready(this.steps.next().map(step_answer)) + } + // The child's own poll, which runs inside the task reading it: what a producer that dies + // under a read does is exactly this, and nothing of the read is told about it. + Behaviour::Panics => panic!("the fragile child partition panicked while being read"), + // Parked on a source that never answers. No waker is registered because nothing can make + // this partition answer: the only way it stops is the read aborting it. + Behaviour::Parked => Poll::Pending, + Behaviour::OneAtATime(_) => { + if !this.stayed_pending { + // One pending poll between two chunks is what makes the frame boundary flush, + // and the wake is this stream's own: nothing outside it has to poll again for the + // next chunk to arrive. + this.stayed_pending = true; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + this.stayed_pending = false; + Poll::Ready(this.steps.next().map(step_answer)) + } + } + } +} + +impl Drop for FragileScript { + fn drop(&mut self) { + self.dropped.fetch_add(1, Ordering::SeqCst); + } +} + +struct DropCounter(Arc); + +impl Drop for DropCounter { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +/// A task guard captured into an unpolled producer future still reports cancellation on drop. +/// +/// Constructing `route_partition` synchronously constructs its captured `ReportIfUnreported` guard; +/// dropping the returned future before first poll therefore runs the failure bookkeeping. A parked +/// sibling then stops, all output sender clones close, and the output streams report the same error. +#[tokio::test] +async fn dropping_unpolled_route_future_ends_outputs_and_cancels_peer() { + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let ending = Arc::new(Mutex::new(Ending::default())); + let inputs = Inputs::new(Arc::clone(&ending), 2, Weak::::new()); + let mut senders = Vec::new(); + let mut receivers = Vec::new(); + for _ in 0..2 { + let (sender, receiver) = mpsc::unbounded_channel(); + senders.push(sender); + receivers.push(Some(receiver)); + } + let shared = Arc::new(SharedExecution { + running: Mutex::new(Running { + receivers, + inputs: Arc::clone(&inputs), + claimed: 0, + remaining: 2, + }), + ending, + }); + let series_schema = Arc::clone(&schema()); + let mut streams = (0..2) + .map(|partition| { + let (receiver, lease) = + SharedExecution::claim(&shared, partition).expect("output claim"); + let frames = pack_evaluated_series( + Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + )), + Arc::clone(&series_schema), + ) + .expect("output frame builds"); + unpack_evaluated_series(frames, None).expect("evaluated output decodes") + }) + .collect::>(); + let peer_polled = Arc::new(AtomicUsize::new(0)); + let peer_dropped = Arc::new(AtomicUsize::new(0)); + let peer_poll_count = Arc::clone(&peer_polled); + let peer_guard = DropCounter(Arc::clone(&peer_dropped)); + let peer_stream = futures::stream::poll_fn(move |_| { + let _ = &peer_guard; + peer_poll_count.fetch_add(1, Ordering::SeqCst); + Poll::Pending + }); + let peer = route_partition( + Box::pin(peer_stream), + 1, + Routing::new(2), + senders.clone(), + MemoryConsumer::new("parked-peer-test") + .register(&(Arc::clone(&pool) as Arc)), + Arc::clone(&inputs), + ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition: 1, + reported: false, + }, + ); + let task = common_runtime::spawn_query(peer); + inputs.attach(1, task); + wait_until("parked sibling is actively holding output senders", || { + peer_polled.load(Ordering::SeqCst) > 0 + }) + .await; + let report = ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition: 0, + reported: false, + }; + let unpolled = route_partition( + Box::pin(futures::stream::pending()), + 0, + Routing::new(2), + senders.clone(), + MemoryConsumer::new("unpolled-route-test") + .register(&(Arc::clone(&pool) as Arc)), + Arc::clone(&inputs), + report, + ); + drop(senders); + drop(unpolled); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(streams.iter_mut().map(until_end)), + ) + .await + .expect("unpolled producer drop closes each output"); + for (chunks, errors) in &outputs { + assert!(chunks.is_empty()); + assert_eq!(1, errors.len()); + assert!(errors[0].contains("input partition 0 stopped without reporting an outcome")); + } + wait_until("parked sibling stream is dropped", || { + peer_dropped.load(Ordering::SeqCst) == 1 + }) + .await; + wait_until("producer tasks stopped", || producer_stopped(&shared)).await; + wait_until("input memory released", || pool.reserved() == 0).await; +} + +#[tokio::test] +async fn unpolled_route_drop_preserves_an_already_elected_failure() { + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + const FIRST: &str = "the child failure was elected first"; + let ending = Arc::new(Mutex::new(Ending::default())); + let inputs = Inputs::new(Arc::clone(&ending), 2, Weak::::new()); + let mut senders = Vec::new(); + let mut streams = Vec::new(); + let shared = { + let mut receivers = Vec::new(); + for _ in 0..2 { + let (sender, receiver) = mpsc::unbounded_channel(); + senders.push(sender); + receivers.push(Some(receiver)); + } + Arc::new(SharedExecution { + running: Mutex::new(Running { + receivers, + inputs: Arc::clone(&inputs), + claimed: 0, + remaining: 2, + }), + ending, + }) + }; + for partition in 0..2 { + let (receiver, lease) = SharedExecution::claim(&shared, partition).expect("output claim"); + let frames = pack_evaluated_series( + Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 2, + )), + Arc::clone(&schema()), + ) + .expect("output frame builds"); + streams.push(unpack_evaluated_series(frames, None).expect("evaluated output decodes")); + } + inputs.failed(1, &Arc::new(DataFusionError::Execution(FIRST.to_string()))); + let unpolled = route_partition( + Box::pin(futures::stream::pending()), + 0, + Routing::new(2), + senders.clone(), + MemoryConsumer::new("unpolled-after-failure-test") + .register(&(Arc::new(GreedyMemoryPool::new(1 << 20)) as Arc)), + Arc::clone(&inputs), + ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition: 0, + reported: false, + }, + ); + drop(senders); + drop(unpolled); + let outputs = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(streams.iter_mut().map(until_end)), + ) + .await + .expect("elected error closes both outputs"); + for (_, errors) in outputs { + assert_eq!(1, errors.len()); + assert!( + errors[0].contains(FIRST), + "unexpected terminal error: {errors:?}" + ); + assert!(!errors[0].contains("partition 0 stopped without reporting")); + } +} + +/// A continuously-ready decoded source is cooperatively budgeted while retaining a partial chain. +#[tokio::test] +async fn ready_chain_yields_after_a_charged_prefix_and_drop_releases_it() { + const TOTAL: usize = 16_384; + let count = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&count); + let long_grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, TOTAL as u64).unwrap()).unwrap(), + ); + let long_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: Some(Arc::clone(&long_grid)), + }); + let source = futures::stream::poll_fn(move |_| { + let position = seen.fetch_add(1, Ordering::SeqCst); + Poll::Ready((position < TOTAL).then(|| { + let time = + TimeAxis::from_grid_spans(long_grid.clone(), vec![GridSpan::new(position, 1)]) + .unwrap(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + position as f64, + ]))); + let series = + Series::try_new(&long_schema, labels("cooperative-prefix"), time, fields).unwrap(); + Ok((series, position + 1 == TOTAL)) + })) + }); + let mut source: SendableEvaluatedSeriesStream = Box::pin(source); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let reservation = MemoryConsumer::new("cooperative-prefix-test") + .register(&(Arc::clone(&pool) as Arc)); + let (sender, receiver) = mpsc::unbounded_channel(); + let routing = Routing::new(1); + let senders = [sender]; + let mut read = Box::pin(read_chains( + &mut source, + 0, + &routing, + &senders, + &reservation, + )); + let result = std::future::poll_fn(|cx| Poll::Ready(read.as_mut().poll(cx))).await; + assert!( + result.is_pending(), + "one poll reaches the cooperative budget, not chain end" + ); + let prefix = count.load(Ordering::SeqCst); + assert!( + prefix > 0 && prefix < TOTAL, + "observed {prefix} of {TOTAL} chunks" + ); + assert!( + pool.reserved() > 0, + "partial chain owns a memory reservation" + ); + drop(read); + assert_eq!( + 0, + pool.reserved(), + "dropping the read releases its partial chain" + ); + drop(receiver); +} + +#[tokio::test(flavor = "current_thread")] +async fn sibling_failure_aborts_budget_yielded_route_and_releases_prefix() { + const TOTAL: usize = 16_384; + let count = Arc::new(AtomicUsize::new(0)); + let dropped = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&count); + let guard = DropCounter(Arc::clone(&dropped)); + let long_grid = Arc::new( + EvalGrid::try_new(GridSpec::try_new(GRID_START, GRID_STEP, TOTAL as u64).unwrap()).unwrap(), + ); + let long_schema = Arc::new(SeriesSchema { + fields: Arc::from([FieldSpec { + name: Arc::from("value"), + }]), + eval_grid: Some(Arc::clone(&long_grid)), + }); + let source = futures::stream::poll_fn(move |_| { + let _ = &guard; + let position = seen.fetch_add(1, Ordering::SeqCst); + Poll::Ready((position < TOTAL).then(|| { + let time = + TimeAxis::from_grid_spans(long_grid.clone(), vec![GridSpan::new(position, 1)]) + .unwrap(); + let fields = SeriesFields::one(Samples::from_float(Float64Array::from(vec![ + position as f64, + ]))); + let series = + Series::try_new(&long_schema, labels("aborted-prefix"), time, fields).unwrap(); + Ok((series, position + 1 == TOTAL)) + })) + }); + let source: SendableEvaluatedSeriesStream = Box::pin(source); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let ending = Arc::new(Mutex::new(Ending::default())); + let inputs = Inputs::new(Arc::clone(&ending), 2, Weak::::new()); + let (sender, receiver) = mpsc::unbounded_channel(); + let shared = Arc::new(SharedExecution { + running: Mutex::new(Running { + receivers: vec![Some(receiver)], + inputs: Arc::clone(&inputs), + claimed: 0, + remaining: 1, + }), + ending: Arc::clone(&ending), + }); + let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); + let mut output: SendableEvaluatedSeriesStream = Box::pin(RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + )); + let future = route_partition( + source, + 0, + Routing::new(1), + vec![sender], + MemoryConsumer::new("aborted-prefix-test") + .register(&(Arc::clone(&pool) as Arc)), + Arc::clone(&inputs), + ReportIfUnreported { + inputs: Arc::clone(&inputs), + partition: 0, + reported: false, + }, + ); + let task = tokio::spawn(future); + inputs.attach(0, task); + wait_until("route yields after charging a strict chunk prefix", || { + let chunks = count.load(Ordering::SeqCst); + chunks > 0 && chunks < TOTAL && pool.reserved() > 0 + }) + .await; + let prefix = count.load(Ordering::SeqCst); + assert!(prefix < TOTAL, "failure arrives before all {TOTAL} chunks"); + inputs.failed( + 1, + &Arc::new(DataFusionError::Execution("sibling failure".into())), + ); + wait_until("aborted route drops its decoded source", || { + dropped.load(Ordering::SeqCst) == 1 + }) + .await; + assert_eq!( + prefix, + count.load(Ordering::SeqCst), + "no chunk is read after cancellation" + ); + assert_eq!( + 0, + pool.reserved(), + "cancellation releases the retained prefix" + ); + let (chunks, errors) = until_end(&mut output).await; + assert!(chunks.is_empty()); + assert_eq!( + 1, + errors.len(), + "one terminal failure precedes EOF: {errors:?}" + ); + assert!( + errors[0].contains("sibling failure"), + "first failure is preserved: {errors:?}" + ); +} + +/// A partition whose task panics ends every output with the read's error, and stops its peers. +/// +/// Two inputs and one output. Partition 0's child panics inside its `poll_next`, which unwinds the +/// task reading it: no arm of [`route_partition`]'s ending match runs, so nothing is recorded and +/// nothing stops the peer. Partition 1's child never answers, keeping its sender alive and preventing +/// the output from observing an ending. The timeout makes the missing terminal error visible. +/// +/// The panic happens on a task of the read, never on this test's, so the unwind is the runtime's to +/// catch and what the test sees is its effect: the output ends with the read's error instead of +/// waiting, the parked peer's child is dropped because its task was aborted, and nothing the read +/// reserved outlives it. +#[tokio::test] +async fn a_panicking_input_ends_every_output_and_stops_its_peers() { + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let child = Fragile::plan( + vec![Behaviour::Panics, Behaviour::Parked], + Arc::new(Switch::default()), + ); + let plan = node(child.plan, 1); + let mut stream = output_of(&plan, 0, &context); + let execution = execution(&plan); + wait_until("the parked peer is reading its child", || { + child.polled[1].load(Ordering::SeqCst) > 0 + }) + .await; + + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut stream)) + .await + .unwrap_or_else(|_| { + panic!( + "the output never ended although one of its two inputs had died: the death recorded \ + nothing, so the parked peer kept its sender alive" + ) + }); + assert!( + chunks.is_empty(), + "no chunk was routed before the read stopped, so the output owes none" + ); + assert_eq!( + 1, + errors.len(), + "the output reports the death of partition 0 once: {errors:?}" + ); + assert!( + errors[0] + .contains("series repartition input partition 0 stopped without reporting an outcome"), + "the output ended with {} rather than the death of its producer", + errors[0] + ); + // The peer did not go on waiting for a child that never answers: its task was aborted, which is + // what the drop of the stream it was reading says. + wait_until("the parked peer's child is dropped", || { + child.dropped[1].load(Ordering::SeqCst) == 1 + }) + .await; + wait_until("the read's tasks are gone", || producer_stopped(&execution)).await; + wait_until("what the read reserved is released", || { + pool.reserved() == 0 + }) + .await; +} + +/// a partition whose task is aborted on its own ends every output, and stops its peers. +/// +/// The same shape as [`a_panicking_input_ends_every_output_and_stops_its_peers`] with the other way a +/// task disappears: nothing panics, one task is aborted where it stands - as a cancelled query's task +/// is - and the future it was running is dropped with it. Both partitions stay parked on children that +/// never answer, so the only thing that can end this read is one of them dying and the read noticing: +/// the output would otherwise wait for two sources that never speak. +/// +/// The abort reaches the handle of one partition only. [`Inputs::stop`] cannot state this: it stops +/// every task and records the read as stopped, so what its outputs report is the stop rather than one +/// producer that died while its peers were still running. +#[tokio::test] +async fn an_aborted_input_task_ends_every_output_and_stops_its_peers() { + let child = Fragile::plan( + vec![Behaviour::Parked, Behaviour::Parked], + Arc::new(Switch::default()), + ); + let plan = node(child.plan, 1); + let context = SessionContext::default().task_ctx(); + let mut stream = output_of(&plan, 0, &context); + let execution = execution(&plan); + wait_until("both partitions are reading their children", || { + child.polled[0].load(Ordering::SeqCst) > 0 && child.polled[1].load(Ordering::SeqCst) > 0 + }) + .await; + + let inputs = Arc::clone( + &execution + .running + .lock() + .unwrap_or_else(|error| error.into_inner()) + .inputs, + ); + inputs.abort(0); + + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut stream)) + .await + .unwrap_or_else(|_| { + panic!( + "the output never ended although one of its two inputs had been aborted: the abort \ + recorded nothing, so the parked peer kept the queue open" + ) + }); + assert!( + chunks.is_empty(), + "both partitions are parked, so nothing was routed before the read stopped" + ); + assert_eq!( + 1, + errors.len(), + "the output reports the death of partition 0 once: {errors:?}" + ); + assert!( + errors[0] + .contains("series repartition input partition 0 stopped without reporting an outcome"), + "the output ended with {} rather than the death of its producer", + errors[0] + ); + wait_until("the other parked partition's child is dropped", || { + child.dropped[1].load(Ordering::SeqCst) == 1 + }) + .await; + wait_until("the read's tasks are gone", || producer_stopped(&execution)).await; +} + +/// A dropped destination's oversized chain is discarded while a live destination continues. +#[tokio::test] +async fn dropped_output_discards_oversized_chain_and_keeps_live_output() { + let doomed = labels_routed_to(0, 1).remove(0); + let survivor = labels_routed_to(1, 1).remove(0); + let many_chunks = (0..GRID_LEN as usize) + .map(|position| { + Step::chunk(chunk( + doomed.clone(), + position, + 1, + position + 1 == GRID_LEN as usize, + )) + }) + .collect::>(); + let many_chunk_count = many_chunks.len(); + let live_chunk = chunk(survivor.clone(), 0, 1, true); + // Independent workspace allowance for the two live base chunks and queued-chain metadata. + // Keep this finite and far below retaining every old whole-chain slice. + let live_chunk_budget = operand_chunk_bytes(&live_chunk.0) * 2; + let byte_budget = live_chunk_budget + 16 * 1024; + assert!( + operand_chunk_bytes(&chunk(doomed.clone(), 0, 1, false).0) * many_chunk_count > byte_budget, + "discarding the oversized chain must be necessary under this finite budget" + ); + let pool = Arc::new(GreedyMemoryPool::new(byte_budget)); + let context = task_context(&pool); + let switch = Arc::new(Switch::default()); + let child = Fragile::plan( + vec![Behaviour::Held({ + let mut steps = many_chunks; + steps.push(Step::chunk(live_chunk)); + steps + })], + Arc::clone(&switch), + ); + let plan = node(child.plan, 2); + let dropped_output = output_of(&plan, 0, &context); + let mut kept_output = output_of(&plan, 1, &context); + let first_chunk_bytes = operand_chunk_bytes(&chunk(doomed.clone(), 0, 1, false).0); + assert_eq!(0, pool.reserved(), "held child has not produced a chunk"); + drop(dropped_output); + assert_eq!( + 0, + pool.reserved(), + "closed destination was dropped before first chunk" + ); + switch.open(); + assert!(first_chunk_bytes > 0); + wait_until( + "held child processes discarded chain and live chain", + || child.polled[0].load(Ordering::SeqCst) >= many_chunk_count + 2, + ) + .await; + + let (chunks, errors) = + tokio::time::timeout(Duration::from_secs(10), until_end(&mut kept_output)) + .await + .expect("discarded oversized chain does not block the live destination"); + assert!(errors.is_empty(), "live output reports {errors:?}"); + assert_eq!( + vec![named(&survivor)], + chunks + .iter() + .map(|chunk| named(&chunk.0.metric)) + .collect::>() + ); + assert_eq!(1, chunks.len()); + wait_until("discarded and delivered chain charges are released", || { + pool.reserved() == 0 + }) + .await; +} + +#[tokio::test] +async fn discarded_open_chain_is_validated_before_live_chain_routes() { + const INVALID: &str = "evaluated series positions are not increasing across chunks"; + let doomed = labels_routed_to(0, 1).remove(0); + let first_chunk_bytes = operand_chunk_bytes(&chunk(doomed.clone(), 0, 1, false).0); + let invalid = chunk(doomed.clone(), 0, 1, true); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let switch = Arc::new(Switch::default()); + let child = Fragile::plan( + vec![Behaviour::Gated(vec![ + Step::chunk(chunk(doomed, 0, 1, false)), + Step::chunk(invalid), + ])], + Arc::clone(&switch), + ); + let plan = node(child.plan, 2); + let dropped_output = output_of(&plan, 0, &context); + let mut kept_output = output_of(&plan, 1, &context); + wait_until("valid prefix is flushed and child parks at gate", || { + pool.reserved() >= first_chunk_bytes && child.polled[0].load(Ordering::SeqCst) >= 2 + }) + .await; + drop(dropped_output); + wait_until( + "discarding destination releases prefix while gate stays closed", + || pool.reserved() == 0, + ) + .await; + switch.open(); + let (chunks, errors) = + tokio::time::timeout(Duration::from_secs(10), until_end(&mut kept_output)) + .await + .expect("malformed continuation fails the execution"); + assert!( + chunks.is_empty(), + "no incomplete chain was routed: {chunks:?}" + ); + assert_eq!( + 1, + errors.len(), + "terminal error is reported exactly once: {errors:?}" + ); + assert!( + errors[0].contains(INVALID), + "expected position validation error: {errors:?}" + ); + assert_eq!( + 0, + pool.reserved(), + "failed discarded chain retains no memory" + ); +} + +#[tokio::test] +async fn discarded_open_chain_eof_is_still_an_error() { + let (sender, receiver) = mpsc::unbounded_channel(); + drop(receiver); + let series = chunk(labels("unclosed"), 0, 1, false).0; + let mut input: SendableEvaluatedSeriesStream = + Box::pin(futures::stream::iter(vec![Ok((series, false))])); + let pool: Arc = Arc::new(GreedyMemoryPool::new(1 << 20)); + let reservation = MemoryConsumer::new("discarded-open-eof-test").register(&pool); + let result = read_chains(&mut input, 0, &Routing::new(1), &[sender], &reservation).await; + assert!( + matches!(result, Stop::Failed(_)), + "open discarded chain EOF is not clean" + ); + assert_eq!(0, pool.reserved()); +} + +/// A closed output is skipped while the sibling output continues to receive its chains. +#[tokio::test] +async fn closing_destination_releases_an_open_chain_before_child_resumes() { + let doomed = labels_routed_to(0, 1).remove(0); + let survivor = labels_routed_to(1, 1).remove(0); + let switch = Arc::new(Switch::default()); + let pool = Arc::new(GreedyMemoryPool::new(1 << 20)); + let context = task_context(&pool); + let held_bytes = operand_chunk_bytes(&chunk(doomed.clone(), 0, 1, false).0); + let child = Fragile::plan( + vec![Behaviour::Gated(vec![ + Step::chunk(chunk(doomed.clone(), 0, 1, false)), + Step::chunk(chunk(doomed, 1, 1, true)), + Step::chunk(chunk(survivor.clone(), 4, 1, true)), + ])], + Arc::clone(&switch), + ); + let plan = node(child.plan, 2); + let dropped_output = output_of(&plan, 0, &context); + let mut kept_output = output_of(&plan, 1, &context); + wait_until("open chain is retained and child is gated", || { + pool.reserved() >= held_bytes + }) + .await; + drop(dropped_output); + wait_until( + "open chain reservation is released while child remains gated", + || pool.reserved() < held_bytes, + ) + .await; + switch.open(); + let (chunks, errors) = + tokio::time::timeout(Duration::from_secs(10), until_end(&mut kept_output)) + .await + .expect("the live output completes after gate opens"); + assert!(errors.is_empty()); + assert_eq!( + vec![named(&survivor)], + chunks + .iter() + .map(|chunk| named(&chunk.0.metric)) + .collect::>() + ); + assert_eq!(1, chunks.len()); + assert_eq!(0, pool.reserved()); +} + +#[tokio::test] +async fn dropping_an_output_leaves_surviving_outputs_running() { + let switch = Arc::new(Switch::default()); + let doomed = labels_routed_to(0, 1).remove(0); + let survivor = labels_routed_to(1, 1).remove(0); + let child = Fragile::plan( + vec![ + Behaviour::OneAtATime(vec![ + Step::chunk(chunk(doomed.clone(), 0, 1, true)), + Step::chunk(chunk(survivor.clone(), 1, 1, true)), + ]), + Behaviour::Held(vec![]), + ], + Arc::clone(&switch), + ); + let plan = node(child.plan, 2); + let context = SessionContext::default().task_ctx(); + // Both outputs are claimed, so the one that goes away releases a claim on a read that still has + // another output to serve instead of stopping it. + let dropped_output = output_of(&plan, 0, &context); + let mut kept_output = output_of(&plan, 1, &context); + let execution = execution(&plan); + + // Dropped while its queue is still empty and its producer parked on the switch, so the chain held + // for it is what finds it gone. + drop(dropped_output); + switch.open(); + + let (chunks, errors) = tokio::time::timeout( + Duration::from_secs(10), + until_end(&mut kept_output), + ) + .await + .unwrap_or_else(|_| { + panic!( + "the output that remains claimed did not finish after its sibling output was dropped" + ) + }); + assert_eq!( + vec![named(&survivor)], + chunks + .iter() + .map(|chunk| named(&chunk.0.metric)) + .collect::>(), + "the surviving output keeps its own chains after the sibling is dropped" + ); + assert!( + errors.is_empty(), + "a dropped sibling is not a read failure: {errors:?}" + ); + wait_until("the parked partition's child is dropped", || { + child.dropped[1].load(Ordering::SeqCst) == 1 + }) + .await; + wait_until("the read's tasks are gone", || producer_stopped(&execution)).await; + assert!(matches!(execution.ending(), Ending::Finished)); +} + +#[tokio::test] +async fn sliced_chunk_memory_is_charged_by_rows_and_released_after_success_drop_and_failure() { + let slices = sliced_chunks(labels("sliced")); + let probe_context = task_context(&Arc::new(GreedyMemoryPool::new(usize::MAX))); + let probe_child = Scripted::plan(vec![ + slices + .iter() + .cloned() + .map(|(series, is_last)| Step::Chunk(series, is_last)) + .collect(), + ]); + let probe_stream = probe_child.execute(0, Arc::clone(&probe_context)).unwrap(); + let mut probe = unpack_evaluated_series(probe_stream, None).unwrap(); + let (decoded_slices, errors) = until_end(&mut probe).await; + assert!( + errors.is_empty(), + "decoded slice probe succeeds: {errors:?}" + ); + assert_eq!(2, decoded_slices.len()); + let slice_costs = decoded_slices + .iter() + .map(|(series, _)| operand_chunk_bytes(series)) + .collect::>(); + let expected_bytes = slice_costs[0]; + assert_eq!( + expected_bytes, slice_costs[1], + "same-sized slices cost equally" + ); + let first = &slices[0].0; + assert!( + series_bytes(first) > expected_bytes, + "a sliced chunk charges its row payload rather than its backing Arrow allocation" + ); + let success_budget = expected_bytes * 2 + 8 * 1024; + let pool = Arc::new(GreedyMemoryPool::new(success_budget)); + let context = task_context(&pool); + let plan = node( + Scripted::plan(vec![ + slices + .iter() + .cloned() + .map(|(s, l)| Step::Chunk(s, l)) + .collect(), + ]), + 1, + ); + let mut output = output_of(&plan, 0, &context); + let (chunks, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .unwrap(); + assert!( + errors.is_empty(), + "expected {expected_bytes} bytes per slice; got {errors:?}" + ); + assert_eq!(2, chunks.len()); + assert_eq!( + 0, + pool.reserved(), + "successful delivery releases reservations" + ); + + let slices = sliced_chunks(labels("sliced-drop")); + let pool = Arc::new(GreedyMemoryPool::new(expected_bytes * 3 + 8 * 1024)); + let context = task_context(&pool); + let plan = node( + Scripted::plan(vec![ + slices + .iter() + .cloned() + .map(|(s, l)| Step::Chunk(s, l)) + .collect(), + ]), + 1, + ); + let mut output = output_of(&plan, 0, &context); + let first = futures::future::poll_fn(|cx| output.as_mut().poll_next(cx)) + .await + .expect("first slice arrives") + .expect("first slice succeeds"); + drop(first); + drop(output); + drop(plan); + wait_until("dropping the output releases slice memory", || { + pool.reserved() == 0 + }) + .await; + assert_eq!( + 0, + pool.reserved(), + "dropping public output releases all slice charges" + ); + + let slices = sliced_chunks(labels("sliced")); + // This deliberately undersized pool refuses the first slice once its independently-accounted + // routing workspace is included; the test covers terminal error and reservation release. + let budget = expected_bytes * 2 - 1; + assert!(expected_bytes <= budget); + assert!(2 * expected_bytes > budget); + assert!( + series_bytes(&slices[0].0) + series_bytes(&slices[1].0) > 2 * expected_bytes, + "full backing arrays cost more than the row-based charges" + ); + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let context = task_context(&pool); + let plan = node( + Scripted::plan(vec![ + slices + .iter() + .cloned() + .map(|(s, l)| Step::Chunk(s, l)) + .collect(), + ]), + 1, + ); + let mut output = output_of(&plan, 0, &context); + let (_, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .unwrap(); + assert_eq!( + 1, + errors.len(), + "the second same-sized slice is refused while the first charge is retained" + ); + assert_eq!( + 0, + pool.reserved(), + "the failed chain releases its slice reservation" + ); +} + +#[tokio::test] +async fn internal_stream_releases_each_chunk_charge_only_when_polled() { + let slices = sliced_chunks(labels("internal-lifetime")); + let costs = slices + .iter() + .map(|(series, _)| operand_chunk_bytes(series)) + .collect::>(); + assert_eq!(costs[0], costs[1]); + let expected = costs.iter().sum::(); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let child = Scripted::plan(vec![ + slices + .into_iter() + .map(|(series, is_last)| Step::Chunk(series, is_last)) + .collect(), + ]); + let shared = + SharedExecution::start(&child, Routing::new(1), Arc::clone(&schema()), &context).unwrap(); + let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); + let mut stream = RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + ); + wait_until("internal producer finishes", || producer_stopped(&shared)).await; + // Queue, EOF, and identity-group metadata are charged alongside the retained chunks; the + // fixture payload is more than four times this allowance, so the bound still proves payload + // retention rather than control-only accounting. + const CONTROL_ALLOWANCE: usize = 256; + assert!( + expected > CONTROL_ALLOWANCE * 4, + "the payload bound is materially larger than the control allowance" + ); + assert!( + (expected..=expected + CONTROL_ALLOWANCE).contains(&pool.reserved()), + "retained bytes {} cover the payload bound {expected} and stay within the {CONTROL_ALLOWANCE}-byte control allowance", + pool.reserved() + ); + let merged = futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) + .await + .unwrap() + .unwrap(); + assert!(merged.1, "the logical owner closes once after global EOF"); + assert_eq!( + 4, + merged.0.time.len(), + "both slices retain their two positions in the merged logical owner" + ); + assert!( + pool.reserved() >= expected, + "the internal handoff retains payload charges for the returned owner" + ); + assert!( + futures::future::poll_fn(|cx| Pin::new(&mut stream).poll_next(cx)) + .await + .is_none(), + "the logical owner has one stable end" + ); + drop(merged); + drop(stream); + assert_eq!( + 0, + pool.reserved(), + "dropping handoff releases retained charges" + ); +} + +#[tokio::test] +async fn public_repartition_interleaved_histogram_tables_stay_linear() { + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + + for table_count in [1usize, 2] { + let metric = labels("interleaved-histograms"); + let mut inputs = Vec::new(); + let mut hist_positions = Vec::new(); + for table in 0..table_count { + let positions = (0..128 / table_count) + .map(|index| index * 2 * table_count + table * 2) + .collect::>(); + hist_positions.extend(positions.iter().copied()); + inputs.push(vec![Step::chunk(sparse_histogram_rows_chunk( + metric.clone(), + positions, + 128 / table_count, + 10.0 + table as f64, + ))]); + } + let float_positions = (0..128).map(|index| index * 2 + 1).collect::>(); + let float_values = float_positions + .iter() + .map(|position| Some(*position as f64)) + .collect::>(); + let float_time = TimeAxis::from_grid_spans( + schema().eval_grid.clone().unwrap(), + float_positions + .iter() + .map(|position| GridSpan::new(*position, 1)) + .collect(), + ) + .unwrap(); + let float_series = Series::try_new( + &schema(), + metric.clone(), + float_time, + SeriesFields::one(Samples::from_float(Float64Array::from(float_values))), + ) + .unwrap(); + inputs.push(vec![Step::chunk((float_series, true))]); + let child = Scripted::plan(inputs); + let context = unbounded_context(); + let plan = node(child, 1); + let mut frames = plan.execute(0, Arc::clone(&context)).unwrap(); + let mut batches = Vec::new(); + let mut referenced_histogram_rows = 0usize; + while let Some(batch) = frames.try_next().await.unwrap() { + let fields = batch + .column(3) + .as_any() + .downcast_ref::() + .unwrap(); + let payload = fields + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let hist = payload + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let tables = hist.column(1).as_any().downcast_ref::().unwrap(); + for row in 0..tables.len() { + referenced_histogram_rows += tables.value_length(row) as usize; + } + batches.push(batch); + } + assert_eq!(128, referenced_histogram_rows, "table_count={table_count}"); + + let frame_schema = batches[0].schema(); + let frames = Box::pin(RecordBatchStreamAdapter::new( + frame_schema, + futures::stream::iter(batches.into_iter().map(Ok)), + )); + let mut decoded = unpack_evaluated_series(frames, Some(schema())).unwrap(); + let mut chunks = Vec::new(); + while let Some(chunk) = decoded.try_next().await.unwrap() { + chunks.push(chunk); + } + let mut next_position = 0; + let mut final_count = 0; + for (series, is_last) in chunks { + assert_eq!(metric, series.metric, "full input identity survives"); + final_count += usize::from(is_last); + assert_eq!(series.time.len(), series.fields.get(0).unwrap().len()); + for (position, row) in series.time.positions() { + assert_eq!( + next_position, position, + "positions stay strictly increasing and complete" + ); + if position % 2 == 0 { + let table = (position / 2) % table_count; + match series.fields.get(0).unwrap().sample(row) { + datatypes::series::samples::SampleRef::Histogram(histogram) => { + assert_eq!(10.0 + table as f64, histogram.sum); + assert_eq!(4.0, histogram.count); + assert_eq!(0.0, histogram.zero_threshold); + assert_eq!(0.0, histogram.zero_count); + assert_eq!(0, histogram.schema); + assert_eq!(CounterResetHint::Unknown, histogram.reset_hint); + assert_eq!(&[0.5, 1.5, 2.5, 3.5], histogram.custom_values); + assert_eq!(&[1.0, 2.0, 3.0, 4.0], histogram.positive.buckets); + assert_eq!(&[1.0, 2.0, 3.0, 4.0], histogram.negative.buckets); + } + _ => panic!("histogram at {position}"), + } + } else { + match series.fields.get(0).unwrap().sample(row) { + datatypes::series::samples::SampleRef::Float(value) => { + assert_eq!(position as f64, value); + } + _ => panic!("float at {position}"), + } + } + next_position += 1; + } + } + assert_eq!(256, next_position, "the full grid survives decoding"); + assert_eq!(1, final_count, "the merged identity closes exactly once"); + assert_eq!(128, hist_positions.len()); + } +} + +#[tokio::test] +async fn packed_child_frame_decodes_all_256_closing_chunks() { + let metric = labels("packed-child"); + let child = Scripted::plan(vec![ + (0..256) + .map(|position| Step::chunk(chunk(metric.clone(), position, 1, position == 255))) + .collect(), + ]); + let mut frames = child + .execute( + 0, + task_context(&Arc::new(GreedyMemoryPool::new(usize::MAX))), + ) + .unwrap(); + let frame = frames + .try_next() + .await + .unwrap() + .expect("child packs all scripted chunks into one frame"); + assert_eq!(256, frame.num_rows(), "one child frame contains 256 chunks"); + assert!(frames.try_next().await.unwrap().is_none()); + + let plan = node(child, 1); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let mut output = output_of(&plan, 0, &context); + let (chunks, errors) = until_end(&mut output).await; + assert!( + errors.is_empty(), + "decoded packed child succeeds: {errors:?}" + ); + assert_eq!(256, chunks.len()); + for (position, (series, is_last)) in chunks.iter().enumerate() { + assert_eq!(named(&metric), named(&series.metric)); + assert_eq!(Some(position), series.time.first_position()); + assert_eq!(1, series.time.len()); + assert_eq!(position == 255, *is_last); + let Some(datatypes::series::samples::DenseSamplesRef::Float(values)) = + series.fields.get(0).unwrap().segments()[0].dense() + else { + panic!("decoded packed chunk has float samples"); + }; + assert_eq!(position as f64, values.values()[0]); + } +} + +#[tokio::test] +async fn histogram_chain_slices_fit_a_budget_below_the_full_table_charge() { + let metric = labels("histogram-sliced-chain"); + let steps = (0..256) + .map(|position| { + let (series, _) = histogram_chunk(metric.clone(), position); + Step::Chunk(series, position == 255) + }) + .collect::>(); + let child = Scripted::plan(vec![steps]); + let probe_context = task_context(&Arc::new(GreedyMemoryPool::new(usize::MAX))); + let mut frames = child.execute(0, Arc::clone(&probe_context)).unwrap(); + let batch = frames + .try_next() + .await + .unwrap() + .expect("one histogram frame"); + assert_eq!( + 256, + batch.num_rows(), + "all histogram chunks share one frame" + ); + assert!(frames.try_next().await.unwrap().is_none()); + drop(batch); + let mut probe = + unpack_evaluated_series(child.execute(0, Arc::clone(&probe_context)).unwrap(), None) + .unwrap(); + let mut full_table_bytes = 0usize; + let mut slice_bytes = 0usize; + let mut base_bytes = 0usize; + for _ in 0..256 { + let (series, _) = futures::future::poll_fn(|cx| probe.as_mut().poll_next(cx)) + .await + .unwrap() + .unwrap(); + for segment in series.fields.iter().flat_map(|field| field.segments()) { + if let Some(datatypes::series::samples::DenseSamplesRef::Histogram(histograms)) = + segment.dense() + { + full_table_bytes = + full_table_bytes.max(histograms.batch().as_arrow().get_array_memory_size()); + slice_bytes += histogram_table_slice_bytes(histograms.batch()).unwrap(); + } + } + base_bytes += operand_chunk_bytes(&series); + } + // Each histogram row has 48 scalar bytes and 112 list-value bytes (4 custom f64 values, + // 2 i32 offsets, 2 i32 lengths and 8 f64 buckets), independently of this accounting helper. + const HISTOGRAM_ROW_PAYLOAD_BYTES: usize = 160; + let independent_slice_lower_bound = HISTOGRAM_ROW_PAYLOAD_BYTES * 256; + let independent_slice_upper_bound = + 256 * (std::mem::size_of::() + 168 + 7 * 8); + assert!( + (independent_slice_lower_bound..=independent_slice_upper_bound).contains(&slice_bytes), + "measured slice charge {slice_bytes} is outside the independent bounds {independent_slice_lower_bound}..={independent_slice_upper_bound}" + ); + let old_charge = full_table_bytes * 256; + // Include independent bounded per-row/per-fragment bookkeeping workspace; the budget remains + // below charging the full shared backing table once for every physical slice. + let budget = base_bytes + independent_slice_upper_bound + 512 * 1024; + assert!( + budget < base_bytes + old_charge, + "slice budget {budget} is below duplicated full-table charge {}", + base_bytes + old_charge + ); + assert!( + budget > base_bytes + independent_slice_lower_bound, + "finite budget must allow the base chunks and independently bounded histogram payload" + ); + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let context = task_context(&pool); + let shared = + SharedExecution::start(&child, Routing::new(1), Arc::clone(&schema()), &context).unwrap(); + let (receiver, lease) = SharedExecution::claim(&shared, 0).expect("claim output"); + let mut output = RepartitionStream::new( + receiver, + lease, + Arc::clone(&schema()), + &(Arc::clone(&pool) as Arc), + 1, + ); + wait_until("histogram producer queues the whole chain", || { + producer_stopped(&shared) + }) + .await; + assert!( + pool.reserved() >= base_bytes + independent_slice_lower_bound, + "queued charge {} must include chunk bases and the independently-derived histogram payload bound {}", + pool.reserved(), + base_bytes + independent_slice_lower_bound + ); + let mut chunks = Vec::new(); + while let Some(item) = tokio::time::timeout( + Duration::from_secs(20), + futures::future::poll_fn(|cx| Pin::new(&mut output).poll_next(cx)), + ) + .await + .unwrap() + { + chunks.push(item.unwrap()); + } + assert_eq!( + 1, + chunks.len(), + "one owner series spans all physical slices" + ); + assert!(chunks[0].1, "the complete owner closes once"); + assert_eq!( + 256, + chunks[0].0.time.len(), + "all histogram rows survive merge" + ); + assert!(matches!(shared.ending(), Ending::Finished)); + assert!( + pool.reserved() >= base_bytes + independent_slice_lower_bound, + "internal handoff retains payload and histogram slice charges" + ); + drop(chunks); + drop(output); + assert_eq!( + 0, + pool.reserved(), + "dropping internal output releases every retained slice charge" + ); +} + +#[tokio::test] +async fn histogram_table_memory_is_charged_and_released_with_retained_chain() { + let metric = labels("histogram"); + let (series, is_last) = histogram_chunk(metric.clone(), 0); + let measured_table = match series.fields.get(0).unwrap().segments()[0].dense() { + Some(datatypes::series::samples::DenseSamplesRef::Histogram(histograms)) => { + histograms.batch().as_arrow().get_array_memory_size() + } + _ => panic!("histogram sample expected"), + }; + let child = Scripted::plan(vec![vec![Step::Chunk(series, is_last)]]); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let plan = node(child, 1); + let mut output = output_of(&plan, 0, &context); + let result = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .expect("histogram output completes"); + assert!(result.1.is_empty()); + assert_eq!(0, pool.reserved(), "success releases queued reservation"); + assert!(measured_table > 0, "test table payload is non-empty"); + + let (first, _) = histogram_chunk(metric.clone(), 0); + let mut second = first.clone(); + second.time = TimeAxis::from_grid_spans( + schema().eval_grid.clone().unwrap(), + vec![GridSpan::new(1, 1)], + ) + .expect("second histogram chunk advances on the grid"); + let pool = Arc::new(GreedyMemoryPool::new(usize::MAX)); + let context = task_context(&pool); + let plan = node( + Scripted::plan(vec![vec![ + Step::Chunk(first.clone(), false), + Step::Chunk(second.clone(), true), + ]]), + 1, + ); + let mut output = output_of(&plan, 0, &context); + let current_execution = execution(&plan); + wait_until("producer queued both shared-table chunks", || { + producer_stopped(¤t_execution) + }) + .await; + let probe_child = Scripted::plan(vec![vec![ + Step::Chunk(first.clone(), false), + Step::Chunk(second.clone(), true), + ]]); + let probe_stream = probe_child + .execute( + 0, + task_context(&Arc::new(GreedyMemoryPool::new(usize::MAX))), + ) + .unwrap(); + let mut probe_output = unpack_evaluated_series(probe_stream, None).unwrap(); + let mut decoded_costs = Vec::new(); + for _ in 0..2 { + let (series, _) = futures::future::poll_fn(|cx| probe_output.as_mut().poll_next(cx)) + .await + .unwrap() + .unwrap(); + let histogram_segments = series + .fields + .iter() + .flat_map(|field| field.segments()) + .filter_map(|segment| match segment.dense() { + Some(datatypes::series::samples::DenseSamplesRef::Histogram(histograms)) => { + Some(histograms) + } + _ => None, + }) + .collect::>(); + assert!( + !histogram_segments.is_empty(), + "decoded chunk contains histogram segments" + ); + decoded_costs.push( + operand_chunk_bytes(&series) + + histogram_segments + .iter() + .map(|histograms| histogram_table_slice_bytes(histograms.batch()).unwrap()) + .sum::(), + ); + } + let decoded_chunk_cost = decoded_costs[0]; + assert_eq!(decoded_chunk_cost, decoded_costs[1]); + let expected = 2 * decoded_chunk_cost; + drop(probe_output); + assert!( + pool.reserved() >= expected, + "reservation includes both chunk payloads and independent shared-table slice charges" + ); + let first = futures::future::poll_fn(|cx| output.as_mut().poll_next(cx)) + .await + .unwrap() + .unwrap(); + drop(first); + drop(output); + drop(plan); + assert_eq!( + 0, + pool.reserved(), + "dropping internal output releases all shared-table reservations" + ); + + let (many, _) = histogram_rows_chunk(labels("histogram-many"), 0, 256); + let probe_child = Scripted::plan(vec![vec![Step::Chunk(many.clone(), true)]]); + let probe_context = task_context(&Arc::new(GreedyMemoryPool::new(usize::MAX))); + let decoded = unpack_evaluated_series( + probe_child.execute(0, Arc::clone(&probe_context)).unwrap(), + None, + ); + let mut probe = decoded.unwrap(); + let (decoded, _) = futures::future::poll_fn(|cx| probe.as_mut().poll_next(cx)) + .await + .unwrap() + .unwrap(); + let base = operand_chunk_bytes(&decoded); + let table_cost = decoded + .fields + .iter() + .flat_map(|field| field.segments()) + .map(|segment| match segment.dense() { + Some(datatypes::series::samples::DenseSamplesRef::Histogram(histograms)) => { + histogram_table_slice_bytes(histograms.batch()).unwrap() + } + _ => 0, + }) + .sum::(); + let budget = base + table_cost - 1; + let pool = Arc::new(GreedyMemoryPool::new(budget)); + let context = task_context(&pool); + let plan = node(Scripted::plan(vec![vec![Step::Chunk(many, true)]]), 1); + let mut output = output_of(&plan, 0, &context); + let (_, errors) = tokio::time::timeout(Duration::from_secs(10), until_end(&mut output)) + .await + .unwrap(); + assert_eq!(1, errors.len(), "underbudget histogram slice is refused"); + assert_eq!(0, pool.reserved(), "refusal releases slice charge"); +} + +/// a chain that outgrows the pool ends every output, cancels its peers and releases the memory. +/// +/// This is what the whole-chain handover costs unless it is charged: a chain is buffered whole before +/// it is routed, so a child that keeps one series open and never closes it would grow the buffer +/// without bound. Here the chain is charged to the pool as it grows, and the pool is sized from the +/// chunks rather than from a number picked by hand. +/// +/// The arithmetic is measured rather than assumed, and both sides of it matter. `per_chunk` is what +/// one chunk costs the read, measured through the same pack/unpack path. +/// The chain is `CHUNKS` of those, the pool holds `CHUNKS / 2` of them, so one chunk fits with room +/// to spare and the chain as a whole does not: the growth that is refused is the one that would have +/// made the buffer exceed the pool, at the `CHUNKS / 2 + 1`-th chunk. A change to what a chunk costs +/// moves the pool with it instead of silently making the refusal happen at the first chunk or not at +/// all - which is what a hard-coded byte count did here. +/// +/// The child hands one chunk over per poll ([`Behaviour::OneAtATime`]) for a reason that is not +/// cosmetic: what a producer hands over in one poll reaches the read as one chunk, so back-to-back +/// chunks would coalesce into a single one and there would be no growing chain to refuse. +/// +/// What is asserted is the whole of it: every output ends with the refusal rather than a clean end, +/// the partition parked on a child that never answers is cancelled, and the pool is back to nothing +/// once the read is over - the chain it refused to grow is not left charged for. +#[tokio::test] +async fn a_chain_that_outgrows_the_pool_ends_every_output_and_releases_it() { + const CHUNKS: usize = 8; + let shared = labels("web-1"); + // One open chain: every chunk continues the series and none of them closes it, so the read goes on + // reading until the pool refuses to hold more of it. + let steps = (0..CHUNKS) + .map(|index| Step::chunk(chunk(shared.clone(), index, 1, false))) + .collect::>(); + // What one chunk of this chain costs, measured through the read's own path on an unbounded pool. + // The probe closes its chain, so what it measures is a whole stream's worth of chunks rather than + // the error a producer that never closes one leaves behind. + let per_chunk = { + let probe = Fragile::plan( + vec![Behaviour::OneAtATime(vec![ + Step::chunk(chunk(shared.clone(), 0, 1, false)), + Step::chunk(chunk(shared.clone(), 1, 1, true)), + ])], + Arc::new(Switch::default()), + ); + let probe_context = unbounded_context(); + let mut probe_chunks = unpack_evaluated_series( + probe + .plan + .execute(0, Arc::clone(&probe_context)) + .expect("the probe partition executes"), + None, + ) + .expect("the probe declares an evaluated frame"); + let (chunks, errors) = + tokio::time::timeout(Duration::from_secs(10), until_end(&mut probe_chunks)) + .await + .expect("the probe chain ends, so its stream ends"); + assert!( + errors.is_empty(), + "the probe chain is a whole one, so nothing about it fails: {errors:?}" + ); + assert_eq!( + 2, + chunks.len(), + "one chunk per poll is what keeps the frame boundary from coalescing two chunks into one" + ); + chunks + .iter() + .map(|(series, _)| operand_chunk_bytes(series)) + .max() + .expect("the probe chain has chunks of its own") + }; + let pool_size = per_chunk * (CHUNKS / 2); + assert!( + per_chunk >= 1 && pool_size < per_chunk * CHUNKS, + "the pool has to hold at least one chunk and fewer than all of them, or this test refuses \ + nothing: {pool_size} bytes against {per_chunk} per chunk" + ); + let pool = Arc::new(GreedyMemoryPool::new(pool_size)); + let context = task_context(&pool); + let child = Fragile::plan( + vec![Behaviour::OneAtATime(steps), Behaviour::Parked], + Arc::new(Switch::default()), + ); + let plan = node(child.plan, 2); + + let outputs = tokio::time::timeout( + Duration::from_secs(10), + every_output_until_end(&plan, &context), + ) + .await + .expect("a reservation that is refused ends the read instead of growing past the pool"); + for (partition, (chunks, errors)) in outputs.iter().enumerate() { + assert!( + chunks.is_empty(), + "output partition {partition} was given the first {n} chunks of a chain that was never \ + routed", + n = chunks.len() + ); + assert_eq!( + 1, + errors.len(), + "output partition {partition} reports the refusal once: {errors:?}" + ); + assert!( + errors[0].contains("Failed to allocate additional"), + "output partition {partition} ended with {} rather than the refusal", + errors[0] + ); + } + wait_until("the parked partition's child is dropped", || { + child.dropped[1].load(Ordering::SeqCst) == 1 + }) + .await; + let execution = execution(&plan); + wait_until("the read's tasks are gone", || producer_stopped(&execution)).await; + wait_until("what the refused chain reserved is released", || { + pool.reserved() == 0 + }) + .await; +} diff --git a/src/promql/src/native/execute.rs b/src/promql/src/native/execute.rs index d2aedc369f1..bf0f4e40de7 100644 --- a/src/promql/src/native/execute.rs +++ b/src/promql/src/native/execute.rs @@ -155,7 +155,7 @@ pub(super) fn evaluated_bytes(chunk: &EvaluatedChunk) -> usize { /// allocation however many chunks index into it, so it is not a per-chunk cost at all. Dropping it /// here is therefore only half the story -- it must not become a figure the operand never charges. /// [`unbilled_tables`] is the other half, and the operand bills each table once through it. -fn operand_chunk_bytes(series: &Series) -> usize { +pub(crate) fn operand_chunk_bytes(series: &Series) -> usize { series_bytes(series) .saturating_sub(fields_bytes(&series.fields)) .saturating_add(sample_slice_bytes(&series.fields)) @@ -243,7 +243,7 @@ fn unbilled_tables(series: &Series, billed: &HashSet) -> (Vec, usi /// /// Existing values and histogram tables remain shared with the fragments, whose estimates already /// account for their payloads. -fn merge_workspace_bytes(rows: usize, fragments: usize, fields: usize) -> usize { +pub(crate) fn merge_workspace_bytes(rows: usize, fragments: usize, fields: usize) -> usize { let sorted_rows = rows.saturating_mul( std::mem::size_of::<(usize, usize, usize)>() + std::mem::size_of::(), );