From 6f3e66d4fcf89733a21d730d3e90fcb493a72dd0 Mon Sep 17 00:00:00 2001 From: waralexrom <108349432+waralexrom@users.noreply.github.com> Date: Tue, 25 Aug 2026 12:03:34 +0200 Subject: [PATCH 1/2] fix(cubestore): inline aggregate dropped rows past the first partition (#11631) * fix(cubestore): recompute inline aggregate plan properties on input swap InlineAggregateExec carried its PlanProperties over to a new child instead of recomputing them. The node is converted from an AggregateExec while a CoalescePartitionsExec sits under it, so the cached output partitioning is 1; EnforceSorting then strips that coalesce as an avoidable bottleneck and reattaches the multi-partition subtree. The stale count stays 1, the parent executes only partition 0, and every row in the remaining partitions is dropped with no error. Delegate the input swap to the AggregateExec the node was built from, then re-convert. DataFusion recomputes the properties there and preserves the output schema, which it derives from aggregate expression names its own rules may rewrite. When the new input is no longer sorted on the group keys the plain hash aggregate is kept, without the limit: on the streaming path it counts complete groups off a sorted input, while a hash aggregate reads it as a cap on distinct groups and would truncate arbitrary ones. Reachable only when every group-by column is pinned to a single value by the filter (which is what makes the aggregate sorted), the leading sort key column is outside the filter set so partition pruning cannot collapse the scan, and the matching rows live outside the first partition. A rollup partitioned by a time dimension that the query does not filter is exactly that shape. Also adds an exhaustive check that MinMaxCondition never prunes a range holding a matching row -- partition pruning was the first suspect and is clean. Co-Authored-By: Claude Opus 5 * refactor(cubestore): derive group-by-limit aggregate properties from the new child GroupByLimitAggregateExec refreshed only the output partitioning when its input was swapped, leaving the rest of the plan properties -- orderings and constants projected from the previous input -- behind. Delegate the swap to the aggregate it was built from, the same way the inline aggregate does, so DataFusion recomputes all of them and preserves the output schema. If the rebuilt aggregate no longer fits the trimming path it is returned as is. The other two nodes that carry their properties over, RollingWindowAggExec and AggregateTopKExec, build them from their own output schema plus constants and read the input's partition count live, so nothing there can go stale; say so at both sites, since the same pattern loses rows in a node whose partitioning follows the input. Co-Authored-By: Claude Opus 5 * test(cubestore): pin the aggregate node behind the partitioning fix Review follow-ups on the inline aggregate fix. The end-to-end regression test asserted only on rows, so a planner change that moved the query off the streaming aggregate, or collapsed the scan to one partition, would have kept it green while covering nothing. Assert the plan shape too, and add the unit-level counterpart the sibling node already has: build the exec over a one-partition input, re-child it onto three, and require both the reported partitioning and the node type to follow. It runs in milliseconds and does not depend on compaction timing. Both fallbacks -- to the hash aggregate here, out of the trimming path in GroupByLimitAggregateExec -- now name the group by and the new input's order mode, so a degradation is diagnosable from the log rather than showing up only as wrong-looking timings. The group-by-limit test also pins that re-childing keeps the trimming exec, which the new fallback would otherwise let it pass without. Co-Authored-By: Claude Opus 5 --------- Co-authored-by: Claude Opus 5 --- .../group_by_limit_aggregate/mod.rs | 90 +++++++++---- .../src/queryplanner/inline_aggregate/mod.rs | 116 +++++++++++++--- .../src/queryplanner/partition_filter.rs | 125 ++++++++++++++++++ .../cubestore/src/queryplanner/rolling.rs | 3 + .../src/queryplanner/topk/execute.rs | 3 + rust/cubestore/cubestore/src/sql/mod.rs | 113 ++++++++++++++++ 6 files changed, 410 insertions(+), 40 deletions(-) diff --git a/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs b/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs index bd5a937239bc9..cc37a2b620ff2 100644 --- a/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs +++ b/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs @@ -4,7 +4,7 @@ use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::stats::Precision; use datafusion::common::Statistics; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; use datafusion::physical_expr::{Distribution, LexRequirement}; @@ -16,7 +16,7 @@ use datafusion::physical_plan::{ PlanProperties, SendableRecordBatchStream, }; use std::any::Any; -use std::fmt::Debug; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; /// Worker-side partial hash aggregate that trims its output to the top-k groups by a total order, @@ -32,7 +32,7 @@ use std::sync::Arc; /// columns), expressed as `(partial-output column index, sort options)`. A total order is required /// for correctness: the same group key can live on multiple workers, and a consistent cut across /// workers guarantees every partial state the router selects reaches it. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct GroupByLimitAggregateExec { group_by: PhysicalGroupBy, aggr_expr: Vec>, @@ -48,6 +48,28 @@ pub struct GroupByLimitAggregateExec { factor: usize, /// Total order over the partial output columns. order: Vec<(usize, SortOptions)>, + /// The aggregate this node replaced. Kept so that swapping the input can be delegated to it: + /// DataFusion recomputes the plan properties there and preserves the output schema. + source: Arc, +} + +/// Skips [GroupByLimitAggregateExec::source]: it carries a copy of the same input subtree the node +/// already prints, so deriving would double the output of every nested aggregate. +impl Debug for GroupByLimitAggregateExec { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GroupByLimitAggregateExec") + .field("group_by", &self.group_by) + .field("aggr_expr", &self.aggr_expr) + .field("filter_expr", &self.filter_expr) + .field("input", &self.input) + .field("schema", &self.schema) + .field("input_schema", &self.input_schema) + .field("cache", &self.cache) + .field("k", &self.k) + .field("factor", &self.factor) + .field("order", &self.order) + .finish_non_exhaustive() + } } impl GroupByLimitAggregateExec { @@ -77,10 +99,8 @@ impl GroupByLimitAggregateExec { return None; } let input = aggregate.input().clone(); - // A partial aggregate preserves its input's partitioning (it runs once per input partition). - // Derive the output partitioning from the input rather than copying the wrapped aggregate's - // cached value, which can be stale: a later pass may swap our input for one with a different - // partition count via `with_new_children` without the cache following, and a too-low count + // A partial aggregate runs once per input partition and so preserves its partitioning. + // `aggregate.cache()` already says that; spell it out anyway, since a wrong count here // makes the parent coalesce read only some partitions and silently drop the rest. let cache = aggregate .cache() @@ -97,6 +117,7 @@ impl GroupByLimitAggregateExec { k, factor, order, + source: Arc::new(aggregate.clone()), }) } @@ -177,24 +198,39 @@ impl ExecutionPlan for GroupByLimitAggregateExec { self: Arc, children: Vec>, ) -> DFResult> { - let input = children[0].clone(); - // Track the (possibly changed) input's partitioning; a partial aggregate preserves it. - let cache = self - .cache - .clone() - .with_partitioning(input.output_partitioning().clone()); - Ok(Arc::new(Self { - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), - input, - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - cache, - k: self.k, - factor: self.factor, - order: self.order.clone(), - })) + // Swapping the input is delegated to the aggregate this node was built from, so that + // DataFusion recomputes the whole of the plan properties -- output partitioning above all, + // a stale count there makes the parent read only some of the input's partitions and + // silently drop the rest -- and preserves the output schema. + let rebuilt = Arc::clone(&self.source).with_new_children(children)?; + let Some(aggregate) = rebuilt.as_any().downcast_ref::() else { + return Err(DataFusionError::Internal(format!( + "AggregateExec::with_new_children returned {}", + rebuilt.name() + ))); + }; + Ok( + match Self::try_new_from_partial(aggregate, self.k, self.factor, self.order.clone()) { + Some(trimmed) => Arc::new(trimmed), + None => { + // Correct but a perf cliff: the worker goes from `factor * k` rows per + // partition to the full group cardinality. Say it out loud so the loss of the + // trim is diagnosable instead of showing up only as a slow query. + log::warn!( + "Rebuilt aggregate no longer fits the group-by-limit trim (input order \ + mode {:?}), dropping it for group by [{}]", + aggregate.input_order_mode(), + self.group_by + .expr() + .iter() + .map(|(_, name)| name.as_str()) + .collect::>() + .join(", ") + ); + rebuilt + } + }, + ) } fn execute( @@ -453,6 +489,10 @@ mod tests { let three = MemorySourceConfig::try_new(&three_parts, input_schema(), None).unwrap(); let three_input: Arc = Arc::new(DataSourceExec::new(Arc::new(three))); let exec3 = exec.with_new_children(vec![three_input]).unwrap(); + assert!( + exec3.as_any().is::(), + "re-childing must keep the trimming exec, otherwise this test no longer pins it" + ); assert_eq!( exec3.output_partitioning().partition_count(), 3, diff --git a/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs b/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs index 8a58d1a8c0dba..e0f031c99435d 100644 --- a/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs +++ b/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs @@ -9,7 +9,7 @@ pub use sorted_group_values_rows::SortedGroupValuesRows; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::common::stats::Precision; use datafusion::common::Statistics; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; use datafusion::physical_expr::{Distribution, LexRequirement}; @@ -22,7 +22,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use std::any::Any; -use std::fmt::Debug; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -31,7 +31,7 @@ pub enum InlineAggregateMode { Final, } -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct InlineAggregateExec { mode: InlineAggregateMode, /// Group by expressions @@ -55,6 +55,29 @@ pub struct InlineAggregateExec { pub input_schema: SchemaRef, cache: PlanProperties, required_input_ordering: Vec>, + /// The aggregate this node replaced. Kept so that swapping the input can be delegated to it: + /// DataFusion recomputes the plan properties there and preserves the output schema, which it + /// derives from aggregate expression names that its own rules are free to rewrite. + source: Arc, +} + +/// Skips [InlineAggregateExec::source]: it carries a copy of the same input subtree the node +/// already prints, so deriving would double the output of every nested aggregate. +impl Debug for InlineAggregateExec { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InlineAggregateExec") + .field("mode", &self.mode) + .field("group_by", &self.group_by) + .field("aggr_expr", &self.aggr_expr) + .field("filter_expr", &self.filter_expr) + .field("limit", &self.limit) + .field("input", &self.input) + .field("schema", &self.schema) + .field("input_schema", &self.input_schema) + .field("cache", &self.cache) + .field("required_input_ordering", &self.required_input_ordering) + .finish_non_exhaustive() + } } impl InlineAggregateExec { @@ -101,6 +124,7 @@ impl InlineAggregateExec { input_schema, cache, required_input_ordering, + source: Arc::new(aggregate.clone()), }) } @@ -192,19 +216,46 @@ impl ExecutionPlan for InlineAggregateExec { self: Arc, children: Vec>, ) -> DFResult> { - let result = Self { - mode: self.mode, - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), - limit: self.limit.clone(), - input: children[0].clone(), - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - cache: self.cache.clone(), - required_input_ordering: self.required_input_ordering.clone(), + // Swapping the input is delegated to the aggregate this node was built from, so that + // DataFusion recomputes the plan properties -- output partitioning above all. Carrying + // them over leaves a stale partition count, and the parent then executes only that many + // of the input's partitions and silently drops the rows of the rest. + let rebuilt = Arc::clone(&self.source).with_new_children(children)?; + let Some(aggregate) = rebuilt.as_any().downcast_ref::() else { + return Err(DataFusionError::Internal(format!( + "AggregateExec::with_new_children returned {}", + rebuilt.name() + ))); }; - Ok(Arc::new(result)) + + // The streaming implementation is only valid while the new input stays sorted on the + // group keys; if it no longer is, keep the plain hash aggregate. The limit stays behind + // with it: here it counts complete groups off a sorted input, while a hash aggregate + // reads it as a cap on distinct groups and would stop mid-way through arbitrary ones. + // Dropping it only costs the early exit. + Ok(match Self::try_new_from_aggregate(aggregate) { + Some(inline) => Arc::new(inline.with_limit(self.limit)), + None => { + // No pass in the pipeline is known to de-sort an inline aggregate's input, and a + // parent merge built around the old ordering would not notice the switch, so say + // it out loud rather than degrade quietly. + log::warn!( + "Input of an inline {:?} aggregate is no longer sorted on the group keys, \ + falling back to a hash aggregate. Group by [{}], new input order mode {:?}, \ + new input ordering {:?}", + self.mode, + self.group_by + .expr() + .iter() + .map(|(_, name)| name.as_str()) + .collect::>() + .join(", "), + aggregate.input_order_mode(), + aggregate.input().properties().output_ordering() + ); + Arc::new(aggregate.clone().with_limit(None)) + } + }) } fn execute( @@ -408,6 +459,41 @@ mod tests { } /// A group continuing in the next input batch must not be emitted early with a partial sum. + /// A partial aggregate runs once per input partition, so its reported partitioning must + /// follow a re-childed input. A stale single-partition count makes the parent coalesce + /// execute only partition 0 and silently drop the rows of the rest -- the row loss this node + /// was fixed for. Built over a 1-partition input so its cache says 1, then re-childed onto 3. + #[test] + fn output_partitioning_follows_rechilded_input() { + let schema = test_schema(); + let one = sorted_source( + &schema, + vec![vec![make_batch(&schema, &[(1, 10), (2, 20)])]], + ); + let exec = partial_sum_inline_aggregate(one, None); + assert_eq!(exec.properties().output_partitioning().partition_count(), 1); + + let three: Vec> = (0..3) + .map(|i| vec![make_batch(&schema, &[(i, 10), (i + 10, 20)])]) + .collect(); + let rechilded = exec + .with_new_children(vec![sorted_source(&schema, three)]) + .unwrap(); + + assert!( + rechilded.as_any().is::(), + "re-childing a still-sorted input must keep the streaming exec" + ); + assert_eq!( + rechilded + .properties() + .output_partitioning() + .partition_count(), + 3, + "exec must report its re-childed input's partition count, not a stale 1" + ); + } + #[test] fn limit_emits_only_closed_groups() { let schema = test_schema(); diff --git a/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs b/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs index d59b2f396fcf8..a619a4fe2ad6d 100644 --- a/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs +++ b/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs @@ -1531,6 +1531,131 @@ mod tests { Vec::new() } } + + /// Exhaustive check that [MinMaxCondition] never prunes a range that holds a matching row. + /// + /// Enumerates every per-column bound combination over a small integer domain against every + /// ordered pair of min/max rows, and compares the verdict with a direct search for a tuple + /// that both satisfies the bounds and lies inside the range. A false negative here means a + /// query silently returning fewer rows than it should. + fn assert_no_false_negatives(n: usize, d: i64) { + // Rows and bounds use 0..d; candidate tuples reach one step further on both sides to + // stand in for the unbounded parts of the domain. + let rows = cross(n, &(0..d).collect::>()); + let candidates = cross(n, &(-1..=d).collect::>()); + + for cond in column_bounds(n, d) { + let min_max = MinMaxCondition { + min: cond.iter().map(|(mn, _)| mn.map(TableValue::Int)).collect(), + max: cond.iter().map(|(_, mx)| mx.map(TableValue::Int)).collect(), + }; + let matching = candidates + .iter() + .filter(|t| satisfies(&cond, t)) + .collect::>(); + + for mn in &rows { + for mx in &rows { + if mn > mx { + continue; + } + let min_row = mn.iter().cloned().map(TableValue::Int).collect::>(); + let max_row = mx.iter().cloned().map(TableValue::Int).collect::>(); + + if matching.iter().any(|t| *t >= mn && *t <= mx) { + assert!( + min_max.can_match(&min_row, &max_row), + "pruned a matching range: cond {:?}, min_row {:?}, max_row {:?}", + cond, + mn, + mx + ); + } + if matching.iter().any(|t| *t >= mn) { + assert!( + min_max.can_match_min(&min_row), + "pruned a matching range with unbounded max: cond {:?}, min_row {:?}", + cond, + mn + ); + } + if matching.iter().any(|t| *t <= mx) { + assert!( + min_max.can_match_max(&max_row), + "pruned a matching range with unbounded min: cond {:?}, max_row {:?}", + cond, + mx + ); + } + } + } + } + } + + /// All tuples of length `n` over `vals`. + fn cross(n: usize, vals: &[i64]) -> Vec> { + let mut r = vec![vec![]]; + for _ in 0..n { + r = r + .iter() + .flat_map(|t| { + vals.iter().map(move |v| { + let mut t = t.clone(); + t.push(*v); + t + }) + }) + .collect(); + } + r + } + + /// All per-column `(min, max)` bound combinations over `0..d`, `None` meaning unbounded. + fn column_bounds(n: usize, d: i64) -> Vec, Option)>> { + let mut per_col = Vec::new(); + for mn in 0..=d { + for mx in 0..=d { + let mn = (mn != d).then_some(mn); + let mx = (mx != d).then_some(mx); + if let (Some(mn), Some(mx)) = (mn, mx) { + if mn > mx { + continue; + } + } + per_col.push((mn, mx)); + } + } + let mut r = vec![vec![]]; + for _ in 0..n { + r = r + .iter() + .flat_map(|t| { + per_col.iter().map(move |c| { + let mut t = t.clone(); + t.push(*c); + t + }) + }) + .collect(); + } + r + } + + fn satisfies(cond: &[(Option, Option)], t: &[i64]) -> bool { + cond.iter() + .zip(t) + .all(|((mn, mx), v)| mn.map_or(true, |mn| mn <= *v) && mx.map_or(true, |mx| *v <= mx)) + } + + #[test] + fn no_false_negatives_two_columns() { + assert_no_false_negatives(2, 3); + } + + #[test] + fn no_false_negatives_three_columns() { + assert_no_false_negatives(3, 3); + } } struct ColumnStat { diff --git a/rust/cubestore/cubestore/src/queryplanner/rolling.rs b/rust/cubestore/cubestore/src/queryplanner/rolling.rs index 7eac8f08dc4aa..61c2223819bb9 100644 --- a/rust/cubestore/cubestore/src/queryplanner/rolling.rs +++ b/rust/cubestore/cubestore/src/queryplanner/rolling.rs @@ -709,6 +709,9 @@ impl ExecutionPlan for RollingWindowAggExec { ) -> Result, DataFusionError> { assert_eq!(children.len(), 1); Ok(Arc::new(RollingWindowAggExec { + // Safe to carry over: these are built from this node's own output schema plus + // constants, nothing in them is derived from the input. A node whose properties do + // follow the input -- output partitioning above all -- must recompute them here. properties: self.properties.clone(), sorted_input: children.remove(0), group_key: self.group_key.clone(), diff --git a/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs b/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs index eca963b705d41..67bfee7cfb1bc 100644 --- a/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs +++ b/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs @@ -190,6 +190,9 @@ impl ExecutionPlan for AggregateTopKExec { having: self.having.clone(), cluster, schema: self.schema.clone(), + // Safe to carry over: built from this node's own output schema plus constants, and + // the input's partition count is read live in `execute`. A node whose properties do + // follow the input -- output partitioning above all -- must recompute them here. cache: self.cache.clone(), sort_requirement: self.sort_requirement.clone(), merge_version: self.merge_version, diff --git a/rust/cubestore/cubestore/src/sql/mod.rs b/rust/cubestore/cubestore/src/sql/mod.rs index 0ce3abad3b156..3c777294e136c 100644 --- a/rust/cubestore/cubestore/src/sql/mod.rs +++ b/rust/cubestore/cubestore/src/sql/mod.rs @@ -6826,6 +6826,119 @@ mod tests { .await; Ok(()) } + /// A grouped aggregate over an index whose sort key does not start with the filtered + /// columns: every group column is pinned to a single value by the filter, which makes the + /// aggregate sorted and hands it to the streaming implementation. Its output partitioning + /// must follow the scan's, otherwise only the first partition is read and the rows living + /// in the others are silently dropped. + #[tokio::test] + async fn single_value_equals_scans_every_partition() -> Result<(), CubeError> { + Config::test("single_value_equals_scans_every_partition") + .update_config(|mut c| { + c.compaction_chunks_count_threshold = 0; + c.partition_split_threshold = 400; + c.max_partition_split_threshold = 400; + c + }) + .start_test(async move |services| { + let service = services.sql_service; + service.exec_query("CREATE SCHEMA pa").await?.collect().await?; + // The time dimension leads the sort key, so the equality-filtered columns are + // not an index prefix -- the shape of a partitioned rollup. + service + .exec_query( + "CREATE TABLE pa.rollup (date_code timestamp, equipment_type text, \ + month_code text, site_id text, hours int)", + ) + .await? + .collect() + .await?; + + let mut rows = Vec::new(); + for month in 1..=12 { + for day in 1..=8 { + for equipment_type in &["QC", "RTG", "STS"] { + for site_id in &["ABC", "PSE", "XYZ"] { + rows.push(format!( + "('2026-{:02}-{:02}T00:00:00.000', '{}', '2026{:02}', '{}', {})", + month, day, equipment_type, month, site_id, 10 + )); + } + } + } + } + for chunk in rows.chunks(300) { + service + .exec_query(&format!( + "INSERT INTO pa.rollup (date_code, equipment_type, month_code, \ + site_id, hours) VALUES {}", + chunk.join(", ") + )) + .await? + .collect() + .await?; + } + // Rows reachable only through the first of several scanned partitions is what + // the regression is about, so wait for the data to be split across partitions + // and fail loudly rather than silently testing a single-partition scan. + // Compaction settles in a few iterations; the bound is a ceiling, not a wait. + let mut scanned_partitions = TableValue::Int(0); + for _ in 0..1200 { + tokio::time::sleep(Duration::from_millis(50)).await; + let partitions = service + .exec_query( + "SELECT count(*) FROM system.partitions \ + WHERE active = true AND main_table_row_count > 0", + ) + .await? + .collect() + .await?; + scanned_partitions = partitions.get_rows()[0].values()[0].clone(); + if scanned_partitions > TableValue::Int(1) { + break; + } + } + assert!( + scanned_partitions > TableValue::Int(1), + "the scan must span several partitions, got {:?}", + scanned_partitions + ); + + // The last month of the range, so the matching rows sit in the last partition + // the scan visits and never in the first one -- reading only the first partition + // has to show up as a missing row rather than by chance returning the right one. + let query = "SELECT site_id, CAST(month_code AS INT) mc, sum(hours) v \ + FROM pa.rollup \ + WHERE (equipment_type = 'QC') \ + AND (CAST(month_code AS INT) = '202612') \ + AND (site_id = 'PSE') GROUP BY 1, 2"; + + // Pinning every group column to a single value is what hands the aggregate to + // the streaming implementation, and only that one carried a stale partition + // count. Without this the test would keep passing on a plan that never exercises + // the fix. + let worker_plan = pp_phys_plan(service.plan_query(query).await?.worker.as_ref()); + assert!( + worker_plan.contains("InlinePartialAggregate"), + "expected a streaming partial aggregate, got:\n{}", + worker_plan + ); + + let single_value = service.exec_query(query).await?.collect().await?; + assert_eq!( + single_value.get_rows(), + &vec![Row::new(vec![ + TableValue::String("PSE".to_string()), + TableValue::Int(202612), + TableValue::Int(80) + ])] + ); + + Ok::<(), CubeError>(()) + }) + .await; + Ok(()) + } } impl SqlServiceImpl { From cf660485c945ae5c78ed9f704445e85247320032 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Tue, 25 Aug 2026 14:38:07 +0200 Subject: [PATCH 2/2] refactor(query-orchestrator): Queue - remove processingId locks (#11638) --- .../src/queue-driver.interface.ts | 21 ++- .../src/CubeStoreQueueDriver.ts | 11 +- .../cubejs-query-orchestrator/DEVELOPMENT.md | 18 +-- .../LocalQueueDriverConnection.ts | 69 ++++----- .../src/orchestrator/QueryQueue.ts | 41 ++---- .../test/benchmarks/QueueBench.abstract.ts | 1 - .../test/unit/QueryQueue.abstract.ts | 139 +++++++++++------- 7 files changed, 150 insertions(+), 150 deletions(-) diff --git a/packages/cubejs-base-driver/src/queue-driver.interface.ts b/packages/cubejs-base-driver/src/queue-driver.interface.ts index bbded5bb199fb..5980031cd59b2 100644 --- a/packages/cubejs-base-driver/src/queue-driver.interface.ts +++ b/packages/cubejs-base-driver/src/queue-driver.interface.ts @@ -1,8 +1,6 @@ export type QueryDef = any; // Primary key of Queue item export type QueueId = string | number | bigint; -// The lock token of a retrieval, always the item's queueId. Only the memory driver compares it. -export type ProcessingId = string | number | bigint; export type QueryKey = (string | [string, any[]]) & { persistent?: true, }; @@ -13,21 +11,21 @@ export type GetActiveAndToProcessResponse = [active: QueryKeysTuple[], toProcess export type QueryStageStateResponse = [active: string[], toProcess: string[]] | [active: string[], toProcess: string[], defs: Record]; export type RetrieveForProcessingSuccess = [ added: unknown, - // QueueId is required for Cube Store, other providers don't support it + // Identifies the retrieved generation of the queue item. queueId: QueueId | null, active: QueryKeyHash[], pending: number, def: QueryDef, - lockAquired: true + retrieved: true ]; export type RetrieveForProcessingFail = [ added: unknown, - // QueueId is required for Cube Store, other providers don't support it + // Null when no queue item was retrieved. queueId: QueueId | null, active: QueryKeyHash[], pending: number, def: null, - lockAquired: false + retrieved: false ]; export type RetrieveForProcessingResponse = RetrieveForProcessingSuccess | RetrieveForProcessingFail | null; export type AddToQueueResponse = [ @@ -106,14 +104,13 @@ export interface QueueDriverConnectionInterface { getStalledQueries(): Promise; getQueryStageState(onlyKeys: boolean): Promise; updateHeartBeat(hash: QueryKeyHash, queueId: QueueId | null): Promise; - // Trying to acquire a lock for processing a queue item, this method can return null when - // multiple nodes tries to process the same query - retrieveForProcessing(hash: QueryKeyHash, processingId: ProcessingId): Promise; - freeProcessingLock(hash: QueryKeyHash, processingId: ProcessingId, activated: unknown): Promise; - optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, processingId: ProcessingId, queueId: QueueId | null): Promise; + // Atomically moves a queue item to active. Returns null when another node is already + // processing the query or the concurrency budget is full. + retrieveForProcessing(hash: QueryKeyHash, queueId: QueueId): Promise; + optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise; cancelQuery(queryKey: QueryKey, queueId: QueueId | null): Promise; getQueryAndRemove(hash: QueryKeyHash, queueId: QueueId | null): Promise<[QueryDef]>; - setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, processingId: ProcessingId, queueId: QueueId | null): Promise; + setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise; release(): void; // getQueriesToCancel(): Promise diff --git a/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts b/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts index 609fe1f3fef93..95563eb392041 100644 --- a/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts +++ b/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts @@ -12,7 +12,6 @@ import { AddToQueueResponse, QueryKey, QueryKeyHash, - ProcessingId, QueueId, GetActiveAndToProcessResponse, QueryKeysTuple, @@ -182,10 +181,6 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async freeProcessingLock(_hash: QueryKeyHash, _processingId: string, _activated: unknown): Promise { - // nothing to do - } - public async getActiveQueries(): Promise { const rows = await this.driver.query('QUEUE ACTIVE ?', [ this.options.redisQueuePrefix @@ -336,7 +331,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, _processingId: ProcessingId, queueId: QueueId): Promise { + public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise { await this.driver.query('QUEUE MERGE_EXTRA ? ?', [ // queryKeyHash as compatibility fallback queueId || this.prefixKey(hash), @@ -372,7 +367,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte ]; } - public async retrieveForProcessing(hash: QueryKeyHash, _processingId: string): Promise { + public async retrieveForProcessing(hash: QueryKeyHash, _queueId: QueueId): Promise { const rows = await this.driver.query('QUEUE RETRIEVE EXTENDED CONCURRENCY ? ?', [ this.options.concurrency, this.prefixKey(hash), @@ -404,7 +399,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, _processingId: ProcessingId, queueId: QueueId): Promise { + public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, queueId: QueueId): Promise { const rows = await this.driver.query('QUEUE ACK ? ?', [ // queryKeyHash as compatibility fallback queueId || this.prefixKey(hash), diff --git a/packages/cubejs-query-orchestrator/DEVELOPMENT.md b/packages/cubejs-query-orchestrator/DEVELOPMENT.md index e8141fe59ebc8..f51d405843f47 100644 --- a/packages/cubejs-query-orchestrator/DEVELOPMENT.md +++ b/packages/cubejs-query-orchestrator/DEVELOPMENT.md @@ -182,12 +182,13 @@ sequenceDiagram ## Background execution: `processQuery` → `executeQuery` `processQuery` retrieves the item and nothing else. What it hands to `sendProcessMessageFn` is a -`RetrievedQuery` — `{ queryKeyHash, queueId, processingId, queueSize, query }`, plain data on +`RetrievedQuery` — `{ queryKeyHash, queueId, queueSize, query }`, plain data on purpose, so a custom implementation can serialize it and let another process run `executeQuery`. The default implementation calls `executeQuery` in-process. -`processingId` is the lock token of the retrieval and always carries the `queueId`. Only the -memory driver compares it, Cube Store ignores it and keys every command off the `queueId`. +`queueId` identifies the specific generation of a queue item. The memory driver compares it +with the active entry before updating or acknowledging a query, while Cube Store keys those +commands directly off the `queueId`. `sendProcessMessageFn` must resolve once the hand-off is done, **not** once the query is executed: reconcile awaits it, and `executeQuery` ends with `reconcileQueue`, which is @@ -200,8 +201,8 @@ Two consequences of retrieving before the hand-off: `@` suffix matches, so `sendProcessMessageFn` is always called on the owning process for them — it just must not route them elsewhere. - A retrieved item is already active. If a custom hand-off loses the message, the item is only - recovered by the stalled-heartbeat / `TO_CANCEL` path; `freeProcessingLock` is a no-op on - Cube Store, so the retrieval cannot be cheaply undone. + recovered by the stalled-heartbeat / `TO_CANCEL` path; a successful retrieval is not undone + when the hand-off fails. A stream query is dispatched while `executeInQueue` is still running, so `waitForQueryStream` subscribes to `streamStarted` *before* the dispatch — a handler which starts fast would @@ -222,10 +223,10 @@ sequenceDiagram QueryQueue->>QueueDriver: retrieveForProcessing QueueDriver->>CubeStore: QUEUE RETRIEVE EXTENDED CONCURRENCY ?n ?path CubeStore-->>QueueDriver: RetrieveResponse - QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, lockAcquired] + QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, retrieved] Note over QueueDriver,CubeStore: The retrieval is atomic in Cube Store:
only one node moves the item to active - alt def && added && activeKeys includes our key && lockAcquired + alt def && added && activeKeys includes our key && retrieved QueryQueue-)Background: sendProcessMessageFn(RetrievedQuery) Note over QueryQueue,Background: Detached from here on: the hand-off returns,
the execution keeps running @@ -254,8 +255,7 @@ sequenceDiagram Background->>Background: reconcileQueue Note over Background: The freed concurrency slot is
immediately given to the next query else the retrieval did not succeed - QueryQueue->>QueueDriver: freeProcessingLock - Note over QueryQueue,QueueDriver: Another node is running it, or the
concurrency budget is full. No-op for Cube Store + Note over QueryQueue,QueueDriver: Another node is running it, or the
concurrency budget is full. Queue state is unchanged end ``` diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts b/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts index f2c0a14a542c3..55d4211440364 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts @@ -4,7 +4,6 @@ import { QueryKey, QueryKeyHash, QueueId, - ProcessingId, QueryDef, AddToQueueQuery, AddToQueueOptions, @@ -54,8 +53,6 @@ export class LocalQueueDriverConnectionState { public active: Record = {}; public heartBeat: Record = {}; - - public processingLocks: Record = {}; } export class LocalQueueDriverConnection implements QueueDriverConnectionInterface { @@ -223,7 +220,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac delete this.state.toProcess[queryKeyHash]; delete this.state.recent[queryKeyHash]; delete this.state.queryDef[queryKeyHash]; - delete this.state.processingLocks[queryKeyHash]; return [query]; } @@ -233,8 +229,8 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac return query; } - public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise { - if (this.state.processingLocks[queryKeyHash] !== processingId) { + public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise { + if (this.state.active[queryKeyHash]?.queueId !== queueId) { return false; } @@ -245,7 +241,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac delete this.state.toProcess[queryKeyHash]; delete this.state.recent[queryKeyHash]; delete this.state.queryDef[queryKeyHash]; - delete this.state.processingLocks[queryKeyHash]; promise.resolved = true; if (promise.resolve) { @@ -277,48 +272,44 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac } } - public async retrieveForProcessing(queryKeyHash: QueryKeyHash, processingId: ProcessingId): Promise { - let lockAcquired = false; - - if (!this.state.processingLocks[queryKeyHash]) { - this.state.processingLocks[queryKeyHash] = processingId; - lockAcquired = true; - } else { - return null; + public async retrieveForProcessing(queryKeyHash: QueryKeyHash, queueId: QueueId): Promise { + const query = this.state.queryDef[queryKeyHash]; + const activeKeys = this.queueArray(this.state.active) as QueryKeyHash[]; + + if ( + !query || + query.queueId !== queueId || + this.state.toProcess[queryKeyHash]?.queueId !== queueId || + this.state.active[queryKeyHash] || + activeKeys.length >= this.concurrency + ) { + return [ + 0, + null, + activeKeys, + Object.keys(this.state.toProcess).length, + null, + false + ]; } - let added = 0; - - if (Object.keys(this.state.active).length < this.concurrency && !this.state.active[queryKeyHash]) { - this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(processingId), queueId: processingId }; - delete this.state.toProcess[queryKeyHash]; - - added = 1; - } + this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(queueId), queueId }; + delete this.state.toProcess[queryKeyHash]; - this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId: processingId }; + this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId }; return [ - added, - this.state.queryDef[queryKeyHash]?.queueId ?? null, + 1, + query.queueId, this.queueArray(this.state.active) as QueryKeyHash[], Object.keys(this.state.toProcess).length, - this.state.queryDef[queryKeyHash], - lockAcquired + query, + true ]; } - public async freeProcessingLock(queryKeyHash: QueryKeyHash, processingId: ProcessingId, activated: any): Promise { - if (this.state.processingLocks[queryKeyHash] === processingId) { - delete this.state.processingLocks[queryKeyHash]; - if (activated) { - delete this.state.active[queryKeyHash]; - } - } - } - - public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise { - if (this.state.processingLocks[queryKeyHash] !== processingId) { + public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, queueId: QueueId): Promise { + if (this.state.active[queryKeyHash]?.queueId !== queueId || !this.state.queryDef[queryKeyHash]) { return false; } diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts index 13d39e518e0f3..d36558735f70c 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts @@ -9,7 +9,6 @@ import { QueryStageStateResponse, AddToQueueOptions, QueuePriority, - ProcessingId, RetrieveForProcessingSuccess } from '@cubejs-backend/base-driver'; import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver'; @@ -29,7 +28,6 @@ export type QueryHandlersMap = Record; export type RetrievedQuery = { queryKeyHash: QueryKeyHash; queueId: QueueId; - processingId: ProcessingId; queueSize: number; query: QueryDef; }; @@ -373,7 +371,6 @@ export class QueryQueue { return { queryKeyHash, queueId, - processingId: queueId, queueSize, query, }; @@ -817,9 +814,8 @@ export class QueryQueue { } /** - * Acquires the processing lock for the query specified by the `queryKeyHashed` and moves it to - * the active set. Returns `null` when the retrieval didn't succeed, which means another node is - * already running the query or the concurrency budget is full. + * Atomically moves the query specified by `queryKeyHashed` to the active set. Returns `null` + * when another node is already running the query or the concurrency budget is full. */ protected async retrieveQueryForProcessing(queryKeyHashed: QueryKeyHash, queueId: QueueId): Promise { const queueConnection = await this.queueDriver.createConnection(); @@ -828,16 +824,12 @@ export class QueryQueue { let activeKeys; let queueSize; let query; - let processingLockAcquired; + let retrievalSucceeded; try { - // The lock token is the queueId, every call which releases the lock has to be handed - // the same value retrieveForProcessing got - const processingId = queueId; - - const retrieveResult = await queueConnection.retrieveForProcessing(queryKeyHashed, processingId); + const retrieveResult = await queueConnection.retrieveForProcessing(queryKeyHashed, queueId); if (retrieveResult) { - [insertedCount, , activeKeys, queueSize, query, processingLockAcquired] = retrieveResult; + [insertedCount, , activeKeys, queueSize, query, retrievalSucceeded] = retrieveResult; } const activated = activeKeys && activeKeys.indexOf(queryKeyHashed) !== -1; @@ -845,7 +837,7 @@ export class QueryQueue { query = await queueConnection.getQueryDef(queryKeyHashed, null); } - if (!query || !insertedCount || !activated || !processingLockAcquired) { + if (!query || !insertedCount || !activated || !retrievalSucceeded) { // TODO Ideally streaming queries should reconcile queue here after waiting on open slot however in practice continue wait timeout reconciles faster CPU-wise // if (query?.queryHandler === 'stream') { // const [active] = await queueConnection.getQueryStageState(true); @@ -857,26 +849,22 @@ export class QueryQueue { this.logger('Skip processing', { queueId, - processingId, queryKey: query && query.queryKey || queryKeyHashed, requestId: query && query.requestId, queuePrefix: this.redisQueuePrefix, - processingLockAcquired, + retrievalSucceeded, query, insertedCount, activeKeys, activated, queryExists: !!query }); - await queueConnection.freeProcessingLock(queryKeyHashed, processingId, activated); - return null; } return { queryKeyHash: queryKeyHashed, queueId, - processingId, queueSize, query, }; @@ -901,7 +889,7 @@ export class QueryQueue { * implementation which hands the query over to another process. */ public async executeQuery(retrieved: RetrievedQuery): Promise { - const { queryKeyHash: queryKeyHashed, queueId, processingId, queueSize, query } = retrieved; + const { queryKeyHash: queryKeyHashed, queueId, queueSize, query } = retrieved; const queueConnection = await this.queueDriver.createConnection(); @@ -915,7 +903,6 @@ export class QueryQueue { const timeInQueue = (new Date()).getTime() - query.addedToQueueTime; this.logger('Performing query', { queueId, - processingId, queueSize, queryKey: query.queryKey, queuePrefix: this.redisQueuePrefix, @@ -927,7 +914,7 @@ export class QueryQueue { preAggregation: query.query?.preAggregation, addedToQueueTime: query.addedToQueueTime, }); - await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, processingId, queueId); + await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, queueId); let queryProcessHeartbeat = Date.now(); const heartBeatTimer = setInterval( @@ -1011,7 +998,7 @@ export class QueryQueue { async (cancelHandler) => { localCancelHandler = cancelHandler; try { - await queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, processingId, queueId); + await queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, queueId); } catch (e: any) { this.logger('Error while query update', { queueId, @@ -1035,7 +1022,6 @@ export class QueryQueue { this.logger('Performing query completed', { queueId, - processingId, queueSize, duration: ((new Date()).getTime() - startQueryTime), queryKey: query.queryKey, @@ -1054,7 +1040,6 @@ export class QueryQueue { }; this.logger('Error while querying', { queueId, - processingId, queueSize, duration: ((new Date()).getTime() - startQueryTime), queryKey: query.queryKey, @@ -1073,7 +1058,6 @@ export class QueryQueue { if (queryWithCancelHandle) { this.logger('Cancelling query due to timeout', { queueId, - processingId, queryKey: queryWithCancelHandle.queryKey, queuePrefix: this.redisQueuePrefix, requestId: queryWithCancelHandle.requestId, @@ -1092,11 +1076,10 @@ export class QueryQueue { clearInterval(heartBeatTimer); } - if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, processingId, queueId))) { + if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, queueId))) { this.logger('Orphaned execution result', { queueId, - processingId, - warn: 'Result for query was not set due to processing lock wasn\'t acquired', + warn: 'Result for query was not set because the queue item is no longer active', queryKey: query.queryKey, queuePrefix: this.redisQueuePrefix, requestId: query.requestId, diff --git a/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts b/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts index 2a3b0e63adecb..f558043c5e217 100644 --- a/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts @@ -42,7 +42,6 @@ function patchQueueDriverConnectionForTrack(connection: QueueDriverConnectionInt setResultAndRemoveQuery: wrapAsyncMethod('setResultAndRemoveQuery'), getQueryStageState: wrapAsyncMethod('getQueryStageState'), getResultBlocking: wrapAsyncMethod('getResultBlocking'), - freeProcessingLock: wrapAsyncMethod('freeProcessingLock'), optimisticQueryUpdate: wrapAsyncMethod('optimisticQueryUpdate'), getQueryAndRemove: wrapAsyncMethod('getQueryAndRemove'), release: connection.release, diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts index 92fdcee8f8bb5..2dc205f4c1be7 100644 --- a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts @@ -465,76 +465,111 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => } }); - onlyLocalTest('queue driver lock obtain race condition', async () => { - const connection: any = await queue.queueDriver.createConnection(); - const connection2: any = await queue.queueDriver.createConnection(); + onlyLocalTest('an active query cannot be retrieved twice', async () => { + const connection = await queue.queueDriver.createConnection(); + const connection2 = await queue.queueDriver.createConnection(); const priority = 10; + const key = 'active-retrieval' as any; - await queue.reconcileQueue(); - - const [, raceQueueId] = await connection.addToQueue( - 'race', 'handler', ['select'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race' } - ); - - const [, race2QueueId] = await connection.addToQueue( - 'race2', 'handler2', ['select2'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race2' } - ); - - // Neither is locked yet, so both releases are no-ops - await connection.freeProcessingLock('race', raceQueueId, true); - await connection.freeProcessingLock('race2', race2QueueId, true); - - await connection2.retrieveForProcessing('race2', race2QueueId); - - const retrieve6 = await connection.retrieveForProcessing('race', raceQueueId); - console.log(retrieve6); - expect(!!retrieve6[5]).toBe(true); + try { + const [, queueId] = await connection.addToQueue( + key, 'handler', ['select'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '1' + } + ); - console.log(await connection.getQueryAndRemove('race')); - console.log(await connection.getQueryAndRemove('race2')); + const firstRetrieval = await connection.retrieveForProcessing(key, queueId); + expect(firstRetrieval?.[5]).toBe(true); - await queue.queueDriver.release(connection); - await queue.queueDriver.release(connection2); + const secondRetrieval = await connection2.retrieveForProcessing(key, queueId); + expect(secondRetrieval).toStrictEqual([0, null, [key], 0, null, false]); + } finally { + await connection.getQueryAndRemove(key, null); + queue.queueDriver.release(connection); + queue.queueDriver.release(connection2); + } }); - onlyLocalTest('activated but lock is not acquired', async () => { + onlyLocalTest('a failed retrieval does not reserve a pending query', async () => { const connection = await queue.queueDriver.createConnection(); const connection2 = await queue.queueDriver.createConnection(); const priority = 10; + const firstKey = 'concurrency-first' as any; + const secondKey = 'concurrency-second' as any; - await queue.reconcileQueue(); + try { + const [, firstQueueId] = await connection.addToQueue( + firstKey, 'handler', ['select'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: firstKey, requestId: '1' + } + ); + const [, secondQueueId] = await connection.addToQueue( + secondKey, 'handler2', ['select2'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: secondKey, requestId: '1' + } + ); - const [, activated1QueueId] = await connection.addToQueue( - 'activated1', 'handler', ['select'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race', requestId: '1' } - ); + expect((await connection.retrieveForProcessing(firstKey, firstQueueId))?.[5]).toBe(true); + expect(await connection2.retrieveForProcessing(secondKey, secondQueueId)).toStrictEqual([ + 0, null, [firstKey], 1, null, false + ]); + expect(await connection.getToProcessQueries()).toStrictEqual([[secondKey, secondQueueId]]); - const [, activated2QueueId] = await connection.addToQueue( - 'activated2', 'handler2', ['select2'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race2', requestId: '1' } - ); + await connection.getQueryAndRemove(firstKey, firstQueueId); - const retrieve1 = await connection.retrieveForProcessing('activated1' as any, activated1QueueId); - console.log(retrieve1); - const retrieve2 = await connection2.retrieveForProcessing('activated2' as any, activated2QueueId); - console.log(retrieve2); - console.log(await connection.freeProcessingLock('activated1' as any, activated1QueueId, retrieve1 && retrieve1[2].indexOf('activated1' as any) !== -1)); + const secondRetrieval = await connection2.retrieveForProcessing(secondKey, secondQueueId); + expect(secondRetrieval?.[0]).toBe(1); + expect(secondRetrieval?.[5]).toBe(true); + } finally { + await connection.getQueryAndRemove(firstKey, null); + await connection.getQueryAndRemove(secondKey, null); + queue.queueDriver.release(connection); + queue.queueDriver.release(connection2); + } + }); - // Another node reaches the same item, so it comes with the same lock token and loses - const retrieve3 = await connection.retrieveForProcessing('activated2' as any, activated2QueueId); - expect(retrieve3).toBeNull(); + test('stale queueId cannot update or acknowledge a requeued query', async () => { + const connection = await queue.queueDriver.createConnection(); + const queryKey = 'requeued-query' as QueryKey; + const key = connection.redisHash(queryKey); + const priority = 10; - console.log(retrieve2[2].indexOf('activated2' as any) !== -1); - console.log(await connection2.freeProcessingLock('activated2' as any, activated2QueueId, retrieve2 && retrieve2[2].indexOf('activated2' as any) !== -1)); + try { + const [, staleQueueId] = await connection.addToQueue( + queryKey, 'handler', ['old'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '1' + } + ); + expect((await connection.retrieveForProcessing(key, staleQueueId))?.[5]).toBe(true); + await connection.getQueryAndRemove(key, staleQueueId); - const retrieve4 = await connection.retrieveForProcessing('activated2' as any, activated2QueueId); - console.log(retrieve4); - expect(retrieve4[0]).toBe(1); - expect(!!retrieve4[5]).toBe(true); + const [, currentQueueId] = await connection.addToQueue( + queryKey, 'handler', ['new'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '2' + } + ); - console.log(await connection.getQueryAndRemove('activated1' as any, null)); - console.log(await connection.getQueryAndRemove('activated2' as any, null)); + if (options.cacheAndQueueDriver !== 'cubestore') { + expect(await connection.retrieveForProcessing(key, staleQueueId)).toStrictEqual([ + 0, null, [], 1, null, false + ]); + expect(await connection.getToProcessQueries()).toStrictEqual([[key, currentQueueId]]); + } + expect((await connection.retrieveForProcessing(key, currentQueueId))?.[5]).toBe(true); - await queue.queueDriver.release(connection); - await queue.queueDriver.release(connection2); + const staleUpdateResult = await connection.optimisticQueryUpdate(key, { stale: true }, staleQueueId); + if (options.cacheAndQueueDriver !== 'cubestore') { + expect(staleUpdateResult).toBe(false); + } + expect(await connection.setResultAndRemoveQuery(key, { result: 'stale' }, staleQueueId)).toBe(false); + const currentQuery = await connection.getQueryDef(key, currentQueueId); + expect(currentQuery).toMatchObject({ query: ['new'] }); + expect(currentQuery).not.toHaveProperty('stale'); + expect(await connection.getActiveQueries()).toStrictEqual([[key, currentQueueId]]); + } finally { + await connection.getQueryAndRemove(key, null); + queue.queueDriver.release(connection); + } }); // eslint-disable-next-line no-unused-expressions