From 15debdab5401cb0f3170797166a85af933065ffb Mon Sep 17 00:00:00 2001 From: Dennis Felsing Date: Wed, 5 Aug 2026 13:23:24 +0200 Subject: [PATCH 1/6] sql-parser: only start a map literal when a bracket follows MAP (#37973) `parse_option_map` committed on the `MAP` keyword alone and then demanded `[`, so a bare `map` in an option-value position was a hard parse error rather than falling through to the item name it is. `map` is a legal bare identifier everywhere else. `SELECT "map"` prints as `SELECT map` and reparses fine, so `AstDisplay` prints an option value that is the item name `map` unquoted, and that output failed to reparse: CREATE SINK s FROM t INTO KAFKA CONNECTION c (TOPIC = "map") printed `TOPIC = map` and then failed with "Expected left square bracket, found right parenthesis". Every generic option value was affected, a materialized view's `PARTITION BY` among them. Look ahead for the bracket instead. The alternative, quoting `map` in `can_be_printed_bare`, would change the printed form of every `map` identifier in every position, and would change what `quote_ident('map')` returns, to work around one non-backtracking call site. The sibling branches of `parse_option_value` all backtrack already. Expression position never needed that quoting, but the comment in `can_be_printed_bare` gave the wrong reason: the map-literal dispatch keys off the `[`, not the `=>` the comment cited, so `map[1]` is not "unambiguously a subscript" on its own. What keeps a bare `map` safe there is `write_subscript_receiver`, which parenthesizes a context-sensitive keyword receiver into `(map)[1]`. Corrected in place, no behavior change. Found by the sql_roundtrip cargo-fuzz target. The datadriven case pins the printed form and AST. The added parser test covers the sink, source and materialized-view option positions, a genuine `MAP['a' => 'b']` option value that must still take the map branch, and the expression forms the corrected comment relies on. Co-authored-by: Claude Opus 5 (1M context) --- src/sql-parser/src/ast/defs/name.rs | 9 ++++++--- src/sql-parser/src/parser.rs | 9 +++++++++ src/sql-parser/tests/sqlparser_common.rs | 23 +++++++++++++++++++++++ src/sql-parser/tests/testdata/ddl | 10 ++++++++++ 4 files changed, 48 insertions(+), 3 deletions(-) diff --git a/src/sql-parser/src/ast/defs/name.rs b/src/sql-parser/src/ast/defs/name.rs index f732b52ff8a06..c6559ebdfcc4d 100644 --- a/src/sql-parser/src/ast/defs/name.rs +++ b/src/sql-parser/src/ast/defs/name.rs @@ -333,9 +333,12 @@ impl Ident { // (`list[1]` is a valid one-element list), so a bare `list` // identifier that gets subscripted — `"list"[1]` — would // reparse as a list literal instead of a subscript. (`ARRAY` - // is reserved-in-scalar-expression and so already quoted; - // `MAP[...]` requires `=>`, so `map[1]` is unambiguously a - // subscript.) + // is reserved-in-scalar-expression and so already quoted. + // `MAP` needs no clause here even though `map[…]` is a + // literal too: `MAP` is a context-sensitive keyword, so + // `write_subscript_receiver` parenthesizes the receiver into + // `(map)[1]`, and every grammar that starts a map literal + // requires a `[` right after the keyword.) || kw == LIST // `DEALLOCATE [PREPARE] ` accepts an optional // `PREPARE` keyword before the name, so a bare `prepare` diff --git a/src/sql-parser/src/parser.rs b/src/sql-parser/src/parser.rs index e0c2a39f779eb..34763fdb55465 100644 --- a/src/sql-parser/src/parser.rs +++ b/src/sql-parser/src/parser.rs @@ -5721,6 +5721,15 @@ impl<'a> Parser<'a> { fn parse_option_map( &mut self, ) -> Result>>, ParserError> { + // `MAP` only begins a map literal when a `[` follows it. Committing on + // the keyword alone makes a bare `map` in an option-value position a hard + // error on the missing bracket, rather than falling through to the item + // name it is: `TOPIC CONFIG = "map"` prints as `TOPIC CONFIG = map`, + // since `map` is a legal bare identifier everywhere else, and that output + // then failed to reparse. + if !(self.peek_keyword(MAP) && self.peek_nth_token(1) == Some(Token::LBracket)) { + return Ok(None); + } Ok(if self.parse_keyword(MAP) { self.expect_token(&Token::LBracket)?; let mut map = BTreeMap::new(); diff --git a/src/sql-parser/tests/sqlparser_common.rs b/src/sql-parser/tests/sqlparser_common.rs index 1e67fb0c313fc..27326a43fb58d 100644 --- a/src/sql-parser/tests/sqlparser_common.rs +++ b/src/sql-parser/tests/sqlparser_common.rs @@ -878,6 +878,29 @@ fn test_list_keyword_bare_identifier_subscript_display_roundtrip() { } } +#[mz_ore::test] +#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux` +fn test_map_keyword_bare_identifier_option_value_display_roundtrip() { + // A `MAP` that no `[` follows is the item name `map`, which is how a quoted + // `"map"` prints. `parse_option_map` used to commit on the keyword alone, so + // the printed form failed to reparse in every generic option-value position. + // Regression for the sql_roundtrip fuzz finding `(TOPIC = "map")`. + for sql in [ + r#"CREATE SINK s FROM t INTO KAFKA CONNECTION c (TOPIC = "map") FORMAT BYTES ENVELOPE DEBEZIUM"#, + r#"CREATE SOURCE s FROM KAFKA CONNECTION c (TOPIC = "map") FORMAT BYTES"#, + r#"CREATE MATERIALIZED VIEW v WITH (PARTITION BY = "map") AS SELECT 1"#, + // A real map-literal option value still takes the map branch. + r#"CREATE SINK s FROM t INTO KAFKA CONNECTION c (TOPIC = 't', TOPIC CONFIG = MAP['a' => 'b']) FORMAT BYTES ENVELOPE DEBEZIUM"#, + // Expression position needs no quoting: a subscripted receiver prints as + // `(map)[1]`, which is why `can_be_printed_bare` has no `MAP` clause. + r#"SELECT "map""#, + r#"SELECT "map"[1]"#, + r#"SELECT map['a' => 1]"#, + ] { + assert_display_roundtrips(sql); + } +} + #[mz_ore::test] #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux` fn test_table_function_special_name_display_roundtrip() { diff --git a/src/sql-parser/tests/testdata/ddl b/src/sql-parser/tests/testdata/ddl index 62597942f6bf5..6ee7e4a46002d 100644 --- a/src/sql-parser/tests/testdata/ddl +++ b/src/sql-parser/tests/testdata/ddl @@ -3781,3 +3781,13 @@ CREATE SINK snk FROM t INTO KAFKA CONNECTION conn1 (TOPIC 'topic') FORMAT AVRO U CREATE SINK snk FROM t INTO KAFKA CONNECTION conn1 (TOPIC = 'topic') FORMAT AVRO USING AWS GLUE SCHEMA REGISTRY CONNECTION glue_conn (KEY SCHEMA NAME = 'k', VALUE SCHEMA NAME = 'v', KEY COMPATIBILITY LEVEL = 'BACKWARD', VALUE COMPATIBILITY LEVEL = 'FULL') ENVELOPE UPSERT => CreateSink(CreateSinkStatement { name: Some(UnresolvedItemName([Ident("snk")])), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("t")])), connection: Kafka { connection: Name(UnresolvedItemName([Ident("conn1")])), options: [KafkaSinkConfigOption { name: Topic, value: Some(Value(String("topic"))) }], key: None, headers: None }, format: Some(Bare(Avro(Glue { connection: Name(UnresolvedItemName([Ident("glue_conn")])), with_options: [GlueAvroOption { name: KeySchemaName, value: Some(Value(String("k"))) }, GlueAvroOption { name: ValueSchemaName, value: Some(Value(String("v"))) }, GlueAvroOption { name: KeyCompatibilityLevel, value: Some(Value(String("BACKWARD"))) }, GlueAvroOption { name: ValueCompatibilityLevel, value: Some(Value(String("FULL"))) }], seed: None }))), envelope: Some(Upsert), mode: None, with_options: [] }) + +# An option value that is the bare identifier `map`. `MAP` only starts a map +# literal when a `[` follows, so this is the item name `map`, which is what +# `AstDisplay` prints for a quoted `"map"`. +parse-statement +CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC CONFIG = "map") +---- +CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC CONFIG = map) +=> +CreateSink(CreateSinkStatement { name: Some(UnresolvedItemName([Ident("foo")])), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaSinkConfigOption { name: TopicConfig, value: Some(UnresolvedItemName(UnresolvedItemName([Ident("map")]))) }], key: None, headers: None }, format: None, envelope: None, mode: None, with_options: [] }) From dd328dbda5b6677427e903f846519482307cd7ee Mon Sep 17 00:00:00 2001 From: Dennis Felsing Date: Wed, 5 Aug 2026 14:10:01 +0200 Subject: [PATCH 2/6] pgwire: decode extended protocol count fields as unsigned (#37934) The Int16 counts that precede the repeated groups of Parse and Bind (parameter types, format codes, parameter values) were decoded as signed, but PostgreSQL decodes them as unsigned, so a client may send up to 65535 parameters. A count above 32767 read as negative, so the group was skipped and the cursor was left misaligned for the rest of the message. Closes: [SQL-491](https://linear.app/materializeinc/issue/SQL-491) Co-authored-by: Claude Opus 5 (1M context) --- src/environmentd/tests/pgwire.rs | 151 ++++++++++++++++++++++++++++++- src/pgwire-common/src/codec.rs | 36 +++++--- src/pgwire-common/src/format.rs | 4 +- src/pgwire/src/codec.rs | 20 ++-- 4 files changed, 189 insertions(+), 22 deletions(-) diff --git a/src/environmentd/tests/pgwire.rs b/src/environmentd/tests/pgwire.rs index fd6fcbb9886d5..5ad74336c2408 100644 --- a/src/environmentd/tests/pgwire.rs +++ b/src/environmentd/tests/pgwire.rs @@ -16,7 +16,7 @@ use std::net::TcpStream; use std::path::Path; use std::time::Duration; -use bytes::BytesMut; +use bytes::{BufMut, BytesMut}; use fallible_iterator::FallibleIterator; use mz_adapter::session::DEFAULT_DATABASE_NAME; use mz_environmentd::test_util::{self, PostgresErrorExt}; @@ -1125,3 +1125,152 @@ fn test_many_columns() { } } } + +/// Writes a frontend message with type byte `typ`, filling in its length prefix. +/// +/// `postgres_protocol`'s builders cap repeated groups at `i16::MAX` entries, so +/// a message with more parameters than that must be assembled by hand. +fn put_frontend_message(buf: &mut BytesMut, typ: u8, body: F) { + buf.put_u8(typ); + let base = buf.len(); + buf.put_u32(0); + body(buf); + let len = u32::try_from(buf.len() - base).unwrap(); + buf[base..base + 4].copy_from_slice(&len.to_be_bytes()); +} + +/// Reads backend messages up to and including the `ReadyForQuery` that answers +/// `Sync`. +fn read_until_ready(stream: &mut TcpStream) -> Vec { + use postgres_protocol::message::backend::Message; + + let mut buf = BytesMut::new(); + let mut chunk = [0; 1 << 13]; + let mut msgs = Vec::new(); + loop { + if let Some(msg) = Message::parse(&mut buf).unwrap() { + let ready = matches!(msg, Message::ReadyForQuery(_)); + msgs.push(msg); + if ready { + return msgs; + } + } else { + let n = stream.read(&mut chunk).unwrap(); + assert!(n > 0, "connection closed after {} messages", msgs.len()); + buf.extend_from_slice(&chunk[..n]); + } + } +} + +// The counts that precede the repeated groups of Parse and Bind (parameter +// types, format codes, parameter values) are decoded as unsigned, like +// PostgreSQL does, so a client may send more than `i16::MAX` parameters. +// Reading such a count as signed made it negative, which dropped the group and +// left the decoder misaligned for the rest of the message. See SQL-491. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_many_bind_params() { + use postgres_protocol::message::backend::Message; + use postgres_protocol::message::frontend; + + // One past `i16::MAX`. The wire format allows 65535, but a bind that large + // runs into the unrelated `MAX_REQUEST_SIZE` limit. + const PARAMS: usize = 32768; + const INT4_OID: u32 = 23; + // Bound to the last parameter, so the inserted row matches only if the + // decoder stayed aligned to the end of the parameter array. + const LAST_VALUE: i32 = 4242; + + let server = test_util::TestHarness::default().start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE many_params (a int)") + .unwrap(); + + let mut stream = TcpStream::connect(server.sql_local_addr()).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(120))) + .unwrap(); + let mut buf = BytesMut::new(); + frontend::startup_message( + vec![ + ("user", "materialize"), + ("database", DEFAULT_DATABASE_NAME), + ("options", "--welcome_message=off"), + ], + &mut buf, + ) + .unwrap(); + stream.write_all(&buf).unwrap(); + read_until_ready(&mut stream); + + // The shape from the report: pgjdbc's reWriteBatchedInserts rewrites a batch + // into one insert whose placeholder count exceeds `i16::MAX`. + let placeholders = (1..=PARAMS) + .map(|i| format!("${i}")) + .collect::>() + .join(","); + let sql = format!("INSERT INTO many_params VALUES (coalesce({placeholders}))"); + + // The statement and portal are unnamed throughout, hence the bare `0`s. + buf.clear(); + put_frontend_message(&mut buf, b'P', |buf| { + buf.put_u8(0); + buf.put_slice(sql.as_bytes()); + buf.put_u8(0); + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 0..PARAMS { + buf.put_u32(INT4_OID); + } + }); + put_frontend_message(&mut buf, b'D', |buf| { + buf.put_u8(b'S'); // the statement, to get a ParameterDescription back + buf.put_u8(0); + }); + put_frontend_message(&mut buf, b'B', |buf| { + buf.put_u8(0); + buf.put_u8(0); + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 0..PARAMS { + buf.put_u16(1); // binary + } + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 1..PARAMS { + buf.put_i32(-1); // NULL + } + buf.put_i32(4); + buf.put_i32(LAST_VALUE); + buf.put_u16(0); // results in text + }); + put_frontend_message(&mut buf, b'E', |buf| { + buf.put_u8(0); + buf.put_i32(0); // no row limit + }); + frontend::sync(&mut buf); + stream.write_all(&buf).unwrap(); + + let mut described_params = None; + let mut tags = Vec::new(); + for msg in read_until_ready(&mut stream) { + match msg { + Message::ErrorResponse(body) => { + let fields: Vec<_> = body + .fields() + .map(|f| Ok(String::from_utf8_lossy(f.value_bytes()).into_owned())) + .collect() + .unwrap(); + panic!("unexpected error response: {fields:?}"); + } + Message::ParameterDescription(body) => { + described_params = Some(body.parameters().count().unwrap()); + } + Message::CommandComplete(body) => tags.push(body.tag().unwrap().to_string()), + _ => (), + } + } + assert_eq!(described_params, Some(PARAMS)); + assert_eq!(tags, ["INSERT 0 1"]); + + let row = client.query_one("SELECT a FROM many_params", &[]).unwrap(); + assert_eq!(row.get::<_, i32>(0), LAST_VALUE); +} diff --git a/src/pgwire-common/src/codec.rs b/src/pgwire-common/src/codec.rs index 84fac7a1ea3c2..04a9763a322d7 100644 --- a/src/pgwire-common/src/codec.rs +++ b/src/pgwire-common/src/codec.rs @@ -52,6 +52,7 @@ impl fmt::Display for CodecError { pub trait Pgbuf: BufMut { fn put_string(&mut self, s: &str); fn put_length_i16(&mut self, len: usize) -> Result<(), io::Error>; + fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error>; fn put_format_i8(&mut self, format: Format); fn put_format_i16(&mut self, format: Format); } @@ -70,6 +71,17 @@ impl Pgbuf for B { Ok(()) } + /// Writes a count field as unsigned, so it may exceed 32767. The protocol + /// calls these fields `Int16`, but PostgreSQL clients decode them as + /// unsigned. + fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error> { + let len = u16::try_from(len).map_err(|_| { + io::Error::new(io::ErrorKind::InvalidData, "length does not fit in a u16") + })?; + self.put_u16(len); + Ok(()) + } + fn put_format_i8(&mut self, format: Format) { self.put_i8(format.into()) } @@ -285,17 +297,6 @@ impl<'a> Cursor<'a> { } } - /// Reads the next 16-bit signed integer, advancing the cursor by two - /// bytes. - pub fn read_i16(&mut self) -> Result { - if self.buf.len() < 2 { - return Err(input_err("not enough buffer for an Int16")); - } - let val = NetworkEndian::read_i16(self.buf); - self.advance(2); - Ok(val) - } - /// Reads the next 32-bit signed integer, advancing the cursor by four /// bytes. pub fn read_i32(&mut self) -> Result { @@ -307,6 +308,17 @@ impl<'a> Cursor<'a> { Ok(val) } + /// Reads the next 16-bit unsigned integer, advancing the cursor by two + /// bytes. + pub fn read_u16(&mut self) -> Result { + if self.buf.len() < 2 { + return Err(input_err("not enough buffer for an Int16")); + } + let val = NetworkEndian::read_u16(self.buf); + self.advance(2); + Ok(val) + } + /// Reads the next 32-bit unsigned integer, advancing the cursor by four /// bytes. pub fn read_u32(&mut self) -> Result { @@ -320,7 +332,7 @@ impl<'a> Cursor<'a> { /// Reads the next 16-bit format code, advancing the cursor by two bytes. pub fn read_format(&mut self) -> Result { - Format::try_from(self.read_i16()?) + Format::try_from(self.read_u16()?) } /// Advances the cursor by `n` bytes. diff --git a/src/pgwire-common/src/format.rs b/src/pgwire-common/src/format.rs index 0502c3829c617..ebc72ee20dbef 100644 --- a/src/pgwire-common/src/format.rs +++ b/src/pgwire-common/src/format.rs @@ -26,10 +26,10 @@ pub enum Format { Binary, } -impl TryFrom for Format { +impl TryFrom for Format { type Error = io::Error; - fn try_from(value: i16) -> Result { + fn try_from(value: u16) -> Result { match value { 0 => Ok(Format::Text), 1 => Ok(Format::Binary), diff --git a/src/pgwire/src/codec.rs b/src/pgwire/src/codec.rs index 91311c07986ab..e54583bb0e4fb 100644 --- a/src/pgwire/src/codec.rs +++ b/src/pgwire/src/codec.rs @@ -430,17 +430,17 @@ impl Codec { dst.put_u32(secret_key); } BackendMessage::ParameterDescription(params) => { - if params.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") { + if params.len() > usize::from(u16::MAX) { return Err(io::Error::new( io::ErrorKind::InvalidData, format!( "{} params in parameter description, which exceeds {}", params.len(), - i16::MAX + u16::MAX ), )); } - dst.put_length_i16(params.len())?; + dst.put_length_u16(params.len())?; for param in params { dst.put_u32(param.oid()); } @@ -792,8 +792,12 @@ fn decode_parse(mut buf: Cursor) -> Result { let name = buf.read_cstr()?; let sql = buf.read_cstr()?; + // NOTE: the protocol calls the counts that precede repeated groups `Int16`, + // but PostgreSQL decodes them as unsigned, so they reach 65535. Reading one + // as signed makes it negative, which drops the group and leaves the cursor + // misaligned for the rest of the message. let mut param_types = vec![]; - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { param_types.push(buf.read_u32()?); } @@ -833,13 +837,15 @@ fn decode_bind(mut buf: Cursor) -> Result { let portal_name = buf.read_cstr()?.to_string(); let statement_name = buf.read_cstr()?.to_string(); + // The three counts below are `Int16` in the protocol but decoded as + // unsigned, for the reason spelled out in `decode_parse`. let mut param_formats = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { param_formats.push(buf.read_format()?); } let mut raw_params = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { let len = buf.read_i32()?; if len == -1 { raw_params.push(None); // NULL @@ -854,7 +860,7 @@ fn decode_bind(mut buf: Cursor) -> Result { } let mut result_formats = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { result_formats.push(buf.read_format()?); } From 26305d8cb00f7a17b09234f540b4da7c283a284d Mon Sep 17 00:00:00 2001 From: Dennis Felsing Date: Wed, 5 Aug 2026 14:46:23 +0200 Subject: [PATCH 3/6] ci: Nightly fixes (2026-08-05) (#38056) Based on https://buildkite.com/materialize/nightly/builds/17698 Test run: https://buildkite.com/materialize/nightly/builds/17700 & https://buildkite.com/materialize/nightly/builds/17716 --- src/adapter/src/config/frontend.rs | 2 +- src/dyncfg-launchdarkly/src/lib.rs | 2 +- src/ore/src/iter.rs | 3 +++ src/ore/src/pool.rs | 10 +++++++++- src/repr/src/adt/datetime.rs | 1 + src/repr/src/adt/interval.rs | 1 + src/repr/src/adt/timestamp.rs | 1 + src/timely-util/src/columnar.rs | 1 + src/timely-util/src/columnar/chunk.rs | 4 ++++ 9 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/adapter/src/config/frontend.rs b/src/adapter/src/config/frontend.rs index a65a637a9f179..53d42f08ca482 100644 --- a/src/adapter/src/config/frontend.rs +++ b/src/adapter/src/config/frontend.rs @@ -471,7 +471,7 @@ async fn ld_client( let mut backoff = Duration::from_secs(5); let timeout = Duration::from_secs(10); - // TODO(materialize#32030): fix retry logic + // TODO(materialize#38055): fix retry logic loop { match ld_client.wait_for_initialization(timeout).await { Some(true) => break, diff --git a/src/dyncfg-launchdarkly/src/lib.rs b/src/dyncfg-launchdarkly/src/lib.rs index 261a5cf695211..0894d305b210e 100644 --- a/src/dyncfg-launchdarkly/src/lib.rs +++ b/src/dyncfg-launchdarkly/src/lib.rs @@ -82,7 +82,7 @@ where let max_backoff = Duration::from_secs(60); let mut backoff = Duration::from_secs(5); - // TODO(materialize#32030): fix retry logic + // TODO(materialize#38055): fix retry logic loop { match client.wait_for_initialization(config_sync_timeout).await { Some(true) => break, diff --git a/src/ore/src/iter.rs b/src/ore/src/iter.rs index 9ee9020741aed..9ac5c9f98926c 100644 --- a/src/ore/src/iter.rs +++ b/src/ore/src/iter.rs @@ -259,6 +259,7 @@ mod tests { #[crate::test] #[cfg(feature = "differential-dataflow")] + #[cfg_attr(miri, ignore)] // too slow fn test_consolidate_sorted() { proptest!(|(mut data in any::>())| { data.sort(); @@ -270,6 +271,7 @@ mod tests { #[crate::test] #[cfg(feature = "differential-dataflow")] + #[cfg_attr(miri, ignore)] // too slow fn test_consolidate() { proptest!(|(mut data in any::>())| { let streamed: Vec<_> = consolidate_iter(data.iter().copied()).collect(); @@ -279,6 +281,7 @@ mod tests { } #[crate::test] + #[cfg_attr(miri, ignore)] // too slow fn test_merge() { proptest!(|(mut data in vec(vec(0usize..100usize, 0..10), 0..10))| { let mut expected: Vec<_> = data.iter().flatten().copied().collect(); diff --git a/src/ore/src/pool.rs b/src/ore/src/pool.rs index 74185bf24cdda..23a1a5920e898 100644 --- a/src/ore/src/pool.rs +++ b/src/ore/src/pool.rs @@ -94,7 +94,9 @@ use crate::pool::region::{Region, SIZE_CLASSES}; /// (address space costs nothing, and touched pages are bounded by peak /// residency) so that no realistic budget, on any machine size, reaches the /// heap-fallback path. -const CLASS_CAPACITY_BYTES: usize = 1 << 40; +/// +/// NOTE: Seen OoMs with Miri since it actually allocates the capacity. +const CLASS_CAPACITY_BYTES: usize = if cfg!(miri) { 16 << 20 } else { 1 << 40 }; /// A chunk-provided transform between a chunk's body bytes and the stored /// bytes its extent holds. The pool owns scheduling: spill threads, the @@ -2987,6 +2989,7 @@ mod tests { /// the same victims (the pool's only two-chunk lock edge). Contents are /// asserted on every read. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_admits_exercise_the_steal_path() { const CHUNKS: u64 = 8; let pool = test_pool(usize::MAX); @@ -3069,6 +3072,7 @@ mod tests { /// contents stay correct, contended steals degrade to skips, and the /// accounting identity settles to zero. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_admits_race_cleanly() { let pool = test_pool(64 << 10); let per_thread = rounds(50, 3); @@ -3364,6 +3368,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn multithreaded_smoke() { // Budget of one small chunk: four inserting threads keep the pool // over budget, so every insert's enforcement pass selects victims @@ -3399,6 +3404,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_read_enforce_churn() { // Races the three actors that can touch one chunk's slot: readers // copying shared chunks out and verifying them, an enforcer evicting @@ -3480,6 +3486,7 @@ mod tests { /// drops) extent-queue entries, and pruning alone must keep the queue /// proportional to the live resident extents. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn extent_queue_stays_bounded_under_cap() { let pool = test_pool(256 << 20); pool.set_rss_target(1 << 40); @@ -3636,6 +3643,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn spill_threads_end_to_end() { let pool = test_pool(128 << 10); pool.set_spill_threads(2); diff --git a/src/repr/src/adt/datetime.rs b/src/repr/src/adt/datetime.rs index 64eb838384536..c1bf1a93b3078 100644 --- a/src/repr/src/adt/datetime.rs +++ b/src/repr/src/adt/datetime.rs @@ -3667,6 +3667,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_naive_time_sort_order() { let time = add_arb_duration(NaiveTime::from_hms_opt(0, 0, 0).unwrap()); let strat = proptest::collection::vec(time, 0..128); diff --git a/src/repr/src/adt/interval.rs b/src/repr/src/adt/interval.rs index a61d64b1253e5..eb5a1dbc32c97 100644 --- a/src/repr/src/adt/interval.rs +++ b/src/repr/src/adt/interval.rs @@ -1325,6 +1325,7 @@ mod test { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_interval_sorts() { fn sort_intervals(mut og: Vec) { let mut packed: Vec<_> = og.iter().copied().map(PackedInterval::from_value).collect(); diff --git a/src/repr/src/adt/timestamp.rs b/src/repr/src/adt/timestamp.rs index 019d657567df2..0a47c25c29a4a 100644 --- a/src/repr/src/adt/timestamp.rs +++ b/src/repr/src/adt/timestamp.rs @@ -1220,6 +1220,7 @@ mod test { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_naive_date_time_sort_order() { let strat = proptest::collection::vec(arb_naive_date_time(), 0..128); proptest!(|(mut times in strat)| { diff --git a/src/timely-util/src/columnar.rs b/src/timely-util/src/columnar.rs index e66c66b07198e..e8d9134ab03f1 100644 --- a/src/timely-util/src/columnar.rs +++ b/src/timely-util/src/columnar.rs @@ -373,6 +373,7 @@ mod tests { /// The ship signal is monotone: once it fires it stays fired, even when /// a single wide record steps far past the 2 MiB boundary in one push. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn ship_threshold_monotone() { use columnar::Push; let mut container = as Columnar>::Container::default(); diff --git a/src/timely-util/src/columnar/chunk.rs b/src/timely-util/src/columnar/chunk.rs index 002c0f012c072..4d693af116b86 100644 --- a/src/timely-util/src/columnar/chunk.rs +++ b/src/timely-util/src/columnar/chunk.rs @@ -1459,6 +1459,7 @@ mod tests { /// A body large enough to spill round-trips through the pool with resident /// metadata intact, and the batcher produces spilled sealed output. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn spill_round_trip() { set_spill_override(Some(test_pool())); @@ -1490,6 +1491,7 @@ mod tests { /// Merging spilled chains loads bodies call-scoped and consolidates /// correctly, and an untouched survivor keeps its spilled body. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn merge_spilled_chains() { set_spill_override(Some(test_pool())); @@ -1604,6 +1606,7 @@ mod tests { /// commit large enough to spill carries the depth into its spilled /// metadata (and thus into the pool hints). #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn settle_commits_at_accumulated_depth() { set_spill_override(Some(test_pool())); let big: Vec = (0..100_000u64).map(|i| ((i, 0), 0, 1i64)).collect(); @@ -1627,6 +1630,7 @@ mod tests { /// periodic ship window, so mid-window chunk sizes cannot make it grow /// past the target unbounded. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn settle_carry_commits_at_target() { // ~1.5 MiB per chunk: inside the dead zone of the periodic window // check (see `at_commit_size`). From b9097f8a3dc1837354fd06038b89866eb10056a1 Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Wed, 5 Aug 2026 15:05:41 +0200 Subject: [PATCH 4/6] adapter: optimizer and dyncfg cleanups for OCC read-then-write (#37919) ### Motivation Part 2 of 7 in a stack that moves `DELETE`, `UPDATE` and `INSERT ... SELECT` off the coordinator onto the session task, using optimistic concurrency control. Two independent cleanups that the later parts need, both of which stand on their own. The first is a live bug. Closes [SQL-588](https://linear.app/materializeinc/issue/SQL-588) ### Description **`ConfigSet` sync at catalog open.** `apply_updates` only mirrors `SystemVars` into the dyncfg `ConfigSet` when a durable `SystemConfiguration` update is present. A deployment that configures a parameter purely through `system_parameter_default` has no such update, so nothing syncs, and anything reading that dyncfg during bootstrap sees the compile-time default instead of the configured value. `ENABLE_EXPRESSION_CACHE` is read a few lines below and is affected today. Syncing the effective values once, after both the defaults and the durable configuration are in the map, fixes every such read. **One subscribe optimization pipeline.** Subscribe optimization had two entry points that differed only in whether the source was a query or an existing id, and each rebuilt the same pipeline around that one decision. Both rebuilt the same `SinkDesc` verbatim, down to the same three comments, and differed only in whether the source needed HIR lowering first. Taking the source as a `SubscribeSource` and the output shape as a parameter leaves the pipeline in one place. The view optimizer's expression-preparation variants collapse the same way, into one constructor plus a builder method for the constant folding limit. The dyncfg commit also splits `dyncfg_updates`, which computed the updates and applied them to its own `ConfigSet` in one method. The only two callers of the applying behavior are the ones that discard the result, so the call sites read like a getter whose value is thrown away, and a later change making the getter pure would compile cleanly while silently reintroducing the bug. `sync_dyncfgs` names that intent. The six callers that forward the updates to metrics or the controllers keep the pure `dyncfg_updates`. ### Verification Existing subscribe and `COPY FROM` coverage. No behavior change is intended for the optimizer part. The dyncfg change is observable as `ENABLE_EXPRESSION_CACHE` now taking effect when set only via `system_parameter_default`. --- src/adapter/src/catalog/apply.rs | 2 +- src/adapter/src/catalog/open.rs | 11 + .../src/coord/sequencer/inner/copy_from.rs | 8 +- src/adapter/src/optimize.rs | 5 +- src/adapter/src/optimize/subscribe.rs | 235 ++++++++++-------- src/adapter/src/optimize/view.rs | 27 +- src/sql/src/session/vars.rs | 16 ++ 7 files changed, 177 insertions(+), 127 deletions(-) diff --git a/src/adapter/src/catalog/apply.rs b/src/adapter/src/catalog/apply.rs index 0b70c7d72c847..8f964df74c402 100644 --- a/src/adapter/src/catalog/apply.rs +++ b/src/adapter/src/catalog/apply.rs @@ -281,7 +281,7 @@ impl CatalogState { } if update_system_config { - self.system_configuration.dyncfg_updates(); + self.system_configuration.sync_dyncfgs(); } Ok((builtin_table_updates, catalog_updates)) diff --git a/src/adapter/src/catalog/open.rs b/src/adapter/src/catalog/open.rs index c946b5ef5c87a..57e99d313dda8 100644 --- a/src/adapter/src/catalog/open.rs +++ b/src/adapter/src/catalog/open.rs @@ -316,6 +316,17 @@ impl Catalog { .await; builtin_table_updates.extend(builtin_table_update); + // Both the `system_parameter_defaults` set above and the durable + // `SystemConfiguration` values applied via `pre_item_updates` live in + // the `SystemVars` value map by now. Mirror their effective values into + // the dyncfg `ConfigSet`, so that startup-only reads observe configured + // values rather than compile-time defaults, `ENABLE_EXPRESSION_CACHE` + // just below being one of them. `apply_updates` only + // syncs the `ConfigSet` when a `SystemConfiguration` update is present, + // and a deployment configured purely via `system_parameter_default` has + // none, so sync explicitly here. + state.system_config().sync_dyncfgs(); + // Ensure mz_system has a password if configured to have one. // It's important we do this after the `pre_item_updates` so that // the mz_system role exists in the catalog. diff --git a/src/adapter/src/coord/sequencer/inner/copy_from.rs b/src/adapter/src/coord/sequencer/inner/copy_from.rs index e89212807a4cc..0c7c5a8d3e892 100644 --- a/src/adapter/src/coord/sequencer/inner/copy_from.rs +++ b/src/adapter/src/coord/sequencer/inner/copy_from.rs @@ -362,11 +362,9 @@ impl Coordinator { session: &session_meta, catalog_state, }; - let mut optimizer = optimize::view::Optimizer::new_with_prep_no_limit( - optimizer_config.clone(), - None, - prep, - ); + let mut optimizer = + optimize::view::Optimizer::new_with_prep(optimizer_config.clone(), None, prep) + .without_fold_constants_limit(); let hir = mz_sql::plan::plan_copy_from( &pcx, diff --git a/src/adapter/src/optimize.rs b/src/adapter/src/optimize.rs index 1c842d06c0336..6138c1c6feeb0 100644 --- a/src/adapter/src/optimize.rs +++ b/src/adapter/src/optimize.rs @@ -90,9 +90,8 @@ use crate::TimestampContext; /// A type for a [`DataflowDescription`] backed by `Mir~` plans. Used internally /// by the optimizer implementations. type MirDataflowDescription = DataflowDescription; -/// A type for a [`DataflowDescription`] backed by `Lir~` plans. Used internally -/// by the optimizer implementations. -type LirDataflowDescription = DataflowDescription; +/// A type for a [`DataflowDescription`] backed by `Lir~` plans. +pub type LirDataflowDescription = DataflowDescription; // Core API // -------- diff --git a/src/adapter/src/optimize/subscribe.rs b/src/adapter/src/optimize/subscribe.rs index 96f7839a253c4..19a082a1610ed 100644 --- a/src/adapter/src/optimize/subscribe.rs +++ b/src/adapter/src/optimize/subscribe.rs @@ -17,8 +17,9 @@ use differential_dataflow::lattice::Lattice; use mz_compute_types::ComputeInstanceId; use mz_compute_types::plan::LirRelationExpr; use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc, SubscribeSinkConnection}; +use mz_expr::{ColumnOrder, MirRelationExpr}; use mz_ore::soft_assert_or_log; -use mz_repr::{GlobalId, Timestamp}; +use mz_repr::{GlobalId, RelationDesc, Timestamp}; use mz_sql::optimizer_metrics::OptimizerMetrics; use mz_sql::plan::{HirToMirConfig, SubscribeFrom, SubscribePlan}; use mz_transform::TransformCtx; @@ -114,6 +115,118 @@ impl Optimizer { pub fn sink_id(&self) -> GlobalId { self.sink_id } + + /// The single subscribe optimization pipeline. Every subscribe, whatever it + /// reads from, goes through here, so a prep or metainfo step added here + /// applies to all of them. + fn optimize_inner( + &mut self, + source: SubscribeSource, + output: Vec, + ) -> Result, OptimizerError> { + let time = Instant::now(); + + let mut df_builder = { + let compute = self.compute_instance.clone(); + DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config) + }; + let mut df_desc = MirDataflowDescription::new(self.debug_name.clone()); + let mut df_meta = DataflowMetainfo::default(); + + let (from, from_desc) = match source { + SubscribeSource::Id(from_id) => { + let from_desc = self + .catalog + .get_entry(&from_id) + .relation_desc() + .expect("subscribes can only be run on items with descs") + .into_owned(); + + df_builder.import_into_dataflow(&from_id, &mut df_desc, &self.config.features)?; + + (from_id, from_desc) + } + SubscribeSource::Query { expr, from_desc } => { + // MIR ⇒ MIR optimization (local) + let mut transform_ctx = TransformCtx::local( + &self.config.features, + &self.typecheck_ctx, + &mut df_meta, + Some(&mut self.metrics), + Some(self.view_id), + ); + let expr = optimize_mir_local(expr, &mut transform_ctx)?; + + df_builder.import_view_into_dataflow( + &self.view_id, + &expr, + &mut df_desc, + &self.config.features, + )?; + + (self.view_id, from_desc) + } + }; + df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; + + // Make SinkDesc + let sink_description = ComputeSinkDesc { + from, + from_desc, + connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { output }), + with_snapshot: self.with_snapshot, + up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), + // No `FORCE NOT NULL` for subscribes + non_null_assertions: vec![], + // No `REFRESH` for subscribes + refresh_schedule: None, + }; + df_desc.export_sink(self.sink_id, sink_description); + + // Prepare expressions in the assembled dataflow. + let style = ExprPrepMaintained; + df_desc.visit_children( + |r| style.prep_relation_expr(r), + |s| style.prep_scalar_expr(s), + )?; + + // Construct TransformCtx for global optimization. + let mut transform_ctx = TransformCtx::global( + &df_builder, + &mz_transform::EmptyStatisticsOracle, // TODO: wire proper stats + &self.config.features, + &self.typecheck_ctx, + &mut df_meta, + Some(&mut self.metrics), + ); + // Run global optimization. + mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?; + + if self.config.mode == OptimizeMode::Explain { + // Collect the list of indexes used by the dataflow at this point. + trace_plan!(at: "global", &df_meta.used_indexes(&df_desc)); + } + + self.duration += time.elapsed(); + + // Return the (sealed) plan at the end of this optimization step. + Ok(GlobalMirPlan { + df_desc, + df_meta, + phantom: PhantomData::, + }) + } +} + +/// What a subscribe reads from, with any HIR ⇒ MIR lowering already done. +enum SubscribeSource { + /// An existing collection, imported by id. + Id(GlobalId), + /// A query, imported into the dataflow as a view under `view_id`. + Query { + expr: MirRelationExpr, + from_desc: RelationDesc, + }, } /// The (sealed intermediate) result after: @@ -171,122 +284,32 @@ impl Optimize for Optimizer { type To = GlobalMirPlan; fn optimize(&mut self, plan: SubscribePlan) -> Result { - let output = plan.output; - let plan = plan.from; - let time = Instant::now(); + let output = plan.output.row_order().to_vec(); - let mut df_builder = { - let compute = self.compute_instance.clone(); - DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config) - }; - let mut df_desc = MirDataflowDescription::new(self.debug_name.clone()); - let mut df_meta = DataflowMetainfo::default(); - - match plan { - SubscribeFrom::Id(from_id) => { - let from = self.catalog.get_entry(&from_id); - let from_desc = from - .relation_desc() - .expect("subscribes can only be run on items with descs") - .into_owned(); - - df_builder.import_into_dataflow(&from_id, &mut df_desc, &self.config.features)?; - df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; - - // Make SinkDesc - let sink_description = ComputeSinkDesc { - from: from_id, - from_desc, - connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { - output: output.row_order().to_vec(), - }), - with_snapshot: self.with_snapshot, - up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), - // No `FORCE NOT NULL` for subscribes - non_null_assertions: vec![], - // No `REFRESH` for subscribes - refresh_schedule: None, - }; - df_desc.export_sink(self.sink_id, sink_description); - } + match plan.from { + SubscribeFrom::Id(from_id) => self.optimize_inner(SubscribeSource::Id(from_id), output), SubscribeFrom::Query { expr, desc } => { // TODO: Change the `expr` type to be `HirRelationExpr` and run // HIR ⇒ MIR lowering and decorrelation here. This would allow // us implement something like `EXPLAIN RAW PLAN FOR SUBSCRIBE.` // // let typ = expr.top_level_typ(); - // let expr = expr.lower(&self.config)?; - - // MIR ⇒ MIR optimization (local) - let mut transform_ctx = TransformCtx::local( - &self.config.features, - &self.typecheck_ctx, - &mut df_meta, - Some(&mut self.metrics), - Some(self.view_id), - ); + // Lowering counts towards the reported optimization time, so we + // time it here and `optimize_inner` times the rest. + let time = Instant::now(); let expr = expr.lower(HirToMirConfig::from(&self.config), None)?; - let expr = optimize_mir_local(expr, &mut transform_ctx)?; - - df_builder.import_view_into_dataflow( - &self.view_id, - &expr, - &mut df_desc, - &self.config.features, - )?; - df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; - - // Make SinkDesc - let sink_description = ComputeSinkDesc { - from: self.view_id, - from_desc: desc.clone(), - connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { - output: output.row_order().to_vec(), - }), - with_snapshot: self.with_snapshot, - up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), - // No `FORCE NOT NULL` for subscribes - non_null_assertions: vec![], - // No `REFRESH` for subscribes - refresh_schedule: None, - }; - df_desc.export_sink(self.sink_id, sink_description); + self.duration += time.elapsed(); + + self.optimize_inner( + SubscribeSource::Query { + expr, + from_desc: desc, + }, + output, + ) } - }; - - // Prepare expressions in the assembled dataflow. - let style = ExprPrepMaintained; - df_desc.visit_children( - |r| style.prep_relation_expr(r), - |s| style.prep_scalar_expr(s), - )?; - - // Construct TransformCtx for global optimization. - let mut transform_ctx = TransformCtx::global( - &df_builder, - &mz_transform::EmptyStatisticsOracle, // TODO: wire proper stats - &self.config.features, - &self.typecheck_ctx, - &mut df_meta, - Some(&mut self.metrics), - ); - // Run global optimization. - mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?; - - if self.config.mode == OptimizeMode::Explain { - // Collect the list of indexes used by the dataflow at this point. - trace_plan!(at: "global", &df_meta.used_indexes(&df_desc)); } - - self.duration += time.elapsed(); - - // Return the (sealed) plan at the end of this optimization step. - Ok(GlobalMirPlan { - df_desc, - df_meta, - phantom: PhantomData::, - }) } } diff --git a/src/adapter/src/optimize/view.rs b/src/adapter/src/optimize/view.rs index 96476cfc44b83..16dcecb8bcc13 100644 --- a/src/adapter/src/optimize/view.rs +++ b/src/adapter/src/optimize/view.rs @@ -53,23 +53,19 @@ pub struct Optimizer { impl Optimizer { /// Creates an optimizer instance that does not perform any expression - /// preparation. Additionally, this instance calls constant folding with a size limit. + /// preparation. pub fn new(config: OptimizerConfig, metrics: Option) -> Self { - Self { - typecheck_ctx: empty_typechecking_context(), - config, - metrics, - expr_prep_style: ExprPrepNoop, - fold_constants_limit: true, - } + Self::new_with_prep(config, metrics, ExprPrepNoop) } } impl Optimizer { /// Creates an optimizer instance that takes an [`ExprPrep`] to handle - /// unmaterializable functions. Additionally, this instance calls constant - /// folding without a size limit. - pub fn new_with_prep_no_limit( + /// unmaterializable functions. + /// + /// Constant folding runs with the usual size limit, see + /// [`Self::without_fold_constants_limit`]. + pub fn new_with_prep( config: OptimizerConfig, metrics: Option, expr_prep_style: S, @@ -79,9 +75,16 @@ impl Optimizer { config, metrics, expr_prep_style, - fold_constants_limit: false, + fold_constants_limit: true, } } + + /// Folds constants of any size, instead of giving up above the configured + /// limit. + pub fn without_fold_constants_limit(mut self) -> Self { + self.fold_constants_limit = false; + self + } } impl Optimize for Optimizer { diff --git a/src/sql/src/session/vars.rs b/src/sql/src/session/vars.rs index 485168d6593d8..1cdc5dd702056 100644 --- a/src/sql/src/session/vars.rs +++ b/src/sql/src/session/vars.rs @@ -2053,6 +2053,10 @@ impl SystemVars { *self.expect_value(&SCRAM_ITERATIONS) } + /// Computes an update for every dyncfg from its configured value. + /// + /// This does not touch our own [`ConfigSet`]. Callers that need this + /// process's dyncfgs to reflect the result want [`Self::sync_dyncfgs`]. pub fn dyncfg_updates(&self) -> ConfigUpdates { let mut updates = ConfigUpdates::default(); for entry in self.dyncfgs.entries() { @@ -2080,6 +2084,18 @@ impl SystemVars { }; updates.add_dynamic(entry.name(), val); } + updates + } + + /// Applies the configured dyncfg values to our own [`ConfigSet`], and + /// returns them. + /// + /// Two callers own keeping this process's dyncfgs in step with the + /// catalog: catalog open, and every durable system-config change. Everyone + /// else only forwards the updates elsewhere and wants + /// [`Self::dyncfg_updates`] instead. + pub fn sync_dyncfgs(&self) -> ConfigUpdates { + let updates = self.dyncfg_updates(); updates.apply(&self.dyncfgs); updates } From ff684283953b7fea45ff17fc6888963b39625629 Mon Sep 17 00:00:00 2001 From: Aljoscha Krettek Date: Wed, 5 Aug 2026 15:05:41 +0200 Subject: [PATCH 5/6] adapter: make statement logging single-owner for frontend sequencing (#37920) **NOTE: This looks big, but a large chunk is moving tests around, in a separate commit.** ### Motivation Part 3 of 7 in a stack that moves `DELETE`, `UPDATE` and `INSERT ... SELECT` off the coordinator onto the session task, using optimistic concurrency control. This part is a restructuring, not a bugfix. A statement sequenced from the session task has its end-of-execution logging obligation handed between several places, and each of them can be the last one to hold it. Whether a log row appears, and what reason it carries, depends on which path the statement takes out of the session task. That is survivable while the only such statement is a peek. It stops being survivable in part 6. Read-then-write is a long-running operation, so part 6 bounds it with cancellation and `statement_timeout` in a `tokio::select!` that drops the operation's future when either fires. The obligation must not live inside that future. If it does, dropping the future runs the guard's `Drop`, which reports `aborted` with no message, and a user who just saw "canceling statement due to statement timeout" finds a log row recording neither the error nor its text. Hoisting the obligation into the frame beside the `select!` is what makes that drop safe, which is why this lands before the path that needs it rather than with it. Closes [SQL-589](https://linear.app/materializeinc/issue/SQL-589) ### Description `ExecutionLogging` is a slot that owns the obligation for one `SessionClient::execute` call, with a single retirement site that maps the call's outcome to an end reason. Handing the statement back to the coordinator passes the obligation as data in `Command::Execute`, and `handle_execute` arms it again on receipt, which is safe because the command travels a channel only the coordinator loop drains. Three defects follow from the old structure and go away with it: - an `EXECUTE` that unrolled to its inner statement got a guard whose retire channel was already dropped, so if that guard ever ran its `Drop` the end event went nowhere and the entry stayed open forever - a statement whose future is dropped mid-flight was recorded as `aborted` rather than carrying its error - a plan that passes the AST allowlist but then fails the plan-kind check bailed out after logging had begun, closing the frontend's entry as errored while the coordinator started its own from scratch, so one statement produced two log rows and counted `query_total` twice None of the three is reachable from SQL on the peek path, which is why none of them gets a test here. The first needs a drop in the few microseconds between installing the guard and the takeover, or a shutdown. The second needs an inner future for something to drop, and the peek path has none. The third sits behind `soft_panic_or_log`. They are latent, which is the point: the caller that reaches the second one arrives in part 6, which carries its tests. Frontend-sequenced subscribes also start counting towards `mz_subscribe_outputs`. The only site bumping it was `handle_execute`, which a statement the frontend took over never reaches, and `SessionMetrics` had no accessor for it, so those subscribes went uncounted. That is a behavior change beyond the restructuring. End events also move from `Client::send`, which panics when the coordinator is gone, to a `try_send` that drops them. A guard can outlive the coordinator during shutdown, and losing one end event is cheaper than panicking in `Drop` and taking the connection with it. Every other logging event keeps the asserting send. `end_statement_execution` takes the end timestamp from its caller. The frontend already recorded when execution finished and the handler discarded it, so `finished_at` and the `ExecutionFinished` lifecycle event charged the statement for however long its message waited in the coordinator's queue. The client never waited for that queue, because the response leaves the session task while the end event is a fire-and-forget send on an unbounded channel. `began_at` already came from the frontend's clock, so the two ends of a statement's duration were being read from different clocks. ### Verification `test_statement_logging_finished_at_excludes_coordinator_queue` covers the `finished_at` change, which is the one thing here observable from SQL today. It stalls the coordinator inside `Catalog::transact` and then runs a constant `SELECT` from a session whose catalog snapshot is already warm, so the statement needs nothing from the coordinator. The logged `finished_at` has to fall before the client saw the statement return. On the old code it lands about three seconds later. All existing `test_statement_logging_*` tests are unchanged and pass. The remaining tests for these paths are DML-specific and arrive with part 6, which adds a harness asserting the frontend and coordinator paths produce the same `mz_statement_execution_history` rows for the same statement, plus the cancellation and statement-timeout tests that reach the latent defects above. The second commit moves the existing tests out of `server.rs` into `statement_logging.rs`. It is a pure move: every moved item is byte-identical to its previous form, and the only other edits are the imports each file now needs. `server.rs` was 7585 lines, of which the statement logging tests and their harness were a contiguous 994-line block. --- src/adapter/src/client.rs | 241 ++-- src/adapter/src/command.rs | 6 +- src/adapter/src/coord.rs | 3 +- src/adapter/src/coord/command_handler.rs | 11 +- src/adapter/src/coord/sequencer.rs | 8 +- src/adapter/src/coord/statement_logging.rs | 12 +- src/adapter/src/frontend_peek.rs | 158 +-- src/adapter/src/metrics.rs | 6 + src/adapter/src/peek_client.rs | 351 ++++-- src/environmentd/tests/server.rs | 1002 +---------------- src/environmentd/tests/statement_logging.rs | 1091 +++++++++++++++++++ 11 files changed, 1547 insertions(+), 1342 deletions(-) create mode 100644 src/environmentd/tests/statement_logging.rs diff --git a/src/adapter/src/client.rs b/src/adapter/src/client.rs index 705b6056195ed..b6c6298969c47 100644 --- a/src/adapter/src/client.rs +++ b/src/adapter/src/client.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use anyhow::bail; use chrono::{DateTime, Utc}; use derivative::Derivative; -use futures::{Stream, StreamExt}; +use futures::{FutureExt, Stream, StreamExt}; use itertools::Itertools; use mz_adapter_types::connection::{ConnectionId, ConnectionIdType}; use mz_auth::password::Password; @@ -60,7 +60,8 @@ use crate::command::{ use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend}; use crate::coord::{Coordinator, ExecuteContextGuard}; use crate::error::AdapterError; -use crate::metrics::{self, Metrics}; +use crate::metrics::Metrics; +use crate::peek_client::{ExecutionLogging, TakeOver}; use crate::session::{ EndTransactionAction, PreparedStatement, Session, SessionConfig, StateRevision, TransactionId, }; @@ -609,10 +610,15 @@ Issue a SQL query to get started. Need help? } #[instrument(level = "debug")] - pub(crate) fn send(&self, cmd: Command) { + pub(crate) fn try_send(&self, cmd: Command) -> bool { self.inner_cmd_tx .send((OpenTelemetryContext::obtain(), cmd)) - .expect("coordinator unexpectedly gone"); + .is_ok() + } + + #[instrument(level = "debug")] + pub(crate) fn send(&self, cmd: Command) { + assert!(self.try_send(cmd), "coordinator unexpectedly gone"); } } @@ -788,80 +794,94 @@ impl SessionClient { outer_ctx_extra: Option, ) -> Result<(ExecuteResponse, Instant), AdapterError> { let execute_started = Instant::now(); + let cancel_future = cancel_future.map(|_| ()).shared(); + + // Owning the end-of-execution obligation in this frame is what lets + // cancellation report an error: the inner future can be dropped without + // the obligation going with it. See `ExecutionLogging`. + let mut logging = ExecutionLogging::adopt(outer_ctx_extra, &self.peek_client); - let mut outer_ctx_extra = outer_ctx_extra; + let result = self + .execute_attempts(portal_name, &mut logging, cancel_future) + .await; + + logging.retire(&result); + + result.map(|response| (response, execute_started)) + } + /// Runs the execution paths in order of preference: frontend peek + /// sequencing, frontend read-then-write sequencing, and finally the + /// coordinator via `Command::Execute`. + async fn execute_attempts( + &mut self, + portal_name: String, + logging: &mut ExecutionLogging, + cancel_future: impl Future + Send + Clone, + ) -> Result { // Unroll SQL `EXECUTE (...)` so the inner statement - // flows through `try_frontend_peek` below, rather than being + // flows through `try_frontend_peek` / + // `try_frontend_read_then_write` below, rather than being // re-dispatched via `Command::Execute` from the coordinator's // `Plan::Execute` handler. Without this, a prepared statement // would route differently from the same statement issued // directly. // - // On a successful unroll, `unroll_sql_execute` also begins - // EXECUTE-level statement logging on the outer portal, so that + // On a successful unroll, `unroll_sql_execute` also takes over + // statement logging on the outer portal, so that // `mz_statement_execution_history` records `EXECUTE foo (...)` - // rather than the inner SQL, and installs the resulting - // `ExecuteContextGuard` into `outer_ctx_extra`. - let portal_name = self - .unroll_sql_execute(portal_name, &mut outer_ctx_extra) - .await?; + // rather than the inner SQL. + let portal_name = self.unroll_sql_execute(portal_name, logging).await?; // Attempt peek sequencing in the session task. // If unsupported, fall back to the Coordinator path. // TODO(peek-seq): wire up cancel_future - let peek_result = self - .try_frontend_peek(&portal_name, &mut outer_ctx_extra) - .await?; + let peek_result = self.try_frontend_peek(&portal_name, logging).await?; if let Some(resp) = peek_result { debug!("frontend peek succeeded"); - // Frontend peek handled the execution and retired outer_ctx_extra if it existed. - // No additional work needed here. - return Ok((resp, execute_started)); - } else { - debug!("frontend peek did not happen, falling back to `Command::Execute`"); - // If we bailed out, outer_ctx_extra is still present (if it was originally). - // `Command::Execute` will handle it. - // (This is not true if we bailed out _after_ the frontend peek sequencing has already - // begun its own statement logging. That case would be a bug.) + return Ok(resp); } - - let response = self - .send_with_cancel( - |tx, session| Command::Execute { - portal_name, - session, - tx, - outer_ctx_extra, - }, - cancel_future, - ) - .await?; - Ok((response, execute_started)) + debug!("frontend peek did not happen, falling back to `Command::Execute`"); + + // No frontend path took the statement over, so the coordinator retires + // whatever entry we hold, or begins its own if we hold none. + let outer_ctx_extra = logging.release(); + self.send_with_cancel( + |tx, session| Command::Execute { + portal_name, + session, + tx, + outer_ctx_extra, + }, + cancel_future, + ) + .await } /// If the named portal binds a SQL `EXECUTE `, resolve the /// prepared statement, install a fresh portal for the inner statement /// (carrying the EXECUTE's actual parameter values), and return that - /// portal's name so the caller can run `try_frontend_peek` against it. + /// portal's name so the caller can run `try_frontend_peek` / + /// `try_frontend_read_then_write` against it. /// /// Only ever unrolls one level: the parser rejects /// `PREPARE foo AS EXECUTE bar` (matching Postgres), so the inner /// statement is guaranteed not to be another `EXECUTE`. A failsafe below /// surfaces an internal error if that invariant is ever violated. /// - /// When the portal does not bind an `EXECUTE` — the common case — - /// returns the original portal name, costing only a portal lookup. + /// When the portal does not bind an `EXECUTE`, the common case, returns the + /// original portal name, costing only a portal lookup. async fn unroll_sql_execute( &mut self, portal_name: String, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result { let (stmt, params, outer_logging, outer_lifecycle_timestamps) = { let session = self.session.as_ref().expect("SessionClient invariant"); let portal = match session.get_portal_unverified(&portal_name) { Some(p) => p, - // No portal: let `try_frontend_peek` surface the + // No portal: let `try_frontend_peek` / + // `try_frontend_read_then_write` surface the // standard "missing portal" error. None => return Ok(portal_name), }; @@ -892,107 +912,31 @@ impl SessionClient { Coordinator::verify_portal(&catalog, session, &portal_name)?; } - // Bump query_total for the outer EXECUTE itself. The inner - // statement gets its own increment inside `try_frontend_peek_inner` - // (or, on bailout, in the coordinator's `handle_execute`). - { - let session = self.session.as_ref().expect("SessionClient invariant"); - session - .metrics() - .query_total(&[ - metrics::session_type_label_value(session.user()), - metrics::statement_type_label_value(&stmt), - ]) - .inc(); - } - - // Begin EXECUTE-level statement logging up front, so that planning - // errors below produce an `Errored` end-event in - // `mz_statement_execution_history` rather than no entry at all. - // - // We pass the *outer* portal's `logging` and pgwire-bound `params` - // so the recorded entry shows the user-visible `EXECUTE foo (...)`, - // not the inner SQL. The id (if any) moves into `outer_ctx_extra` - // below for `try_frontend_peek` to retire; on planning error we - // explicitly emit an `Errored` end-event below. - let began_outer_logging = outer_ctx_extra.is_none(); - let logging_id: Option = - if began_outer_logging { - let session = self.session.as_mut().expect("SessionClient invariant"); - let result = self - .peek_client - .statement_logging_frontend - .begin_statement_execution( - session, - ¶ms, - &outer_logging, - catalog.system_config(), - outer_lifecycle_timestamps, - ); - if let Some((id, began_execution, mseh_update, prepared_statement)) = result { - self.peek_client.log_began_execution( - began_execution, - mseh_update, - prepared_statement, - ); - Some(id) - } else { - None - } - } else { - None - }; - - let new_portal_name = match self.install_inner_portal_for_execute(&catalog, &stmt, ¶ms) - { - Ok(name) => name, - Err(err) => { - if let Some(id) = logging_id { - self.peek_client.log_ended_execution( - id, - StatementEndedExecutionReason::Errored { - error: err.to_string(), - }, - ); - } - return Err(err); - } - }; - - // Hand off to `outer_ctx_extra` whenever we entered the begin path - // for the outer EXECUTE — even if `begin_statement_execution` - // returned `None` (sampling decided not to sample, or logging is - // disabled for the user). This mirrors the original coord path, - // which always installs a guard via - // `ExecuteContextGuard::new(maybe_uuid, ...)`. Without this, the - // inner portal would be treated as a fresh statement by - // `try_frontend_peek` (or the fallback `Command::Execute` path) - // and re-account its bytes against - // `mz_statement_logging_unsampled_bytes`, double-counting the - // inner SQL. - if began_outer_logging { - // Soft invariant: `try_frontend_peek` takes ownership of - // `outer_ctx_extra` immediately, so this guard's `Drop` is - // unreachable on the normal flow and the dummy channel is - // never used. If a panic does fire `Drop` between here and - // that takeover, the `Aborted` end-event is silently lost - // — an acceptable trade given the panic implies the - // connection is going down anyway. - let (dummy_tx, _dummy_rx) = mpsc::unbounded_channel(); - *outer_ctx_extra = Some(ExecuteContextGuard::new(logging_id, dummy_tx)); - } + // Take over the EXECUTE itself, so that a planning error below produces + // an errored end-event in `mz_statement_execution_history` rather than + // no entry at all. We pass the *outer* portal's `logging` and + // pgwire-bound `params` so the recorded entry shows the user-visible + // `EXECUTE foo (...)`, not the inner SQL. The inner statement gets its + // own `query_total` increment when a frontend path takes it over, or, + // on bailout, in the coordinator's `handle_execute`. + logging.take_over( + &self.peek_client, + self.session.as_mut().expect("SessionClient invariant"), + Some(&stmt), + ¶ms, + &outer_logging, + &catalog, + outer_lifecycle_timestamps, + TakeOver::UnrolledExecute, + ); - Ok(new_portal_name) + self.install_inner_portal_for_execute(&catalog, &stmt, ¶ms) } /// Helper for [`Self::unroll_sql_execute`]: plans the outer /// `Statement::Execute`, verifies the referenced prepared statement, and /// installs a fresh portal carrying the inner statement plus the /// EXECUTE's bound parameter values. Returns the new portal's name. - /// - /// Split out so [`Self::unroll_sql_execute`] can wrap the fallible work - /// in a single error-handling site that emits an `Errored` end-event - /// for the EXECUTE-level statement-logging entry. fn install_inner_portal_for_execute( &mut self, catalog: &Arc, @@ -1028,11 +972,11 @@ impl SessionClient { } }; - // Verify and install the inner portal. Mirrors - // `Coordinator::sequence_execute`. The new portal carries the inner - // prepared statement's `logging`, but `try_frontend_peek` will see - // `outer_ctx_extra=Some(...)` and inherit the EXECUTE-level logging - // instead of starting fresh from this portal. + // Verify and install the inner portal. The new portal carries the + // inner prepared statement's `logging`, but the logging slot already + // holds the EXECUTE-level entry, so a frontend path taking the inner + // statement over inherits that entry instead of starting fresh from + // this portal. let session = self.session.as_mut().expect("SessionClient invariant"); Coordinator::verify_prepared_statement(catalog, session, &execute_plan.name)?; let ps = session @@ -1046,7 +990,7 @@ impl SessionClient { // Failsafe: `PREPARE foo AS EXECUTE bar` is rejected by the parser, // so the resolved inner statement must not be another `EXECUTE`. If // that ever changes, we'd silently skip frontend sequencing for the - // deeper EXECUTEs — surface it as an internal error instead. + // deeper EXECUTEs. Surface it as an internal error instead. if let Some(inner) = inner_stmt.as_ref() { if matches!(inner, Statement::Execute(_)) { return Err(AdapterError::Internal(format!( @@ -1325,7 +1269,7 @@ impl SessionClient { async fn send_with_cancel( &mut self, f: F, - cancel_future: impl Future + Send, + cancel_future: impl Future + Send, ) -> Result where F: FnOnce(oneshot::Sender>, Session) -> Command, @@ -1430,7 +1374,7 @@ impl SessionClient { *client_session = Some(res.session); return res.result; }, - _err = &mut cancel_future, if !cancelled => { + _ = &mut cancel_future, if !cancelled => { cancelled = true; inner_client.send(Command::PrivilegedCancelRequest { conn_id: conn_id.clone(), @@ -1476,18 +1420,15 @@ impl SessionClient { /// /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` to fall back to the /// Coordinator's sequencing. If it returns an error, it should be returned to the user. - /// - /// `outer_ctx_extra` is Some when we are executing as part of an outer statement, e.g., a FETCH - /// triggering the execution of the underlying query. pub(crate) async fn try_frontend_peek( &mut self, portal_name: &str, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { if self.enable_frontend_peek_sequencing { let session = self.session.as_mut().expect("SessionClient invariant"); self.peek_client - .try_frontend_peek(portal_name, session, outer_ctx_extra) + .try_frontend_peek(portal_name, session, logging) .await } else { Ok(None) diff --git a/src/adapter/src/command.rs b/src/adapter/src/command.rs index 3174ff359e99f..1665e1035491c 100644 --- a/src/adapter/src/command.rs +++ b/src/adapter/src/command.rs @@ -130,7 +130,11 @@ pub enum Command { portal_name: String, session: Session, tx: oneshot::Sender>, - outer_ctx_extra: Option, + /// The end-of-execution obligation of the statement this execution + /// serves, if any. `Coordinator::handle_execute` arms it again on + /// receipt. `None` means no log entry exists yet, so the coordinator + /// begins one. + outer_ctx_extra: Option, }, /// Attempts to commit or abort the session's transaction. Guarantees that the Coordinator's diff --git a/src/adapter/src/coord.rs b/src/adapter/src/coord.rs index 0796e4b4b7b56..6356d9d3ed0b9 100644 --- a/src/adapter/src/coord.rs +++ b/src/adapter/src/coord.rs @@ -4279,7 +4279,8 @@ impl Coordinator { ctx_extra: ExecuteContextExtra, ) { if let Some(uuid) = ctx_extra.retire() { - self.end_statement_execution(uuid, reason); + let ended_at = self.now(); + self.end_statement_execution(uuid, reason, ended_at); } } diff --git a/src/adapter/src/coord/command_handler.rs b/src/adapter/src/coord/command_handler.rs index 41bbde3c9b7bf..45f582cf42cee 100644 --- a/src/adapter/src/coord/command_handler.rs +++ b/src/adapter/src/coord/command_handler.rs @@ -87,7 +87,7 @@ use crate::webhook::{ }; use crate::{AppendWebhookError, ExecuteContext, catalog, metrics}; -use super::ExecuteContextGuard; +use super::{ExecuteContextExtra, ExecuteContextGuard}; /// The login status of a role, used by authentication handlers to check role /// existence and login permission before proceeding to credential verification. @@ -1107,8 +1107,15 @@ impl Coordinator { // then `outer_context` should be `Some`. // This instructs the coordinator that the // outer execute should be considered finished once the inner one is. - outer_context: Option, + outer_context: Option, ) { + // Arm the obligation before anything can return: from here on, every + // exit either retires it or hands it to a guard whose `Drop` does. + // Nothing could have dropped it on the way here either, the command + // travels a channel that only this loop drains. + let outer_context = outer_context + .map(|extra| ExecuteContextGuard::new(extra.retire(), self.internal_cmd_tx.clone())); + // A new statement is starting, so discard any cancellation that was signaled while no // statement was running. Such a cancellation targeted an earlier statement and must not // cancel the new one. (Like in PostgreSQL, a cancel request that arrives when nothing is diff --git a/src/adapter/src/coord/sequencer.rs b/src/adapter/src/coord/sequencer.rs index 1c1f19303d774..c2c22fc996c83 100644 --- a/src/adapter/src/coord/sequencer.rs +++ b/src/adapter/src/coord/sequencer.rs @@ -624,13 +624,19 @@ impl Coordinator { match self.sequence_execute(ctx.session_mut(), plan) { Ok(portal_name) => { let (tx, _, session, extra, response_barriers) = ctx.into_parts(); + // The obligation travels as data and + // `handle_execute` arms it again. It cannot be + // dropped in between: the command goes to the + // channel this loop drains, and a spawned barrier + // task is only dropped when the process is going + // down, at which point nothing records anything. let command = Message::Command( OpenTelemetryContext::obtain(), Command::Execute { portal_name, session, tx: tx.take(), - outer_ctx_extra: Some(extra), + outer_ctx_extra: Some(extra.defuse()), }, ); if response_barriers.is_empty() { diff --git a/src/adapter/src/coord/statement_logging.rs b/src/adapter/src/coord/statement_logging.rs index 3cf13bbfaa175..0b1c35993237c 100644 --- a/src/adapter/src/coord/statement_logging.rs +++ b/src/adapter/src/coord/statement_logging.rs @@ -170,6 +170,7 @@ impl Coordinator { self.end_statement_execution( StatementLoggingId(ended_record.id), ended_record.reason, + ended_record.ended_at, ); } FrontendStatementLoggingEvent::SetCluster { @@ -394,17 +395,22 @@ impl Coordinator { /// Record the end of statement execution for a statement whose beginning /// was logged. Ends are idempotent: the first end wins and later ends for /// the same statement are ignored. + /// + /// `ended_at` is when execution finished, which is not the same as when this + /// runs. A statement sequenced off the coordinator loop finishes in its + /// session task and reports the end as a message, so taking the clock here + /// would charge it for however long that message sat in the queue. pub(crate) fn end_statement_execution( &mut self, id: StatementLoggingId, reason: StatementEndedExecutionReason, + ended_at: EpochMillis, ) { let StatementLoggingId(uuid) = id; - let now = self.now(); let ended_record = StatementEndedExecutionRecord { id: uuid, reason, - ended_at: now, + ended_at, }; let Some(began_record) = self.statement_logging.executions_begun.remove(&uuid) else { @@ -439,7 +445,7 @@ impl Coordinator { self.record_statement_lifecycle_event( &id, &StatementLifecycleEvent::ExecutionFinished, - now, + ended_at, ); } diff --git a/src/adapter/src/frontend_peek.rs b/src/adapter/src/frontend_peek.rs index 016365b994cc6..fc05dab424522 100644 --- a/src/adapter/src/frontend_peek.rs +++ b/src/adapter/src/frontend_peek.rs @@ -45,27 +45,24 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::catalog::Catalog; use crate::command::Command; +use crate::coord; use crate::coord::peek::{FastPathPlan, PeekPlan}; use crate::coord::sequencer::{eval_copy_to_uri, statistics_oracle}; use crate::coord::timeline::timedomain_for; use crate::coord::timestamp_selection::TimestampDetermination; -use crate::coord::{ - Coordinator, CopyToContext, ExecuteContextGuard, ExplainContext, ExplainPlanContext, - TargetCluster, -}; +use crate::coord::{Coordinator, CopyToContext, ExplainContext, ExplainPlanContext, TargetCluster}; use crate::explain::insights::PlanInsightsContext; use crate::explain::optimizer_trace::OptimizerTrace; use crate::optimize::Optimize; use crate::optimize::dataflows::{ComputeInstanceSnapshot, DataflowBuilder}; -use crate::peek_client::StatementLoggingGuard; +use crate::peek_client::{ExecutionLogging, TakeOver}; use crate::session::{Session, TransactionOps, TransactionStatus}; +use crate::statement_logging::StatementLifecycleEvent; use crate::statement_logging::WatchSetCreation; -use crate::statement_logging::{StatementEndedExecutionReason, StatementLifecycleEvent}; use crate::{ AdapterError, AdapterNotice, CollectionIdBundle, ExecuteResponse, PeekClient, ReadHolds, TimelineContext, TimestampContext, TimestampProvider, optimize, }; -use crate::{coord, metrics}; impl PeekClient { /// Attempt to sequence a peek from the session task. @@ -73,13 +70,14 @@ impl PeekClient { /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` to fall back to the /// Coordinator's sequencing. If it returns an error, it should be returned to the user. /// - /// `outer_ctx_extra` is Some when we are executing as part of an outer statement, e.g., a FETCH - /// triggering the execution of the underlying query. + /// `logging` holds the end-of-execution obligation for this statement. The + /// caller retires it, this function only takes the statement over and, at + /// the dispatch sites that hand execution off, defuses the slot. pub(crate) async fn try_frontend_peek( &mut self, portal_name: &str, session: &mut Session, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { // # From handle_execute @@ -98,26 +96,12 @@ impl PeekClient { let catalog = self.catalog_snapshot("try_frontend_peek").await; - // Extract things from the portal. - let (stmt, params, logging, lifecycle_timestamps) = { - if let Err(err) = Coordinator::verify_portal(&*catalog, session, portal_name) { - // An inherited outer statement (e.g. EXECUTE) ends here and we - // own its end, so we log it. If we discarded the id instead, - // the statement would stay "running" forever in - // mz_statement_execution_history. - if let Some(id) = outer_ctx_extra - .take() - .and_then(|guard| guard.defuse().retire()) - { - self.log_ended_execution( - id, - StatementEndedExecutionReason::Errored { - error: err.to_string(), - }, - ); - } - return Err(err); - } + // Extract things from the portal. A failed verification does not begin + // an entry, mirroring the coordinator: the portal is what statement + // logging draws its record from. An adopted outer entry still gets its + // `Errored` end from the retirement site in `SessionClient::execute`. + let (stmt, params, logging_info, lifecycle_timestamps) = { + Coordinator::verify_portal(&*catalog, session, portal_name)?; let portal = session .get_portal_unverified(portal_name) // The portal is a session-level thing, so it couldn't have concurrently disappeared @@ -125,9 +109,9 @@ impl PeekClient { .expect("called verify_portal above"); let params = portal.parameters.clone(); let stmt = portal.stmt.clone(); - let logging = Arc::clone(&portal.logging); + let logging_info = Arc::clone(&portal.logging); let lifecycle_timestamps = portal.lifecycle_timestamps.clone(); - (stmt, params, logging, lifecycle_timestamps) + (stmt, params, logging_info, lifecycle_timestamps) }; // Before planning, check if this is a statement type we can handle. @@ -201,70 +185,39 @@ impl PeekClient { } } - // Set up statement logging, and log the beginning of execution. - // (But only if we're not executing in the context of another statement.) - let mut logging_guard = self.begin_statement_logging( + // Past this point the coordinator never sees this statement, so every + // exit has to produce an outcome for it. The `Ok(None)` bailouts are all + // above. + logging.take_over( + self, session, + stmt.as_deref(), ¶ms, - &logging, + &logging_info, &catalog, lifecycle_timestamps, - outer_ctx_extra, + TakeOver::StatementToRun, ); - let result = self - .try_frontend_peek_inner(session, catalog, stmt, params, &mut logging_guard) - .await; - - // If we still own end-of-execution logging, retire it with the - // execution's outcome. We don't own it when a dispatch site in - // `try_frontend_peek_inner` handed it off: for streaming responses the - // end is logged asynchronously, by the coordinator for registered - // peeks and by the protocol layer for subscribes. - if logging_guard.id().is_some() { - let reason = match &result { - // Bailout case, which should not happen. - Ok(None) => { - soft_panic_or_log!( - "Bailed out from `try_frontend_peek_inner` after we already logged the beginning of statement execution." - ); - // The old peek sequencing would start its own statement - // logging from scratch; close out this one as errored. - StatementEndedExecutionReason::Errored { - error: "Internal error: bailed out from `try_frontend_peek_inner`" - .to_string(), - } - } - // Streaming responses cannot reach this arm: their dispatch - // sites hand off the guard. The `From` impl panics on them. - Ok(Some(resp)) => resp.into(), - Err(e) => StatementEndedExecutionReason::Errored { - error: e.to_string(), - }, - }; - - logging_guard.retire(reason); - } - - result + self.try_frontend_peek_inner(session, catalog, stmt, params, logging) + .await } /// This is encapsulated in an inner function so that the outer function can still do statement /// logging after the `?` returns of the inner function. /// - /// `logging_guard` owns end-of-execution logging for this statement. /// Dispatch sites that hand the statement to the coordinator for /// asynchronous completion (registered peeks, subscribes) `defuse` the - /// guard at the point where the coordinator takes over. Everywhere else - /// the guard stays armed and the caller logs the end from the returned - /// result. + /// logging slot at the point where the coordinator takes over. Everywhere + /// else the slot stays armed and `SessionClient::execute` logs the end from + /// the returned result. async fn try_frontend_peek_inner( &mut self, session: &mut Session, catalog: Arc, stmt: Option>>, params: Params, - logging_guard: &mut StatementLoggingGuard, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { let stmt = match stmt { Some(stmt) => stmt, @@ -274,14 +227,6 @@ impl PeekClient { } }; - session - .metrics() - .query_total(&[ - metrics::session_type_label_value(session.user()), - metrics::statement_type_label_value(&stmt), - ]) - .inc(); - // # From handle_execute_inner let conn_catalog = catalog.for_session(session); @@ -377,15 +322,17 @@ impl PeekClient { } _ => { // This shouldn't happen because we already checked for this at the AST - // level before calling `try_frontend_peek_inner`. + // level before calling `try_frontend_peek_inner`. The logging takeover + // has already happened, so falling back to the coordinator would count + // the statement twice. Report the inconsistency instead. soft_panic_or_log!( "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing: {:?}", explainee ); - debug!( - "Bailing out from try_frontend_peek_inner, because EXPLAIN FILTER PUSHDOWN is not for a SELECT query or is EXPLAIN BROKEN" - ); - return Ok(None); + return Err(AdapterError::Internal( + "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing" + .into(), + )); } } } @@ -445,15 +392,16 @@ impl PeekClient { Plan::Subscribe(subscribe) => (QueryPlan::Subscribe(subscribe), ExplainContext::None), _ => { // This shouldn't happen because we already checked for this at the AST - // level before calling `try_frontend_peek_inner`. + // level before calling `try_frontend_peek_inner`. The logging takeover has + // already happened, so falling back to the coordinator would count the + // statement twice. Report the inconsistency instead. soft_panic_or_log!( "Unexpected plan kind in frontend peek sequencing: {:?}", plan ); - debug!( - "Bailing out from try_frontend_peek_inner, because the Plan is not a SELECT, side-effecting SELECT, EXPLAIN SELECT, EXPLAIN FILTER PUSHDOWN, or COPY TO S3" - ); - return Ok(None); + return Err(AdapterError::Internal( + "unexpected plan kind in frontend peek sequencing".into(), + )); } }; @@ -496,7 +444,7 @@ impl PeekClient { }; // Log cluster selection - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_set_cluster(logging_id, target_cluster_id, target_cluster_name.clone()); } @@ -1185,7 +1133,7 @@ impl PeekClient { }; // Log optimization finished - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished); } @@ -1293,7 +1241,7 @@ impl PeekClient { // # Now back to peek_finish - let watch_set = logging_guard.id().map(|logging_id| { + let watch_set = logging.id().map(|logging_id| { WatchSetCreation::new( logging_id, catalog.state(), @@ -1314,7 +1262,7 @@ impl PeekClient { let response = match peek_plan { PeekPlan::FastPath(fast_path_plan) => { - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { // TODO(peek-seq): Actually, we should log it also for // FastPathPlan::Constant. The only reason we are not doing so at the // moment is to match the old peek sequencing, so that statement logging @@ -1362,12 +1310,12 @@ impl PeekClient { session.conn_id().clone(), source_ids, watch_set, - logging_guard, + logging, ) .await? } PeekPlan::SlowPath(dataflow_plan) => { - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_set_transient_index_id(logging_id, dataflow_plan.id); } @@ -1392,7 +1340,7 @@ impl PeekClient { // coordinator logs nothing (see // `implement_slow_path_peek`), so the guard stays armed // and the caller logs the error. - logging_guard.defuse(); + logging.defuse(); response } }; @@ -1454,7 +1402,7 @@ impl PeekClient { session_uuid: session.uuid(), read_holds, plan: subscribe_plan, - statement_logging_id: logging_guard.id(), + statement_logging_id: logging.id(), tx, }) .await?; @@ -1464,7 +1412,7 @@ impl PeekClient { // coordinator logs nothing (see the `ExecuteSubscribe` // handler), so the guard stays armed and the caller logs the // error. - logging_guard.defuse(); + logging.defuse(); Ok(Some(response)) } Execution::CopyToS3 { @@ -1511,7 +1459,7 @@ impl PeekClient { .await?; // Preflight succeeded, now execute the actual COPY TO dataflow - let watch_set = logging_guard.id().map(|logging_id| { + let watch_set = logging.id().map(|logging_id| { WatchSetCreation::new( logging_id, catalog.state(), diff --git a/src/adapter/src/metrics.rs b/src/adapter/src/metrics.rs index 30abf69c635e2..b5168a4916b98 100644 --- a/src/adapter/src/metrics.rs +++ b/src/adapter/src/metrics.rs @@ -308,6 +308,7 @@ impl Metrics { row_set_finishing_seconds: self.row_set_finishing_seconds(), session_startup_table_writes_seconds: self.session_startup_table_writes_seconds.clone(), query_total: self.query_total.clone(), + subscribe_outputs: self.subscribe_outputs.clone(), determine_timestamp: self.determine_timestamp.clone(), timestamp_difference_for_strict_serializable_ms: self .timestamp_difference_for_strict_serializable_ms @@ -329,6 +330,7 @@ pub struct SessionMetrics { row_set_finishing_seconds: Histogram, session_startup_table_writes_seconds: Histogram, query_total: IntCounterVec, + subscribe_outputs: IntCounterVec, determine_timestamp: IntCounterVec, timestamp_difference_for_strict_serializable_ms: HistogramVec, timestamp_difference_for_bounded_staleness_ms: HistogramVec, @@ -351,6 +353,10 @@ impl SessionMetrics { self.query_total.with_label_values(label_values) } + pub(crate) fn subscribe_outputs(&self, label_values: &[&str]) -> GenericCounter { + self.subscribe_outputs.with_label_values(label_values) + } + pub(crate) fn determine_timestamp(&self, label_values: &[&str]) -> GenericCounter { self.determine_timestamp.with_label_values(label_values) } diff --git a/src/adapter/src/peek_client.rs b/src/adapter/src/peek_client.rs index d489ca9a4eb09..62321a9c5d469 100644 --- a/src/adapter/src/peek_client.rs +++ b/src/adapter/src/peek_client.rs @@ -18,13 +18,17 @@ use mz_compute_client::protocol::command::PeekTarget; use mz_compute_types::ComputeInstanceId; use mz_expr::row::RowCollection; use mz_ore::cast::CastFrom; +use mz_ore::soft_panic_or_log; use mz_persist_client::PersistClient; use mz_repr::GlobalId; use mz_repr::Timestamp; use mz_repr::global_id::TransientIdGen; use mz_repr::{RelationDesc, Row}; +use mz_sql::ast::{Raw, Statement}; use mz_sql::optimizer_metrics::OptimizerMetrics; use mz_sql::plan::Params; +use mz_sql::session::metadata::SessionMetadata; +use mz_sql_parser::ast::{CopyRelation, CopyStatement, SubscribeStatement}; use mz_storage_types::sources::Timeline; use mz_timestamp_oracle::TimestampOracle; use prometheus::Histogram; @@ -35,15 +39,15 @@ use tokio::sync::oneshot; use uuid::Uuid; use crate::catalog::Catalog; -use crate::command::{CatalogSnapshot, Command}; +use crate::command::{CatalogSnapshot, Command, ExecuteResponse}; use crate::coord::peek::FastPathPlan; -use crate::coord::{Coordinator, ExecuteContextGuard}; +use crate::coord::{Coordinator, ExecuteContextExtra, ExecuteContextGuard}; use crate::session::{LifecycleTimestamps, Session}; use crate::statement_logging::{ FrontendStatementLoggingEvent, PreparedStatementEvent, PreparedStatementLoggingInfo, StatementLoggingFrontend, StatementLoggingId, WatchSetCreation, }; -use crate::{AdapterError, Client, CollectionIdBundle, ReadHolds, statement_logging}; +use crate::{AdapterError, Client, CollectionIdBundle, ReadHolds, metrics, statement_logging}; /// Storage collections trait alias we need to consult for since/frontiers. pub type StorageCollectionsHandle = @@ -271,11 +275,10 @@ impl PeekClient { /// peek target. For slow-path peeks (to be implemented later), we'll need to additionally call /// into the Controller to acquire a hold on the peek target after we create the dataflow. /// - /// `logging_guard` owns end-of-execution logging for this statement. For a - /// constant peek it stays armed and the caller logs the end from the - /// returned result. For a `PeekExisting`/`PeekPersist` peek, successful - /// registration with the coordinator hands ownership of the end to the - /// coordinator and the guard is defused here. That holds even when the + /// For a constant peek the logging slot stays armed and the caller logs the + /// end from the returned result. For a `PeekExisting`/`PeekPersist` peek, + /// successful registration with the coordinator hands ownership of the end + /// to the coordinator and the slot is defused here. That holds even when the /// subsequent `client.peek()` fails to issue. pub(crate) async fn implement_fast_path_peek_plan( &mut self, @@ -294,7 +297,7 @@ impl PeekClient { conn_id: mz_adapter_types::connection::ConnectionId, depends_on: std::collections::BTreeSet, watch_set: Option, - logging_guard: &mut StatementLoggingGuard, + logging: &mut ExecutionLogging, ) -> Result { // If the dataflow optimizes to a constant expression, we can immediately return the result. if let FastPathPlan::Constant(rows_res, _) = fast_path { @@ -431,7 +434,7 @@ impl PeekClient { // cancellation, concurrent teardown (e.g. a DROP CLUSTER), or the // unregistration below. We defuse the guard so the frontend doesn't // also log the end. - logging_guard.defuse(); + logging.defuse(); // Test-only synchronization point: parks a peek between registration // and issue, so a test can land a concurrent DROP CLUSTER in this @@ -494,51 +497,32 @@ impl PeekClient { }) } - /// Set up statement logging for a frontend-sequenced operation. + /// Begins a new statement execution log entry, sampling permitting. /// - /// If `outer_ctx_extra` is `None`, begins a new statement execution log - /// entry. If `outer_ctx_extra` is `Some` (e.g. EXECUTE/FETCH), reuses and - /// retires the existing logging context. - /// - /// Returns a [`StatementLoggingGuard`]. Callers must either - /// [`retire`](StatementLoggingGuard::retire) the guard on the execution's - /// terminal outcome, or [`defuse`](StatementLoggingGuard::defuse) it at - /// the point where end-of-execution logging is handed off to another - /// component. Dropping the guard without retiring it emits an `Aborted` - /// end-execution event. - pub(crate) fn begin_statement_logging( + /// Only [`ExecutionLogging::take_over`] may call this: an entry that exists + /// without the session task owning its end would stay unfinished forever. + fn begin_statement_logging( &self, session: &mut Session, params: &Params, logging: &Arc>, catalog: &Catalog, lifecycle_timestamps: Option, - outer_ctx_extra: &mut Option, ) -> StatementLoggingGuard { - let id = if outer_ctx_extra.is_none() { - // This is a new statement, so begin statement logging. - let result = self.statement_logging_frontend.begin_statement_execution( - session, - params, - logging, - catalog.system_config(), - lifecycle_timestamps, - ); + let result = self.statement_logging_frontend.begin_statement_execution( + session, + params, + logging, + catalog.system_config(), + lifecycle_timestamps, + ); - if let Some((logging_id, began_execution, mseh_update, prepared_statement)) = result { + let id = result.map( + |(logging_id, began_execution, mseh_update, prepared_statement)| { self.log_began_execution(began_execution, mseh_update, prepared_statement); - Some(logging_id) - } else { - None - } - } else { - // We're executing in the context of another statement (e.g. FETCH), - // so take ownership of the outer context and inherit its logging id - // (if any). The end of execution will be logged by the caller. - outer_ctx_extra - .take() - .and_then(|guard| guard.defuse().retire()) - }; + logging_id + }, + ); StatementLoggingGuard { id, @@ -616,38 +600,16 @@ impl PeekClient { FrontendStatementLoggingEvent::Lifecycle { id, event, when }, )); } - - /// Emit a `FrontendStatementLoggingEvent::EndedExecution` for the given - /// logging id. Used by the few callers that hold a bare - /// [`StatementLoggingId`] rather than a [`StatementLoggingGuard`], e.g. - /// error paths of the EXECUTE unrolling where the id is not wrapped in a - /// guard. - pub(crate) fn log_ended_execution( - &self, - id: StatementLoggingId, - reason: statement_logging::StatementEndedExecutionReason, - ) { - let ended_at = (self.statement_logging_frontend.now)(); - let record = statement_logging::StatementEndedExecutionRecord { - id: id.0, - reason, - ended_at, - }; - self.coordinator_client - .send(Command::FrontendStatementLogging( - FrontendStatementLoggingEvent::EndedExecution(record), - )); - } } /// RAII guard owning a frontend statement-logging lifecycle. /// -/// Created by [`PeekClient::begin_statement_logging`]. Unless logging -/// responsibility is handed off via [`defuse`](StatementLoggingGuard::defuse), -/// the guard ensures that every statement for which `BeganExecution` was logged -/// also receives a corresponding `EndedExecution`, even on early-return, panic, -/// or mid-flight drop of the enclosing future: if the guard is dropped without -/// being defused, it emits `StatementEndedExecutionReason::Aborted`. +/// Unless logging responsibility is handed off via +/// [`defuse`](StatementLoggingGuard::defuse), the guard ensures that every +/// statement for which `BeganExecution` was logged also receives a +/// corresponding `EndedExecution`, even on early-return, panic, or mid-flight +/// drop of the enclosing future: if the guard is dropped without being defused, +/// it emits `StatementEndedExecutionReason::Aborted`. /// /// When the guard is `defuse`d, some other component (e.g. the coordinator, for /// streaming peek / subscribe responses) takes over and logs `EndedExecution` @@ -657,7 +619,7 @@ impl PeekClient { /// retirement / drop are no-ops. #[must_use = "StatementLoggingGuard must be explicitly retired or handed off; \ otherwise `Drop` will log the statement as Aborted"] -pub(crate) struct StatementLoggingGuard { +struct StatementLoggingGuard { /// `None` if the statement was not sampled for logging. id: Option, coordinator_client: Client, @@ -665,21 +627,37 @@ pub(crate) struct StatementLoggingGuard { } impl StatementLoggingGuard { + /// Arms a guard for the obligation the coordinator armed for `outer`, the + /// statement whose execution the one we are about to run serves. + fn adopt(outer: ExecuteContextGuard, peek_client: &PeekClient) -> Self { + Self { + id: outer.defuse().retire(), + coordinator_client: peek_client.coordinator_client.clone(), + now: peek_client.statement_logging_frontend.now.clone(), + } + } + /// Returns the logging id, if this statement is being logged. - pub(crate) fn id(&self) -> Option { + fn id(&self) -> Option { self.id } /// Retires the guard with an explicit end-execution reason. /// A no-op if the guard was defused or the statement is not sampled. - pub(crate) fn retire(mut self, reason: statement_logging::StatementEndedExecutionReason) { + fn retire(mut self, reason: statement_logging::StatementEndedExecutionReason) { self.emit(reason); } + /// Turns the obligation back into its transferable form, disarming this + /// guard. + fn release(mut self) -> ExecuteContextExtra { + ExecuteContextExtra::new(self.id.take()) + } + /// Hands off logging responsibility without emitting an end-execution /// event. Call this at the point where another component takes over /// end-of-execution logging. Afterwards the guard is inert. - pub(crate) fn defuse(&mut self) { + fn defuse(&mut self) { self.id = None; } @@ -693,8 +671,12 @@ impl StatementLoggingGuard { reason, ended_at, }; - self.coordinator_client - .send(Command::FrontendStatementLogging( + // A guard can outlive the coordinator during shutdown. Failing to send + // costs us one end event, panicking in `Drop` would cost the whole + // connection. + let _ = self + .coordinator_client + .try_send(Command::FrontendStatementLogging( FrontendStatementLoggingEvent::EndedExecution(record), )); } @@ -708,6 +690,217 @@ impl Drop for StatementLoggingGuard { } } +/// The session task's slot for the end-of-execution obligation of the statement +/// it is running. One slot is held for the whole `SessionClient::execute` call +/// and there is exactly one retirement site. +/// +/// An empty slot means no log entry exists for this execution, so a fallback to +/// the coordinator lets it begin its own. An occupied slot means the session +/// task owes an end event: [`Self::retire`] pays it, [`Self::release`] transfers +/// it to the coordinator. [`Self::id`] returning `None` covers both "not +/// sampled" and "the end is logged elsewhere", which want identical treatment +/// everywhere. +pub(crate) struct ExecutionLogging { + guard: Option, + /// Whether the coordinator must not run this statement, because the session + /// task has already counted it in the metrics `Coordinator::handle_execute` + /// maintains. Handing it over afterwards would count it twice. + coordinator_must_not_run: bool, +} + +/// Which statement the session task is taking the log entry over for. +pub(crate) enum TakeOver { + /// The statement that will run here, so the coordinator must not run it. + StatementToRun, + /// A SQL `EXECUTE` that unrolls into an inner statement, for a session task + /// that will go on to run that inner statement. + /// + /// The entry stays armed, so a failure to unroll is recorded against the + /// `EXECUTE`. Only the `EXECUTE` itself is counted here. The inner + /// statement is a statement in its own right and is counted wherever it + /// ends up running, exactly as the coordinator does when it re-dispatches a + /// `Plan::Execute`, which is why the slot stays releasable. + UnrolledExecute, +} + +impl ExecutionLogging { + /// Adopts the end-of-execution obligation of an outer statement (a FETCH or + /// an EXECUTE running its inner statement), or starts out empty when there + /// is no outer statement. + pub(crate) fn adopt(outer: Option, peek_client: &PeekClient) -> Self { + Self { + guard: outer.map(|outer| StatementLoggingGuard::adopt(outer, peek_client)), + coordinator_must_not_run: false, + } + } + + /// Returns the logging id, if an end event is owed and the statement is + /// being logged. + pub(crate) fn id(&self) -> Option { + self.guard.as_ref().and_then(|guard| guard.id()) + } + + /// Records that the session task, not the coordinator, is executing `stmt`. + /// Bumps the counters `Coordinator::handle_execute` would have bumped and + /// makes sure a log entry exists, inheriting the adopted one if there is + /// one. `stmt` is `None` for an empty portal, which is logged but not + /// counted. + /// + /// For [`TakeOver::StatementToRun`], every exit after this call must produce + /// an outcome for the statement: the coordinator will not see it. + pub(crate) fn take_over( + &mut self, + peek_client: &PeekClient, + session: &mut Session, + stmt: Option<&Statement>, + params: &Params, + logging: &Arc>, + catalog: &Catalog, + lifecycle_timestamps: Option, + taking_over: TakeOver, + ) -> Option { + self.begin_or_inherit( + peek_client, + session, + params, + logging, + catalog, + lifecycle_timestamps, + ); + count_statement(session, stmt); + if matches!(taking_over, TakeOver::StatementToRun) { + self.coordinator_must_not_run = true; + } + self.id() + } + + /// Makes sure a log entry exists for this execution, keeping an adopted one + /// if there is one. + fn begin_or_inherit( + &mut self, + peek_client: &PeekClient, + session: &mut Session, + params: &Params, + logging: &Arc>, + catalog: &Catalog, + lifecycle_timestamps: Option, + ) { + if self.guard.is_none() { + self.guard = Some(peek_client.begin_statement_logging( + session, + params, + logging, + catalog, + lifecycle_timestamps, + )); + } + } + + /// Hands the obligation to the coordinator, which retires it once the + /// statement it dispatches finishes. `None` tells the coordinator to begin + /// its own entry, `Some` with no id inside tells it that an entry already + /// exists or that sampling declined. + #[must_use] + pub(crate) fn release(&mut self) -> Option { + if self.coordinator_must_not_run { + soft_panic_or_log!( + "statement handed to the coordinator after the session task took it over: \ + its per-statement metrics are counted twice" + ); + } + self.guard.take().map(|guard| guard.release()) + } + + /// Emits the end event for `result`, if we still owe one. + pub(crate) fn retire(self, result: &Result) { + let Some(guard) = self.guard else { + return; + }; + // A defused or released slot owes no end event. Bail before mapping + // `result` to an end reason, which soft-panics for the responses whose + // end is logged elsewhere. + if guard.id().is_none() { + return; + } + guard.retire(end_reason(result)); + } + + /// Leaves the slot inert, for dispatch sites that hand the end of execution + /// to the coordinator or the protocol layer (registered peeks, subscribes). + pub(crate) fn defuse(&mut self) { + if let Some(guard) = self.guard.as_mut() { + guard.defuse(); + } + } +} + +/// Maps an execution outcome to the reason to record for it. +/// +/// The responses whose end is logged elsewhere are filtered out first: their +/// `StatementEndedExecutionReason` conversion panics, and this runs for every +/// statement the session task executes, so that panic would take down the +/// connection. +fn end_reason( + result: &Result, +) -> statement_logging::StatementEndedExecutionReason { + if let Ok(response) = result { + if terminates_elsewhere(response) { + soft_panic_or_log!( + "frontend-sequenced statement still owed an end event while returning {:?}", + crate::command::ExecuteResponseKind::from(response) + ); + return statement_logging::StatementEndedExecutionReason::Aborted; + } + } + result.into() +} + +/// Bumps the per-statement counters `Coordinator::handle_execute` maintains. +/// +/// `stmt` is `None` for an empty portal, which is logged but not counted. The +/// coordinator skips it the same way, and matching that is the point of this +/// function. +fn count_statement(session: &Session, stmt: Option<&Statement>) { + let Some(stmt) = stmt else { + return; + }; + let session_type = metrics::session_type_label_value(session.user()); + session + .metrics() + .query_total(&[session_type, metrics::statement_type_label_value(stmt)]) + .inc(); + if let Statement::Subscribe(SubscribeStatement { output, .. }) + | Statement::Copy(CopyStatement { + relation: CopyRelation::Subscribe(SubscribeStatement { output, .. }), + .. + }) = stmt + { + session + .metrics() + .subscribe_outputs(&[session_type, metrics::subscribe_output_label_value(output)]) + .inc(); + } +} + +/// Whether someone else logs the end of execution for `response`: the +/// coordinator for a registered peek, the protocol layer for a subscribe, a +/// FETCH or a COPY FROM. The dispatch sites that produce these defuse the slot, +/// so an armed slot alongside one of them means a dispatch site did not. +fn terminates_elsewhere(response: &ExecuteResponse) -> bool { + match response { + ExecuteResponse::SendingRowsStreaming { .. } + | ExecuteResponse::Subscribing { .. } + | ExecuteResponse::Fetch { .. } + | ExecuteResponse::CopyFrom { .. } => true, + // COPY TO STDOUT of an immediate result terminates here. Anything else + // it can wrap does not. + ExecuteResponse::CopyTo { resp, .. } => { + !matches!(**resp, ExecuteResponse::SendingRowsImmediate { .. }) + } + _ => false, + } +} + /// Errors arising during collection lookup in peek client operations. #[derive(Error, Debug)] pub enum CollectionLookupError { diff --git a/src/environmentd/tests/server.rs b/src/environmentd/tests/server.rs index 9ea6298a9d86b..7e3c0dd58f669 100644 --- a/src/environmentd/tests/server.rs +++ b/src/environmentd/tests/server.rs @@ -41,19 +41,15 @@ use mz_frontegg_auth::{ }; use mz_frontegg_mock::{FronteggMockServer, models::ApiToken, models::UserConfig}; use mz_ore::cast::CastFrom; -use mz_ore::cast::CastLossy; -use mz_ore::cast::TryCastFrom; -use mz_ore::collections::CollectionExt; use mz_ore::error::ErrorExt; use mz_ore::metrics::MetricsRegistry; -use mz_ore::now::{NowFn, SYSTEM_TIME, to_datetime}; +use mz_ore::now::{NowFn, SYSTEM_TIME}; use mz_ore::retry::Retry; use mz_ore::{assert_contains, task::RuntimeExt}; -use mz_ore::{assert_err, assert_none, assert_ok, task}; +use mz_ore::{assert_err, assert_ok, task}; use mz_pgrepr::UInt8; use mz_repr::UNKNOWN_COLUMN_NAME; use mz_sql::session::user::{ANALYTICS_USER, HTTP_DEFAULT_USER, SYSTEM_USER}; -use mz_sql_parser::ast::display::AstDisplay; use openssl::ssl::{SslConnectorBuilder, SslVerifyMode}; use openssl::x509::X509; use postgres::config::SslMode; @@ -172,978 +168,6 @@ fn test_persistence() { ); } -/// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped. -/// -/// This guard ensures that all statements have finished executing (have non-NULL `finished_at` -/// and `finished_status` in `mz_internal.mz_recent_activity_log`) before the test completes. -struct TestServerWithStatementLoggingChecks { - server: test_util::TestServerWithRuntime, -} - -impl TestServerWithStatementLoggingChecks { - /// Connect to the __internal__ SQL port of the running `environmentd` server. - pub fn connect_internal(&self, tls: T) -> Result - where - T: postgres::tls::MakeTlsConnect + Send + 'static, - T::TlsConnect: Send, - T::Stream: Send, - >::Future: Send, - { - self.server.connect_internal(tls) - } - - /// Returns the metrics registry for the test server. - pub fn metrics_registry(&self) -> &MetricsRegistry { - self.server.metrics_registry() - } -} - -/// Helper to get statement logging record counts from the metrics registry. -/// Returns (sampled_true_count, sampled_false_count). -#[allow(clippy::disallowed_methods)] -fn get_statement_logging_record_counts( - server: &TestServerWithStatementLoggingChecks, -) -> (u64, u64) { - let metrics = server.metrics_registry().gather(); - let record_count_metric = metrics - .into_iter() - .find(|m| m.name() == "mz_statement_logging_record_count") - .expect("mz_statement_logging_record_count metric should exist"); - - let metric_entries = record_count_metric.get_metric(); - let sampled_true = metric_entries - .iter() - .find(|m| { - m.get_label() - .iter() - .any(|l| l.name() == "sample" && l.value() == "true") - }) - .map(|m| u64::cast_lossy(m.get_counter().value())) - .unwrap_or(0); - let sampled_false = metric_entries - .iter() - .find(|m| { - m.get_label() - .iter() - .any(|l| l.name() == "sample" && l.value() == "false") - }) - .map(|m| u64::cast_lossy(m.get_counter().value())) - .unwrap_or(0); - - (sampled_true, sampled_false) -} - -impl Drop for TestServerWithStatementLoggingChecks { - #[allow(clippy::disallowed_methods)] - fn drop(&mut self) { - // Don't run checks if we're already panicking, as this could mask the original error. - if std::thread::panicking() { - return; - } - - let mut mz_client = self - .server - .connect_internal(postgres::NoTls) - .expect("Failed to connect to internal SQL port for statement logging check"); - - // Disable RBAC checks so we can query mz_internal tables. - // (We don't need to restore this afterwards, since no more tests run in the same system.) - mz_client - .batch_execute("ALTER SYSTEM SET enable_rbac_checks = false") - .expect("Failed to disable RBAC checks"); - - // The statement log has a 5-second buffer flush interval, so allow sufficient time. - Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let result = mz_client.query_one( - "SELECT count(*) - FROM mz_internal.mz_recent_activity_log - WHERE - (finished_at IS NULL OR finished_status IS NULL) - AND sql NOT LIKE '%__FILTER-OUT-THIS-QUERY__%' - AND finished_status IS DISTINCT FROM 'aborted'", - &[], - ); - - match result { - Ok(row) => { - let count: i64 = row.get(0); - if count == 0 { - Ok(()) - } else { - Err(format!("{} statements have not finished", count)) - } - } - Err(e) => Err(format!("Query failed: {}", e)), - } - }) - .expect("All statements should have finished executing"); - } -} - -fn setup_statement_logging_core( - max_sample_rate: f64, - sample_rate: f64, - target_data_rate: &str, - test_harness: test_util::TestHarness, -) -> (TestServerWithStatementLoggingChecks, postgres::Client) { - let server = test_harness - .with_system_parameter_default( - "statement_logging_max_sample_rate".to_string(), - max_sample_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_default_sample_rate".to_string(), - sample_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_max_data_credit".to_string(), - "".to_string(), - ) - .with_system_parameter_default( - "statement_logging_target_data_rate".to_string(), - target_data_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_use_reproducible_rng".to_string(), - "true".to_string(), - ) - .start_blocking(); - let client = server.connect(postgres::NoTls).unwrap(); - let server = TestServerWithStatementLoggingChecks { server }; - (server, client) -} - -fn setup_statement_logging( - max_sample_rate: f64, - sample_rate: f64, - target_data_rate: &str, -) -> (TestServerWithStatementLoggingChecks, postgres::Client) { - setup_statement_logging_core( - max_sample_rate, - sample_rate, - target_data_rate, - test_util::TestHarness::default(), - ) -} - -// Test that we log various kinds of statement whose execution terminates in the coordinator. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_immediate() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - - let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET enable_statement_lifecycle_logging = false") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_sample_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_default_sample_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET enable_load_generator_counter = true") - .unwrap(); - - let successful_immediates: &[&str] = &[ - "CREATE VIEW v AS SELECT 1;", - "CREATE DEFAULT INDEX i ON v;", - "CREATE TABLE t (x bigint);", - "INSERT INTO t VALUES (1), (2), (3)", - "UPDATE t SET x=x+1", - "DELETE FROM t;", - "CREATE SECRET s AS 'hunter2';", - "DROP SECRET s;", - "", - "CREATE SOURCE s FROM LOAD GENERATOR COUNTER", - "PREPARE foo AS SELECT * FROM t", - "EXECUTE foo", - "BEGIN", - "DECLARE c CURSOR FOR SELECT * FROM t", - "FETCH FORWARD ALL FROM c", - "COMMIT", - "BEGIN", - "ROLLBACK", - "SET application_name='my_application'", - "SHOW ALL", - "SHOW application_name", - ]; - let constants: &[&str] = &["1", "2", "3", "hunter2", "my_application"]; - - for &statement in successful_immediates { - client.execute(statement, &[]).unwrap(); - - // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering - // of logged statements non-deterministic when we retrieve them below. - thread::sleep(Duration::from_millis(10)); - } - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mst.sql, - mpsh.prepared_at, - mst.redacted_sql - FROM mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst - ON mpsh.sql_hash = mst.sql_hash - WHERE - mst.sql !~~ '%mz_statement_execution_history%' AND - mseh.finished_at IS NOT NULL - ORDER BY mseh.began_at"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = client.query(seh_query, &[]).unwrap(); - if sl.len() >= successful_immediates.len() { - break; - } - } - - assert_eq!(sl.len(), successful_immediates.len()); - - #[derive(Debug)] - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - sql: String, - prepared_at: DateTime, - redacted_sql: String, - } - for (r, stmt) in std::iter::zip(sl.iter(), successful_immediates) { - let r = Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - sql: r.get(4), - prepared_at: r.get(5), - redacted_sql: r.get(6), - }; - assert_eq!(r.sample_rate, 1.0); - - let expected_sql = if r.sql.contains("SECRET") - || r.sql.contains("INSERT") - || r.sql.contains("UPDATE") - || r.sql.contains("EXECUTE") - { - mz_sql::parse::parse(&r.sql) - .unwrap() - .into_element() - .ast - .to_ast_string_redacted() - } else { - stmt.chars().filter(|&ch| ch != ';').collect::() - }; - assert_eq!(r.sql, expected_sql); - assert_eq!(r.finished_status, "success"); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - // NB[btv] -- It would be a bit nicer if we could separately mock - // both the start and end time, but the `NowFn` mechanism doesn't - // appear to give us any way to do that. Instead, let's just check - // that none of these statements took longer than 5s wall-clock time. - assert!(r.finished_at - r.began_at <= chrono::Duration::try_seconds(5).unwrap()); - if !r.sql.is_empty() { - let expected_redacted = mz_sql::parse::parse(&r.sql) - .unwrap() - .into_element() - .ast - .to_ast_string_redacted(); - assert_eq!(r.redacted_sql, expected_redacted); - for constant in constants { - assert!(!r.redacted_sql.contains(constant)); - } - } - } -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_basic() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - client.execute("SELECT 1", &[]).unwrap(); - // We test that queries of this view execute on a cluster. - // If we ever change the threshold for constant folding such that - // this gets to run on environmentd, change this query. - client - .execute( - "CREATE VIEW v AS SELECT * FROM generate_series(1, 10001)", - &[], - ) - .unwrap(); - client.execute("SELECT * FROM v", &[]).unwrap(); - client.execute("CREATE DEFAULT INDEX i ON v", &[]).unwrap(); - client.execute("SELECT * FROM v", &[]).unwrap(); - let _ = client.execute("SELECT 1/0", &[]); - client.execute("CREATE TABLE t (x int)", &[]).unwrap(); - client.execute("SELECT * FROM t", &[]).unwrap(); - - #[derive(Debug)] - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - error_message: Option, - prepared_at: DateTime, - execution_strategy: Option, - result_size: Option, - rows_returned: Option, - execution_timestamp: Option, - } - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - - let result = Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mseh.error_message, - mpsh.prepared_at, - mseh.execution_strategy, - mseh.result_size, - mseh.rows_returned, - mseh.execution_timestamp -FROM - mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst - ON mpsh.sql_hash = mst.sql_hash -WHERE (mst.sql ~~ 'SELECT%' -AND mst.sql !~~ '%unique string to prevent this query showing up in results after retries%' -AND mst.sql !~~ '%pg_catalog.pg_type%' --this gets executed behind the scenes by tokio-postgres -OR mst.sql ~~ 'CREATE TABLE%') -AND mseh.finished_at IS NOT NULL -ORDER BY mseh.began_at", - &[], - ) - .unwrap(); - - if sl_results.len() == 6 { - Ok(sl_results) - } else { - Err(sl_results.len()) - } - }); - let sl_results = match result { - Ok(rows) => rows - .into_iter() - .map(|r| Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - error_message: r.get(4), - prepared_at: r.get(5), - execution_strategy: r.get(6), - result_size: r.get(7), - rows_returned: r.get(8), - execution_timestamp: r.get::<_, Option>(9).map(|UInt8(val)| val), - }) - .collect::>(), - Err(rows) => { - panic!("number of results never became correct: {rows}"); - } - }; - // The two queries on generate_series(1,10001) execute at the maximum timestamp - assert_eq!( - sl_results - .iter() - .filter(|r| r.execution_timestamp == Some(u64::MAX)) - .count(), - 2 - ); - // The two queries that can be satisfied by envd (SELECT 1 and SELECT 1/0) have no execution timestamp - assert_eq!( - sl_results - .iter() - .filter(|r| r.execution_timestamp.is_none()) - .count(), - 2 - ); - // All other queries have an execution timestamp, in particular, including `CREATE TABLE`. - assert_eq!(sl_results.len(), 6); - for r in &sl_results { - assert_eq!(r.sample_rate, 1.0); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - // It would be nice to be able to control - // execution timestamp via a `NowFn`, but - // that is hard to get right and interferes with our logic - // about when to flush to persist. So instead, just check that they're sane. - if let Some(ts) = r.execution_timestamp { - if ts != u64::MAX { - let ts = to_datetime(ts); - assert!((ts - r.prepared_at).abs() < chrono::Duration::try_seconds(5).unwrap()) - } - } - } - assert!(sl_results[0].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[0].rows_returned, Some(1)); - assert_eq!(sl_results[0].finished_status, "success"); - assert_eq!( - sl_results[0].execution_strategy.as_ref().unwrap(), - "constant" - ); - assert!(sl_results[1].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[1].rows_returned, Some(10001)); - assert_eq!(sl_results[1].finished_status, "success"); - assert_eq!( - sl_results[1].execution_strategy.as_ref().unwrap(), - "standard" - ); - assert!(sl_results[2].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[2].rows_returned, Some(10001)); - assert_eq!(sl_results[2].finished_status, "success"); - assert_eq!( - sl_results[2].execution_strategy.as_ref().unwrap(), - "fast-path" - ); - assert_eq!(sl_results[3].finished_status, "error"); - assert!( - sl_results[3] - .error_message - .as_ref() - .unwrap() - .contains("division by zero") - ); - assert_none!(sl_results[3].result_size); - assert_none!(sl_results[3].rows_returned); - - // Verify metrics show all statements were sampled (100% sample rate means no unsampled). - let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); - assert!( - sampled_true > 0, - "some statements should be sampled with 100% rate" - ); - assert_eq!( - sampled_false, 0, - "no statements should be unsampled with 100% rate" - ); - - // Verify statement_logging_actual_bytes metric is being tracked. - // With 100% sample rate, actual_bytes should equal unsampled_bytes. - let metrics = server.metrics_registry().gather(); - let actual_bytes = metrics - .iter() - .find(|m| m.name() == "mz_statement_logging_actual_bytes") - .expect("mz_statement_logging_actual_bytes metric should exist") - .get_metric()[0] - .get_counter() - .value(); - let unsampled_bytes = metrics - .iter() - .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") - .expect("mz_statement_logging_unsampled_bytes metric should exist") - .get_metric()[0] - .get_counter() - .value(); - assert!( - actual_bytes > 0.0, - "actual_bytes should be > 0 with 100% sample rate" - ); - assert_eq!( - actual_bytes, unsampled_bytes, - "with 100% sample rate, actual_bytes should equal unsampled_bytes" - ); -} - -#[allow(clippy::disallowed_methods)] -fn run_throttling_test(use_prepared_statement: bool) { - // The `target_data_rate` should be - // - high enough so that the `SELECT 1` queries get throttled (even with high CPU load due to - // other tests running in parallel), - // - but low enough that the `SELECT 2` query after the sleep doesn't get throttled. - let (server, mut client) = setup_statement_logging(1.0, 1.0, "200"); - thread::sleep(Duration::from_secs(2)); - - if use_prepared_statement { - let statement = client.prepare("SELECT 1").unwrap(); - for _ in 0..100 { - client.execute(&statement, &[]).unwrap(); - } - } else { - for _ in 0..100 { - client.execute("SELECT 1", &[]).unwrap(); - } - } - - thread::sleep(Duration::from_secs(4)); - client.execute("SELECT 2", &[]).unwrap(); - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let logs = Retry::default() - .max_duration(Duration::from_secs(60)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - sql, - throttled_count -FROM mz_internal.mz_statement_execution_history mseh -JOIN mz_internal.mz_prepared_statement_history mpsh -ON mseh.prepared_statement_id = mpsh.id -JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst -ON mpsh.sql_hash = mst.sql_hash -WHERE sql IN ('SELECT 1', 'SELECT 2')", - &[], - ) - .unwrap(); - - if sl_results.iter().any(|stmt| { - let sql: String = stmt.get(0); - sql == "SELECT 2" - }) { - Ok(sl_results) - } else { - Err(()) - } - }) - .expect("Never saw last statement (`SELECT 2`)"); - let throttled_count = logs - .iter() - .map(|log| { - let UInt8(throttled_count) = log.get(1); - throttled_count - }) - .sum::(); - assert!( - throttled_count > 0, - "at least some statements should have been throttled" - ); - - assert_eq!(logs.len() + usize::cast_from(throttled_count), 101); -} - -#[mz_ore::test] -fn test_statement_logging_throttling() { - run_throttling_test(false); -} - -#[mz_ore::test] -fn test_statement_logging_prepared_statement_throttling() { - run_throttling_test(true); -} - -/// throttling the first execution of a prepared statement must not make -// every subsequent execution of that statement invisible in -// `mz_recent_activity_log`. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_throttled_prepared_statement_stays_visible() { - const APPLICATION_NAME: &str = "throttled_prepared_repro"; - const NUM_SUBSEQUENT_STATEMENT_EXECUTIONS: i64 = 5; - - let (server, mut client) = setup_statement_logging(1.0, 1.0, "1"); - let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); - - // cap the credit at a single byte such that the first execution of our - // prepared statement is guaranteed to be throttled. - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1") - .unwrap(); - - client - .batch_execute(&format!("SET application_name = '{APPLICATION_NAME}'")) - .unwrap(); - - // First execution of the prepared statement. This one gets throttled, and so - // is never logged. - let stmt = client.prepare("SELECT $1::text").unwrap(); - assert_eq!( - client - .query_one(&stmt, &[&"first"]) - .unwrap() - .get::<_, String>(0), - "first" - ); - - // Lift the throttle so subsequent executions are logged. Sleep afterwards so - // the token bucket has time to refill before the next execution - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1000000000") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1000000000") - .unwrap(); - thread::sleep(Duration::from_secs(2)); - - // Subsequent executions of the same prepared statement. - for _ in 0..NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { - assert_eq!( - client - .query_one(&stmt, &[&"again"]) - .unwrap() - .get::<_, String>(0), - "again" - ); - } - - // Statement logging flushes asynchronously; wait until the - // subsequent executions have landed in `mz_statement_execution_history`. - Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let mseh_count: i64 = mz_client - .query_one( - &format!( - "SELECT count(*) - FROM mz_internal.mz_statement_execution_history - WHERE application_name = '{APPLICATION_NAME}'" - ), - &[], - ) - .unwrap() - .get(0); - if mseh_count >= NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { - Ok(()) - } else { - Err(format!( - "only {mseh_count} of {NUM_SUBSEQUENT_STATEMENT_EXECUTIONS} executions logged so far" - )) - } - }) - .expect("subsequent executions never appeared in mz_statement_execution_history"); - - // The subsequent executions in mz_statement_execution_history should - // have a matching entry in mz_prepared_statement_history - let orphan_count: i64 = mz_client - .query_one( - &format!( - "SELECT count(*) - FROM mz_internal.mz_statement_execution_history mseh - LEFT JOIN mz_internal.mz_prepared_statement_history mpsh - ON mseh.prepared_statement_id = mpsh.id - WHERE mseh.application_name = '{APPLICATION_NAME}' - AND mpsh.id IS NULL" - ), - &[], - ) - .unwrap() - .get(0); - assert_eq!( - orphan_count, 0, - "{orphan_count} mz_statement_execution_history rows reference a \ - prepared_statement_id with no matching mz_prepared_statement_history row" - ); -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_subscribes() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - let cancel_token = client.cancel_token(); - - // This should finish - client - .execute( - "SUBSCRIBE TO (SELECT * FROM generate_series(1, 10001))", - &[], - ) - .unwrap(); - - let handle = thread::spawn(move || { - client.execute("CREATE TABLE t (x int)", &[]).unwrap(); - // This should not finish until it's canceled. - let _ = client.execute("SUBSCRIBE TO (SELECT * FROM t)", &[]); - }); - - while !handle.is_finished() { - thread::sleep(Duration::from_secs(1)); - cancel_token.cancel_query(postgres::NoTls).unwrap(); - } - handle.join().unwrap(); - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mpsh.prepared_at, - mseh.execution_strategy - FROM mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - mz_internal.mz_sql_text AS mst - ON mpsh.sql_hash = mst.sql_hash - WHERE - mst.sql ~~ 'SUBSCRIBE%' AND - mseh.finished_at IS NOT NULL - ORDER BY mseh.began_at"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = client.query(seh_query, &[]).unwrap(); - if sl.len() >= 2 { - break; - } - } - - assert_eq!(sl.len(), 2); - - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - prepared_at: DateTime, - execution_strategy: Option, - } - - let sl_subscribes = sl - .into_iter() - .map(|r| Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - prepared_at: r.get(4), - execution_strategy: r.get(5), - }) - .collect::>(); - for r in &sl_subscribes { - assert_eq!(r.sample_rate, 1.0); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - assert_none!(r.execution_strategy); - } - assert_eq!(sl_subscribes[0].finished_status, "success"); - assert_eq!(sl_subscribes[1].finished_status, "canceled"); -} - -/// Test that we are sampling approximately 50% of statements. -/// Relies on two assumptions: -/// (1) that the effective sampling rate for the session is 50%, -/// (2) that we are using the deterministic testing RNG. -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_sampling_inner( - server: TestServerWithStatementLoggingChecks, - mut client: postgres::Client, -) { - for i in 0..50 { - client.execute(&format!("SELECT {i}"), &[]).unwrap(); - - // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering - // of logged statements non-deterministic when we retrieve them below. - thread::sleep(Duration::from_millis(10)); - } - - // 23 randomly sampled out of 50 with 50% sampling. Seems legit! - let expected_sqls = [ - 2, 4, 5, 6, 9, 15, 17, 18, 19, 20, 21, 23, 24, 25, 31, 32, 33, 36, 37, 42, 46, - ] - .into_iter() - .map(|i| format!("SELECT {i}")) - .collect::>(); - - let mut internal_client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT mst.sql - FROM mz_internal.mz_statement_execution_history AS mseh - JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - mz_internal.mz_sql_text AS mst - ON mpsh.sql_hash = mst.sql_hash - WHERE mst.sql ~~ 'SELECT%' AND mst.sql !~~ '%mz_statement_execution_history%' - ORDER BY mseh.began_at ASC"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = internal_client.query(seh_query, &[]).unwrap(); - if sl.len() >= expected_sqls.len() { - break; - } - } - - let sqls: Vec = sl.into_iter().map(|r| r.get(0)).collect(); - assert_eq!(sqls, expected_sqls); - - // Verify the statement_logging_record_count metric correctly tracks sampled vs unsampled. - // With 50% sampling and deterministic RNG, exactly 21 of 50 statements should be sampled. - let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); - assert_eq!( - sampled_true, 21, - "expected 21 statements to be sampled with 50% rate and deterministic RNG" - ); - assert_eq!( - sampled_false, 29, - "expected 29 statements to not be sampled with 50% rate and deterministic RNG" - ); -} - -#[mz_ore::test] -fn test_statement_logging_sampling() { - let (server, client) = setup_statement_logging(1.0, 0.5, ""); - test_statement_logging_sampling_inner(server, client); -} - -/// Test that we are not allowed to set `statement_logging_sample_rate` -/// arbitrarily high, but that it is constrained by `statement_logging_max_sample_rate`. -#[mz_ore::test] -fn test_statement_logging_sampling_constrained() { - let (server, client) = setup_statement_logging(0.5, 1.0, ""); - test_statement_logging_sampling_inner(server, client); -} - -/// Test that the `mz_statement_logging_unsampled_bytes` metric tracks the total bytes -/// of SQL text that would have been logged if statement logging were fully enabled. -/// We set `sample_rate=0.0` so no statements are actually sampled/logged, but the -/// unsampled_bytes metric still gets incremented for every executed statement. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_unsampled_metrics() { - // Use sample_rate=0.0 so statements are not sampled, but unsampled_bytes metric is still tracked. - let (server, mut client) = setup_statement_logging(1.0, 0.0, ""); - - let batch_queries = [ - "SELECT 'Hello, world!';SELECT 1;;", - "SELECT 'Hello, world again!'", - ]; - let batch_total: usize = batch_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - let single_queries = ["SELECT 'foo'", "SELECT 'bar';;;"]; - let single_total: usize = single_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - let prepared_queries = ["SELECT 'baz';;;", "SELECT 'quux';"]; - let prepared_total: usize = prepared_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - - let named_prepared_inner = "SELECT 42"; - let named_prepared_outer = format!("PREPARE p AS {named_prepared_inner};EXECUTE p;"); - let named_prepared_outer_len = named_prepared_outer - .as_bytes() - .iter() - .filter(|&&ch| ch != b';') - .count(); - - for q in batch_queries { - client.batch_execute(q).unwrap(); - } - - for q in single_queries { - client.execute(q, &[]).unwrap(); - } - - for q in prepared_queries { - let s = client.prepare(q).unwrap(); - client.execute(&s, &[]).unwrap(); - } - - client.batch_execute(&named_prepared_outer).unwrap(); - - // This should NOT be logged, since we never actually execute it. - client.prepare("SELECT 'Hello, not counted!'").unwrap(); - - let expected_total = batch_total + single_total + prepared_total + named_prepared_outer_len; - let metric_value = server - .metrics_registry() - .gather() - .into_iter() - .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") - .unwrap() - .take_metric()[0] - .get_counter() - .value(); - let metric_value = usize::cast_from(u64::try_cast_from(metric_value).unwrap()); - assert_eq!(expected_total, metric_value); - - // Also verify that statement_logging_record_count shows all statements as not sampled - // (since we're using 0% sample rate). - let (sampled_true, _sampled_false) = get_statement_logging_record_counts(&server); - assert_eq!( - sampled_true, 0, - "no statements should be sampled with 0% sample rate" - ); -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_enable_internal_statement_logging() { - let (server, mut client) = setup_statement_logging_core( - 1.0, - 1.0, - "", - test_util::TestHarness::default().with_system_parameter_default( - "enable_internal_statement_logging".to_string(), - "true".to_string(), - ), - ); - - client.execute("SELECT 1", &[]).unwrap(); - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let num_mz_system_statements = Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - count(*) -FROM mz_internal.mz_prepared_statement_history mpsh -JOIN mz_internal.mz_session_history USING (session_id) -WHERE authenticated_user='mz_system'", - &[], - ) - .unwrap(); - - let count: i64 = sl_results[0].get(0); - - if count > 0 { Ok(count) } else { Err(()) } - }) - .expect("at least some statements from mz_system should have been logged"); - - assert!( - num_mz_system_statements > 0, - "statements executed by mz_system should have been logged" - ); -} - // Test the POST and WS server endpoints. #[mz_ore::test] #[allow(clippy::disallowed_methods)] @@ -2246,28 +1270,6 @@ fn test_ws_passes_options() { } } -#[mz_ore::test] -// Test that statement logging -// doesn't cause a crash with subscribes over web sockets, -// which was previously happening (in staging) due to us -// dropping the `ExecuteContext` on the floor in that case. -fn test_statement_logging_ws_subscribe_no_crash() { - let (server, _client) = setup_statement_logging(1.0, 1.0, ""); - - // Create our WebSocket. - let ws_url = server.server.ws_addr(); - let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap(); - test_util::auth_with_ws(&mut ws, Default::default()).unwrap(); - - let query = "SUBSCRIBE (SELECT 1)"; - let json = format!("{{\"query\":\"{query}\"}}"); - let json: serde_json::Value = serde_json::from_str(&json).unwrap(); - ws.send(Message::Text(json.to_string().into())).unwrap(); - - // Give the server time to crash, if it's going to. - std::thread::sleep(Duration::from_secs(1)) -} - #[mz_ore::test] #[cfg_attr(miri, ignore)] // too slow fn test_ws_notifies_for_bad_options() { diff --git a/src/environmentd/tests/statement_logging.rs b/src/environmentd/tests/statement_logging.rs new file mode 100644 index 0000000000000..859ef1dd493a1 --- /dev/null +++ b/src/environmentd/tests/statement_logging.rs @@ -0,0 +1,1091 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Integration tests for statement logging. + +use std::thread; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use mz_environmentd::test_util; +use mz_ore::assert_none; +use mz_ore::cast::{CastFrom, CastLossy, TryCastFrom}; +use mz_ore::collections::CollectionExt; +use mz_ore::metrics::MetricsRegistry; +use mz_ore::now::to_datetime; +use mz_ore::retry::Retry; +use mz_pgrepr::UInt8; +use mz_sql_parser::ast::display::AstDisplay; +use tungstenite::Message; + +/// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped. +/// +/// This guard ensures that all statements have finished executing (have non-NULL `finished_at` +/// and `finished_status` in `mz_internal.mz_recent_activity_log`) before the test completes. +struct TestServerWithStatementLoggingChecks { + server: test_util::TestServerWithRuntime, +} + +impl TestServerWithStatementLoggingChecks { + /// Connect to the __internal__ SQL port of the running `environmentd` server. + pub fn connect_internal(&self, tls: T) -> Result + where + T: postgres::tls::MakeTlsConnect + Send + 'static, + T::TlsConnect: Send, + T::Stream: Send, + >::Future: Send, + { + self.server.connect_internal(tls) + } + + /// Returns the metrics registry for the test server. + pub fn metrics_registry(&self) -> &MetricsRegistry { + self.server.metrics_registry() + } +} + +/// Helper to get statement logging record counts from the metrics registry. +/// Returns (sampled_true_count, sampled_false_count). +#[allow(clippy::disallowed_methods)] +fn get_statement_logging_record_counts( + server: &TestServerWithStatementLoggingChecks, +) -> (u64, u64) { + let metrics = server.metrics_registry().gather(); + let record_count_metric = metrics + .into_iter() + .find(|m| m.name() == "mz_statement_logging_record_count") + .expect("mz_statement_logging_record_count metric should exist"); + + let metric_entries = record_count_metric.get_metric(); + let sampled_true = metric_entries + .iter() + .find(|m| { + m.get_label() + .iter() + .any(|l| l.name() == "sample" && l.value() == "true") + }) + .map(|m| u64::cast_lossy(m.get_counter().value())) + .unwrap_or(0); + let sampled_false = metric_entries + .iter() + .find(|m| { + m.get_label() + .iter() + .any(|l| l.name() == "sample" && l.value() == "false") + }) + .map(|m| u64::cast_lossy(m.get_counter().value())) + .unwrap_or(0); + + (sampled_true, sampled_false) +} + +impl Drop for TestServerWithStatementLoggingChecks { + #[allow(clippy::disallowed_methods)] + fn drop(&mut self) { + // Don't run checks if we're already panicking, as this could mask the original error. + if std::thread::panicking() { + return; + } + + let mut mz_client = self + .server + .connect_internal(postgres::NoTls) + .expect("Failed to connect to internal SQL port for statement logging check"); + + // Disable RBAC checks so we can query mz_internal tables. + // (We don't need to restore this afterwards, since no more tests run in the same system.) + mz_client + .batch_execute("ALTER SYSTEM SET enable_rbac_checks = false") + .expect("Failed to disable RBAC checks"); + + // The statement log has a 5-second buffer flush interval, so allow sufficient time. + Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let result = mz_client.query_one( + "SELECT count(*) + FROM mz_internal.mz_recent_activity_log + WHERE + (finished_at IS NULL OR finished_status IS NULL) + AND sql NOT LIKE '%__FILTER-OUT-THIS-QUERY__%' + AND finished_status IS DISTINCT FROM 'aborted'", + &[], + ); + + match result { + Ok(row) => { + let count: i64 = row.get(0); + if count == 0 { + Ok(()) + } else { + Err(format!("{} statements have not finished", count)) + } + } + Err(e) => Err(format!("Query failed: {}", e)), + } + }) + .expect("All statements should have finished executing"); + } +} + +fn setup_statement_logging_core( + max_sample_rate: f64, + sample_rate: f64, + target_data_rate: &str, + test_harness: test_util::TestHarness, +) -> (TestServerWithStatementLoggingChecks, postgres::Client) { + let server = test_harness + .with_system_parameter_default( + "statement_logging_max_sample_rate".to_string(), + max_sample_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_default_sample_rate".to_string(), + sample_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_max_data_credit".to_string(), + "".to_string(), + ) + .with_system_parameter_default( + "statement_logging_target_data_rate".to_string(), + target_data_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_use_reproducible_rng".to_string(), + "true".to_string(), + ) + .start_blocking(); + let client = server.connect(postgres::NoTls).unwrap(); + let server = TestServerWithStatementLoggingChecks { server }; + (server, client) +} + +fn setup_statement_logging( + max_sample_rate: f64, + sample_rate: f64, + target_data_rate: &str, +) -> (TestServerWithStatementLoggingChecks, postgres::Client) { + setup_statement_logging_core( + max_sample_rate, + sample_rate, + target_data_rate, + test_util::TestHarness::default(), + ) +} + +// Test that we log various kinds of statement whose execution terminates in the coordinator. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_immediate() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET enable_statement_lifecycle_logging = false") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_sample_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_default_sample_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET enable_load_generator_counter = true") + .unwrap(); + + let successful_immediates: &[&str] = &[ + "CREATE VIEW v AS SELECT 1;", + "CREATE DEFAULT INDEX i ON v;", + "CREATE TABLE t (x bigint);", + "INSERT INTO t VALUES (1), (2), (3)", + "UPDATE t SET x=x+1", + "DELETE FROM t;", + "CREATE SECRET s AS 'hunter2';", + "DROP SECRET s;", + "", + "CREATE SOURCE s FROM LOAD GENERATOR COUNTER", + "PREPARE foo AS SELECT * FROM t", + "EXECUTE foo", + "BEGIN", + "DECLARE c CURSOR FOR SELECT * FROM t", + "FETCH FORWARD ALL FROM c", + "COMMIT", + "BEGIN", + "ROLLBACK", + "SET application_name='my_application'", + "SHOW ALL", + "SHOW application_name", + ]; + let constants: &[&str] = &["1", "2", "3", "hunter2", "my_application"]; + + for &statement in successful_immediates { + client.execute(statement, &[]).unwrap(); + + // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering + // of logged statements non-deterministic when we retrieve them below. + thread::sleep(Duration::from_millis(10)); + } + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mst.sql, + mpsh.prepared_at, + mst.redacted_sql + FROM mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst + ON mpsh.sql_hash = mst.sql_hash + WHERE + mst.sql !~~ '%mz_statement_execution_history%' AND + mseh.finished_at IS NOT NULL + ORDER BY mseh.began_at"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = client.query(seh_query, &[]).unwrap(); + if sl.len() >= successful_immediates.len() { + break; + } + } + + assert_eq!(sl.len(), successful_immediates.len()); + + #[derive(Debug)] + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + sql: String, + prepared_at: DateTime, + redacted_sql: String, + } + for (r, stmt) in std::iter::zip(sl.iter(), successful_immediates) { + let r = Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + sql: r.get(4), + prepared_at: r.get(5), + redacted_sql: r.get(6), + }; + assert_eq!(r.sample_rate, 1.0); + + let expected_sql = if r.sql.contains("SECRET") + || r.sql.contains("INSERT") + || r.sql.contains("UPDATE") + || r.sql.contains("EXECUTE") + { + mz_sql::parse::parse(&r.sql) + .unwrap() + .into_element() + .ast + .to_ast_string_redacted() + } else { + stmt.chars().filter(|&ch| ch != ';').collect::() + }; + assert_eq!(r.sql, expected_sql); + assert_eq!(r.finished_status, "success"); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + // NB[btv] -- It would be a bit nicer if we could separately mock + // both the start and end time, but the `NowFn` mechanism doesn't + // appear to give us any way to do that. Instead, let's just check + // that none of these statements took longer than 5s wall-clock time. + assert!(r.finished_at - r.began_at <= chrono::Duration::try_seconds(5).unwrap()); + if !r.sql.is_empty() { + let expected_redacted = mz_sql::parse::parse(&r.sql) + .unwrap() + .into_element() + .ast + .to_ast_string_redacted(); + assert_eq!(r.redacted_sql, expected_redacted); + for constant in constants { + assert!(!r.redacted_sql.contains(constant)); + } + } + } +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_basic() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + client.execute("SELECT 1", &[]).unwrap(); + // We test that queries of this view execute on a cluster. + // If we ever change the threshold for constant folding such that + // this gets to run on environmentd, change this query. + client + .execute( + "CREATE VIEW v AS SELECT * FROM generate_series(1, 10001)", + &[], + ) + .unwrap(); + client.execute("SELECT * FROM v", &[]).unwrap(); + client.execute("CREATE DEFAULT INDEX i ON v", &[]).unwrap(); + client.execute("SELECT * FROM v", &[]).unwrap(); + let _ = client.execute("SELECT 1/0", &[]); + client.execute("CREATE TABLE t (x int)", &[]).unwrap(); + client.execute("SELECT * FROM t", &[]).unwrap(); + + #[derive(Debug)] + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + error_message: Option, + prepared_at: DateTime, + execution_strategy: Option, + result_size: Option, + rows_returned: Option, + execution_timestamp: Option, + } + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + + let result = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mseh.error_message, + mpsh.prepared_at, + mseh.execution_strategy, + mseh.result_size, + mseh.rows_returned, + mseh.execution_timestamp +FROM + mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE (mst.sql ~~ 'SELECT%' +AND mst.sql !~~ '%unique string to prevent this query showing up in results after retries%' +AND mst.sql !~~ '%pg_catalog.pg_type%' --this gets executed behind the scenes by tokio-postgres +OR mst.sql ~~ 'CREATE TABLE%') +AND mseh.finished_at IS NOT NULL +ORDER BY mseh.began_at", + &[], + ) + .unwrap(); + + if sl_results.len() == 6 { + Ok(sl_results) + } else { + Err(sl_results.len()) + } + }); + let sl_results = match result { + Ok(rows) => rows + .into_iter() + .map(|r| Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + error_message: r.get(4), + prepared_at: r.get(5), + execution_strategy: r.get(6), + result_size: r.get(7), + rows_returned: r.get(8), + execution_timestamp: r.get::<_, Option>(9).map(|UInt8(val)| val), + }) + .collect::>(), + Err(rows) => { + panic!("number of results never became correct: {rows}"); + } + }; + // The two queries on generate_series(1,10001) execute at the maximum timestamp + assert_eq!( + sl_results + .iter() + .filter(|r| r.execution_timestamp == Some(u64::MAX)) + .count(), + 2 + ); + // The two queries that can be satisfied by envd (SELECT 1 and SELECT 1/0) have no execution timestamp + assert_eq!( + sl_results + .iter() + .filter(|r| r.execution_timestamp.is_none()) + .count(), + 2 + ); + // All other queries have an execution timestamp, in particular, including `CREATE TABLE`. + assert_eq!(sl_results.len(), 6); + for r in &sl_results { + assert_eq!(r.sample_rate, 1.0); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + // It would be nice to be able to control + // execution timestamp via a `NowFn`, but + // that is hard to get right and interferes with our logic + // about when to flush to persist. So instead, just check that they're sane. + if let Some(ts) = r.execution_timestamp { + if ts != u64::MAX { + let ts = to_datetime(ts); + assert!((ts - r.prepared_at).abs() < chrono::Duration::try_seconds(5).unwrap()) + } + } + } + assert!(sl_results[0].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[0].rows_returned, Some(1)); + assert_eq!(sl_results[0].finished_status, "success"); + assert_eq!( + sl_results[0].execution_strategy.as_ref().unwrap(), + "constant" + ); + assert!(sl_results[1].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[1].rows_returned, Some(10001)); + assert_eq!(sl_results[1].finished_status, "success"); + assert_eq!( + sl_results[1].execution_strategy.as_ref().unwrap(), + "standard" + ); + assert!(sl_results[2].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[2].rows_returned, Some(10001)); + assert_eq!(sl_results[2].finished_status, "success"); + assert_eq!( + sl_results[2].execution_strategy.as_ref().unwrap(), + "fast-path" + ); + assert_eq!(sl_results[3].finished_status, "error"); + assert!( + sl_results[3] + .error_message + .as_ref() + .unwrap() + .contains("division by zero") + ); + assert_none!(sl_results[3].result_size); + assert_none!(sl_results[3].rows_returned); + + // Verify metrics show all statements were sampled (100% sample rate means no unsampled). + let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); + assert!( + sampled_true > 0, + "some statements should be sampled with 100% rate" + ); + assert_eq!( + sampled_false, 0, + "no statements should be unsampled with 100% rate" + ); + + // Verify statement_logging_actual_bytes metric is being tracked. + // With 100% sample rate, actual_bytes should equal unsampled_bytes. + let metrics = server.metrics_registry().gather(); + let actual_bytes = metrics + .iter() + .find(|m| m.name() == "mz_statement_logging_actual_bytes") + .expect("mz_statement_logging_actual_bytes metric should exist") + .get_metric()[0] + .get_counter() + .value(); + let unsampled_bytes = metrics + .iter() + .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") + .expect("mz_statement_logging_unsampled_bytes metric should exist") + .get_metric()[0] + .get_counter() + .value(); + assert!( + actual_bytes > 0.0, + "actual_bytes should be > 0 with 100% sample rate" + ); + assert_eq!( + actual_bytes, unsampled_bytes, + "with 100% sample rate, actual_bytes should equal unsampled_bytes" + ); +} + +#[allow(clippy::disallowed_methods)] +fn run_throttling_test(use_prepared_statement: bool) { + // The `target_data_rate` should be + // - high enough so that the `SELECT 1` queries get throttled (even with high CPU load due to + // other tests running in parallel), + // - but low enough that the `SELECT 2` query after the sleep doesn't get throttled. + let (server, mut client) = setup_statement_logging(1.0, 1.0, "200"); + thread::sleep(Duration::from_secs(2)); + + if use_prepared_statement { + let statement = client.prepare("SELECT 1").unwrap(); + for _ in 0..100 { + client.execute(&statement, &[]).unwrap(); + } + } else { + for _ in 0..100 { + client.execute("SELECT 1", &[]).unwrap(); + } + } + + thread::sleep(Duration::from_secs(4)); + client.execute("SELECT 2", &[]).unwrap(); + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let logs = Retry::default() + .max_duration(Duration::from_secs(60)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + sql, + throttled_count +FROM mz_internal.mz_statement_execution_history mseh +JOIN mz_internal.mz_prepared_statement_history mpsh +ON mseh.prepared_statement_id = mpsh.id +JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst +ON mpsh.sql_hash = mst.sql_hash +WHERE sql IN ('SELECT 1', 'SELECT 2')", + &[], + ) + .unwrap(); + + if sl_results.iter().any(|stmt| { + let sql: String = stmt.get(0); + sql == "SELECT 2" + }) { + Ok(sl_results) + } else { + Err(()) + } + }) + .expect("Never saw last statement (`SELECT 2`)"); + let throttled_count = logs + .iter() + .map(|log| { + let UInt8(throttled_count) = log.get(1); + throttled_count + }) + .sum::(); + assert!( + throttled_count > 0, + "at least some statements should have been throttled" + ); + + assert_eq!(logs.len() + usize::cast_from(throttled_count), 101); +} + +#[mz_ore::test] +fn test_statement_logging_throttling() { + run_throttling_test(false); +} + +#[mz_ore::test] +fn test_statement_logging_prepared_statement_throttling() { + run_throttling_test(true); +} + +/// throttling the first execution of a prepared statement must not make +// every subsequent execution of that statement invisible in +// `mz_recent_activity_log`. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_throttled_prepared_statement_stays_visible() { + const APPLICATION_NAME: &str = "throttled_prepared_repro"; + const NUM_SUBSEQUENT_STATEMENT_EXECUTIONS: i64 = 5; + + let (server, mut client) = setup_statement_logging(1.0, 1.0, "1"); + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + + // cap the credit at a single byte such that the first execution of our + // prepared statement is guaranteed to be throttled. + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1") + .unwrap(); + + client + .batch_execute(&format!("SET application_name = '{APPLICATION_NAME}'")) + .unwrap(); + + // First execution of the prepared statement. This one gets throttled, and so + // is never logged. + let stmt = client.prepare("SELECT $1::text").unwrap(); + assert_eq!( + client + .query_one(&stmt, &[&"first"]) + .unwrap() + .get::<_, String>(0), + "first" + ); + + // Lift the throttle so subsequent executions are logged. Sleep afterwards so + // the token bucket has time to refill before the next execution + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1000000000") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1000000000") + .unwrap(); + thread::sleep(Duration::from_secs(2)); + + // Subsequent executions of the same prepared statement. + for _ in 0..NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { + assert_eq!( + client + .query_one(&stmt, &[&"again"]) + .unwrap() + .get::<_, String>(0), + "again" + ); + } + + // Statement logging flushes asynchronously; wait until the + // subsequent executions have landed in `mz_statement_execution_history`. + Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let mseh_count: i64 = mz_client + .query_one( + &format!( + "SELECT count(*) + FROM mz_internal.mz_statement_execution_history + WHERE application_name = '{APPLICATION_NAME}'" + ), + &[], + ) + .unwrap() + .get(0); + if mseh_count >= NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { + Ok(()) + } else { + Err(format!( + "only {mseh_count} of {NUM_SUBSEQUENT_STATEMENT_EXECUTIONS} executions logged so far" + )) + } + }) + .expect("subsequent executions never appeared in mz_statement_execution_history"); + + // The subsequent executions in mz_statement_execution_history should + // have a matching entry in mz_prepared_statement_history + let orphan_count: i64 = mz_client + .query_one( + &format!( + "SELECT count(*) + FROM mz_internal.mz_statement_execution_history mseh + LEFT JOIN mz_internal.mz_prepared_statement_history mpsh + ON mseh.prepared_statement_id = mpsh.id + WHERE mseh.application_name = '{APPLICATION_NAME}' + AND mpsh.id IS NULL" + ), + &[], + ) + .unwrap() + .get(0); + assert_eq!( + orphan_count, 0, + "{orphan_count} mz_statement_execution_history rows reference a \ + prepared_statement_id with no matching mz_prepared_statement_history row" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_subscribes() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + let cancel_token = client.cancel_token(); + + // This should finish + client + .execute( + "SUBSCRIBE TO (SELECT * FROM generate_series(1, 10001))", + &[], + ) + .unwrap(); + + let handle = thread::spawn(move || { + client.execute("CREATE TABLE t (x int)", &[]).unwrap(); + // This should not finish until it's canceled. + let _ = client.execute("SUBSCRIBE TO (SELECT * FROM t)", &[]); + }); + + while !handle.is_finished() { + thread::sleep(Duration::from_secs(1)); + cancel_token.cancel_query(postgres::NoTls).unwrap(); + } + handle.join().unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mpsh.prepared_at, + mseh.execution_strategy + FROM mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + mz_internal.mz_sql_text AS mst + ON mpsh.sql_hash = mst.sql_hash + WHERE + mst.sql ~~ 'SUBSCRIBE%' AND + mseh.finished_at IS NOT NULL + ORDER BY mseh.began_at"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = client.query(seh_query, &[]).unwrap(); + if sl.len() >= 2 { + break; + } + } + + assert_eq!(sl.len(), 2); + + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + prepared_at: DateTime, + execution_strategy: Option, + } + + let sl_subscribes = sl + .into_iter() + .map(|r| Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + prepared_at: r.get(4), + execution_strategy: r.get(5), + }) + .collect::>(); + for r in &sl_subscribes { + assert_eq!(r.sample_rate, 1.0); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + assert_none!(r.execution_strategy); + } + assert_eq!(sl_subscribes[0].finished_status, "success"); + assert_eq!(sl_subscribes[1].finished_status, "canceled"); +} + +/// Test that we are sampling approximately 50% of statements. +/// Relies on two assumptions: +/// (1) that the effective sampling rate for the session is 50%, +/// (2) that we are using the deterministic testing RNG. +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_sampling_inner( + server: TestServerWithStatementLoggingChecks, + mut client: postgres::Client, +) { + for i in 0..50 { + client.execute(&format!("SELECT {i}"), &[]).unwrap(); + + // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering + // of logged statements non-deterministic when we retrieve them below. + thread::sleep(Duration::from_millis(10)); + } + + // 23 randomly sampled out of 50 with 50% sampling. Seems legit! + let expected_sqls = [ + 2, 4, 5, 6, 9, 15, 17, 18, 19, 20, 21, 23, 24, 25, 31, 32, 33, 36, 37, 42, 46, + ] + .into_iter() + .map(|i| format!("SELECT {i}")) + .collect::>(); + + let mut internal_client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT mst.sql + FROM mz_internal.mz_statement_execution_history AS mseh + JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + mz_internal.mz_sql_text AS mst + ON mpsh.sql_hash = mst.sql_hash + WHERE mst.sql ~~ 'SELECT%' AND mst.sql !~~ '%mz_statement_execution_history%' + ORDER BY mseh.began_at ASC"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = internal_client.query(seh_query, &[]).unwrap(); + if sl.len() >= expected_sqls.len() { + break; + } + } + + let sqls: Vec = sl.into_iter().map(|r| r.get(0)).collect(); + assert_eq!(sqls, expected_sqls); + + // Verify the statement_logging_record_count metric correctly tracks sampled vs unsampled. + // With 50% sampling and deterministic RNG, exactly 21 of 50 statements should be sampled. + let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); + assert_eq!( + sampled_true, 21, + "expected 21 statements to be sampled with 50% rate and deterministic RNG" + ); + assert_eq!( + sampled_false, 29, + "expected 29 statements to not be sampled with 50% rate and deterministic RNG" + ); +} + +#[mz_ore::test] +fn test_statement_logging_sampling() { + let (server, client) = setup_statement_logging(1.0, 0.5, ""); + test_statement_logging_sampling_inner(server, client); +} + +/// Test that we are not allowed to set `statement_logging_sample_rate` +/// arbitrarily high, but that it is constrained by `statement_logging_max_sample_rate`. +#[mz_ore::test] +fn test_statement_logging_sampling_constrained() { + let (server, client) = setup_statement_logging(0.5, 1.0, ""); + test_statement_logging_sampling_inner(server, client); +} + +/// Test that the `mz_statement_logging_unsampled_bytes` metric tracks the total bytes +/// of SQL text that would have been logged if statement logging were fully enabled. +/// We set `sample_rate=0.0` so no statements are actually sampled/logged, but the +/// unsampled_bytes metric still gets incremented for every executed statement. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_unsampled_metrics() { + // Use sample_rate=0.0 so statements are not sampled, but unsampled_bytes metric is still tracked. + let (server, mut client) = setup_statement_logging(1.0, 0.0, ""); + + let batch_queries = [ + "SELECT 'Hello, world!';SELECT 1;;", + "SELECT 'Hello, world again!'", + ]; + let batch_total: usize = batch_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + let single_queries = ["SELECT 'foo'", "SELECT 'bar';;;"]; + let single_total: usize = single_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + let prepared_queries = ["SELECT 'baz';;;", "SELECT 'quux';"]; + let prepared_total: usize = prepared_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + + let named_prepared_inner = "SELECT 42"; + let named_prepared_outer = format!("PREPARE p AS {named_prepared_inner};EXECUTE p;"); + let named_prepared_outer_len = named_prepared_outer + .as_bytes() + .iter() + .filter(|&&ch| ch != b';') + .count(); + + for q in batch_queries { + client.batch_execute(q).unwrap(); + } + + for q in single_queries { + client.execute(q, &[]).unwrap(); + } + + for q in prepared_queries { + let s = client.prepare(q).unwrap(); + client.execute(&s, &[]).unwrap(); + } + + client.batch_execute(&named_prepared_outer).unwrap(); + + // This should NOT be logged, since we never actually execute it. + client.prepare("SELECT 'Hello, not counted!'").unwrap(); + + let expected_total = batch_total + single_total + prepared_total + named_prepared_outer_len; + let metric_value = server + .metrics_registry() + .gather() + .into_iter() + .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") + .unwrap() + .take_metric()[0] + .get_counter() + .value(); + let metric_value = usize::cast_from(u64::try_cast_from(metric_value).unwrap()); + assert_eq!(expected_total, metric_value); + + // Also verify that statement_logging_record_count shows all statements as not sampled + // (since we're using 0% sample rate). + let (sampled_true, _sampled_false) = get_statement_logging_record_counts(&server); + assert_eq!( + sampled_true, 0, + "no statements should be sampled with 0% sample rate" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_enable_internal_statement_logging() { + let (server, mut client) = setup_statement_logging_core( + 1.0, + 1.0, + "", + test_util::TestHarness::default().with_system_parameter_default( + "enable_internal_statement_logging".to_string(), + "true".to_string(), + ), + ); + + client.execute("SELECT 1", &[]).unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let num_mz_system_statements = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + count(*) +FROM mz_internal.mz_prepared_statement_history mpsh +JOIN mz_internal.mz_session_history USING (session_id) +WHERE authenticated_user='mz_system'", + &[], + ) + .unwrap(); + + let count: i64 = sl_results[0].get(0); + + if count > 0 { Ok(count) } else { Err(()) } + }) + .expect("at least some statements from mz_system should have been logged"); + + assert!( + num_mz_system_statements > 0, + "statements executed by mz_system should have been logged" + ); +} + +#[mz_ore::test] +// Test that statement logging +// doesn't cause a crash with subscribes over web sockets, +// which was previously happening (in staging) due to us +// dropping the `ExecuteContext` on the floor in that case. +fn test_statement_logging_ws_subscribe_no_crash() { + let (server, _client) = setup_statement_logging(1.0, 1.0, ""); + + // Create our WebSocket. + let ws_url = server.server.ws_addr(); + let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap(); + test_util::auth_with_ws(&mut ws, Default::default()).unwrap(); + + let query = "SUBSCRIBE (SELECT 1)"; + let json = format!("{{\"query\":\"{query}\"}}"); + let json: serde_json::Value = serde_json::from_str(&json).unwrap(); + ws.send(Message::Text(json.to_string().into())).unwrap(); + + // Give the server time to crash, if it's going to. + std::thread::sleep(Duration::from_secs(1)) +} + +/// `finished_at` must record when execution finished, not when the coordinator +/// got around to the end event. A statement the session task retires itself +/// reports its own end timestamp, so a busy coordinator cannot inflate it. +/// +/// The coordinator is stalled inside `Catalog::transact` while the measured +/// statement runs. That failpoint sits ahead of every catalog mutation, so the +/// catalog revision does not move while it sleeps and the session's cached +/// snapshot stays valid, which is what keeps the measured statement from +/// needing the coordinator at all. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_finished_at_excludes_coordinator_queue() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + + // Populate this session's catalog snapshot cache, so the measured statement + // below needs nothing from the stalled coordinator. + client.execute("SELECT 1", &[]).unwrap(); + + let mut ddl = server.connect_internal(postgres::NoTls).unwrap(); + fail::cfg("catalog_transact", "sleep(3000)").unwrap(); + let stall = thread::spawn(move || { + let _ = ddl.batch_execute("CREATE TABLE stalls_the_coordinator (x int)"); + }); + // Give the DDL time to reach the failpoint before we measure. + thread::sleep(Duration::from_millis(500)); + + // Logged timestamps are epoch milliseconds and truncate, so floor the lower + // bound and ceil the upper one rather than comparing against the + // sub-millisecond instants `Utc::now` reports. + let began_bound = Utc::now().timestamp_millis(); + client + .execute("SELECT 1 AS finished_at_probe", &[]) + .unwrap(); + let finished_bound = Utc::now().timestamp_millis() + 1; + + stall.join().unwrap(); + fail::remove("catalog_transact"); + + let mut internal = server.connect_internal(postgres::NoTls).unwrap(); + let query = " + SELECT mseh.began_at, mseh.finished_at + FROM mz_internal.mz_statement_execution_history AS mseh + JOIN mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN mz_internal.mz_sql_text AS mst ON mpsh.sql_hash = mst.sql_hash + WHERE mst.sql ~~ '%finished_at_probe%' AND mseh.finished_at IS NOT NULL"; + + let mut rows = Vec::new(); + for _ in 0..30 { + rows = internal.query(query, &[]).unwrap(); + if !rows.is_empty() { + break; + } + thread::sleep(Duration::from_secs(1)); + } + assert_eq!(rows.len(), 1, "expected exactly one logged probe statement"); + + let began_at: DateTime = rows[0].get(0); + let finished_at: DateTime = rows[0].get(1); + assert!( + began_at.timestamp_millis() >= began_bound, + "began_at precedes the client-observed start of the statement by {} ms", + began_bound - began_at.timestamp_millis() + ); + assert!( + finished_at.timestamp_millis() <= finished_bound, + "finished_at is {} ms past the client-observed end of the statement, so it \ + charges the statement for time its end event spent queued for the coordinator", + finished_at.timestamp_millis() - finished_bound + ); +} From 72a8db1fd695088624eab0c34d8e6b3a4e10a035 Mon Sep 17 00:00:00 2001 From: Dennis Felsing Date: Wed, 5 Aug 2026 15:08:52 +0200 Subject: [PATCH 6/6] parallel-workload: broaden SQL-surface and action coverage (#37590) Substantially broadens what the parallel-workload stress test exercises, from a doc and surface-area audit, so nightly runs cover far more of the SQL and DDL surface, at larger scale, and find more product bugs. Every bug this coverage surfaced is worked around (the triggering action is disabled or narrowed, each with a TODO naming the tracked issue) so the nightly stays green until the bug is fixed. ## New coverage Adds ~34 actions and query-generation breadth: - **SELECT surface**: GROUP BY, DISTINCT / DISTINCT ON, outer and lateral joins, subqueries, CTEs, WITH MUTUALLY RECURSIVE, table functions, temporal filters, and the aggregate families. Generated `WHERE` conjuncts are parenthesized so an `OR` inside one part stays a single top-level conjunct. Boundary literals (float4 `NaN`/`Infinity`) are injected as expression leaves to stress overflow and eval paths. - **DDL**: `CREATE OR REPLACE`, `DROP ... CASCADE`, `CREATE TYPE`, network policies, `ALTER CLUSTER SET`, `ALTER ROLE SET/RESET` session defaults, and broader `COMMENT`. - **Transactions**: parameterized prepared statements, bounded staleness, read-only and DDL transactions. - **Introspection**: system-catalog reads, `EXPLAIN ANALYZE`, and the full `EXPLAIN` matrix. - **Sources and sinks**: load generator sources, sink options, Iceberg `MODE APPEND`, and `CREATE REPLACEMENT MATERIALIZED VIEW` with both the `APPLY REPLACEMENT` and discard paths. - **Other**: RBAC, connection and secret lifecycle, `SHOW`. ## Incident 1136 coverage: replacement lifecycle and a sealed-shard oracle Incident 1136 (a dropped replacement finalizing its live target's shard after an envd restart) needed `CREATE REPLACEMENT` and the resolving `DROP` to happen on different connections across a restart. The old `ReplaceMaterializedViewAction` resolved a replacement on the connection that created it, so it structurally could not hit this. The lifecycle is now split into create/apply/drop actions that pick live replacements from `mz_internal.mz_replacements`, so replacements created by other workers or leaked across kills and 0dt deploys are found and resolved too. Verified to go red in one kill cycle on a build without the #37705 fix. Detection does not rely on the (much later) tombstone panic but on a client-side data-loss oracle, `SealedCollectionCheckAction`: a live user table or materialized view whose write frontier is the empty antichain has had its shard sealed out from under it. Legitimate seals are classified away: REFRESH AT views and bounded (UP TO) load generators seal when they complete, sealing is transitive through readers (a catalog dependency walk), and a view whose optimized plan reads no live storage collections (constant folding, e.g. a join `ON false`, drops inputs the catalog still records) completes on hydration, whatever its refresh schedule says (checked via `EXPLAIN OPTIMIZED PLAN AS JSON`). Candidates are re-checked after 5s since the catalog and `mz_frontiers` are not updated atomically. A guaranteed final oracle pass runs at the end of every run, including on the wedged-worker exit path, so damage from a run's last seconds cannot escape. ## Scale Raises the per-object caps (databases, schemas, tables, views, indexes, roles, sources, sinks, types, network policies, and columns) by roughly 10x, so a run holds many more concurrent objects and exercises deeper dependency graphs. ## Product bugs surfaced The new coverage found several pre-existing bugs. Each is worked around with a TODO pointing at the tracked issue so the action can be re-enabled once fixed: - **[CLO-143](https://linear.app/materializeinc/issue/CLO-143)**: `ALTER NETWORK POLICY` cannot resolve quoted names. - **[SQL-517](https://linear.app/materializeinc/issue/SQL-517)**: coordinator panic when a connection re-alter (via `ALTER CONNECTION`, `GRANT`/`REVOKE`, or `ALTER OWNER`) reshapes a dependent sink's export connection while its key indices are momentarily out of sync. - **[SQL-518](https://linear.app/materializeinc/issue/SQL-518)**: `get_database` panic on a concurrent `DROP DATABASE CASCADE` during a staged create's finish. - **[SQL-519](https://linear.app/materializeinc/issue/SQL-519)**: `EXPLAIN FILTER PUSHDOWN` missing-collection panic when a referenced index/MV is dropped between planning and read-hold acquisition. - **[SQL-520](https://linear.app/materializeinc/issue/SQL-520)**: frontend-peek read-hold vs compaction soft-panic (`enable_frontend_peek_sequencing` pinned off, peeks run the classic path). - **[SS-344](https://linear.app/materializeinc/issue/SS-344)**: `ALTER SINK SET FROM` hang. - **[SS-345](https://linear.app/materializeinc/issue/SS-345)**: `COPY FROM` CSV cannot decode a large-year date that `COPY TO` CSV wrote. - **[SS-347](https://linear.app/materializeinc/issue/SS-347)**: `VALIDATE CONNECTION` transiently reads an empty password during a concurrent `ALTER SECRET` rotation. - **[CLU-169](https://linear.app/materializeinc/issue/CLU-169)**: `SUBSCRIBE ... UP TO` optimizer soft-panic when the object's `as_of` has advanced to `MAX` (e.g. a finished bounded load generator). - **[SQL-595](https://linear.app/materializeinc/issue/SQL-595)**: `CREATE REPLACEMENT MATERIALIZED VIEW` accepts a sealed target (e.g. a constant-folded view, which seals on hydration whatever its refresh schedule says) that `APPLY REPLACEMENT` then refuses forever, leaving an unappliable pending replacement. Found by the sealed-shard oracle. Known negative-accumulation errors (database-issues#9308) from `DELETE ... USING` and `repeat_row` are tolerated in the scenarios that generate them and stay strict elsewhere, so a genuine finding still fails. Coverage for two further directions is landed but left disabled, ready to flip on once the bugs are fixed: a dependency-consistency oracle ([SQL-521](https://linear.app/materializeinc/issue/SQL-521)) and a source read-hold sweep across restart ([SS-346](https://linear.app/materializeinc/issue/SS-346), [PER-49](https://linear.app/materializeinc/issue/PER-49)). Closes: [QAR-139](https://linear.app/materializeinc/issue/QAR-139) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- .../materialize/parallel_workload/action.py | 2786 +++++++++++++++-- .../materialize/parallel_workload/column.py | 74 +- .../materialize/parallel_workload/database.py | 338 +- .../materialize/parallel_workload/executor.py | 66 +- .../parallel_workload/expression.py | 31 +- .../negative_accumulation_errors.py | 13 + .../parallel_workload/parallel_workload.py | 204 +- .../materialize/parallel_workload/settings.py | 24 + 8 files changed, 3191 insertions(+), 345 deletions(-) diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 44d6f9887ab6c..d314cc10db9c9 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -28,10 +28,22 @@ import materialize.parallel_workload.column from materialize.data_ingest.data_type import ( DATA_TYPES, + DATA_TYPES_FOR_COLUMNS, NUMBER_TYPES, + RANGE_TYPES, + UUID, Boolean, + Bytea, + Char, + IntArray, + IntList, + Jsonb, + Oid, Text, TextTextMap, + Timestamp, + TimestampTz, + VarChar, ) from materialize.data_ingest.query_error import QueryError from materialize.data_ingest.row import Operation @@ -43,6 +55,8 @@ Materialized, ) from materialize.mzcompose.services.minio import minio_blob_uri +from materialize.mzcompose.services.mysql import MySql +from materialize.mzcompose.services.sql_server import SqlServer from materialize.parallel_workload.database import ( DB, MAX_CLUSTER_REPLICAS, @@ -53,13 +67,16 @@ MAX_INDEXES, MAX_KAFKA_SINKS, MAX_KAFKA_SOURCES, + MAX_LOADGEN_SOURCES, MAX_MYSQL_SOURCES, + MAX_NETWORK_POLICIES, MAX_POSTGRES_SOURCES, MAX_ROLES, MAX_ROWS, MAX_SCHEMAS, MAX_SQL_SERVER_SOURCES, MAX_TABLES, + MAX_TYPES, MAX_VIEWS, MAX_WEBHOOK_SOURCES, Cluster, @@ -71,14 +88,18 @@ Index, KafkaSink, KafkaSource, + LoadGeneratorSource, + MultiLoadGeneratorSource, MySqlSource, MzTempSchema, + NetworkPolicy, PostgresSource, Role, S3Object, Schema, SqlServerSource, Table, + Type, View, WebhookSource, ) @@ -139,6 +160,67 @@ def ws_connect(ws: websocket.WebSocket, host, port, user: str) -> tuple[int, int return (ws_conn_id, ws_secret_key) +def untrack_objects_in_schemas(exe: Executor, schemas: set[Schema]) -> None: + """Remove every tracked object living in one of the given schemas. + + Called after a CASCADE drop of a schema or database. Source executor + connections are closed. Cross-schema dependents are cascade-dropped + server-side but not pruned here, they later surface as "does not exist", + which is why the CASCADE actions only run in DDL complexity.""" + with exe.db.lock: + exe.db.tables[:] = [t for t in exe.db.tables if t.schema not in schemas] + exe.db.views[:] = [v for v in exe.db.views if v.schema not in schemas] + # NOTE: Removed in place, like the lists above. DropIndexAction + # discards from this set holding only the index's own lock, so + # rebinding the attribute would resurrect an index that was dropped + # between building the replacement set and assigning it. + exe.db.indexes.difference_update( + {i for i in exe.db.indexes if i.schema in schemas} + ) + # Close the executor connections of the ingestion sources being + # untracked so they don't leak. + connected_sources = ( + exe.db.kafka_sources + + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + for src in connected_sources: + if src.schema in schemas: + try: + src.executor.mz_conn.close() + except Exception: + pass + exe.db.kafka_sources[:] = [ + s for s in exe.db.kafka_sources if s.schema not in schemas + ] + exe.db.postgres_sources[:] = [ + s for s in exe.db.postgres_sources if s.schema not in schemas + ] + exe.db.mysql_sources[:] = [ + s for s in exe.db.mysql_sources if s.schema not in schemas + ] + exe.db.sql_server_sources[:] = [ + s for s in exe.db.sql_server_sources if s.schema not in schemas + ] + exe.db.loadgen_sources[:] = [ + s for s in exe.db.loadgen_sources if s.schema not in schemas + ] + exe.db.multi_loadgen_sources[:] = [ + s for s in exe.db.multi_loadgen_sources if s.schema not in schemas + ] + exe.db.webhook_sources[:] = [ + s for s in exe.db.webhook_sources if s.schema not in schemas + ] + exe.db.kafka_sinks[:] = [ + s for s in exe.db.kafka_sinks if s.schema not in schemas + ] + exe.db.iceberg_sinks[:] = [ + s for s in exe.db.iceberg_sinks if s.schema not in schemas + ] + exe.db.types[:] = [t for t in exe.db.types if t.schema not in schemas] + + # TODO: CASCADE in DROPs, keep track of what will be deleted class Action: rng: random.Random @@ -201,7 +283,14 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # catalog apply, which surfaces as an unexpected failure. "non-temporary items cannot depend on temporary item", "is not of expected type SqlColumnType", # TODO: Remove when SQL-566 is fixed + # generate_select_query occasionally emits a WITH MUTUALLY + # RECURSIVE body with ERROR AT RECURSION LIMIT, which errors on + # purpose when the iteration outruns the limit. Only our generated + # queries produce this, so ignoring it globally is safe. + "exceeded the recursion limit", ] + if exe.statement_timeout_set: + result.append("canceling statement due to statement timeout") if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): result.extend( [ @@ -214,6 +303,12 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # matching "unknown cluster replica size" errors. "unknown cluster '", "unknown schema", # schema was dropped + # The database was dropped concurrently: DropDatabaseAction + # can land on an emptied database, and + # DropDatabaseCascadeAction (disabled until SQL-518 is + # fixed) drops non-empty ones. + "unknown database", + "invalid database", # CREATE SCHEMA wording for a vanished database "the transaction's active cluster has been dropped", # cluster was dropped "was removed", # dependency was removed, started with moving optimization off main thread, see database-issues#7285 "real-time source dropped before ingesting the upstream system's visible frontier", # Expected, see https://buildkite.com/materialize/nightly/builds/9399#0191be17-1f4c-4321-9b51-edc4b08b71c5 @@ -285,21 +380,83 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: ) if exe.db.scenario == Scenario.Rename: result.extend(["unknown schema", "ambiguous reference to schema name"]) + # RepeatRow is the only scenario that deliberately produces negative + # multiplicities (`repeat_row` with a negative count), and the + # corruption is observed by any later reader of the affected + # collection, not just by the statement that created it. So the whole + # class is tolerated for the entire scenario, not per action. Every + # other scenario stays strict, a negative-accumulation error there is a + # genuine finding. The central list lives in + # `negative_accumulation_errors.py`. if exe.db.scenario == Scenario.RepeatRow: - # Views that use `repeat_row` with a negative count can surface - # negative-accumulation errors both at DDL time (constant folding) - # and at read time (various operators checking multiplicities). - # The central list lives in `negative_accumulation_errors.py` so - # we can update it in one place when the error messages change. result.extend(NEGATIVE_ACCUMULATION_ERRORS) if materialize.parallel_workload.column.NAUGHTY_IDENTIFIERS: result.extend(["identifier length exceeds 255 bytes"]) return result + # MIN/MAX exist for every type except these (bytea, jsonb, map, list, + # array, uuid, oid, and the range types). + _MINMAX_EXCLUDED = ( + Bytea, + Jsonb, + TextTextMap, + IntList, + IntArray, + UUID, + Oid, + ) + tuple(RANGE_TYPES) + + def aggregate_fns(self, column: Column) -> list[str]: + """Aggregate function templates valid for the column's type. + + Used both in window position (OVER ..) and in GROUP BY position. The + collection aggregates (array_agg/list_agg/jsonb_agg/string_agg) + exercise the "collection" reduce rendering, distinct from the + accumulable sum/count path. Type exclusions are empirically derived, + e.g. array_agg rejects char and cannot nest map/list/array.""" + dt = column.data_type + fns = ["COUNT({})"] + if dt in NUMBER_TYPES: + fns.extend( + [ + "SUM({})", + "AVG({})", + "STDDEV({})", + "STDDEV_POP({})", + "STDDEV_SAMP({})", + "VAR_SAMP({})", + "VAR_POP({})", + ] + ) + if dt == Boolean: + fns.extend(["BOOL_AND({})", "BOOL_OR({})"]) + if dt not in self._MINMAX_EXCLUDED: + fns.extend(["MAX({})", "MIN({})"]) + # Collection aggregates. + fns.append("jsonb_agg({})") + if dt != Char: + fns.append("list_agg({})") + if dt not in (Char, TextTextMap, IntList, IntArray): + fns.append("array_agg({})") + if dt in (Text, VarChar, Char): + fns.append("string_agg({}, ',')") + return fns + def generate_select_query(self, exe: Executor, expr_kind: ExprKind) -> str: - obj = self.rng.choice(exe.db.db_objects()) + objects = exe.db.db_objects() + if not objects: + # A concurrent CASCADE drop removed every object. Nothing to query. + return "SELECT 1" + obj = self.rng.choice(objects) column = self.rng.choice(obj.columns) - obj2 = self.rng.choice(exe.db.db_objects_without_views()) + # db_objects_without_views() can momentarily be empty if a CASCADE drop + # removed the last table/source/MV, leaving only plain views. Fall back + # to obj (only used for an optional join, which is skipped for views). + # The list is taken once and reused below: recomputing it can return an + # empty list even though it was non-empty here, and rng.choice on that + # raises IndexError, which fails the whole run. + objs_without_views = exe.db.db_objects_without_views() + obj2 = self.rng.choice(objs_without_views) if objs_without_views else obj obj_name = str(obj) obj2_name = str(obj2) columns = [ @@ -312,105 +469,219 @@ def generate_select_query(self, exe: Executor, expr_kind: ExprKind) -> str: all_columns = list(obj.columns) + list(obj2.columns) if join else obj.columns - column_types = [] - if self.rng.random() < 0.9: - column_types = [ - self.rng.choice(list(DATA_TYPES)) - for i in range(self.rng.randint(1, 10)) + # Self-contained iterative dataflow, cross joined with a real object. + # The iteration bound and the recursion limit are drawn independently, + # so the limit sometimes fires, exercising the RETURN/ERROR AT paths + # on purpose ("exceeded the recursion limit" is an expected error). + if self.rng.random() < 0.05: + limit_kind = self.rng.choice(["RETURN AT", "ERROR AT"]) + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind + ) + return ( + f"WITH MUTUALLY RECURSIVE ({limit_kind} RECURSION LIMIT {self.rng.randint(2, 200)}) " + f"cnt (i int8) AS (" + f"SELECT 1 UNION ALL SELECT i + 1 FROM cnt WHERE i < {self.rng.randint(1, 100)}" + f") SELECT i, {expr} FROM cnt, {obj_name}" + f" LIMIT {self.rng.randint(0, 100)}" + ) + + # Explicit LATERAL derived table correlated on a matching-type column. + # Exercises correlated-subquery decorrelation. + if objs_without_views and self.rng.random() < 0.08: + outer_col = self.rng.choice(obj.columns) + inner_obj = self.rng.choice(objs_without_views) + match_cols = [ + c + for c in inner_obj.columns + if c.data_type == outer_col.data_type and c.data_type != TextTextMap ] - expressions = ", ".join( - [ - expression( - column_type, - all_columns, - self.rng, - expr_kind, - ) - for column_type in column_types - ] + corr = "" + if match_cols: + corr = f" WHERE {self.rng.choice(match_cols)} = {outer_col}" + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind + ) + return ( + f"SELECT {expr} FROM {obj_name}, LATERAL (" + f"SELECT count(*) AS cnt FROM {inner_obj}{corr}" + f") AS lat LIMIT {self.rng.randint(0, 100)}" + ) + + # Table function in FROM. unnest of a column reference is an + # implicitly lateral call. + if self.rng.random() < 0.1: + int_list_columns = [c for c in obj.columns if c.data_type == IntList] + if int_list_columns and self.rng.choice([True, False]): + func = f"unnest({self.rng.choice(int_list_columns)})" + else: + func = f"generate_series(1, {self.rng.randint(1, 100)})" + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind ) if self.rng.choice([True, False]): - column_types = [] - column1 = self.rng.choice(all_columns) - column2 = self.rng.choice(all_columns) - column3 = self.rng.choice(all_columns) - fns = [ - "COUNT({})", - # "LIST_AGG({})", - # "JSONB_AGG({})", - ] - # if column1.data_type == Text: - # fns.extend(["STRING_AGG({}, ',')"]) - # if column1.data_type not in [TextTextMap, IntArray, IntList]: - # fns.extend(["ARRAY_AGG({})"]) - if column1.data_type in NUMBER_TYPES: - fns.extend( - [ - "SUM({})", - "AVG({})", - "MAX({})", - "MIN({})", - "STDDEV({})", - "STDDEV_POP({})", - "STDDEV_SAMP({})", - "VAR_SAMP({})", - "VAR_POP({})", - ] - ) - elif column1.data_type == Boolean: - fns.extend(["BOOL_AND({})", "BOOL_OR({})"]) - window_fn = self.rng.choice(fns) - expressions += f", {window_fn.format(column1)} OVER (PARTITION BY {column2} ORDER BY {column3})" - else: - expressions = "*" + tf = f"{func} WITH ORDINALITY AS tf(x, ord)" + else: + tf = f"{func} AS tf(x)" + return ( + f"SELECT tf.x, {expr} FROM {obj_name}, {tf}" + f" LIMIT {self.rng.randint(0, 100)}" + ) - query = f"SELECT {expressions} FROM {obj_name}" + star = self.rng.random() >= 0.9 + exprs: list[tuple[type, str]] = [] + if not star: + for i in range(self.rng.randint(1, 10)): + dt = self.rng.choice(list(DATA_TYPES)) + exprs.append((dt, expression(dt, all_columns, self.rng, expr_kind))) if join: + join_kind = self.rng.choice( + ["JOIN", "JOIN", "JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN"] + ) column2 = self.rng.choice(columns) - query += f" JOIN {obj2_name} ON {column} = {column2}" + join_clause = f" {join_kind} {obj2_name} ON {column} = {column2}" + if self.rng.random() < 0.2: + join_clause += ( + f" AND {expression(Boolean, all_columns, self.rng, expr_kind)}" + ) + else: + join_clause = "" - if self.rng.choice([True, False]): - query += f" WHERE {expression(Boolean, all_columns, self.rng, expr_kind)}" - - if bool(column_types) and self.rng.choice([True, False]): - obj3 = self.rng.choice(exe.db.db_objects()) - obj3_name = str(obj3) - column3 = self.rng.choice(obj3.columns) - obj4 = self.rng.choice(exe.db.db_objects()) - obj4_name = str(obj4) - columns_union = [ - c - for c in obj4.columns - if c.data_type == column3.data_type and c.data_type != TextTextMap + def where_clause() -> str: + parts = [] + if self.rng.choice([True, False]): + parts.append(expression(Boolean, all_columns, self.rng, expr_kind)) + if objs_without_views and self.rng.random() < 0.2: + obj3 = self.rng.choice(objs_without_views) + sub_columns = [ + c + for c in obj3.columns + if c.data_type == column.data_type and c.data_type != TextTextMap + ] + sub_kind = self.rng.choice(["exists", "not exists", "in", "scalar"]) + if sub_kind in ("exists", "not exists"): + cond = expression(Boolean, obj3.columns, self.rng, expr_kind) + parts.append( + f"{'NOT ' if sub_kind == 'not exists' else ''}EXISTS (SELECT 1 FROM {obj3} WHERE {cond})" + ) + elif sub_columns: + sub_column = self.rng.choice(sub_columns) + if sub_kind == "in": + parts.append(f"{column} IN (SELECT {sub_column} FROM {obj3})") + else: + parts.append( + f"{column} = (SELECT {sub_column} FROM {obj3} LIMIT 1)" + ) + # Deliberate temporal filter, mz_now() has to be a top-level + # conjunct for the temporal filter machinery to apply. + ts_columns = [ + c for c in all_columns if c.data_type in (Timestamp, TimestampTz) ] - join_union = ( - obj3_name != obj4_name and obj3 not in exe.db.views and columns_union + if ts_columns and expr_kind != ExprKind.WRITE and self.rng.random() < 0.2: + ts_column = self.rng.choice(ts_columns) + op = self.rng.choice(["<=", ">="]) + parts.append( + f"mz_now() {op} {ts_column} + INTERVAL '{self.rng.randint(0, 100000)} seconds'" + ) + if not parts: + return "" + # Generated Boolean expressions can contain OR, so parenthesize + # each part to keep every part as a top-level conjunct. + return " WHERE " + " AND ".join(f"({part})" for part in parts) + + distinct_on_expr = None + group_by = bool(exprs) and self.rng.random() < 0.2 + if group_by: + num_group = self.rng.randint(1, len(exprs)) + select_list = [expr for _, expr in exprs[:num_group]] + # Only the group keys come from exprs. The aggregates that follow + # get None, their result type is not exprs[i]'s. + select_types: list[type | None] = [dt for dt, _ in exprs[:num_group]] + for i in range(self.rng.randint(1, 3)): + agg_column = self.rng.choice(all_columns) + fn = self.rng.choice(self.aggregate_fns(agg_column)) + select_list.append(fn.format(agg_column)) + select_types.append(None) + group_clause = ( + f" GROUP BY {', '.join(str(i + 1) for i in range(num_group))}" ) - all_columns_union = ( - list(obj3.columns) + list(obj4.columns) if join_union else obj3.columns + # Clause order is fixed: HAVING before OPTIONS. + if self.rng.random() < 0.3: + group_clause += f" HAVING count(*) >= {self.rng.randint(0, 10)}" + if self.rng.random() < 0.3: + group_clause += f" OPTIONS (AGGREGATE INPUT GROUP SIZE = {self.rng.randint(1, 1000)})" + distinct_clause = "" + else: + select_list = [expr for _, expr in exprs] if exprs else ["*"] + # `*` expands to every column, so the output types are the columns'. + select_types = ( + [dt for dt, _ in exprs] if exprs else [c.data_type for c in all_columns] ) - expressions3 = ", ".join( - [ - expression( - column_type, - all_columns_union, - self.rng, - expr_kind, - ) - for column_type in column_types - ] + group_clause = "" + distinct_clause = "" + distinct_on_expr = None + if exprs and self.rng.random() < 0.15: + if exprs[0][0] == TextTextMap or self.rng.choice([True, False]): + distinct_clause = "DISTINCT " + else: + # The DISTINCT ON expressions must be a prefix of the + # ORDER BY expressions. The leading ORDER BY item repeats + # this expression verbatim (not its ordinal) so the match + # holds even for bare literals. + distinct_on_expr = exprs[0][1] + distinct_clause = f"DISTINCT ON ({distinct_on_expr}) " + if exprs and self.rng.choice([True, False]): + column1 = self.rng.choice(all_columns) + column2 = self.rng.choice(all_columns) + column3 = self.rng.choice(all_columns) + window_fn = self.rng.choice(self.aggregate_fns(column1)) + select_list.append( + f"{window_fn.format(column1)} OVER (PARTITION BY {column2} ORDER BY {column3})" + ) + select_types.append(None) + + expressions = ", ".join(select_list) + query = f"SELECT {distinct_clause}{expressions} FROM {obj_name}" + query += join_clause + query += where_clause() + query += group_clause + + if not distinct_clause and self.rng.choice([True, False]): + set_op = self.rng.choice( + ["UNION ALL"] * 4 + + ["UNION", "INTERSECT", "INTERSECT ALL", "EXCEPT", "EXCEPT ALL"] ) - query += f" UNION ALL SELECT {expressions3} FROM {obj3_name}" - - if join_union: - column4 = self.rng.choice(columns_union) - query += f" JOIN {obj4_name} ON {column3} = {column4}" - - if self.rng.choice([True, False]): - query += f" WHERE {expression(Boolean, all_columns_union, self.rng, expr_kind)}" + query += f" {set_op} SELECT {expressions} FROM {obj_name}" + query += join_clause + query += where_clause() + query += group_clause + + # Ordinals keep ORDER BY valid across set operations. Map-typed + # output columns are skipped, same as in the join column selection. + # NOTE: select_types is indexed by output position, which is not + # exprs' indexing once a GROUP BY appends aggregates after the keys. + orderable = [i + 1 for i, dt in enumerate(select_types) if dt != TextTextMap] + if distinct_on_expr is not None: + order_by = [distinct_on_expr] + for i in self.rng.sample(orderable, self.rng.randint(0, len(orderable))): + if i != 1: + order_by.append(str(i)) + query += f" ORDER BY {', '.join(order_by)}" + elif exprs and orderable and self.rng.random() < 0.3: + order_by = [] + for i in self.rng.sample(orderable, self.rng.randint(1, len(orderable))): + direction = self.rng.choice(["", " ASC", " DESC"]) + nulls = self.rng.choice(["", " NULLS FIRST", " NULLS LAST"]) + order_by.append(f"{i}{direction}{nulls}") + query += f" ORDER BY {', '.join(order_by)}" query += f" LIMIT {self.rng.randint(0, 100)}" + if self.rng.random() < 0.2: + query += f" OFFSET {self.rng.randint(0, 100)}" + + if self.rng.random() < 0.15: + query = f"WITH cte0 AS ({query}) SELECT * FROM cte0" return query def exe_prepared(self, query: str, stmt_name: str, exe: Executor) -> None: @@ -460,7 +731,13 @@ def run(self, exe: Executor) -> bool: if self.rng.choice([True, False]) else exe.commit(http=Http.NO) ) + # NOTE: A bounded SUBSCRIBE (UP TO) over an object whose as_of has + # advanced to the end of time (e.g. a finished bounded load generator + # source) soft-panics the optimizer + # (https://linear.app/materializeinc/issue/CLU-169). Left out until that + # is fixed. AS OF AT LEAST 0 below is safe (empty until). query = "SUBSCRIBE " + envelope_used = False if self.rng.choice([True, False]): obj = self.rng.choice(exe.db.db_objects()) query += f"{obj}" @@ -470,12 +747,21 @@ def run(self, exe: Executor) -> bool: columns = self.rng.sample(obj.columns, len(obj.columns)) key = ", ".join(column.name(True) for column in columns) query += f" ENVELOPE {envelope} (KEY ({key}))" - - if self.rng.choice([True, False]): - query += " WITH (SNAPSHOT = false)" + envelope_used = True else: query += f"({self.generate_select_query(exe, ExprKind.MATERIALIZABLE)})" + options = [] + if self.rng.choice([True, False]): + options.append("SNAPSHOT = false") + if not envelope_used and self.rng.random() < 0.3: + options.append("PROGRESS") + if options: + query += f" WITH ({', '.join(options)})" + if self.rng.random() < 0.2: + # AT LEAST always plans, no matter how far the since advanced. + query += " AS OF AT LEAST 0" + exe.execute(f"DECLARE c{self.i} CURSOR FOR {query}", http=Http.NO) while True: rows = self.rng.choice(["ALL", self.rng.randrange(1000)]) @@ -690,6 +976,19 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # roundtrip can produce NULLs for NOT NULL columns. "violates not-null constraint", "timeout: error trying to connect", + # COPY TO CSV writes a large-year date that COPY FROM CSV then + # fails to parse back. + # See https://linear.app/materializeinc/issue/SS-345 + "expected_dur_like_tokens can only be called with", + # TODO: Reenable when SS-361 is fixed. A COPY FROM without a + # column list plans every target column as its DEFAULT, and + # that projection only acts as identity by accident. Expression + # memoization collapses two identical default literals, so any + # table with two same-type nullable columns (their defaults are + # identical typed nulls) decodes a shifted row. Shifts that are + # type-incompatible fail here, type-compatible ones are + # inserted silently. + "failed to decode Row from a record batch", ] ) if exe.db.complexity == Complexity.DDL: @@ -912,6 +1211,17 @@ def run(self, exe: Executor) -> bool: class InsertReturningAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + # The RETURNING expressions re-render the fully-qualified table and + # column names, so a concurrent schema or table rename landing between + # the INSERT target and the RETURNING clause leaves the two referring to + # different names and the query fails with "does not exist". Same class + # as UpdateAction/DeleteAction. + if exe.db.complexity == Complexity.DDL or exe.db.scenario == Scenario.Rename: + result.extend(["does not exist"]) + return result + def run(self, exe: Executor) -> bool: table = None if exe.insert_table is not None: @@ -975,6 +1285,39 @@ def run(self, exe: Executor) -> bool: return True +class CopyToStdoutAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + # A prior statement in the read transaction routed it to + # mz_catalog_server, where this COPY of a user object is + # rejected. + 'is not allowed from the "mz_catalog_server" cluster', + # BINARY format is deliberately kept in the mix even though + # some types (e.g. map) have no binary output function. + "no binary output function available for type", + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend( + [ + "does not exist", + ] + ) + return result + + def run(self, exe: Executor) -> bool: + obj = self.rng.choice(exe.db.db_objects()) + query = f"COPY (SELECT * FROM {obj} LIMIT {self.rng.randint(0, 100)}) TO STDOUT" + if self.rng.choice([True, False]): + format = self.rng.choice(["TEXT", "CSV", "BINARY"]) + query += f" WITH (FORMAT {format})" + exe.copy_to_stdout(query) + return True + + class SourceInsertAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -1051,8 +1394,14 @@ def run(self, exe: Executor) -> bool: return False table = self.rng.choice(tables) - column2 = self.rng.choice(table.columns) - query = f"UPDATE {table} SET {column2.name(True)} = {expression(column2.data_type, table.columns, self.rng, kind=ExprKind.WRITE)} WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" + set_columns = self.rng.sample( + table.columns, self.rng.randint(1, len(table.columns)) + ) + set_clause = ", ".join( + f"{c.name(True)} = {expression(c.data_type, table.columns, self.rng, kind=ExprKind.WRITE)}" + for c in set_columns + ) + query = f"UPDATE {table} SET {set_clause} WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" if self.rng.choice([True, False]): self.stmt_id += 1 self.exe_prepared(query, f"update{self.stmt_id}", exe) @@ -1115,7 +1464,28 @@ def run(self, exe: Executor) -> bool: return False table = self.rng.choice(tables) query = f"DELETE FROM {table}" - if self.rng.random() < 0.95: + using_tables = [ + t + for t in exe.db.tables + if t != table and (not t.temp or t in exe.temp_objects) + ] + # TODO: Drop the RepeatRow gate once database-issues#9308 is fixed. + # DELETE .. USING lowers to a semijoin whose DistinctBy can leave the + # target table with a net-negative row, and every later reader of that + # table then surfaces the corruption. Tolerating that class outside + # RepeatRow would blind the DML and DDL complexities to every genuine + # negative-accumulation finding, so the variant only runs in the one + # scenario that already expects negative multiplicities. + if ( + using_tables + and exe.db.scenario == Scenario.RepeatRow + and self.rng.random() < 0.2 + ): + using_table = self.rng.choice(using_tables) + all_columns = list(table.columns) + list(using_table.columns) + query += f" USING {using_table}" + query += f" WHERE {expression(Boolean, all_columns, self.rng, kind=ExprKind.WRITE)}" + elif self.rng.random() < 0.95: query += f" WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" if self.rng.choice([True, False]): self.stmt_id += 1 @@ -1133,15 +1503,83 @@ def run(self, exe: Executor) -> bool: class CommentAction(Action): - def run(self, exe: Executor) -> bool: - table = self.rng.choice(exe.db.tables) + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + ] + super().errors_to_ignore(exe) - if self.rng.choice([True, False]): - column = self.rng.choice(table.columns) - query = f"COMMENT ON COLUMN {column} IS '{Text.random_value(self.rng)}'" - else: - query = f"COMMENT ON TABLE {table} IS '{Text.random_value(self.rng)}'" + def run(self, exe: Executor) -> bool: + # Only the snapshot needs the global lock. Rendering the candidate + # names is O(tracked objects) with a naughtify each, and holding the + # lock across that serializes every other worker behind this action. + with exe.db.lock: + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + indexes = list(exe.db.indexes) + schemas = list(exe.db.schemas) + dbs = list(exe.db.dbs) + clusters = list(exe.db.clusters) + roles = list(exe.db.roles) + + candidates: list[tuple[str, str]] = [] + for table in tables: + candidates.append(("TABLE", str(table))) + # Columns are only ever appended, never removed, so picking one + # off the snapshot cannot index past the end. + candidates.append(("COLUMN", str(self.rng.choice(table.columns)))) + for view in views: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) + ) + # Kafka and webhook source objects are readable directly, the + # others follow the source-table model: str(obj) names the table, + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ("SOURCE", f"{source.schema}.{identifier(source.source_name())}") + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for index in indexes: + candidates.append(("INDEX", str(index))) + for schema in schemas: + candidates.append(("SCHEMA", str(schema))) + for db in dbs: + candidates.append(("DATABASE", str(db))) + for cluster in clusters: + candidates.append(("CLUSTER", str(cluster))) + for role in roles: + candidates.append(("ROLE", str(role))) + candidates.append(("SECRET", "materialize.public.pgpass")) + candidates.append(("CONNECTION", "materialize.public.kafka_conn")) + if not candidates: + return False + kind, name = self.rng.choice(candidates) + comment = self.rng.choice([f"'{Text.random_value(self.rng)}'", "NULL"]) + query = f"COMMENT ON {kind} {name} IS {comment}" exe.execute(query, http=Http.RANDOM) return True @@ -1196,10 +1634,12 @@ def run(self, exe: Executor) -> bool: # The indexed object or its schema may have been dropped # concurrently, taking the index with it. Untrack the index # either way so stale entries don't fill up the set and choke - # off CreateIndexAction. - exe.db.indexes.remove(index) + # off CreateIndexAction. Use discard, not remove: a concurrent + # CASCADE drop's untrack_objects_in_schemas may have already + # removed it, and remove would raise KeyError. + exe.db.indexes.discard(index) raise - exe.db.indexes.remove(index) + exe.db.indexes.discard(index) return True @@ -1263,7 +1703,12 @@ def run(self, exe: Executor) -> bool: query = f"DROP TABLE {table}" exe.execute(query, http=Http.RANDOM) - exe.db.tables.remove(table) + # A concurrent CASCADE drop's untrack_objects_in_schemas may have + # already filtered this table out of the list, tolerate that. + try: + exe.db.tables.remove(table) + except ValueError: + pass return True @@ -1425,11 +1870,28 @@ def run(self, exe: Executor) -> bool: return True -class ReplaceMaterializedViewAction(Action): +class CreateReplacementMaterializedViewAction(Action): + """CREATE REPLACEMENT MATERIALIZED VIEW for a random materialized view. + + Deliberately only creates. The replacement stays alive until an + ApplyReplacementMaterializedViewAction or + DropReplacementMaterializedViewAction finds it in the catalog, from any + worker and possibly only after a kill or 0dt deploy. Resolving the + replacement on the creating connection would tie its lifetime to that + connection, so it could never span an envd restart. Incident 1136 + (bootstrap loses the replacement collection's primary link to the + target's shard, a post-restart DROP of the replacement then finalizes the + target's live shard) lives exactly in that gap. + """ + def errors_to_ignore(self, exe: Executor) -> list[str]: errors = [ - # A concurrent or leaked replacement of the same view + # A concurrent replacement of the same view "because it already has a replacement", + # The target's query constant-folded to a plan without live + # inputs and completed, which can_seal cannot know. The + # SealedCollectionCheckAction oracle arbitrates whether a sealed + # shard is legitimate. "is sealed and thus cannot be replaced", ] + super().errors_to_ignore(exe) if exe.db.scenario == Scenario.Rename: @@ -1440,7 +1902,15 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: def run(self, exe: Executor) -> bool: with exe.db.lock: - mvs = [v for v in exe.db.views if v.materialized] + # Skip views whose shard is known to seal legitimately. Sealing + # is transitive: REFRESH AT views seal after their last refresh, + # repeat_row constant views seal on hydration, and any view that + # reads from a sealing input seals too (View.can_seal). A sealed + # target cannot be replaced, so skipping these avoids wasted + # work. can_seal cannot know about constant-folded queries + # though, so "is sealed" errors are still ignored and + # SealedCollectionCheckAction arbitrates damage. + mvs = [v for v in exe.db.views if v.materialized and not v.can_seal] if not mvs: return False view = self.rng.choice(mvs) @@ -1456,17 +1926,298 @@ def run(self, exe: Executor) -> bool: exe.execute( f"CREATE REPLACEMENT MATERIALIZED VIEW {tmp_mv} FOR {view} AS {view.get_select()}", ) - time.sleep(self.rng.random()) + return True + + +class ResolveReplacementMaterializedViewAction(Action): + """Base for actions resolving live replacements. + + Replacements are picked from the catalog instead of tracked state so that + replacements created by other workers or left behind by dead connections, + kills, and 0dt deploys are found too. + """ + + def pick_replacement(self, exe: Executor) -> tuple[str, str] | None: + exe.execute( + "SELECT rd.name, rs.name, r_mv.name, td.name, ts.name, t_mv.name" + " FROM mz_internal.mz_replacements r" + " JOIN mz_catalog.mz_materialized_views r_mv ON r_mv.id = r.id" + " JOIN mz_catalog.mz_schemas rs ON rs.id = r_mv.schema_id" + " JOIN mz_catalog.mz_databases rd ON rd.id = rs.database_id" + " JOIN mz_catalog.mz_materialized_views t_mv ON t_mv.id = r.target_id" + " JOIN mz_catalog.mz_schemas ts ON ts.id = t_mv.schema_id" + " JOIN mz_catalog.mz_databases td ON td.id = ts.database_id", + http=Http.NO, + ) + rows = exe.cur.fetchall() + if not rows: + return None + r_db, r_schema, r_name, t_db, t_schema, t_name = self.rng.choice(rows) + replacement = f"{identifier(r_db)}.{identifier(r_schema)}.{identifier(r_name)}" + target = f"{identifier(t_db)}.{identifier(t_schema)}.{identifier(t_name)}" + return replacement, target + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # Reading catalog relations inside a txn that already touched + # user objects crosses timedomains. + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + +class ApplyReplacementMaterializedViewAction(ResolveReplacementMaterializedViewAction): + def errors_to_ignore(self, exe: Executor) -> list[str]: + errors = [ + # A concurrent apply or drop won the race for this replacement + "replacement materialized view does not exist", + # The target's shard sealed, which is legitimate for targets + # whose query constant-folded to a plan without live inputs. + # SealedCollectionCheckAction arbitrates whether a sealed shard + # is damage, so it is not treated as a failure here. + "is sealed and thus cannot be replaced", + ] + super().errors_to_ignore(exe) + if exe.db.scenario == Scenario.Rename: + # The target was renamed between picking and applying + errors += ["does not exist"] + return errors + + def run(self, exe: Executor) -> bool: + picked = self.pick_replacement(exe) + if picked is None: + return False + replacement, target = picked + # APPLY REPLACEMENT waits for the target's write frontier to catch up + # to the replacement's. The wait never finishes if the target is + # concurrently cascade-dropped (database-issues#9820), so bound it + # with a statement timeout. + exe.execute("SET statement_timeout = '60s'") + # A timeout is expected for as long as this one is configured. + exe.statement_timeout_set = True try: - exe.execute(f"ALTER MATERIALIZED VIEW {view} APPLY REPLACEMENT {tmp_mv}") - except QueryError: - # Clean up, a leaked replacement blocks all future replacements - # of this view. + exe.execute( + f"ALTER MATERIALIZED VIEW {target} APPLY REPLACEMENT {replacement}" + ) + finally: try: - exe.execute(f"DROP MATERIALIZED VIEW IF EXISTS {tmp_mv}") + exe.execute("RESET statement_timeout") except QueryError: + # The timeout is still configured, so keep tolerating it. + pass + else: + # Back at the default timeout, so a timeout is a genuine + # finding again. The worker's statements are sequential, + # nothing from the window above can still be in flight. + exe.statement_timeout_set = False + return True + + +class DropReplacementMaterializedViewAction(ResolveReplacementMaterializedViewAction): + """Drop a live replacement without applying it. + + This is the incident 1136 trigger statement when it runs after an envd + restart: bootstrap loses the replacement collection's primary link, so + the drop wrongly finalizes the target's live shard. + SealedCollectionCheckAction detects the resulting damage. + """ + + def run(self, exe: Executor) -> bool: + picked = self.pick_replacement(exe) + if picked is None: + return False + replacement, _ = picked + # IF EXISTS, a concurrent apply or drop may have resolved it already. + exe.execute(f"DROP MATERIALIZED VIEW IF EXISTS {replacement}") + return True + + +def plan_user_input_ids(plan: str) -> list[str]: + """User collection ids that an `EXPLAIN .. AS JSON` dataflow plan reads. + + Raises on a plan whose shape it cannot read, rather than returning an empty + list. Constant folding legitimately produces plans that read nothing, so + "reads nothing" and "no longer understands the plan" are indistinguishable + to a caller, and treating the second as the first retires whatever the + caller checks without a trace. Only accepts dataflow (MIR) plans, a + fast-path peek plan carries its inputs elsewhere and is rejected. + """ + try: + parsed = json.loads(plan) + except json.JSONDecodeError as e: + raise ValueError(f"EXPLAIN .. AS JSON output does not parse ({e}): {plan}") + + # `Constant` and `Get` are the only leaves of MirRelationExpr, so a plan + # with neither is not a shape this can read. Global ids outside `Get` are + # not inputs, e.g. the index an `access_strategy` picked. + get_ids: list = [] + constant = False + + def walk(node: object) -> None: + nonlocal constant + if isinstance(node, dict): + get = node.get("Get") + if isinstance(get, dict) and "id" in get: + get_ids.append(get["id"]) + if "Constant" in node: + constant = True + for value in node.values(): + walk(value) + elif isinstance(node, list): + for value in node: + walk(value) + + walk(parsed) + if not get_ids and not constant: + raise ValueError(f"plan has no Get or Constant leaf: {plan}") + + user_ids = set() + for get_id in get_ids: + match get_id: + case {"Global": {"User": int(gid)}}: + user_ids.add(f"u{gid}") + case {"Local": _} | { + "Global": {"System": _} + | {"Transient": _} + | {"IntrospectionSourceIndex": _} + }: + # A Let binding or a non-user namespace, neither is a user + # collection with a frontier to check. pass + case _: + raise ValueError(f"unrecognized Get id {get_id!r} in plan: {plan}") + return sorted(user_ids) + + +class SealedCollectionCheckAction(Action): + """Client-side data-loss oracle for wrongly finalized persist shards. + + A shard's upper only becomes empty when the shard is sealed. That is + legitimate when the writing dataflow completed (a constant plan, inputs + that all sealed, a REFRESH AT view past its last refresh, a bounded load + generator that finished) and destructive when shard finalization ran + against a live collection. Incident 1136 is the known destructive + trigger: dropping a replacement after an envd restart finalizes the + target's shard because bootstrap dropped the replacement collection's + primary link. Tables are covered too because versioned tables + (ALTER TABLE ADD COLUMN) carry the same primary ownership chain. + + A live user table or materialized view with an empty write frontier and + at least one live plan input has no legitimate explanation and is + reported as damage. This oracle is the single arbiter of wrongly sealed + shards, the replacement actions ignore "is sealed" errors instead of + treating them as failures. + """ + + def applicable(self, exe: Executor) -> bool: + # repeat_row constant views seal legitimately on hydration and are + # indistinguishable from damage in the catalog. + return exe.db.scenario != Scenario.RepeatRow + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def completed_legitimately( + self, exe: Executor, database: str, schema: str, name: str + ) -> bool: + """Whether a sealed materialized view's dataflow completed on its own. + + A dataflow completes, sealing its shard, once it can never produce + more output: its optimized plan reads no storage collections at all + (constant folding, e.g. a join ON false, drops inputs the catalog + still records as dependencies), or every storage collection it reads + is itself sealed. Only a sealed view with at least one live plan + input is damage. + """ + try: + exe.execute( + "EXPLAIN OPTIMIZED PLAN AS JSON FOR MATERIALIZED VIEW" + f" {identifier(database)}.{identifier(schema)}.{identifier(name)}", + http=Http.NO, + ) + plan = "\n".join(str(row[0]) for row in exe.cur.fetchall()) + except QueryError as e: + if "unknown catalog item" in str(e) or "does not exist" in str(e): + # Dropped concurrently, so its shard is allowed to seal. + return True raise + input_ids = plan_user_input_ids(plan) + if not input_ids: + return True + id_list = ", ".join(f"'{i}'" for i in input_ids) + exe.execute( + "SELECT count(*) FROM mz_internal.mz_frontiers" + f" WHERE object_id IN ({id_list}) AND write_frontier IS NOT NULL", + http=Http.NO, + ) + return exe.cur.fetchall()[0][0] == 0 + + def run(self, exe: Executor) -> bool: + # Sealing is transitive: a view reading from a REFRESH AT view seals + # once that input seals, even though its own refresh strategy is not + # 'at', and likewise for anything reading a bounded load generator's + # table. Walk the dependency graph up from every REFRESH AT view and + # every load generator source and exclude everything that transitively + # depends on one. The COUNTER sources are UP TO bounded and seal when + # they finish. The multi-subsource generators never seal, so sweeping + # them in only widens the exclusion, erring toward tolerating. repeat_row + # constant views also seal transitively but are not marked in the + # catalog, which is why this oracle is disabled in the RepeatRow + # scenario, see applicable. Plain user tables never seal and need no + # exclusions, and source-fed tables are covered by the walk through + # their source. + def sealed_collections() -> list: + exe.execute( + "WITH MUTUALLY RECURSIVE" + " sealing(id text) AS (" + "SELECT rs.materialized_view_id" + " FROM mz_internal.mz_materialized_view_refresh_strategies rs" + " WHERE rs.type = 'at'" + " UNION" + " SELECT s.id" + " FROM mz_catalog.mz_sources s" + " WHERE s.type = 'load-generator'" + " UNION" + " SELECT d.object_id" + " FROM mz_internal.mz_object_dependencies d" + " JOIN sealing s ON d.referenced_object_id = s.id)" + " SELECT o.id, d.name, sc.name, o.name, o.type" + " FROM mz_catalog.mz_objects o" + " JOIN mz_catalog.mz_schemas sc ON sc.id = o.schema_id" + " JOIN mz_catalog.mz_databases d ON d.id = sc.database_id" + " JOIN mz_internal.mz_frontiers f ON f.object_id = o.id" + " WHERE o.type IN ('materialized-view', 'table')" + " AND o.id LIKE 'u%'" + " AND f.write_frontier IS NULL" + " AND o.id NOT IN (SELECT id FROM sealing)", + http=Http.NO, + ) + return exe.cur.fetchall() + + sealed = sealed_collections() + if not sealed: + return True + # The catalog relations and mz_frontiers are not updated atomically, + # so a query timestamp can land inside a concurrent CREATE or DROP + # where the object is cataloged but its frontier reads as the empty + # antichain. A wrongly finalized shard is permanent, so only report + # objects that are still sealed on a re-check. + time.sleep(5) + sealed_ids = {row[0] for row in sealed} + damaged = [ + row + for row in sealed_collections() + if row[0] in sealed_ids + and not ( + row[4] == "materialized-view" + and self.completed_legitimately(exe, row[1], row[2], row[3]) + ) + ] + if damaged: + raise ValueError( + "Collections with wrongly sealed persist shards" + f" (incident 1136 class): {damaged}" + ) return True @@ -1495,13 +2246,13 @@ def run(self, exe: Executor) -> bool: # Iceberg sinks always have a key, so only allow a conservative # case: all names, types, and nullabilities match, which also # guarantees the key columns exist in the new object. - # TODO: Switch back when SS-324 is fixed to make sure it errors + # TODO: Switch back when SS-344 is fixed to make sure it errors # instead of causing a stall objs = [] old_cols = { c.name(True): (c.data_type, c.nullable) for c in old_object.columns } - for o in exe.db.db_objects_without_views(): + for o in exe.db.db_objects_for_sinks(): if isinstance(old_object, WebhookSource): continue if isinstance(o, WebhookSource): @@ -1557,14 +2308,14 @@ def run(self, exe: Executor) -> bool: # single column formats objs = [ o - for o in exe.db.db_objects_without_views() + for o in exe.db.db_objects_for_sinks() if len(o.columns) == 1 and o.columns[0].data_type == old_object.columns[0].data_type ] elif sink.format in ["FORMAT JSON"]: # We should be able to format all data types as JSON, and they have no # particular backwards-compatiblility requirements. - objs = [o for o in exe.db.db_objects_without_views()] + objs = [o for o in exe.db.db_objects_for_sinks()] else: # Avro schema migration checking can be quite strict, and we need to be not only # compatible with the latest object's schema but all previous schemas. @@ -1577,7 +2328,7 @@ def run(self, exe: Executor) -> bool: old_cols = { c.name(True): (c.data_type, c.nullable) for c in old_object.columns } - for o in exe.db.db_objects_without_views(): + for o in exe.db.db_objects_for_sinks(): if isinstance(old_object, WebhookSource): continue if isinstance(o, WebhookSource): @@ -1805,43 +2556,241 @@ def run(self, exe: Executor) -> bool: return True -class CommitRollbackAction(Action): - def run(self, exe: Executor) -> bool: - if not exe.action_run_since_last_commit_rollback: - return False +class ParameterizedQueryAction(Action): + """PREPARE a query with $1..$n placeholders, then EXECUTE it with values. - if self.rng.random() < 0.7: - exe.commit() - else: - exe.rollback() - exe.action_run_since_last_commit_rollback = False + Exercises the parameter-type-inference and bind/assignment-cast path that + the workload's other prepared statements (no parameters) never reach.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + obj = self.rng.choice(exe.db.db_objects()) + n = self.rng.randint(1, 4) + # Record and record list are expression-only pseudo types, a + # parameter cast to them fails with "cannot reference pseudo type". + param_types = [self.rng.choice(list(DATA_TYPES_FOR_COLUMNS)) for _ in range(n)] + projection = ", ".join( + f"${i + 1}::{t.name()}" for i, t in enumerate(param_types) + ) + self.stmt_id += 1 + name = f"pq{self.stmt_id}" + query = f"SELECT {projection} FROM {obj} LIMIT {self.rng.randint(0, 10)}" + # Each argument is cast to its parameter's declared type so the + # assignment cast on EXECUTE always succeeds (e.g. a bytea parameter + # rejects a bare text literal). + values = ", ".join( + f"({t.random_value(self.rng, in_query=True)})::{t.name()}" + for t in param_types + ) + # Run sequentially, not in a try/finally: if EXECUTE fails it aborts + # the transaction, and a DEALLOCATE in a finally would then fail with + # "current transaction is aborted", masking the real error. On failure + # the statement leaks on the session, which is harmless: statement + # names are never reused, and the session's prepared statements die + # with it. + exe.execute(f"PREPARE {name} AS {query}", http=Http.NO) + exe.execute(f"EXECUTE {name} ({values})", http=Http.NO, fetch=True) + exe.execute(f"DEALLOCATE {name}", http=Http.NO) return True -class FlipFlagsAction(Action): - def __init__( - self, - rng: random.Random, - composition: Composition | None, - ): - super().__init__(rng, composition) +class BoundedStalenessReadAction(Action): + """A read under `bounded staleness` isolation, a distinct timestamp- + selection path that never blocks and returns 40001 when the freshness + bound cannot be met. The isolation is set transiently around the read and + restored afterwards, since bounded staleness is read-only and would break + writes.""" - BOOLEAN_FLAG_VALUES = ["TRUE", "FALSE"] + def applicable(self, exe: Executor) -> bool: + return exe.db.flags.get("enable_bounded_staleness_isolation", "FALSE") == "TRUE" - self.flags_with_values: dict[str, list[str]] = dict() - self.flags_with_values["persist_blob_target_size"] = ( - # 1 MiB, 16 MiB, 128 MiB - ["1048576", "16777216", "134217728"] - ) - for flag in ["catalog", "source", "snapshot", "txn"]: - self.flags_with_values[f"persist_use_critical_since_{flag}"] = ( - BOOLEAN_FLAG_VALUES - ) - self.flags_with_values["persist_claim_unclaimed_compactions"] = ( - BOOLEAN_FLAG_VALUES + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + # The flag was flipped off between applicable() and run(). + "is not available", + # The freshness bound could not be met. Bounded staleness + # never blocks, it errors instead. + "not been materialized", + "could not find a valid timestamp for the query", + "cannot serve query under bounded staleness", + # A leaked real_time_recency SET on the session (its own reset + # was discarded on a prior query's error path) conflicts with + # bounded staleness. + "cannot be combined with bounded staleness", + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] ) - self.flags_with_values["persist_optimize_ignored_data_fetch"] = ( - BOOLEAN_FLAG_VALUES + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + bound = self.rng.choice(["1s", "5s", "30s"]) + restore = f"SET TRANSACTION_ISOLATION TO '{exe.isolation}'" + exe.execute( + f"SET TRANSACTION_ISOLATION TO 'bounded staleness {bound}'", + explainable=False, + ) + try: + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute(query, http=Http.NO, fetch=True) + except QueryError: + # The restore must not raise on top of the read's error, e.g. on a + # killed or cancelled session, that error would replace the read's + # and hide the real failure. A leaked bounded staleness isolation + # makes later actions fail with errors only this action ignores, so + # a failed restore forces a reconnect rather than being swallowed. + try: + exe.execute(restore, explainable=False) + except QueryError: + exe.reconnect_next = True + raise + # No error is being masked here, so a failing restore is a genuine + # failure and propagates. + exe.execute(restore, explainable=False) + return True + + +class ReadOnlyTransactionAction(Action): + """A multi-statement `BEGIN READ ONLY` transaction. All reads run at one + pinned timestamp (one timedomain), and holding the read pins compaction + for the transaction's lifetime.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + isolation = self.rng.choice( + [ + "", + " ISOLATION LEVEL SERIALIZABLE", + " ISOLATION LEVEL STRICT SERIALIZABLE", + ] + ) + exe.execute(f"BEGIN{isolation} READ ONLY", http=Http.NO) + try: + for _ in range(self.rng.randint(1, 3)): + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute(query, http=Http.NO, fetch=True) + except QueryError: + # The transaction still has to end, an open one wedges every later + # action on this session, but ending it must not raise over the + # read's error and hide the real failure. COMMIT on a failed + # transaction is turned into a rollback by the coordinator anyway, + # so always rolling back here loses no coverage. + try: + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + exe.reconnect_next = True + raise + # Outside the error path both endings are worth exercising, and a + # failure of the end statement itself is a genuine one. + end = "COMMIT" if self.rng.choice([True, False]) else "ROLLBACK" + exe.execute(end, http=Http.NO) + return True + + +class DDLTransactionAction(Action): + """A DDL statement inside an explicit `BEGIN`/`COMMIT`. Materialize allows + only a single statement in a DDL transaction, so the value over autocommit + DDL is the open commit window: racing concurrent DDL against it exercises + the "another session modified the catalog while this DDL transaction was + open" serialization path.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "another session modified the catalog", + "unknown schema", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + if len([t for t in exe.db.tables if not t.temp]) >= MAX_TABLES: + return False + try: + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + table_id = exe.db.table_id + exe.db.table_id += 1 + table = Table(self.rng, table_id, schema) + exe.execute("BEGIN", http=Http.NO) + try: + table.create(exe) + exe.execute("COMMIT", http=Http.NO) + except QueryError: + try: + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + pass + raise + with exe.db.lock: + exe.db.tables.append(table) + return True + + +class CommitRollbackAction(Action): + def run(self, exe: Executor) -> bool: + if not exe.action_run_since_last_commit_rollback: + return False + + if self.rng.random() < 0.7: + exe.commit() + else: + exe.rollback() + exe.action_run_since_last_commit_rollback = False + return True + + +class FlipFlagsAction(Action): + def __init__( + self, + rng: random.Random, + composition: Composition | None, + ): + super().__init__(rng, composition) + + BOOLEAN_FLAG_VALUES = ["TRUE", "FALSE"] + + self.flags_with_values: dict[str, list[str]] = dict() + self.flags_with_values["persist_blob_target_size"] = ( + # 1 MiB, 16 MiB, 128 MiB + ["1048576", "16777216", "134217728"] + ) + for flag in ["catalog", "source", "snapshot", "txn"]: + self.flags_with_values[f"persist_use_critical_since_{flag}"] = ( + BOOLEAN_FLAG_VALUES + ) + self.flags_with_values["persist_claim_unclaimed_compactions"] = ( + BOOLEAN_FLAG_VALUES + ) + self.flags_with_values["persist_optimize_ignored_data_fetch"] = ( + BOOLEAN_FLAG_VALUES ) self.flags_with_values["persist_source_fetch_concurrency"] = [ "1", @@ -1960,6 +2909,9 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["enable_alter_table_add_column"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["enable_bounded_staleness_isolation"] = ( + BOOLEAN_FLAG_VALUES + ) self.flags_with_values["enable_arrangement_dictionary_compression_alpha"] = ( BOOLEAN_FLAG_VALUES ) @@ -1975,10 +2927,10 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["cluster"] = ["quickstart", "dont_exist"] - self.flags_with_values["enable_frontend_peek_sequencing"] = [ - "true", - "false", - ] + # NOTE: enable_frontend_peek_sequencing is pinned off in + # ADDITIONAL_SYSTEM_PARAMETER_DEFAULTS (frontend-peek read-hold vs + # compaction race, https://linear.app/materializeinc/issue/SQL-520), so + # it is not flipped here. self.flags_with_values["enable_frontend_subscribes"] = [ "true", "false", @@ -1995,11 +2947,17 @@ def __init__( self.flags_with_values["enable_column_paged_batcher_spill"] = ( BOOLEAN_FLAG_VALUES ) + # Fractions of the *cgroup* memory limit, which under mzcompose's + # process orchestrator is the whole container budget shared by + # environmentd and every replica, not one replica's allowance. A + # process-singleton pool at 0.25 is 6 GiB per clusterd on the 24 GiB CI + # agent, so the values stay small enough that all replicas together + # cannot claim the budget. Small pools also fill up sooner, which + # exercises the paging and compression paths more, not less. self.flags_with_values["column_paged_batcher_budget_fraction"] = [ "0.0", "0.01", - "0.05", - "0.25", + "0.02", ] self.flags_with_values["column_paged_batcher_lz4"] = BOOLEAN_FLAG_VALUES self.flags_with_values["column_paged_batcher_swap_pageout"] = ( @@ -2013,10 +2971,14 @@ def __init__( self.flags_with_values["column_paged_batcher_eager_backing"] = ( BOOLEAN_FLAG_VALUES ) + # Same shared-budget reasoning as the budget fraction above. Sharing its + # value set keeps all three orderings against the budget covered: a + # target of 0 collapses the compressed tier, a target below the budget + # pages out above it, and a target above the budget pages nothing out. self.flags_with_values["column_paged_batcher_pool_rss_target_fraction"] = [ "0.0", - "0.25", - "0.5", + "0.01", + "0.02", ] self.flags_with_values["enable_upsert_paged_spill"] = BOOLEAN_FLAG_VALUES # 0 forces the estimated-size path for every table, the default forces @@ -2329,12 +3291,21 @@ def run(self, exe: Executor) -> bool: flag_value = self.rng.choice(self.flags_with_values[flag_name]) + # Occasionally restore the default instead, a distinct path from + # SET-to-value. + reset = self.rng.random() < 0.1 + conn = None try: conn = self.create_system_connection(exe) if cluster is not None: self.set_cluster_compression(conn, cluster) + elif reset: + self.reset_flag(conn, flag_name) + # Gates reading exe.db.flags fall back to their conservative + # default when the key is absent. + exe.db.flags.pop(flag_name, None) else: self.flip_flag(conn, flag_name, flag_value) exe.db.flags[flag_name] = flag_value @@ -2366,6 +3337,12 @@ def set_cluster_compression(self, conn: Connection, cluster: Cluster) -> None: with conn.cursor() as cur: cur.execute(f"ALTER CLUSTER {cluster} {option};".encode()) + def reset_flag(self, conn: Connection, flag_name: str) -> None: + with conn.cursor() as cur: + cur.execute( + f"ALTER SYSTEM RESET {flag_name};".encode(), + ) + class CreateViewAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: @@ -2469,7 +3446,46 @@ def run(self, exe: Executor) -> bool: else: query = f"DROP VIEW {view}" exe.execute(query, http=Http.RANDOM) - exe.db.views.remove(view) + # A concurrent CASCADE drop's untrack_objects_in_schemas may have + # already filtered this view out of the list, tolerate that. + try: + exe.db.views.remove(view) + except ValueError: + pass + return True + + +class CreateOrReplaceViewAction(Action): + """In-place swap of an existing view's definition via CREATE OR REPLACE. + + The body is unchanged, so dependents stay valid, but the coordinator still + tears down and rebuilds the item (and the dataflow, for a materialized + view). Racing that swap against reads and concurrent DDL is the point.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + errors = [ + # A dependent references a column the replacement would drop. The + # body is unchanged here, but a concurrent replacement may have + # changed it. + "still depended upon by", + "replica-targeted materialized views is not supported", + "unknown cluster replica", + ] + super().errors_to_ignore(exe) + if exe.db.scenario == Scenario.Rename: + # A base object was renamed, invalidating the captured body. + errors += ["does not exist", "ambiguous reference to schema name"] + return errors + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + views = [view for view in exe.db.views if not view.temp] + if not views: + return False + view = self.rng.choice(views) + with view.lock: + if view not in exe.db.views: + return False + view.create(exe, or_replace=True) return True @@ -2517,6 +3533,48 @@ def run(self, exe: Executor) -> bool: return True +class AlterRoleAction(Action): + """ALTER ROLE ... SET / RESET a default session variable. + + Exercises the per-role session-default path (applied at session init) and + role-name resolution on the ALTER path, neither of which CREATE/DROP ROLE + touch. Part of broadening ALTER coverage (the ALTER paths are where several + catalog bugs have hidden vs the well-worn CREATE/DROP).""" + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + with role.lock: + # Was dropped while we were acquiring the lock. + if role not in exe.db.roles: + return False + var, value = self.rng.choice( + [ + ("cluster", "'quickstart'"), + ("transaction_isolation", "'serializable'"), + ("statement_timeout", "'120s'"), + ("search_path", "public"), + ] + ) + query = ( + f"ALTER ROLE {role} RESET {var}" + if self.rng.choice([True, False]) + else f"ALTER ROLE {role} SET {var} = {value}" + ) + try: + exe.execute(query, http=Http.RANDOM) + except QueryError as e: + # Concurrent DROP ROLE, expected as with DropRoleAction. + if ( + exe.db.scenario not in (Scenario.Kill, Scenario.ZeroDowntimeDeploy) + or "unknown role" not in e.msg + ): + raise e + return True + + class CreateClusterAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -2527,9 +3585,7 @@ def run(self, exe: Executor) -> bool: cluster = Cluster( cluster_id, managed=self.rng.choice([True, False]), - size=self.rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), replication_factor=self.rng.choice([1, 2]), introspection_interval="1s", ) @@ -2681,9 +3737,7 @@ def run(self, exe: Executor) -> bool: replica = ClusterReplica( replica_id, - size=self.rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), cluster=cluster, ) replica.create(exe) @@ -2729,8 +3783,10 @@ def run(self, exe: Executor) -> bool: class ReconfigureClusterAction(Action): - """Gracefully resize a random managed cluster and assert convergence. + """Gracefully resize a random managed cluster and watch its record settle. + The asserted invariant is that a reconfiguration record never sits + in-progress past its own deadline, whether it cuts over or rolls back. Regression coverage for SQL-530: readiness must not wait for single-replica sources (Postgres, MySQL, SQL Server) hosted on the cluster, which never hydrate on the reconfiguration's target replicas @@ -2781,40 +3837,85 @@ def run(self, exe: Executor) -> bool: # e.g. when another worker put a REFRESH EVERY materialized view # with a far-future refresh on the cluster, which cannot hydrate # on the target replicas before its refresh time. What must never - # happen is the record staying "in-progress" past its deadline, - # which is how a wedged readiness check (SQL-530) manifests. Only - # trust a terminal status once the ALTER's own deadline has - # passed: builtin table reads can lag the catalog, so before that - # a terminal status can only be a stale record from an earlier - # reconfiguration of this cluster. + # happen is a record staying "in-progress" past its own deadline, + # which is how a wedged readiness check (SQL-530) manifests. + # + # `mz_now()` is the timestamp the read itself ran at, and comparing + # it against the record's `deadline` is what makes that assertion + # sound. `mz_clusters` is a builtin table and + # `mz_cluster_reconfigurations` a builtin materialized view + # maintained by mz_catalog_server, and a SERIALIZABLE read of either + # takes the freshest timestamp that does not block, so under load it + # can lag wall clock by minutes. Judging the deadline against the + # test's wall clock therefore calls a correctly rolled-back record + # wedged as soon as the read lags. Both sides coming from the same + # read keeps them in the same frame, stale or not. name_literal = cluster.name().replace("'", "''") query = ( - "SELECT c.size, r.status FROM mz_clusters c " + "SELECT c.size, r.status, r.deadline::text, mz_now()::text " + "FROM mz_clusters c " "LEFT JOIN mz_internal.mz_cluster_reconfigurations r ON r.cluster_id = c.id " f"WHERE c.name = '{name_literal}'" ) - status = None + # The deadline our own ALTER wrote is `now() + 120s` in the same + # millisecond epoch `mz_now()` reports, so a read past this is + # guaranteed to show our record or a later one, never a stale + # terminal record from an earlier reconfiguration of this cluster. + our_deadline_ms = int((alter_started + 120) * 1000) seen = False - deadline = time.time() + 180 - while time.time() < deadline: + # Must stay comfortably under the 300s the run gives all workers to + # join after `end_time` (`parallel_workload.run`). A poll that + # outlives that budget makes a worker doing its job look wedged, and + # the run then hard-exits 0 without its final checks. + poll_until = time.time() + 240 + while time.time() < poll_until: exe.execute(query) rows = exe.cur.fetchall() if not rows: # `mz_clusters` is a builtin table, so the read can lag a # concurrent rename or swap of this cluster and not show - # the polled name yet. Keep polling and let the deadline + # the polled name yet. Keep polling and let the check # below decide, rather than indexing an empty result. time.sleep(1) continue seen = True - size, status = rows[0] + size, status, record_deadline, read_ts = rows[0] + read_ts_ms = int(read_ts) if size == new_size: cluster.size = new_size return True - if time.time() > alter_started + 120 and status in ( - "timed-out", - "resource-exhausted", + # 60s of slack over the controller's 5s tick, which is what + # turns a reached deadline into the terminal status. + if ( + status == "in-progress" + and read_ts_ms > int(record_deadline) + 60_000 ): + raise ValueError( + f"Graceful reconfiguration of cluster {cluster} to size " + f"{new_size} sat in-progress {read_ts_ms - int(record_deadline)}ms " + f"past its deadline {record_deadline}" + ) + if read_ts_ms > our_deadline_ms and status is None: + raise ValueError( + f"Reconfiguration record of cluster {cluster} is gone as of " + f"{read_ts_ms}, past the deadline {our_deadline_ms} of the " + f"reconfiguration to size {new_size} that wrote it" + ) + # Any settled record is a legitimate end state, so stop on every + # status but "in-progress". A rollback at the deadline + # ("timed-out"), a controller that could not get the replicas + # ("resource-exhausted"), a concurrent shape-touching ALTER + # calling the transition off ("cancelled"), or one that + # retargeted it to a shape that then cut over ("finalized" on a + # size that is not ours) all mean the reconfiguration settled. + # FlipFlagsAction alters EXPERIMENTAL ARRANGEMENT COMPRESSION + # without holding the cluster lock, and that is a shape + # dimension (`alter_changes_replica_shape`), so it reaches the + # record this poll is watching. Once ours settles, such an ALTER + # even starts a fresh record, in-progress with the 24h default + # deadline, which is why running out of the budget below on an + # in-progress record is not by itself a failure. + if read_ts_ms > our_deadline_ms and status != "in-progress": return True time.sleep(1) if not seen: @@ -2823,10 +3924,7 @@ def run(self, exe: Executor) -> bool: f"name the reconfiguration to size {new_size} polled for, so the " f"name the workload holds does not match the catalog" ) - raise ValueError( - f"Graceful reconfiguration of cluster {cluster} to size {new_size} " - f"did not complete (reconfiguration status: {status})" - ) + return True class GrantPrivilegesAction(Action): @@ -2835,7 +3933,7 @@ def run(self, exe: Executor) -> bool: if not exe.db.roles: return False role = self.rng.choice(exe.db.roles) - privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "ALL"]) + privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"]) tables_views: list[DBObject] = [*exe.db.tables, *exe.db.views] table = self.rng.choice(tables_views) with table.lock, role.lock: @@ -2863,7 +3961,7 @@ def run(self, exe: Executor) -> bool: if not exe.db.roles: return False role = self.rng.choice(exe.db.roles) - privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "ALL"]) + privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"]) tables_views: list[DBObject] = [*exe.db.tables, *exe.db.views] table = self.rng.choice(tables_views) with table.lock, role.lock: @@ -2885,68 +3983,597 @@ def run(self, exe: Executor) -> bool: return True -# TODO: Should factor this out so can easily use it without action -class ReconnectAction(Action): - def __init__( - self, - rng: random.Random, - composition: Composition | None, - random_role: bool = True, - ): - super().__init__(rng, composition) - self.random_role = random_role +class GrantRoleAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + # Concurrent memberships can close a cycle, which is rejected. + "is a member of role", + ] + super().errors_to_ignore(exe) def run(self, exe: Executor) -> bool: - exe.mz_service = "materialized" - exe.log("reconnecting") - # The connection's temp objects die with it, drop them from the - # tracked state so other workers stop querying them. - if exe.temp_objects: - with exe.db.lock: - exe.db.tables[:] = [ - t for t in exe.db.tables if t not in exe.temp_objects - ] - exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] - exe.temp_objects.clear() - host = exe.db.host + with exe.db.lock: + if len(exe.db.roles) < 2: + return False + role1, role2 = self.rng.sample(exe.db.roles, 2) + # Both locks are held across the round trip below, so take them in + # a fixed order. Two workers that sampled the same pair in + # opposite orders would otherwise deadlock. The GRANT operands + # keep the sampled order, so memberships still form in both + # directions, including the cycles this action provokes. + first, second = sorted((role1, role2), key=lambda role: role.role_id) + with first.lock, second.lock: + if role1 not in exe.db.roles or role2 not in exe.db.roles: + return False + exe.execute(f"GRANT {role1} TO {role2}", http=Http.RANDOM) + return True - def pg_port() -> int: - # System workers (e.g. the Cancel worker) live on the internal - # port, everyone else on the external one of the current service. - if exe.user == "mz_system": - return exe.db.ports[ - "mz_system" if exe.mz_service == "materialized" else "mz_system2" - ] - return exe.db.ports[exe.mz_service] +class RevokeRoleAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: with exe.db.lock: - if self.random_role and exe.db.roles: - user = self.rng.choice( - ["materialize", str(self.rng.choice(exe.db.roles))] + if len(exe.db.roles) < 2: + return False + role1, role2 = self.rng.sample(exe.db.roles, 2) + # Fixed lock order, as in GrantRoleAction: the sampled order + # deadlocks two workers that picked the same pair, and the REVOKE + # operands have to stay in the sampled order. + first, second = sorted((role1, role2), key=lambda role: role.role_id) + with first.lock, second.lock: + if role1 not in exe.db.roles or role2 not in exe.db.roles: + return False + exe.execute(f"REVOKE {role1} FROM {role2}", http=Http.RANDOM) + return True + + +class AlterOwnerAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = [ + "unknown role", + "must be a member of", + ] + super().errors_to_ignore(exe) + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + # Only the snapshot needs the global lock, see CommentAction. + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + schemas = list(exe.db.schemas) + dbs = list(exe.db.dbs) + clusters = list(exe.db.clusters) + + candidates: list[tuple[str, str]] = [] + # Temp objects cannot change owner, they die with the session. + for table in tables: + if not table.temp: + candidates.append(("TABLE", str(table))) + for view in views: + if not view.temp: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) ) - else: - # Keep the executor's original user, e.g. the Cancel worker - # must stay mz_system or its cancels fail with "must be a - # member of" - user = exe.user - conn = exe.cur.connection + # Kafka and webhook source objects are readable directly, the + # others follow the source-table model: str(obj) names the table, + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ("SOURCE", f"{source.schema}.{identifier(source.source_name())}") + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for schema in schemas: + candidates.append(("SCHEMA", str(schema))) + for db in dbs: + candidates.append(("DATABASE", str(db))) + for cluster in clusters: + candidates.append(("CLUSTER", str(cluster))) + candidates.append(("SECRET", "materialize.public.pgpass")) + # NOTE: No CONNECTION target. Changing a connection's owner emits a + # Connection(Altered) implication, which re-alters every dependent + # sink's export connection. That re-alter can fail with InvalidAlter + # and panic the coordinator. + # See https://linear.app/materializeinc/issue/SQL-517 + kind, name = self.rng.choice(candidates) + with role.lock: + if role not in exe.db.roles: + return False + exe.execute(f"ALTER {kind} {name} OWNER TO {role}", http=Http.RANDOM) + return True - if exe.ws and exe.use_ws: - try: - exe.ws.close() - except: - pass - try: - exe.cur.close() - except: - pass - try: - conn.close() - except: - pass +class AlterDefaultPrivilegesAction(Action): + # Privileges per object type. The bool pair is (allows IN SCHEMA, allows IN + # DATABASE): schema-scoped objects accept both, SCHEMAS only IN DATABASE, + # DATABASES and CLUSTERS neither. Mixing e.g. ON DATABASES with IN DATABASE + # is a plan error, not a race, so it is generated out. + OBJECT_TYPES = { + "TABLES": (["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"], True, True), + "TYPES": (["USAGE", "ALL"], True, True), + "SECRETS": (["USAGE", "ALL"], True, True), + "CONNECTIONS": (["USAGE", "ALL"], True, True), + "SCHEMAS": (["USAGE", "CREATE", "ALL"], False, True), + "DATABASES": (["USAGE", "CREATE", "ALL"], False, False), + "CLUSTERS": (["USAGE", "CREATE", "ALL"], False, False), + } - NUM_ATTEMPTS = 20 + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = [ + "unknown role", + "unknown schema", + "unknown database", + "must be a member of", + # FOR ALL ROLES and system-adjacent grants require privileges the + # (possibly reconnected-as-a-random-role) session may lack. + "permission denied to", + ] + super().errors_to_ignore(exe) + return result + + def run(self, exe: Executor) -> bool: + object_type = self.rng.choice(list(self.OBJECT_TYPES.keys())) + privileges, allows_in_schema, allows_in_database = self.OBJECT_TYPES[ + object_type + ] + privilege = self.rng.choice(privileges) + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + for_clause = self.rng.choice( + ["FOR ALL ROLES"] + [f"FOR ROLE {r}" for r in exe.db.roles] + ) + in_clause = "" + if allows_in_schema and exe.db.schemas and self.rng.random() < 0.3: + in_clause = f" IN SCHEMA {self.rng.choice(exe.db.schemas)}" + elif allows_in_database and exe.db.dbs and self.rng.random() < 0.3: + in_clause = f" IN DATABASE {self.rng.choice(exe.db.dbs)}" + with role.lock: + if role not in exe.db.roles: + return False + if self.rng.choice([True, False]): + query = f"ALTER DEFAULT PRIVILEGES {for_clause}{in_clause} GRANT {privilege} ON {object_type} TO {role}" + else: + query = f"ALTER DEFAULT PRIVILEGES {for_clause}{in_clause} REVOKE {privilege} ON {object_type} FROM {role}" + exe.execute(query, http=Http.RANDOM) + return True + + +class BroadPrivilegesAction(Action): + """GRANT/REVOKE on object classes beyond the tables and views covered by + Grant/RevokePrivilegesAction.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + "unknown schema", + "unknown database", + "unknown cluster", + # System privileges require superuser, which a session reconnected + # as a random role does not have. + "permission denied to", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + targets: list[tuple[str, list[str]]] = [ + ("SYSTEM", ["CREATEDB", "CREATECLUSTER", "CREATEROLE", "ALL"]), + ("SECRET materialize.public.pgpass", ["USAGE", "ALL"]), + # NOTE: No CONNECTION target. GRANT/REVOKE on a connection emits + # a Connection(Altered) implication, which re-alters every + # dependent sink's export connection. That re-alter can fail + # with InvalidAlter and panic the coordinator. + # See https://linear.app/materializeinc/issue/SQL-517 + ] + if exe.db.schemas: + targets.append( + ( + f"SCHEMA {self.rng.choice(exe.db.schemas)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + if exe.db.dbs: + targets.append( + ( + f"DATABASE {self.rng.choice(exe.db.dbs)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + if exe.db.clusters: + targets.append( + ( + f"CLUSTER {self.rng.choice(exe.db.clusters)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + target, privileges = self.rng.choice(targets) + privilege = self.rng.choice(privileges) + with role.lock: + if role not in exe.db.roles: + return False + if self.rng.choice([True, False]): + exe.execute( + f"GRANT {privilege} ON {target} TO {role}", http=Http.RANDOM + ) + else: + exe.execute( + f"REVOKE {privilege} ON {target} FROM {role}", http=Http.RANDOM + ) + return True + + +class ShowAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + # SHOW CREATE CLUSTER only works for managed clusters. We only + # target managed ones, this covers a managed->unmanaged race. + "SHOW CREATE for unmanaged clusters not yet supported", + # With auto_route_catalog_queries off, SHOW compiles to a + # catalog query on the active cluster, so it shares the read + # transaction's timedomain. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + if self.rng.choice([True, False]): + schema_scoped = [ + "SHOW TABLES", + "SHOW VIEWS", + "SHOW MATERIALIZED VIEWS", + "SHOW SOURCES", + "SHOW SINKS", + "SHOW INDEXES", + "SHOW OBJECTS", + "SHOW SECRETS", + "SHOW CONNECTIONS", + "SHOW TYPES", + ] + other = [ + "SHOW CLUSTERS", + "SHOW CLUSTER REPLICAS", + "SHOW DATABASES", + "SHOW SCHEMAS", + "SHOW ROLES", + "SHOW PRIVILEGES", + "SHOW DEFAULT PRIVILEGES", + "SHOW ROLE MEMBERSHIP", + "SHOW ALL", + ] + if self.rng.choice([True, False]): + query = self.rng.choice(schema_scoped) + with exe.db.lock: + if exe.db.schemas and self.rng.choice([True, False]): + query += f" FROM {self.rng.choice(exe.db.schemas)}" + if self.rng.random() < 0.2: + query += " LIKE '%1%'" + else: + query = self.rng.choice(other) + else: + # Only the snapshot needs the global lock, see CommentAction. + with exe.db.lock: + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + indexes = list(exe.db.indexes) + clusters = list(exe.db.clusters) + + candidates: list[tuple[str, str]] = [ + ("CONNECTION", "materialize.public.kafka_conn"), + ("CONNECTION", "materialize.public.csr_conn"), + ] + for table in tables: + candidates.append(("TABLE", str(table))) + for view in views: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) + ) + # Kafka and webhook sources are readable directly. The others + # follow the source-table model where str() is the table and + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.source_name())}", + ) + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for index in indexes: + candidates.append(("INDEX", str(index))) + for cluster in clusters: + # SHOW CREATE CLUSTER is not supported for unmanaged + # clusters. + if cluster.managed: + candidates.append(("CLUSTER", str(cluster))) + kind, name = self.rng.choice(candidates) + # SHOW REDACTED CREATE CLUSTER is not supported + redacted = ( + "REDACTED " + if kind != "CLUSTER" and self.rng.choice([True, False]) + else "" + ) + query = f"SHOW {redacted}CREATE {kind} {name}" + exe.execute(query, http=Http.RANDOM, fetch=True) + return True + + +class SetSessionVariableAction(Action): + def __init__(self, rng: random.Random, composition: Composition | None): + super().__init__(rng, composition) + self.vars_with_values: dict[str, list[str]] = { + "statement_timeout": ["'30s'", "'60s'", "'0s'"], + "application_name": ["'parallel-workload'", "''"], + "client_min_messages": ["debug1", "info", "notice", "warning", "error"], + "max_query_result_size": ["100000", "1000000", "1000000000"], + "emit_timestamp_notice": ["true", "false"], + "emit_trace_id_notice": ["true", "false"], + # Only UTC is accepted, the rejection of other time zones is + # deliberate error path coverage. + "timezone": ["'UTC'", "'America/New_York'"], + } + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "invalid value for parameter", + "cannot have value", + "unrecognized configuration parameter", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + var = self.rng.choice(list(self.vars_with_values.keys())) + # statement_timeout is tracked on the executor, and only a statement on + # the worker's own session can change what that session times out at. + # Over the HTTP endpoint the SET or RESET lands on a one-shot session + # instead, leaving the tracking describing a value nothing has. + http = Http.NO if var == "statement_timeout" else Http.RANDOM + if self.rng.random() < 0.2: + exe.execute(f"RESET {var}", http=http) + if var == "statement_timeout": + # Back at the default timeout, so a timeout is a genuine + # finding again. + exe.statement_timeout_set = False + return True + value = self.rng.choice(self.vars_with_values[var]) + local = "LOCAL " if self.rng.random() < 0.1 else "" + exe.execute(f"SET {local}{var} = {value}", http=http) + if var == "statement_timeout": + # A timeout is expected for as long as this one is configured. A + # SET LOCAL keeps this set past the transaction that discards the + # value, which only errs towards tolerating more. + exe.statement_timeout_set = True + return True + + +class DiscardAction(Action): + def run(self, exe: Executor) -> bool: + # The session's temp objects die with DISCARD, drop them from the + # tracked state so other workers stop querying them, mirroring + # ReconnectAction. + if exe.temp_objects: + with exe.db.lock: + exe.db.tables[:] = [ + t for t in exe.db.tables if t not in exe.temp_objects + ] + exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] + exe.temp_objects.clear() + # Only DISCARD TEMP, not DISCARD ALL. DISCARD ALL (like DEALLOCATE ALL) + # deallocates every prepared statement, including the ones psycopg + # transparently auto-prepares. psycopg's client-side cache would then + # be stale and the next reuse of such a statement fails with + # "prepared statement ... does not exist". DISCARD TEMP still exercises + # the temp-object teardown, which is the interesting path here. + exe.execute("DISCARD TEMP", http=Http.NO) + return True + + +class ValidateConnectionAction(Action): + # aws_conn is excluded, MinIO's STS support for validation is unclear. + CONNECTIONS = [ + "kafka_conn", + "csr_conn", + "postgres_conn", + "mysql_conn", + "sql_server_conn", + "polaris_conn", + ] + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "timeout: error trying to connect", + # A concurrent ALTER SECRET rotation (the workload only rotates a + # secret to its own value) can transiently expose an empty secret, + # so VALIDATE CONNECTION sends an empty password and the upstream + # system rejects it (secret-rotation atomicity). The same race + # surfaces with each upstream's own wording. + # TODO: Remove when https://linear.app/materializeinc/issue/SS-347 + # is fixed. + "empty password returned by client", # Postgres + "Access denied for user", # MySQL + "Login failed for user", # SQL Server + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + name = self.rng.choice(self.CONNECTIONS) + exe.execute(f"VALIDATE CONNECTION materialize.public.{name}", http=Http.NO) + return True + + +class AlterConnectionAction(Action): + # The SET clause per connection, setting the option to the value the + # connection already has. That still exercises the full reconfiguration + # path (restarting dependent sources and sinks) without breaking them. + # NOTE: BROKER takes no `=` (it is parsed specially), HOST/URL do. + SET_CLAUSES = { + "kafka_conn": "BROKER 'kafka:9092'", + "csr_conn": "URL = 'http://schema-registry:8081'", + "postgres_conn": "HOST = 'postgres'", + "mysql_conn": "HOST = 'mysql'", + "sql_server_conn": "HOST = 'sql-server'", + } + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "timeout: error trying to connect", + # The storage controller can refuse an in-place connection change + # depending on the connection's current state or dependents. + "cannot be altered in the requested way", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + name = self.rng.choice(list(self.SET_CLAUSES.keys())) + set_clause = self.SET_CLAUSES[name] + query = f"ALTER CONNECTION materialize.public.{name} SET ({set_clause})" + if self.rng.choice([True, False]): + validate = self.rng.choice(["true", "false"]) + query += f" WITH (VALIDATE = {validate})" + exe.execute(query, http=Http.RANDOM) + return True + + +class AlterSecretAction(Action): + def run(self, exe: Executor) -> bool: + # Rotate to the same value, exercising the rotation path (including + # dependent connections picking up the new secret version) without + # breaking the credentials. + name, value = self.rng.choice( + [ + ("pgpass", "postgres"), + ("mypass", MySql.DEFAULT_ROOT_PASSWORD), + ("sql_server_pass", SqlServer.DEFAULT_SA_PASSWORD), + ("minio", "minioadmin"), + ] + ) + exe.execute( + f"ALTER SECRET materialize.public.{name} AS '{value}'", http=Http.RANDOM + ) + return True + + +# TODO: Should factor this out so can easily use it without action +class ReconnectAction(Action): + def __init__( + self, + rng: random.Random, + composition: Composition | None, + random_role: bool = True, + ): + super().__init__(rng, composition) + self.random_role = random_role + + def run(self, exe: Executor) -> bool: + exe.mz_service = "materialized" + exe.log("reconnecting") + # The connection's temp objects die with it, drop them from the + # tracked state so other workers stop querying them. + if exe.temp_objects: + with exe.db.lock: + exe.db.tables[:] = [ + t for t in exe.db.tables if t not in exe.temp_objects + ] + exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] + exe.temp_objects.clear() + host = exe.db.host + + def pg_port() -> int: + # System workers (e.g. the Cancel worker) live on the internal + # port, everyone else on the external one of the current service. + if exe.user == "mz_system": + return exe.db.ports[ + "mz_system" if exe.mz_service == "materialized" else "mz_system2" + ] + return exe.db.ports[exe.mz_service] + + with exe.db.lock: + if self.random_role and exe.db.roles: + user = self.rng.choice( + ["materialize", str(self.rng.choice(exe.db.roles))] + ) + else: + # Keep the executor's original user, e.g. the Cancel worker + # must stay mz_system or its cancels fail with "must be a + # member of" + user = exe.user + conn = exe.cur.connection + + if exe.ws and exe.use_ws: + try: + exe.ws.close() + except: + pass + + try: + exe.cur.close() + except: + pass + try: + conn.close() + except: + pass + + NUM_ATTEMPTS = 20 if exe.ws: for i in range( NUM_ATTEMPTS @@ -2999,6 +4626,10 @@ def pg_port() -> int: exe.set_isolation("SERIALIZABLE") cur.execute("SET auto_route_catalog_queries TO false") conn.autocommit = exe.autocommit + # A statement_timeout set on the old session died with it, so + # timeouts are genuine findings again unless a role default + # carries one over. + exe.statement_timeout_set = False try: cur.execute("SELECT pg_backend_pid()") except Exception as e: @@ -3073,7 +4704,11 @@ def __init__( def run(self, exe: Executor) -> bool: pid = self.rng.choice( - [worker.exe.pg_pid for worker in self.workers if worker.exe and worker.exe.pg_pid != -1] # type: ignore + [ + worker.exe.pg_pid + for worker in self.workers + if worker.exe and worker.exe.pg_pid != -1 + ] # type: ignore ) worker = None for i in range(len(self.workers)): @@ -3303,6 +4938,125 @@ def run(self, exe: Executor) -> bool: return True +class CreateLoadGeneratorSourceAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.loadgen_sources) >= MAX_LOADGEN_SOURCES: + return False + source_id = exe.db.loadgen_source_id + exe.db.loadgen_source_id += 1 + try: + cluster = self.rng.choice(exe.db.clusters) + schema = self.rng.choice(exe.db.schemas) + except IndexError: + # We mostly prevent index errors, but we don't want to lock too + # much since that would reduce our chance of finding race + # conditions in production code, so ignore the rare case where + # we accidentally removed all objects. + return False + with schema.lock, cluster.lock: + if schema not in exe.db.schemas: + return False + if cluster not in exe.db.clusters: + return False + + source = LoadGeneratorSource(source_id, cluster, schema, self.rng) + source.create(exe) + exe.db.loadgen_sources.append(source) + return True + + +class DropLoadGeneratorSourceAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.loadgen_sources: + return False + try: + source = self.rng.choice(exe.db.loadgen_sources) + except IndexError: + # We mostly prevent index errors, but we don't want to lock too + # much since that would reduce our chance of finding race + # conditions in production code, so ignore the rare case where + # we accidentally removed all objects. + return False + with source.lock: + # Was dropped while we were acquiring lock + if source not in exe.db.loadgen_sources: + return False + + exe.execute(f"DROP TABLE IF EXISTS {source}") + exe.execute( + f"DROP SOURCE {source.schema}.{identifier(source.source_name())} CASCADE", + http=Http.RANDOM, + ) + exe.db.loadgen_sources.remove(source) + return True + + +class CreateMultiLoadGeneratorSourceAction(Action): + GENERATORS = ["AUCTION", "TPCH", "MARKETING"] + + def errors_to_ignore(self, exe: Executor) -> list[str]: + # A rare cross-type subsource-name collision, or a concurrent create of + # the same generator type. + return ["already exists"] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + present = {s.generator for s in exe.db.multi_loadgen_sources} + available = [g for g in self.GENERATORS if g not in present] + if not available: + return False + generator = self.rng.choice(available) + source_id = exe.db.multi_loadgen_source_id + exe.db.multi_loadgen_source_id += 1 + try: + cluster = self.rng.choice(exe.db.clusters) + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + with schema.lock, cluster.lock: + if schema not in exe.db.schemas: + return False + if cluster not in exe.db.clusters: + return False + source = MultiLoadGeneratorSource( + source_id, cluster, schema, generator, self.rng + ) + source.create(exe) + # NOTE: No db.lock around the append. Taking it here would nest it + # inside the schema and cluster locks, inverting the db.lock-first + # order every other action uses. list.append is atomic, and a + # concurrent same-generator create is caught server-side ("already + # exists", tolerated above). + exe.db.multi_loadgen_sources.append(source) + return True + + +class DropMultiLoadGeneratorSourceAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.multi_loadgen_sources: + return False + source = self.rng.choice(exe.db.multi_loadgen_sources) + with source.lock: + if source not in exe.db.multi_loadgen_sources: + return False + exe.execute(f"DROP SOURCE {source} CASCADE", http=Http.RANDOM) + exe.db.multi_loadgen_sources.remove(source) + return True + + class CreateKafkaSourceAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -3713,7 +5467,7 @@ def run(self, exe: Executor) -> bool: sink_id, cluster, schema, - self.rng.choice(exe.db.db_objects_without_views()), + self.rng.choice(exe.db.db_objects_for_sinks()), self.rng, ) sink.create(exe) @@ -3781,7 +5535,7 @@ def run(self, exe: Executor) -> bool: sink_id, cluster, schema, - self.rng.choice(exe.db.db_objects_without_views()), + self.rng.choice(exe.db.db_objects_for_sinks()), self.rng, ) sink.create(exe) @@ -3823,6 +5577,11 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: result = super().errors_to_ignore(exe) if exe.db.scenario == Scenario.Rename: result.extend(["404: no object was found at the path"]) + # DropSchemaCascadeAction / DropDatabaseCascadeAction can drop a + # webhook source concurrently without taking its per-object lock, so a + # POST that already picked the source can 404. + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["404: no object was found at the path"]) return result def run(self, exe: Executor) -> bool: @@ -3887,6 +5646,383 @@ def run(self, exe: Executor) -> bool: return True +class AlterClusterSetAction(Action): + """Live reconfigure of a managed cluster (SIZE / REPLICATION FACTOR). + + Resizing or changing the replica count of a cluster hosting indexes, MVs, + sources, and sinks forces rehydration and replica teardown/spin-up under + concurrent DDL and DML.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # A SET (SIZE) here or a ReconfigureCluster on the same cluster + # leaves a reconfiguration record in flight past the statement + # that started it. Replication factor is folded in at cut-over, + # so changing it meanwhile is refused. + "cannot change replication factor while a reconfiguration is in progress", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + # Cluster 0 stays fixed, it hosts sources and sinks. + managed = [c for c in exe.db.clusters[1:] if c.managed] + if not managed: + return False + cluster = self.rng.choice(managed) + with cluster.lock: + if cluster not in exe.db.clusters or not cluster.managed: + return False + choice = self.rng.choice(["size", "replication_factor", "reset_rf"]) + if choice == "size": + new_size = self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]) + exe.execute( + f"ALTER CLUSTER {cluster} SET (SIZE = '{new_size}')", + http=Http.RANDOM, + ) + cluster.size = new_size + for replica in cluster.replicas: + replica.size = new_size + elif choice == "replication_factor": + rf = self.rng.choice([1, 2]) + exe.execute( + f"ALTER CLUSTER {cluster} SET (REPLICATION FACTOR = {rf})", + http=Http.RANDOM, + ) + self._resize_replicas(cluster, rf) + else: + exe.execute( + f"ALTER CLUSTER {cluster} RESET (REPLICATION FACTOR)", + http=Http.RANDOM, + ) + self._resize_replicas(cluster, 1) + return True + + def _resize_replicas(self, cluster: Cluster, count: int) -> None: + # Managed cluster replicas are server-named (r1..rN), so the tracked + # list is rebuilt from scratch to match: the ids are what + # `ClusterReplica.name()` derives those names from. + cluster.replicas = [ + ClusterReplica(i, cluster.size, cluster) for i in range(count) + ] + cluster.replica_id = count + + +class DropSchemaCascadeAction(Action): + """DROP SCHEMA .. CASCADE, an atomic multi-object catalog mutation. + + Only enabled in DDL complexity: cross-schema dependents are cascade-dropped + server-side but stay tracked until they surface as "does not exist", which + DDL complexity ignores.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.schemas) <= 1: + return False + schema = self.rng.choice(exe.db.schemas) + # Keep at least two non-temp tables alive outside the dropped + # schema. Query generation picks from the non-view objects, and a + # CASCADE that emptied that set would crash it. This mirrors + # DropTableAction's minimum. + if ( + len([t for t in exe.db.tables if not t.temp and t.schema is not schema]) + < 2 + ): + return False + with schema.lock: + if schema not in exe.db.schemas: + return False + if len(exe.db.schemas) <= 1: + return False + exe.execute(f"DROP SCHEMA {schema} CASCADE", http=Http.RANDOM) + exe.db.schemas.remove(schema) + untrack_objects_in_schemas(exe, {schema}) + return True + + +class DropDatabaseCascadeAction(Action): + """DROP DATABASE .. CASCADE, an atomic multi-object catalog mutation. + + DDL-complexity only, for the same reason as DropSchemaCascadeAction.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.dbs) <= 1: + return False + db = self.rng.choice(exe.db.dbs) + # Keep at least two non-temp tables alive outside the dropped + # database, so query generation always has a non-view object. + if ( + len([t for t in exe.db.tables if not t.temp and t.schema.db is not db]) + < 2 + ): + return False + with db.lock: + if db not in exe.db.dbs: + return False + if len(exe.db.dbs) <= 1: + return False + exe.execute(f"DROP DATABASE {db} CASCADE", http=Http.RANDOM) + exe.db.dbs.remove(db) + with exe.db.lock: + dropped = {s for s in exe.db.schemas if s.db is db} + exe.db.schemas[:] = [s for s in exe.db.schemas if s.db is not db] + untrack_objects_in_schemas(exe, dropped) + return True + + +class CreateTypeAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.types) >= MAX_TYPES: + return False + type_id = exe.db.type_id + exe.db.type_id += 1 + try: + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + with schema.lock: + if schema not in exe.db.schemas: + return False + typ = Type(type_id, schema, self.rng) + typ.create(exe) + exe.db.types.append(typ) + return True + + +class DropTypeAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + "cannot be dropped", + # Another worker (or a CASCADE drop of the schema/database) can + # drop the type first. + "does not exist", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.types: + return False + typ = self.rng.choice(exe.db.types) + with typ.lock: + if typ not in exe.db.types: + return False + exe.execute(f"DROP TYPE {typ}", http=Http.RANDOM) + exe.db.types.remove(typ) + return True + + +class CreateNetworkPolicyAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.network_policies) >= MAX_NETWORK_POLICIES: + return False + policy_id = exe.db.network_policy_id + exe.db.network_policy_id += 1 + policy = NetworkPolicy(policy_id, self.rng) + policy.create(exe) + with exe.db.lock: + exe.db.network_policies.append(policy) + return True + + +class AlterNetworkPolicyAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.network_policies: + return False + policy = self.rng.choice(exe.db.network_policies) + with policy.lock: + if policy not in exe.db.network_policies: + return False + exe.execute( + f"ALTER NETWORK POLICY {policy} SET ({policy.rules_clause()})", + http=Http.RANDOM, + ) + policy.num_rules = self.rng.randint(1, 3) + return True + + +class DropNetworkPolicyAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # The policy is installed as a default somewhere (should not happen, + # we never install ours, but be safe). + "cannot be dropped", + # Another worker dropped the same policy first. The error carries + # the raw name ('netpol-N', not the quoted form), so DROP resolves + # the name correctly. This is a concurrency race, not the ALTER + # NETWORK POLICY quoted-name bug. + "unknown network policy", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.network_policies: + return False + policy = self.rng.choice(exe.db.network_policies) + with policy.lock: + if policy not in exe.db.network_policies: + return False + exe.execute(f"DROP NETWORK POLICY {policy}", http=Http.RANDOM) + exe.db.network_policies.remove(policy) + return True + + +class SystemCatalogReadAction(Action): + """Read a random system-catalog / introspection relation while DDL churns. + + Exercises catalog-read-vs-write consistency and the auto-route path. + Relations that the `materialize` user cannot read (mz_notices, + mz_recent_activity_log_redacted) are left out.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # An introspection view over many objects can outrun + # statement_timeout, never a real bug in a stress run. + "canceling statement due to statement timeout", + # Reading a system-catalog relation inside a read transaction that + # already touched user objects crosses timedomains. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + super().errors_to_ignore(exe) + + RELATIONS = [ + "mz_catalog.mz_objects", + "mz_catalog.mz_columns", + "mz_catalog.mz_indexes", + "mz_catalog.mz_sources", + "mz_catalog.mz_sinks", + "mz_catalog.mz_materialized_views", + "mz_catalog.mz_views", + "mz_catalog.mz_tables", + "mz_catalog.mz_audit_events", + "mz_catalog.mz_databases", + "mz_catalog.mz_schemas", + "mz_catalog.mz_roles", + "mz_catalog.mz_clusters", + "mz_catalog.mz_cluster_replicas", + "mz_internal.mz_frontiers", + "mz_internal.mz_hydration_statuses", + "mz_internal.mz_compute_dependencies", + "mz_internal.mz_source_statuses", + "mz_internal.mz_sink_statuses", + "mz_internal.mz_materialization_lag", + "mz_internal.mz_wallclock_global_lag_recent_history", + "mz_internal.mz_cluster_replica_statuses", + "mz_internal.mz_object_dependencies", + "mz_internal.mz_object_transitive_dependencies", + "mz_internal.mz_show_all_objects", + "mz_internal.mz_comments", + ] + + def run(self, exe: Executor) -> bool: + relation = self.rng.choice(self.RELATIONS) + exe.execute( + f"SELECT * FROM {relation} LIMIT {self.rng.randint(1, 100)}", + http=Http.RANDOM, + fetch=True, + ) + return True + + +class ExplainAnalyzeAction(Action): + """EXPLAIN ANALYZE against a live materialized view or index dataflow. + + Runs generated introspection queries on the active cluster. Racing it + against drop/replace of the target probes the introspection path.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # The active cluster has more than one replica. + "log source reads must target a replica", + "does not exist", + "not been hydrated", + "not been materialized", + # A concurrent DROP/reconfigure of the targeted replica retires the + # introspection query. No panic in services.log, just a race. + "target replica failed or was dropped", + # Introspection over a large dataflow can outrun statement_timeout. + "canceling statement due to statement timeout", + # The generated introspection queries can cross timedomains or be + # routed to mz_catalog_server. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + mvs = [v for v in exe.db.views if v.materialized] + with exe.db.lock: + indexes = list(exe.db.indexes) + candidates: list[tuple[str, DBObject | Index]] = [] + for v in mvs: + candidates.append(("MATERIALIZED VIEW", v)) + for i in indexes: + candidates.append(("INDEX", i)) + if not candidates: + return False + kind, obj = self.rng.choice(candidates) + analysis = self.rng.choice( + ["MEMORY", "CPU", "MEMORY WITH SKEW", "CPU WITH SKEW", "HINTS"] + ) + with obj.lock: + if kind == "MATERIALIZED VIEW" and obj not in exe.db.views: + return False + if kind == "INDEX" and obj not in exe.db.indexes: + return False + exe.execute( + f"EXPLAIN ANALYZE {analysis} FOR {kind} {obj}", + http=Http.NO, + fetch=True, + ) + return True + + +class ExplainFilterPushdownAction(Action): + """EXPLAIN FILTER PUSHDOWN, which inspects durable persist state to compute + which parts a query's filters would read.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + # Scanning persist part stats can outrun statement_timeout. + "canceling statement due to statement timeout", + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + mvs = [v for v in exe.db.views if v.materialized] + if mvs and self.rng.choice([True, False]): + view = self.rng.choice(mvs) + with view.lock: + if view not in exe.db.views: + return False + exe.execute( + f"EXPLAIN FILTER PUSHDOWN FOR MATERIALIZED VIEW {view}", + http=Http.NO, + fetch=True, + ) + else: + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute( + f"EXPLAIN FILTER PUSHDOWN FOR {query}", http=Http.NO, fetch=True + ) + return True + + class SourceSinkStallCheckAction(Action): def applicable(self, exe: Executor) -> bool: return exe.db.scenario not in ( @@ -3924,6 +6060,7 @@ def run(self, exe: Executor) -> bool: ("postgres_sources", exe.db.postgres_sources), ("mysql_sources", exe.db.mysql_sources), ("sql_server_sources", exe.db.sql_server_sources), + ("loadgen_sources", exe.db.loadgen_sources), ("webhook_sources", exe.db.webhook_sources), ]: counts = [] @@ -3935,6 +6072,85 @@ def run(self, exe: Executor) -> bool: return True +class DependencyConsistencyAction(Action): + """Client-side catalog dependency oracle: flag any dependency edge whose + endpoints are not both live objects (a dangling uses/used_by edge). + + PREPARED BUT DISABLED (commented out of read_action_list). This targets the + SQL-521 class: a sink left pointing at a materialized view that was dropped + under cancel, i.e. the same MissingUses inconsistency the coordinator's own + check_consistency asserts on. While SQL-521 is open the workload + legitimately produces such danglers, so enabling this now would re-detect + the known corruption rather than find new bugs. + TODO: enable in read_action_list once SQL-521 is fixed. Verify the + mz_internal.mz_object_dependencies column names (object_id, + referenced_object_id) still hold when enabling.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # Reading a catalog relation inside a read txn that already touched + # user objects crosses timedomains. + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + exe.execute( + "SELECT d.object_id, d.referenced_object_id " + "FROM mz_internal.mz_object_dependencies d " + "LEFT JOIN mz_objects o1 ON d.object_id = o1.id " + "LEFT JOIN mz_objects o2 ON d.referenced_object_id = o2.id " + "WHERE o1.id IS NULL OR o2.id IS NULL", + http=Http.NO, + ) + dangling = exe.cur.fetchall() + if dangling: + raise ValueError( + f"dangling catalog dependency edges (SQL-521 class): {dangling[:5]}" + ) + return True + + +class SourceReadHoldSweepAction(Action): + """Read a source-backed relation to force read-hold acquisition on the + source's remap shard. + + PREPARED BUT DISABLED (commented out of read_action_list). Intended for the + kill scenario, where racing envd/clusterd restarts stresses read-hold + reinstatement: the SS-346 class (a dependent's read hold on a source's remap + shard not upheld across a restart, so its since advances past the + dependent's upper) and PER-49 (a compute import as_of behind the compacted + since after ALTER TABLE ADD COLUMN + kill). Enabling it now just re-triggers + those known coordinator/compute panics. + TODO: enable in read_action_list once SS-346 and PER-49 are fixed.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + sources = [ + o + for o in exe.db.db_objects() + if isinstance( + o, + LoadGeneratorSource + | KafkaSource + | PostgresSource + | MySqlSource + | SqlServerSource + | WebhookSource, + ) + ] + if not sources: + return False + obj = self.rng.choice(sources) + exe.execute(f"SELECT count(*) FROM {obj}", http=Http.RANDOM) + exe.cur.fetchall() + return True + + class ActionList: action_classes: list[type[Action]] weights: list[float] @@ -3952,15 +6168,31 @@ def __init__( [ (SelectAction, 100), (SelectOneAction, 1), + (ParameterizedQueryAction, 20), # (SQLsmithAction, 30), # Questionable use ( CopyToS3Action, 100, ), + (CopyToStdoutAction, 20), + (ShowAction, 10), + (SystemCatalogReadAction, 10), + # TODO: Reenable once EXPLAIN FILTER PUSHDOWN can no longer panic the + # coordinator when a referenced compute collection is concurrently + # dropped. sequence_explain_pushdown -> acquire_read_holds().expect( + # "missing compute collection") at read_policy.rs:389 (normal peeks and + # EXPLAIN ANALYZE handle the drop gracefully). + # See https://linear.app/materializeinc/issue/SQL-519 + # (ExplainFilterPushdownAction, 5), + # PREPARED BUT DISABLED (see class docstrings): enabling these now just + # re-detects known-unfixed coordinator bugs rather than finding new ones. + # (DependencyConsistencyAction, 5), # TODO: enable once SQL-521 fixed + # (SourceReadHoldSweepAction, 5), # TODO: enable once SS-346 & PER-49 fixed (SetClusterAction, 1), (CommitRollbackAction, 30), (ReconnectAction, 1), (FlipFlagsAction, 2), + (SetSessionVariableAction, 2), ], autocommit=False, ) @@ -4006,6 +6238,8 @@ def __init__( (SetClusterAction, 1), (ReconnectAction, 1), (FlipFlagsAction, 2), + (SetSessionVariableAction, 2), + (DiscardAction, 2), # TODO: Reenable when SS-193 and SS-325 are fixed # (SourceSinkStallCheckAction, 4), # (TransactionIsolationAction, 1), @@ -4021,10 +6255,13 @@ def __init__( (DropTableAction, 2), (CreateViewAction, 8), (DropViewAction, 8), + (CreateOrReplaceViewAction, 4), (CreateRoleAction, 2), (DropRoleAction, 2), + (AlterRoleAction, 2), (CreateClusterAction, 1), (DropClusterAction, 1), + (AlterClusterSetAction, 3), (SwapClusterAction, 10), (CreateClusterReplicaAction, 2), (DropClusterReplicaAction, 2), @@ -4038,8 +6275,13 @@ def __init__( (DropIcebergSinkAction, 4), (CreateKafkaSourceAction, 4), (DropKafkaSourceAction, 4), - (CreateMySqlSourceAction, 4), - (DropMySqlSourceAction, 4), + (CreateLoadGeneratorSourceAction, 4), + (DropLoadGeneratorSourceAction, 4), + (CreateMultiLoadGeneratorSourceAction, 2), + (DropMultiLoadGeneratorSourceAction, 2), + # TODO: Reenable when https://linear.app/materializeinc/issue/SS-307 is fixed + # (CreateMySqlSourceAction, 4), + # (DropMySqlSourceAction, 4), (CreatePostgresSourceAction, 4), (DropPostgresSourceAction, 4), # TODO: Reenable when https://linear.app/materializeinc/issue/SS-290 is fixed @@ -4047,18 +6289,64 @@ def __init__( # (DropSqlServerSourceAction, 4), (GrantPrivilegesAction, 4), (RevokePrivilegesAction, 1), + (GrantRoleAction, 2), + (RevokeRoleAction, 1), + (AlterOwnerAction, 2), + (AlterDefaultPrivilegesAction, 2), + (BroadPrivilegesAction, 2), + (ShowAction, 4), + (ValidateConnectionAction, 2), + # TODO: Reenable once altering a connection that sinks or sources depend + # on can no longer panic the coordinator. Re-altering a dependent sink's + # export connection after the txn fails with InvalidAlter, which + # unwrap_or_terminate turns into a panic. + # See https://linear.app/materializeinc/issue/SQL-517 + # (AlterConnectionAction, 2), + (AlterSecretAction, 2), (ReconnectAction, 1), (CreateDatabaseAction, 1), (DropDatabaseAction, 1), + # TODO: Reenable once a concurrent DROP DATABASE CASCADE can no longer + # panic the coordinator. A staged create (e.g. the source executor's + # CREATE SECRET) whose target database is dropped between staging and + # finish hits resolve_full_name -> get_database (panicking OrdMap index) + # in catalog transact_op. Only CASCADE can drop a non-empty database, + # so this is the precise trigger. + # See https://linear.app/materializeinc/issue/SQL-518 + # (DropDatabaseCascadeAction, 1), (CreateSchemaAction, 1), (DropSchemaAction, 1), + (DropSchemaCascadeAction, 1), + (CreateTypeAction, 2), + (DropTypeAction, 2), + (CreateNetworkPolicyAction, 1), + # TODO: Reenable once ALTER NETWORK POLICY resolves quoted (e.g. + # hyphenated) names. It looks the policy up by its quoted display form, + # so it fails with "unknown network policy" for any name that requires + # quoting, even though CREATE and DROP work. + # See https://linear.app/materializeinc/issue/CLO-143 + # (AlterNetworkPolicyAction, 1), + (DropNetworkPolicyAction, 1), (RenameSchemaAction, 10), (RenameTableAction, 10), (RenameViewAction, 10), (RenameKafkaSinkAction, 10), (RenameIcebergSinkAction, 10), (SwapSchemaAction, 10), - (ReplaceMaterializedViewAction, 20), + (CreateReplacementMaterializedViewAction, 10), + (ApplyReplacementMaterializedViewAction, 5), + (DropReplacementMaterializedViewAction, 5), + (SealedCollectionCheckAction, 2), + (TransactionIsolationAction, 1), + (BoundedStalenessReadAction, 2), + (ReadOnlyTransactionAction, 3), + (DDLTransactionAction, 2), + (SystemCatalogReadAction, 4), + (ExplainAnalyzeAction, 4), + # TODO: Reenable with EXPLAIN FILTER PUSHDOWN's coordinator panic on a + # concurrently-dropped compute collection (read_policy.rs:389). + # See https://linear.app/materializeinc/issue/SQL-519 + # (ExplainFilterPushdownAction, 2), (FlipFlagsAction, 2), # TODO: Reenable when https://linear.app/materializeinc/issue/SQL-405 is fixed. # (AlterTableAddColumnAction, 10), diff --git a/misc/python/materialize/parallel_workload/column.py b/misc/python/materialize/parallel_workload/column.py index 9a874d71443c6..9ec25e15de454 100644 --- a/misc/python/materialize/parallel_workload/column.py +++ b/misc/python/materialize/parallel_workload/column.py @@ -88,7 +88,11 @@ def create(self) -> str: return result -class WebhookColumn(Column): +class SourceColumn(Column): + """Column of a source, taking its name from the upstream schema instead of + generating one. The per-source subclasses only exist to keep the column + lists of the different source kinds distinct in type annotations.""" + def __init__( self, name: str, @@ -105,69 +109,25 @@ def name(self, in_query: bool = False) -> str: return identifier(self.raw_name) if in_query else self.raw_name -class KafkaColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object +class WebhookColumn(SourceColumn): + pass - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class KafkaColumn(SourceColumn): + pass -class MySqlColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class LoadGeneratorColumn(SourceColumn): + pass -class PostgresColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object +class MySqlColumn(SourceColumn): + pass - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class PostgresColumn(SourceColumn): + pass -class SqlServerColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class SqlServerColumn(SourceColumn): + pass diff --git a/misc/python/materialize/parallel_workload/database.py b/misc/python/materialize/parallel_workload/database.py index 8920be86fbad6..b3d72d2fbb281 100644 --- a/misc/python/materialize/parallel_workload/database.py +++ b/misc/python/materialize/parallel_workload/database.py @@ -47,6 +47,7 @@ from materialize.parallel_workload.column import ( Column, KafkaColumn, + LoadGeneratorColumn, MySqlColumn, PostgresColumn, SqlServerColumn, @@ -60,7 +61,18 @@ MAX_COLUMNS = 50 MAX_INCLUDE_HEADERS = 5 -MAX_ROWS = 500 +# The row count a view's join squares. A view over two tables joins them on a +# random boolean predicate, so the join is a cross product and every read of +# that view builds MAX_ROWS**2 wide intermediate rows, whatever LIMIT the +# reading query carries. Measured at 500 with only four columns, one +# `COPY (SELECT * FROM view WHERE .. LIMIT 100) TO 's3://..'` cost ~450 MiB of +# replica RSS, and pw tables carry up to MAX_COLUMNS of them. A handful of such +# reads at once is what grew a clusterd to 19 GiB and had the kernel OOM-killer +# take it, plus environmentd and the Kafka broker, out of the one cgroup all of +# a run's containers share (nightlies 17660-17701). 100 keeps enough rows for +# the DML, persist and index paths to be interesting while cutting the +# intermediate 25-fold. +MAX_ROWS = 100 MAX_CLUSTERS = 4 MAX_CLUSTER_REPLICAS = 2 MAX_DBS = 50 @@ -74,8 +86,11 @@ MAX_MYSQL_SOURCES = 50 MAX_SQL_SERVER_SOURCES = 50 MAX_POSTGRES_SOURCES = 50 +MAX_LOADGEN_SOURCES = 50 MAX_KAFKA_SINKS = 50 MAX_ICEBERG_SINKS = 50 +MAX_TYPES = 50 +MAX_NETWORK_POLICIES = 30 MAX_INITIAL_DBS = 1 MAX_INITIAL_SCHEMAS = 1 @@ -88,6 +103,7 @@ MAX_INITIAL_MYSQL_SOURCES = 1 MAX_INITIAL_SQL_SERVER_SOURCES = 1 MAX_INITIAL_POSTGRES_SOURCES = 1 +MAX_INITIAL_LOADGEN_SOURCES = 1 class BodyFormat(Enum): @@ -165,6 +181,13 @@ def name(self) -> str: class DBObject: columns: list[Column] lock: threading.Lock + # Whether reading from this object can legitimately reach the empty + # (sealed) frontier: its own shard seals in normal operation, or a + # dataflow reading it sees an input frontier that becomes empty. Plain + # tables and unbounded sources never seal, so the default is False. + # Bounded (UP TO) load generators seal once they finish, and views + # propagate sealing from their inputs, materialized or not. + can_seal: bool = False def __init__(self): self.lock = threading.Lock() @@ -283,8 +306,9 @@ def __init__( rng.choice( [ "ON COMMIT", - f"EVERY '{rng.randint(1, 60)} seconds {rng.randint(0, 60)} minutes'", - f"EVERY '{rng.randint(1, 60)} seconds {rng.randint(0, 60)} minutes' ALIGNED TO (mz_now())", + # TODO: Restore minute-scale intervals when CPU-196 is fixed + f"EVERY '{rng.randint(1, 15)} seconds'", + f"EVERY '{rng.randint(1, 15)} seconds' ALIGNED TO (mz_now())", # Always in the future of all refreshes of previously generated MVs "AT mz_now()::string::int8 + 1000", ] @@ -293,6 +317,21 @@ def __init__( else None ) + # A materialized view's shard seals (its write frontier advances to the + # empty frontier) once it can never produce more output. That happens + # for REFRESH AT views after their last refresh, for repeat_row(-1) + # constant views on hydration, and transitively for any view that reads + # from an input that itself seals. The replacement and sealed-shard + # oracles key off this to tell legitimate seals from wrongly finalized + # shards. Unmaterialized views have no shard, but a dataflow reading + # one inlines its inputs, so sealing must propagate through them too. + self.can_seal = ( + (self.refresh or "").startswith("AT") + or self.repeat_row_const + or base_object.can_seal + or (base_object2 is not None and base_object2.can_seal) + ) + if base_object2: self.on_expr = expression( Boolean, all_columns, rng, kind=ExprKind.MATERIALIZABLE @@ -370,8 +409,12 @@ def select_str(exprs: str) -> str: return query - def create(self, exe: Executor) -> None: + def create(self, exe: Executor, or_replace: bool = False) -> None: query = "CREATE " + # OR REPLACE keeps the same catalog item, swapping its definition and + # rebuilding the dataflow. It is incompatible with TEMP. + if or_replace and not self.temp: + query += "OR REPLACE " if self.temp: query += "TEMP " if self.materialized: @@ -576,14 +619,19 @@ def __init__( self.cluster = cluster self.schema = schema self.base_object = base_object - key_cols = [ - column - for column in rng.sample( - base_object.columns, k=rng.randint(1, len(base_object.columns)) - ) - ] - key_col_names = [column.name(True) for column in key_cols] - self.key = f"KEY ({', '.join(key_col_names)}) NOT ENFORCED" + self.mode = rng.choice(["UPSERT", "APPEND"]) + if self.mode == "UPSERT": + key_cols = [ + column + for column in rng.sample( + base_object.columns, k=rng.randint(1, len(base_object.columns)) + ) + ] + key_col_names = [column.name(True) for column in key_cols] + self.key = f"KEY ({', '.join(key_col_names)}) NOT ENFORCED" + else: + # APPEND mode does not permit a KEY. + self.key = "" self.table_name = f"icesink_topic{self.sink_id}_{uuid.uuid4().hex[:8]}" self.rename = 0 @@ -596,7 +644,7 @@ def __str__(self) -> str: return f"{self.schema}.{identifier(self.name())}" def create(self, exe: Executor) -> None: - query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO ICEBERG CATALOG CONNECTION polaris_conn (NAMESPACE 'default_namespace', TABLE '{self.table_name}') USING AWS CONNECTION aws_conn {self.key} MODE UPSERT WITH (COMMIT INTERVAL '1s')" + query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO ICEBERG CATALOG CONNECTION polaris_conn (NAMESPACE 'default_namespace', TABLE '{self.table_name}') USING AWS CONNECTION aws_conn {self.key} MODE {self.mode} WITH (COMMIT INTERVAL '1s')" exe.execute(query) @@ -608,6 +656,8 @@ class KafkaSink(DBObject): base_object: DBObject envelope: str key: str + connection_options: list[str] + no_snapshot: bool def __init__( self, @@ -622,6 +672,19 @@ def __init__( self.cluster = cluster self.schema = schema self.base_object = base_object + self.connection_options = [] + if rng.random() < 0.3: + compression = rng.choice(["none", "gzip", "lz4", "zstd", "snappy"]) + self.connection_options.append(f"COMPRESSION TYPE = '{compression}'") + if rng.random() < 0.2: + self.connection_options.append( + f"TRANSACTIONAL ID PREFIX 'pw-txn-{self.sink_id}'" + ) + if rng.random() < 0.2: + self.connection_options.append( + f"PROGRESS GROUP ID PREFIX 'pw-progress-{self.sink_id}'" + ) + self.no_snapshot = rng.random() < 0.2 universal_formats = [ "FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION csr_conn", "FORMAT JSON", @@ -681,7 +744,10 @@ def create(self, exe: Executor) -> None: if self.partition_count else "" ) - query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO KAFKA CONNECTION kafka_conn (TOPIC {topic}{maybe_partition}) {self.key} {self.format} ENVELOPE {self.envelope}" + options = "".join(f", {option}" for option in self.connection_options) + query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO KAFKA CONNECTION kafka_conn (TOPIC {topic}{maybe_partition}{options}) {self.key} {self.format} ENVELOPE {self.envelope}" + if self.no_snapshot: + query += " WITH (SNAPSHOT = false)" exe.execute(query) @@ -861,6 +927,110 @@ def create(self, exe: Executor) -> None: self.executor.create(logging_exe=exe) +class LoadGeneratorSource(DBObject): + """A COUNTER load generator source. Always bounded by UP TO so it cannot + grow without limit, which also exercises the finished-source lifecycle + state. The readable object (str(self)) is the table created from the + source, matching the source-table model the other sources use.""" + + # A finished (UP TO reached) counter advances the source table's frontier + # to the empty antichain, so anything reading it seals legitimately. + can_seal = True + + source_id: int + cluster: "Cluster" + schema: Schema + columns: list[LoadGeneratorColumn] + tick_interval: str + up_to: int + + def __init__( + self, + source_id: int, + cluster: "Cluster", + schema: Schema, + rng: random.Random, + ): + super().__init__() + self.source_id = source_id + self.cluster = cluster + self.schema = schema + self.tick_interval = rng.choice(["10ms", "100ms", "1s"]) + # Bounded near MAX_ROWS: this source-table is a readable relation that + # views/MVs join over, and unlike user tables it is not capped by + # MAX_ROWS. A large counter feeding a maintained join is a clusterd + # memory risk, so keep it in the same size class as the other data. + self.up_to = rng.randint(1, MAX_ROWS) + self.columns = [LoadGeneratorColumn("counter", Long, False, self)] + + def name(self) -> str: + return naughtify(f"lg-{self.source_id}") + + def source_name(self) -> str: + return naughtify(f"lg-src-{self.source_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + source = f"{self.schema}.{identifier(self.source_name())}" + exe.execute( + f"CREATE SOURCE {source} IN CLUSTER {self.cluster} FROM LOAD GENERATOR COUNTER (TICK INTERVAL '{self.tick_interval}', UP TO {self.up_to})" + ) + exe.execute(f"CREATE TABLE {self} FROM SOURCE {source}") + + +class MultiLoadGeneratorSource: + """A multi-subsource load generator (AUCTION / TPCH / MARKETING) created + with FOR ALL TABLES. Tracked create/drop-only, not as a readable relation: + its subsources have fixed names and hardcoded schemas, so wiring them into + the general read/expression machinery is out of scope. The value is the + multi-subsource source lifecycle and CASCADE teardown under concurrency. + + At most one source per generator type exists at a time, because FOR ALL + TABLES names its subsources fixedly (a second AUCTION would collide on + `auctions`, `bids`, ...).""" + + source_id: int + cluster: "Cluster" + schema: Schema + generator: str + tick_interval: str + lock: threading.Lock + + def __init__( + self, + source_id: int, + cluster: "Cluster", + schema: Schema, + generator: str, + rng: random.Random, + ): + self.source_id = source_id + self.cluster = cluster + self.schema = schema + self.generator = generator + self.tick_interval = rng.choice(["100ms", "1s"]) + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"mlg-{self.source_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + options = [f"TICK INTERVAL '{self.tick_interval}'"] + if self.generator == "TPCH": + options.append("SCALE FACTOR 0.0001") + query = ( + f"CREATE SOURCE {self} IN CLUSTER {self.cluster} " + f"FROM LOAD GENERATOR {self.generator} ({', '.join(options)}) " + f"FOR ALL TABLES" + ) + exe.execute(query) + + class S3Object(DBObject): """A COPY TO dump of `table` that CopyFromS3Action can load back into it. @@ -1066,6 +1236,11 @@ def __str__(self) -> str: def create(self, exe: Executor) -> None: exe.execute(f"CREATE ROLE {self}") + # Make the workload's own user (materialize) a member, so it can + # manage the role, e.g. transfer object ownership to it: ALTER .. + # OWNER TO requires the executor to be a member of . + # The creating session has admin on the role it just created. + exe.execute(f"GRANT {self} TO materialize") class ClusterReplica: @@ -1083,6 +1258,12 @@ def __init__(self, replica_id: int, size: str, cluster: "Cluster"): self.lock = threading.Lock() def name(self) -> str: + # A managed cluster's replicas are named by the controller as r1..rN, + # never by us, so neither the rename counter nor naughtify applies. + # Rendering our own name there would make every reference to the + # replica, e.g. a replica-targeted materialized view, fail to resolve. + if self.cluster.managed: + return f"r{self.replica_id+1}" if self.rename: return naughtify(f"r-{self.replica_id+1}-{self.rename}") return naughtify(f"r-{self.replica_id+1}") @@ -1147,6 +1328,78 @@ def create(self, exe: Executor) -> None: exe.execute(query) +class Type: + """A user-defined type (row, list, or map). Schema-scoped.""" + + type_id: int + schema: Schema + kind: str + lock: threading.Lock + rng: random.Random + + def __init__(self, type_id: int, schema: Schema, rng: random.Random): + self.type_id = type_id + self.schema = schema + self.kind = rng.choice(["row", "list", "map"]) + self.rng = rng + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"type-{self.type_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + if self.kind == "row": + scalar_types = ["int4", "text", "bool", "float8", "timestamp"] + fields = ", ".join( + f"f{i} {self.rng.choice(scalar_types)}" + for i in range(self.rng.randint(1, 4)) + ) + query = f"CREATE TYPE {self} AS ({fields})" + elif self.kind == "list": + element = self.rng.choice(["int4", "text", "float8"]) + query = f"CREATE TYPE {self} AS LIST (ELEMENT TYPE = {element})" + else: + value = self.rng.choice(["int4", "text", "float8"]) + query = f"CREATE TYPE {self} AS MAP (KEY TYPE = text, VALUE TYPE = {value})" + exe.execute(query) + + +class NetworkPolicy: + """A network policy. Top-level (not schema-scoped), like clusters. + + Rules are always allow-all so the workload can never lock itself out of + its own connections. The policy is never installed as the active + `network_policy` system parameter for the same reason.""" + + policy_id: int + num_rules: int + lock: threading.Lock + + def __init__(self, policy_id: int, rng: random.Random): + self.policy_id = policy_id + self.num_rules = rng.randint(1, 3) + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"netpol-{self.policy_id}") + + def __str__(self) -> str: + return identifier(self.name()) + + def rules_clause(self) -> str: + rules = ", ".join( + f"r{i} (action='allow', direction='ingress', address='0.0.0.0/0')" + for i in range(self.num_rules) + ) + return f"RULES ({rules})" + + def create(self, exe: Executor) -> None: + exe.execute(f"CREATE NETWORK POLICY {self} ({self.rules_clause()})") + + # TODO: Can access both databases from same connection! class Database: complexity: Complexity @@ -1176,10 +1429,18 @@ class Database: postgres_source_id: int sql_server_sources: list[SqlServerSource] sql_server_source_id: int + loadgen_sources: list[LoadGeneratorSource] + loadgen_source_id: int + multi_loadgen_sources: list[MultiLoadGeneratorSource] + multi_loadgen_source_id: int iceberg_sinks: list[IcebergSink] iceberg_sink_id: int kafka_sinks: list[KafkaSink] kafka_sink_id: int + types: list[Type] + type_id: int + network_policies: list[NetworkPolicy] + network_policy_id: int s3_path: int s3_objects: list[S3Object] read_then_write_counter: ReadThenWriteCounter @@ -1243,9 +1504,7 @@ def __init__( Cluster( i, managed=rng.choice([True, False]), - size=rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), replication_factor=1, introspection_interval="1s", ) @@ -1262,6 +1521,12 @@ def __init__( self.mysql_sources = [] self.postgres_sources = [] self.sql_server_sources = [] + self.loadgen_sources = [ + LoadGeneratorSource( + i, rng.choice(self.clusters), rng.choice(self.schemas), rng + ) + for i in range(rng.randint(0, MAX_INITIAL_LOADGEN_SOURCES)) + ] self.iceberg_sinks = [] self.kafka_sinks = [] self.s3_objects = [] @@ -1269,9 +1534,16 @@ def __init__( self.mysql_source_id = len(self.mysql_sources) self.postgres_source_id = len(self.postgres_sources) self.sql_server_source_id = len(self.sql_server_sources) + self.loadgen_source_id = len(self.loadgen_sources) + self.multi_loadgen_sources = [] + self.multi_loadgen_source_id = 0 self.iceberg_sink_id = len(self.iceberg_sinks) self.kafka_sink_id = len(self.kafka_sinks) self.read_then_write_counter = ReadThenWriteCounter() + self.types = [] + self.type_id = 0 + self.network_policies = [] + self.network_policy_id = 0 self.lock = threading.Lock() self.sqlsmith_state = "" self.flags = {} @@ -1284,6 +1556,7 @@ def db_objects( | PostgresSource | SqlServerSource | KafkaSource + | LoadGeneratorSource | View | Table ]: @@ -1300,6 +1573,7 @@ def db_objects( + self.mysql_sources + self.postgres_sources + self.sql_server_sources + + self.loadgen_sources + self.webhook_sources ) @@ -1311,6 +1585,7 @@ def db_objects_without_views( | PostgresSource | SqlServerSource | KafkaSource + | LoadGeneratorSource | View | Table ]: @@ -1318,6 +1593,28 @@ def db_objects_without_views( obj for obj in self.db_objects() if type(obj) != View or obj.materialized ] + def db_objects_for_sinks( + self, + ) -> list[ + WebhookSource + | MySqlSource + | PostgresSource + | SqlServerSource + | KafkaSource + | View + | Table + ]: + """Objects usable as a sink's input (base object or ALTER SINK SET + FROM target). Load generator source tables are excluded: an + ALTER SINK .. SET FROM one can trigger the sink stall of + https://linear.app/materializeinc/issue/SS-344, which is worse for a + continuously-producing input.""" + return [ + obj + for obj in self.db_objects_without_views() + if not isinstance(obj, LoadGeneratorSource) + ] + def __iter__(self): """Returns all relations""" return ( @@ -1342,6 +1639,13 @@ def create(self, exe: Executor, composition: Composition) -> None: for row in exe.cur.fetchall(): exe.execute(f"DROP ROLE {identifier(row[0])}") + # Network policies survive restarts and are top-level, so leftovers + # from a killed run have to be swept before recreating. + exe.execute("SELECT name FROM mz_internal.mz_network_policies") + for row in exe.cur.fetchall(): + if row[0].startswith("netpol-"): + exe.execute(f"DROP NETWORK POLICY {identifier(row[0])}") + print("Creating connections") exe.execute( diff --git a/misc/python/materialize/parallel_workload/executor.py b/misc/python/materialize/parallel_workload/executor.py index 052f9c0ab8dc9..11bbaf6c052e0 100644 --- a/misc/python/materialize/parallel_workload/executor.py +++ b/misc/python/materialize/parallel_workload/executor.py @@ -53,10 +53,16 @@ class Executor: reconnect_next: bool rollback_next: bool last_log: str + # "running" exactly while this session waits on the server. Every path that + # makes a round trip has to set it, the end-of-run wedge check reads it to + # tell a server-side hang from a worker stuck in the workload's own code. last_status: str action_run_since_last_commit_rollback: bool autocommit: bool user: str + # The session's transaction isolation, as last set through set_isolation. + # Actions that set an isolation transiently restore this one afterwards. + isolation: str def __init__( self, @@ -84,9 +90,17 @@ def __init__( self.use_ws = self.rng.choice([True, False]) if self.ws else False self.autocommit = cur.connection.autocommit self.mz_service = "materialized" + # Set while a non-default statement_timeout is configured on this + # session, statement timeouts are expected errors then. Cleared again + # on RESET and on reconnect, otherwise a hang in this worker stays + # invisible for the rest of the run. + self.statement_timeout_set = False + # Materialize's default, until the worker sets one on connect. + self.isolation = "STRICT SERIALIZABLE" def set_isolation(self, level: str) -> None: self.execute(f"SET TRANSACTION_ISOLATION TO '{level}'") + self.isolation = level def commit(self, http: Http = Http.RANDOM) -> None: self._end_transaction("commit", http) @@ -97,6 +111,7 @@ def rollback(self, http: Http = Http.RANDOM) -> None: def _end_transaction(self, command: str, http: Http) -> None: self.insert_table = None self.log(command) + self.last_status = "running" ws_error = None try: # When this executor uses the WS session, statements executed with @@ -118,6 +133,8 @@ def _end_transaction(self, command: str, http: Http) -> None: raise except Exception as e: raise QueryError(str(e), command) + finally: + self.last_status = "finished" if ws_error is not None: raise ws_error # TODO(def-): Enable when things are stable @@ -182,6 +199,7 @@ def copy( ) -> None: query += ";" self.log(f"{query} ({rows})") + self.last_status = "running" try: try: @@ -195,6 +213,22 @@ def copy( finally: self.last_status = "finished" + def copy_to_stdout(self, query: str) -> None: + query += ";" + self.log(query) + self.last_status = "running" + try: + try: + with self.cur.copy(query.encode()) as copy: + for _ in copy: + pass + except Exception as e: + raise QueryError(str(e), query) + + self.action_run_since_last_commit_rollback = True + finally: + self.last_status = "finished" + def execute( self, query: str, @@ -207,7 +241,37 @@ def execute( http == Http.RANDOM and self.rng.choice([True, False]) ) or http == Http.YES if explainable and self.rng.choice([True, False]): - query = f"EXPLAIN OPTIMIZED PLAN AS VERBOSE TEXT FOR {query}" + if self.rng.random() < 0.1: + as_json = " AS JSON" if self.rng.choice([True, False]) else "" + query = f"EXPLAIN TIMESTAMP{as_json} FOR {query}" + else: + stage = self.rng.choice( + [ + "RAW PLAN", + "DECORRELATED PLAN", + "LOCALLY OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "PHYSICAL PLAN", + ] + ) + modifiers = "" + if stage == "OPTIMIZED PLAN" and self.rng.random() < 0.3: + mods = self.rng.sample( + [ + "arity", + "join implementations", + "keys", + "types", + "humanized expressions", + "redacted", + ], + self.rng.randint(1, 3), + ) + modifiers = f" WITH ({', '.join(mods)})" + format = self.rng.choice(["VERBOSE TEXT", "TEXT", "JSON"]) + query = f"EXPLAIN {stage}{modifiers} AS {format} FOR {query}" query += ";" extra_info_str = f" ({extra_info})" if extra_info else "" use_ws = self.use_ws and http != Http.NO diff --git a/misc/python/materialize/parallel_workload/expression.py b/misc/python/materialize/parallel_workload/expression.py index 72daf7e02d252..0d75239f02f16 100644 --- a/misc/python/materialize/parallel_workload/expression.py +++ b/misc/python/materialize/parallel_workload/expression.py @@ -45,6 +45,7 @@ from materialize.parallel_workload.column import ( Column, KafkaColumn, + LoadGeneratorColumn, MySqlColumn, PostgresColumn, SqlServerColumn, @@ -340,11 +341,33 @@ def __init__(self, text: str, params: list, unsupported: ExprKind = ExprKind.ALL # ] +# Edge-case literals injected at low probability as leaves of read/filter +# expressions (ExprKind.ALL only), to surface NaN/Infinity handling in +# arithmetic, comparisons, and aggregation. The generator's `Float` is float4, +# so these MUST be float4: a wider literal (float8/int8) injected where a float4 +# is expected breaks function overload resolution (e.g. round(numeric, bigint)), +# producing spurious failures rather than findings. +# +# Scope is deliberately narrow: NaN/Infinity are the genuinely-new coverage +# (random_value never produces them). Numeric/int boundary and overflow values +# are NOT injected here, because the existing LARGE record size already +# generates overflow-inducing magnitudes and mis-widthed literals only break +# overload resolution. Extreme-year date/timestamp literals are also omitted: +# a 6-digit year deterministically trips the known SS-345 date-parser bug in +# every context, spamming a known-unfixed issue rather than finding new ones. +EDGE_VALUES: dict[type, list[str]] = { + Float: ["'NaN'::float4", "'Infinity'::float4", "'-Infinity'::float4"], +} + + def expression( data_type: type[DataType], columns: list[Column] | ( list[MySqlColumn] - | (list[PostgresColumn] | (list[SqlServerColumn] | list[KafkaColumn])) + | ( + list[PostgresColumn] + | (list[SqlServerColumn] | (list[KafkaColumn] | list[LoadGeneratorColumn])) + ) ), rng: random.Random, kind: ExprKind = ExprKind.ALL, @@ -371,6 +394,12 @@ def expression( if col.data_type == data_type: return str(col) + # Only in read/filter contexts: overflow/eval errors are tolerated there, + # but not in write (INSERT value) or materialized-view-body contexts. + edges = EDGE_VALUES.get(data_type) + if edges and kind == ExprKind.ALL and rng.random() < 0.2: + return rng.choice(edges) + record_size = rng.choice( [RecordSize.TINY, RecordSize.SMALL, RecordSize.MEDIUM, RecordSize.LARGE] ) diff --git a/misc/python/materialize/parallel_workload/negative_accumulation_errors.py b/misc/python/materialize/parallel_workload/negative_accumulation_errors.py index 715ce5ba2704e..6f60be9596676 100644 --- a/misc/python/materialize/parallel_workload/negative_accumulation_errors.py +++ b/misc/python/materialize/parallel_workload/negative_accumulation_errors.py @@ -31,8 +31,18 @@ "Negative multiplicities in TopK", # Reduce "Net-zero records with non-zero accumulation in ReduceAccumulable", + # Client-facing variant of the above (reduce.rs:1514, EvalError), distinct + # from the internal ReduceAccumulable log text. Seen in repeat_row (#8106). + "with non-zero accumulation in accumulable aggregate", "Non-positive multiplicity in DistinctBy", + # Covers the `ReduceInaccumulable`, `ReduceInaccumulable DISTINCT` and + # `ReduceMinsMaxes` sites, which surface their internal log text verbatim. "Non-positive accumulation", + # The hierarchical min/max stage words its client-facing error differently + # from its "Non-positive accumulation in MinsMaxesHierarchical" log line + # (reduce.rs `build_bucketed_stage`), so the entry above does not cover it. + # Seen in repeat_row (build 17664). + "saw non-positive accumulation", "Invalid negative unsigned aggregation in ReduceAccumulable", "saw negative accumulation", # Peek handling @@ -43,6 +53,9 @@ "S3 oneshot sink encountered negative multiplicities", # Constant folding "Negative multiplicity in constant result", + # Constant folding a DISTINCT/INTERSECT/reduce over a repeat_row collection + # with negative diffs. Seen in repeat_row (builds 17205, 17214). + "constant folding encountered reduce on collection with non-positive multiplicities", # Scalar subquery guard "negative number of rows produced in subquery", ] diff --git a/misc/python/materialize/parallel_workload/parallel_workload.py b/misc/python/materialize/parallel_workload/parallel_workload.py index fb24101ca83af..4a2740fc70842 100644 --- a/misc/python/materialize/parallel_workload/parallel_workload.py +++ b/misc/python/materialize/parallel_workload/parallel_workload.py @@ -9,6 +9,7 @@ import argparse import datetime +import faulthandler import gc import os import random @@ -24,12 +25,17 @@ from materialize.parallel_workload.action import ( Action, ActionList, + ApplyReplacementMaterializedViewAction, BackupRestoreAction, CancelAction, + CreateReplacementMaterializedViewAction, DropClusterAction, DropDatabaseAction, + DropRoleAction, DropSchemaAction, + ExplainAnalyzeAction, KillAction, + SealedCollectionCheckAction, StatisticsAction, ZeroDowntimeDeployAction, action_lists, @@ -45,6 +51,7 @@ MAX_KAFKA_SINKS, MAX_KAFKA_SOURCES, MAX_MYSQL_SOURCES, + MAX_NETWORK_POLICIES, MAX_POSTGRES_SOURCES, MAX_ROLES, MAX_SCHEMAS, @@ -67,6 +74,23 @@ REPORT_TIME = 10 +def run_final_sealed_check( + rng: random.Random, database: Database, cur: psycopg.Cursor +) -> None: + """Guaranteed final oracle pass, before the objects are dropped. + + The finalize task seals a wrongly finalized shard only ~5s after the + triggering DROP, so damage from the run's last actions is not visible + immediately, and a run whose random scheduling never picked + SealedCollectionCheckAction would otherwise end without any pass. + """ + exe = Executor(rng, cur, None, database) + sealed_check = SealedCollectionCheckAction(rng, None) + if sealed_check.applicable(exe): + time.sleep(10) + sealed_check.run(exe) + + def run( host: str, ports: dict[str, int], @@ -123,6 +147,9 @@ def run( system_exe.execute( f"ALTER SYSTEM SET max_roles = {MAX_ROLES * 1000 + num_threads}" ) + system_exe.execute( + f"ALTER SYSTEM SET max_network_policies = {MAX_NETWORK_POLICIES + num_threads}" + ) system_exe.execute( f"ALTER SYSTEM SET max_clusters = {MAX_CLUSTERS * 40 + num_threads}" ) @@ -163,7 +190,7 @@ def run( system_exe.execute("DROP CLUSTER quickstart CASCADE") replica_names = [f"r{replica_id}" for replica_id in range(0, replicas)] replica_string = ",".join( - f"{replica_name} (SIZE 'scale=1,workers=4')" + f"{replica_name} (SIZE 'scale=1,workers=1')" for replica_name in replica_names ) system_exe.execute( @@ -228,10 +255,17 @@ def run( ) workers.append(worker) + # Daemon threads, so that a worker still stuck in a long-running query + # cannot block interpreter shutdown. The clean paths below join + # explicitly, and a WorkerFailedException propagates out of `run` past + # those joins. Without this a single wedged worker keeps the process + # alive until the CI step timeout kills it, which reports the timeout + # instead of the failure that actually happened. thread = threading.Thread( name=thread_name, target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -252,6 +286,7 @@ def run( name="cancel", target=worker.run, args=(host, ports["mz_system"], ports["http"], "mz_system", database), + daemon=True, ) thread.start() threads.append(thread) @@ -272,6 +307,7 @@ def run( name="kill", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -299,6 +335,7 @@ def run( name="zero-downtime-deploy", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -319,6 +356,7 @@ def run( name="kill", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -343,6 +381,7 @@ def run( name="statistics", target=worker.run, args=(host, ports["mz_system"], ports["http"], "mz_system", database), + daemon=True, ) thread.start() threads.append(thread) @@ -379,13 +418,79 @@ def run( if all([not thread.is_alive() for thread in threads]): break else: - for worker, thread in zip(workers, threads): - if thread.is_alive(): - print( - f"{thread.name} still running ({worker.exe.mz_service}): {worker.exe.last_log} ({worker.exe.last_status})" - ) + alive = [(w, t) for w, t in zip(workers, threads) if t.is_alive()] + for worker, thread in alive: + if worker.exe is None: + # Still in the initial connect retry loop. + print(f"{thread.name} still connecting") + continue + print( + f"{thread.name} still running ({worker.exe.mz_service}): {worker.exe.last_log} ({worker.exe.last_status})" + ) + + # Workers are daemon threads, so a wedged one cannot keep the process + # alive and has to be caught here or it silently costs the run its + # remaining workload. A worker waiting on the server is a server-side + # hang (DB-118) and is tolerated, as is one still in its connect retry + # loop (exe is None), which can only be hanging inside a connect call + # to the server. One that never waits on the server is stuck in the + # workload's own code, e.g. deadlocked on two object locks taken in + # opposite orders, which no timeout ever clears. Nothing legitimate + # pauses without a round trip for the 300s above (the longest such + # pause is BackupRestoreAction's 240s sleep), but a single sample can + # catch a worker between round trips, so only call a worker wedged if + # it stays off the server for a whole confirmation window. + # + # A tolerated server-side hang also explains every other worker: the + # hung statement holds its actions' object locks (an `ALTER SINK ... + # SET FROM` waits unboundedly for the sink's frontier to catch up, + # database-issues#9820, while holding both the sink's and the new base + # object's lock), and a worker queued on such a lock can no more reach + # the server than a deadlocked one can. Telling the two apart from + # outside is not possible, so a workload-side deadlock is only + # diagnosed when no worker is on the server at all. Waiters and + # deadlocks both resolve when the statement does, and calling a waiter + # deadlocked costs the run its final checks. + def not_on_server() -> set[str]: + return { + t.name + for w, t in alive + if t.is_alive() and w.exe is not None and w.exe.last_status != "running" + } + + def on_server() -> set[str]: + return { + t.name + for w, t in alive + if t.is_alive() and w.exe is not None and w.exe.last_status == "running" + } + + wedged = not_on_server() + confirm_until = time.time() + 60 + while wedged and time.time() < confirm_until: + time.sleep(1) + wedged &= not_on_server() + if on_server(): + wedged = set() + break merge_num_queries(num_queries, workers) - print_stats(num_queries, workers, num_threads, scenario) + print_stats(num_queries, workers, num_threads, complexity, scenario) + + # A wedged worker must not mask damage, so run the final oracle pass + # on a fresh connection before exiting. Skipped for 0dt deploys, + # where connecting to the fenced-out environmentd can hang forever. + if scenario != Scenario.ZeroDowntimeDeploy: + try: + check_conn = psycopg.connect( + host=host, port=ports["materialized"], user="materialize" + ) + except Exception as e: + print(f"Skipping final sealed-collection check: {e}") + else: + check_conn.autocommit = True + with check_conn.cursor() as cur: + run_final_sealed_check(rng, database, cur) + check_conn.close() if num_threads >= 50: # Under high load some queries can't finish quickly, especially UPDATE/DELETE @@ -395,9 +500,13 @@ def run( # environmentd will be stuck forever, the promoted environmentd can # take > 10 minutes to become responsive as well os._exit(0) - # TODO: Reenable when https://linear.app/materializeinc/issue/DB-118 is fixed - # print("Threads have not stopped within 5 minutes, exiting hard") - # os._exit(1) + if wedged: + faulthandler.dump_traceback() + print( + "^^^ +++ Threads have not stopped within 5 minutes and are not" + f" waiting on the server, exiting hard: {', '.join(sorted(wedged))}" + ) + os._exit(1) os._exit(0) try: @@ -422,6 +531,8 @@ def run( # increment can still be in flight. database.read_then_write_counter.validate(exe) + run_final_sealed_check(rng, database, cur) + # Dropping the database also releases the long running connections # used by database objects. database.drop(exe) @@ -446,7 +557,7 @@ def run( conn.close() merge_num_queries(num_queries, workers) - print_stats(num_queries, workers, num_threads, scenario) + print_stats(num_queries, workers, num_threads, complexity, scenario) def merge_num_queries( @@ -474,6 +585,7 @@ def print_stats( num_queries: defaultdict[ActionList, Counter[type[Action]]], workers: list[Worker], num_threads: int, + complexity: Complexity, scenario: Scenario, ) -> None: ignored_errors: defaultdict[str, Counter[type[Action]]] = defaultdict(Counter) @@ -527,7 +639,11 @@ def print_stats( # (ReplaceMaterializedView, whose replacements pile up unfinalized in the # workload). "unknown cluster 'dont_exist'" comes from FlipFlags setting the # default cluster to a nonexistent one to hunt for panics, so any action - # needing a cluster fails with it while that flag is live. Tracked + # needing a cluster fails with it while that flag is live. "cannot be + # dropped because some objects depend on it" is the same rejection for + # DROP ROLE: AlterOwnerAction reassigns object ownership to random roles, + # so a role usually owns something and a short run can see DropRoleAction + # never land a dependency-free role. Tracked # separately so such actions don't trip the broken-action assertion below. noise = { "must be owner of", @@ -537,14 +653,26 @@ def print_stats( "because it already has a replacement", "is sealed and thus cannot be replaced", "unknown cluster 'dont_exist'", + "cannot be dropped because some objects depend on it", } + if complexity in (Complexity.DDL, Complexity.DDLOnly): + # CreateTableAction and CreateViewAction make half their objects TEMP, + # and those land in the shared table and view lists every worker samples + # from. A temporary schema belongs to the session that created it, so + # any other worker resolving "db"."mz_temp"."name" fails with "unknown + # schema 'mz_temp'". That is roughly half the attempts of an action + # picking a random table or view, which an action as low-weight as + # RevokePrivilegesAction can spend a whole run on. Concurrent schema + # drops and renames land here too. Only DDL complexity creates temp + # objects or touches schemas, so elsewhere this stays a genuine signal. + noise.add("unknown schema") if scenario == Scenario.Rename: # Concurrent renames invalidate the qualified names an action captured # earlier (a stored SELECT it re-renders, a target it resolved before # the rename), so a rare action can fail on every attempt with a # name-resolution race. Expected only in this scenario. In Regression, # where nothing renames, these would be a genuine signal. - noise |= {"does not exist", "unknown schema", "unknown catalog item"} + noise |= {"does not exist", "unknown catalog item"} num_errored_real: Counter[type[Action]] = Counter() for worker in workers: num_successes.update(worker.num_successes) @@ -558,11 +686,42 @@ def print_stats( for action_list in action_lists for action_class in action_list.action_classes } - # These use RESTRICT and their targets practically always contain - # objects (schemas and databases their items, clusters their sources, - # sinks, and indexes), so they exercise the rejection path and are not - # expected to ever succeed. - action_classes -= {DropClusterAction, DropDatabaseAction, DropSchemaAction} + # A given churny seed can legitimately see zero successes for these, so the + # broken-action assertion must not fire on them: + # * DropCluster/DropDatabase/DropSchema: RESTRICT rejections. Their + # targets practically always contain objects, and the rejection message + # is class-specific, so `noise` above cannot cover it without excusing + # it for every other action too. + # * DropRoleAction: needs a dependency-free role, but AlterOwnerAction + # keeps reassigning object ownership to random roles. + # * ExplainAnalyzeAction: needs a hydrated MV/index on the active cluster, + # which renames/drops keep retiring. + # These succeed in normal runs, so they aren't broken. + # + # NOTE: listed by name rather than matched on a Drop* name prefix. The other + # Drop* actions pick an existing object and are expected to succeed + # constantly, and a broken one (wrong quoting or qualification) fails with + # `unknown catalog item`, which the base errors_to_ignore tolerates and + # `noise` does not list. That is exactly what the assertion is meant to + # catch, so a prefix would silently exempt the failure mode it exists for. + excluded: set[type[Action]] = { + DropClusterAction, + DropDatabaseAction, + DropRoleAction, + DropSchemaAction, + ExplainAnalyzeAction, + } + if scenario == Scenario.Rename: + # CreateReplacementMaterializedViewAction re-renders the view's SELECT + # with the object names captured at creation. Renames invalidate them, + # so CREATE REPLACEMENT fails with a tolerated "does not exist"/"unknown + # schema", and a churny rename seed can legitimately never land a clean + # attempt. With no replacement ever created, + # ApplyReplacementMaterializedViewAction then has nothing to apply + # either. Both succeed in the other scenarios, so they stay checked + # there. + excluded.add(CreateReplacementMaterializedViewAction) + excluded.add(ApplyReplacementMaterializedViewAction) never_succeeded = [] for action_class in sorted(action_classes, key=lambda cls: cls.__name__): successes = num_successes[action_class] @@ -596,7 +755,7 @@ def print_stats( always_erroring = [ action_class.__name__ for action_class, skips, errored in never_succeeded - if num_errored_real[action_class] > 0 + if num_errored_real[action_class] > 0 and action_class not in excluded ] assert ( not always_erroring @@ -636,7 +795,12 @@ def parse_common_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--azurite", action="store_true", help="Use Azurite as blob store instead of S3" ) - parser.add_argument("--replicas", type=int, default=2, help="use multiple replicas") + parser.add_argument( + "--replicas", + type=int, + default=1, + help="default replica number for quickstart cluster", + ) def main() -> int: diff --git a/misc/python/materialize/parallel_workload/settings.py b/misc/python/materialize/parallel_workload/settings.py index 79196897d521a..4efdd32a3876e 100644 --- a/misc/python/materialize/parallel_workload/settings.py +++ b/misc/python/materialize/parallel_workload/settings.py @@ -58,4 +58,28 @@ def _missing_(cls, value): # it on outside that scenario is harmless: no Parallel Workload codegen # emits `repeat_row` unless the scenario is active. "enable_repeat_row": "true", + # TODO: Reenable once the frontend-peek path stops acquiring a read hold + # whose `since` has compacted past the chosen `as_of`. The soft-assert + # assert_read_holds_correct (frontend_peek.rs:1759) fires for peeks in + # multi-statement transactions ("... read hold at .. is not enough for + # as_of .."), panicking the coordinator. Peeks still work via the classic + # coordinator path with this off. + # See https://linear.app/materializeinc/issue/SQL-520 + "enable_frontend_peek_sequencing": "false", + # 64 MiB, down from the 1 GiB default. A peek's result is materialized in + # memory on the replica serving it before this bound errors it + # (`to_error_if_exceeds`, compute_state.rs), and measured on a workload-like + # relation a 300 MiB result costs ~1 GiB of replica RSS plus ~0.5 GiB in + # environmentd. The ceiling is per in-flight statement and every worker + # holds a pg and a WebSocket session, so at the default a handful of + # concurrent wide reads can claim tens of GiB on the default cluster, where + # every peek lands. All containers of a run share one cgroup budget (24 GiB + # on the CI agent, environmentd plus every replica process plus Kafka, + # Postgres, MySQL, SQL Server), so that ends in the kernel OOM-killer + # picking a victim and taking unrelated processes down with it. Runs do + # reach the cap ("result exceeds max size of" shows up in the ignored-error + # statistics), which is what makes the tail dangerous. Query shapes are + # unaffected: an oversized result errors instead of being buffered, and that + # error is already tolerated for every action. + "max_result_size": "67108864", }