diff --git a/doc/developer/design/20260210_incremental_occ_read_then_write.md b/doc/developer/design/20260210_incremental_occ_read_then_write.md index db6286fb2f210..52a75c53788b6 100644 --- a/doc/developer/design/20260210_incremental_occ_read_then_write.md +++ b/doc/developer/design/20260210_incremental_occ_read_then_write.md @@ -41,10 +41,13 @@ a subscribe that continually tracks the current state of the data. - Removing the in-process locks immediately. During rollout, the old lock-based path and the new OCC path coexist behind a feature flag. The locks can be removed once the OCC path is fully rolled out. -- Multi-statement transactions. The OCC approach as described here applies to - single-statement implicit transactions. Explicit multi-statement write - transactions continue to use the existing path. And there are not plans to - support mixed read/write transactions. +- Mixed read/write transactions. A write that reads persisted state commits at + the frontier it observed, which it cannot postpone until COMMIT, so it runs + only as a single statement. A write that reads nothing does compose with + transactions: its diffs are frontier-independent, so they are buffered as + session write ops and land when the transaction commits. That covers, for + example, `INSERT INTO t SELECT generate_series(1, 20000)`, whose values are + constant but too large to fold into a literal. ## Overview @@ -128,7 +131,7 @@ Session Task Coordinator | | |-- acquire OCC semaphore | | | - |-- CreateReadThenWriteSubscribe ----> | + |-- CreateInternalSubscribe ---------> | | <------------ subscribe channel -----| | | | +-- OCC Loop ------------------+ | @@ -141,7 +144,7 @@ Session Task Coordinator | | if Success: break | | | +------------------------------+ | | | - |-- DropReadThenWriteSubscribe ------> | + |-- DropInternalSubscribe -----------> | | | ``` @@ -193,7 +196,7 @@ subscribe. The subscribes created for read-then-write are internal: they do not appear in `mz_subscriptions` or other introspection tables, and they don't increment the active subscribes metric. They are created and dropped via dedicated `Command` -variants (`CreateReadThenWriteSubscribe`, `DropReadThenWriteSubscribe`). +variants (`CreateInternalSubscribe`, `DropInternalSubscribe`). ## Correctness @@ -295,10 +298,15 @@ Benchmarking a PoC-level implementation of the OCC approach against `main` for The benchmark varies concurrency (number of workers) on the x-axis and shows throughput (left) and latency (right). Key observations: -- At low concurrency (1-7 workers), the OCC approach is comparable or _better_ - than `main`. This is because the OCC path begins preparing the write (opening - the subscribe, receiving the snapshot) before the write timestamp is claimed, - whereas the old path only starts the peek after acquiring the lock. +- At low concurrency (1-7 workers), the result depends on write size. A single + large `UPDATE`/`DELETE` is comparable or _better_ than `main`, because the + subscribe streams the mutation diffs directly whereas the old path peeks every + matched row and then recomputes the diffs. Small writes, however, _regress_: + every operation installs a subscribe dataflow, waits for its snapshot, and + tears it down, where the old path uses a cheap fast-path peek. This + per-operation subscribe overhead makes tiny `UPDATE`s roughly 1.5-2x slower at + low/no concurrency (observed in the nightly feature benchmark + `ManySmallUpdates` and the scalability `UpdateWorkload`). - At higher concurrency, performance degrades as expected due to the O(N^2) retry behavior: with more concurrent writers, more retries are needed. The concurrency semaphore (default 4 permits) bounds this in practice. @@ -306,6 +314,11 @@ throughput (left) and latency (right). Key observations: table). Real workloads with writes to different tables won't experience the contention. +The chart above is from the PoC, which benchmarked `UPDATE t SET x = x + 1` over +a larger table (the regime where OCC wins). It does not capture the small-write +regression noted above, which is an accepted cost: high write throughput is a +non-goal (see Non-Goals). + ## Rollout The new path is controlled by a `enable_adapter_frontend_occ_read_then_write` diff --git a/doc/developer/design/20260210_incremental_occ_read_then_write_qa_findings.md b/doc/developer/design/20260210_incremental_occ_read_then_write_qa_findings.md new file mode 100644 index 0000000000000..c97c72f1a56ba --- /dev/null +++ b/doc/developer/design/20260210_incremental_occ_read_then_write_qa_findings.md @@ -0,0 +1,119 @@ +# QA findings: incremental OCC read-then-write + +Adversarial QA pass over the `adapter-read-then-write-occ-incremental` branch +(frontend OCC sequencing for DELETE / UPDATE / INSERT...SELECT, gated by +`enable_adapter_frontend_occ_read_then_write`). + +Tests live in `src/environmentd/tests/server.rs`, prefixed `qa_occ_` / +`qa_legacy_`. Run them with the OCC binary already built: + +``` +METADATA_BACKEND_URL=postgres://root@localhost:26257/materialize \ + cargo nextest run -p mz-environmentd -E 'test(/qa_occ_/)' +``` + +All `qa_occ_` tests — including the two former bug repros +(`qa_occ_far_future_refresh_mv_respects_statement_timeout` and +`qa_occ_far_future_rtw_starves_permit_pool`) — are now active and passing; they +verify the fix described below. The only `#[ignore]`d test is +`qa_legacy_far_future_refresh_mv_statement_timeout`, which documents the +pre-existing far-future hang on the legacy coordinator path (out of scope for +this OCC-only fix). + +## What was verified correct (passing tests) + +These exercise behaviors the design must preserve; all pass under OCC: + +- `qa_occ_concurrent_delete_no_overdelete` — N concurrent `DELETE`s of a + multiset row of multiplicity M sum to exactly M deleted, table ends empty, + multiplicity never goes negative. OCC retry prevents over-delete. +- `qa_occ_duplicate_row_multiplicity_counts` — DELETE/UPDATE/INSERT...SELECT + affected-row counts respect multiset multiplicity. +- `qa_occ_not_null_constraint_enforced` — `NOT NULL` violations via UPDATE and + INSERT...SELECT error out and leave the table unchanged (no partial commit). +- `qa_occ_returning_values_correct` — INSERT RETURNING (constant and + read-dependent INSERT...SELECT) returns the right rows/expressions. (Note: + Materialize only parses `RETURNING` for INSERT, so the DELETE/UPDATE + RETURNING branches in `build_success_response` are effectively dead code.) +- `qa_occ_update_moves_overlapping_rows` — row-moving UPDATEs + (`SET id = id + 1`) produce the correct final set through the + Let/Negate/map MIR transform and consolidation. +- `qa_occ_insert_select_from_mv` — INSERT...SELECT reading a materialized view + (TimestampDependent timeline; linearization defaults to EpochMilliseconds). +- `qa_occ_concurrent_mixed_dml_no_internal_error` — concurrent + UPDATE/DELETE/INSERT mix produces no internal errors / coordinator panics. + +The core OCC retry/consolidation logic, the timestamped-write/oracle +interaction, and the snapshot/progress `NoRowsMatched` reasoning were also +reviewed statically and found sound. + +## Fixed: OCC read-then-write now honors `statement_timeout` everywhere + +**Original bug.** `frontend_read_then_write::ensure_read_linearized` loops +`tokio::time::sleep` until the oracle reaches `as_of`, with no +`statement_timeout` check. When a read-then-write's `as_of` was far in the +future — e.g. the read depends on a `REFRESH AT ` materialized +view — the session hung indefinitely; only client cancellation / disconnect +freed it. `statement_timeout` was only enforced inside `run_occ_loop` (via +`tokio::time::timeout` on `recv`), which is reached *after* linearization, so +the pre-loop phases (planning, OCC permit acquisition, timestamp determination, +linearization) were all unbounded. This contradicted the call-site comment, +which claimed a pathological RTW "hits `statement_timeout` and returns". + +The OCC-specific amplification was the real concern: the parked op holds an OCC +semaphore permit (`max_concurrent_occ_writes`, default 4) for its entire +lifetime. A handful of such ops exhausted the global permit pool and wedged +**all** read-then-writes process-wide — including ones on unrelated tables — +because those victims block on permit acquisition *before* the OCC loop, so +they too ignored `statement_timeout`. The legacy path only blocks writes to the +*target* table (per-table write lock), so OCC widened the blast radius from one +table to the whole RTW pipeline. + +**The fix.** `statement_timeout` is now enforced **centrally**, in the +`tokio::select!` of `SessionClient::try_frontend_read_then_write_with_cancel` +(`src/adapter/src/client.rs`) — the one place that already owns the whole +operation's lifetime and handles cancellation. A new select arm sleeps for +`*session.vars().statement_timeout()` (treating `Duration::ZERO` as "off / +wait forever" via `futures::future::pending`, mirroring `run_occ_loop`'s +`effective_timeout`). When it fires it returns `AdapterError::StatementTimeout` +and forwards `Command::PrivilegedCancelRequest` to the coordinator (mirroring +the existing `cancel_future` arm) to clean up any in-flight coordinator-owned +work. Because the whole `try_frontend_read_then_write` future is dropped, the +OCC permit, read holds, and `SubscribeHandle` (whose `Drop` sends +`DropInternalSubscribe`) are all released. + +This single deadline now bounds *every* phase: planning, OCC permit +acquisition, timestamp determination, `ensure_read_linearized`, **and** the OCC +loop. The far-future op times out promptly and releases its permit, so victims +behind a starved pool also honor their own `statement_timeout`. The in-loop +`run_occ_loop` timeout is kept as defense-in-depth (it also bounds a single +blocking `recv`); the call-site comment in `frontend_read_then_write.rs` was +corrected to point at the central enforcement. + +Verified by `qa_occ_far_future_refresh_mv_respects_statement_timeout` (the +far-future op returns a timeout error well within budget) and +`qa_occ_far_future_rtw_starves_permit_pool` (the unrelated victim times out on +permit acquisition within its own `statement_timeout` instead of hanging). +Both are now active (no longer `#[ignore]`d) and passing. + +### Out of scope / unchanged + +- **Legacy path.** The legacy (lock-based) coordinator path also hangs on a + far-future read past `statement_timeout` + (`qa_legacy_far_future_refresh_mv_statement_timeout`, still `#[ignore]`d). + This pre-existing hang is intentionally left unchanged; the fix is OCC-only. + +### Corollary: `max_concurrent_occ_writes` is constrained to `>= 1` + +The same "block before the timeout is enforced" structure meant that +`max_concurrent_occ_writes = 0` sizes the semaphore to zero permits and bricks +every read-then-write after restart. With the central timeout such operations +at least *time out* instead of hanging forever, but zero permits is still a +footgun, so the parameter now carries a domain constraint requiring at least +1. That covers both ways of setting it, `ALTER SYSTEM SET` and +`system_parameter_default`. + +`ALTER SYSTEM SET`/`RESET` of the parameter is allowed, since the value is +sampled once at boot and the running process cannot observe a later change. +The statement emits a warning that the change only takes effect when +`environmentd` restarts. diff --git a/doc/user/data/metrics.yml b/doc/user/data/metrics.yml index ee66eaf5b7b16..ee00d6e273203 100644 --- a/doc/user/data/metrics.yml +++ b/doc/user/data/metrics.yml @@ -1417,6 +1417,20 @@ metrics: visibility: public tags: - environment +- name: mz_occ_read_then_write_retry_count_bucket + help: Number of OCC retries per read-then-write operation. + labels: + - le + source: src/adapter/src/metrics.rs + visibility: internal +- name: mz_occ_read_then_write_retry_count_count + help: Number of OCC retries per read-then-write operation. + source: src/adapter/src/metrics.rs + visibility: internal +- name: mz_occ_read_then_write_retry_count_sum + help: Number of OCC retries per read-then-write operation. + source: src/adapter/src/metrics.rs + visibility: internal - name: mz_optimization_notices help: Number of optimization notices per notice type. labels: diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index f622f03de3c8d..c87303fada0df 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -240,6 +240,11 @@ def get_variable_system_parameters( "true", ["true", "false"], ), + VariableSystemParameter( + "enable_adapter_frontend_occ_read_then_write", + "true", + ["true", "false"], + ), VariableSystemParameter( "enable_cast_elimination", "true", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index cf2ae87276c00..aa1a025a3c319 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -1931,6 +1931,10 @@ def __init__( # behavior, you should add it. Feature flags which turn on/off # externally visible features should not be flipped. self.uninteresting_flags: list[str] = [ + # Read once at environmentd startup, so an ALTER SYSTEM SET only + # takes effect after a restart. Flipping it here would be a no-op + # for the running process. + "enable_adapter_frontend_occ_read_then_write", "enable_compute_half_join2", "enable_mz_join_core", "enable_compute_correction_v2", diff --git a/src/adapter-types/src/dyncfgs.rs b/src/adapter-types/src/dyncfgs.rs index 2a370901583f8..086834c155b35 100644 --- a/src/adapter-types/src/dyncfgs.rs +++ b/src/adapter-types/src/dyncfgs.rs @@ -422,6 +422,13 @@ pub const DEFAULT_HYDRATION_BURST_LINGER: Config = Config::new( "The burst-replica linger duration written when an AUTO SCALING STRATEGY omits LINGER DURATION.", ); +pub const FRONTEND_READ_THEN_WRITE: Config = Config::new( + "enable_adapter_frontend_occ_read_then_write", + false, + "Use frontend sequencing (with optimistic concurrency control) for \ + DELETE, UPDATE, and INSERT operations.", +); + /// Adds the full set of all adapter `Config`s. pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { configs @@ -474,4 +481,5 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&CATALOG_INFO_METRICS_RECONCILE_INTERVAL) .add(&PG_TIMESTAMP_ORACLE_STATEMENT_TIMEOUT) .add(&ENABLE_SCOPED_SYSTEM_PARAMETERS) + .add(&FRONTEND_READ_THEN_WRITE) } diff --git a/src/adapter/src/catalog/open.rs b/src/adapter/src/catalog/open.rs index 693515ccca026..d9a8d1cef604c 100644 --- a/src/adapter/src/catalog/open.rs +++ b/src/adapter/src/catalog/open.rs @@ -313,6 +313,17 @@ impl Catalog { .await; builtin_table_updates.extend(builtin_table_update); + // Both the `system_parameter_defaults` set above and the durable + // `SystemConfiguration` values just applied via `pre_item_updates` now + // live in the `SystemVars` value map. Mirror their effective values into + // the dyncfg `ConfigSet` so startup-only reads — the + // `ENABLE_EXPRESSION_CACHE` read just below and `FRONTEND_READ_THEN_WRITE` + // in coord bootstrap — observe configured values, not compile-time + // defaults. `apply_updates` only syncs the `ConfigSet` when a + // `SystemConfiguration` update is present; deployments configured purely + // via `system_parameter_default` have none, so sync explicitly here. + state.system_config().dyncfg_updates(); + // 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/client.rs b/src/adapter/src/client.rs index a0719ede9ac36..75ef1fe49d252 100644 --- a/src/adapter/src/client.rs +++ b/src/adapter/src/client.rs @@ -18,13 +18,14 @@ 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; use mz_auth::{Authenticated, AuthenticatorKind}; use mz_build_info::BuildInfo; use mz_compute_types::ComputeInstanceId; +use mz_expr::UnmaterializableFunc; use mz_ore::channel::OneshotReceiverExt; use mz_ore::collections::CollectionExt; use mz_ore::id_gen::{IdAllocator, IdAllocatorInnerBitSet, MAX_ORG_ID, org_id_conn_bits}; @@ -60,9 +61,14 @@ use crate::command::{ use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend}; use crate::coord::{Coordinator, ExecuteContextGuard}; use crate::error::AdapterError; +use crate::frontend_read_then_write::{FrontendWriteAttemptState, FrontendWriteCancellation}; use crate::metrics::{self, Metrics}; +use crate::optimize::dataflows::{EvalTime, ExprPrepOneShot}; +use crate::optimize::{self, Optimize, OptimizerError}; +use crate::peek_client::StatementLoggingGuard; use crate::session::{ EndTransactionAction, PreparedStatement, Session, SessionConfig, StateRevision, TransactionId, + TransactionStatus, }; use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy}; use crate::telemetry::{self, EventDetails, SegmentClientExt, StatementFailureType}; @@ -286,6 +292,9 @@ impl Client { persist_client, statement_logging_frontend, superuser_attribute, + occ_write_semaphore, + frontend_read_then_write_enabled, + read_only, } = response; let peek_client = PeekClient::new( @@ -296,6 +305,9 @@ impl Client { optimizer_metrics, persist_client, statement_logging_frontend, + occ_write_semaphore, + frontend_read_then_write_enabled, + read_only, ); let mut client = SessionClient { @@ -602,10 +614,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"); } } @@ -632,6 +649,25 @@ pub struct SessionClient { pub enable_frontend_peek_sequencing: bool, } +/// Keeps a connection cancel watch installed in the coordinator for the +/// duration of a frontend read-then-write attempt. +struct FrontendConnectionCancelWatchGuard { + conn_id: ConnectionId, + operation_id: Uuid, + client: Option, +} + +impl Drop for FrontendConnectionCancelWatchGuard { + fn drop(&mut self) { + if let Some(client) = self.client.take() { + client.try_send(Command::UnregisterConnectionCancelWatch { + conn_id: self.conn_id.clone(), + operation_id: self.operation_id, + }); + } + } +} + impl SessionClient { /// Parses a SQL expression, reporting failures as a telemetry event if /// possible. @@ -781,11 +817,13 @@ impl SessionClient { outer_ctx_extra: Option, ) -> Result<(ExecuteResponse, Instant), AdapterError> { let execute_started = Instant::now(); + let cancel_future = cancel_future.map(|_| ()).shared(); let mut outer_ctx_extra = outer_ctx_extra; // 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 @@ -812,11 +850,45 @@ impl SessionClient { // No additional work needed here. return Ok((resp, execute_started)); } else { - debug!("frontend peek did not happen, falling back to `Command::Execute`"); + debug!("frontend peek did not happen, trying frontend read-then-write"); + } + + // Attempt read-then-write sequencing in the session task. + let rtw_result = self + .try_frontend_read_then_write_with_cancel( + &portal_name, + &mut outer_ctx_extra, + cancel_future.clone(), + ) + .await; + let rtw_result = match rtw_result { + Ok(result) => result, + Err(err) => { + // If we still hold EXECUTE-level logging (it is only consumed + // once the frontend commits to handling the write), retire it + // as errored. Letting the guard drop would misrecord the + // statement as aborted. + if let Some(id) = outer_ctx_extra + .take() + .and_then(|guard| guard.defuse().retire()) + { + self.peek_client.log_ended_execution( + id, + StatementEndedExecutionReason::Errored { + error: err.to_string(), + }, + ); + } + return Err(err); + } + }; + if let Some(resp) = rtw_result { + debug!("frontend read-then-write succeeded"); + return Ok((resp, execute_started)); + } else { + debug!("frontend read-then-write 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.) } let response = self @@ -827,7 +899,7 @@ impl SessionClient { tx, outer_ctx_extra, }, - cancel_future, + cancel_future.clone(), ) .await?; Ok((response, execute_started)) @@ -836,7 +908,8 @@ impl SessionClient { /// 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 @@ -854,7 +927,8 @@ impl SessionClient { 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), }; @@ -905,66 +979,45 @@ impl SessionClient { // // 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 - }; + // not the inner SQL. On success the id moves into `outer_ctx_extra` + // for `try_frontend_peek` / `try_frontend_read_then_write` to retire; + // on planning error we explicitly retire this guard with `Errored`. + let mut held_guard: Option = if outer_ctx_extra.is_none() { + let session = self.session.as_mut().expect("SessionClient invariant"); + let mut none_guard: Option = None; + Some(self.peek_client.begin_statement_logging( + session, + ¶ms, + &outer_logging, + &catalog, + outer_lifecycle_timestamps, + &mut none_guard, + )) + } 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(), - }, - ); + if let Some(g) = held_guard.take() { + g.retire(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 + // Hand off the logging id to `outer_ctx_extra`. `try_frontend_peek` + // / `try_frontend_read_then_write` retire it themselves once they + // take ownership. + if let Some(mut g) = held_guard.take() { + let id = g.id(); + g.defuse(); + // Soft invariant: the next caller (`try_frontend_peek` / + // `try_frontend_read_then_write`) 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 @@ -972,7 +1025,7 @@ impl SessionClient { // — 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)); + *outer_ctx_extra = Some(ExecuteContextGuard::new(id, dummy_tx)); } Ok(new_portal_name) @@ -984,8 +1037,8 @@ impl SessionClient { /// 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. + /// in a single error-handling site that retires the EXECUTE-level + /// statement-logging guard with `Errored`. fn install_inner_portal_for_execute( &mut self, catalog: &Arc, @@ -1021,11 +1074,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 `try_frontend_peek` / + // `try_frontend_read_then_write` will see `outer_ctx_extra=Some(...)` + // and inherit the EXECUTE-level logging 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 @@ -1038,8 +1091,8 @@ 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. + // that ever changes, we'd silently skip OCC routing for the 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!( @@ -1308,7 +1361,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, @@ -1382,7 +1435,12 @@ impl SessionClient { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => {} + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::UnregisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => {} }; cmd }); @@ -1413,7 +1471,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,6 +1534,491 @@ impl SessionClient { Ok(None) } } + + /// Runs frontend read-then-write while reacting to both local/session + /// cancellation and coordinator-issued connection cancellation. + async fn try_frontend_read_then_write_with_cancel( + &mut self, + portal_name: &str, + outer_ctx_extra: &mut Option, + cancel_future: impl Future + Send, + ) -> Result, AdapterError> { + let conn_id = self.session().conn_id().clone(); + let statement_timeout = *self.session().vars().statement_timeout(); + let inner_client = self.inner().clone(); + let operation_id = Uuid::new_v4(); + let attempt_state = Arc::new(FrontendWriteAttemptState::new()); + let _connection_cancel_guard = FrontendConnectionCancelWatchGuard { + conn_id: conn_id.clone(), + operation_id, + client: Some(inner_client.clone()), + }; + + let mut cancel_future = pin::pin!(cancel_future); + let statement_timeout = async move { + if statement_timeout.is_zero() { + futures::future::pending::<()>().await; + } else { + tokio::time::sleep(statement_timeout).await; + } + }; + tokio::pin!(statement_timeout); + + // Registration is part of the statement lifetime. The operation ID + // prevents this guard from removing a newer operation's watch. + let mut connection_cancel_rx = { + let register = + self.peek_client + .call_coordinator(|tx| Command::RegisterConnectionCancelWatch { + conn_id: conn_id.clone(), + operation_id, + tx, + }); + tokio::pin!(register); + tokio::select! { + rx = &mut register => rx, + _ = &mut cancel_future => { + inner_client.try_send(Command::PrivilegedCancelRequest { + conn_id: conn_id.clone(), + }); + return Err(AdapterError::Canceled); + } + _ = &mut statement_timeout => { + inner_client.try_send(Command::PrivilegedCancelRequest { + conn_id: conn_id.clone(), + }); + return Err(AdapterError::StatementTimeout); + } + } + }; + if *connection_cancel_rx.borrow() { + return Err(AdapterError::Canceled); + } + let connection_cancel = async move { + if connection_cancel_rx.wait_for(|v| *v).await.is_err() { + futures::future::pending::<()>().await; + } + }; + tokio::pin!(connection_cancel); + + let frontend_read_then_write = self.try_frontend_read_then_write( + portal_name, + outer_ctx_extra, + Arc::clone(&attempt_state), + ); + tokio::pin!(frontend_read_then_write); + + let requested = tokio::select! { + response = &mut frontend_read_then_write => return response, + _ = &mut cancel_future => FrontendWriteCancellation::Canceled, + _ = &mut connection_cancel => FrontendWriteCancellation::Canceled, + _ = &mut statement_timeout => FrontendWriteCancellation::StatementTimeout, + }; + + attempt_state.request(requested); + inner_client.try_send(Command::PrivilegedCancelRequest { + conn_id: conn_id.clone(), + }); + + if !attempt_state.write_submitted() { + return Err(match requested { + FrontendWriteCancellation::Canceled => AdapterError::Canceled, + FrontendWriteCancellation::StatementTimeout => AdapterError::StatementTimeout, + }); + } + + // A submitted write can already be durable. Await its definitive result + // rather than reporting cancellation or timeout incorrectly. + frontend_read_then_write.await + } + + /// Attempt to sequence a read-then-write (DELETE/UPDATE/INSERT INTO .. + /// SELECT .. FROM) from the session task. + /// + /// Returns `Ok(Some(response))` if we handled the operation, or `Ok(None)` + /// to fall back to the Coordinator's sequencing. If it returns an error, it + /// should be returned to the user. + pub(crate) async fn try_frontend_read_then_write( + &mut self, + portal_name: &str, + outer_ctx_extra: &mut Option, + attempt_state: Arc, + ) -> Result, AdapterError> { + use mz_expr::{CollectionPlan, RowSetFinishing}; + use mz_sql::ast::ConstantVisitor; + use mz_sql::plan::{MutationKind, Plan, ReadThenWritePlan}; + use mz_sql_parser::ast::{InsertStatement, Statement}; + + // Check if frontend read-then-write is enabled (determined once at process startup + // to avoid a mixed-mode window where both the old lock-based path and the new OCC + // path are active concurrently). + if !self.peek_client.frontend_read_then_write_enabled { + return Ok(None); + } + + let catalog = self.catalog_snapshot("try_frontend_read_then_write").await; + + let stmt = { + let session = self.session.as_ref().expect("SessionClient invariant"); + let portal = match session.get_portal_unverified(portal_name) { + Some(portal) => portal, + None => return Ok(None), // Portal doesn't exist, fall back + }; + portal.stmt.clone() + }; + + let stmt = match stmt { + Some(stmt) + if matches!( + &*stmt, + Statement::Delete(_) | Statement::Update(_) | Statement::Insert(_) + ) => + { + stmt + } + Some(_stmt) => { + return Ok(None); + } + None => { + return Ok(None); + } + }; + + // Mirror the coordinator's transaction-state gate in `handle_execute`: + // in a multi-statement transaction (an implicit batch or an explicit + // block), the only DML allowed is an AST-constant INSERT without + // RETURNING, which joins the transaction's write ops and commits at + // transaction end. All other DML is prohibited because writes on this + // path commit immediately and cannot be rolled back at transaction + // end. `Failed` transactions pass through, pgwire only admits + // COMMIT/ROLLBACK in that state. + // + // An AST-constant source can still plan to a read, so the check on the + // planned selection further down narrows this. + { + let session = self.session.as_ref().expect("SessionClient invariant"); + match session.transaction() { + TransactionStatus::Default + | TransactionStatus::Started(_) + | TransactionStatus::Failed(_) => {} + TransactionStatus::InTransactionImplicit(_) + | TransactionStatus::InTransaction(_) => { + let constant_insert = matches!( + &*stmt, + Statement::Insert(InsertStatement { + source, returning, .. + }) if returning.is_empty() && ConstantVisitor::insert_source(source) + ); + if !constant_insert { + return Err(prohibited_in_transaction(&stmt)); + } + } + } + } + + if self + .session + .as_ref() + .expect("SessionClient invariant") + .vars() + .transaction_isolation() + .is_bounded_staleness() + { + return Err(AdapterError::BoundedStalenessReadOnly); + } + + // Verify and plan against one catalog snapshot. Pairing a stale plan + // with a newer target generation could direct a write incorrectly. + Coordinator::verify_portal( + &catalog, + self.session.as_mut().expect("SessionClient invariant"), + portal_name, + )?; + + let (params, logging, lifecycle_timestamps) = { + let portal = self + .session + .as_ref() + .expect("SessionClient invariant") + .get_portal_unverified(portal_name) + .expect("verified above"); + ( + portal.parameters.clone(), + Arc::clone(&portal.logging), + portal.lifecycle_timestamps.clone(), + ) + }; + + // Reject mutations in read-only mode (e.g. during 0dt upgrades). Done + // early, before any planning or fast-path dispatch, so every sub-path + // (constant INSERT, OCC INSERT/UPDATE/DELETE) is covered uniformly. + if self.peek_client.read_only { + return Err(AdapterError::ReadOnly); + } + + let (plan, target_cluster, resolved_ids, sql_impl_ids) = { + let session = self.session.as_mut().expect("SessionClient invariant"); + let conn_catalog = catalog.for_session(session); + let (stmt, resolved_ids) = mz_sql::names::resolve(&conn_catalog, (*stmt).clone())?; + let pcx = session.pcx(); + let (plan, sql_impl_ids) = + mz_sql::plan::plan(Some(pcx), &conn_catalog, stmt, ¶ms, &resolved_ids)?; + + let target_cluster = match session.transaction().cluster() { + Some(cluster_id) => crate::coord::TargetCluster::Transaction(cluster_id), + None => crate::coord::catalog_serving::auto_run_on_catalog_server( + &conn_catalog, + session, + &plan, + ), + }; + + (plan, target_cluster, resolved_ids, sql_impl_ids) + }; + + // Cluster restrictions and RBAC, mirroring the coordinator's checks + // in sequencer.rs. Resolution may fail if the target cluster doesn't + // exist — that gets reported later (with the correct error) by + // `validate_read_then_write`; for the purposes of these checks we + // treat it as "no cluster known", consistent with the coordinator. + let (target_cluster_id, target_cluster_name) = { + let session = self.session.as_ref().expect("SessionClient invariant"); + match catalog.resolve_target_cluster(target_cluster.clone(), session) { + Ok(cluster) => (Some(cluster.id), Some(cluster.name.clone())), + Err(_) => (None, None), + } + }; + { + let session = self.session.as_ref().expect("SessionClient invariant"); + let conn_catalog = catalog.for_session(session); + if let Some(cluster_name) = &target_cluster_name { + crate::coord::catalog_serving::check_cluster_restrictions( + cluster_name, + &conn_catalog, + &plan, + )?; + } + if let Err(e) = mz_sql::rbac::check_plan( + &conn_catalog, + None, + session, + &plan, + target_cluster_id, + &resolved_ids, + &sql_impl_ids, + ) { + return Err(e.into()); + } + } + + // Wait for any in-flight startup builtin-table appends that this plan + // depends on. Mirrors the frontend_peek and coordinator sequencer + // paths; no-op for plans that don't depend on builtin tables. + { + let session = self.session.as_mut().expect("SessionClient invariant"); + if let Some((_, wait_future)) = + crate::coord::appends::waiting_on_startup_appends(&catalog, session, &plan) + { + wait_future.await; + } + } + + // Handle ReadThenWrite plans or Insert plans. + let rtw_plan = match plan { + Plan::ReadThenWrite(rtw_plan) => rtw_plan, + Plan::Insert(insert_plan) => { + // For INSERT, we need to check if it's a constant insert + // without RETURNING. Constant inserts use a fast path in the + // coordinator, so we fall back. + // + // We need to lower HIR to MIR to check for constants because + // VALUES statements are planned as Wrap calls at the HIR level. + // + // Only take the constant-INSERT fast path when the HIR + // names no persisted collections (no Get nodes on tables + // or MVs). The MIR optimizer can fold an MV reference + // into a literal when the MV's plan happens to be + // constant, but "plan is constant" is NOT the same as + // "content is visible at the current oracle_ts". A + // `REFRESH AT year 30000` MV has a constant plan but no + // durable content until the refresh fires; folding it + // and blind-writing the literal would skip timestamp + // selection and linearization, producing data that was + // never observable. + // + // Preserving HIR-level Get nodes routes the INSERT + // through the RTW path, where timestamp selection + // handles REFRESH and other time-dependent reads + // correctly. + let has_read_deps = !insert_plan.values.depends_on().is_empty(); + + if !has_read_deps { + let optimized_mir = if insert_plan.values.as_const().is_some() { + // Already constant at HIR level - just lower without optimization + let expr = insert_plan + .values + .clone() + .lower(catalog.system_config(), None)?; + mz_expr::OptimizedMirRelationExpr(expr) + } else { + // Need to optimize to check if it becomes constant. + // Use one-shot expression prep so unmaterializable + // functions like current_user() are resolved before we + // decide whether this can use the blind-write path. + let optimizer_config = + optimize::OptimizerConfig::from(catalog.system_config()); + let session = self.session.as_ref().expect("SessionClient invariant"); + let prep = ExprPrepOneShot { + logical_time: EvalTime::NotAvailable, + session, + catalog_state: catalog.state(), + }; + let mut optimizer = + optimize::view::Optimizer::new_with_prep(optimizer_config, None, prep); + match optimizer.optimize(insert_plan.values.clone()) { + Ok(expr) => expr, + Err(OptimizerError::UncallableFunction { + func: UnmaterializableFunc::MzNow, + .. + }) => { + // Preserve the established user-facing `mz_now()` + // error by falling back to the RTW validator. + let expr = insert_plan + .values + .clone() + .lower(catalog.system_config(), None)?; + mz_expr::OptimizedMirRelationExpr(expr) + } + Err(e) => return Err(e.into()), + } + }; + + // Constant INSERT without RETURNING are blind-writes. Add to + // the transaction and let those code paths handle it. + let inner_mir = optimized_mir.into_inner(); + if inner_mir.as_const().is_some() && insert_plan.returning.is_empty() { + let session = self.session.as_mut().expect("SessionClient invariant"); + + let logging_guard = self.peek_client.begin_statement_logging( + session, + ¶ms, + &logging, + &catalog, + lifecycle_timestamps, + outer_ctx_extra, + ); + if let (Some(logging_id), Some(cluster_id)) = + (logging_guard.id(), target_cluster_id) + { + let cluster_name = catalog.get_cluster(cluster_id).name.clone(); + self.peek_client + .log_set_cluster(logging_id, cluster_id, cluster_name); + } + + let result = Coordinator::insert_constant( + &catalog, + session, + insert_plan.id, + inner_mir, + ); + + logging_guard.retire_with_result(&result); + + return Ok(Some(result?)); + } + } + + let desc_arity = match catalog.try_get_entry(&insert_plan.id) { + Some(table) => { + let desc = table + .relation_desc_latest() + .ok_or_else(|| AdapterError::Internal("table has no desc".into()))?; + desc.arity() + } + None => { + return Err(AdapterError::Catalog(mz_catalog::memory::error::Error { + kind: mz_catalog::memory::error::ErrorKind::Sql( + mz_sql::catalog::CatalogError::UnknownItem( + insert_plan.id.to_string(), + ), + ), + })); + } + }; + + let finishing = RowSetFinishing { + order_by: vec![], + limit: None, + offset: 0, + project: (0..desc_arity).collect(), + }; + + ReadThenWritePlan { + id: insert_plan.id, + selection: insert_plan.values, + finishing, + assignments: BTreeMap::new(), + kind: MutationKind::Insert, + returning: insert_plan.returning, + } + } + _ => { + return Err(AdapterError::Internal( + "unexpected plan type for mutation".into(), + )); + } + }; + + // Inside a transaction, only a write that reads no persisted state can + // run on this path: its diffs are frontier-independent, so they can sit + // in the session's write ops until COMMIT. A write that reads persisted + // state commits at the frontier it observed and cannot wait. + // + // The AST gate above is not enough to establish this. It admits + // INSERTs whose source is constant in the AST, and such a statement can + // still plan to a selection with `Get` nodes, because SQL-implemented + // builtins (`pg_get_viewdef`, `text` to `reg*` casts, ...) read system + // relations. So decide on the planned selection, and do it before we + // execute a dataflow for a statement we would then refuse. + { + let session = self.session.as_ref().expect("SessionClient invariant"); + if session.transaction().is_in_multi_statement_transaction() + && !rtw_plan.selection.depends_on().is_empty() + { + return Err(prohibited_in_transaction(&stmt)); + } + } + + let session = self.session.as_mut().expect("SessionClient invariant"); + self.peek_client + .frontend_read_then_write( + session, + rtw_plan, + target_cluster, + Arc::clone(&catalog), + ¶ms, + &logging, + lifecycle_timestamps, + outer_ctx_extra, + attempt_state, + ) + .await + .map(Some) + } +} + +/// Builds the error for DML that cannot run in a transaction block, mirroring +/// the coordinator's redaction in `handle_execute`: statements that can carry +/// sensitive literals are redacted because the error message is persisted in +/// `mz_statement_execution_history`. +fn prohibited_in_transaction(stmt: &Statement) -> AdapterError { + use mz_sql_parser::ast::StatementKind; + let op = if StatementKind::from(stmt).is_sensitive() { + stmt.to_ast_string_redacted() + } else { + stmt.to_string() + }; + AdapterError::OperationProhibitsTransaction(op) } impl Drop for SessionClient { diff --git a/src/adapter/src/command.rs b/src/adapter/src/command.rs index 3174ff359e99f..855d0ad146704 100644 --- a/src/adapter/src/command.rs +++ b/src/adapter/src/command.rs @@ -31,7 +31,7 @@ use mz_persist_client::PersistClient; use mz_pgcopy::CopyFormatParams; use mz_repr::global_id::TransientIdGen; use mz_repr::role_id::RoleId; -use mz_repr::{CatalogItemId, ColumnIndex, GlobalId, RowIterator, SqlRelationType}; +use mz_repr::{CatalogItemId, ColumnIndex, Diff, GlobalId, Row, RowIterator, SqlRelationType}; use mz_sql::ast::{FetchDirection, Raw, Statement}; use mz_sql::catalog::ObjectType; use mz_sql::optimizer_metrics::OptimizerMetrics; @@ -42,17 +42,18 @@ use mz_sql::session::vars::{OwnedVarInput, SystemVars}; use mz_sql_parser::ast::{AlterObjectRenameStatement, AlterOwnerStatement, DropObjectsStatement}; use mz_storage_types::sources::Timeline; use mz_timestamp_oracle::TimestampOracle; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Semaphore, mpsc, oneshot, watch}; use uuid::Uuid; use crate::catalog::Catalog; use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend}; -use crate::coord::appends::BuiltinTableAppendNotify; +use crate::coord::appends::{BuiltinTableAppendNotify, WriteResult}; use crate::coord::consistency::CoordinatorInconsistencies; use crate::coord::peek::{PeekDataflowPlan, PeekResponseUnary}; use crate::coord::timestamp_selection::TimestampDetermination; use crate::coord::{ExecuteContextExtra, ExecuteContextGuard}; use crate::error::AdapterError; +use crate::optimize::LirDataflowDescription; use crate::session::{EndTransactionAction, RowBatchStream, Session}; use crate::statement_logging::{ FrontendStatementLoggingEvent, StatementEndedExecutionReason, StatementExecutionStrategy, @@ -391,6 +392,73 @@ pub enum Command { /// Statement logging event from frontend peek sequencing. /// No response channel needed - this is fire-and-forget. FrontendStatementLogging(FrontendStatementLoggingEvent), + + /// Registers a connection-scoped cancellation watch and returns a receiver + /// that becomes `true` when cancellation is requested for the connection. + /// + /// This is shared by coordinator staged sequencing and frontend + /// read-then-write execution. `operation_id` identifies the registering + /// operation, a matching unregister only removes a watch registered with + /// the same id. + RegisterConnectionCancelWatch { + conn_id: ConnectionId, + operation_id: Uuid, + tx: oneshot::Sender>, + }, + + /// Unregisters the connection-scoped cancellation watch registered with + /// the same `operation_id`. A no-op if a newer operation has since + /// registered its own watch. + UnregisterConnectionCancelWatch { + conn_id: ConnectionId, + operation_id: Uuid, + }, + + /// Creates an internal subscribe (not visible in introspection) and returns + /// the response channel. Initially used for frontend-sequenced + /// read-then-write (DELETE/UPDATE/INSERT...SELECT) operations via OCC. + CreateInternalSubscribe { + df_desc: Box, + cluster_id: ComputeInstanceId, + replica_id: Option, + depends_on: BTreeSet, + as_of: mz_repr::Timestamp, + arity: usize, + sink_id: GlobalId, + conn_id: ConnectionId, + session_uuid: Uuid, + start_time: mz_ore::now::EpochMillis, + read_holds: ReadHolds, + tx: oneshot::Sender, AdapterError>>, + }, + + /// Submits a write attempt. Carries the accumulated diffs to write. + /// + /// `write_ts` selects between two modes: + /// - `Some(ts)`: timestamped write at a specific timestamp, fails when the + /// table timestamp is already past that. + /// - `None`: blind write where the coordinator picks the timestamp via the + /// oracle during group commit. This does not fail and will be retried + /// until the write succeeds. + AttemptWrite { + /// Connection originating the write. Used so the coordinator can + /// cancel this pending write if the connection is cancelled before + /// the write commits. + conn_id: ConnectionId, + target_id: CatalogItemId, + target_global_id: GlobalId, + diffs: Vec<(Row, Diff)>, + write_ts: Option, + tx: oneshot::Sender, + }, + + /// Drops an internal subscribe. + /// + /// Used for cleanup after the subscribe's purpose is fulfilled or on error. + /// Fire-and-forget — the caller doesn't wait for completion. + DropInternalSubscribe { + sink_id: GlobalId, + }, } impl Command { @@ -431,7 +499,12 @@ impl Command { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => None, + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::UnregisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => None, } } @@ -472,7 +545,12 @@ impl Command { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => None, + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::UnregisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => None, } } } @@ -510,6 +588,15 @@ pub struct StartupResponse { pub optimizer_metrics: OptimizerMetrics, pub persist_client: PersistClient, pub statement_logging_frontend: StatementLoggingFrontend, + /// Semaphore for limiting concurrent OCC (optimistic concurrency control) + /// write operations. + pub occ_write_semaphore: Arc, + /// Whether frontend OCC read-then-write is enabled (determined once at + /// process startup). + pub frontend_read_then_write_enabled: bool, + /// Whether the coordinator is in read-only mode (e.g. during 0dt upgrades). + /// The frontend path must reject mutations when this is true. + pub read_only: bool, } #[derive(Derivative)] diff --git a/src/adapter/src/coord.rs b/src/adapter/src/coord.rs index 29ada71ba3e6d..9040deb4c4920 100644 --- a/src/adapter/src/coord.rs +++ b/src/adapter/src/coord.rs @@ -176,7 +176,7 @@ use thiserror::Error; use timely::progress::{Antichain, Timestamp as _}; use tokio::runtime::Handle as TokioHandle; use tokio::select; -use tokio::sync::{Notify, OwnedMutexGuard, mpsc, oneshot, watch}; +use tokio::sync::{Notify, OwnedMutexGuard, Semaphore, mpsc, oneshot, watch}; use tokio::time::{Interval, MissedTickBehavior}; use tracing::{Instrument, Level, Span, debug, info, info_span, span, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; @@ -362,6 +362,8 @@ pub enum Message { responses: Vec, /// Statement executions associated with this commit. statement_logging_ids: Vec, + /// Frontend-sequenced writes to complete after local timestamp bookkeeping. + internal_results: Vec, /// The applied write timestamp. write_ts: Timestamp, }, @@ -507,6 +509,13 @@ impl Message { Command::FrontendStatementLogging(..) => "frontend-statement-logging", Command::StartCopyFromStdin { .. } => "start-copy-from-stdin", Command::InjectAuditEvents { .. } => "inject-audit-events", + Command::RegisterConnectionCancelWatch { .. } => "register-connection-cancel-watch", + Command::UnregisterConnectionCancelWatch { .. } => { + "unregister-connection-cancel-watch" + } + Command::CreateInternalSubscribe { .. } => "create-internal-subscribe", + Command::AttemptWrite { .. } => "attempt-write", + Command::DropInternalSubscribe { .. } => "drop-internal-subscribe", }, Message::ControllerReady { controller: ControllerReadiness::Compute, @@ -2060,8 +2069,12 @@ pub struct Coordinator { /// is requested for that connection, at which point it is set to `true`. /// /// Consumers install/remove these watches while they have cancellable work - /// in flight. - connection_cancel_watches: BTreeMap, watch::Receiver)>, + /// in flight. The `Option` scopes removal: frontend operations + /// register with a fresh operation id and unregister with that same id, so + /// a late unregister cannot remove a newer registration. Coordinator + /// staged sequencing registers with `None` and removes its watch itself. + connection_cancel_watches: + BTreeMap, watch::Sender, watch::Receiver)>, /// Active introspection subscribes. introspection_subscribes: BTreeMap, @@ -2073,6 +2086,24 @@ pub struct Coordinator { /// Pending writes waiting for a group commit. pending_writes: Vec, + /// Semaphore to limit concurrent OCC (optimistic concurrency control) + /// read-then-write operations. + /// + /// Each operation maintains a subscribe that continually receives and + /// consolidates updates; with N concurrent loops, every successful write + /// forces the other N-1 to redo work, so total work scales as `O(n^2)`. + /// The semaphore caps concurrency to keep that bounded. + /// + /// NOTE: The number of permits is read from `max_concurrent_occ_writes` at + /// coordinator startup; runtime changes require an `environmentd` restart. + occ_write_semaphore: Arc, + + /// Whether frontend OCC read-then-write is enabled. Read once at startup + /// from the `FRONTEND_READ_THEN_WRITE` dyncfg and fixed for the lifetime of + /// this process; see the module-level docs on `frontend_read_then_write` + /// for why mixed-mode operation is not allowed. + frontend_read_then_write_enabled: bool, + /// For the realtime timeline, an explicit SELECT or INSERT on a table will bump the /// table's timestamps, but there are cases where timestamps are not bumped but /// we expect the closed timestamps to advance (`AS OF X`, SUBSCRIBing views over @@ -3965,8 +3996,8 @@ impl Coordinator { // and make it follow from all the Spans in the pending // writes. let user_write_spans = self.pending_writes.iter().flat_map(|x| match x { - PendingWriteTxn::User{span, ..} => Some(span), - PendingWriteTxn::System{..} => None, + PendingWriteTxn::User { span, .. } => Some(span), + PendingWriteTxn::System { .. } => None, }); let span = match user_write_spans.exactly_one() { Ok(span) => span.clone(), @@ -5075,6 +5106,15 @@ pub fn serve( } let catalog = Arc::new(catalog); + // Read once at startup; changing this system variable requires + // an environmentd restart to take effect (see field doc on + // `occ_write_semaphore`). + let max_concurrent_occ_writes = + usize::cast_from(catalog.system_config().max_concurrent_occ_writes()); + let frontend_read_then_write_enabled = { + use mz_adapter_types::dyncfgs::FRONTEND_READ_THEN_WRITE; + FRONTEND_READ_THEN_WRITE.get(catalog.system_config().dyncfgs()) + }; let caching_secrets_reader = CachingSecretsReader::new(secrets_controller.reader()); let (group_committer_tx, group_committer_rx) = mpsc::unbounded_channel(); @@ -5103,6 +5143,8 @@ pub fn serve( write_locks: BTreeMap::new(), deferred_write_ops: BTreeMap::new(), pending_writes: Vec::new(), + occ_write_semaphore: Arc::new(Semaphore::new(max_concurrent_occ_writes)), + frontend_read_then_write_enabled, advance_timelines_interval, secrets_controller, caching_secrets_reader, diff --git a/src/adapter/src/coord/appends.rs b/src/adapter/src/coord/appends.rs index a0c482c179938..55c6bf590952d 100644 --- a/src/adapter/src/coord/appends.rs +++ b/src/adapter/src/coord/appends.rs @@ -50,6 +50,7 @@ use mz_ore::assert_none; use mz_ore::halt; use mz_ore::instrument; use mz_ore::now::NowFn; +use mz_ore::soft_panic_or_log; use mz_ore::task; use mz_repr::{CatalogItemId, GlobalId, Timestamp}; use mz_sql::names::ResolvedIds; @@ -164,12 +165,79 @@ pub(crate) enum BuiltinTableUpdateSource { Background(oneshot::Sender<()>), } +/// Result of a write submitted by frontend sequencing. +#[derive(Debug, Clone)] +pub enum WriteResult { + /// The write committed at this timestamp. + Success { timestamp: Timestamp }, + /// The requested timestamp was no longer eligible. + TimestampPassed { + target_timestamp: Timestamp, + next_eligible_timestamp: Timestamp, + }, + /// The write was canceled before it entered the committer. + Canceled, + /// The coordinator cannot accept writes. + ReadOnly, + /// The target table was dropped or changed after planning. + TargetChanged, + /// The committer shut down with the write's outcome unknown. + Indeterminate, +} + +/// Delivers an internal write result, including on task shutdown. +#[derive(Debug)] +pub struct InternalWriteResponder { + tx: Option>, + expected_target_global_id: GlobalId, +} + +impl InternalWriteResponder { + pub(crate) fn new( + tx: oneshot::Sender, + expected_target_global_id: GlobalId, + ) -> Self { + Self { + tx: Some(tx), + expected_target_global_id, + } + } + + pub(crate) fn send(mut self, result: WriteResult) { + if let Some(tx) = self.tx.take() { + let _ = tx.send(result); + } + } +} + +impl Drop for InternalWriteResponder { + fn drop(&mut self) { + if let Some(tx) = self.tx.take() { + let _ = tx.send(WriteResult::Indeterminate); + } + } +} + /// Where to deliver the result of a [`PendingWriteTxn::User`] write. #[derive(Debug)] pub(crate) enum UserWriteResponder { /// Session-bound write. The coordinator retires the session's /// `ExecuteContext` once the write commits. Session(PendingTxn), + /// Frontend-sequenced blind write. + Internal { + conn_id: ConnectionId, + result: InternalWriteResponder, + }, +} + +impl UserWriteResponder { + pub(crate) fn conn_id(&self) -> &ConnectionId { + match self { + UserWriteResponder::Session(pending) => pending.ctx.session().conn_id(), + UserWriteResponder::Internal { conn_id, .. } => conn_id, + } + } } /// A pending write transaction that will be committing during the next group commit. @@ -209,6 +277,7 @@ impl PendingWriteTxn { pub(crate) enum TableWriteCmd { GroupCommit(GroupCommitRequest), + TimestampedWrite(TimestampedWriteRequest), Register { tables: Vec, result: oneshot::Sender, @@ -219,6 +288,14 @@ pub(crate) enum TableWriteCmd { }, } +/// An OCC write whose diffs are valid only at `target_timestamp`. +pub(crate) struct TimestampedWriteRequest { + pub(crate) appends: Vec<(GlobalId, Vec)>, + pub(crate) target_timestamp: Timestamp, + pub(crate) result: InternalWriteResponder, + pub(crate) span: Span, +} + /// A group commit staged on the coordinator loop for the [`GroupCommitter`]. pub(crate) struct GroupCommitRequest { /// Appends resolved to their latest [`GlobalId`]. Empty for a keepalive. @@ -226,6 +303,7 @@ pub(crate) struct GroupCommitRequest { responses: Vec, statement_logging_ids: Vec, notifies: Vec>, + internal_results: Vec, write_locks: GroupCommitWriteLocks, /// In-progress permits held until the commit is applied. permits: Vec, @@ -240,6 +318,7 @@ impl GroupCommitRequest { responses, statement_logging_ids, notifies, + internal_results, write_locks, permits, contains_internal_system_write, @@ -249,6 +328,7 @@ impl GroupCommitRequest { self.responses.extend(responses); self.statement_logging_ids.extend(statement_logging_ids); self.notifies.extend(notifies); + self.internal_results.extend(internal_results); self.write_locks.extend(write_locks); self.permits.extend(permits); self.contains_internal_system_write |= contains_internal_system_write; @@ -285,6 +365,12 @@ impl GroupCommitter { ControlFlow::Break(()) => return, } } + TableWriteCmd::TimestampedWrite(request) => { + let span = request.span.clone(); + if !self.commit_timestamped(request).instrument(span).await { + return; + } + } TableWriteCmd::Register { tables, result } => { let Some(write_ts) = self .write_to_txns(None, |ts, _advance_to| { @@ -312,6 +398,85 @@ impl GroupCommitter { } } + /// Attempts an OCC write exactly once at its requested timestamp. + /// + /// Unlike blind writes, an `InvalidUppers` conflict must return to the + /// subscribe loop. Retrying the same diffs at a fresh timestamp would apply + /// a mutation to state it was not computed from. + async fn commit_timestamped(&self, request: TimestampedWriteRequest) -> bool { + let TimestampedWriteRequest { + appends, + target_timestamp, + result, + span: _, + } = request; + + let oracle_write_ts = self.oracle.peek_write_ts().await; + if target_timestamp <= oracle_write_ts { + result.send(WriteResult::TimestampPassed { + target_timestamp, + next_eligible_timestamp: oracle_write_ts.step_forward(), + }); + return true; + } + + let advance_to = target_timestamp.step_forward(); + let catalog_upper_start = Instant::now(); + self.catalog_upper + .advance_upper(advance_to) + .await + .unwrap_or_terminate("unable to advance catalog upper"); + self.metrics + .group_commit_catalog_upper_seconds + .observe(catalog_upper_start.elapsed().as_secs_f64()); + + let append_start = Instant::now(); + let append_result = self + .table_write_handle + .append(target_timestamp, advance_to, appends) + .await; + self.metrics + .append_table_duration_seconds + .observe(append_start.elapsed().as_secs_f64()); + + match append_result { + Ok(Ok(())) => {} + Ok(Err(StorageError::InvalidUppers(_))) => { + result.send(WriteResult::TimestampPassed { + target_timestamp, + next_eligible_timestamp: advance_to, + }); + return true; + } + Ok(Err(other)) => { + Err::<(), _>(other).unwrap_or_terminate("cannot fail to write to txns shard"); + unreachable!("unwrap_or_terminate does not return on Err"); + } + Err(_recv) => { + warn!("table write worker gone with a timestamped write outstanding"); + return false; + } + } + + let now: Timestamp = (self.now)().into(); + crate::coord::timeline::check_runaway_write_ts(&now, target_timestamp); + self.oracle.apply_write(target_timestamp).await; + + if self + .internal_cmd_tx + .send(Message::GroupCommitApplied { + responses: Vec::new(), + statement_logging_ids: Vec::new(), + internal_results: vec![result], + write_ts: target_timestamp, + }) + .is_err() + { + warn!("coordinator shut down before a timestamped write could be finalized"); + } + true + } + /// Writes at a fresh oracle timestamp and applies a successful write to the oracle. /// /// Returns `None` when the table write worker shuts down. @@ -439,6 +604,7 @@ impl GroupCommitter { responses, statement_logging_ids, notifies, + internal_results, write_locks, permits, contains_internal_system_write: _, @@ -487,6 +653,7 @@ impl GroupCommitter { .send(Message::GroupCommitApplied { responses, statement_logging_ids, + internal_results, write_ts: timestamp, }) .is_err() @@ -640,90 +807,116 @@ impl Coordinator { // Validate, merge, and possibly acquire write locks for as many pending writes as possible. for pending_write in pending_writes { match pending_write { - // We always allow system writes to proceed. PendingWriteTxn::System { .. } => validated_writes.push(pending_write), - // We have a set of locks! Validate they're correct (expected). PendingWriteTxn::User { span, write_locks: Some(write_locks), writes, - responder: UserWriteResponder::Session(pending_txn), + responder, } => match write_locks.validate(writes.keys().copied()) { Ok(validated_locks) => { - // Merge all of our write locks together since we can allow concurrent - // writes at the same timestamp. group_write_locks.merge(validated_locks); - - let validated_write = PendingWriteTxn::User { + validated_writes.push(PendingWriteTxn::User { span, writes, write_locks: None, - responder: UserWriteResponder::Session(pending_txn), - }; - validated_writes.push(validated_write); + responder, + }); } - // This is very unexpected since callers of this method should be validating. - // - // We cannot allow these write to occur since if the correct set of locks was - // not taken we could violate serializability. Err(missing) => { let writes: Vec<_> = writes.keys().collect(); panic!( - "got to group commit with partial set of locks!\nmissing: {:?}, writes: {:?}, txn: {:?}", - missing, writes, pending_txn, + "got to group commit with partial set of locks!\nmissing: {:?}, writes: {:?}, conn_id: {}", + missing, + writes, + responder.conn_id(), ); } }, - // If we don't have any locks, try to acquire them, otherwise defer the write. PendingWriteTxn::User { span, writes, write_locks: None, - responder: UserWriteResponder::Session(pending_txn), + responder, } => { let missing = group_write_locks.missing_locks(writes.keys().copied()); - if missing.is_empty() { - // We have all the locks! Queue the pending write. - let validated_write = PendingWriteTxn::User { + validated_writes.push(PendingWriteTxn::User { span, writes, write_locks: None, - responder: UserWriteResponder::Session(pending_txn), - }; - validated_writes.push(validated_write); - } else { - // Try to acquire the locks we're missing. - let mut just_in_time_locks = WriteLocks::builder(missing.clone()); - for collection in missing { - if let Some(lock) = self.try_grant_object_write_lock(collection) { - just_in_time_locks.insert_lock(collection, lock); + responder, + }); + continue; + } + + match responder { + UserWriteResponder::Session(pending_txn) => { + let mut just_in_time_locks = WriteLocks::builder(missing.clone()); + for collection in missing { + if let Some(lock) = self.try_grant_object_write_lock(collection) { + just_in_time_locks.insert_lock(collection, lock); + } + } + match just_in_time_locks + .all_or_nothing(pending_txn.ctx.session().conn_id()) + { + Ok(locks) => { + group_write_locks.merge(locks); + validated_writes.push(PendingWriteTxn::User { + span, + writes, + write_locks: None, + responder: UserWriteResponder::Session(pending_txn), + }); + } + Err(missing) => { + let acquire_future = + self.grant_object_write_lock(missing).map(Option::Some); + deferred_writes.push(( + acquire_future, + DeferredWrite { + span, + writes, + pending_txn, + }, + )); + } } } - - match just_in_time_locks.all_or_nothing(pending_txn.ctx.session().conn_id()) - { - // We acquired all of the locks! Proceed with the write. - Ok(locks) => { - group_write_locks.merge(locks); - let validated_write = PendingWriteTxn::User { + UserWriteResponder::Internal { conn_id, result } => { + let acquired = missing + .into_iter() + .map(|id| { + self.try_grant_object_write_lock(id).map(|lock| (id, lock)) + }) + .collect::>>(); + if let Some(acquired) = acquired { + for (id, lock) in acquired { + group_write_locks.insert_lock(id, lock); + } + validated_writes.push(PendingWriteTxn::User { span, writes, write_locks: None, - responder: UserWriteResponder::Session(pending_txn), - }; - validated_writes.push(validated_write); - } - // Darn. We couldn't acquire the locks, defer the write. - Err(missing) => { - let acquire_future = - self.grant_object_write_lock(missing).map(Option::Some); - let write = DeferredWrite { + responder: UserWriteResponder::Internal { conn_id, result }, + }); + } else { + // Retry by riding the next group commit + // initiate, at the latest the periodic + // timeline advancement tick. Internal writes + // have no `ExecuteContext`, so they can't use + // `defer_op` like session writes. Lock hold + // times are short while frontend OCC + // sequencing is enabled because the + // coordinator's lock-based read-then-write + // path is disabled. + self.pending_writes.push(PendingWriteTxn::User { span, writes, - pending_txn, - }; - deferred_writes.push((acquire_future, write)); + write_locks: None, + responder: UserWriteResponder::Internal { conn_id, result }, + }); } } } @@ -744,6 +937,7 @@ impl Coordinator { let mut responses = Vec::with_capacity(validated_writes.len()); let mut statement_logging_ids = Vec::new(); let mut notifies = Vec::new(); + let mut internal_results = Vec::new(); for validated_write_txn in validated_writes { match validated_write_txn { @@ -775,6 +969,35 @@ impl Coordinator { responses.push(CompletedClientTransmitter::new(ctx, response, action)); } + PendingWriteTxn::User { + span: _, + writes, + write_locks, + responder: UserWriteResponder::Internal { result, .. }, + } => { + assert_none!(write_locks, "should have merged together all locks above"); + // Frontend internal writes target exactly one table, + // enforced at enqueue time in `handle_attempt_write`. + let mut writes = writes.into_iter(); + let (target_id, table_data) = match (writes.next(), writes.next()) { + (Some(write), None) => write, + _ => { + soft_panic_or_log!("internal write must target exactly one table"); + result.send(WriteResult::TargetChanged); + continue; + } + }; + let current_global_id = self + .catalog() + .try_get_entry(&target_id) + .map(|entry| entry.latest_global_id()); + if current_global_id != Some(result.expected_target_global_id) { + result.send(WriteResult::TargetChanged); + continue; + } + appends.entry(target_id).or_default().extend(table_data); + internal_results.push(result); + } PendingWriteTxn::System { updates, source } => { for update in updates { appends.entry(update.id).or_default().push(update.data); @@ -821,6 +1044,7 @@ impl Coordinator { responses, statement_logging_ids, notifies, + internal_results, write_locks: group_write_locks, permits: permit.into_iter().collect(), contains_internal_system_write, diff --git a/src/adapter/src/coord/command_handler.rs b/src/adapter/src/coord/command_handler.rs index e6046a8a501e9..cbb98f08d444d 100644 --- a/src/adapter/src/coord/command_handler.rs +++ b/src/adapter/src/coord/command_handler.rs @@ -62,7 +62,7 @@ use mz_sql_parser::ast::{ }; use mz_storage_types::sources::Timeline; use opentelemetry::trace::TraceContextExt; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{mpsc, oneshot, watch}; use tracing::{Instrument, debug_span, info, warn}; use tracing_opentelemetry::OpenTelemetrySpanExt; use uuid::Uuid; @@ -71,7 +71,7 @@ use crate::command::{ CatalogSnapshot, Command, ExecuteResponse, Response, SASLChallengeResponse, SASLVerifyProofResponse, StartupResponse, SuperuserAttribute, }; -use crate::coord::appends::{PendingWriteTxn, UserWriteResponder}; +use crate::coord::appends::{PendingWriteTxn, UserWriteResponder, WriteResult}; use crate::coord::peek::PendingPeek; use crate::coord::{ ConnMeta, Coordinator, DeferredPlanStatement, Message, PendingTxn, PlanStatement, PlanValidity, @@ -614,6 +614,85 @@ impl Coordinator { Command::FrontendStatementLogging(event) => { self.handle_frontend_statement_logging_event(event); } + Command::RegisterConnectionCancelWatch { + conn_id, + operation_id, + tx, + } => { + // Always replace any existing entry. Another code path + // (e.g. `sequence_staged`) may have left a stale watch + // here, possibly already signaled `true` from a prior + // cancel. Reusing it via `or_insert_with` would hand out + // a `Receiver` that already reads `true`, causing the new + // operation to immediately return `Canceled` even though + // it hasn't been cancelled. + let (watch_tx, watch_rx) = watch::channel(false); + self.connection_cancel_watches + .insert(conn_id, (Some(operation_id), watch_tx, watch_rx.clone())); + let _ = tx.send(watch_rx); + } + Command::UnregisterConnectionCancelWatch { + conn_id, + operation_id, + } => { + if self + .connection_cancel_watches + .get(&conn_id) + .is_some_and(|(registered_id, _, _)| *registered_id == Some(operation_id)) + { + self.connection_cancel_watches.remove(&conn_id); + } + } + Command::CreateInternalSubscribe { + df_desc, + cluster_id, + replica_id, + depends_on, + as_of, + arity, + sink_id, + conn_id, + session_uuid, + start_time, + read_holds, + tx, + } => { + self.handle_create_internal_subscribe( + *df_desc, + cluster_id, + replica_id, + depends_on, + as_of, + arity, + sink_id, + conn_id, + session_uuid, + start_time, + read_holds, + tx, + ) + .await; + } + Command::AttemptWrite { + conn_id, + target_id, + target_global_id, + diffs, + write_ts, + tx, + } => { + self.handle_attempt_write( + conn_id, + target_id, + target_global_id, + diffs, + write_ts, + tx, + ); + } + Command::DropInternalSubscribe { sink_id } => { + self.drop_internal_subscribe(sink_id).await; + } } } .instrument(debug_span!("handle_command")) @@ -902,6 +981,9 @@ impl Coordinator { persist_client: self.persist_client.clone(), statement_logging_frontend, superuser_attribute, + occ_write_semaphore: Arc::clone(&self.occ_write_semaphore), + frontend_read_then_write_enabled: self.frontend_read_then_write_enabled, + read_only: self.controller.read_only(), }); if tx.send(resp).is_err() { // Failed to send to adapter, but everything is setup so we can terminate @@ -1902,20 +1984,41 @@ impl Coordinator { pub(crate) async fn handle_privileged_cancel(&mut self, conn_id: ConnectionId) { let mut maybe_ctx = None; - // Cancel pending writes. There is at most one pending write per session. - let pending_write_idx = self.pending_writes.iter().position(|pending_write_txn| { - matches!(pending_write_txn, PendingWriteTxn::User { - responder: UserWriteResponder::Session(PendingTxn { ctx, .. }), - .. - } if *ctx.session().conn_id() == conn_id) - }); - if let Some(idx) = pending_write_idx { - if let PendingWriteTxn::User { - responder: UserWriteResponder::Session(PendingTxn { ctx, .. }), - .. - } = self.pending_writes.remove(idx) - { - maybe_ctx = Some(ctx); + // Cancel all pending writes for this connection: + // - At most one session-bound write (`UserWriteResponder::Session`); + // retired via its `ExecuteContext` below. + // - Any frontend blind write that has not entered the committer. The + // waiter receives `WriteResult::Canceled`. Timestamped writes already + // in the committer have an indeterminate outcome until it responds. + let (cancelled, kept): (Vec<_>, Vec<_>) = std::mem::take(&mut self.pending_writes) + .into_iter() + .partition(|pending_write_txn| match pending_write_txn { + PendingWriteTxn::User { + responder: UserWriteResponder::Session(PendingTxn { ctx, .. }), + .. + } => *ctx.session().conn_id() == conn_id, + PendingWriteTxn::User { + responder: UserWriteResponder::Internal { conn_id: c, .. }, + .. + } => *c == conn_id, + PendingWriteTxn::System { .. } => false, + }); + self.pending_writes = kept; + for pending in cancelled { + match pending { + PendingWriteTxn::User { + responder: UserWriteResponder::Session(PendingTxn { ctx, .. }), + .. + } => { + maybe_ctx = Some(ctx); + } + PendingWriteTxn::User { + responder: UserWriteResponder::Internal { result, .. }, + .. + } => { + result.send(WriteResult::Canceled); + } + PendingWriteTxn::System { .. } => unreachable!("filtered out above"), } } @@ -1957,7 +2060,7 @@ impl Coordinator { self.cancel_cluster_reconfigurations_for_conn(&conn_id) .await; self.cancel_pending_copy(&conn_id); - if let Some((tx, _rx)) = self.connection_cancel_watches.get_mut(&conn_id) { + if let Some((_operation_id, tx, _rx)) = self.connection_cancel_watches.get_mut(&conn_id) { let _ = tx.send(true); } } diff --git a/src/adapter/src/coord/message_handler.rs b/src/adapter/src/coord/message_handler.rs index c9eb4e0e17c95..988000c1abfb4 100644 --- a/src/adapter/src/coord/message_handler.rs +++ b/src/adapter/src/coord/message_handler.rs @@ -114,6 +114,7 @@ impl Coordinator { Message::GroupCommitApplied { responses, statement_logging_ids, + internal_results, write_ts, } => { // Record statement timestamps before retiring, since retiring ends the statement @@ -130,6 +131,11 @@ impl Coordinator { // that and we can downgrade the local read holds without an oracle round trip. self.downgrade_local_read_holds(write_ts); self.advance_custom_timelines().boxed_local().await; + for result in internal_results { + result.send(crate::coord::appends::WriteResult::Success { + timestamp: write_ts, + }); + } } Message::AdvanceTimelines => { // Only sent by the periodic tick in read-only mode, where group commits (which diff --git a/src/adapter/src/coord/read_then_write.rs b/src/adapter/src/coord/read_then_write.rs index ff2bdd336d69b..6afa9d886aa62 100644 --- a/src/adapter/src/coord/read_then_write.rs +++ b/src/adapter/src/coord/read_then_write.rs @@ -8,14 +8,25 @@ // by the Apache License, Version 2.0. //! Coordinator-side support machinery for (frontend) read-then write. +//! +//! N.B. It's a bit annoying that we still have the write submission go through +//! the coordinator. We can imagine in the long run we want a group-commit task +//! that runs independently, and we can directly submit write requests there. use std::collections::BTreeSet; use mz_catalog::memory::objects::CatalogItem; use mz_repr::CatalogItemId; +use mz_repr::{Diff, GlobalId, Row, Timestamp}; use mz_sql::catalog::CatalogItemType; +use mz_sql::plan::SubscribeOutput; +use tokio::sync::mpsc; +use crate::PeekResponseUnary; +use crate::active_compute_sink::{ActiveComputeSink, ActiveSubscribe}; use crate::catalog::Catalog; +use crate::coord::Coordinator; +use crate::coord::appends::WriteResult; use crate::error::AdapterError; /// Adds `id` to the worklist the first time it is seen, enforcing the @@ -41,6 +52,151 @@ fn enqueue( Ok(()) } +impl Coordinator { + /// Handle a Command to create an internal subscribe. + /// + /// Internal subscribes are not visible in introspection collections. They + /// are initially used for frontend-sequenced read-then-write + /// (DELETE/UPDATE/INSERT ...SELECT) via OCC. + /// + /// This is called from the frontend OCC implementation after it has + /// acquired the semaphore permit. We create the subscribe here (on the + /// coordinator) and return the channel to the caller. + /// + /// The `read_holds` parameter contains the read holds for this specific + /// operation. They are passed directly through the stages (not via the + /// connection-keyed txn_read_holds map) to avoid issues where multiple + /// operations on the same connection could interfere with each other's + /// holds. + #[allow(clippy::too_many_arguments)] + pub(crate) async fn handle_create_internal_subscribe( + &mut self, + df_desc: crate::optimize::LirDataflowDescription, + cluster_id: mz_compute_types::ComputeInstanceId, + replica_id: Option, + depends_on: BTreeSet, + as_of: Timestamp, + arity: usize, + sink_id: GlobalId, + conn_id: mz_adapter_types::connection::ConnectionId, + session_uuid: uuid::Uuid, + start_time: mz_ore::now::EpochMillis, + read_holds: crate::ReadHolds, + response_tx: tokio::sync::oneshot::Sender< + Result, AdapterError>, + >, + ) { + // Client disconnected while waiting for the semaphore. + if !self.active_conns.contains_key(&conn_id) { + let _ = response_tx.send(Err(AdapterError::Canceled)); + return; + } + + let (tx, rx) = mpsc::unbounded_channel(); + + let active_subscribe = ActiveSubscribe { + conn_id: conn_id.clone(), + session_uuid, + channel: tx, + emit_progress: true, // We need progress updates for OCC + as_of, + arity, + cluster_id, + depends_on, + start_time, + output: SubscribeOutput::Diffs, + internal: true, // skip builtin table updates and metrics + }; + active_subscribe.initialize(); + + let write_notify_fut = + self.add_active_compute_sink(sink_id, ActiveComputeSink::Subscribe(active_subscribe)); + let ship_dataflow_fut = self.ship_dataflow(df_desc, cluster_id, replica_id); + let ((), ()) = futures::future::join(write_notify_fut, ship_dataflow_fut).await; + + let _ = response_tx.send(Ok(rx)); + + // Drop read holds only after `ship_dataflow` returns, so the since + // can't advance past `as_of` before the dataflow is running. + drop(read_holds); + } + + /// Enqueues a write attempt from the frontend OCC loop. + pub(crate) fn handle_attempt_write( + &mut self, + conn_id: mz_adapter_types::connection::ConnectionId, + target_id: mz_repr::CatalogItemId, + target_global_id: GlobalId, + diffs: Vec<(Row, Diff)>, + write_ts: Option, + result_tx: tokio::sync::oneshot::Sender, + ) { + use crate::coord::appends::{ + InternalWriteResponder, PendingWriteTxn, TableWriteCmd, TimestampedWriteRequest, + UserWriteResponder, + }; + use mz_storage_client::client::TableData; + use smallvec::smallvec; + use std::collections::BTreeMap; + use tracing::Span; + + let result = InternalWriteResponder::new(result_tx, target_global_id); + if !self.active_conns.contains_key(&conn_id) { + result.send(WriteResult::Canceled); + return; + } + if self.controller.read_only() { + result.send(WriteResult::ReadOnly); + return; + } + + let current_global_id = self + .catalog() + .try_get_entry(&target_id) + .map(|entry| entry.latest_global_id()); + if current_global_id != Some(target_global_id) { + result.send(WriteResult::TargetChanged); + return; + } + + let table_data = TableData::Rows(diffs); + match write_ts { + Some(target_timestamp) => { + let request = TimestampedWriteRequest { + appends: vec![(target_global_id, vec![table_data])], + target_timestamp, + result, + span: Span::current(), + }; + if self + .group_committer_tx + .send(TableWriteCmd::TimestampedWrite(request)) + .is_err() + { + tracing::warn!("group committer task gone, dropping timestamped write"); + } + } + None => { + let writes = BTreeMap::from([(target_id, smallvec![table_data])]); + self.pending_writes.push(PendingWriteTxn::User { + span: Span::current(), + writes, + write_locks: None, + responder: UserWriteResponder::Internal { conn_id, result }, + }); + self.trigger_group_commit(); + } + } + } + + /// Drop an internal subscribe. + pub(crate) async fn drop_internal_subscribe(&mut self, sink_id: GlobalId) { + // Use drop_compute_sink instead of remove_active_compute_sink to also + // cancel the dataflow on the compute side, not just remove bookkeeping. + let _ = self.drop_compute_sink(sink_id).await; + } +} + /// Validates that all dependencies are valid for read-then-write operations. /// /// Ensures all objects the selection transitively depends on (seeded by `ids`) are valid for diff --git a/src/adapter/src/coord/sequencer.rs b/src/adapter/src/coord/sequencer.rs index 26eec2c888289..5a836300ee265 100644 --- a/src/adapter/src/coord/sequencer.rs +++ b/src/adapter/src/coord/sequencer.rs @@ -895,6 +895,13 @@ impl Coordinator { // Consolidate rows. This is useful e.g. for an UPDATE where the row // doesn't change, and we need to reflect that in the number of // affected rows. + // + // NOTE: This behavior differs from PostgreSQL. In PostgreSQL, + // `UPDATE t SET x = x` reports the number of rows matching the WHERE + // clause, even if no values actually change. In Materialize, because + // we use differential dataflow, the +1 and -1 diffs for unchanged rows + // cancel out during consolidation, resulting in 0 affected rows. + // This has been Materialize's behavior since early versions. differential_dataflow::consolidation::consolidate(&mut plan.updates); affected_rows = Diff::ZERO; diff --git a/src/adapter/src/coord/sequencer/inner.rs b/src/adapter/src/coord/sequencer/inner.rs index 4b62df13d3edb..d56c0e488b76d 100644 --- a/src/adapter/src/coord/sequencer/inner.rs +++ b/src/adapter/src/coord/sequencer/inner.rs @@ -20,7 +20,9 @@ use futures::{Future, StreamExt, future}; use itertools::Itertools; use mz_adapter_types::compaction::CompactionWindow; use mz_adapter_types::connection::ConnectionId; -use mz_adapter_types::dyncfgs::{ENABLE_PASSWORD_AUTH, READ_THEN_WRITE_MAX_DEPENDENCIES}; +use mz_adapter_types::dyncfgs::{ + ENABLE_PASSWORD_AUTH, FRONTEND_READ_THEN_WRITE, READ_THEN_WRITE_MAX_DEPENDENCIES, +}; use mz_catalog::memory::error::ErrorKind; use mz_catalog::memory::objects::{ CatalogItem, Connection, DataSourceDesc, Sink, Source, Table, TableDataSource, Type, @@ -75,7 +77,7 @@ use mz_sql::plan::{ use mz_sql::session::metadata::SessionMetadata; use mz_sql::session::user::UserKind; use mz_sql::session::vars::{ - self, IsolationLevel, NETWORK_POLICY, OwnedVarInput, SCHEMA_ALIAS, + self, IsolationLevel, MAX_CONCURRENT_OCC_WRITES, NETWORK_POLICY, OwnedVarInput, SCHEMA_ALIAS, TRANSACTION_ISOLATION_VAR_NAME, Var, VarError, VarInput, }; use mz_sql::{plan, rbac}; @@ -210,9 +212,12 @@ impl Coordinator { if cancel_enabled { // Channel to await cancellation. Insert a new channel, but check if the previous one // was already canceled. - if let Some((_prev_tx, prev_rx)) = self + if let Some((_prev_operation_id, _prev_tx, prev_rx)) = self .connection_cancel_watches - .insert(session.conn_id().clone(), watch::channel(false)) + .insert(session.conn_id().clone(), { + let (tx, rx) = watch::channel(false); + (None, tx, rx) + }) { let was_canceled = *prev_rx.borrow(); if was_canceled { @@ -269,7 +274,7 @@ impl Coordinator { T: Send + 'static, F: FnOnce(C, T) + Send + 'static, { - let rx: BoxFuture<()> = if let Some((_tx, rx)) = ctx + let rx: BoxFuture<()> = if let Some((_operation_id, _tx, rx)) = ctx .session() .and_then(|session| self.connection_cancel_watches.get(session.conn_id())) { @@ -2713,6 +2718,17 @@ impl Coordinator { return; } + // The lock-based and OCC paths do not synchronize with each other, so + // reaching this path while frontend sequencing is enabled is a routing + // bug that could corrupt data. + if self.frontend_read_then_write_enabled { + ctx.retire(Err(AdapterError::Internal( + "coordinator read-then-write reached while frontend OCC sequencing is enabled" + .into(), + ))); + return; + } + let mut source_ids: BTreeSet<_> = plan .selection .depends_on() @@ -4197,6 +4213,7 @@ impl Coordinator { }; self.catalog_transact(Some(session), vec![op]).await?; + Self::notice_if_startup_only(session, &name); session.add_notice(AdapterNotice::VarDefaultUpdated { role: None, var_name: Some(name), @@ -4213,6 +4230,7 @@ impl Coordinator { self.is_user_allowed_to_alter_system(session, Some(&name))?; let op = catalog::Op::ResetSystemConfiguration { name: name.clone() }; self.catalog_transact(Some(session), vec![op]).await?; + Self::notice_if_startup_only(session, &name); session.add_notice(AdapterNotice::VarDefaultUpdated { role: None, var_name: Some(name), @@ -4227,8 +4245,16 @@ impl Coordinator { _: plan::AlterSystemResetAllPlan, ) -> Result { self.is_user_allowed_to_alter_system(session, None)?; + // Which parameters `RESET ALL` changes has to be read before the + // transaction applies it, afterwards they all read as their default. + let startup_only_changed = self.startup_only_vars_changed_by_reset_all(); let op = catalog::Op::ResetAllSystemConfiguration; self.catalog_transact(Some(session), vec![op]).await?; + for name in startup_only_changed { + session.add_notice(AdapterNotice::StartupOnlyVarUpdated { + var_name: name.to_string(), + }); + } session.add_notice(AdapterNotice::VarDefaultUpdated { role: None, var_name: None, @@ -4236,6 +4262,63 @@ impl Coordinator { Ok(ExecuteResponse::AlteredSystemConfiguration) } + /// System parameters whose value `environmentd` samples once at startup. + /// + /// `enable_adapter_frontend_occ_read_then_write` selects between the + /// lock-based and the OCC read-then-write path. Both are never live in one + /// process, so the choice is fixed at boot and every session inherits it. + /// `max_concurrent_occ_writes` sizes the OCC semaphore at boot. + fn startup_only_vars() -> [&'static str; 2] { + [ + FRONTEND_READ_THEN_WRITE.name(), + MAX_CONCURRENT_OCC_WRITES.name(), + ] + } + + /// Warns that the parameter an `ALTER SYSTEM SET`/`RESET` names is only + /// read at startup, so the running process keeps its sampled value. + /// + /// We let the change through: the catalog value is what the next process + /// start reads, and the running process cannot observe it, so there is no + /// window where two code paths are live at once. + /// + /// The operator named the parameter, so restating the restart requirement + /// is appropriate whether or not the value actually changed. + /// `startup_only_vars_changed_by_reset_all` is value-based instead, because + /// `RESET ALL` names every parameter. + fn notice_if_startup_only(session: &Session, name: &str) { + if Self::startup_only_vars() + .iter() + .any(|n| n.eq_ignore_ascii_case(name)) + { + session.add_notice(AdapterNotice::StartupOnlyVarUpdated { + var_name: name.to_string(), + }); + } + } + + /// The startup-only parameters whose value `ALTER SYSTEM RESET ALL` would + /// change. Parameters already at their effective default are untouched, so + /// they are not reported. + fn startup_only_vars_changed_by_reset_all(&self) -> Vec<&'static str> { + let config = self.catalog().system_config(); + let defaults = config.defaults(); + Self::startup_only_vars() + .into_iter() + .filter(|name| { + // Both names are registered system vars, a lookup failure + // would mean the definitions and this list have drifted apart. + let current = config + .get(name) + .expect("startup-only parameter is a registered system var") + .value(); + defaults + .get(*name) + .is_some_and(|default| default != ¤t) + }) + .collect() + } + // TODO(jkosh44) Move this into rbac.rs once RBAC is always on. fn is_user_allowed_to_alter_system( &self, diff --git a/src/adapter/src/frontend_read_then_write.rs b/src/adapter/src/frontend_read_then_write.rs new file mode 100644 index 0000000000000..ef5e906948660 --- /dev/null +++ b/src/adapter/src/frontend_read_then_write.rs @@ -0,0 +1,1504 @@ +// 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. + +//! Frontend sequencing for read-then-write operations. +//! +//! This module implements INSERT [...] SELECT FROM [...], DELETE and UPDATE +//! operations using a subscribe with optimistic concurrency control (OCC), +//! sequenced from the session task rather than the Coordinator. This reduces +//! coordinator bottlenecking. +//! +//! The approach is: +//! 1. Validate and optimize MIR locally +//! 2. Determine timestamp via coordinator +//! 3. Optimize LIR locally +//! 4. Acquire OCC semaphore +//! 5. Create subscribe via Coordinator Command +//! 6. Run OCC loop (receive diffs, attempt write, retry on conflict) +//! 7. Return the result. The write either went through the coordinator during +//! the loop, or, for a selection that reads no persisted state inside a +//! transaction, we buffer the diffs as session write ops that land at COMMIT +//! +//! ## Rollout note +//! +//! The `FRONTEND_READ_THEN_WRITE` dyncfg is read once at process startup and +//! fixed for the lifetime of the `environmentd` process. This avoids a +//! mixed-mode window where both the lock-based coordinator path and this OCC +//! path are active concurrently — the coordinator path acquires write locks to +//! prevent concurrent writes between its read and write phases, but this OCC +//! path does not use write locks, so concurrent operation of both paths could +//! allow an OCC write to slip between a coordinator-path reader's read and +//! write. + +use std::collections::BTreeMap; +use std::collections::BTreeSet; +use std::num::{NonZeroI64, NonZeroUsize}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU8, Ordering}; +use std::time::Duration; + +use bytesize::ByteSize; +use differential_dataflow::consolidation; +use itertools::Itertools; +use mz_catalog::memory::error::ErrorKind; +use mz_cluster_client::ReplicaId; +use mz_compute_types::ComputeInstanceId; +use mz_expr::Eval; +use mz_expr::row::RowCollection; +use mz_expr::{CollectionPlan, Id, LocalId, MirRelationExpr, MirScalarExpr, RowSetFinishing}; +use mz_ore::cast::CastFrom; +use mz_ore::soft_panic_or_log; +use mz_repr::optimize::OverrideFrom; +use mz_repr::{ + CatalogItemId, Diff, GlobalId, IntoRowIterator, RelationDesc, Row, RowArena, Timestamp, +}; +use mz_sql::catalog::CatalogError; +use mz_sql::plan::{self, MutationKind, Params, QueryWhen}; +use mz_sql::session::metadata::SessionMetadata; +use mz_storage_client::client::TableData; +use prometheus::Histogram; +use qcell::QCell; +use timely::progress::Antichain; +use tokio::sync::mpsc; +use uuid::Uuid; + +use crate::catalog::Catalog; +use crate::command::{Command, ExecuteResponse}; +use crate::coord::appends::WriteResult; +use crate::coord::read_then_write::validate_read_then_write_dependencies; +use crate::coord::timestamp_selection::TimestampProvider; +use crate::coord::{Coordinator, ExecuteContextGuard, TargetCluster}; +use crate::error::AdapterError; +use crate::optimize::Optimize; +use crate::optimize::dataflows::{ComputeInstanceSnapshot, EvalTime, ExprPrep, ExprPrepOneShot}; +use crate::session::{LifecycleTimestamps, Session, TransactionOps, WriteOp}; +use crate::statement_logging::{ + PreparedStatementLoggingInfo, StatementLifecycleEvent, StatementLoggingId, +}; +use crate::{PeekClient, PeekResponseUnary, TimelineContext, optimize}; + +/// Reason a frontend write attempt is being torn down early. +#[derive(Clone, Copy)] +pub(crate) enum FrontendWriteCancellation { + Canceled, + StatementTimeout, +} + +const CANCELLATION_NONE: u8 = 0; +const CANCELLATION_CANCELED: u8 = 1; +const CANCELLATION_STATEMENT_TIMEOUT: u8 = 2; + +/// State shared between an in-flight frontend write attempt and its +/// cancellation wrapper, +/// `SessionClient::try_frontend_read_then_write_with_cancel`. +/// +/// The contract: `write_submitted` is true from just before the +/// `AttemptWrite` command is sent until the attempt resolves as definitively +/// not committed. While it is true, cancellation and statement timeout must +/// not synthesize an error but await the definitive write result instead, +/// because the write may already be durable. +pub(crate) struct FrontendWriteAttemptState { + write_submitted: AtomicBool, + /// One of the `CANCELLATION_*` codes, set at most once. + cancellation: AtomicU8, +} + +impl FrontendWriteAttemptState { + pub(crate) fn new() -> Self { + Self { + write_submitted: AtomicBool::new(false), + cancellation: AtomicU8::new(CANCELLATION_NONE), + } + } + + pub(crate) fn mark_write_submitted(&self) { + self.write_submitted.store(true, Ordering::Release); + } + + /// Marks the submitted write as definitively not committed. + /// + /// NOTE: This must only be called for outcomes where the write is known + /// to not have landed (`TimestampPassed`). Terminal outcomes leave + /// `write_submitted` set so a concurrent cancellation path can never + /// fabricate an error for a write that may have committed. + fn mark_write_resolved(&self) { + self.write_submitted.store(false, Ordering::Release); + } + + pub(crate) fn write_submitted(&self) -> bool { + self.write_submitted.load(Ordering::Acquire) + } + + pub(crate) fn request(&self, cancellation: FrontendWriteCancellation) { + let code = match cancellation { + FrontendWriteCancellation::Canceled => CANCELLATION_CANCELED, + FrontendWriteCancellation::StatementTimeout => CANCELLATION_STATEMENT_TIMEOUT, + }; + let _ = self.cancellation.compare_exchange( + CANCELLATION_NONE, + code, + Ordering::AcqRel, + Ordering::Acquire, + ); + } + + fn requested_error(&self) -> Option { + match self.cancellation.load(Ordering::Acquire) { + CANCELLATION_CANCELED => Some(AdapterError::Canceled), + CANCELLATION_STATEMENT_TIMEOUT => Some(AdapterError::StatementTimeout), + _ => None, + } + } +} + +/// What the OCC loop produced. +enum OccOutcome { + /// The write is durable at `write_ts`, or there was nothing to write. + Committed { + response: ExecuteResponse, + write_ts: Option, + }, + /// These diffs are a blind write: the subscribe ran to completion, which + /// only a dataflow over no persisted inputs does, so the diffs do not + /// depend on any read frontier. They belong to the session's transaction + /// and must not become visible before it commits, so the caller buffers + /// them as write ops. + Deferred { + response: ExecuteResponse, + diffs: Vec<(Row, Diff)>, + }, +} + +/// A handle to an internal subscribe (not visible in introspection collections +/// like `mz_subscriptions`). A `Drop` impl ensures the subscribe's dataflow is +/// cleaned up when dropped. +pub(crate) struct SubscribeHandle { + rx: mpsc::UnboundedReceiver, + sink_id: GlobalId, + /// Wrapped in `Option` so we can move it out in `Drop`. + client: Option, +} + +impl SubscribeHandle { + /// Receive the next message from the subscribe, waiting if necessary. + pub async fn recv(&mut self) -> Option { + self.rx.recv().await + } + + /// Try to receive a message without waiting. + pub fn try_recv(&mut self) -> Result { + self.rx.try_recv() + } +} + +impl Drop for SubscribeHandle { + fn drop(&mut self) { + if let Some(client) = self.client.take() { + // Fire-and-forget: if the coordinator is gone, the subscribe will + // be cleaned up when the process exits anyway. + client.try_send(Command::DropInternalSubscribe { + sink_id: self.sink_id, + }); + } + } +} + +impl PeekClient { + /// Execute a read-then-write operation using frontend sequencing. + /// + /// Called by session code when the frontend_read_then_write dyncfg is + /// enabled. + pub(crate) async fn frontend_read_then_write( + &mut self, + session: &mut Session, + plan: plan::ReadThenWritePlan, + target_cluster: TargetCluster, + catalog: Arc, + params: &Params, + logging: &Arc>, + lifecycle_timestamps: Option, + outer_ctx_extra: &mut Option, + attempt_state: Arc, + ) -> Result { + // The caller verified and planned the portal against this snapshot. + // Keep it through optimization and write-target generation capture. + // The guard's `Drop` impl emits `Aborted` if the inner future is + // dropped mid-flight, so the end-execution event is never skipped. + let logging_guard = self.begin_statement_logging( + session, + params, + logging, + &catalog, + lifecycle_timestamps, + outer_ctx_extra, + ); + + let result = self + .frontend_read_then_write_inner( + session, + plan, + target_cluster, + &catalog, + logging_guard.id(), + attempt_state, + ) + .await; + + logging_guard.retire_with_result(&result); + + result + } + + /// Separated from the outer function so the statement-logging guard always + /// retires on the same return path. + async fn frontend_read_then_write_inner( + &mut self, + session: &mut Session, + mut plan: plan::ReadThenWritePlan, + target_cluster: TargetCluster, + catalog: &Arc, + statement_logging_id: Option, + attempt_state: Arc, + ) -> Result { + // A transaction that has taken a timestamped read, was opened READ + // ONLY, or is committed to some other kind of operation cannot take a + // write. Check up front, mirroring `sequence_insert`: the marker op + // below rejects only some of those states, and only with its own + // errors, so without this check the reported error and SQLSTATE would + // depend on which path sequenced the statement. + // + // Both this and the marker op require an open transaction. The + // frontends start one before they execute anything, and a `Failed` + // transaction only ever admits COMMIT/ROLLBACK, so DML never arrives + // in a state where these panic. + if !session.transaction().allows_writes() { + return Err(AdapterError::ReadOnlyTransaction); + } + + let validation_result = + self.validate_read_then_write(catalog, session, &plan, target_cluster)?; + + let ValidationResult { + cluster_id, + replica_id, + timeline, + depends_on, + table_desc, + } = validation_result; + + if let Some(logging_id) = statement_logging_id { + let cluster_name = catalog.get_cluster(cluster_id).name.clone(); + self.log_set_cluster(logging_id, cluster_id, cluster_name); + } + + // Mark this as a write transaction in the session state machine. For a + // single statement that lets auto-commit handle the write correctly. + // The rows are added later: either the coordinator's group commit + // applies them directly, or, in a transaction, they are buffered as + // write ops once we know them. + session.add_transaction_ops(TransactionOps::Writes(vec![]))?; + + // A write that reads persisted state commits at the frontier it + // observed, which cannot be deferred to COMMIT, so inside a + // transaction only writes that read nothing get here. See the + // read-dependency check in + // `SessionClient::try_frontend_read_then_write`. + let defer_write = session.transaction().is_in_multi_statement_transaction(); + + // Prepare expressions (resolve unmaterializable functions like + // current_user()) + let style = ExprPrepOneShot { + logical_time: EvalTime::NotAvailable, // We already errored out on mz_now above. + session, + catalog_state: catalog.state(), + }; + for expr in plan + .assignments + .values_mut() + .chain(plan.returning.iter_mut()) + { + style.prep_scalar_expr(expr)?; + } + + let (optimizer, global_mir_plan) = + self.optimize_mir_read_then_write(catalog, session, &plan, cluster_id)?; + + // Acquire the OCC semaphore permit *before* acquiring read holds in + // `frontend_determine_timestamp`. Under contention, waiters will + // otherwise sit on read holds on the RTW's read dependencies for the + // entire time they are queued, pinning compaction on those + // collections. Waiting on the permit first keeps queued operations + // hold-free; once we have a permit we proceed to acquire the read + // holds needed for the rest of the operation. + // + // The semaphore is owned by the coordinator and outlives every + // session task, so `acquire_owned` cannot return `Err` in practice. + let permit = Arc::clone(&self.occ_write_semaphore) + .acquire_owned() + .await + .expect("semaphore is never closed during coordinator lifetime"); + + // Determine timestamp and acquire read holds. + let oracle_read_ts = self.oracle_read_ts(&timeline).await?; + let bundle = global_mir_plan.id_bundle(cluster_id); + let (determination, read_holds) = self + .frontend_determine_timestamp( + session, + &bundle, + &QueryWhen::FreshestTableWrite, + cluster_id, + &timeline, + oracle_read_ts, + None, + ) + .await?; + + let as_of = determination.timestamp_context.timestamp_or_default(); + + let global_lir_plan = + self.optimize_lir_read_then_write(optimizer, global_mir_plan, as_of)?; + + // Log optimization finished + if let Some(logging_id) = statement_logging_id { + self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished); + } + + let sink_id = global_lir_plan.sink_id(); + let target_id = plan.id; + let target_global_id = catalog.get_entry(&target_id).latest_global_id(); + let kind = plan.kind.clone(); + let returning = plan.returning.clone(); + + let (df_desc, _df_meta) = global_lir_plan.unapply(); + + let arity = df_desc + .sink_exports + .values() + .next() + .expect("has sink") + .from_desc + .arity(); + + let conn_id = session.conn_id().clone(); + let session_uuid = session.uuid(); + let start_time = (self.statement_logging_frontend.now)(); + let max_result_size = catalog.system_config().max_result_size(); + let max_query_result_size = session.vars().max_query_result_size(); + let row_set_finishing_seconds = session.metrics().row_set_finishing_seconds().clone(); + let max_occ_retries = usize::cast_from(catalog.system_config().max_occ_retries()); + + // Linearize the read BEFORE subscribing or writing: block until + // the oracle for this query's timeline has advanced to `as_of`. + // + // The up-front ordering is load-bearing. If `as_of` is in the far + // future (e.g. reading from a `REFRESH AT` MV with a far-future + // since), submitting a write at `chosen_ts >= as_of` would have + // the group commit bump the oracle to that far-future value, + // stalling every subsequent write on the `EpochMilliseconds` + // timeline until then. So a pathological far-future RTW must park + // here without ever touching the oracle. This park is unbounded on + // its own; the caller bounds it by `statement_timeout`. + self.ensure_read_linearized(&timeline, as_of).await?; + + let subscribe_handle = self + .create_internal_subscribe( + Box::new(df_desc), + cluster_id, + replica_id, + depends_on.clone(), + as_of, + arity, + sink_id, + conn_id.clone(), + session_uuid, + start_time, + read_holds, + ) + .await?; + + let (retry_count, result) = self + .run_occ_loop( + subscribe_handle, + target_id, + target_global_id, + kind, + returning, + max_result_size, + max_query_result_size, + row_set_finishing_seconds, + max_occ_retries, + table_desc, + conn_id, + statement_logging_id, + as_of, + attempt_state, + defer_write, + ) + .await; + + self.coordinator_client() + .metrics() + .occ_retry_count + .observe(f64::from(u32::try_from(retry_count).unwrap_or(u32::MAX))); + + // Release the OCC permit only after the OCC loop has fully completed + // (success, failure, or timeout). Holding it for the entire operation + // is what bounds concurrency; an early drop would let a waiter start + // its subscribe while we are still consolidating diffs and retrying. + drop(permit); + + let outcome = result?; + match outcome { + OccOutcome::Committed { response, write_ts } => { + if let Some(write_ts) = write_ts { + session.apply_write(write_ts); + } + Ok(response) + } + OccOutcome::Deferred { response, diffs } => { + // NOTE: A buffered session write carries no target-generation + // guard. The immediate path pins `target_global_id` and group + // commit re-validates it, but a `WriteOp` only names the + // `CatalogItemId` and commit staging resolves whatever global + // id is current then. So an `ALTER TABLE ... ADD COLUMN` that + // lands between here and COMMIT appends rows of the old arity + // under the new schema. This holds for every buffered write, + // not just ours. + session.add_transaction_ops(TransactionOps::Writes(vec![WriteOp { + id: target_id, + rows: TableData::Rows(diffs), + }]))?; + Ok(response) + } + } + } + + /// Validate a read-then-write operation. + fn validate_read_then_write( + &self, + catalog: &Arc, + session: &Session, + plan: &plan::ReadThenWritePlan, + target_cluster: TargetCluster, + ) -> Result { + // Disallow mz_now in any position because read time and write time differ. + let contains_temporal = plan.selection.contains_temporal() + || plan.assignments.values().any(|e| e.contains_temporal()) + || plan.returning.iter().any(|e| e.contains_temporal()); + if contains_temporal { + return Err(AdapterError::Unsupported( + "calls to mz_now in write statements", + )); + } + + // Validate read dependencies. The plan was built against an earlier + // catalog snapshot; an item it depends on may have been dropped by + // concurrent DDL before we got here. + let dependency_ids = plan + .selection + .depends_on() + .into_iter() + .map(|gid| { + catalog.try_resolve_item_id(&gid).ok_or_else(|| { + AdapterError::Catalog(mz_catalog::memory::error::Error { + kind: ErrorKind::Sql(CatalogError::UnknownItem(gid.to_string())), + }) + }) + }) + .collect::, _>>()?; + let max_rw_dependencies = mz_adapter_types::dyncfgs::READ_THEN_WRITE_MAX_DEPENDENCIES + .get(catalog.system_config().dyncfgs()); + validate_read_then_write_dependencies(catalog, dependency_ids, max_rw_dependencies)?; + + let cluster = catalog.resolve_target_cluster(target_cluster, session)?; + let cluster_id = cluster.id; + + if cluster.replicas().next().is_none() { + return Err(AdapterError::NoClusterReplicasAvailable { + name: cluster.name.clone(), + is_managed: cluster.is_managed(), + }); + } + + let replica_id = session + .vars() + .cluster_replica() + .map(|name| { + cluster + .replica_id(name) + .ok_or(AdapterError::UnknownClusterReplica { + cluster_name: cluster.name.clone(), + replica_name: name.to_string(), + }) + }) + .transpose()?; + + let depends_on = plan.selection.depends_on(); + let timeline = catalog.validate_timeline_context(depends_on.iter().copied())?; + + // Get the table descriptor for constraint validation. The plan's + // target table may have been dropped by concurrent DDL between + // planning and here, so tolerate a missing entry. + let table_desc = match catalog.try_get_entry(&plan.id) { + Some(entry) => entry + .relation_desc_latest() + .expect("table has desc") + .into_owned(), + None => { + return Err(AdapterError::Catalog(mz_catalog::memory::error::Error { + kind: ErrorKind::Sql(CatalogError::UnknownItem(plan.id.to_string())), + })); + } + }; + + Ok(ValidationResult { + cluster_id, + replica_id, + timeline, + depends_on, + table_desc, + }) + } + + /// Optimize MIR for a read-then-write operation. + fn optimize_mir_read_then_write( + &self, + catalog: &Arc, + session: &dyn SessionMetadata, + plan: &plan::ReadThenWritePlan, + cluster_id: ComputeInstanceId, + ) -> Result< + ( + optimize::subscribe::Optimizer, + optimize::subscribe::GlobalMirPlan, + ), + AdapterError, + > { + // `finishing` is unused: the OCC path emits raw diffs and + // `apply_mutation_to_mir` handles update projection. + let plan::ReadThenWritePlan { + id: _, + selection, + finishing: _, + assignments, + kind, + returning: _, + } = plan; + + let expr = selection.clone().lower(catalog.system_config(), None)?; + let mut expr = apply_mutation_to_mir(expr, kind, assignments); + + // Resolve unmaterializable functions (now(), current_user, ...) before + // the subscribe optimizer sees them: it uses `ExprPrepMaintained`, + // which rejects them, but our subscribe is a one-shot read so we can + // resolve them to constants. `mz_now()` is rejected upstream by + // `validate_read_then_write`. + let style = ExprPrepOneShot { + logical_time: EvalTime::NotAvailable, + session, + catalog_state: catalog.state(), + }; + expr.try_visit_scalars_mut(&mut |s| style.prep_scalar_expr(s))?; + + let compute_instance = ComputeInstanceSnapshot::new_without_collections(cluster_id); + let (_, view_id) = self.transient_id_gen.allocate_id(); + let (_, sink_id) = self.transient_id_gen.allocate_id(); + let debug_name = format!("frontend-read-then-write-subscribe-{}", sink_id); + let optimizer_config = optimize::OptimizerConfig::from(catalog.system_config()) + .override_from(&catalog.get_cluster(cluster_id).config.features()) + .override_from( + &catalog + .state() + .cluster_scoped_optimizer_overrides(cluster_id), + ); + + let mut optimizer = optimize::subscribe::Optimizer::new( + Arc::::clone(catalog), + compute_instance, + view_id, + sink_id, + true, // with_snapshot + None, // up_to + debug_name, + optimizer_config, + self.optimizer_metrics.clone(), + ); + + let expr_typ = expr.typ(); + let sql_typ = mz_repr::SqlRelationType::from_repr(&expr_typ); + let column_names: Vec = (0..sql_typ.column_types.len()) + .map(|i| format!("column{}", i)) + .collect(); + let relation_desc = RelationDesc::new(sql_typ, column_names.iter().map(|s| s.as_str())); + + // MIR => MIR optimization (global). We use optimize_mir because the + // mutation has already been applied in MIR, so we bypass the normal + // SubscribePlan path which expects HIR. + let global_mir_plan = optimizer.optimize_mir(expr, relation_desc)?; + + Ok((optimizer, global_mir_plan)) + } + + /// Optimize LIR for a read-then-write operation. + fn optimize_lir_read_then_write( + &self, + mut optimizer: optimize::subscribe::Optimizer, + global_mir_plan: optimize::subscribe::GlobalMirPlan, + as_of: Timestamp, + ) -> Result { + let global_mir_plan = global_mir_plan.resolve(Antichain::from_elem(as_of)); + let global_lir_plan = optimizer.optimize(global_mir_plan)?; + Ok(global_lir_plan) + } + + /// Get the oracle read timestamp hint for the timeline of this query. + async fn oracle_read_ts( + &mut self, + timeline: &TimelineContext, + ) -> Result, AdapterError> { + // See `ensure_read_linearized` for why `get_timeline` is the right + // function here: the write target lives on `EpochMilliseconds`, so + // we want that oracle's read_ts as the hint for timestamp + // selection even when the read side is MV-only + // (`TimestampDependent`). + let timeline = ::get_timeline(timeline); + + match timeline { + Some(timeline) => { + let oracle = self.ensure_oracle(timeline).await?; + Ok(Some(oracle.read_ts().await)) + } + None => Ok(None), + } + } + + /// Block until the oracle for this query's timeline has advanced to + /// `as_of`. Returns immediately if it already has. + /// + /// This implements the strict-serializable read guarantee for RTW: + /// once this returns, any session observing the oracle sees a read + /// timestamp at least as large as `as_of`, so reads at `as_of` (and + /// writes derived from them) cannot appear to "go backwards" relative + /// to subsequent queries. + async fn ensure_read_linearized( + &mut self, + timeline: &TimelineContext, + as_of: Timestamp, + ) -> Result<(), AdapterError> { + // Pick the oracle this RTW operates against. `timeline` is derived + // from the read side (`plan.selection.depends_on()`), so an + // MV-only read produces `TimestampDependent` — an MV itself + // doesn't pin the query to any source timeline. The write target, + // however, is always a Table living on `EpochMilliseconds`, and + // future readers of that table will consult the + // `EpochMilliseconds` oracle, so linearization must target + // `EpochMilliseconds` regardless of the read side. + // + // `get_timeline` encodes that defaulting (`TimestampDependent` → + // `Some(EpochMilliseconds)`). `TimelineContext::timeline()` + // answers a different question ("is there a source-forced + // timeline?") and would return `None` for MV-only reads, + // silently skipping linearization. + let tl = match ::get_timeline(timeline) { + Some(tl) => tl, + None => return Ok(()), + }; + + let oracle = self.ensure_oracle(tl).await?; + + loop { + let oracle_ts = oracle.read_ts().await; + if as_of <= oracle_ts { + return Ok(()); + } + + // Sleep for roughly the difference between as_of and the current + // oracle timestamp. Since timestamps are epoch milliseconds, the + // difference is the approximate wall-clock time we need to wait. + // Cap at 1s to avoid very long sleeps if clocks are skewed, + // matching the cap in `message_linearize_reads`. + let wait_ms = u64::from(as_of.saturating_sub(oracle_ts)); + let wait = Duration::from_millis(wait_ms).min(Duration::from_secs(1)); + tokio::time::sleep(wait).await; + } + } + + /// Creates an internal subscribe that does not appear in introspection + /// tables. Returns a [`SubscribeHandle`] that ensures cleanup on drop. + async fn create_internal_subscribe( + &self, + df_desc: Box, + cluster_id: ComputeInstanceId, + replica_id: Option, + depends_on: BTreeSet, + as_of: Timestamp, + arity: usize, + sink_id: GlobalId, + conn_id: mz_adapter_types::connection::ConnectionId, + session_uuid: Uuid, + start_time: mz_ore::now::EpochMillis, + read_holds: crate::ReadHolds, + ) -> Result { + let rx: mpsc::UnboundedReceiver = self + .call_coordinator(|tx| Command::CreateInternalSubscribe { + df_desc, + cluster_id, + replica_id, + depends_on, + as_of, + arity, + sink_id, + conn_id, + session_uuid, + start_time, + read_holds, + tx, + }) + .await?; + + Ok(SubscribeHandle { + rx, + sink_id, + client: Some(self.coordinator_client().clone()), + }) + } + + /// Run the OCC loop: drain the subscribe at `as_of`, apply the + /// mutation, and submit the resulting diffs as a write. + /// + /// Semantically this is a SELECT at `as_of` followed by an INSERT. + /// Because we hold no write lock, a concurrent writer may bump the + /// target table's upper past our chosen write timestamp, in which + /// case the coordinator returns `WriteResult::TimestampPassed`; we + /// then wait for the subscribe to advance and retry, up to + /// `max_occ_retries` times. + /// + /// Read linearization is the caller's responsibility: `as_of` must + /// already be linearized (oracle read_ts >= `as_of`) on entry. See + /// `ensure_read_linearized` at the call site. + /// + /// Returns `(retry_count, result)` so the caller can record OCC retry + /// metrics regardless of whether the operation succeeded or failed. + async fn run_occ_loop( + &self, + mut subscribe_handle: SubscribeHandle, + target_id: CatalogItemId, + target_global_id: GlobalId, + kind: MutationKind, + returning: Vec, + max_result_size: u64, + max_query_result_size: u64, + row_set_finishing_seconds: Histogram, + max_occ_retries: usize, + table_desc: RelationDesc, + conn_id: mz_adapter_types::connection::ConnectionId, + statement_logging_id: Option, + as_of: Timestamp, + attempt_state: Arc, + defer_write: bool, + ) -> (usize, Result) { + let mut state = OccState::new(); + + // Correctness invariant for retries: + // + // `all_diffs` accumulates *all* rows ever received from the subscribe, + // across retries. The subscribe emits a snapshot (at the as_of + // timestamp) followed by incremental updates. We consolidate on every + // progress message (flattening timestamps to MIN first), so after + // consolidation `all_diffs` always represents "what the query returns + // as of the latest progress timestamp" — old snapshot rows that were + // retracted by newer updates cancel out, and new rows appear. This is + // exactly the set of diffs we want to write. + // + // Consolidating on every progress also means the NoRowsMatched check + // works correctly across retries: if the consolidated result becomes + // logically empty (all diffs cancel out), `all_diffs` will be empty + // and we early-return without attempting a write. + let result = loop { + if let Some(error) = attempt_state.requested_error() { + break Err(error); + } + let msg = match subscribe_handle.recv().await { + Some(msg) => msg, + None => { + // Channel closed cleanly: the SELECT is constant (no + // table dependency). Submit the accumulated diffs as a + // blind write — the oracle picks the timestamp at group + // commit, so we just flatten to `Timestamp::MIN` for + // `consolidate_updates`. + state.consolidate(Timestamp::MIN); + if state.all_diffs.is_empty() { + break build_no_rows_response(&kind, &returning).map(|response| { + OccOutcome::Committed { + response, + write_ts: None, + } + }); + } + let success_response = match self.build_success_response( + &kind, + &returning, + &state.all_diffs, + max_result_size, + max_query_result_size, + &row_set_finishing_seconds, + ) { + Ok(response) => response, + Err(e) => break Err(e), + }; + let diffs = state + .all_diffs + .iter() + .map(|(row, _ts, diff)| (row.clone(), *diff)) + .collect_vec(); + + // Inside a transaction the rows must not become visible + // before COMMIT, and they have to disappear on ROLLBACK, + // so we hand them back to the caller to buffer as session + // write ops. The transaction's commit submits them to + // group commit like any other buffered write. + if defer_write { + break Ok(OccOutcome::Deferred { + response: success_response, + diffs, + }); + } + + attempt_state.mark_write_submitted(); + let result = self + .call_coordinator(|tx| Command::AttemptWrite { + conn_id: conn_id.clone(), + target_id, + target_global_id, + diffs, + write_ts: None, + tx, + }) + .await; + // All arms below terminate the attempt, so `write_submitted` + // stays set per its contract. + match result { + WriteResult::Success { timestamp } => { + if let Some(id) = statement_logging_id { + self.log_set_timestamp(id, timestamp); + } + break Ok(OccOutcome::Committed { + response: success_response, + write_ts: Some(timestamp), + }); + } + WriteResult::TimestampPassed { .. } => { + // Unreachable: blind writes use + // `UserWriteResponder::Internal`, which group + // commit never resolves to `TimestampPassed`. + soft_panic_or_log!( + "blind read-then-write unexpectedly got TimestampPassed" + ); + break Err(AdapterError::Internal( + "blind write unexpectedly got TimestampPassed".into(), + )); + } + WriteResult::Canceled => { + break Err(attempt_state + .requested_error() + .unwrap_or(AdapterError::Canceled)); + } + WriteResult::ReadOnly => break Err(AdapterError::ReadOnly), + WriteResult::TargetChanged => { + break Err(AdapterError::Unstructured(anyhow::anyhow!( + "target table changed while read-then-write was executing", + ))); + } + WriteResult::Indeterminate => break Err(AdapterError::Internal( + "write outcome is indeterminate because the group committer shut down" + .into(), + )), + } + } + }; + + match process_message(msg, &mut state, as_of, max_result_size, &table_desc) { + ProcessResult::Continue { ready_to_write } => { + if !ready_to_write { + continue; + } + + // Drain pending messages before attempting write + let drain_err = loop { + match subscribe_handle.try_recv() { + Ok(msg) => { + match process_message( + msg, + &mut state, + as_of, + max_result_size, + &table_desc, + ) { + ProcessResult::Continue { .. } => {} + ProcessResult::NoRowsMatched => { + break Some(build_no_rows_response(&kind, &returning).map( + |response| OccOutcome::Committed { + response, + write_ts: None, + }, + )); + } + ProcessResult::Error(e) => { + break Some(Err(e)); + } + } + } + Err(mpsc::error::TryRecvError::Empty) => break None, + // The subscribe can finish (coordinator drops the + // sender after `process_response` returns true) + // between our last recv() and this drain. This is + // benign — all buffered messages have already been + // consumed via the Ok(msg) arm above. + Err(mpsc::error::TryRecvError::Disconnected) => break None, + } + }; + if let Some(result) = drain_err { + break result; + } + + // A write that reads persisted state cannot be deferred: + // its diffs are only valid at the frontier we observed. + // Inside a transaction such statements are rejected before + // we execute anything, by the read-dependency check in + // `SessionClient::try_frontend_read_then_write`, so this is + // defense in depth. + if defer_write { + soft_panic_or_log!( + "read-dependent read-then-write reached the OCC write path \ + inside a transaction" + ); + break Err(AdapterError::Internal( + "read-then-write cannot be run inside a transaction block".into(), + )); + } + + let write_ts = state + .current_upper + .expect("must have seen progress to be ready to write"); + + // Consolidate any rows received during the drain + // (the bulk was already consolidated on the last progress). + state.consolidate(write_ts); + + let success_response = match self.build_success_response( + &kind, + &returning, + &state.all_diffs, + max_result_size, + max_query_result_size, + &row_set_finishing_seconds, + ) { + Ok(response) => response, + Err(e) => break Err(e), + }; + + // Submit write. + // + // perf: clones every row on each attempt. Under contention + // we retry up to `max_occ_retries` times (default 1000), + // so a large DELETE/UPDATE under heavy contention can do a + // lot of row-cloning work. If this shows up in profiles, + // consider storing `Arc` in `all_diffs` to make the + // per-attempt copy cheap. + attempt_state.mark_write_submitted(); + let result = self + .call_coordinator(|tx| Command::AttemptWrite { + conn_id: conn_id.clone(), + target_id, + target_global_id, + diffs: state + .all_diffs + .iter() + .map(|(row, _ts, diff)| (row.clone(), *diff)) + .collect_vec(), + write_ts: Some(write_ts), + tx, + }) + .await; + + match result { + WriteResult::Success { timestamp } => { + if let Some(id) = statement_logging_id { + self.log_set_timestamp(id, timestamp); + } + // N.B. subscribe_handle is dropped here, which + // fires off the cleanup message. + break Ok(OccOutcome::Committed { + response: success_response, + write_ts: Some(timestamp), + }); + } + WriteResult::TimestampPassed { + next_eligible_timestamp, + .. + } => { + // The write definitively did not land, so the + // attempt is resolved. Clearing `write_submitted` + // lets a cancel or statement timeout that fires + // during the upcoming subscribe wait resolve + // promptly instead of awaiting a write result. + attempt_state.mark_write_resolved(); + // Do not advance `state.current_upper` (and + // therefore `write_ts`) from `next_eligible_timestamp`. + // The diffs in `all_diffs` are only known to be + // correct as of subscribe progress we have actually + // observed. Retrying at a newer oracle timestamp + // before subscribe progress catches up would risk + // applying stale diffs at the wrong timestamp. So + // on `TimestampPassed` we wait for the subscribe to + // progress and retry using that observed frontier. + state.retry_count += 1; + if let Some(error) = attempt_state.requested_error() { + break Err(error); + } + if state.retry_count > max_occ_retries { + // High contention is a user-visible + // condition, not an internal invariant + // violation. Surface it as + // `Unstructured` so it doesn't trip + // internal-error alerts. + break Err(AdapterError::Unstructured(anyhow::anyhow!( + "read-then-write exceeded maximum retry attempts under contention", + ))); + } + tracing::debug!( + retry_count = state.retry_count, + write_ts = %write_ts, + next_eligible_timestamp = %next_eligible_timestamp, + "OCC write conflict, retrying" + ); + continue; + } + WriteResult::Canceled => { + break Err(attempt_state + .requested_error() + .unwrap_or(AdapterError::Canceled)); + } + WriteResult::ReadOnly => break Err(AdapterError::ReadOnly), + WriteResult::TargetChanged => { + break Err(AdapterError::Unstructured(anyhow::anyhow!( + "target table changed while read-then-write was executing", + ))); + } + WriteResult::Indeterminate => break Err(AdapterError::Internal( + "write outcome is indeterminate because the group committer shut down" + .into(), + )), + } + } + ProcessResult::NoRowsMatched => { + break build_no_rows_response(&kind, &returning).map(|response| { + OccOutcome::Committed { + response, + write_ts: None, + } + }); + } + ProcessResult::Error(e) => { + break Err(e); + } + } + }; + + (state.retry_count, result) + } + + /// Build the success response after a successful write. + fn build_success_response( + &self, + kind: &MutationKind, + returning: &[MirScalarExpr], + all_diffs: &[(Row, Timestamp, Diff)], + max_result_size: u64, + max_query_result_size: u64, + row_set_finishing_seconds: &Histogram, + ) -> Result { + if returning.is_empty() { + // For UPDATE each changed row produces a retraction (-1) and an + // insertion (+1), so we divide by 2 below. + let row_count = all_diffs + .iter() + .map(|(_, _, diff)| diff.into_inner().unsigned_abs()) + .sum::(); + let row_count = + usize::try_from(row_count).expect("positive row count must fit in usize"); + + return Ok(match kind { + MutationKind::Delete => ExecuteResponse::Deleted(row_count), + MutationKind::Update => ExecuteResponse::Updated(row_count / 2), + MutationKind::Insert => ExecuteResponse::Inserted(row_count), + }); + } + + let mut returning_rows = Vec::new(); + let arena = RowArena::new(); + // RETURNING expressions are evaluated row-by-row in this loop, so an + // expression like `RETURNING repeat('x', 10_000_000)` will allocate + // unbounded data unless we bail mid-loop. The post-loop + // `RowSetFinishing::finish` below would also reject this, but only + // after we've materialized everything. The early-bail caps the + // temporary allocation. We pick the lower of the two configured + // caps — whichever fires first wins. + let mut projected_byte_size: u64 = 0; + let early_cap = std::cmp::min(max_result_size, max_query_result_size); + + for (row, _ts, diff) in all_diffs { + let include = match kind { + MutationKind::Delete => diff.is_negative(), + MutationKind::Update | MutationKind::Insert => diff.is_positive(), + }; + + if !include { + continue; + } + + let mut returning_row = Row::with_capacity(returning.len()); + let mut packer = returning_row.packer(); + let datums: Vec<_> = row.iter().collect(); + + for expr in returning { + match expr.eval(&datums, &arena) { + Ok(datum) => packer.push(datum), + Err(err) => return Err(err.into()), + } + } + + let multiplicity = NonZeroUsize::try_from( + NonZeroI64::try_from(diff.into_inner().abs()).expect("diff is non-zero"), + ) + .map_err(AdapterError::from)?; + + let row_bytes = u64::cast_from(returning_row.byte_len()) + .saturating_mul(u64::cast_from(multiplicity.get())); + projected_byte_size = projected_byte_size.saturating_add(row_bytes); + if projected_byte_size > early_cap { + return Err(AdapterError::ResultSize(format!( + "result exceeds max size of {}", + ByteSize::b(early_cap) + ))); + } + + returning_rows.push((returning_row, multiplicity)); + } + + // Run the canonical finish to enforce both caps with full precision + // (including the sorted-view memory overhead) and to register the + // row-set-finishing duration histogram, mirroring the legacy + // `send_diffs` path. + let finishing = RowSetFinishing { + order_by: Vec::new(), + limit: None, + offset: 0, + project: (0..returning.len()).collect(), + }; + match finishing.finish( + RowCollection::new(returning_rows, &finishing.order_by), + max_result_size, + Some(max_query_result_size), + row_set_finishing_seconds, + ) { + Ok((rows, _size_bytes)) => Ok(ExecuteResponse::SendingRowsImmediate { + rows: Box::new(rows), + }), + Err(e) => Err(AdapterError::ResultSize(e)), + } + } +} + +/// Result of validating a read-then-write operation. +struct ValidationResult { + cluster_id: ComputeInstanceId, + replica_id: Option, + timeline: TimelineContext, + depends_on: BTreeSet, + /// The table descriptor, used for constraint validation. + table_desc: RelationDesc, +} + +/// Accumulated state for the OCC loop in `run_occ_loop`. +struct OccState { + all_diffs: Vec<(Row, Timestamp, Diff)>, + current_upper: Option, + retry_count: usize, + byte_size: u64, +} + +impl OccState { + fn new() -> Self { + Self { + all_diffs: Vec::new(), + current_upper: None, + retry_count: 0, + byte_size: 0, + } + } + + /// Forward all diff timestamps to `target_ts` and consolidate. + /// + /// After consolidation, `all_diffs` represents the net state of the + /// query as of `target_ts`. Rows that were retracted by newer updates + /// cancel out, and `byte_size` is recomputed to reflect the + /// consolidated data. + fn consolidate(&mut self, target_ts: Timestamp) { + for (_, ts, _) in self.all_diffs.iter_mut() { + *ts = target_ts; + } + consolidation::consolidate_updates(&mut self.all_diffs); + self.byte_size = self + .all_diffs + .iter() + .map(|(row, _, _)| u64::cast_from(row.byte_len())) + .sum(); + } +} + +/// Result of processing a single subscribe message in the OCC loop. +enum ProcessResult { + Continue { ready_to_write: bool }, + NoRowsMatched, + Error(AdapterError), +} + +/// Process one subscribe message, updating `state` in place. +/// +/// Data rows are accumulated into `state.all_diffs` (with per-row constraint +/// and max-result-size checks). Progress messages trigger consolidation and +/// can promote the accumulated diffs to "ready to write". +fn process_message( + response: PeekResponseUnary, + state: &mut OccState, + as_of: Timestamp, + max_result_size: u64, + table_desc: &RelationDesc, +) -> ProcessResult { + match response { + PeekResponseUnary::Rows(mut rows) => { + let mut saw_progress = false; + + while let Some(row) = rows.next() { + let mut datums = row.iter(); + + // Extract mz_timestamp (SubscribeOutput::Diffs format: + // mz_timestamp, mz_progressed, mz_diff, ...data columns...). + // + // Format drift would mean we'd silently commit an incorrect + // write, so surface every shape mismatch as an internal + // error rather than panicking the process. + let Some(ts_datum) = datums.next() else { + return ProcessResult::Error(AdapterError::Internal( + "missing mz_timestamp in subscribe output".into(), + )); + }; + let ts = match ts_datum { + mz_repr::Datum::Numeric(n) => match n.0.try_into() { + Ok(ts_u64) => Timestamp::new(ts_u64), + Err(_) => { + return ProcessResult::Error(AdapterError::Internal(format!( + "mz_timestamp in subscribe output is not a valid u64: {n}" + ))); + } + }, + other => { + return ProcessResult::Error(AdapterError::Internal(format!( + "unexpected mz_timestamp datum: {other:?}" + ))); + } + }; + + // Extract mz_progressed + let Some(progressed_datum) = datums.next() else { + return ProcessResult::Error(AdapterError::Internal( + "missing mz_progressed in subscribe output".into(), + )); + }; + let is_progress = matches!(progressed_datum, mz_repr::Datum::True); + + if is_progress { + state.current_upper = Some(ts); + saw_progress = true; + + // Consolidate incrementally on each progress + // message. This keeps memory bounded by the + // consolidated size and makes the byte_size check + // below accurate (except for rows received between + // two progress messages, which is a small window). + state.consolidate(ts); + + // The very first progress message we receive is + // always at `as_of`, emitted synchronously by + // `ActiveSubscribe::initialize` *before* any data + // batch is processed. At that point `all_diffs` is + // empty by construction, regardless of whether the + // snapshot is actually empty, so we must not + // conclude `NoRowsMatched` from it. Progress + // messages emitted later from `process_response` + // are gated on `batch.upper > as_of`, so any + // progress with `ts > as_of` is past the initial + // one and an empty `all_diffs` then genuinely + // means no rows matched. See + // `src/adapter/src/active_compute_sink.rs` for + // the emission order. + if ts > as_of && state.all_diffs.is_empty() { + return ProcessResult::NoRowsMatched; + } + } else { + let Some(diff_datum) = datums.next() else { + return ProcessResult::Error(AdapterError::Internal( + "missing mz_diff in subscribe output".into(), + )); + }; + let diff = match diff_datum { + mz_repr::Datum::Int64(d) => Diff::from(d), + other => { + return ProcessResult::Error(AdapterError::Internal(format!( + "unexpected mz_diff datum while processing read-then-write: {other:?}" + ))); + } + }; + + let data_row = Row::pack(datums); + + // Validate constraints for rows being added (positive diff) + if diff.is_positive() { + for (idx, datum) in data_row.iter().enumerate() { + if let Err(e) = table_desc.constraints_met(idx, &datum) { + return ProcessResult::Error(e.into()); + } + } + } + + state.byte_size = state + .byte_size + .saturating_add(u64::cast_from(data_row.byte_len())); + if state.byte_size > max_result_size { + return ProcessResult::Error(AdapterError::ResultSize(format!( + "result exceeds max size of {}", + max_result_size + ))); + } + state.all_diffs.push((data_row, ts, diff)); + } + } + + // We're ready to write once we've seen a progress + // message and have accumulated any diffs. Data rows can + // only arrive *after* the initial progress at `as_of` + // (see the note in the progress branch), so a non-empty + // `all_diffs` here implies we're past the initial + // progress. + let ready_to_write = saw_progress && !state.all_diffs.is_empty(); + ProcessResult::Continue { ready_to_write } + } + PeekResponseUnary::Error(e) => { + ProcessResult::Error(AdapterError::Unstructured(anyhow::anyhow!(e))) + } + PeekResponseUnary::DependencyDropped(dep) => ProcessResult::Error( + AdapterError::Unstructured(anyhow::anyhow!(dep.query_terminated_error())), + ), + PeekResponseUnary::Canceled => ProcessResult::Error(AdapterError::Canceled), + } +} + +/// Build the response returned when no rows matched the selection. +fn build_no_rows_response( + kind: &MutationKind, + returning: &[MirScalarExpr], +) -> Result { + if !returning.is_empty() { + let rows: Vec = vec![]; + return Ok(ExecuteResponse::SendingRowsImmediate { + rows: Box::new(rows.into_row_iter()), + }); + } + Ok(match kind { + MutationKind::Delete => ExecuteResponse::Deleted(0), + MutationKind::Update => ExecuteResponse::Updated(0), + MutationKind::Insert => ExecuteResponse::Inserted(0), + }) +} + +/// Transform a MIR expression to produce the appropriate diffs for a mutation. +/// +/// - DELETE: Negates the expression to produce `(row, -1)` diffs +/// - UPDATE: Unions negated old rows with mapped new rows to produce both +/// `(old_row, -1)` and `(new_row, +1)` diffs +fn apply_mutation_to_mir( + expr: MirRelationExpr, + kind: &MutationKind, + assignments: &BTreeMap, +) -> MirRelationExpr { + match kind { + MutationKind::Delete => MirRelationExpr::Negate { + input: Box::new(expr), + }, + MutationKind::Update => { + let arity = expr.arity(); + + // Find a fresh LocalId that won't conflict with any in the expression. + // + // Invariant: `Let` and `LetRec` are the only MIR nodes that *bind* + // LocalIds; `Get` references them but does not introduce new ones. + // So scanning just those two node kinds and picking `max + 1` is + // guaranteed to produce an id unused by the subtree. + let mut max_id = 0_u64; + expr.visit_pre(|e| match e { + MirRelationExpr::Let { id, .. } => { + max_id = std::cmp::max(max_id, id.into()); + } + MirRelationExpr::LetRec { ids, .. } => { + for id in ids { + max_id = std::cmp::max(max_id, id.into()); + } + } + _ => {} + }); + let binding_id = LocalId::new(max_id + 1); + + let get_binding = MirRelationExpr::Get { + id: Id::Local(binding_id), + typ: expr.typ(), + access_strategy: mz_expr::AccessStrategy::UnknownOrLocal, + }; + + // Build map expressions + let map_scalars: Vec = (0..arity) + .map(|i| { + assignments + .get(&i) + .cloned() + .unwrap_or_else(|| MirScalarExpr::column(i)) + }) + .collect(); + + let new_rows = get_binding + .clone() + .map(map_scalars) + .project((arity..2 * arity).collect()); + + let old_rows = MirRelationExpr::Negate { + input: Box::new(get_binding), + }; + + let body = new_rows.union(old_rows); + + MirRelationExpr::Let { + id: binding_id, + value: Box::new(expr), + body: Box::new(body), + } + } + // INSERT: rows pass through unchanged; the subscribe emits them with diff +1. + MutationKind::Insert => expr, + } +} diff --git a/src/adapter/src/lib.rs b/src/adapter/src/lib.rs index 6c31e5058fc70..38daee9c7ddfe 100644 --- a/src/adapter/src/lib.rs +++ b/src/adapter/src/lib.rs @@ -44,6 +44,7 @@ mod coord; mod error; mod explain; mod frontend_peek; +mod frontend_read_then_write; mod notice; mod optimize; mod util; diff --git a/src/adapter/src/metrics.rs b/src/adapter/src/metrics.rs index 446f73e1f8905..8595be3d3f351 100644 --- a/src/adapter/src/metrics.rs +++ b/src/adapter/src/metrics.rs @@ -59,6 +59,7 @@ pub struct Metrics { pub catalog_transact_seconds: HistogramVec, pub apply_catalog_implications_seconds: Histogram, pub group_commit_catalog_upper_seconds: Histogram, + pub occ_retry_count: Histogram, } impl Metrics { @@ -289,6 +290,13 @@ impl Metrics { help: "The time it takes to advance the catalog shard upper for a txns-shard write (group commits and table register/forget).", buckets: histogram_seconds_buckets(0.001, 32.0), )), + occ_retry_count: registry.register(metric!( + name: "mz_occ_read_then_write_retry_count", + help: "Number of OCC retries per read-then-write operation.", + buckets: vec![ + 0., 1., 2., 3., 5., 10., 25., 50., 100., 200., 300., 500., 750., 1000., + ], + )) } } diff --git a/src/adapter/src/notice.rs b/src/adapter/src/notice.rs index e6bed59c3f599..f801dd3411e6a 100644 --- a/src/adapter/src/notice.rs +++ b/src/adapter/src/notice.rs @@ -132,6 +132,12 @@ pub enum AdapterNotice { role: Option, var_name: Option, }, + /// An `ALTER SYSTEM` statement named a system parameter that + /// `environmentd` only reads at startup, so the running process keeps the + /// value it sampled at boot. + StartupOnlyVarUpdated { + var_name: String, + }, Welcome(String), PlanInsights(String), IntrospectionClusterUsage, @@ -211,6 +217,7 @@ impl AdapterNotice { AdapterNotice::DroppedInUseIndex { .. } => Severity::Notice, AdapterNotice::PerReplicaLogRead { .. } => Severity::Notice, AdapterNotice::VarDefaultUpdated { .. } => Severity::Notice, + AdapterNotice::StartupOnlyVarUpdated { .. } => Severity::Warning, AdapterNotice::Welcome(_) => Severity::Notice, AdapterNotice::PlanInsights(_) => Severity::Notice, AdapterNotice::IntrospectionClusterUsage => Severity::Warning, @@ -323,6 +330,7 @@ impl AdapterNotice { AdapterNotice::WebhookSourceCreated { .. } => SqlState::SUCCESSFUL_COMPLETION, AdapterNotice::PerReplicaLogRead { .. } => SqlState::SUCCESSFUL_COMPLETION, AdapterNotice::VarDefaultUpdated { .. } => SqlState::SUCCESSFUL_COMPLETION, + AdapterNotice::StartupOnlyVarUpdated { .. } => SqlState::WARNING, AdapterNotice::Welcome(_) => SqlState::SUCCESSFUL_COMPLETION, AdapterNotice::PlanInsights(_) => SqlState::from_code("MZ001"), AdapterNotice::IntrospectionClusterUsage => SqlState::WARNING, @@ -508,6 +516,11 @@ impl fmt::Display for AdapterNotice { "{vars} updated for {target}, this will have no effect on the current session" ) } + AdapterNotice::StartupOnlyVarUpdated { var_name } => write!( + f, + "changes to {} only take effect when environmentd restarts", + var_name.quoted() + ), AdapterNotice::Welcome(message) => message.fmt(f), AdapterNotice::PlanInsights(message) => message.fmt(f), AdapterNotice::IntrospectionClusterUsage => write!( 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..33b77cbebaaf1 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::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,86 @@ impl Optimizer { pub fn sink_id(&self) -> GlobalId { self.sink_id } + + /// Optimize a subscribe dataflow starting from a pre-lowered MIR + /// expression. Used by the frontend read-then-write path which applies + /// mutation transformations in MIR before optimization. + pub fn optimize_mir( + &mut self, + expr: MirRelationExpr, + desc: RelationDesc, + ) -> 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(); + + // 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, + )?; + df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; + + let sink_description = ComputeSinkDesc { + from: self.view_id, + from_desc: desc, + connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { + // Read-then-write subscribes use raw diffs. + output: vec![], + }), + with_snapshot: self.with_snapshot, + up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), + non_null_assertions: vec![], + 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, + &self.config.features, + &self.typecheck_ctx, + &mut df_meta, + Some(&mut self.metrics), + ); + mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?; + + if self.config.mode == OptimizeMode::Explain { + trace_plan!(at: "global", &df_meta.used_indexes(&df_desc)); + } + + self.duration += time.elapsed(); + + Ok(GlobalMirPlan { + df_desc, + df_meta, + phantom: PhantomData::, + }) + } } /// The (sealed intermediate) result after: diff --git a/src/adapter/src/optimize/view.rs b/src/adapter/src/optimize/view.rs index 96476cfc44b83..9c86ca70243fe 100644 --- a/src/adapter/src/optimize/view.rs +++ b/src/adapter/src/optimize/view.rs @@ -66,6 +66,23 @@ impl Optimizer { } impl Optimizer { + /// Creates an optimizer instance that takes an [`ExprPrep`] to handle + /// unmaterializable functions while preserving the usual constant-folding + /// size limit. + pub fn new_with_prep( + config: OptimizerConfig, + metrics: Option, + expr_prep_style: S, + ) -> Optimizer { + Self { + typecheck_ctx: empty_typechecking_context(), + config, + metrics, + expr_prep_style, + fold_constants_limit: true, + } + } + /// Creates an optimizer instance that takes an [`ExprPrep`] to handle /// unmaterializable functions. Additionally, this instance calls constant /// folding without a size limit. diff --git a/src/adapter/src/peek_client.rs b/src/adapter/src/peek_client.rs index d489ca9a4eb09..74e4011baa11e 100644 --- a/src/adapter/src/peek_client.rs +++ b/src/adapter/src/peek_client.rs @@ -31,11 +31,11 @@ use prometheus::Histogram; use qcell::QCell; use thiserror::Error; use timely::progress::Antichain; -use tokio::sync::oneshot; +use tokio::sync::{Semaphore, 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::session::{LifecycleTimestamps, Session}; @@ -75,6 +75,12 @@ pub struct PeekClient { persist_client: PersistClient, /// Statement logging state for frontend peek sequencing. pub statement_logging_frontend: StatementLoggingFrontend, + /// Semaphore for limiting concurrent OCC (optimistic concurrency control) write operations. + pub occ_write_semaphore: Arc, + /// Whether frontend OCC read-then-write is enabled (determined once at process startup). + pub frontend_read_then_write_enabled: bool, + /// Whether the coordinator is in read-only mode. Mutations must be rejected. + pub read_only: bool, } impl PeekClient { @@ -90,6 +96,9 @@ impl PeekClient { optimizer_metrics: OptimizerMetrics, persist_client: PersistClient, statement_logging_frontend: StatementLoggingFrontend, + occ_write_semaphore: Arc, + frontend_read_then_write_enabled: bool, + read_only: bool, ) -> Self { Self { coordinator_client, @@ -101,6 +110,9 @@ impl PeekClient { statement_logging_frontend, oracles: Default::default(), // lazily populated persist_client, + occ_write_semaphore, + frontend_read_then_write_enabled, + read_only, } } @@ -200,6 +212,12 @@ impl PeekClient { .expect("if the coordinator is still alive, it shouldn't have dropped our call") } + /// Returns a clone of the coordinator client, for use in cleanup guards + /// that need to send fire-and-forget commands. + pub(crate) fn coordinator_client(&self) -> &crate::Client { + &self.coordinator_client + } + /// Acquire read holds on the given compute/storage collections, and /// determine the smallest common valid write frontier among the specified collections. /// @@ -500,12 +518,11 @@ impl PeekClient { /// 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. + /// Returns a [`StatementLoggingGuard`]. Callers must retire the guard + /// (e.g. via [`StatementLoggingGuard::retire_with_result`]) on the + /// execution's terminal outcome, or [`defuse`](StatementLoggingGuard::defuse) + /// it when handing off logging responsibility. Dropping the guard without + /// retiring it emits an `Aborted` end-execution event. pub(crate) fn begin_statement_logging( &self, session: &mut Session, @@ -676,6 +693,17 @@ impl StatementLoggingGuard { self.emit(reason); } + /// Retires the guard using the standard mapping from an execution result. + pub(crate) fn retire_with_result(self, result: &Result) { + let reason = match result { + Ok(resp) => resp.into(), + Err(e) => statement_logging::StatementEndedExecutionReason::Errored { + error: e.to_string(), + }, + }; + self.retire(reason); + } + /// 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. diff --git a/src/adapter/src/session.rs b/src/adapter/src/session.rs index 9482dac94296c..b8c9e91711102 100644 --- a/src/adapter/src/session.rs +++ b/src/adapter/src/session.rs @@ -1841,6 +1841,14 @@ impl GroupCommitWriteLocks { self.locks.extend(std::mem::take(&mut other.locks)); } + /// Insert a single lock into this collection. + /// + /// Useful when a lock is acquired directly during group commit rather + /// than handed off from a session via [`Self::merge`]. + pub fn insert_lock(&mut self, id: CatalogItemId, lock: tokio::sync::OwnedMutexGuard<()>) { + self.locks.insert(id, lock); + } + /// Returns the collections we're missing locks for, if any. pub fn missing_locks( &self, diff --git a/src/environmentd/tests/server.rs b/src/environmentd/tests/server.rs index 5fc3b83acf6d6..f81e7de40ba2e 100644 --- a/src/environmentd/tests/server.rs +++ b/src/environmentd/tests/server.rs @@ -672,6 +672,122 @@ ORDER BY mseh.began_at", ); } +#[mz_ore::test] +fn test_statement_logging_frontend_constant_insert_sets_cluster() { + let harness = test_util::TestHarness::default().with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ); + let (server, mut client) = setup_statement_logging_core(1.0, 1.0, "", harness); + + client.execute("SET CLUSTER TO quickstart", &[]).unwrap(); + client + .execute( + "CREATE TABLE statement_logging_constant_insert_t (x INT)", + &[], + ) + .unwrap(); + client + .execute( + "INSERT INTO statement_logging_constant_insert_t VALUES (1)", + &[], + ) + .unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let row = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let rows = client + .query( + "SELECT mseh.cluster_name, mseh.finished_status +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 FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE mst.sql ~~ 'INSERT INTO statement_logging_constant_insert_t%' + AND mseh.finished_at IS NOT NULL +ORDER BY mseh.began_at DESC", + &[], + ) + .unwrap(); + + if let Some(row) = rows.into_iter().next() { + Ok(row) + } else { + Err(()) + } + }) + .expect("constant INSERT statement log entry should be recorded"); + + let cluster_name: Option = row.get(0); + let finished_status: String = row.get(1); + assert_eq!(cluster_name.as_deref(), Some("quickstart")); + assert_eq!(finished_status, "success"); +} + +// Regression test: the frontend OCC read-then-write path must set +// `execution_timestamp` on the statement's log entry. The old, coordinator +// path does this through `set_statement_execution_timestamp` during group +// commit; the frontend path needs to emit the equivalent signal once its +// write commits. +#[mz_ore::test] +fn test_statement_logging_frontend_read_then_write_sets_execution_timestamp() { + let harness = test_util::TestHarness::default().with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ); + let (server, mut client) = setup_statement_logging_core(1.0, 1.0, "", harness); + + client.execute("SET CLUSTER TO quickstart", &[]).unwrap(); + client + .execute("CREATE TABLE statement_logging_rtw_t (x INT)", &[]) + .unwrap(); + client + .execute("INSERT INTO statement_logging_rtw_t VALUES (1), (2)", &[]) + .unwrap(); + // DELETE goes through the frontend OCC read-then-write path. + client + .execute("DELETE FROM statement_logging_rtw_t WHERE x = 1", &[]) + .unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let row = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let rows = client + .query( + "SELECT mseh.execution_timestamp, mseh.finished_status +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 FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE mst.sql ~~ 'DELETE FROM statement_logging_rtw_t%' + AND mseh.finished_at IS NOT NULL +ORDER BY mseh.began_at DESC", + &[], + ) + .unwrap(); + + if let Some(row) = rows.into_iter().next() { + Ok(row) + } else { + Err(()) + } + }) + .expect("DELETE statement log entry should be recorded"); + + let execution_timestamp: Option = row.get(0); + let finished_status: String = row.get(1); + assert_eq!(finished_status, "success"); + assert!( + execution_timestamp.is_some(), + "frontend OCC read-then-write DELETE must set execution_timestamp, got NULL" + ); +} + #[allow(clippy::disallowed_methods)] fn run_throttling_test(use_prepared_statement: bool) { // The `target_data_rate` should be @@ -1337,6 +1453,683 @@ fn test_cancel_long_running_query() { .expect("simple query succeeds after cancellation"); } +// Test that frontend-sequenced read-then-write statements honor pgwire cancel +// requests and do not run to completion after cancellation. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_cancel_frontend_read_then_write_long_running_query() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + server.enable_feature_flags(&["unsafe_enable_unsafe_functions"]); + + let mut client = server.connect(postgres::NoTls).unwrap(); + let cancel_token = client.cancel_token(); + + client + .batch_execute("CREATE TABLE t (a TEXT, ts INT)") + .unwrap(); + client + .batch_execute("INSERT INTO t VALUES ('hello', 10)") + .unwrap(); + + let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel(); + let cancel_thread = thread::spawn(move || { + loop { + thread::sleep(Duration::from_millis(200)); + match shutdown_rx.try_recv() { + Ok(()) => return, + Err(std::sync::mpsc::TryRecvError::Empty) => { + let _ = cancel_token.cancel_query(postgres::NoTls); + } + Err(std::sync::mpsc::TryRecvError::Disconnected) => return, + } + } + }); + + match client.batch_execute( + "INSERT INTO t SELECT a, CASE WHEN mz_unsafe.mz_sleep(ts) > 0 THEN 0 END AS ts FROM t", + ) { + Err(e) if e.code() == Some(&SqlState::QUERY_CANCELED) => {} + Err(e) => panic!("expected error SqlState::QUERY_CANCELED, but got {e:?}"), + Ok(_) => panic!("expected error SqlState::QUERY_CANCELED, but query succeeded"), + } + + shutdown_tx.send(()).unwrap(); + cancel_thread.join().unwrap(); + + let rows = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get::<_, i64>(0); + assert_eq!( + rows, 1, + "cancelled statement should not have committed writes" + ); + + // NOTE: mz_sleep with a constant ts gets evaluated differently. This gives + // us additional coverage for cancelling at different moments in the + // processing pipeline. + let cancel_token = client.cancel_token(); + let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel(); + let cancel_thread = thread::spawn(move || { + loop { + thread::sleep(Duration::from_millis(200)); + match shutdown_rx.try_recv() { + Ok(()) => return, + Err(std::sync::mpsc::TryRecvError::Empty) => { + let _ = cancel_token.cancel_query(postgres::NoTls); + } + Err(std::sync::mpsc::TryRecvError::Disconnected) => return, + } + } + }); + + match client.batch_execute( + "INSERT INTO t SELECT a, CASE WHEN mz_unsafe.mz_sleep(10) > 0 THEN 0 END AS ts FROM t", + ) { + Err(e) if e.code() == Some(&SqlState::QUERY_CANCELED) => {} + Err(e) => panic!("expected error SqlState::QUERY_CANCELED, but got {e:?}"), + Ok(_) => panic!("expected error SqlState::QUERY_CANCELED, but query succeeded"), + } + + shutdown_tx.send(()).unwrap(); + cancel_thread.join().unwrap(); + + let rows = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get::<_, i64>(0); + assert_eq!( + rows, 1, + "cancelled statement should not have committed writes" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_constant_insert_prepares_unmaterializable_functions() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client.batch_execute("CREATE TABLE t (u text)").unwrap(); + client.batch_execute("BEGIN").unwrap(); + client + .execute("INSERT INTO t VALUES (current_user())", &[]) + .unwrap(); + client.batch_execute("COMMIT").unwrap(); + + let inserted_matches_current_user = client + .query_one("SELECT u = current_user() FROM t", &[]) + .unwrap() + .get::<_, bool>(0); + assert!(inserted_matches_current_user); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_rejected_in_multi_statement_batch() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client.batch_execute("CREATE TABLE t (a int)").unwrap(); + client.batch_execute("INSERT INTO t VALUES (1)").unwrap(); + + // Non-constant DML in a multi-statement implicit transaction (a simple + // query batch) is prohibited, matching the coordinator's transaction + // gate. Allowing it into the OCC path would commit the write durably + // mid-batch, breaking the batch's atomicity. + let err = client + .batch_execute("INSERT INTO t SELECT * FROM t; SELECT 1") + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("cannot be run inside a transaction block"), + "unexpected error: {err:?}" + ); + + // Constant INSERTs join the implicit transaction's write ops, so a later + // error in the batch rolls them back. + let err = client + .batch_execute("INSERT INTO t VALUES (2); SELECT 1/0") + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err.message().contains("division by zero"), + "unexpected error: {err:?}" + ); + + let count = client + .query_one("SELECT count(*)::int4 FROM t", &[]) + .unwrap() + .get::<_, i32>(0); + assert_eq!(count, 1, "no batch write may have committed"); +} + +/// An INSERT whose values read no persisted state may run in a transaction, +/// even when the values are too large to fold into a literal. The rows are +/// buffered as session write ops, so they commit with the transaction and +/// disappear if it does not commit. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_nonconstant_insert_in_transaction() { + // Above `FOLD_CONSTANTS_LIMIT`, so the planned values stay a dataflow + // instead of folding into a constant. + const BIG_INSERT: &str = "INSERT INTO t SELECT generate_series(1, 20000)"; + + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE t (a int)").unwrap(); + + let count = |client: &mut postgres::Client| { + client + .query_one("SELECT count(*)::int4 FROM t", &[]) + .unwrap() + .get::<_, i32>(0) + }; + + client + .batch_execute(&format!("BEGIN; {BIG_INSERT}; COMMIT;")) + .unwrap(); + assert_eq!(count(&mut client), 20000); + + client + .batch_execute(&format!("BEGIN; {BIG_INSERT}; ROLLBACK;")) + .unwrap(); + assert_eq!(count(&mut client), 20000, "rolled back write must not land"); + + // A later error in the same implicit batch aborts the transaction, so the + // write must not be visible. + let err = client + .batch_execute(&format!("{BIG_INSERT}; SELECT 1/0;")) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err.message().contains("division by zero"), + "unexpected error: {err:?}" + ); + let _ = client.batch_execute("ROLLBACK"); + assert_eq!(count(&mut client), 20000, "aborted write must not land"); + + // Outside a transaction the same statement commits on its own. + client.batch_execute(BIG_INSERT).unwrap(); + assert_eq!(count(&mut client), 40000); + + // The command tag counts diffs, not distinct rows: 20000 copies of the + // same row consolidate into one row with diff 20000. + client.batch_execute("BEGIN").unwrap(); + let inserted = client + .execute("INSERT INTO t SELECT 1 FROM generate_series(1, 20000)", &[]) + .unwrap(); + assert_eq!(inserted, 20000); + client.batch_execute("COMMIT").unwrap(); + assert_eq!(count(&mut client), 60000); + + // Deferred and constant writes mix freely in one transaction and all land + // at COMMIT. + client.batch_execute("BEGIN").unwrap(); + assert_eq!(client.execute(BIG_INSERT, &[]).unwrap(), 20000); + assert_eq!(client.execute(BIG_INSERT, &[]).unwrap(), 20000); + assert_eq!(client.execute("INSERT INTO t VALUES (1)", &[]).unwrap(), 1); + client.batch_execute("COMMIT").unwrap(); + assert_eq!(count(&mut client), 100001); + + // Constraint violations are caught while the diffs are being collected, so + // the statement fails and the transaction buffers nothing. + client + .batch_execute("CREATE TABLE nn (a int NOT NULL)") + .unwrap(); + client.batch_execute("BEGIN").unwrap(); + let err = client + .execute( + "INSERT INTO nn SELECT CASE WHEN g = 5 THEN NULL ELSE g END \ + FROM generate_series(1, 20000) g", + &[], + ) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err.message().contains("violates not-null constraint"), + "unexpected error: {err:?}" + ); + let _ = client.batch_execute("ROLLBACK"); + let nn_count = client + .query_one("SELECT count(*)::int4 FROM nn", &[]) + .unwrap() + .get::<_, i32>(0); + assert_eq!(nn_count, 0, "failed write must not land"); + + // RETURNING needs the rows to be visible now, so it cannot wait for + // COMMIT. The transaction gate rejects it whether or not the values fold. + client.batch_execute("BEGIN").unwrap(); + let err = client + .query(&format!("{BIG_INSERT} RETURNING a"), &[]) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("cannot be run inside a transaction block"), + "unexpected error: {err:?}" + ); + let _ = client.batch_execute("ROLLBACK"); +} + +/// A transaction that has committed itself to DDL or to a subscribe cannot +/// take a write. The reported error must be the one the coordinator reports, +/// not whatever the write-op merge in the session state machine happens to +/// produce. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_rejected_in_non_writable_transaction() { + const BIG_INSERT: &str = "INSERT INTO t SELECT generate_series(1, 20000)"; + + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE t (a int)").unwrap(); + client.batch_execute("CREATE TABLE z (a int)").unwrap(); + client.batch_execute("CREATE TABLE u (a int)").unwrap(); + client.batch_execute("INSERT INTO u VALUES (1)").unwrap(); + + let assert_read_only = |client: &mut postgres::Client, setup: &[&str]| { + client.batch_execute("BEGIN").unwrap(); + for stmt in setup { + client.batch_execute(stmt).unwrap(); + } + let err = client.execute(BIG_INSERT, &[]).unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err.message().contains("transaction in read-only mode"), + "after {setup:?}, unexpected error: {err:?}" + ); + // The failed statement aborts the transaction, so the rollback is what + // hands the session back in a usable state for the next case. + let _ = client.batch_execute("ROLLBACK"); + }; + + // Each case rolls back, so the catalog changes never apply and the next + // case starts from the same state. + assert_read_only(&mut client, &["ALTER TABLE z RENAME TO z2"]); + assert_read_only(&mut client, &["CREATE TABLE zz (i int)"]); + // The FETCH is what pins the transaction to the subscribe, the DECLARE + // alone leaves it undecided. + assert_read_only( + &mut client, + &["DECLARE c CURSOR FOR SUBSCRIBE u", "FETCH 1 c"], + ); + + let count = client + .query_one("SELECT count(*)::int4 FROM t", &[]) + .unwrap() + .get::<_, i32>(0); + assert_eq!(count, 0, "no rejected write may have landed"); +} + +/// Changing a startup-only parameter is allowed and warns that it only takes +/// effect after a restart. The running process keeps its sampled value, so the +/// routing decision cannot change underneath open sessions. A change that is +/// rejected, or a `RESET ALL` that leaves the parameter where it was, must not +/// warn. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_startup_only_system_var_warns() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .start_blocking(); + + let (tx, mut rx) = futures::channel::mpsc::unbounded(); + let mut client = server + .pg_config_internal() + .notice_callback(move |notice| tx.unbounded_send(notice).expect("send notice")) + .connect(postgres::NoTls) + .unwrap(); + + const WARNING: &str = "only take effect when environmentd restarts"; + let drain = |rx: &mut futures::channel::mpsc::UnboundedReceiver<_>| -> Vec { + std::iter::from_fn(|| rx.try_recv().ok()) + .map(|notice: postgres::error::DbError| notice.message().to_string()) + .collect() + }; + + for stmt in [ + "ALTER SYSTEM SET enable_adapter_frontend_occ_read_then_write = true", + "ALTER SYSTEM RESET enable_adapter_frontend_occ_read_then_write", + ] { + client.batch_execute(stmt).unwrap(); + let notices = drain(&mut rx); + assert!( + notices.iter().any(|message| message.contains(WARNING)), + "{stmt} did not warn, notices: {notices:?}" + ); + } + + // A rejected change must not warn: the catalog value does not move, so a + // restart would not pick anything new up. + let err = client + .batch_execute("ALTER SYSTEM SET max_concurrent_occ_writes = 'not-a-number'") + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("requires a \"unsigned integer\" value"), + "unexpected error: {err:?}" + ); + let notices = drain(&mut rx); + assert!( + !notices.iter().any(|message| message.contains(WARNING)), + "rejected change warned, notices: {notices:?}" + ); + + // Zero permits would block every read-then-write until its statement + // timeout, so the parameter is constrained to at least 1. + let err = client + .batch_execute("ALTER SYSTEM SET max_concurrent_occ_writes = 0") + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("only supports values in range 1.."), + "unexpected error: {err:?}" + ); + + // `RESET ALL` also goes through, and warns for the parameter it changes. + client + .batch_execute("ALTER SYSTEM SET enable_adapter_frontend_occ_read_then_write = true") + .unwrap(); + let _ = drain(&mut rx); + client.batch_execute("ALTER SYSTEM RESET ALL").unwrap(); + let notices = drain(&mut rx); + assert!( + notices.iter().any(|message| message.contains(WARNING)), + "RESET ALL did not warn, notices: {notices:?}" + ); + + // Everything is at its effective default now, so a second `RESET ALL` + // changes no startup-only parameter and must stay quiet about them. + client.batch_execute("ALTER SYSTEM RESET ALL").unwrap(); + let notices = drain(&mut rx); + assert!( + !notices.iter().any(|message| message.contains(WARNING)), + "RESET ALL warned without changing anything, notices: {notices:?}" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_constant_insert_respects_max_result_size() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .with_system_parameter_default("max_result_size".to_string(), "1MB".to_string()) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client + .batch_execute("CREATE TABLE t2 (a int4, b text)") + .unwrap(); + + let err = client + .execute( + "INSERT INTO t2 SELECT * FROM generate_series(1, 10001), repeat('a', 100)", + &[], + ) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("result exceeds max size of 1.0 MiB"), + "unexpected error: {err:?}" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_constant_insert_mz_now_uses_legacy_error() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client + .batch_execute("CREATE TABLE dec (d mz_timestamp)") + .unwrap(); + + let err = client + .execute("INSERT INTO dec VALUES (mz_now())", &[]) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err + .message() + .contains("calls to mz_now in write statements"), + "unexpected error: {err:?}" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_returning_error_does_not_commit_write() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client + .batch_execute("CREATE TABLE t (a INT, b INT)") + .unwrap(); + + let err = client + .query("INSERT INTO t VALUES (7, 8) RETURNING 1/0", &[]) + .unwrap_err(); + let db_err = err.as_db_error().expect("expected db error"); + assert!( + db_err.message().contains("division by zero"), + "unexpected error message: {:?}", + db_err.message() + ); + + let rows = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get::<_, i64>(0); + assert_eq!(rows, 0, "failing RETURNING must not commit the write"); +} + +// Regression test for the empty-snapshot branch of the OCC loop. +// +// `ActiveSubscribe::initialize` emits a progress message at `as_of` before +// any data batch is processed, so the OCC loop must not conclude +// `NoRowsMatched` on that first progress — the snapshot hasn't been +// delivered yet. The check that distinguishes "initial progress" from +// "snapshot complete and empty" is `ts > as_of`; this test exercises both +// empty-match cases and asserts the operations return zero without +// hanging or writing. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_empty_snapshot_returns_zero() { + let server = test_util::TestHarness::default() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + + client.batch_execute("CREATE TABLE t (x INT)").unwrap(); + + // DELETE on a completely empty table. + let deleted = client + .execute("DELETE FROM t", &[]) + .expect("DELETE on empty table should return 0 rows"); + assert_eq!(deleted, 0, "DELETE on empty table must report 0 rows"); + + // DELETE with a WHERE clause that matches no rows against a non-empty + // table. The snapshot is non-empty (contains row (1)) but the selection + // is empty after filtering. + client.batch_execute("INSERT INTO t VALUES (1)").unwrap(); + let deleted = client + .execute("DELETE FROM t WHERE x = 999", &[]) + .expect("DELETE with no matches should return 0 rows"); + assert_eq!(deleted, 0, "DELETE with no matches must report 0 rows"); + + // UPDATE with a WHERE clause that matches no rows. + let updated = client + .execute("UPDATE t SET x = 2 WHERE x = 999", &[]) + .expect("UPDATE with no matches should return 0 rows"); + assert_eq!(updated, 0, "UPDATE with no matches must report 0 rows"); + + // The original row is still there. + let rows = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get::<_, i64>(0); + assert_eq!(rows, 1); +} + +// End-to-end coverage of the OCC retry path: +// +// N concurrent connections each issue M `UPDATE counter SET v = v + 1` +// statements against the same single-row table. Without a working +// `TimestampPassed` retry loop this would lose updates (two writers reading +// `v = k` and both committing `v = k + 1`); the final value pinning down at +// `N * M` proves retries actually re-read fresh state and re-apply the diff. +// +// Also asserts the `mz_occ_read_then_write_retry_count` histogram observes +// every UPDATE and that at least one observation reports a retry, so the +// retry-count metric stays wired up to the OCC loop. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_read_then_write_concurrent_updates_retry() { + const NUM_WORKERS: usize = 4; + const UPDATES_PER_WORKER: usize = 25; + + let server = test_util::TestHarness::default() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut setup = server.connect(postgres::NoTls).unwrap(); + setup + .batch_execute("CREATE TABLE counter (id INT, v INT)") + .unwrap(); + setup + .batch_execute("INSERT INTO counter VALUES (1, 0)") + .unwrap(); + + let mut handles = Vec::with_capacity(NUM_WORKERS); + for _ in 0..NUM_WORKERS { + let mut client = server.connect(postgres::NoTls).unwrap(); + handles.push(thread::spawn(move || { + for _ in 0..UPDATES_PER_WORKER { + client + .execute("UPDATE counter SET v = v + 1 WHERE id = 1", &[]) + .expect("UPDATE under contention should succeed via OCC retry"); + } + })); + } + for handle in handles { + handle.join().expect("worker thread panicked"); + } + + let final_v: i32 = setup + .query_one("SELECT v FROM counter WHERE id = 1", &[]) + .unwrap() + .get(0); + let expected = i32::try_from(NUM_WORKERS * UPDATES_PER_WORKER).unwrap(); + assert_eq!( + final_v, expected, + "concurrent OCC UPDATEs lost updates: expected {expected}, got {final_v}", + ); + + // Inspect the OCC retry-count histogram. Every UPDATE that took the + // frontend OCC path should produce exactly one observation, so + // sample_count must be >= NUM_WORKERS * UPDATES_PER_WORKER. Same-row + // contention essentially guarantees at least one observation lands above + // the 0-retry bucket, so we assert that too. + let metrics = server.metrics_registry().gather(); + let retry_metric = metrics + .iter() + .find(|m| m.name() == "mz_occ_read_then_write_retry_count") + .expect("mz_occ_read_then_write_retry_count metric should be registered"); + let metric = retry_metric.get_metric(); + assert_eq!(metric.len(), 1, "expected a single histogram series"); + let histogram = metric[0].get_histogram(); + + let total_updates = u64::try_from(NUM_WORKERS * UPDATES_PER_WORKER).unwrap(); + assert!( + histogram.get_sample_count() >= total_updates, + "expected at least {} OCC observations, got {}", + total_updates, + histogram.get_sample_count(), + ); + + let zero_retry_bucket = histogram + .get_bucket() + .iter() + .find(|b| b.upper_bound() == 0.0) + .expect("histogram should have a 0-retry bucket"); + assert!( + zero_retry_bucket.cumulative_count() < histogram.get_sample_count(), + "expected at least one UPDATE to retry under contention; \ + all {} observations landed in the 0-retry bucket", + histogram.get_sample_count(), + ); +} + #[allow(clippy::disallowed_methods)] fn test_cancellation_cancels_dataflows(query: &str) { // Query that returns how many dataflows are currently installed. @@ -3573,6 +4366,45 @@ fn test_cancel_read_then_write() { .unwrap(); } +// A frontend OCC read-then-write that times out must not commit its writes. +#[mz_ore::test] +#[cfg_attr(miri, ignore)] +#[allow(clippy::disallowed_methods)] +fn test_timeout_frontend_read_then_write() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking(); + server.enable_feature_flags(&["unsafe_enable_unsafe_functions"]); + + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE frontend_timeout (a TEXT, ts INT);") + .unwrap(); + client + .batch_execute("INSERT INTO frontend_timeout VALUES ('hello', 10)") + .unwrap(); + client + .batch_execute("SET statement_timeout = '5s'") + .unwrap(); + + let err = client + .batch_execute( + "INSERT INTO frontend_timeout SELECT a, CASE WHEN mz_unsafe.mz_sleep(ts) > 0 THEN 0 END AS ts FROM frontend_timeout", + ) + .unwrap_err(); + assert_contains!(err.to_string_with_causes(), "statement timeout"); + + let rows: i64 = client + .query_one("SELECT count(*) FROM frontend_timeout", &[]) + .unwrap() + .get(0); + assert_eq!(rows, 1, "timed-out statement committed writes"); +} + #[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))] #[cfg_attr(miri, ignore)] // too slow async fn test_http_metrics() { @@ -7290,3 +8122,649 @@ fn test_shutdown_with_inflight_writes() { tracing::info!("round {round} survived"); } } + +/// Helper: start a server with frontend OCC read-then-write enabled. +fn qa_occ_server() -> test_util::TestServerWithRuntime { + test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .start_blocking() +} + +/// Concurrent DELETEs of the same multiset rows must not over-delete. +/// With a row of multiplicity M and N concurrent deleters, exactly the +/// committed deletes should sum to M, the table must end empty, and the +/// stored multiplicity must never go negative. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_concurrent_delete_no_overdelete() { + const MULTIPLICITY: i32 = 7; + const NUM_WORKERS: usize = 6; + + let server = qa_occ_server(); + let mut setup = server.connect(postgres::NoTls).unwrap(); + setup.batch_execute("CREATE TABLE t (id INT)").unwrap(); + + for _round in 0..10 { + setup.batch_execute("DELETE FROM t").unwrap(); + setup + .execute( + "INSERT INTO t SELECT 1 FROM generate_series(1, $1)", + &[&MULTIPLICITY], + ) + .unwrap(); + + let total_deleted = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::new(); + for _ in 0..NUM_WORKERS { + let mut client = server.connect(postgres::NoTls).unwrap(); + let total_deleted = Arc::clone(&total_deleted); + handles.push(thread::spawn(move || { + let n = client + .execute("DELETE FROM t WHERE id = 1", &[]) + .expect("DELETE under contention should succeed via OCC"); + total_deleted.fetch_add(usize::try_from(n).unwrap(), Ordering::SeqCst); + })); + } + for h in handles { + h.join().expect("worker panicked"); + } + + let remaining: i64 = setup + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get(0); + assert_eq!( + remaining, 0, + "table should be empty after concurrent deletes" + ); + assert_eq!( + total_deleted.load(Ordering::SeqCst), + usize::try_from(MULTIPLICITY).unwrap(), + "sum of reported deletes must equal initial multiplicity (no over/under-delete)", + ); + } +} + +/// Multiset multiplicity must be reflected in affected-row counts for +/// DELETE / UPDATE / INSERT...SELECT. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_duplicate_row_multiplicity_counts() { + let server = qa_occ_server(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE t (a INT, b INT)") + .unwrap(); + + // 3 identical rows. + client + .execute("INSERT INTO t SELECT 1, 10 FROM generate_series(1, 3)", &[]) + .unwrap(); + + // UPDATE all 3 -> Updated(3). + let n = client + .execute("UPDATE t SET b = 20 WHERE a = 1", &[]) + .unwrap(); + assert_eq!( + n, 3, + "UPDATE should report 3 affected rows for multiplicity 3" + ); + + // INSERT INTO t SELECT * FROM t -> doubles, returns 3. + let n = client + .execute("INSERT INTO t SELECT a, b FROM t", &[]) + .unwrap(); + assert_eq!(n, 3, "INSERT...SELECT should report 3 inserted rows"); + let cnt: i64 = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get(0); + assert_eq!(cnt, 6); + + // DELETE all 6 -> Deleted(6). + let n = client.execute("DELETE FROM t WHERE a = 1", &[]).unwrap(); + assert_eq!(n, 6, "DELETE should report 6 affected rows"); + let cnt: i64 = client + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get(0); + assert_eq!(cnt, 0); +} + +/// NOT NULL constraint violations via UPDATE and INSERT...SELECT must error +/// and leave the table unchanged (no partial commit). +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_not_null_constraint_enforced() { + let server = qa_occ_server(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE t (a INT NOT NULL, b INT)") + .unwrap(); + client.execute("INSERT INTO t VALUES (1, 10)", &[]).unwrap(); + + // UPDATE setting NOT NULL column to NULL must fail. + let err = client + .execute("UPDATE t SET a = NULL WHERE b = 10", &[]) + .unwrap_err(); + let msg = err + .as_db_error() + .expect("expected db error") + .message() + .to_lowercase(); + assert!( + msg.contains("null"), + "expected null-constraint error, got: {msg}" + ); + + // INSERT...SELECT producing a NULL into a NOT NULL column must fail. + let err = client + .execute("INSERT INTO t SELECT NULL::INT, 99", &[]) + .unwrap_err(); + let msg = err + .as_db_error() + .expect("expected db error") + .message() + .to_lowercase(); + assert!( + msg.contains("null"), + "expected null-constraint error, got: {msg}" + ); + + // Table must be unchanged: still exactly (1, 10). + let rows = client.query("SELECT a, b FROM t", &[]).unwrap(); + assert_eq!(rows.len(), 1, "no rows should have been added/removed"); + let a: i32 = rows[0].get(0); + let b: i32 = rows[0].get(1); + assert_eq!( + (a, b), + (1, 10), + "row must be unchanged after failed mutations" + ); +} + +/// RETURNING must produce new values for UPDATE, old values for DELETE, +/// and inserted values for INSERT. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_returning_values_correct() { + let server = qa_occ_server(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE t (id INT, v INT)") + .unwrap(); + client + .batch_execute("INSERT INTO t VALUES (1, 100), (2, 200)") + .unwrap(); + + // Materialize only supports RETURNING on INSERT (the parser rejects it for + // UPDATE/DELETE), so we exercise INSERT...RETURNING in both the constant + // and the read-dependent (INSERT...SELECT) shapes. + + // Constant INSERT RETURNING with an expression over the inserted row. + let rows = client + .query("INSERT INTO t VALUES (3, 300) RETURNING id, v, v * 2", &[]) + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::<_, i32>(0), 3); + assert_eq!(rows[0].get::<_, i32>(1), 300); + assert_eq!( + rows[0].get::<_, i32>(2), + 600, + "RETURNING expression mis-evaluated" + ); + + // INSERT...SELECT RETURNING reading existing rows (goes through OCC). + let mut rows: Vec<(i32, i32)> = client + .query( + "INSERT INTO t SELECT id + 100, v + 1 FROM t WHERE id <= 2 RETURNING id, v", + &[], + ) + .unwrap() + .iter() + .map(|r| (r.get(0), r.get(1))) + .collect(); + rows.sort(); + assert_eq!( + rows, + vec![(101, 101), (102, 201)], + "INSERT...SELECT RETURNING returned wrong inserted rows" + ); + + // Final state: original (1,100),(2,200),(3,300) plus (101,101),(102,201). + let mut got: Vec<(i32, i32)> = client + .query("SELECT id, v FROM t", &[]) + .unwrap() + .iter() + .map(|r| (r.get(0), r.get(1))) + .collect(); + got.sort(); + assert_eq!( + got, + vec![(1, 100), (2, 200), (3, 300), (101, 101), (102, 201)] + ); +} + +/// UPDATEs that move rows into a range that overlaps existing rows must +/// produce the correct final set (exercises the Let/Negate/map MIR +/// transform with consolidation overlap). +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_update_moves_overlapping_rows() { + let server = qa_occ_server(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE t (id INT)").unwrap(); + client + .execute("INSERT INTO t SELECT generate_series(1, 5)", &[]) + .unwrap(); + + // Overlapping shift: {1,2,3,4,5} -> {2,3,4,5,6}. + client.execute("UPDATE t SET id = id + 1", &[]).unwrap(); + let mut got: Vec = client + .query("SELECT id FROM t ORDER BY id", &[]) + .unwrap() + .iter() + .map(|r| r.get(0)) + .collect(); + got.sort(); + assert_eq!( + got, + vec![2, 3, 4, 5, 6], + "overlapping +1 shift produced wrong set" + ); + + // Non-overlapping shift: {2..6} -> {12..16}. + client.execute("UPDATE t SET id = id + 10", &[]).unwrap(); + let mut got: Vec = client + .query("SELECT id FROM t ORDER BY id", &[]) + .unwrap() + .iter() + .map(|r| r.get(0)) + .collect(); + got.sort(); + assert_eq!(got, vec![12, 13, 14, 15, 16]); +} + +/// INSERT INTO t SELECT FROM a materialized view must read the MV's content +/// correctly. Exercises the TimestampDependent timeline + linearization +/// defaulting to EpochMilliseconds. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_insert_select_from_mv() { + let server = qa_occ_server(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE src (a INT)").unwrap(); + client + .batch_execute("INSERT INTO src VALUES (1), (2), (3)") + .unwrap(); + client + .batch_execute("CREATE MATERIALIZED VIEW mv AS SELECT a * 10 AS a FROM src") + .unwrap(); + client.batch_execute("CREATE TABLE dst (a INT)").unwrap(); + + let n = client + .execute("INSERT INTO dst SELECT a FROM mv", &[]) + .unwrap(); + assert_eq!(n, 3); + let mut got: Vec = client + .query("SELECT a FROM dst ORDER BY a", &[]) + .unwrap() + .iter() + .map(|r| r.get(0)) + .collect(); + got.sort(); + assert_eq!(got, vec![10, 20, 30]); +} + +/// Concurrent mixed DML (UPDATE / DELETE / INSERT...SELECT) on one table must +/// not panic the coordinator or produce internal errors, and the server must +/// remain responsive afterward. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_concurrent_mixed_dml_no_internal_error() { + const NUM_WORKERS: usize = 8; + const ITERS: usize = 30; + + let server = qa_occ_server(); + let mut setup = server.connect(postgres::NoTls).unwrap(); + setup + .batch_execute("CREATE TABLE t (id INT, v INT)") + .unwrap(); + setup + .execute("INSERT INTO t SELECT generate_series(1, 20), 0", &[]) + .unwrap(); + + let internal_errors = Arc::new(Mutex::new(Vec::::new())); + let mut handles = Vec::new(); + for w in 0..NUM_WORKERS { + let mut client = server.connect(postgres::NoTls).unwrap(); + let internal_errors = Arc::clone(&internal_errors); + handles.push(thread::spawn(move || { + for i in 0..ITERS { + let id = (w * 7 + i) % 20 + 1; + let stmt = match i % 4 { + 0 => format!("UPDATE t SET v = v + 1 WHERE id = {id}"), + 1 => format!("DELETE FROM t WHERE id = {id} AND v < 0"), + 2 => format!("INSERT INTO t SELECT {id}, 0 WHERE NOT EXISTS (SELECT 1 FROM t WHERE id = {id})"), + _ => format!("UPDATE t SET v = v WHERE id = {id}"), + }; + if let Err(e) = client.execute(stmt.as_str(), &[]) { + let msg = e.to_string(); + // Contention exhaustion is acceptable; internal errors are not. + if msg.contains("internal error") + || msg.contains("reached the coordinator") + || msg.contains("unexpectedly got") + { + internal_errors.lock().unwrap().push(msg); + } + } + } + })); + } + for h in handles { + h.join().expect("worker panicked"); + } + + let errs = internal_errors.lock().unwrap(); + assert!(errs.is_empty(), "internal errors observed: {:?}", *errs); + + // Server still responsive. + let cnt: i64 = setup + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get(0); + assert!(cnt >= 0); +} + +/// FIX VERIFICATION: A read-then-write whose read resolves to a far-future +/// timestamp (here, a `REFRESH AT ` materialized view) used to hang +/// in `ensure_read_linearized`'s sleep loop, ignoring `statement_timeout`. The +/// fix bounds the *entire* OCC read-then-write operation by `statement_timeout` +/// centrally, in `SessionClient::try_frontend_read_then_write_with_cancel`'s +/// `select!`, so the parked far-future op now hits the timeout and returns +/// promptly with a statement-timeout error. +/// +/// NOTE: the fix is OCC-path-only. The legacy path still hangs on the same +/// scenario — see `qa_legacy_far_future_refresh_mv_statement_timeout`. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_far_future_refresh_mv_respects_statement_timeout() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .with_system_parameter_default("enable_refresh_every_mvs".to_string(), "true".to_string()) + .start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE src (a INT)").unwrap(); + client + .batch_execute("INSERT INTO src VALUES (1), (2), (3)") + .unwrap(); + // Refresh only at a far-future instant: the MV holds no readable content + // until then, so a freshest-table-write read must pick a far-future as_of. + client + .batch_execute( + "CREATE MATERIALIZED VIEW mv \ + WITH (REFRESH AT '3000-01-01 00:00:00') AS SELECT a FROM src", + ) + .unwrap(); + client.batch_execute("CREATE TABLE dst (a INT)").unwrap(); + + let mut worker = server.connect(postgres::NoTls).unwrap(); + // Grab a cancel token so we can free the parked statement for a clean + // teardown if it hangs. + let cancel = worker.cancel_token(); + let (tx, rx) = std::sync::mpsc::channel(); + let handle = thread::spawn(move || { + worker + .batch_execute("SET statement_timeout = '3s'") + .unwrap(); + let started = Instant::now(); + let res = worker.batch_execute("INSERT INTO dst SELECT a FROM mv"); + // `tokio_postgres::Error::to_string()` is just "db error"; preserve the + // SqlState code and server message so the assertion can inspect them. + let res = res.map_err(|e| { + ( + e.code().cloned(), + e.as_db_error().map(|d| d.message().to_string()), + ) + }); + let _ = tx.send((started.elapsed(), res)); + }); + + let outcome = rx.recv_timeout(Duration::from_secs(45)); + if outcome.is_err() { + // Free the parked statement so the server can shut down cleanly. + let _ = cancel.cancel_query(postgres::NoTls); + } + let _ = handle.join(); + + match outcome { + Ok((elapsed, res)) => { + // The fix bounds the whole operation by `statement_timeout`, so the + // far-future op must error out (not silently succeed) and do so + // promptly — well within the 45s recv budget. + let (code, message) = res.expect_err( + "far-future RTW should have failed with a statement-timeout error, \ + but it returned successfully", + ); + // `StatementTimeout` surfaces as QUERY_CANCELED with the standard + // "canceling statement due to statement timeout" message. + assert_eq!( + code.as_ref(), + Some(&SqlState::QUERY_CANCELED), + "far-future RTW failed with unexpected SqlState {code:?} (message: {message:?})" + ); + assert!( + message + .as_deref() + .is_some_and(|m| m.to_lowercase().contains("timeout")), + "far-future RTW error did not mention a timeout: {message:?}" + ); + assert!( + elapsed < Duration::from_secs(15), + "statement_timeout was 3s but the op took {elapsed:?} to return", + ); + } + Err(recv_err) => { + panic!( + "INSERT...SELECT from a far-future REFRESH AT MV did not return \ + within 45s despite statement_timeout = '3s'; the central \ + statement_timeout enforcement in \ + try_frontend_read_then_write_with_cancel did not fire ({recv_err})." + ); + } + } +} + +/// CONTEXT (ignored): the same far-future scenario on the LEGACY (non-OCC) +/// path also hangs past `statement_timeout`. This documents that the +/// indefinite hang is pre-existing behavior, not a regression introduced by +/// the OCC path. +/// +/// Stays `#[ignore]`d: the statement-timeout fix is OCC-path-only (it lives in +/// `try_frontend_read_then_write_with_cancel`). The legacy coordinator path is +/// intentionally left unchanged and out of scope for this fix, so this would +/// still hang in CI. +#[mz_ore::test] +#[ignore = "legacy path far-future hang is out of scope for the OCC statement_timeout fix"] +#[allow(clippy::disallowed_methods)] +fn qa_legacy_far_future_refresh_mv_statement_timeout() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default("enable_refresh_every_mvs".to_string(), "true".to_string()) + .start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE src (a INT)").unwrap(); + client + .batch_execute("INSERT INTO src VALUES (1), (2), (3)") + .unwrap(); + client + .batch_execute( + "CREATE MATERIALIZED VIEW mv \ + WITH (REFRESH AT '3000-01-01 00:00:00') AS SELECT a FROM src", + ) + .unwrap(); + client.batch_execute("CREATE TABLE dst (a INT)").unwrap(); + + let mut worker = server.connect(postgres::NoTls).unwrap(); + let cancel = worker.cancel_token(); + let (tx, rx) = std::sync::mpsc::channel(); + let handle = thread::spawn(move || { + worker + .batch_execute("SET statement_timeout = '3s'") + .unwrap(); + let started = Instant::now(); + let res = worker.batch_execute("INSERT INTO dst SELECT a FROM mv"); + let _ = tx.send((started.elapsed(), res.map_err(|e| e.to_string()))); + }); + + let outcome = rx.recv_timeout(Duration::from_secs(45)); + // Cancel the parked statement (if any) for a clean teardown, then join. + let _ = cancel.cancel_query(postgres::NoTls); + let _ = handle.join(); + + match outcome { + Ok((elapsed, res)) => { + eprintln!("LEGACY far-future RTW returned in {elapsed:?} with result {res:?}"); + assert!( + elapsed < Duration::from_secs(40), + "legacy far-future RTW also appears to hang ({elapsed:?})" + ); + } + Err(recv_err) => { + panic!( + "LEGACY INSERT...SELECT from a far-future REFRESH AT MV also did not \ + return within 45s; the far-future hang is pre-existing, not an OCC \ + regression ({recv_err})." + ); + } + } +} + +/// FIX VERIFICATION: a far-future RTW parked in `ensure_read_linearized` holds +/// an OCC semaphore permit while parked. With a bounded permit pool, a single +/// such op used to starve *all* other read-then-writes — even on unrelated +/// tables — because those victims block on permit acquisition *before* the OCC +/// loop, and `statement_timeout` was only enforced inside the loop. +/// +/// The fix moves `statement_timeout` enforcement up into +/// `try_frontend_read_then_write_with_cancel`'s `select!`, which bounds the +/// *whole* operation — including the permit-acquisition wait. So a victim +/// blocked behind the starved pool now honors its own `statement_timeout` and +/// returns promptly with a statement-timeout error instead of hanging. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn qa_occ_far_future_rtw_starves_permit_pool() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) + .with_system_parameter_default("enable_refresh_every_mvs".to_string(), "true".to_string()) + // One permit: a single hung op exhausts the pool. + .with_system_parameter_default("max_concurrent_occ_writes".to_string(), "1".to_string()) + .start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client.batch_execute("CREATE TABLE src (a INT)").unwrap(); + client.batch_execute("INSERT INTO src VALUES (1)").unwrap(); + client + .batch_execute( + "CREATE MATERIALIZED VIEW mv \ + WITH (REFRESH AT '3000-01-01 00:00:00') AS SELECT a FROM src", + ) + .unwrap(); + client.batch_execute("CREATE TABLE dst (a INT)").unwrap(); + client + .batch_execute("CREATE TABLE other (id INT, v INT)") + .unwrap(); + client + .batch_execute("INSERT INTO other VALUES (1, 0)") + .unwrap(); + + // Launch the hung far-future RTW; it will grab the single OCC permit and + // park in ensure_read_linearized. + let mut hung = server.connect(postgres::NoTls).unwrap(); + let hung_cancel = hung.cancel_token(); + let hung_handle = thread::spawn(move || { + let _ = hung.batch_execute("INSERT INTO dst SELECT a FROM mv"); + }); + // Give it time to acquire the permit and reach the linearize loop. + thread::sleep(Duration::from_secs(3)); + + // A completely unrelated UPDATE, with a short statement_timeout, should be + // able to make progress. Run it on a worker thread with a wall-clock guard. + let mut victim = server.connect(postgres::NoTls).unwrap(); + let victim_cancel = victim.cancel_token(); + let (tx, rx) = std::sync::mpsc::channel(); + let victim_handle = thread::spawn(move || { + victim + .batch_execute("SET statement_timeout = '3s'") + .unwrap(); + let started = Instant::now(); + let res = victim.execute("UPDATE other SET v = v + 1 WHERE id = 1", &[]); + // Preserve the SqlState code and server message (`to_string()` is just + // "db error"). + let res = res.map_err(|e| { + ( + e.code().cloned(), + e.as_db_error().map(|d| d.message().to_string()), + ) + }); + let _ = tx.send((started.elapsed(), res)); + }); + + let outcome = rx.recv_timeout(Duration::from_secs(25)); + // Free both parked statements for a clean teardown. + let _ = victim_cancel.cancel_query(postgres::NoTls); + let _ = hung_cancel.cancel_query(postgres::NoTls); + let _ = victim_handle.join(); + let _ = hung_handle.join(); + + match outcome { + Ok((elapsed, res)) => { + // The far-future op holds the sole permit for its (default 60s) + // lifetime, so the victim cannot acquire a permit. With the fix, + // the victim's own `statement_timeout = '3s'` now bounds the + // permit-acquisition wait, so it returns a timeout error rather + // than hanging. (Before the fix this never returned within 25s.) + let (code, message) = res.expect_err( + "victim UPDATE should have timed out waiting on the starved permit pool, \ + but it returned successfully", + ); + assert_eq!( + code.as_ref(), + Some(&SqlState::QUERY_CANCELED), + "victim UPDATE failed with unexpected SqlState {code:?} (message: {message:?})" + ); + assert!( + message + .as_deref() + .is_some_and(|m| m.to_lowercase().contains("timeout")), + "victim UPDATE error did not mention a timeout: {message:?}" + ); + // It should time out on its own 3s budget, well within the 25s + // recv guard. + assert!( + elapsed < Duration::from_secs(15), + "victim statement_timeout was 3s but it took {elapsed:?} to return", + ); + } + Err(recv_err) => { + panic!( + "an unrelated UPDATE on a different table did not return within 25s \ + (statement_timeout = '3s'). A single far-future read-then-write holds the \ + sole OCC permit while parked in ensure_read_linearized; with the central \ + statement_timeout fix the victim should time out on permit acquisition \ + within ~3s instead of hanging ({recv_err})." + ); + } + } +} diff --git a/src/environmentd/tests/sql.rs b/src/environmentd/tests/sql.rs index 6e891c8b1b85b..a187df7cf1683 100644 --- a/src/environmentd/tests/sql.rs +++ b/src/environmentd/tests/sql.rs @@ -1716,8 +1716,12 @@ fn test_subscribe_outlive_cluster() { .batch_execute("CREATE CLUSTER newcluster REPLICAS (r1 (size 'scale=1,workers=1'))") .unwrap(); client2_cancel.cancel_query(postgres::NoTls).unwrap(); - client2 - .batch_execute("ROLLBACK; SET CLUSTER = default") + // The cancel is asynchronous and might race with subsequent commands. + // Retry ROLLBACK in a loop in case it gets canceled. + Retry::default() + .max_tries(5) + .clamp_backoff(Duration::from_millis(100)) + .retry(|_| client2.batch_execute("ROLLBACK; SET CLUSTER = default")) .unwrap(); assert_eq!( client2 @@ -1731,7 +1735,28 @@ fn test_subscribe_outlive_cluster() { #[mz_ore::test] #[allow(clippy::disallowed_methods)] fn test_read_then_write_serializability() { - let server = test_util::TestHarness::default().start_blocking(); + test_read_then_write_serializability_inner(false); +} + +// Same as `test_read_then_write_serializability`, but exercising the frontend +// OCC read-then-write path. Concurrent `INSERT INTO t SELECT * FROM t` must +// still double the row count exactly, i.e. OCC retries must prevent lost +// updates. +#[mz_ore::test] +fn test_read_then_write_serializability_frontend_occ() { + test_read_then_write_serializability_inner(true); +} + +#[allow(clippy::disallowed_methods)] +fn test_read_then_write_serializability_inner(frontend_occ: bool) { + let mut harness = test_util::TestHarness::default(); + if frontend_occ { + harness = harness.with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ); + } + let server = harness.start_blocking(); // Create table with initial value { diff --git a/src/sql/src/session/vars.rs b/src/sql/src/session/vars.rs index a34499861cc15..6fcbc3ce16808 100644 --- a/src/sql/src/session/vars.rs +++ b/src/sql/src/session/vars.rs @@ -1197,6 +1197,8 @@ impl SystemVars { &MAX_RESULT_SIZE, &MAX_COPY_FROM_ROW_SIZE, &ALLOWED_CLUSTER_REPLICA_SIZES, + &MAX_CONCURRENT_OCC_WRITES, + &MAX_OCC_RETRIES, &upsert_rocksdb::UPSERT_ROCKSDB_COMPACTION_STYLE, &upsert_rocksdb::UPSERT_ROCKSDB_OPTIMIZE_COMPACTION_MEMTABLE_BUDGET, &upsert_rocksdb::UPSERT_ROCKSDB_LEVEL_COMPACTION_DYNAMIC_LEVEL_BYTES, @@ -1750,6 +1752,16 @@ impl SystemVars { .collect() } + /// Returns the value of the `max_concurrent_occ_writes` configuration parameter. + pub fn max_concurrent_occ_writes(&self) -> u32 { + *self.expect_value(&MAX_CONCURRENT_OCC_WRITES) + } + + /// Returns the value of the `max_occ_retries` configuration parameter. + pub fn max_occ_retries(&self) -> u32 { + *self.expect_value(&MAX_OCC_RETRIES) + } + /// Returns the value of the `default_cluster_replication_factor` configuration parameter. pub fn default_cluster_replication_factor(&self) -> u32 { *self.expect_value::(&DEFAULT_CLUSTER_REPLICATION_FACTOR) diff --git a/src/sql/src/session/vars/constraints.rs b/src/sql/src/session/vars/constraints.rs index f373fc559c430..bbea900536360 100644 --- a/src/sql/src/session/vars/constraints.rs +++ b/src/sql/src/session/vars/constraints.rs @@ -28,6 +28,8 @@ pub static NUMERIC_BOUNDED_0_1_INCLUSIVE: NumericInRange> = pub static BYTESIZE_AT_LEAST_1MB: ByteSizeInRange> = ByteSizeInRange(ByteSize::mb(1)..); +pub static U32_AT_LEAST_1: U32InRange> = U32InRange(1..); + #[derive(Debug)] pub enum ValueConstraint { /// Variable is read-only and cannot be updated. @@ -197,3 +199,25 @@ where } } } + +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct U32InRange(pub R); + +impl DomainConstraint for U32InRange +where + R: RangeBounds + std::fmt::Debug + Send + Sync + 'static, +{ + type Value = u32; + + fn check(&self, var: &dyn Var, n: &u32) -> Result<(), VarError> { + if self.0.contains(n) { + Ok(()) + } else { + Err(VarError::InvalidParameterValue { + name: var.name(), + invalid_values: vec![n.to_string()], + reason: format!("only supports values in range {:?}", self.0), + }) + } + } +} diff --git a/src/sql/src/session/vars/definitions.rs b/src/sql/src/session/vars/definitions.rs index d8c2c2a0bdf6f..6d117122fe662 100644 --- a/src/sql/src/session/vars/definitions.rs +++ b/src/sql/src/session/vars/definitions.rs @@ -41,7 +41,7 @@ use uncased::UncasedStr; use crate::session::user::{SUPPORT_USER, SYSTEM_USER, User}; use crate::session::vars::constraints::{ BYTESIZE_AT_LEAST_1MB, DomainConstraint, NON_ZERO_DURATION, NUMERIC_BOUNDED_0_1_INCLUSIVE, - NUMERIC_NON_NEGATIVE, ValueConstraint, + NUMERIC_NON_NEGATIVE, U32_AT_LEAST_1, ValueConstraint, }; use crate::session::vars::errors::VarError; use crate::session::vars::polyfill::{LazyValueFn, lazy_value, value}; @@ -645,6 +645,24 @@ pub static ALLOWED_CLUSTER_REPLICA_SIZES: VarDefinition = VarDefinition::new( true, ); +/// Sizes the OCC write semaphore at boot. Zero permits would block every +/// read-then-write until its `statement_timeout`, so the value must be at +/// least 1. +pub static MAX_CONCURRENT_OCC_WRITES: VarDefinition = VarDefinition::new( + "max_concurrent_occ_writes", + value!(u32; 4), + "Maximum number of concurrent read-then-write (DELETE/UPDATE) operations using OCC. Read at startup; changes require an environmentd restart (Materialize).", + false, +) +.with_constraint(&U32_AT_LEAST_1); + +pub static MAX_OCC_RETRIES: VarDefinition = VarDefinition::new( + "max_occ_retries", + value!(u32; 1000), + "Maximum number of OCC retry attempts per read-then-write operation before giving up (Materialize).", + false, +); + pub static PERSIST_FAST_PATH_LIMIT: VarDefinition = VarDefinition::new( "persist_fast_path_limit", value!(usize; 25),