From a42d56ce6616810005e043ce7e86d910cdabb5b0 Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Tue, 11 Aug 2026 13:13:53 +0000 Subject: [PATCH] compute: add configurable peek row iteration limit Compute workers iterate arrangements synchronously while serving index-backed peeks, so a query that scans far more rows than it returns can hold a worker for a long time and delay everything else on the cluster. Persist fast-path peeks have the same shape: filtering happens after the rows have been read. Add an off-by-default failsafe that bounds how many rows a worker may examine for one peek. Two dyncfgs, a feature gate and a threshold that defaults to 1000 rows, both read through handles so that an `UpdateConfiguration` reaches peeks that are already in flight. The budget covers the index result trace, the index error trace and the Persist fast path, and counts rows before literal and MFP filtering, because a row that is read and then discarded costs the same scan time as one that is returned. Exactly the configured number of rows may be examined. A peek fails only when it asks for the row after that. The limit deliberately stops at the peek stash. A stashed peek restarts its scan and produces in bounded bursts, so bounding it needs the count to survive the hand-off, and the restart makes that count charge the same rows twice. Leaving it out keeps this change small. The peeks that motivate the failsafe, large filtered scans, fail before they ever reach the stash threshold. Reporting the limit needs an error type that survives the trip from the worker. `PeekResponse::Error` carried a bare `String`, so every peek failure reached the adapter as `AdapterError::Unstructured` and was reported as XX000. Give it a `PeekError` of `Dataflow`, `Unstructured` or `RowIterationLimitExceeded`, and let `PeekResponseUnary::Error` carry an `AdapterError`, so the conversion happens once instead of once per frontend. The limit then reports SQLSTATE 54000 with a hint naming the threshold parameter, and worker responses merge by error precedence: cancellation, then ordinary errors, then the limit. Carrying the dataflow error structurally also fixes the SQLSTATE of evaluation errors raised while reading a collection: `SELECT a / b FROM t` now reports 22012 like its constant-folded counterpart. Such an error keeps the message `DataflowError` renders, so one that used to come back bare from an index or Persist fast-path MFP now carries the `Evaluation error:` prefix the error-trace path already used. The wire encoding is bincode, which cannot skip a variant it does not know, so `PeekResponse` serializes through a mirror type that keeps `Error(String)` where it was for the unstructured case and appends the structured one. Existing frames, `Canceled` in particular, encode exactly as before. The test and CI configuration enables the feature with a high threshold, so the guarded path is exercised broadly without constraining ordinary queries. --- Cargo.lock | 2 + misc/python/materialize/mzcompose/__init__.py | 5 + .../materialize/parallel_workload/action.py | 4 + src/adapter/src/active_compute_sink.rs | 4 +- src/adapter/src/coord/catalog_implications.rs | 6 +- src/adapter/src/coord/peek.rs | 22 +- src/adapter/src/coord/sequencer/inner.rs | 4 +- src/adapter/src/error.rs | 105 +++++++++- src/compute-client/Cargo.toml | 3 + src/compute-client/src/controller.rs | 2 +- src/compute-client/src/controller/instance.rs | 7 +- src/compute-client/src/protocol/command.rs | 5 +- src/compute-client/src/protocol/response.rs | 162 ++++++++++++++- src/compute-client/src/service.rs | 186 +++++++++++++---- src/compute-types/src/dyncfgs.rs | 27 +++ src/compute/src/compute_state.rs | 192 ++++++++++++++++-- .../src/compute_state/peek_result_iterator.rs | 39 +++- src/compute/src/compute_state/peek_stash.rs | 10 +- src/environmentd/src/http/sql.rs | 9 +- src/environmentd/tests/sql.rs | 45 ++++ src/environmentd/tests/testdata/http/ws | 2 +- src/pgwire/src/protocol.rs | 9 +- src/sqllogictest/Cargo.toml | 1 + src/sqllogictest/src/bin/sqllogictest.rs | 11 + .../mzcompose.py | 2 + test/sqllogictest/max_result_size.slt | 83 ++++++++ test/sqllogictest/persist-fast-path.slt | 33 ++- 27 files changed, 876 insertions(+), 104 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8a488ba4b23ff..a79ea4cf66010 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6785,6 +6785,7 @@ version = "0.0.0" dependencies = [ "anyhow", "async-trait", + "bincode", "bytesize", "chrono", "derivative", @@ -8610,6 +8611,7 @@ dependencies = [ "mz-authenticator", "mz-build-info", "mz-catalog", + "mz-compute-types", "mz-controller", "mz-dyncfgs", "mz-environmentd", diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 18afa831d4b44..2850af6cc6b24 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -119,6 +119,11 @@ def get_minimal_system_parameters( # End of list (ordered by name) } + if version >= MzVersion.parse_mz("v26.38.0-dev"): + # Exercise the row-limit check without constraining normal test queries. + config["compute_peek_row_iteration_limit"] = "1000000000" + config["enable_compute_peek_row_iteration_limit"] = "true" + if version < MzVersion.parse_mz("v0.163.0-dev"): config["enable_compute_active_dataflow_cancelation"] = "true" diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 03f25fd615ca1..2de91cb14b11c 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -2982,6 +2982,10 @@ def __init__( self.flags_with_values["enable_compute_peek_response_stash"] = ( BOOLEAN_FLAG_VALUES ) + self.flags_with_values["enable_compute_peek_row_iteration_limit"] = ( + BOOLEAN_FLAG_VALUES + ) + self.flags_with_values["compute_peek_row_iteration_limit"] = ["1000000000"] self.flags_with_values["compute_peek_response_stash_threshold_bytes"] = [ "0", # "force enabled" "1048576", # 1 MiB, an in-between value diff --git a/src/adapter/src/active_compute_sink.rs b/src/adapter/src/active_compute_sink.rs index 4520868c4617e..74beb9e283976 100644 --- a/src/adapter/src/active_compute_sink.rs +++ b/src/adapter/src/active_compute_sink.rs @@ -178,7 +178,9 @@ impl ActiveSubscribe { mz_ore::iter::consolidate_update_iter(merged) } Err(s) => { - self.send(PeekResponseUnary::Error(s)); + self.send(PeekResponseUnary::Error(AdapterError::Unstructured( + anyhow::Error::msg(s), + ))); return true; } }; diff --git a/src/adapter/src/coord/catalog_implications.rs b/src/adapter/src/coord/catalog_implications.rs index 9f5cda2cdf1ec..9d5b4814d67cc 100644 --- a/src/adapter/src/coord/catalog_implications.rs +++ b/src/adapter/src/coord/catalog_implications.rs @@ -39,7 +39,7 @@ use mz_catalog::memory::objects::{ }; use mz_cloud_resources::VpcEndpointConfig; use mz_compute_client::logging::LogVariant; -use mz_compute_client::protocol::response::PeekResponse; +use mz_compute_client::protocol::response::{PeekError, PeekResponse}; use mz_controller::clusters::{ClusterRole, ReplicaConfig}; use mz_controller_types::{ClusterId, ReplicaId}; use mz_ore::collections::CollectionExt; @@ -992,7 +992,9 @@ impl Coordinator { if !peeks_to_drop.is_empty() { for (dep, uuid) in peeks_to_drop { if let Some(pending_peek) = self.remove_pending_peek(&uuid) { - let cancel_reason = PeekResponse::Error(dep.query_terminated_error()); + let cancel_reason = PeekResponse::Error(PeekError::unstructured( + dep.query_terminated_error(), + )); self.controller .compute .cancel_peek(pending_peek.cluster_id, uuid, cancel_reason) diff --git a/src/adapter/src/coord/peek.rs b/src/adapter/src/coord/peek.rs index f593fa539a696..a8d653495210f 100644 --- a/src/adapter/src/coord/peek.rs +++ b/src/adapter/src/coord/peek.rs @@ -85,13 +85,9 @@ pub(crate) struct PendingPeek { #[derive(Debug)] pub enum PeekResponseUnary { Rows(Box), - Error(String), + Error(AdapterError), Canceled, /// A dependency was dropped during execution. - /// - /// N.B. This is a bit of a workaround for the fact that our Error variant - /// is unstructured and right now we specifically care about this error and - /// need to render differently based on context. DependencyDropped(DroppedDependency), } @@ -1049,7 +1045,7 @@ impl crate::coord::Coordinator { let rows = match result { Ok(rows) => rows, Err(e) => { - yield PeekResponseUnary::Error(e.to_string()); + yield PeekResponseUnary::Error(AdapterError::Unstructured(anyhow::anyhow!(e))); return; } }; @@ -1064,7 +1060,11 @@ impl crate::coord::Coordinator { &duration_histogram, ) { Ok((rows, _size_bytes)) => yield PeekResponseUnary::Rows(Box::new(rows)), - Err(e) => yield PeekResponseUnary::Error(e), + Err(e) => { + yield PeekResponseUnary::Error(AdapterError::Unstructured( + anyhow::Error::msg(e), + )) + } } } PeekResponse::Stashed(response) => { @@ -1221,7 +1221,11 @@ impl crate::coord::Coordinator { match result_rows { Ok(result_rows) => yield PeekResponseUnary::Rows(Box::new(result_rows)), - Err(e) => yield PeekResponseUnary::Error(e), + Err(e) => { + yield PeekResponseUnary::Error(AdapterError::Unstructured( + anyhow::Error::msg(e), + )) + } } } @@ -1236,7 +1240,7 @@ impl crate::coord::Coordinator { yield PeekResponseUnary::Canceled; } PeekResponse::Error(e) => { - yield PeekResponseUnary::Error(e); + yield PeekResponseUnary::Error(e.into()); } } }) diff --git a/src/adapter/src/coord/sequencer/inner.rs b/src/adapter/src/coord/sequencer/inner.rs index f979d56383ef1..75733cf9964f7 100644 --- a/src/adapter/src/coord/sequencer/inner.rs +++ b/src/adapter/src/coord/sequencer/inner.rs @@ -2987,9 +2987,7 @@ impl Coordinator { }; } PeekResponseUnary::Canceled => break Err(AdapterError::Canceled), - PeekResponseUnary::Error(e) => { - break Err(AdapterError::Unstructured(anyhow!(e))); - } + PeekResponseUnary::Error(e) => break Err(e), PeekResponseUnary::DependencyDropped(dep) => { break Err(dep.to_concurrent_dependency_drop()); } diff --git a/src/adapter/src/error.rs b/src/adapter/src/error.rs index ac348a8dd26d3..866c4d692e914 100644 --- a/src/adapter/src/error.rs +++ b/src/adapter/src/error.rs @@ -36,7 +36,7 @@ use mz_sql::rbac; use mz_sql::session::vars::VarError; use mz_storage_types::connections::ConnectionValidationError; use mz_storage_types::controller::StorageError; -use mz_storage_types::errors::CollectionMissing; +use mz_storage_types::errors::{CollectionMissing, DataflowError}; use smallvec::SmallVec; use timely::progress::Antichain; use tokio::sync::oneshot; @@ -69,6 +69,8 @@ pub enum AdapterError { DuplicateCursor(String), /// An error while evaluating an expression. Eval(EvalError), + /// An error produced while executing a dataflow. + Dataflow(Box), /// An error occurred while planning the statement. Explain(ExplainError), /// The ID allocator exhausted all valid IDs. @@ -173,6 +175,11 @@ pub enum AdapterError { }, /// Result size of a query is too large. ResultSize(String), + /// A query exceeded the configured compute peek row iteration limit. + PeekRowIterationLimitExceeded { + /// The configured per-worker limit. + limit: usize, + }, /// The specified feature is not permitted in safe mode. SafeModeViolation(String), /// The current transaction had the wrong set of write locks. @@ -495,6 +502,15 @@ fn eval_error_code(err: &EvalError) -> SqlState { } } +fn dataflow_error_code(error: &DataflowError) -> SqlState { + match error { + DataflowError::EvalError(error) => eval_error_code(error), + DataflowError::DecodeError(_) + | DataflowError::SourceError(_) + | DataflowError::EnvelopeError(_) => SqlState::INTERNAL_ERROR, + } +} + impl AdapterError { pub fn into_response(self, severity: Severity) -> ErrorResponse { ErrorResponse { @@ -522,6 +538,10 @@ impl AdapterError { } AdapterError::Catalog(c) => c.detail(), AdapterError::Eval(e) => e.detail(), + AdapterError::Dataflow(e) => match &**e { + DataflowError::EvalError(e) => e.detail(), + _ => None, + }, AdapterError::RelationOutsideTimeDomain { relations, names } => Some(format!( "The following relations in the query are outside the transaction's time domain:\n{}\n{}", relations @@ -562,6 +582,11 @@ impl AdapterError { objects. Reduce the number of dependencies, or raise the \ read_then_write_max_dependencies system parameter." )), + AdapterError::PeekRowIterationLimitExceeded { limit } => Some(format!( + "The query attempted to examine more than {limit} rows on a single compute \ + worker. This limit prevents long-running SELECT queries from delaying other \ + work on the cluster." + )), AdapterError::SafeModeViolation(_) => Some( "The Materialize server you are connected to is running in \ safe mode, which limits the features that are available." @@ -735,6 +760,10 @@ impl AdapterError { ), AdapterError::Catalog(c) => c.hint(), AdapterError::Eval(e) => e.hint(), + AdapterError::Dataflow(e) => match &**e { + DataflowError::EvalError(e) => e.hint(), + _ => None, + }, AdapterError::AlterClusterUnmanagedWhileReconfiguring => Some( "Cancel the reconfiguration by altering the cluster back to its current \ configuration, or wait for it to settle, then convert." @@ -798,6 +827,13 @@ impl AdapterError { statement_timeout = '120s'`." .into(), ), + AdapterError::PeekRowIterationLimitExceeded { .. } => Some( + "Reduce the number of rows the query must examine, for example by querying an \ + indexed, more selective result. Queries with `LIMIT` and no `ORDER BY` can also \ + stop early. To permit this query, increase \ + `compute_peek_row_iteration_limit`." + .into(), + ), AdapterError::PlanError(e) => e.hint(), AdapterError::UnallowedOnCluster { cluster, .. } => { (cluster != MZ_CATALOG_SERVER_CLUSTER.name).then(|| @@ -864,6 +900,7 @@ impl AdapterError { // exhaustively so the catch-all `INTERNAL_ERROR` no longer applies // to errors that are really the user's fault. See SQL-326. AdapterError::Eval(e) => eval_error_code(e), + AdapterError::Dataflow(e) => dataflow_error_code(e), AdapterError::Explain(_) => SqlState::INTERNAL_ERROR, AdapterError::IdExhaustionError => SqlState::INTERNAL_ERROR, AdapterError::Internal(_) => SqlState::INTERNAL_ERROR, @@ -923,6 +960,7 @@ impl AdapterError { AdapterError::RelationOutsideTimeDomain { .. } => SqlState::INVALID_TRANSACTION_STATE, AdapterError::ResourceExhaustion { .. } => SqlState::INSUFFICIENT_RESOURCES, AdapterError::ResultSize(_) => SqlState::OUT_OF_MEMORY, + AdapterError::PeekRowIterationLimitExceeded { .. } => SqlState::PROGRAM_LIMIT_EXCEEDED, AdapterError::SafeModeViolation(_) => SqlState::INTERNAL_ERROR, AdapterError::SubscribeOnlyTransaction => SqlState::INVALID_TRANSACTION_STATE, AdapterError::Optimizer(e) => match e { @@ -1197,6 +1235,7 @@ impl fmt::Display for AdapterError { write!(f, "cursor {} already exists", name.quoted()) } AdapterError::Eval(e) => e.fmt(f), + AdapterError::Dataflow(e) => e.fmt(f), AdapterError::Explain(e) => e.fmt(f), AdapterError::IdExhaustionError => f.write_str("ID allocator exhausted all valid IDs"), AdapterError::Internal(e) => write!(f, "internal error: {}", e), @@ -1237,6 +1276,12 @@ impl fmt::Display for AdapterError { "selection has too many transitive dependencies to validate (limit {max_rw_dependencies})" ) } + AdapterError::PeekRowIterationLimitExceeded { limit } => { + write!( + f, + "query exceeded the configured row iteration limit of {limit} rows" + ) + } AdapterError::ReplaceMaterializedViewSealed { name } => { write!( f, @@ -1554,6 +1599,20 @@ impl From for AdapterError { } } +impl From for AdapterError { + fn from(error: mz_compute_client::protocol::response::PeekError) -> Self { + use mz_compute_client::protocol::response::PeekError; + + match error { + PeekError::Dataflow(error) => AdapterError::Dataflow(error), + PeekError::Unstructured(error) => AdapterError::Unstructured(anyhow::Error::msg(error)), + PeekError::RowIterationLimitExceeded { limit } => { + AdapterError::PeekRowIterationLimitExceeded { limit } + } + } + } +} + impl From for AdapterError { fn from(e: ExplainError) -> AdapterError { match e { @@ -1688,3 +1747,47 @@ impl From for AdapterError { } impl Error for AdapterError {} + +#[cfg(test)] +mod tests { + use super::*; + + #[mz_ore::test] + fn peek_row_iteration_limit_error_is_user_facing() { + let response = AdapterError::PeekRowIterationLimitExceeded { limit: 1000 } + .into_response(Severity::Error); + + assert_eq!(response.code, SqlState::PROGRAM_LIMIT_EXCEEDED); + assert_eq!( + response.message, + "query exceeded the configured row iteration limit of 1000 rows" + ); + assert_eq!( + response.detail.as_deref(), + Some( + "The query attempted to examine more than 1000 rows on a single compute worker. \ + This limit prevents long-running SELECT queries from delaying other work on the \ + cluster." + ) + ); + assert!( + response + .hint + .as_deref() + .is_some_and(|hint| hint.contains("compute_peek_row_iteration_limit")) + ); + } + + #[mz_ore::test] + fn structured_dataflow_error_preserves_message_and_code() { + use mz_compute_client::protocol::response::PeekError; + + let dataflow_error = DataflowError::from(EvalError::DivisionByZero); + let expected_message = dataflow_error.to_string(); + let response = + AdapterError::from(PeekError::from(dataflow_error)).into_response(Severity::Error); + + assert_eq!(response.code, SqlState::DIVISION_BY_ZERO); + assert_eq!(response.message, expected_message); + } +} diff --git a/src/compute-client/Cargo.toml b/src/compute-client/Cargo.toml index 967413b52ed3f..b4811d10a12ba 100644 --- a/src/compute-client/Cargo.toml +++ b/src/compute-client/Cargo.toml @@ -41,5 +41,8 @@ tokio.workspace = true tracing.workspace = true uuid = { workspace = true, features = ["serde", "v4"] } +[dev-dependencies] +bincode.workspace = true + [features] default = [] diff --git a/src/compute-client/src/controller.rs b/src/compute-client/src/controller.rs index 33c3d9a1a2672..abe39fb317bc0 100644 --- a/src/compute-client/src/controller.rs +++ b/src/compute-client/src/controller.rs @@ -170,7 +170,7 @@ impl PeekNotification { result_size: u64::cast_from(result_size), } } - PeekResponse::Error(err) => Self::Error(err.clone()), + PeekResponse::Error(err) => Self::Error(err.to_string()), PeekResponse::Canceled => Self::Canceled, } } diff --git a/src/compute-client/src/controller/instance.rs b/src/compute-client/src/controller/instance.rs index 04587b9d26019..4ef26f3b2d7f4 100644 --- a/src/compute-client/src/controller/instance.rs +++ b/src/compute-client/src/controller/instance.rs @@ -64,8 +64,8 @@ use crate::protocol::command::{ }; use crate::protocol::history::ComputeCommandHistory; use crate::protocol::response::{ - ComputeResponse, CopyToResponse, FrontiersResponse, PeekResponse, StatusResponse, - SubscribeBatch, SubscribeResponse, + ComputeResponse, CopyToResponse, FrontiersResponse, PeekError as ProtocolPeekError, + PeekResponse, StatusResponse, SubscribeBatch, SubscribeResponse, }; #[derive(Error, Debug)] @@ -1346,7 +1346,8 @@ impl Instance { self.deliver_response(response); } for uuid in to_drop { - let response = PeekResponse::Error(ERROR_TARGET_REPLICA_FAILED.into()); + let response = + PeekResponse::Error(ProtocolPeekError::unstructured(ERROR_TARGET_REPLICA_FAILED)); self.finish_peek(uuid, response); } diff --git a/src/compute-client/src/protocol/command.rs b/src/compute-client/src/protocol/command.rs index 3c3f35980c249..b69bca23fffca 100644 --- a/src/compute-client/src/protocol/command.rs +++ b/src/compute-client/src/protocol/command.rs @@ -233,12 +233,13 @@ pub enum ComputeCommand { /// After receiving a `Peek` command, the replica must eventually produce a single /// [`PeekResponse`]: /// - /// * For peeks that were not cancelled: either [`Rows`] or [`Error`]. - /// * For peeks that were cancelled: either [`Rows`], or [`Error`], or [`Canceled`]. + /// * For peeks that were not cancelled: [`Rows`], [`Stashed`], or [`Error`]. + /// * For peeks that were cancelled: any response above, or [`Canceled`]. /// /// [`PeekResponse`]: super::response::PeekResponse /// [`PeekResponse::Error`]: super::response::PeekResponse::Error /// [`Rows`]: super::response::PeekResponse::Rows + /// [`Stashed`]: super::response::PeekResponse::Stashed /// [`Error`]: super::response::PeekResponse::Error /// [`Canceled`]: super::response::PeekResponse::Canceled Peek(Box), diff --git a/src/compute-client/src/protocol/response.rs b/src/compute-client/src/protocol/response.rs index fd62e744bffac..61f48abd684fc 100644 --- a/src/compute-client/src/protocol/response.rs +++ b/src/compute-client/src/protocol/response.rs @@ -9,13 +9,17 @@ //! Compute protocol responses. +use std::fmt; + +use mz_expr::EvalError; use mz_expr::row::RowCollection; use mz_ore::cast::CastFrom; use mz_ore::tracing::OpenTelemetryContext; use mz_persist_client::batch::ProtoBatch; use mz_persist_types::ShardId; use mz_repr::{GlobalId, RelationDesc, Timestamp, UpdateCollection}; -use serde::{Deserialize, Serialize}; +use mz_storage_types::errors::DataflowError; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; use timely::progress::frontier::Antichain; use uuid::Uuid; @@ -187,18 +191,78 @@ impl FrontiersResponse { /// /// Note that each `Peek` expects to generate exactly one `PeekResponse`, i.e. /// we expect a 1:1 contract between `Peek` and `PeekResponse`. -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +#[derive(Clone, Debug, PartialEq)] pub enum PeekResponse { /// Returned rows of a successful peek. Rows(Vec), /// Results of the peek were stashed in persist batches. Stashed(Box), /// Error of an unsuccessful peek. - Error(String), + Error(PeekError), /// The peek was canceled. Canceled, } +/// Wire shape of [`PeekResponse`], mirrored so that giving `Error` a payload type does not +/// change the bytes we put on the CTP connection. +/// +/// The compute protocol is encoded with bincode, which identifies a variant by its position and +/// has no way to skip one it does not know. So variants may only be appended, and the encoding of +/// an existing variant may not change. `Error(String)` is therefore kept where it was and carries +/// the unstructured case, while structured errors ride along in `StructuredError`. +/// +/// NOTE: A replica that predates `StructuredError` cannot decode it. The CTP handshake compares +/// semantic versions, not builds, so it does not catch that on its own. +#[derive(Serialize)] +enum WirePeekResponseRef<'a> { + Rows(&'a [RowCollection]), + Stashed(&'a StashedPeekResponse), + Error(&'a str), + Canceled, + StructuredError(&'a PeekError), +} + +/// Owned counterpart of [`WirePeekResponseRef`], used for decoding. +#[derive(Deserialize)] +enum WirePeekResponse { + Rows(Vec), + Stashed(Box), + Error(String), + Canceled, + StructuredError(PeekError), +} + +impl Serialize for PeekResponse { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + match self { + Self::Rows(rows) => WirePeekResponseRef::Rows(rows.as_slice()), + Self::Stashed(stashed) => WirePeekResponseRef::Stashed(stashed.as_ref()), + Self::Error(PeekError::Unstructured(error)) => WirePeekResponseRef::Error(error), + Self::Error(error) => WirePeekResponseRef::StructuredError(error), + Self::Canceled => WirePeekResponseRef::Canceled, + } + .serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for PeekResponse { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + Ok(match WirePeekResponse::deserialize(deserializer)? { + WirePeekResponse::Rows(rows) => Self::Rows(rows), + WirePeekResponse::Stashed(stashed) => Self::Stashed(stashed), + WirePeekResponse::Error(error) => Self::Error(PeekError::Unstructured(error)), + WirePeekResponse::Canceled => Self::Canceled, + WirePeekResponse::StructuredError(error) => Self::Error(error), + }) + } +} + impl PeekResponse { /// Return the size of row bytes stored inline in this response. pub fn inline_byte_len(&self) -> usize { @@ -210,6 +274,58 @@ impl PeekResponse { } } +/// The error of an unsuccessful peek. +/// +/// The variant decides what the user sees: a `Dataflow` error keeps the structure the dataflow +/// produced and so gets the same SQLSTATE constant folding would have assigned, a +/// `RowIterationLimitExceeded` names a limit the user can raise, and an `Unstructured` error is +/// reported as an internal error. Prefer the structured variants whenever the source has one. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub enum PeekError { + /// An error produced while executing the dataflow, for example evaluating an expression over + /// a collection. + Dataflow(Box), + /// An error from the peek machinery itself, with no structured form. + Unstructured(String), + /// A worker examined more rows than `compute_peek_row_iteration_limit` allows. + RowIterationLimitExceeded { + /// The limit that was in effect, in rows. + limit: usize, + }, +} + +impl PeekError { + /// Constructs an unstructured peek error. + pub fn unstructured(message: impl Into) -> Self { + Self::Unstructured(message.into()) + } +} + +impl fmt::Display for PeekError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Dataflow(error) => error.fmt(f), + Self::Unstructured(error) => f.write_str(error), + Self::RowIterationLimitExceeded { limit } => write!( + f, + "query exceeded the configured row iteration limit of {limit} rows" + ), + } + } +} + +impl From for PeekError { + fn from(error: DataflowError) -> Self { + Self::Dataflow(Box::new(error)) + } +} + +impl From for PeekError { + fn from(error: EvalError) -> Self { + Self::Dataflow(Box::new(error.into())) + } +} + /// Response from a peek whose results have been stashed into persist. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct StashedPeekResponse { @@ -328,6 +444,46 @@ pub enum StatusResponse { #[cfg(test)] mod tests { use super::*; + use bincode::Options; + + #[allow(dead_code)] + #[derive(Serialize)] + enum LegacyPeekResponse { + Rows(Vec), + Stashed(Box), + Error(String), + Canceled, + } + + fn serialize_ctp(value: &impl Serialize) -> Vec { + bincode::DefaultOptions::new().serialize(value).unwrap() + } + + #[mz_ore::test] + fn peek_response_preserves_legacy_encodings() { + let legacy_error = LegacyPeekResponse::Error("dataflow error".into()); + let error = PeekResponse::Error(PeekError::unstructured("dataflow error")); + assert_eq!(serialize_ctp(&error), serialize_ctp(&legacy_error)); + + let decoded: PeekResponse = bincode::DefaultOptions::new() + .deserialize(&serialize_ctp(&legacy_error)) + .unwrap(); + assert_eq!(decoded, error); + + assert_eq!( + serialize_ctp(&PeekResponse::Canceled), + serialize_ctp(&LegacyPeekResponse::Canceled) + ); + } + + #[mz_ore::test] + fn structured_peek_error_roundtrips() { + let response = PeekResponse::Error(PeekError::RowIterationLimitExceeded { limit: 1000 }); + let decoded: PeekResponse = bincode::DefaultOptions::new() + .deserialize(&serialize_ctp(&response)) + .unwrap(); + assert_eq!(decoded, response); + } /// Test to ensure the size of the `ComputeResponse` enum doesn't regress. #[mz_ore::test] diff --git a/src/compute-client/src/service.rs b/src/compute-client/src/service.rs index 019bcab56fcc6..a7e4b88d77176 100644 --- a/src/compute-client/src/service.rs +++ b/src/compute-client/src/service.rs @@ -27,8 +27,8 @@ use uuid::Uuid; use crate::protocol::command::ComputeCommand; use crate::protocol::response::{ - ComputeResponse, CopyToResponse, FrontiersResponse, PeekResponse, StashedPeekResponse, - SubscribeBatch, SubscribeResponse, + ComputeResponse, CopyToResponse, FrontiersResponse, PeekError, PeekResponse, + StashedPeekResponse, SubscribeBatch, SubscribeResponse, }; /// A client to a compute server. @@ -106,7 +106,7 @@ pub struct PartitionedComputeState { /// The compute protocol requires that exactly one response is emitted for each peek. This /// property ensures that a) we can eventually drop the tracking state maintained for a peek /// and b) we won't re-initialize tracking for a peek we have already served. - peek_responses: BTreeMap)>, + peek_responses: BTreeMap, /// For each in-progress copy-to the response data received so far, and the set of shards that /// provided responses already. /// @@ -223,19 +223,14 @@ impl PartitionedComputeState { response: PeekResponse, otel_ctx: OpenTelemetryContext, ) -> Option { - let (merged, ready_shards) = self.peek_responses.entry(uuid).or_insert(( - PeekResponse::Rows(vec![RowCollection::default()]), - BTreeSet::new(), - )); - - let first = ready_shards.insert(shard_id); - assert!(first, "duplicate peek response"); - - let resp1 = mem::replace(merged, PeekResponse::Canceled); - *merged = merge_peek_responses(resp1, response, self.max_result_size); - - if ready_shards.len() == self.parts { - let (response, _) = self.peek_responses.remove(&uuid).unwrap(); + let pending = self + .peek_responses + .entry(uuid) + .or_insert_with(PendingPeek::new); + pending.absorb(shard_id, response, self.max_result_size); + + if pending.ready_shards.len() == self.parts { + let response = self.peek_responses.remove(&uuid).unwrap().response; Some(ComputeResponse::PeekResponse(uuid, response, otel_ctx)) } else { None @@ -565,32 +560,67 @@ impl PendingSubscribe { } } +/// Accumulates the per-worker responses to one peek into the single response the controller +/// hands upwards. +#[derive(Debug)] +struct PendingPeek { + /// The responses merged so far. + response: PeekResponse, + /// Inline result bytes seen so far, across all shards. + /// + /// Tracked separately from `response` because a worker's rows are dropped as soon as any + /// worker reports an error. Without this the aggregate size check would depend on the order + /// the responses happen to arrive in. + inline_byte_len: usize, + /// The shards that have provided responses. + ready_shards: BTreeSet, +} + +impl PendingPeek { + fn new() -> Self { + Self { + response: PeekResponse::Rows(vec![RowCollection::default()]), + inline_byte_len: 0, + ready_shards: BTreeSet::new(), + } + } + + fn absorb(&mut self, shard_id: usize, response: PeekResponse, max_result_size: u64) { + let first = self.ready_shards.insert(shard_id); + assert!(first, "duplicate peek response"); + + self.inline_byte_len = self + .inline_byte_len + .saturating_add(response.inline_byte_len()); + let current = mem::replace(&mut self.response, PeekResponse::Canceled); + self.response = merge_peek_responses(current, response); + + // Merging eagerly is what keeps the controller's memory bounded, so the size check has to + // happen on every response rather than once at the end. + if self.inline_byte_len > max_result_size.cast_into() { + // NOTE: Tests match on this exact message, so nothing else may produce it. + let error = PeekError::unstructured(format!( + "total result exceeds max size of {}", + ByteSize::b(max_result_size) + )); + let current = mem::replace(&mut self.response, PeekResponse::Canceled); + self.response = merge_peek_responses(current, PeekResponse::Error(error)); + } + } +} + /// Merge two [`PeekResponse`]s. -fn merge_peek_responses( - resp1: PeekResponse, - resp2: PeekResponse, - max_result_size: u64, -) -> PeekResponse { +fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekResponse { use PeekResponse::*; // Cancelations and errors short-circuit. Cancelations take precedence over errors. let (resp1, resp2) = match (resp1, resp2) { (Canceled, _) | (_, Canceled) => return Canceled, + (Error(e1), Error(e2)) => return Error(merge_peek_errors(e1, e2)), (Error(e), _) | (_, Error(e)) => return Error(e), resps => resps, }; - let total_byte_len = resp1.inline_byte_len() + resp2.inline_byte_len(); - if total_byte_len > max_result_size.cast_into() { - // Note: We match on this specific error message in tests so it's important that - // nothing else returns the same string. - let err = format!( - "total result exceeds max size of {}", - ByteSize::b(max_result_size) - ); - return Error(err); - } - match (resp1, resp2) { (Rows(mut rows1), Rows(rows2)) => { rows1.extend(rows2); @@ -625,14 +655,14 @@ fn merge_peek_responses( "shard IDs of stashed responses do not match: \ {shard_id1} != {shard_id2}" ); - return Error("internal error".into()); + return Error(PeekError::unstructured("internal error")); } if relation_desc1 != relation_desc2 { soft_panic_or_log!( "relation descs of stashed responses do not match: \ {relation_desc1:?} != {relation_desc2:?}" ); - return Error("internal error".into()); + return Error(PeekError::unstructured("internal error")); } batches1.append(&mut batches2); @@ -650,3 +680,91 @@ fn merge_peek_responses( _ => unreachable!("handled above"), } } + +/// Merge two [`PeekError`]s into the one we report. +/// +/// A row-iteration-limit error only wins against another one, and then the smaller limit wins so +/// the choice does not depend on which worker answered first. Any other error outranks it: a +/// query that both overran the limit and failed for a real reason should report the real reason. +fn merge_peek_errors(error1: PeekError, error2: PeekError) -> PeekError { + match (error1, error2) { + ( + PeekError::RowIterationLimitExceeded { limit: limit1 }, + PeekError::RowIterationLimitExceeded { limit: limit2 }, + ) => PeekError::RowIterationLimitExceeded { + limit: limit1.min(limit2), + }, + (PeekError::RowIterationLimitExceeded { .. }, error) + | (error, PeekError::RowIterationLimitExceeded { .. }) => error, + (error, _) => error, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::num::NonZeroUsize; + + use mz_repr::Row; + + #[mz_ore::test] + fn pending_peek_response_precedence() { + let rows = PeekResponse::Rows(vec![RowCollection::default()]); + let error = PeekResponse::Error(PeekError::unstructured("dataflow error")); + let row_limit = |limit| PeekResponse::Error(PeekError::RowIterationLimitExceeded { limit }); + + let mut pending = PendingPeek::new(); + pending.absorb(0, rows.clone(), u64::MAX); + pending.absorb(1, row_limit(1000), u64::MAX); + pending.absorb(2, row_limit(500), u64::MAX); + assert_eq!(pending.response, row_limit(500)); + + let mut pending = PendingPeek::new(); + pending.absorb(0, rows, u64::MAX); + pending.absorb(1, row_limit(1000), u64::MAX); + pending.absorb(2, error.clone(), u64::MAX); + assert_eq!(pending.response, error); + + let mut pending = PendingPeek::new(); + pending.absorb( + 0, + PeekResponse::Error(PeekError::unstructured("dataflow error")), + u64::MAX, + ); + pending.absorb(3, PeekResponse::Canceled, u64::MAX); + assert_eq!(pending.response, PeekResponse::Canceled); + } + + #[mz_ore::test] + fn peek_max_size_wins_over_row_iteration_limit_in_every_order() { + let row = RowCollection::new(vec![(Row::default(), NonZeroUsize::new(1).unwrap())], &[]); + let rows = PeekResponse::Rows(vec![row]); + let max_result_size = u64::try_from(rows.inline_byte_len()).unwrap(); + let responses = [ + rows.clone(), + PeekResponse::Error(PeekError::RowIterationLimitExceeded { limit: 1000 }), + rows, + ]; + let permutations = [ + [0, 1, 2], + [0, 2, 1], + [1, 0, 2], + [1, 2, 0], + [2, 0, 1], + [2, 1, 0], + ]; + let expected = PeekResponse::Error(PeekError::unstructured(format!( + "total result exceeds max size of {}", + ByteSize::b(max_result_size) + ))); + + for permutation in permutations { + let mut pending = PendingPeek::new(); + for (shard_id, response_index) in permutation.into_iter().enumerate() { + pending.absorb(shard_id, responses[response_index].clone(), max_result_size); + } + + assert_eq!(pending.response, expected, "{permutation:?}"); + } + } +} diff --git a/src/compute-types/src/dyncfgs.rs b/src/compute-types/src/dyncfgs.rs index 42cd46753beba..150603477956a 100644 --- a/src/compute-types/src/dyncfgs.rs +++ b/src/compute-types/src/dyncfgs.rs @@ -519,6 +519,20 @@ pub const PEEK_STASH_BATCH_SIZE: Config = Config::new( "The size, as number of rows, of each batch pumped from the peek result iterator (in one iteration through the worker loop) when stashing peek responses.", ); +/// Whether compute should stop peeks that iterate over too many rows. +pub const ENABLE_PEEK_ROW_ITERATION_LIMIT: Config = Config::new( + "enable_compute_peek_row_iteration_limit", + false, + "Whether compute should stop peeks that exceed compute_peek_row_iteration_limit.", +); + +/// The maximum number of rows a peek may iterate over on each worker. +pub const PEEK_ROW_ITERATION_LIMIT: Config = Config::new( + "compute_peek_row_iteration_limit", + 1000, + "The maximum number of rows a peek may iterate over on each worker when enable_compute_peek_row_iteration_limit is enabled. Does not apply once a peek's results move to the peek stash.", +); + /// The collection interval for the Prometheus metrics introspection source. /// /// Set to zero to disable scraping and retract any existing data. @@ -590,6 +604,8 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES) .add(&PEEK_STASH_NUM_BATCHES) .add(&PEEK_STASH_BATCH_SIZE) + .add(&ENABLE_PEEK_ROW_ITERATION_LIMIT) + .add(&PEEK_ROW_ITERATION_LIMIT) .add(&COMPUTE_PROMETHEUS_INTROSPECTION_SCRAPE_INTERVAL) .add(&SUBSCRIBE_SNAPSHOT_OPTIMIZATION) .add(&MV_SINK_ADVANCE_PERSIST_FRONTIERS) @@ -602,3 +618,14 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&COLUMN_PAGED_BATCHER_EAGER_BACKING) .add(&COLUMN_PAGED_BATCHER_POOL_RSS_TARGET_FRACTION) } + +#[cfg(test)] +mod tests { + use super::*; + + #[mz_ore::test] + fn peek_row_iteration_limit_defaults() { + assert!(!*ENABLE_PEEK_ROW_ITERATION_LIMIT.default()); + assert_eq!(*PEEK_ROW_ITERATION_LIMIT.default(), 1000); + } +} diff --git a/src/compute/src/compute_state.rs b/src/compute/src/compute_state.rs index e20ef9266ec9f..f97f072f1f2c2 100644 --- a/src/compute/src/compute_state.rs +++ b/src/compute/src/compute_state.rs @@ -26,15 +26,16 @@ use mz_compute_client::protocol::command::{ }; use mz_compute_client::protocol::history::ComputeCommandHistory; use mz_compute_client::protocol::response::{ - ComputeResponse, CopyToResponse, FrontiersResponse, PeekResponse, SubscribeResponse, + ComputeResponse, CopyToResponse, FrontiersResponse, PeekError, PeekResponse, SubscribeResponse, }; use mz_compute_types::dataflows::DataflowDescription; use mz_compute_types::dyncfgs::{ - ENABLE_PEEK_RESPONSE_STASH, PEEK_RESPONSE_STASH_BATCH_MAX_RUNS, - PEEK_RESPONSE_STASH_THRESHOLD_BYTES, PEEK_STASH_BATCH_SIZE, PEEK_STASH_NUM_BATCHES, + ENABLE_PEEK_RESPONSE_STASH, ENABLE_PEEK_ROW_ITERATION_LIMIT, + PEEK_RESPONSE_STASH_BATCH_MAX_RUNS, PEEK_RESPONSE_STASH_THRESHOLD_BYTES, + PEEK_ROW_ITERATION_LIMIT, PEEK_STASH_BATCH_SIZE, PEEK_STASH_NUM_BATCHES, }; use mz_compute_types::plan::render_plan::RenderPlan; -use mz_dyncfg::ConfigSet; +use mz_dyncfg::{ConfigSet, ConfigValHandle}; use mz_expr::row::RowCollection; use mz_expr::{RowComparator, SafeMfpPlan}; use mz_ore::cast::{CastFrom, CastLossy}; @@ -79,6 +80,121 @@ use crate::server::{ComputeInstanceContext, ResponseSender}; mod peek_result_iterator; mod peek_stash; +/// Cheap handles on the dyncfgs that bound how many rows a peek may examine. +/// +/// The limit is read through handles rather than captured once, because `UpdateConfiguration` +/// applies to peeks that are already in flight. +#[derive(Clone, Debug)] +struct PeekRowIterationConfig { + enabled: ConfigValHandle, + limit: ConfigValHandle, +} + +impl PeekRowIterationConfig { + fn new(config: &ConfigSet) -> Self { + Self { + enabled: ENABLE_PEEK_ROW_ITERATION_LIMIT.handle(config), + limit: PEEK_ROW_ITERATION_LIMIT.handle(config), + } + } + + fn current_limit(&self) -> Option { + self.enabled.get().then(|| self.limit.get()) + } +} + +/// Counts the rows a peek has examined on this worker and fails it once that exceeds the limit. +/// +/// A "row" here is a record the worker had to look at, not a record it returned. Records that a +/// literal constraint or the MFP throws away, and records that consolidate to zero, cost scan +/// time all the same, so they count too. +/// +/// Exactly `limit` rows are allowed. The peek only fails when it asks for the row after that. +#[derive(Debug)] +pub(crate) struct PeekRowIterationTracker { + limit: Option, + rows_iterated: usize, +} + +impl PeekRowIterationTracker { + fn new(limit: Option, rows_iterated: usize) -> Self { + Self { + limit, + rows_iterated, + } + } + + /// Adopts a new limit without forgetting the rows already examined. + /// + /// Rows counted while the feature was off still count, so turning it on mid-scan accounts for + /// the work the peek has already caused. + fn set_limit(&mut self, limit: Option) { + self.limit = limit; + } + + fn rows_iterated(&self) -> usize { + self.rows_iterated + } + + fn track_next(&mut self) -> Result<(), PeekError> { + if let Some(limit) = self.limit + && self.rows_iterated >= limit + { + return Err(PeekError::RowIterationLimitExceeded { limit }); + } + + self.rows_iterated = self.rows_iterated.saturating_add(1); + Ok(()) + } +} + +fn peek_row_iteration_limit(config: &ConfigSet) -> Option { + ENABLE_PEEK_ROW_ITERATION_LIMIT + .get(config) + .then(|| PEEK_ROW_ITERATION_LIMIT.get(config)) +} + +#[cfg(test)] +mod tests { + use mz_dyncfg::ConfigUpdates; + + use super::*; + + #[mz_ore::test] + fn row_iteration_limit_observes_updates_and_disabled_rows() { + let config = mz_dyncfgs::all_dyncfgs(); + let row_iteration_config = PeekRowIterationConfig::new(&config); + let mut tracker = PeekRowIterationTracker::new(row_iteration_config.current_limit(), 0); + + tracker.track_next().unwrap(); + tracker.track_next().unwrap(); + + let mut updates = ConfigUpdates::default(); + updates.add(&PEEK_ROW_ITERATION_LIMIT, 3); + updates.add(&ENABLE_PEEK_ROW_ITERATION_LIMIT, true); + updates.apply(&config); + tracker.set_limit(row_iteration_config.current_limit()); + tracker.track_next().unwrap(); + + let mut updates = ConfigUpdates::default(); + updates.add(&ENABLE_PEEK_ROW_ITERATION_LIMIT, false); + updates.apply(&config); + tracker.set_limit(row_iteration_config.current_limit()); + tracker.track_next().unwrap(); + + let mut updates = ConfigUpdates::default(); + updates.add(&PEEK_ROW_ITERATION_LIMIT, 5); + updates.add(&ENABLE_PEEK_ROW_ITERATION_LIMIT, true); + updates.apply(&config); + tracker.set_limit(row_iteration_config.current_limit()); + tracker.track_next().unwrap(); + assert_eq!( + tracker.track_next(), + Err(PeekError::RowIterationLimitExceeded { limit: 5 }) + ); + } +} + /// Worker-local state that is maintained across dataflows. /// /// This state is restricted to the COMPUTE state, the deterministic, idempotent work @@ -756,6 +872,7 @@ impl<'a> ActiveComputeState<'a> { metadata, usize::cast_from(self.compute_state.max_result_size), self.timely_worker, + PeekRowIterationConfig::new(&self.compute_state.worker_config), ) } }; @@ -1006,6 +1123,9 @@ impl<'a> ActiveComputeState<'a> { PendingPeek::Index(peek) => { let start = Instant::now(); + let row_iteration_limit = + peek_row_iteration_limit(&self.compute_state.worker_config); + let peek_stash_eligible = peek .peek .finishing @@ -1056,6 +1176,7 @@ impl<'a> ActiveComputeState<'a> { self.compute_state.max_result_size, peek_stash_enabled && peek_stash_eligible, peek_stash_threshold_bytes, + row_iteration_limit, &metrics, ); @@ -1071,6 +1192,9 @@ impl<'a> ActiveComputeState<'a> { let _span = span!(parent: &peek.span, Level::DEBUG, "process_stash_peek").entered(); + // NOTE: The row iteration limit does not follow a peek into the stash. The + // stash restarts the scan and produces in bounded bursts, so a stashed + // peek may examine any number of rows. let peek_stash_batch_max_runs = PEEK_RESPONSE_STASH_BATCH_MAX_RUNS .get(&self.compute_state.worker_config); @@ -1329,6 +1453,7 @@ impl PendingPeek { metadata: CollectionMetadata, max_result_size: usize, timely_worker: &TimelyWorker, + row_iteration_config: PeekRowIterationConfig, ) -> Self { let active_worker = { // Choose the worker that does the actual peek arbitrarily but consistently. @@ -1366,6 +1491,7 @@ impl PendingPeek { mfp_plan, max_result_size, max_results_needed, + row_iteration_config, ) .await } else { @@ -1373,7 +1499,7 @@ impl PendingPeek { }; let result = match result { Ok(rows) => PeekResponse::Rows(vec![RowCollection::new(rows, &order_by)]), - Err(e) => PeekResponse::Error(e.to_string()), + Err(error) => PeekResponse::Error(error), }; match result_tx.send((result, start.elapsed())) { Ok(()) => {} @@ -1437,11 +1563,12 @@ impl PersistPeek { mfp_plan: SafeMfpPlan, max_result_size: usize, mut limit_remaining: usize, - ) -> Result, String> { + row_iteration_config: PeekRowIterationConfig, + ) -> Result, PeekError> { let client = persist_clients .open(metadata.persist_location) .await - .map_err(|e| e.to_string())?; + .map_err(|e| PeekError::unstructured(e.to_string()))?; let mut reader: ReadHandle = client .open_leased_reader( @@ -1452,7 +1579,7 @@ impl PersistPeek { USE_CRITICAL_SINCE_SNAPSHOT.get(client.dyncfgs()), ) .await - .map_err(|e| e.to_string())?; + .map_err(|e| PeekError::unstructured(e.to_string()))?; // If we are using txn-wal for this collection, then the upper might // be advanced lazily and we have to go through txn-wal for reads. @@ -1478,7 +1605,9 @@ impl PersistPeek { ) .await .map_err(|since| { - format!("attempted to peek at {as_of}, but the since has advanced to {since:?}") + PeekError::unstructured(format!( + "attempted to peek at {as_of}, but the since has advanced to {since:?}" + )) })?; // Re-used state for processing and building rows. @@ -1487,6 +1616,7 @@ impl PersistPeek { let mut row_builder = Row::default(); let arena = RowArena::new(); let mut total_size = 0usize; + let mut row_iteration_tracker = PeekRowIterationTracker::new(None, 0); let literal_len = match &literal_constraint { None => 0, @@ -1498,7 +1628,12 @@ impl PersistPeek { break; }; for (data, _, d) in batch { - let row = data.map_err(|e| e.to_string())?; + // Count before literal and MFP filtering because the Persist row + // has already been read and must still be examined. + row_iteration_tracker.set_limit(row_iteration_config.current_limit()); + row_iteration_tracker.track_next()?; + + let row = data.map_err(PeekError::from)?; if let Some(literal) = &literal_constraint { match row.iter().take(literal_len).cmp(literal.iter()) { @@ -1513,11 +1648,11 @@ impl PersistPeek { shard = %metadata.data_shard, diff = d, ?row, "persist peek encountered negative multiplicities", ); - format!( + PeekError::unstructured(format!( "Invalid data in source, \ saw retractions ({}) for row that does not exist: {:?}", -d, row, - ) + )) })?; let Some(count) = NonZeroUsize::new(count) else { continue; @@ -1526,16 +1661,16 @@ impl PersistPeek { let eval_result = mfp_plan .evaluate_into(&mut datum_local, &arena, &mut row_builder) .map(|row| row.cloned()) - .map_err(|e| e.to_string())?; + .map_err(PeekError::from)?; if let Some(row) = eval_result { total_size = total_size .saturating_add(row.byte_len()) .saturating_add(std::mem::size_of::()); if total_size > max_result_size { - return Err(format!( + return Err(PeekError::unstructured(format!( "result exceeds max size of {}", ByteSize::b(u64::cast_from(max_result_size)) - )); + ))); } result.push((row, count)); limit_remaining = limit_remaining.saturating_sub(count.get()); @@ -1594,6 +1729,7 @@ impl IndexPeek { max_result_size: u64, peek_stash_eligible: bool, peek_stash_threshold_bytes: usize, + row_iteration_limit: Option, metrics: &IndexPeekMetrics<'_>, ) -> PeekStatus { let method_start = Instant::now(); @@ -1614,7 +1750,7 @@ impl IndexPeek { read_frontier.elements(), self.peek.timestamp, ); - return PeekStatus::Ready(PeekResponse::Error(error)); + return PeekStatus::Ready(PeekResponse::Error(PeekError::unstructured(error))); } metrics @@ -1625,6 +1761,7 @@ impl IndexPeek { max_result_size, peek_stash_eligible, peek_stash_threshold_bytes, + row_iteration_limit, metrics, ); @@ -1641,14 +1778,20 @@ impl IndexPeek { max_result_size: u64, peek_stash_eligible: bool, peek_stash_threshold_bytes: usize, + row_iteration_limit: Option, metrics: &IndexPeekMetrics<'_>, ) -> PeekStatus { let error_scan_start = Instant::now(); // Check if there exist any errors and, if so, return whatever one we // find first. + let mut row_iteration_tracker = PeekRowIterationTracker::new(row_iteration_limit, 0); let (mut cursor, storage) = self.trace_bundle.errs_mut().cursor(); while cursor.key_valid(&storage) { + if let Err(error) = row_iteration_tracker.track_next() { + return PeekStatus::Ready(PeekResponse::Error(error)); + } + let mut copies = Diff::ZERO; cursor.map_times(&storage, |time, diff| { if time.less_equal(&self.peek.timestamp) { @@ -1661,14 +1804,15 @@ impl IndexPeek { target = %self.peek.target.id(), diff = %copies, %error, "index peek encountered negative multiplicities in error trace", ); - return PeekStatus::Ready(PeekResponse::Error(format!( + return PeekStatus::Ready(PeekResponse::Error(PeekError::unstructured(format!( "Invalid data in source errors, \ saw retractions ({}) for row that does not exist: {}", -copies, error, - ))); + )))); } if copies.is_positive() { - return PeekStatus::Ready(PeekResponse::Error(cursor.key(&storage).to_string())); + let error = cursor.key(&storage).deserialize(); + return PeekStatus::Ready(PeekResponse::Error(error.into())); } cursor.step_key(&storage); } @@ -1683,6 +1827,8 @@ impl IndexPeek { max_result_size, peek_stash_eligible, peek_stash_threshold_bytes, + row_iteration_limit, + row_iteration_tracker.rows_iterated(), metrics, ) } @@ -1694,6 +1840,8 @@ impl IndexPeek { max_result_size: u64, peek_stash_eligible: bool, peek_stash_threshold_bytes: usize, + row_iteration_limit: Option, + rows_iterated: usize, metrics: &IndexPeekMetrics<'_>, ) -> PeekStatus where @@ -1720,6 +1868,8 @@ impl IndexPeek { peek.timestamp, peek.literal_constraints.clone().as_deref_mut(), oks_handle, + row_iteration_limit, + rows_iterated, ); metrics @@ -1759,10 +1909,10 @@ impl IndexPeek { return PeekStatus::UsePeekStash; } if total_size > max_result_size { - return PeekStatus::Ready(PeekResponse::Error(format!( + return PeekStatus::Ready(PeekResponse::Error(PeekError::unstructured(format!( "result exceeds max size of {}", ByteSize::b(u64::cast_from(max_result_size)) - ))); + )))); } results.push((row, copies)); diff --git a/src/compute/src/compute_state/peek_result_iterator.rs b/src/compute/src/compute_state/peek_result_iterator.rs index 8e3a2967693fb..737f1b99d805c 100644 --- a/src/compute/src/compute_state/peek_result_iterator.rs +++ b/src/compute/src/compute_state/peek_result_iterator.rs @@ -12,18 +12,20 @@ use std::ops::Range; use differential_dataflow::trace::cursor::{BatchCursor, BatchKey, CursorList}; use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::{Cursor, Navigable, TraceReader}; +use mz_compute_client::protocol::response::PeekError; /// The merged cursor a [`TraceReader::cursor`] hands out over all of a trace's batches: a /// [`CursorList`] over the per-batch cursors. type TraceCursor = CursorList>; /// Backing storage for a [`TraceCursor`]: the batches the cursor borrows from. type TraceStorage = Vec<::Batch>; -use mz_ore::result::ResultExt; use mz_repr::fixed_length::ExtendDatums; use mz_repr::{DatumVec, Diff, GlobalId, Row, RowArena}; use timely::order::PartialOrder; -pub struct PeekResultIterator +use super::PeekRowIterationTracker; + +pub(super) struct PeekResultIterator where Tr: TraceReader, { @@ -37,6 +39,8 @@ where datum_vec: DatumVec, literals: Option>, rows_processed: usize, + row_iteration_tracker: PeekRowIterationTracker, + exhausted: bool, } /// Helper to handle literals in peeks @@ -121,12 +125,14 @@ where DiffGat<'a> = &'a Diff, >, { - pub fn new( + pub(super) fn new( target_id: GlobalId, map_filter_project: mz_expr::SafeMfpPlan, peek_timestamp: mz_repr::Timestamp, literal_constraints: Option<&mut [Row]>, trace_reader: &mut Tr, + row_iteration_limit: Option, + rows_iterated: usize, ) -> Self { let (mut cursor, storage) = trace_reader.cursor(); let literals = literal_constraints @@ -142,6 +148,8 @@ where datum_vec: DatumVec::new(), literals, rows_processed: 0, + row_iteration_tracker: PeekRowIterationTracker::new(row_iteration_limit, rows_iterated), + exhausted: false, } } @@ -180,9 +188,13 @@ where DiffGat<'a> = &'a Diff, >, { - type Item = Result<(Row, NonZeroI64), String>; + type Item = Result<(Row, NonZeroI64), PeekError>; fn next(&mut self) -> Option { + if self.exhausted { + return None; + } + let result = loop { if self.literals_exhausted() { return None; @@ -199,6 +211,17 @@ where } } + // Filtered and zero-multiplicity rows still consume worker time, so + // they count against the budget before evaluation. + // + // Failing here leaves the cursor where it is, so latch the iterator + // shut. Otherwise a caller that polls again would get the same error + // forever rather than an end. + if let Err(error) = self.row_iteration_tracker.track_next() { + self.exhausted = true; + return Some(Err(error)); + } + self.rows_processed = self.rows_processed.saturating_add(1); match self.extract_current_row() { Ok(Some(row)) => break Ok(row), @@ -230,7 +253,7 @@ where /// Extracts and returns the row currently pointed at by our cursor. Returns /// `Ok(None)` if our MapFilterProject evaluates to `None`. Also returns any /// errors that arise from evaluating the MapFilterProject. - fn extract_current_row(&mut self) -> Result, String> { + fn extract_current_row(&mut self) -> Result, PeekError> { // TODO: This arena could be maintained and reused for longer, // but it wasn't clear at what interval we should flush // it to ensure we don't accidentally spike our memory use. @@ -260,7 +283,7 @@ where .map_filter_project .evaluate_into(&mut borrow, &arena, &mut self.row_builder) .map(|row| row.cloned()) - .map_err_to_string_with_causes()? + .map_err(PeekError::from)? { let mut copies = Diff::ZERO; self.cursor.map_times(&self.storage, |time, diff| { @@ -274,11 +297,11 @@ where target = %self.target_id, diff = %copies, ?row, "index peek encountered negative multiplicities in ok trace", ); - return Err(format!( + return Err(PeekError::unstructured(format!( "Invalid data in source, \ saw retractions ({}) for row that does not exist: {:?}", -copies, row, - )); + ))); } else { copies.into_inner() }; diff --git a/src/compute/src/compute_state/peek_stash.rs b/src/compute/src/compute_state/peek_stash.rs index 214d9f371f64f..b9f70ce00425d 100644 --- a/src/compute/src/compute_state/peek_stash.rs +++ b/src/compute/src/compute_state/peek_stash.rs @@ -11,7 +11,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use mz_compute_client::protocol::command::Peek; -use mz_compute_client::protocol::response::{PeekResponse, StashedPeekResponse}; +use mz_compute_client::protocol::response::{PeekError, PeekResponse, StashedPeekResponse}; use mz_expr::row::RowCollection; use mz_ore::cast::CastFrom; use mz_ore::task::AbortOnDropHandle; @@ -45,7 +45,7 @@ pub struct StashingPeek { /// We can't give a PeekResultIterator to our async upload task because the /// underlying trace reader is not Send/Sync. So we need to use a channel to /// send result rows from the worker thread to the async background task. - rows_tx: Option, String>>>, + rows_tx: Option, PeekError>>>, /// The result of the background task, eventually. pub result: oneshot::Receiver<(PeekResponse, Duration)>, /// The `tracing::Span` tracking this peek's operation @@ -80,6 +80,8 @@ impl StashingPeek { peek.timestamp, peek.literal_constraints.as_deref_mut(), oks_handle, + None, + 0, ); let rows_needed_by_finishing = peek.finishing.num_rows_needed(); @@ -102,7 +104,7 @@ impl StashingPeek { let result = match result { Ok(peek_response) => peek_response, - Err(e) => PeekResponse::Error(e.to_string()), + Err(e) => PeekResponse::Error(PeekError::unstructured(e.to_string())), }; match result_tx.send((result, start.elapsed())) { Ok(()) => {} @@ -130,7 +132,7 @@ impl StashingPeek { peek_uuid: Uuid, relation_desc: RelationDesc, max_rows: Option, // The number of rows needed by the RowSetFinishing's offset + limit - mut rows_rx: tokio::sync::mpsc::Receiver, String>>, + mut rows_rx: tokio::sync::mpsc::Receiver, PeekError>>, ) -> Result { let client = persist_clients .open(persist_location) diff --git a/src/environmentd/src/http/sql.rs b/src/environmentd/src/http/sql.rs index 9e07c81f27f90..19b7655be8197 100644 --- a/src/environmentd/src/http/sql.rs +++ b/src/environmentd/src/http/sql.rs @@ -647,7 +647,7 @@ impl SqlResult { let mut sql_rows = match peek_response { PeekResponseUnary::Rows(rows) => rows, PeekResponseUnary::Error(e) => { - return Ok(SqlResult::err(client, Error::Unstructured(anyhow!(e)))); + return Ok(SqlResult::err(client, e)); } PeekResponseUnary::DependencyDropped(dep) => { return Ok(SqlResult::err(client, dep.to_concurrent_dependency_drop())); @@ -1063,12 +1063,11 @@ impl ResultSender for WebSocket { } } } - Some(PeekResponseUnary::Error(error)) => { + Some(PeekResponseUnary::Error(err)) => { + let error = err.to_string(); break ( true, - vec![WebSocketResponse::Error( - Error::Unstructured(anyhow!(error.clone())).into(), - )], + vec![WebSocketResponse::Error(err.into())], Some((StatementEndedExecutionReason::Errored { error }, ctx_extra)), ); } diff --git a/src/environmentd/tests/sql.rs b/src/environmentd/tests/sql.rs index 93d8c09c9c5b6..4d77e8adf61c6 100644 --- a/src/environmentd/tests/sql.rs +++ b/src/environmentd/tests/sql.rs @@ -2451,6 +2451,51 @@ fn test_parse_error_codes() { } } +#[mz_ore::test] +fn test_dataflow_error_codes() { + let server = test_util::TestHarness::default().start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + + client + .batch_execute("CREATE TABLE t (a int4, b int4)") + .unwrap(); + client.batch_execute("INSERT INTO t VALUES (1, 0)").unwrap(); + + let cases: &[(&str, &SqlState)] = &[ + ("SELECT a / b FROM t", &SqlState::DIVISION_BY_ZERO), + ( + "SELECT 2147483647 + a FROM t", + &SqlState::NUMERIC_VALUE_OUT_OF_RANGE, + ), + ]; + + for (query, expected) in cases { + let err = client.query_one(*query, &[]).unwrap_err().unwrap_db_error(); + assert_eq!( + err.code(), + *expected, + "unexpected SQLSTATE {} for query `{query}`: {}", + err.code().code(), + err.message(), + ); + } + + client + .batch_execute("CREATE MATERIALIZED VIEW mv AS SELECT a / b AS x FROM t;") + .unwrap(); + let err = client + .query_one("SELECT * FROM mv", &[]) + .unwrap_err() + .unwrap_db_error(); + assert_eq!( + err.code(), + &SqlState::DIVISION_BY_ZERO, + "unexpected SQLSTATE {} reading from materialized view: {}", + err.code().code(), + err.message(), + ); +} + #[mz_ore::test] #[allow(clippy::disallowed_methods)] fn test_emit_timestamp_notice() { diff --git a/src/environmentd/tests/testdata/http/ws b/src/environmentd/tests/testdata/http/ws index 17dbeff38719b..ade66fa8f4e76 100644 --- a/src/environmentd/tests/testdata/http/ws +++ b/src/environmentd/tests/testdata/http/ws @@ -414,7 +414,7 @@ ws-text fixid=true ---- {"type":"Notice","payload":{"message":"{\n \"plans\": {\n \"raw\": {\n \"text\": \"Finish limit=1 output=[#0]\\n Project (#15)\\n Map ((1 / 0))\\n Get mz_catalog.mz_sources\\n\\nTarget cluster: mz_catalog_server\\n\",\n \"json\": {\n \"Project\": {\n \"input\": {\n \"Map\": {\n \"input\": {\n \"Get\": {\n \"id\": {\n \"Global\": {\n \n }\n },\n \"typ\": {\n \"column_types\": [\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"Oid\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": {\n \"Array\": \"MzAclItem\"\n },\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n }\n ],\n \"keys\": [\n [\n 0\n ],\n [\n 1\n ]\n ]\n }\n }\n },\n \"scalars\": [\n {\n \"CallBinary\": {\n \"func\": {\n \"DivInt32\": null\n },\n \"expr1\": {\n \"Literal\": [\n {\n \"data\": [\n 45,\n 1\n ]\n },\n {\n \"scalar_type\": \"Int32\",\n \"nullable\": false\n },\n null\n ]\n },\n \"expr2\": {\n \"Literal\": [\n {\n \"data\": [\n 44\n ]\n },\n {\n \"scalar_type\": \"Int32\",\n \"nullable\": false\n },\n null\n ]\n },\n \"name\": null\n }\n }\n ]\n }\n },\n \"outputs\": [\n 15\n ]\n }\n }\n },\n \"optimized\": {\n \"global\": {\n \"text\": \":\\n Finish limit=1 output=[#0]\\n ArrangeBy keys=[[#0]]\\n ReadGlobalFromSameDataflow \\n\\n:\\n Project (#15)\\n Map (error(\\\"division by zero\\\"))\\n ReadIndex on=mz_sources mz_sources_ind=[*** full scan ***]\\n\\nTarget cluster: mz_catalog_server\\n\",\n \"json\": {\n \"plans\": [\n {\n \"id\": \"\",\n \"plan\": {\n \"ArrangeBy\": {\n \"input\": {\n \"Get\": {\n \"id\": {\n \"Global\": {\n \n }\n },\n \"typ\": {\n \"column_types\": [\n {\n \"scalar_type\": \"Int32\",\n \"nullable\": false\n }\n ],\n \"keys\": []\n },\n \"access_strategy\": \"SameDataflow\"\n }\n },\n \"keys\": [\n [\n {\n \"Column\": [\n 0,\n null\n ]\n }\n ]\n ]\n }\n }\n },\n {\n \"id\": \"\",\n \"plan\": {\n \"Project\": {\n \"input\": {\n \"Map\": {\n \"input\": {\n \"Get\": {\n \"id\": {\n \"Global\": {\n \n }\n },\n \"typ\": {\n \"column_types\": [\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"UInt32\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": false\n },\n {\n \"scalar_type\": {\n \"Array\": \"MzAclItem\"\n },\n \"nullable\": false\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n },\n {\n \"scalar_type\": \"String\",\n \"nullable\": true\n }\n ],\n \"keys\": [\n [\n 0\n ],\n [\n 1\n ]\n ]\n },\n \"access_strategy\": {\n \"Index\": [\n [\n {\n \n },\n \"FullScan\"\n ]\n ]\n }\n }\n },\n \"scalars\": [\n {\n \"Literal\": [\n {\n \"Err\": \"DivisionByZero\"\n },\n {\n \"scalar_type\": \"Int32\",\n \"nullable\": false\n }\n ]\n }\n ]\n }\n },\n \"outputs\": [\n 15\n ]\n }\n }\n }\n ],\n \"sources\": []\n }\n },\n \"fast_path\": {\n \"text\": \"Explained Query (fast path):\\n Finish limit=1 output=[#0]\\n →Map/Filter/Project\\n Project: #15\\n Map: error(\\\"division by zero\\\")\\n →Indexed mz_catalog.mz_sources (using mz_catalog.mz_sources_ind)\\n\\nTarget cluster: mz_catalog_server\\n\",\n \"json\": {\n \"plans\": [\n {\n \"id\": \"Explained Query (fast path)\",\n \"plan\": {\n \"PeekExisting\": [\n {\n \n },\n {\n \n },\n null,\n {\n \"mfp\": {\n \"expressions\": [\n {\n \"Literal\": [\n {\n \"Err\": \"DivisionByZero\"\n },\n {\n \"scalar_type\": \"Int32\",\n \"nullable\": false\n }\n ]\n }\n ],\n \"predicates\": [],\n \"projection\": [\n 15\n ],\n \"input_arity\": 15\n }\n }\n ]\n }\n }\n ],\n \"sources\": []\n }\n }\n }\n },\n \"insights\": {\n \"imports\": {\n \"\": {\n \"name\": {\n \"schema\": \"mz_catalog\",\n \"item\": \"mz_sources_ind\"\n },\n \"type\": \"compute\"\n }\n },\n \"fast_path_clusters\": {},\n \"fast_path_limit\": null,\n \"persist_count\": []\n },\n \"cluster\": {\n \"name\": \"mz_catalog_server\",\n \"id\": {\n \n }\n },\n \"redacted_sql\": \"SELECT '' / '' FROM [ AS mz_catalog.mz_sources] LIMIT ''\"\n}","code":"MZ001","severity":"notice"}} {"type":"CommandStarting","payload":{"has_rows":false,"is_streaming":false}} -{"type":"Error","payload":{"message":"division by zero","code":"XX000"}} +{"type":"Error","payload":{"message":"Evaluation error: division by zero","code":"22012"}} {"type":"ReadyForQuery","payload":"I"} ws-text rows=2 fixtimestamp=true diff --git a/src/pgwire/src/protocol.rs b/src/pgwire/src/protocol.rs index ce39b6cbb9261..0a7149baa4000 100644 --- a/src/pgwire/src/protocol.rs +++ b/src/pgwire/src/protocol.rs @@ -2580,7 +2580,7 @@ where None => FetchResult::Rows(None), Some(PeekResponseUnary::Rows(rows)) => FetchResult::Rows(Some(rows)), Some(PeekResponseUnary::Error(err)) => { - FetchResult::Error(ErrorResponse::error(SqlState::INTERNAL_ERROR, err)) + FetchResult::Error(err.into_response(Severity::Error)) } Some(PeekResponseUnary::DependencyDropped(dep)) => { FetchResult::Error( @@ -2844,11 +2844,10 @@ where e = self.conn.wait_closed() => return Err(e), batch = stream.recv() => match batch { None => break, - Some(PeekResponseUnary::Error(text)) => { - let err = - ErrorResponse::error(SqlState::INTERNAL_ERROR, text.clone()); + Some(PeekResponseUnary::Error(err)) => { + let text = err.to_string(); return self - .send_error_and_get_state(err) + .send_error_and_get_state(err.into_response(Severity::Error)) .await .map(|state| (state, SendRowsEndedReason::Errored { error: text })); } diff --git a/src/sqllogictest/Cargo.toml b/src/sqllogictest/Cargo.toml index da6264f1dca92..a7df0a19b07e4 100644 --- a/src/sqllogictest/Cargo.toml +++ b/src/sqllogictest/Cargo.toml @@ -25,6 +25,7 @@ mz-adapter-types = { path = "../adapter-types" } mz-authenticator = { path = "../authenticator", default-features = false } mz-build-info = { path = "../build-info" } mz-catalog = { path = "../catalog" } +mz-compute-types = { path = "../compute-types" } mz-controller = { path = "../controller" } mz-dyncfgs = { path = "../dyncfgs" } mz-environmentd = { path = "../environmentd", default-features = false } diff --git a/src/sqllogictest/src/bin/sqllogictest.rs b/src/sqllogictest/src/bin/sqllogictest.rs index ec6aecd00d419..45373a18fdf60 100644 --- a/src/sqllogictest/src/bin/sqllogictest.rs +++ b/src/sqllogictest/src/bin/sqllogictest.rs @@ -19,6 +19,7 @@ use std::process::ExitCode; use chrono::Utc; use clap::ArgAction; use mz_adapter_types::dyncfgs::ENABLE_BACKGROUND_ALTER_CLUSTER; +use mz_compute_types::dyncfgs::{ENABLE_PEEK_ROW_ITERATION_LIMIT, PEEK_ROW_ITERATION_LIMIT}; use mz_orchestrator_tracing::{StaticTracingConfig, TracingCliArgs}; use mz_ore::cli::{self, CliConfig, KeyValueArg}; use mz_ore::metrics::MetricsRegistry; @@ -184,6 +185,16 @@ async fn main() -> ExitCode { .entry(ENABLE_BACKGROUND_ALTER_CLUSTER.name().to_string()) .or_insert_with(|| "true".to_string()); + // Keep the guard enabled in the suite without constraining normal test queries. + for (name, value) in [ + (ENABLE_PEEK_ROW_ITERATION_LIMIT.name(), "true"), + (PEEK_ROW_ITERATION_LIMIT.name(), "1000000000"), + ] { + system_parameter_defaults + .entry(name.to_string()) + .or_insert_with(|| value.to_string()); + } + let config = RunConfig { stdout: &OutputStream::new(io::stdout(), args.timestamps), stderr: &OutputStream::new(io::stderr(), args.timestamps), diff --git a/test/launchdarkly-flag-consistency/mzcompose.py b/test/launchdarkly-flag-consistency/mzcompose.py index 094c20f552554..3e08bea538182 100644 --- a/test/launchdarkly-flag-consistency/mzcompose.py +++ b/test/launchdarkly-flag-consistency/mzcompose.py @@ -210,6 +210,7 @@ compute_flat_map_fuel compute_logical_backpressure_max_retained_capabilities compute_mv_sink_advance_persist_frontiers + compute_peek_row_iteration_limit compute_peek_response_stash_batch_max_runs compute_peek_response_stash_read_batch_size_bytes compute_peek_response_stash_read_memory_budget_bytes @@ -241,6 +242,7 @@ enable_bounded_staleness_isolation enable_coalesce_case_transform enable_compute_half_join2 + enable_compute_peek_row_iteration_limit enable_compute_render_fueled_as_specific_collection enable_date_bin_hopping enable_default_connection_validation diff --git a/test/sqllogictest/max_result_size.slt b/test/sqllogictest/max_result_size.slt index e165c3e4f16c9..26878379d9291 100644 --- a/test/sqllogictest/max_result_size.slt +++ b/test/sqllogictest/max_result_size.slt @@ -9,6 +9,89 @@ mode cockroach +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET compute_peek_row_iteration_limit TO 2; +---- +COMPLETE 0 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET enable_compute_peek_row_iteration_limit TO false; +---- +COMPLETE 0 + +statement ok +CREATE CLUSTER peek_limit_cluster SIZE 'scale=1,workers=1'; + +statement ok +SET cluster TO 'peek_limit_cluster'; + +statement ok +CREATE TABLE peek_iteration_limit (a int); + +statement ok +INSERT INTO peek_iteration_limit VALUES (1), (2); + +statement ok +CREATE INDEX peek_iteration_limit_idx ON peek_iteration_limit (a); + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET enable_compute_peek_row_iteration_limit TO true; +---- +COMPLETE 0 + +# A peek may examine exactly the configured number of rows. +query I +SELECT a FROM peek_iteration_limit ORDER BY a; +---- +1 +2 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET enable_compute_peek_row_iteration_limit TO false; +---- +COMPLETE 0 + +statement ok +INSERT INTO peek_iteration_limit VALUES (3); + +query I +SELECT a FROM peek_iteration_limit ORDER BY a; +---- +1 +2 +3 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET enable_compute_peek_row_iteration_limit TO true; +---- +COMPLETE 0 + +query error query exceeded the configured row iteration limit of 2 rows +SELECT a FROM peek_iteration_limit ORDER BY a; + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET compute_peek_row_iteration_limit TO 3; +---- +COMPLETE 0 + +query I +SELECT a FROM peek_iteration_limit ORDER BY a; +---- +1 +2 +3 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM RESET compute_peek_row_iteration_limit; +---- +COMPLETE 0 + +statement ok +SET cluster TO 'quickstart'; + +statement ok +DROP CLUSTER peek_limit_cluster CASCADE; + simple conn=mz_system,user=mz_system ALTER SYSTEM SET max_result_size TO 1; ---- diff --git a/test/sqllogictest/persist-fast-path.slt b/test/sqllogictest/persist-fast-path.slt index a61cf5a061787..bb4c47f46cb2f 100644 --- a/test/sqllogictest/persist-fast-path.slt +++ b/test/sqllogictest/persist-fast-path.slt @@ -244,7 +244,38 @@ ALTER SYSTEM RESET enable_compute_peek_response_stash ---- COMPLETE 0 -# Does not apply when an index exists. +# Persist-backed peeks enforce the same row iteration limit. +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET enable_compute_peek_row_iteration_limit TO true; +---- +COMPLETE 0 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET compute_peek_row_iteration_limit TO 2; +---- +COMPLETE 0 + +query error query exceeded the configured row iteration limit of 2 rows +SELECT a FROM large_rows LIMIT 3; + +simple conn=mz_system,user=mz_system +ALTER SYSTEM SET compute_peek_row_iteration_limit TO 3; +---- +COMPLETE 0 + +query I rowsort +SELECT a FROM large_rows LIMIT 3; +---- +1 +2 +3 + +simple conn=mz_system,user=mz_system +ALTER SYSTEM RESET compute_peek_row_iteration_limit; +---- +COMPLETE 0 + +# The Persist fast path does not apply when an index exists. statement ok CREATE DEFAULT INDEX ON numbers;