diff --git a/ci/nightly/pipeline.template.yml b/ci/nightly/pipeline.template.yml index 4716f7e1ae149..2182770f01e6c 100644 --- a/ci/nightly/pipeline.template.yml +++ b/ci/nightly/pipeline.template.yml @@ -2408,6 +2408,17 @@ steps: agents: queue: hetzner-aarch64-4cpu-8gb + - id: occ-mixed-mode-read-then-write + label: "Read-then-write sequencing split across two environmentd processes" + depends_on: build-aarch64 + timeout_in_minutes: 30 + plugins: + - ./ci/plugins/mzcompose: + composition: txn-wal-fencing + run: mixed-mode-read-then-write + agents: + queue: hetzner-aarch64-4cpu-8gb + - group: "Copy" key: copy steps: diff --git a/ci/test/lint-main/checks/check-mzcompose-files.sh b/ci/test/lint-main/checks/check-mzcompose-files.sh index 645a263b74dd2..9045e3d8934c2 100755 --- a/ci/test/lint-main/checks/check-mzcompose-files.sh +++ b/ci/test/lint-main/checks/check-mzcompose-files.sh @@ -55,6 +55,7 @@ check_default_workflow_references_others() { -not -wholename "./test/workload-replay/mzcompose.py" `# Handled differently` \ -not -wholename "./test/race-condition/mzcompose.py" `# rotate-keys-race workflow is run separately` \ -not -wholename "./test/aws-glue-schema-registry/mzcompose.py" `# 'aws' workflow runs against real AWS, opt-in via nightly only` \ + -not -wholename "./test/txn-wal-fencing/mzcompose.py" `# mixed-mode-read-then-write workflow is run separately` \ ) for file in "${MZCOMPOSE_TEST_FILES[@]}"; do 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..c10ce67f6fa69 100644 --- a/doc/developer/design/20260210_incremental_occ_read_then_write.md +++ b/doc/developer/design/20260210_incremental_occ_read_then_write.md @@ -36,15 +36,14 @@ a subscribe that continually tracks the current state of the data. ## Non-Goals - High-performance writes under heavy contention. The current implementation - serializes writes behind a global lock; the new implementation serializes - them via OCC retries. Neither is designed for high write throughput. + serializes writes behind a global lock. The OCC implementation serializes + them via retries. Neither is designed for high write throughput. - 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 on this path commits at the frontier it + observed, which it cannot postpone until COMMIT, so it runs only as a single + statement. ## Overview @@ -128,7 +127,7 @@ Session Task Coordinator | | |-- acquire OCC semaphore | | | - |-- CreateReadThenWriteSubscribe ----> | + |-- CreateInternalSubscribe ---------> | | <------------ subscribe channel -----| | | | +-- OCC Loop ------------------+ | @@ -141,14 +140,14 @@ Session Task Coordinator | | if Success: break | | | +------------------------------+ | | | - |-- DropReadThenWriteSubscribe ------> | + |-- DropInternalSubscribe -----------> | | | ``` ### Timestamped writes A timestamped write is a write that must be committed at a specific timestamp. -The group commit machinery has to be extended to supports this by: +The group commit machinery has to be extended to support this by: 1. Checking if the target timestamp is still valid (hasn't been passed by the oracle) @@ -193,7 +192,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 @@ -242,7 +241,7 @@ oracle read timestamp. However, actually applying the write bumps the oracle read timestamp to at least the write timestamp, so at write time it holds that `write_ts <= oracle_read_ts`. The linearization invariant is maintained. -### Single timestamped write write per group commit round +### Single timestamped write per group commit round Only one timestamped write is processed per group commit round. This is correct because: @@ -257,15 +256,22 @@ because: ### Timeouts -We have to be careful about bounding the lifetime of the occ loop, both in -wallclock time and number of retries. With the old approach, a read-then-write -could take arbitrarily long, and block the rest of the system. With the new -approach, the occ loop might try arbitrarily long, without ever succeeding. It -will not block the rest of the system, though, which is a big benefit. +The lifetime of the OCC loop has to be bounded, both in wallclock time and in +number of retries. With the lock-based approach, a read-then-write could take +arbitrarily long and block the rest of the system. With OCC it can retry +arbitrarily long without ever succeeding, but it does not block the rest of the +system, which is a big benefit. -As a safety net, we should bound the lifetime of the occ loop with our existing -statement timeout, and potentially add a hard upper limit on the number of -attempts per occ loop. +`statement_timeout` provides the wallclock bound. It is enforced in the session +task, around the whole operation rather than around the loop alone, so it also +covers planning, OCC permit acquisition, timestamp determination, and read +linearization. Any of those can park indefinitely, and a parked operation holds +an OCC permit, so a bound on the loop alone would leave the permit pool +starvable. + +`max_occ_retries` provides the retry bound. A statement that keeps losing the +race for its write timestamp fails with a contention error instead of retrying +forever. ### Comparison with the old approach @@ -284,6 +290,37 @@ The new approach is arguably easier to reason about: there is no global lock state to consider, no deferred operations, no lock merging. The correctness argument is local to the OCC loop and the group commit mechanism. +## Deliberate differences from the lock-based path + +A user must not be able to tell which path sequenced their statement. These are +the places where the two paths do differ, on purpose. They are listed here so +that the next reader does not take them for bugs. + +- **Statement lifecycle events.** The frontend path records an + `optimization-finished` event for a DML, the coordinator path does not, + because it hands the read-then-write's inner peek a trivial logging context + and so logs nothing for it. We keep the extra event, it is real information + about a statement the user did run. +- **`max_result_size` accounting.** The coordinator sums one row length per diff + entry before consolidation. The frontend recomputes the total from the + consolidated set, which counts one row length per distinct row and ignores + multiplicity. So a `DELETE` of a million copies of one row can exceed the + limit on the coordinator path and succeed on the frontend path. We keep the + frontend's accounting: it matches what the write actually appends, one entry + with a large diff. +- **The write-timeline throttle.** A timestamped write does not go through the + throttle that a blind write's group commit applies, because its timestamp + comes from an observed subscribe frontier rather than from the clock. See the + doc comment on `GroupCommitter::commit_timestamped` for the full list of what + that path skips and why. +- **Zero-row `INSERT ... RETURNING`.** Both paths report `INSERT 0 0` with no + result set when no rows match, because the coordinator decides the response + kind from the evaluated RETURNING rows and there are none. Postgres returns an + empty result set here, with a row description. The frontend path is + deliberately bug-compatible with the coordinator rather than correct on its + own: fixing it changes the behavior of the path that ships today, which is a + separate decision from this change. + ## Performance The goal is not to make writes faster, but to not regress significantly. @@ -295,10 +332,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 +348,33 @@ 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). + +Measured on the full implementation, with the OCC path on for every mzcompose +suite, the small-write regression is at the bad end of that range. Across nightly +runs the feature benchmark `ManySmallUpdates` is 1.7-1.9x slower and `Update` +1.4x slower, and the scalability `UpdateWorkload` loses 36-39% throughput at +concurrency 1 and about 22% at 8 and 32. + +`ManySmallUpdates` also steps `memory_clusterd` up by about 56%, from 56.8 MB to +88.5 MB. Same cause as the wallclock step, from the other side: the subscribe +dataflow each operation installs is arranged on the cluster, where the fast-path +peek it replaces holds nothing. The absolute figures stay small because the +dataflow lives only as long as the operation. + +The performance suites run the OCC path, because that is the configuration we +intend to ship. The write benchmarks therefore record a one-time step, which we +accept for the reasons above. Registering it is a follow-up once the change has +landed and has a commit hash: `ManySmallUpdates` and `Update` go in +`get_ancestor_overrides_for_performance_regressions` and `UpdateWorkload` in +`ANCESTOR_OVERRIDES_FOR_SCALABILITY_REGRESSIONS`, both in +`misc/python/materialize/version_ancestor_overrides.py`. That justification only +applies when the comparison is against a released version, so until the step is +inside the baseline these scenarios report a regression against `main`. + ## Rollout The new path is controlled by a `enable_adapter_frontend_occ_read_then_write` @@ -318,6 +387,12 @@ write locks). We therefore must make the flag sticky per `environmentd` process lifetime (check on bootstrap only) to avoid this, and keep the current `confirm_leadership` checks. +In CI the flag defaults to enabled for versions that carry it, so the mzcompose +suites exercise the OCC path even though production keeps it off. The version +gate leaves it disabled for the older versions an upgrade test runs, and +`CI_SYSTEM_PARAMETERS=random` can pick either value, which is how both paths +stay covered. + Once the OCC path is fully rolled out and validated: 1. Remove the old `sequence_read_then_write` code path diff --git a/doc/user/data/metrics.yml b/doc/user/data/metrics.yml index caea74b2909ab..20ad02f1d0d48 100644 --- a/doc/user/data/metrics.yml +++ b/doc/user/data/metrics.yml @@ -1472,6 +1472,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/checks/all_checks/read_then_write.py b/misc/python/materialize/checks/all_checks/read_then_write.py new file mode 100644 index 0000000000000..e09c3a452c94c --- /dev/null +++ b/misc/python/materialize/checks/all_checks/read_then_write.py @@ -0,0 +1,66 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# 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. +from textwrap import dedent + +from materialize.checks.actions import Testdrive +from materialize.checks.checks import Check + + +class ReadThenWriteForeignRead(Check): + """DELETE, UPDATE and INSERT ... SELECT whose selection reads objects other + than the write target. + + Which of the two sequencing paths runs a read-then-write is decided once per + process at startup, so the phases of a restart or upgrade scenario can take + different ones, and the final state has to be the same either way. Reading a + second table and a materialized view is what gives the mutation read + dependencies its write target does not have, the shape that puts a foreign + collection under a statement that writes elsewhere.""" + + def initialize(self) -> Testdrive: + return Testdrive(dedent(""" + > CREATE TABLE rtw_target (key INTEGER, val INTEGER); + > INSERT INTO rtw_target SELECT generate_series, 0 FROM generate_series(1, 1000); + + > CREATE TABLE rtw_filter (key INTEGER); + > INSERT INTO rtw_filter SELECT generate_series * 4 FROM generate_series(1, 250); + + > CREATE MATERIALIZED VIEW rtw_filter_mv AS SELECT key FROM rtw_filter WHERE key % 8 = 0; + """)) + + def manipulate(self) -> list[Testdrive]: + return [ + Testdrive(dedent(s)) + for s in [ + """ + > UPDATE rtw_target SET val = val + 1 WHERE key IN (SELECT key FROM rtw_filter_mv); + + > DELETE FROM rtw_target WHERE key IN (SELECT key FROM rtw_filter WHERE key % 8 = 4); + """, + """ + > UPDATE rtw_target SET val = val + 1 WHERE key IN (SELECT key FROM rtw_filter_mv); + + > INSERT INTO rtw_target SELECT key, 9 FROM rtw_filter_mv; + """, + ] + ] + + def validate(self) -> Testdrive: + # The 250 multiples of 4 split evenly: the odd multiples were deleted, + # the even ones (the materialized view's rows) were updated twice and + # then inserted a second time at val 9. + return Testdrive(dedent(""" + > SELECT val, count(*), count(DISTINCT key) FROM rtw_target GROUP BY val ORDER BY val; + 0 750 750 + 2 125 125 + 9 125 125 + + > SELECT count(*), min(key), max(key) FROM rtw_target; + 1000 1 1000 + """)) diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 9fb80b0a1b169..272bbd24deb86 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" if version >= MzVersion.parse_mz("v26.36.0-dev") else "false", + ["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 c585e9bccb09a..56265c4c73886 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -78,6 +78,7 @@ MAX_TYPES, MAX_VIEWS, MAX_WEBHOOK_SOURCES, + OCC_CONTENTION_EXHAUSTED_ERROR, Cluster, ClusterReplica, Column, @@ -1096,7 +1097,6 @@ def run(self, exe: Executor) -> bool: ) all_column_values = ", ".join(f"({v})" for v in column_values) query = f"INSERT INTO {table} ({column_names}) VALUES {all_column_values}" - # TODO: Use INSERT INTO {} SELECT {} (only works for tables) if self.rng.choice([True, False]): self.stmt_id += 1 self.exe_prepared(query, f"insert{self.stmt_id}", exe) @@ -1107,6 +1107,70 @@ def run(self, exe: Executor) -> bool: return True +def readable_by_read_then_write(obj: DBObject) -> bool: + """Whether a read-then-write's selection may read `obj`. + + Only user tables and the views over them qualify. A selection that reads a + source, or a view that transitively does, is refused with + `InvalidTableMutationSelection`, because such a collection's notion of time + moves differently than the table being written.""" + if isinstance(obj, Table): + return True + if isinstance(obj, View): + return readable_by_read_then_write(obj.base_object) and ( + obj.base_object2 is None or readable_by_read_then_write(obj.base_object2) + ) + return False + + +def foreign_read_objects(exe: Executor, target: DBObject) -> list[DBObject]: + """Objects a read-then-write against `target` may read besides `target` + itself. + + A selection that only touches its own write target gives the mutation a + read dependency no other session can drop, and one that is always the table + it writes. Reading a foreign object instead is what puts a view or + materialized view under the mutation, what makes the dependency traversal + walk more than one object, and what lets a concurrent DROP land on a + collection the statement is reading. + + Another session's temporary objects are excluded: they are invisible from + here, so naming one yields a catalog error instead of a read dependency.""" + return [ + obj + for obj in exe.db.db_objects() + if str(obj) != str(target) + and readable_by_read_then_write(obj) + and not ( + isinstance(obj, Table | View) and obj.temp and obj not in exe.temp_objects + ) + ] + + +def foreign_read_predicate( + rng: random.Random, exe: Executor, table: Table +) -> str | None: + """A `col IN (SELECT col FROM other)` predicate for an UPDATE or DELETE on + `table`, or None when no object offers a type-compatible column. + + The subquery is what makes the statement read something other than its + write target, see `foreign_read_objects`.""" + objects = foreign_read_objects(exe, table) + rng.shuffle(objects) + for obj in objects: + pairs = [ + (column, foreign_column) + for column in table.columns + for foreign_column in obj.columns + if column.data_type == foreign_column.data_type + and column.data_type != TextTextMap + ] + if pairs: + column, foreign_column = rng.choice(pairs) + return f"{column.name(True)} IN (SELECT {foreign_column} FROM {obj})" + return None + + class InsertSelectAction(Action): """INSERT INTO ... SELECT ... FROM ... WHERE ..., the read-dependent INSERT. @@ -1119,6 +1183,7 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: result.extend( [ "canceling statement due to statement timeout", + OCC_CONTENTION_EXHAUSTED_ERROR, # A random expression can evaluate to NULL (e.g. a map-key # miss) even for a NOT NULL column, which is a legitimate # rejection. The base list only ignores it for DDL complexity. @@ -1145,8 +1210,13 @@ def run(self, exe: Executor) -> bool: return False table = self.rng.choice(tables) # Reading the insert target itself makes the target a read dependency - # too, the most contended shape a read-then-write can have. - source = table if self.rng.choice([True, False]) else self.rng.choice(tables) + # too, the most contended shape a read-then-write can have. The other + # half of the time the source is any readable object, see + # `foreign_read_objects`. + source: DBObject = table + if self.rng.choice([True, False]): + objects = foreign_read_objects(exe, table) + source = self.rng.choice(objects) if objects else table column_names = ", ".join(column.name(True) for column in table.columns) # The cast is an identity cast: `expression` returns the requested type @@ -1238,6 +1308,9 @@ def run(self, exe: Executor) -> bool: class InsertReturningAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: result = super().errors_to_ignore(exe) + # A constant INSERT is a blind write, but RETURNING takes it off that + # fast path and makes it a read-then-write. + result.append(OCC_CONTENTION_EXHAUSTED_ERROR) # The RETURNING expressions re-render the fully-qualified table and # column names, so a concurrent schema or table rename landing between # the INSERT target and the RETURNING clause leaves the two referring to @@ -1287,7 +1360,6 @@ def run(self, exe: Executor) -> bool: ) all_column_values = ", ".join(f"({v})" for v in column_values) query = f"INSERT INTO {table} ({column_names}) VALUES {all_column_values}" - # TODO: Use INSERT INTO {} SELECT {} (only works for tables) returning_exprs = [] if self.rng.random() < 0.5: returning_exprs += [ @@ -1391,6 +1463,7 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: result.extend( [ "canceling statement due to statement timeout", + OCC_CONTENTION_EXHAUSTED_ERROR, # A random SET expression can evaluate to NULL (e.g. a map-key # miss) even for a NOT NULL column. That is a legitimate # rejection, not a bug, and the column type can't be coerced @@ -1434,7 +1507,14 @@ def run(self, exe: Executor) -> bool: f"{c.name(True)} = {expression(c.data_type, table.columns, self.rng, kind=ExprKind.WRITE)}" for c in set_columns ) - query = f"UPDATE {table} SET {set_clause} WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" + predicate = None + if self.rng.random() < 0.3: + predicate = foreign_read_predicate(self.rng, exe, table) + if predicate is None: + predicate = expression( + Boolean, table.columns, self.rng, kind=ExprKind.WRITE + ) + query = f"UPDATE {table} SET {set_clause} WHERE {predicate}" if self.rng.choice([True, False]): self.stmt_id += 1 self.exe_prepared(query, f"update{self.stmt_id}", exe) @@ -1455,6 +1535,9 @@ class ReadThenWriteCounterUpdateAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: return [ "canceling statement due to statement timeout", + # Extreme contention on one row is what this action creates, so + # exhausting the retry budget is an expected outcome here. + OCC_CONTENTION_EXHAUSTED_ERROR, ] + super().errors_to_ignore(exe) def run(self, exe: Executor) -> bool: @@ -1477,12 +1560,113 @@ def run(self, exe: Executor) -> bool: return True +class ReadThenWriteNoRowsAction(Action): + """Run a deterministic read-then-write whose selection is empty.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "canceling statement due to statement timeout", + OCC_CONTENTION_EXHAUSTED_ERROR, + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + counter = exe.db.read_then_write_counter + exe.execute(f"UPDATE {counter} SET v = v + 1 WHERE id = 0", http=Http.NO) + return True + + +class BlindWriteTransactionAction(Action): + """Run nonconstant blind writes through transaction outcomes. + + Transactional buffering is not available in every server version that the + workload can run against. A server may reject the first INSERT. When it is + supported, one action covers commit, rollback, an abort after the write was + buffered, mixing buffered and constant writes, and RETURNING rejection. + """ + + # The optimizer folds constant dataflows with at most 10,000 rows. Staying + # just above that bound makes the INSERT use the frontend read-then-write + # path without generating more data than the coverage shape needs. + ROWS = 10_001 + UNSUPPORTED = "cannot be run inside a transaction block" + + def _rollback_after_error(self, exe: Executor) -> None: + try: + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + exe.reconnect_next = True + + def run(self, exe: Executor) -> bool: + table = exe.db.blind_write_table + value = self.rng.randint(-(2**63), 2**63 - 1) + insert = ( + f"INSERT INTO {table} " + f"SELECT {value} FROM generate_series(1, {self.ROWS})" + ) + + # Two write shapes in one transaction exercise merging the deferred + # dataflow result with an ordinary constant write. + exe.execute("BEGIN", http=Http.NO) + try: + exe.execute(insert, http=Http.NO) + exe.execute(f"INSERT INTO {table} VALUES ({value})", http=Http.NO) + exe.execute("COMMIT", http=Http.NO) + except QueryError as e: + self._rollback_after_error(exe) + if self.UNSUPPORTED in e.msg: + return True + raise + + # Remove the committed value before running variants that must leave no + # rows behind. The unique value keeps concurrent action instances from + # deleting each other's rows. + exe.execute(f"DELETE FROM {table} WHERE id = {value}", http=Http.NO) + + exe.execute("BEGIN", http=Http.NO) + try: + exe.execute(insert, http=Http.NO) + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + self._rollback_after_error(exe) + raise + + # In an implicit multi-statement transaction, a later error must abort + # the buffered write with the rest of the batch. + try: + exe.execute(f"{insert}; SELECT 1 / 0", http=Http.NO) + except QueryError as e: + self._rollback_after_error(exe) + if "division by zero" not in e.msg: + raise + else: + raise RuntimeError("blind write batch unexpectedly succeeded") + + # RETURNING needs rows before COMMIT, so deferred writes cannot provide + # it and the statement must be rejected. + exe.execute("BEGIN", http=Http.NO) + try: + exe.execute(f"{insert} RETURNING id", http=Http.NO) + except QueryError as e: + self._rollback_after_error(exe) + if self.UNSUPPORTED not in e.msg: + raise + else: + self._rollback_after_error(exe) + raise RuntimeError("blind write with RETURNING unexpectedly succeeded") + + return True + + class DeleteAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: errors = [ "canceling statement due to statement timeout", + OCC_CONTENTION_EXHAUSTED_ERROR, ] + super().errors_to_ignore(exe) - if exe.db.scenario == Scenario.Rename: + # The selection can name a foreign object, which another session may + # drop between planning and execution. Same tolerance the other + # read-then-write actions carry for the same reason. + if exe.db.complexity == Complexity.DDL or exe.db.scenario == Scenario.Rename: errors += ["does not exist"] return errors @@ -1519,7 +1703,14 @@ def run(self, exe: Executor) -> bool: query += f" USING {using_table}" query += f" WHERE {expression(Boolean, all_columns, self.rng, kind=ExprKind.WRITE)}" elif self.rng.random() < 0.95: - query += f" WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" + predicate = None + if self.rng.random() < 0.3: + predicate = foreign_read_predicate(self.rng, exe, table) + if predicate is None: + predicate = expression( + Boolean, table.columns, self.rng, kind=ExprKind.WRITE + ) + query += f" WHERE {predicate}" if self.rng.choice([True, False]): self.stmt_id += 1 self.exe_prepared(query, f"delete{self.stmt_id}", exe) @@ -3054,6 +3245,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", @@ -3395,6 +3590,46 @@ def reset_flag(self, conn: Connection, flag_name: str) -> None: ) +class StartupOnlySystemVarsAction(Action): + """Set and restore the system parameters sampled only at startup.""" + + PARAMS = ( + ("enable_adapter_frontend_occ_read_then_write", "true"), + ("max_concurrent_occ_writes", "1"), + ) + + def applicable(self, exe: Executor) -> bool: + # A process restart between SET and RESET could boot with the temporary + # value. The regression scenario keeps the process alive throughout. + return exe.db.scenario == Scenario.Regression + + def run(self, exe: Executor) -> bool: + conn = None + pending_reset = None + try: + conn = self.create_system_connection(exe) + with conn.cursor() as cur: + for name, value in self.PARAMS: + cur.execute(f"ALTER SYSTEM SET {name} = {value};".encode()) + pending_reset = name + cur.execute(f"ALTER SYSTEM RESET {name};".encode()) + pending_reset = None + return True + except OperationalError: + return False + except Exception as e: + raise QueryError(str(e), "StartupOnlySystemVars") + finally: + if conn is not None: + if pending_reset is not None: + try: + with conn.cursor() as cur: + cur.execute(f"ALTER SYSTEM RESET {pending_reset};".encode()) + except Exception: + pass + conn.close() + + class CreateViewAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: errors = super().errors_to_ignore(exe) @@ -6297,6 +6532,8 @@ def __init__( # this list, often enough that every worker on it contends on the one # counter row, not often enough to starve the rest of the workload. (ReadThenWriteCounterUpdateAction, 10), + (ReadThenWriteNoRowsAction, 5), + (BlindWriteTransactionAction, 2), # COPY FROM is oneshot ingestion, it can't run inside a transaction (CopyFromS3Action, 10), (CommentAction, 5), @@ -6409,6 +6646,7 @@ def __init__( (ExplainAnalyzeAction, 4), (ExplainFilterPushdownAction, 2), (FlipFlagsAction, 2), + (StartupOnlySystemVarsAction, 1), # TODO: Reenable when https://linear.app/materializeinc/issue/SQL-405 is fixed. # (AlterTableAddColumnAction, 10), (AlterIcebergSinkFromAction, 8), diff --git a/misc/python/materialize/parallel_workload/database.py b/misc/python/materialize/parallel_workload/database.py index 49ccfbabfe48e..de67567d916ab 100644 --- a/misc/python/materialize/parallel_workload/database.py +++ b/misc/python/materialize/parallel_workload/database.py @@ -1101,12 +1101,26 @@ def __str__(self) -> str: # drops that database or schema, so the name resolves for a whole run and the # end-of-run check is guaranteed to find the table. READ_THEN_WRITE_COUNTER_NAME = "materialize.public.pw_rtw_counter" +BLIND_WRITE_TABLE_NAME = "materialize.public.pw_blind_write" + +# The frontend read-then-write path gives up after its OCC retry budget when a +# statement keeps losing the race for the write timestamp. That is a +# user-visible consequence of contention, not a bug, so every action whose +# statement is a read-then-write (DELETE, UPDATE, INSERT ... SELECT, +# INSERT ... RETURNING) has to tolerate it. It is still counted in the error +# statistics. +OCC_CONTENTION_EXHAUSTED_ERROR = ( + "read-then-write exceeded maximum retry attempts under contention" +) # Error texts that prove an increment did not land. # -# A concurrently modified dependency is what the coordinator reports when it -# revalidates a plan before sequencing it, which is before any write. The other -# is a cluster-resolution failure during planning, which the workload provokes +# An exhausted retry budget is checked right after an attempt the group +# committer rejected, so nothing was appended. A concurrently modified +# dependency is what the coordinator reports when it revalidates a plan before +# sequencing it, and what the frontend path reports for a changed write target, +# both before any write. The last is a cluster-resolution failure during +# planning, which the workload provokes # on purpose by pointing the default cluster at a nonexistent one, so the # statement never reaches a write path at all. The trailing quote keeps it from # matching "unknown cluster replica size" errors. @@ -1115,6 +1129,7 @@ def __str__(self) -> str: # included, because either can race a commit that did happen. A wrong entry here # makes healthy runs fail, an unnecessary unknown only widens the upper bound. DEFINITELY_NOT_COMMITTED_ERRORS = ( + OCC_CONTENTION_EXHAUSTED_ERROR, "was concurrently modified", "unknown cluster '", ) @@ -1236,6 +1251,26 @@ def validate(self, exe: Executor) -> None: ) +class BlindWriteTable: + """A stable target for nonconstant blind writes in transactions. + + The table is not registered in the workload's object lists, so random DDL + cannot rename, alter, or drop it while a transaction is open. Its rows are + disposable and each action removes the value it inserted. + """ + + def __str__(self) -> str: + return BLIND_WRITE_TABLE_NAME + + def create(self, exe: Executor) -> None: + """Creates the table once per run after PUBLIC table grants exist.""" + exe.execute(f"DROP TABLE IF EXISTS {self} CASCADE") + exe.execute(f"CREATE TABLE {self} (id bigint)") + + def drop(self, exe: Executor) -> None: + exe.execute(f"DROP TABLE IF EXISTS {self} CASCADE") + + class Index: _name: str schema: Schema @@ -1477,6 +1512,7 @@ class Database: s3_path: int s3_objects: list[S3Object] read_then_write_counter: ReadThenWriteCounter + blind_write_table: BlindWriteTable lock: threading.Lock seed: str sqlsmith_state: str @@ -1598,6 +1634,7 @@ def __init__( self.iceberg_sink_id = len(self.iceberg_sinks) self.kafka_sink_id = len(self.kafka_sinks) self.read_then_write_counter = ReadThenWriteCounter() + self.blind_write_table = BlindWriteTable() self.types = [] self.type_id = 0 self.network_policies = [] @@ -1825,6 +1862,7 @@ def create(self, exe: Executor, composition: Composition) -> None: # Created and seeded exactly once per run: re-seeding mid-run would reset # `v` while the workers' tallies keep growing. self.read_then_write_counter.create(exe) + self.blind_write_table.create(exe) print("Creating relations") @@ -1846,8 +1884,9 @@ def create(self, exe: Executor, composition: Composition) -> None: # self.sqlsmith_state = result.stdout def drop(self, exe: Executor) -> None: - # The counter table lives outside the workload's databases, so dropping - # them does not reclaim it. + # The helper tables live outside the workload's databases, so dropping + # them does not reclaim either table. + self.blind_write_table.drop(exe) self.read_then_write_counter.drop(exe) for db in self.dbs: 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 afe4f8f1729c7..1ba160ae3dd5b 100644 --- a/src/adapter/src/catalog/open.rs +++ b/src/adapter/src/catalog/open.rs @@ -324,8 +324,9 @@ impl Catalog { // `SystemConfiguration` values applied via `pre_item_updates` live in // the `SystemVars` value map by now. Mirror their effective values into // the dyncfg `ConfigSet`, so that startup-only reads observe configured - // values rather than compile-time defaults, `ENABLE_EXPRESSION_CACHE` - // just below being one of them. `apply_updates` only + // values rather than compile-time defaults. Those reads are the + // `ENABLE_EXPRESSION_CACHE` read just below and + // `FRONTEND_READ_THEN_WRITE` in coord bootstrap. `apply_updates` only // syncs the `ConfigSet` when a `SystemConfiguration` update is present, // and a deployment configured purely via `system_parameter_default` has // none, so sync explicitly here. diff --git a/src/adapter/src/client.rs b/src/adapter/src/client.rs index b6c6298969c47..331d6f4581f9b 100644 --- a/src/adapter/src/client.rs +++ b/src/adapter/src/client.rs @@ -25,6 +25,7 @@ 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,10 +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::Metrics; +use crate::optimize::dataflows::{EvalTime, ExprPrepOneShot}; +use crate::optimize::{self, Optimize, OptimizerError}; use crate::peek_client::{ExecutionLogging, TakeOver}; use crate::session::{ EndTransactionAction, PreparedStatement, Session, SessionConfig, StateRevision, TransactionId, + TransactionStatus, }; use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy}; use crate::telemetry::{self, EventDetails, SegmentClientExt, StatementFailureType}; @@ -294,6 +299,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( @@ -304,6 +312,9 @@ impl Client { optimizer_metrics, persist_client, statement_logging_frontend, + occ_write_semaphore, + frontend_read_then_write_enabled, + read_only, ); let mut client = SessionClient { @@ -841,7 +852,17 @@ impl SessionClient { debug!("frontend peek succeeded"); return Ok(resp); } - 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, logging, cancel_future.clone()) + .await?; + if let Some(resp) = rtw_result { + debug!("frontend read-then-write succeeded"); + return Ok(resp); + } + debug!("frontend read-then-write did not happen, falling back to `Command::Execute`"); // No frontend path took the statement over, so the coordinator retires // whatever entry we hold, or begins its own if we hold none. @@ -1343,7 +1364,11 @@ impl SessionClient { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => {} + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => {} }; cmd }); @@ -1434,6 +1459,537 @@ impl SessionClient { Ok(None) } } + + /// Whether the frontend read-then-write path could take this portal over. + /// + /// The gate is deliberately cheap, a flag read and a portal lookup, because + /// every statement that reaches `execute_attempts` without being handled by + /// the peek path is tested against it. Everything expensive, including the + /// coordinator round-trip that registers the connection cancel watch, sits + /// behind it. + fn frontend_read_then_write_applies(&self, portal_name: &str) -> bool { + if !self.peek_client.frontend_read_then_write_enabled { + return false; + } + let session = self.session.as_ref().expect("SessionClient invariant"); + match session.get_portal_unverified(portal_name) { + Some(portal) => portal + .stmt + .as_deref() + .is_some_and(is_read_then_write_statement), + None => false, + } + } + + /// 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, + logging: &mut ExecutionLogging, + cancel_future: impl Future + Send, + ) -> Result, AdapterError> { + // Bail out before the cancel-watch registration below, which is a + // synchronous round-trip through the coordinator's command loop. A + // statement this path will not take over must not pay for it, and must + // not add queueing latency for other sessions either. + if !self.frontend_read_then_write_applies(portal_name) { + return Ok(None); + } + + let conn_id = self.session().conn_id().clone(); + let statement_timeout = *self.session().vars().statement_timeout(); + let inner_client = self.inner().clone(); + let attempt_state = Arc::new(FrontendWriteAttemptState::new()); + + 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); + + // Registering installs a fresh channel, so this cannot observe a + // cancellation aimed at an earlier statement. The entry it leaves behind + // is replaced by the next registration and removed when a statement + // reaches the coordinator or the connection's state is cleared, so there + // is nothing to unregister here. + let mut connection_cancel_rx = { + let register = + self.peek_client + .call_coordinator(|tx| Command::RegisterConnectionCancelWatch { + conn_id: conn_id.clone(), + 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, logging, 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(requested.into()); + } + + // 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, + logging: &mut ExecutionLogging, + 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}; + + // Re-checked here rather than relying on the caller's gate. See the + // module-level docs on `frontend_read_then_write` for why the flag is + // fixed for the lifetime of the process. + 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 is_read_then_write_statement(&stmt) => stmt, + Some(_stmt) => { + return Ok(None); + } + None => { + return Ok(None); + } + }; + + // Verify and plan against one catalog snapshot. Pairing a stale plan + // with a newer target generation could direct a write incorrectly. + // A failed verification is not logged, mirroring the coordinator: the + // portal is what statement logging draws its record from. + Coordinator::verify_portal( + &catalog, + self.session.as_mut().expect("SessionClient invariant"), + portal_name, + )?; + + let (params, logging_info, 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(), + ) + }; + + // Past this point the coordinator never sees this statement, so every + // exit has to produce an outcome for it. The remaining `Ok(None)` + // bailouts are all above. + logging.take_over( + &self.peek_client, + self.session.as_mut().expect("SessionClient invariant"), + Some(&stmt), + ¶ms, + &logging_info, + &catalog, + lifecycle_timestamps, + TakeOver::StatementToRun, + ); + + // 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)); + } + } + } + } + + 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) + }; + + // Reject mutations in read-only mode (e.g. during 0dt upgrades). Placed + // where the coordinator has it, in `sequence_plan`: after planning, so a + // statement that does not plan reports the planning error, and before + // the cluster and RBAC checks below, which the coordinator also reports + // second. Every sub-path from here on writes (constant INSERT and the + // OCC INSERT/UPDATE/DELETE), so one check covers them all. + if self.peek_client.read_only { + return Err(AdapterError::ReadOnly); + } + + // 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), + } + }; + + // Record the cluster before the checks below can fail, so that their + // error rows carry it, as the coordinator's do. + if let (Some(logging_id), Some(cluster_id), Some(cluster_name)) = + (logging.id(), target_cluster_id, target_cluster_name.clone()) + { + self.peek_client + .log_set_cluster(logging_id, cluster_id, cluster_name); + } + + { + 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, and is a 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; + } + } + + // The coordinator's per-plan checks, in the order it applies them: + // `sequence_insert` rejects a transaction that cannot take a write + // before it rejects the isolation level, and both it and + // `sequence_read_then_write` reject bounded staleness before dispatching + // on the plan. So both checks sit here, above the constant-INSERT + // dispatch as well as the read-then-write path. + // + // `allows_writes` is only defined inside a transaction, which is also + // the only place it can be false: outside one the session task opens a + // fresh transaction with no ops. Autocommit statements therefore rely on + // the check in `PeekClient::frontend_read_then_write` instead. + { + let session = self.session.as_ref().expect("SessionClient invariant"); + if session.transaction().is_in_multi_statement_transaction() + && !session.transaction().allows_writes() + { + return Err(AdapterError::ReadOnlyTransaction); + } + if session + .vars() + .transaction_isolation() + .is_bounded_staleness() + { + return Err(AdapterError::BoundedStalenessReadOnly); + } + } + + // Handle ReadThenWrite plans or Insert plans. + let rtw_plan = match plan { + Plan::ReadThenWrite(rtw_plan) => rtw_plan, + Plan::Insert(insert_plan) => { + // A constant INSERT without RETURNING is a blind write, handled + // here through the coordinator's `insert_constant` helper, which + // buffers the rows as session write ops. + // + // Deciding that needs HIR lowered to MIR, because a VALUES list + // is planned as a `Wrap` call at the HIR level. + // + // Only take that 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 the 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()), + } + }; + + 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 result = Coordinator::insert_constant( + &catalog, + session, + insert_plan.id, + inner_mir, + ); + + 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(), + )); + } + }; + + // Only single-statement (`Started`) transactions may enter the OCC + // loop, its writes commit immediately and cannot be rolled back at + // transaction end. Multi-statement transactions reach this point only + // for AST-constant INSERTs whose planned expression turned out + // non-constant. Match the coordinator's error precedence: `mz_now()` + // gets its dedicated error, everything else is prohibited in a + // transaction block. The coordinator's lock-based path additionally + // supports INSERTs of volatile constants (for example `random()`) in + // transaction blocks by buffering the diffs until commit, which the + // OCC path cannot do. + { + let session = self.session.as_ref().expect("SessionClient invariant"); + if !matches!(session.transaction(), TransactionStatus::Started(_)) { + let contains_temporal = rtw_plan.selection.contains_temporal() + || rtw_plan.assignments.values().any(|e| e.contains_temporal()) + || rtw_plan.returning.iter().any(|e| e.contains_temporal()); + if contains_temporal { + return Err(AdapterError::Unsupported( + "calls to mz_now in write statements", + )); + } + 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, + &catalog, + logging.id(), + attempt_state, + ) + .await + .map(Some) + } +} + +/// Whether a statement is one the frontend read-then-write path sequences. +/// +/// These are the statement kinds that plan to a `ReadThenWrite` or an `Insert`. +/// Note that not every one of them ends up on the OCC path: an `INSERT` whose +/// source folds to a constant is dispatched as a blind write instead. +fn is_read_then_write_statement(stmt: &Statement) -> bool { + matches!( + stmt, + Statement::Delete(_) | Statement::Update(_) | Statement::Insert(_) + ) +} + +/// 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 1665e1035491c..a410e2ec423b8 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, @@ -395,6 +396,62 @@ 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. + /// + /// Registration always installs a fresh channel, so the caller cannot + /// observe a cancellation aimed at an earlier statement. + RegisterConnectionCancelWatch { + conn_id: ConnectionId, + tx: oneshot::Sender>, + }, + + /// Creates an internal subscribe (not visible in introspection) and returns + /// the response channel. Used by 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)`: the write must land at exactly `ts`, and reports + /// `WriteResult::TimestampPassed` if the table's timestamp is already past + /// it. The caller decides whether to recompute the diffs and try again. + /// - `None`: the coordinator picks the timestamp from the oracle during + /// group commit, so the timestamp cannot be passed. Every other outcome, + /// including read-only, a changed target and cancellation, is reported the + /// same way in both modes. + 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. Fire-and-forget, the caller does not wait + /// for completion. + DropInternalSubscribe { + sink_id: GlobalId, + }, } impl Command { @@ -435,7 +492,11 @@ impl Command { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => None, + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => None, } } @@ -476,7 +537,11 @@ impl Command { | Command::UnregisterFrontendPeek { .. } | Command::ExplainTimestamp { .. } | Command::FrontendStatementLogging(..) - | Command::InjectAuditEvents { .. } => None, + | Command::InjectAuditEvents { .. } + | Command::RegisterConnectionCancelWatch { .. } + | Command::CreateInternalSubscribe { .. } + | Command::AttemptWrite { .. } + | Command::DropInternalSubscribe { .. } => None, } } } @@ -514,6 +579,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 66650756c64ef..412d4f5590ad8 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; @@ -509,6 +509,10 @@ 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::CreateInternalSubscribe { .. } => "create-internal-subscribe", + Command::AttemptWrite { .. } => "attempt-write", + Command::DropInternalSubscribe { .. } => "drop-internal-subscribe", }, Message::ControllerReady { controller: ControllerReadiness::Compute, @@ -2062,8 +2066,11 @@ pub struct Coordinator { /// Each entry is a watch channel whose value is `false` until cancellation /// 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. + /// Consumers install these watches while they have cancellable work in + /// flight, always as a fresh channel, so nobody can observe a cancellation + /// aimed at an earlier statement. An entry is removed when a statement + /// starts, when a stage runs uncancelable, and when the connection's state + /// is cleared. connection_cancel_watches: BTreeMap, watch::Receiver)>, /// Active introspection subscribes. introspection_subscribes: BTreeMap, @@ -2076,6 +2083,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 @@ -3968,8 +3993,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(), @@ -5079,6 +5104,14 @@ pub fn serve( } let catalog = Arc::new(catalog); + // Both are read once at startup, see the field docs on + // `occ_write_semaphore` and `frontend_read_then_write_enabled`. + 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(); @@ -5107,6 +5140,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 732e50fd1ff3d..9e7183379be58 100644 --- a/src/adapter/src/coord/appends.rs +++ b/src/adapter/src/coord/appends.rs @@ -165,9 +165,6 @@ pub(crate) enum BuiltinTableUpdateSource { } /// Result of a write submitted by frontend sequencing. -// The read-then-write path that submits these writes lands later in this -// stack, this attribute goes away with it. -#[allow(dead_code)] #[derive(Debug, Clone)] pub enum WriteResult { /// The write committed at this timestamp. @@ -200,9 +197,6 @@ pub struct InternalWriteResponder { } impl InternalWriteResponder { - // The read-then-write path that uses this lands later in this stack, this - // attribute goes away with it. - #[allow(dead_code)] pub(crate) fn new(tx: oneshot::Sender) -> Self { Self { tx: Some(tx) } } @@ -229,9 +223,6 @@ pub(crate) enum UserWriteResponder { /// `ExecuteContext` once the write commits. Session(PendingTxn), /// Frontend-sequenced blind write. - // The read-then-write path that uses this lands later in this stack, this - // attribute goes away with it. - #[allow(dead_code)] Internal { conn_id: ConnectionId, /// The table the diffs were computed against, item id and the generation @@ -295,9 +286,6 @@ impl PendingWriteTxn { pub(crate) enum TableWriteCmd { GroupCommit(GroupCommitRequest), - // The read-then-write path that uses this lands later in this stack, this - // attribute goes away with it. - #[allow(dead_code)] TimestampedWrite(TimestampedWriteRequest), Register { tables: Vec, @@ -504,6 +492,13 @@ impl GroupCommitter { } } + // The write is durable and applied to the oracle, and the session task + // that submitted it has not been told yet. Only a timestamped write + // reaches here, so arming this holds one read-then-write's result + // without stalling the keepalives that advance table uppers. Used by + // workflow_test_occ_cancel_of_submitted_write. + fail::fail_point!("timestamped_write_before_result"); + if self .internal_cmd_tx .send(Message::GroupCommitApplied { @@ -605,6 +600,16 @@ impl GroupCommitter { let now: Timestamp = (self.now)().into(); crate::coord::timeline::check_runaway_write_ts(&now, write_ts.timestamp); + // The append above is already readable in Persist and has advanced the + // table's upper, while no oracle-timestamped read can reach it until + // the line below. Anything concluding from a read that follows Persist + // rather than the oracle has to cope with this window, so a test can + // hold it open here. Every txns-shard write parks here while armed, + // including the keepalives that advance table uppers, so arm it with a + // bounded `sleep` rather than a `pause`. Used by + // workflow_test_occ_zero_row_write_linearization. + fail::fail_point!("group_commit_before_apply_write"); + self.oracle.apply_write(write_ts.timestamp).await; TxnsWriteAttempt::Applied diff --git a/src/adapter/src/coord/command_handler.rs b/src/adapter/src/coord/command_handler.rs index 45f582cf42cee..cb20bd0278a16 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, @@ -617,6 +617,69 @@ impl Coordinator { Command::FrontendStatementLogging(event) => { self.handle_frontend_statement_logging_event(event); } + Command::RegisterConnectionCancelWatch { conn_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, (watch_tx, watch_rx.clone())); + let _ = tx.send(watch_rx); + } + 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")) @@ -905,6 +968,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 @@ -1932,20 +1998,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"), } } diff --git a/src/adapter/src/coord/read_then_write.rs b/src/adapter/src/coord/read_then_write.rs index 8dc07a6149dae..29a169ca97c65 100644 --- a/src/adapter/src/coord/read_then_write.rs +++ b/src/adapter/src/coord/read_then_write.rs @@ -17,7 +17,7 @@ use std::collections::BTreeSet; use mz_catalog::memory::objects::CatalogItem; use mz_repr::CatalogItemId; -use mz_repr::{GlobalId, Timestamp}; +use mz_repr::{Diff, GlobalId, Row, Timestamp}; use mz_sql::catalog::CatalogItemType; use mz_sql::plan::SubscribeOutput; use tokio::sync::mpsc; @@ -26,6 +26,7 @@ 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 @@ -51,9 +52,6 @@ fn enqueue( Ok(()) } -// Every handler here is driven by the read-then-write path, which lands later -// in this stack. This attribute goes away with it. -#[allow(dead_code)] impl Coordinator { /// Creates a subscribe that introspection does not see. /// @@ -139,6 +137,86 @@ impl Coordinator { drop(read_holds); } + /// Enqueues a write attempt, answering through `result_tx`. + /// + /// `write_ts` picks the path. `Some` names a timestamp the diffs are only + /// valid at and goes straight to the committer, pinned to the `GlobalId` + /// validated here. `None` is a blind write that rides the next group + /// commit, whose staging re-checks the target generation. + 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, WriteTarget, + }; + use mz_storage_client::client::TableData; + use smallvec::smallvec; + use std::collections::BTreeMap; + use tracing::Span; + + let result = InternalWriteResponder::new(result_tx); + 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, + target: WriteTarget { + item_id: target_id, + global_id: target_global_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 diff --git a/src/adapter/src/coord/sequencer.rs b/src/adapter/src/coord/sequencer.rs index c2c22fc996c83..a1dae4166dc9d 100644 --- a/src/adapter/src/coord/sequencer.rs +++ b/src/adapter/src/coord/sequencer.rs @@ -911,6 +911,14 @@ 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 differs from PostgreSQL, where `UPDATE t SET x = x` + // reports the number of rows matching the WHERE clause even when + // no value changes. Because Materialize works in differential + // dataflow, the +1 and -1 diffs for an unchanged row cancel out + // during consolidation, so it reports 0 affected rows. This is + // longstanding behavior and both read-then-write paths agree on + // it. 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..f58a58d1004e8 100644 --- a/src/adapter/src/coord/sequencer/inner.rs +++ b/src/adapter/src/coord/sequencer/inner.rs @@ -2713,6 +2713,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() 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..9eff060952773 --- /dev/null +++ b/src/adapter/src/frontend_read_then_write.rs @@ -0,0 +1,1587 @@ +// 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. +//! +//! ## Whether the write reads persisted state +//! +//! Two predicates answer that one question, and they have to agree. Before +//! anything runs, `SessionClient::try_frontend_read_then_write` decides it +//! syntactically, from `depends_on()` on the planned selection, because inside +//! a transaction a read-dependent write has to be refused while refusing is +//! still possible. Once the dataflow runs, the subscribe answers it +//! dynamically: a subscribe over a persisted collection never ends, so a +//! channel that closes on its own means the selection read nothing. +//! +//! The answer decides where the diffs go. Diffs from a selection that reads +//! persisted state are only correct at the frontier they were observed at, so +//! they commit inside the OCC loop. Diffs from a selection that reads nothing +//! are frontier-independent, so the caller of the loop either submits them +//! right after it or, inside a multi-statement transaction, buffers them as +//! session write ops that land at COMMIT. +//! +//! Disagreement is caught on both sides, and only one side can still refuse. +//! `frontend_read_then_write` re-checks the syntactic predicate before running a +//! dataflow, which catches a caller that skipped the gate. If the syntactic +//! predicate were laxer than the dynamic one, that check would pass and the +//! write would commit mid-transaction, so the loop's `Committed` arm soft-panics +//! when it has a write timestamp to apply inside a transaction. By then the +//! write is durable, so all that arm can do is make the disagreement loud. +//! +//! ## 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::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +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, RelationDesc, Row, RowArena, Timestamp}; +use mz_sql::catalog::CatalogError; +use mz_sql::plan::{self, MutationKind, QueryWhen}; +use mz_sql::session::metadata::SessionMetadata; +use prometheus::Histogram; +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, TargetCluster}; +use crate::error::AdapterError; +use crate::optimize::Optimize; +use crate::optimize::dataflows::{ComputeInstanceSnapshot, EvalTime, ExprPrep, ExprPrepOneShot}; +use crate::session::{Session, TransactionOps}; +use crate::statement_logging::{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, +} + +impl From for AdapterError { + fn from(cancellation: FrontendWriteCancellation) -> Self { + match cancellation { + FrontendWriteCancellation::Canceled => AdapterError::Canceled, + FrontendWriteCancellation::StatementTimeout => AdapterError::StatementTimeout, + } + } +} + +/// 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. +/// +/// The wrapper and the attempt it wraps are polled by the same task, so the +/// mutex and the atomic are here to satisfy `Send`, not to arbitrate between +/// concurrent writers. There is one writer for each field. +pub(crate) struct FrontendWriteAttemptState { + write_submitted: AtomicBool, + /// Set at most once, by the cancellation wrapper. + cancellation: Mutex>, +} + +impl FrontendWriteAttemptState { + pub(crate) fn new() -> Self { + Self { + write_submitted: AtomicBool::new(false), + cancellation: Mutex::new(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) + } + + /// Records why the attempt is being torn down. The first reason recorded + /// is the one the attempt reports. + pub(crate) fn request(&self, cancellation: FrontendWriteCancellation) { + self.cancellation + .lock() + .expect("cancellation lock poisoned") + .get_or_insert(cancellation); + } + + fn requested_error(&self) -> Option { + self.cancellation + .lock() + .expect("cancellation lock poisoned") + .map(AdapterError::from) + } +} + +/// What the OCC loop produced. +enum OccOutcome { + /// The write is durable at `write_ts`. + Committed { + response: ExecuteResponse, + write_ts: Timestamp, + }, + /// The selection was empty, so there was nothing to write. + /// + /// `observed_ts` is the timestamp emptiness was concluded at, and is + /// `None` only when the selection reads no persisted state. When it is + /// `Some`, the caller must linearize against it before responding: the + /// subscribe follows Persist, which runs ahead of the oracle, so the + /// emptiness can be concluded from state no oracle-timestamped read can + /// reach yet. + NoRowsMatched { + response: ExecuteResponse, + observed_ts: Option, + }, + /// Diffs from a selection that reads no persisted state. The subscribe ran + /// to completion, so they are frontier-independent and the caller chooses + /// whether to submit them now or buffer them into the transaction. + Blind { + response: ExecuteResponse, + diffs: Vec<(Row, Diff)>, + }, +} + +/// What the coordinator's answer to a submitted write means for the statement. +enum WriteOutcome { + /// The write is durable at this timestamp. + Committed(Timestamp), + /// The write did not land, and resubmitting these diffs cannot change + /// that. This is the error to report. + Failed(AdapterError), + /// Another writer advanced the target's upper past the timestamp we asked + /// for. The diffs still describe the mutation, so the OCC loop can + /// resubmit them once the subscribe has caught up. + Conflict { next_eligible_timestamp: Timestamp }, +} + +/// Maps a [`WriteResult`] to the outcome the statement reports, or to the one +/// conflict the OCC loop can retry. +fn classify_write_result( + result: WriteResult, + target_id: CatalogItemId, + attempt_state: &FrontendWriteAttemptState, +) -> WriteOutcome { + match result { + WriteResult::Success { timestamp } => WriteOutcome::Committed(timestamp), + WriteResult::TimestampPassed { + next_eligible_timestamp, + .. + } => WriteOutcome::Conflict { + next_eligible_timestamp, + }, + WriteResult::Canceled => WriteOutcome::Failed( + attempt_state + .requested_error() + .unwrap_or(AdapterError::Canceled), + ), + WriteResult::ReadOnly => WriteOutcome::Failed(AdapterError::ReadOnly), + WriteResult::TargetChanged => { + // A concurrent DDL gave the table a new generation after we + // computed these diffs against the old one. The same error the + // coordinator raises when a dependency changes underneath a + // statement, so clients see one retryable outcome for both. + WriteOutcome::Failed(AdapterError::ConcurrentDependencyMutation { + dependency_id: target_id.to_string(), + }) + } + WriteResult::Indeterminate => WriteOutcome::Failed(AdapterError::Internal( + "write outcome is indeterminate because the group committer shut down".into(), + )), + } +} + +/// 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. The caller owns the end-of-execution logging for + /// `statement_logging_id` and verified and planned the portal against + /// `catalog`, which stays in force through optimization and write-target + /// generation capture. + pub(crate) async fn frontend_read_then_write( + &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; + + // 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 on this path commits immediately and cannot be rolled back at + // transaction end, so only a single-statement transaction may reach it. + // The check lives in `SessionClient::try_frontend_read_then_write`, and + // this is defense in depth for it: rejecting here, before we run a + // dataflow, is the last point where refusing is still possible. + if session.transaction().is_in_multi_statement_transaction() { + soft_panic_or_log!("read-then-write reached the OCC path inside a transaction"); + return Err(AdapterError::Internal( + "read-then-write cannot be run inside a transaction block".into(), + )); + } + + // 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 cost of this ordering is that a permit held by a long-running + // operation stalls every read-then-write in the process, including ones + // on unrelated tables, where the coordinator's write lock would only + // stall writes to the target table. We accept that because the + // statement timeout in + // `SessionClient::try_frontend_read_then_write_with_cancel` covers the + // permit wait, so the stall is bounded for everyone but a session that + // disabled its own timeout. + // + // 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(); + + // The coordinator sequences this statement's read as a real peek, so the + // optimizer's notices and the timestamp notice reach the session there. + // Emit both here for the same statement to look the same on either path. + crate::coord::sequencer::emit_optimizer_notices( + &**catalog, + session, + &df_meta.optimizer_notices, + ); + if session.vars().emit_timestamp_notice() { + let conn_id = session.conn_id().clone(); + let session_wall_time = session.pcx().wall_time; + let explanation = self + .call_coordinator(|tx| Command::ExplainTimestamp { + conn_id, + session_wall_time, + cluster_id, + id_bundle: bundle, + determination, + tx, + }) + .await; + session.add_notice(crate::AdapterNotice::QueryTimestamp { explanation }); + } + + 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.clone(), + statement_logging_id, + as_of, + &attempt_state, + ) + .await; + + self.coordinator_client() + .metrics() + .occ_retry_count + .observe(f64::from(u32::try_from(retry_count).unwrap_or(u32::MAX))); + + // Finish the operation, including a blind write's submission, before + // releasing the OCC permit. 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, retrying, or + // waiting for our write to commit. + let response = match result { + Ok(OccOutcome::Committed { response, write_ts }) => { + session.apply_write(write_ts); + Ok(response) + } + Ok(OccOutcome::NoRowsMatched { + response, + observed_ts, + }) => { + // A write would have linearized this for us, because group + // commit advances the oracle before it answers. With nothing + // to write we have to do it ourselves, or we report an empty + // selection from state a later strict-serializable read cannot + // see yet, and that read finds the rows we said were not + // there. + match observed_ts { + Some(observed_ts) => self + .ensure_read_linearized(&timeline, observed_ts) + .await + .map(|()| response), + None => Ok(response), + } + } + Ok(OccOutcome::Blind { response, diffs }) => { + match self + .submit_blind_write( + conn_id, + target_id, + target_global_id, + diffs, + statement_logging_id, + &attempt_state, + ) + .await + { + Ok(write_ts) => { + session.apply_write(write_ts); + Ok(response) + } + Err(err) => Err(err), + } + } + Err(err) => Err(err), + }; + + drop(permit); + + 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, so 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). The mutation is already applied in + // MIR, so we hand the expression to the subscribe optimizer directly + // instead of going through the `SubscribePlan` path, which expects HIR. + // An empty `output` makes the sink emit raw diffs. + let global_mir_plan = optimizer.optimize_query(expr, relation_desc, vec![])?; + + 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 does not 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. + /// + /// A subscribe that ends on its own reads no persisted state, so its diffs + /// are frontier-independent. Those are returned as [`OccOutcome::Blind`] + /// for the caller to submit or buffer, and this never writes them. + /// + /// Read linearization is the caller's responsibility, on both ends. + /// `as_of` must already be linearized (oracle read_ts >= `as_of`) on + /// entry, and an [`OccOutcome::NoRowsMatched`] carrying an `observed_ts` + /// must be linearized against it before the response goes out. 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: &FrontendWriteAttemptState, + ) -> (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), so the diffs do not depend on any + // read frontier. The caller decides where they go, so we + // flatten to `Timestamp::MIN` for `consolidate_updates`. + state.consolidate(Timestamp::MIN); + if state.all_diffs.is_empty() { + break Ok(OccOutcome::NoRowsMatched { + response: build_no_rows_response(&kind), + observed_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(); + + break Ok(OccOutcome::Blind { + response: success_response, + diffs, + }); + } + }; + + 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 { observed_ts } => { + break Some(Ok(OccOutcome::NoRowsMatched { + response: build_no_rows_response(&kind), + observed_ts: Some(observed_ts), + })); + } + 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; + } + + let write_ts = state + .current_upper + .expect("must have seen progress to be ready to write"); + + // Invariant: every diff we are about to write comes from a + // time strictly below `write_ts`. `consolidate` below + // rewrites all diff timestamps to `write_ts`, so a diff from + // at or after `write_ts` would durably record rows that + // reflect state from after the timestamp they were written + // at. + // + // The drain can get ahead of the frontier: a subscribe sends + // a batch's data and the following progress as two separate + // channel messages, so `try_recv` can pick up data from the + // next batch and then see `Empty`, leaving `current_upper` + // at the older progress. When that happens we do not write. + // Waiting for the next progress message re-establishes the + // invariant, and it is bounded by `statement_timeout` like + // every other wait in this loop. + if state + .max_data_ts + .is_some_and(|max_data_ts| max_data_ts >= write_ts) + { + continue; + } + + // 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. + // + // TODO(aljoscha): Store `Arc` in `all_diffs` if this + // shows up in profiles. Every attempt clones every row, and + // we retry up to `max_occ_retries` times (default 1000). + 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 classify_write_result(result, target_id, attempt_state) { + WriteOutcome::Committed(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: timestamp, + }); + } + WriteOutcome::Failed(err) => break Err(err), + WriteOutcome::Conflict { + 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 a conflict we wait for the subscribe to + // progress and retry using that observed frontier. + state.retry_count += 1; + // Cancellation wins over the retry budget: if both + // apply, the user asked us to stop and that is the + // more truthful answer. + 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; + } + } + } + ProcessResult::NoRowsMatched { observed_ts } => { + break Ok(OccOutcome::NoRowsMatched { + response: build_no_rows_response(&kind), + observed_ts: Some(observed_ts), + }); + } + ProcessResult::Error(e) => { + break Err(e); + } + } + }; + + (state.retry_count, result) + } + + /// Submits frontier-independent diffs to group commit, which picks the + /// write timestamp, and returns the timestamp the write committed at. + /// + /// Only valid for diffs that do not depend on an observed read frontier: + /// the write lands at a timestamp this caller does not choose. + async fn submit_blind_write( + &self, + conn_id: mz_adapter_types::connection::ConnectionId, + target_id: CatalogItemId, + target_global_id: GlobalId, + diffs: Vec<(Row, Diff)>, + statement_logging_id: Option, + attempt_state: &FrontendWriteAttemptState, + ) -> Result { + attempt_state.mark_write_submitted(); + let result = self + .call_coordinator(|tx| Command::AttemptWrite { + conn_id, + target_id, + target_global_id, + diffs, + write_ts: None, + tx, + }) + .await; + + // Every outcome here terminates the attempt, so `write_submitted` + // stays set per its contract. + match classify_write_result(result, target_id, attempt_state) { + WriteOutcome::Committed(timestamp) => { + if let Some(id) = statement_logging_id { + self.log_set_timestamp(id, timestamp); + } + Ok(timestamp) + } + WriteOutcome::Failed(err) => Err(err), + WriteOutcome::Conflict { .. } => { + // Unreachable: a write that requests no timestamp cannot have + // one pass. Group commit resolves it through + // `UserWriteResponder::Internal`, which only reports a conflict + // to a write that asked for a specific timestamp. + soft_panic_or_log!("blind read-then-write unexpectedly got TimestampPassed"); + Err(AdapterError::Internal( + "blind write unexpectedly got TimestampPassed".into(), + )) + } + } + } + + /// Builds the response for a write that is about to be submitted. + /// + /// This runs before the write, so the result-size checks in here reject the + /// statement without having written anything. + 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, + /// The largest timestamp among the data rows accumulated since the last + /// [`Self::consolidate`], which is where the diffs' own timestamps are + /// erased. `None` means every accumulated diff is already known to be from + /// before `current_upper`. + max_data_ts: Option, + retry_count: usize, + byte_size: u64, +} + +impl OccState { + fn new() -> Self { + Self { + all_diffs: Vec::new(), + current_upper: None, + max_data_ts: 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. + /// + /// The caller must have established that every diff comes from a time at or + /// before `target_ts`, otherwise the consolidated set claims to describe + /// `target_ts` while reflecting state from after it. + 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(); + self.max_data_ts = None; + } +} + +/// Result of processing a single subscribe message in the OCC loop. +enum ProcessResult { + Continue { + ready_to_write: bool, + }, + /// The consolidated selection is empty as of `observed_ts`. + NoRowsMatched { + observed_ts: Timestamp, + }, + 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:?}" + ))); + } + }; + + 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 { observed_ts: ts }; + } + } 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.max_data_ts = Some(match state.max_data_ts { + Some(max_ts) => std::cmp::max(max_ts, ts), + None => ts, + }); + 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. +/// +/// Bug-compatible with the coordinator path, which evaluates RETURNING over the +/// diffs and so reports a plain row count when there are none. Postgres returns +/// an empty result set for a zero-row `INSERT ... RETURNING` instead, but +/// changing that is a change to the path that ships today, not to this one. +fn build_no_rows_response(kind: &MutationKind) -> ExecuteResponse { + 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, + }; + + 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 71ae2b17016bf..a45e40ea7048d 100644 --- a/src/adapter/src/metrics.rs +++ b/src/adapter/src/metrics.rs @@ -61,6 +61,7 @@ pub struct Metrics { pub catalog_transact_phase_seconds: HistogramVec, pub apply_catalog_implications_seconds: Histogram, pub group_commit_catalog_upper_seconds: Histogram, + pub occ_retry_count: Histogram, } impl Metrics { @@ -302,6 +303,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/optimize/subscribe.rs b/src/adapter/src/optimize/subscribe.rs index 19a082a1610ed..b16e4263743e2 100644 --- a/src/adapter/src/optimize/subscribe.rs +++ b/src/adapter/src/optimize/subscribe.rs @@ -116,6 +116,22 @@ impl Optimizer { self.sink_id } + /// Optimizes a subscribe over an already-lowered MIR expression, for + /// callers that build their own MIR (such as the frontend read-then-write + /// path, which applies the mutation in MIR). + /// + /// `output` is the sink's row ordering, as produced by + /// [`mz_sql::plan::SubscribeOutput::row_order`]. Empty means the sink emits + /// raw diffs. + pub fn optimize_query( + &mut self, + expr: MirRelationExpr, + from_desc: RelationDesc, + output: Vec, + ) -> Result, OptimizerError> { + self.optimize_inner(SubscribeSource::Query { expr, from_desc }, output) + } + /// The single subscribe optimization pipeline. Every subscribe, whatever it /// reads from, goes through here, so a prep or metainfo step added here /// applies to all of them. diff --git a/src/adapter/src/peek_client.rs b/src/adapter/src/peek_client.rs index 62321a9c5d469..8e2eeb6ce8a2b 100644 --- a/src/adapter/src/peek_client.rs +++ b/src/adapter/src/peek_client.rs @@ -35,7 +35,7 @@ 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; @@ -79,6 +79,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 { @@ -94,6 +100,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, @@ -105,6 +114,9 @@ impl PeekClient { statement_logging_frontend, oracles: Default::default(), // lazily populated persist_client, + occ_write_semaphore, + frontend_read_then_write_enabled, + read_only, } } @@ -204,6 +216,11 @@ impl PeekClient { .expect("if the coordinator is still alive, it shouldn't have dropped our call") } + /// The client for sending commands to the coordinator. + 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. /// diff --git a/src/environmentd/src/test_util.rs b/src/environmentd/src/test_util.rs index 578a60c98b19c..86fe89ebcb24c 100644 --- a/src/environmentd/src/test_util.rs +++ b/src/environmentd/src/test_util.rs @@ -40,6 +40,7 @@ use mz_dyncfg::ConfigUpdates; use mz_license_keys::ValidatedLicenseKey; use mz_orchestrator_process::{ProcessOrchestrator, ProcessOrchestratorConfig}; use mz_orchestrator_tracing::{TracingCliArgs, TracingOrchestrator}; +use mz_ore::cast::CastLossy; use mz_ore::metrics::MetricsRegistry; use mz_ore::now::{EpochMillis, NowFn, SYSTEM_TIME}; use mz_ore::retry::Retry; @@ -1911,3 +1912,26 @@ impl Ca { Ok((cert_path, key_path)) } } + +/// Sums the counter series named `name` whose labels include all of `labels`. +/// +/// Returns 0 when no series matches, which is how a labelled counter reads +/// before its first increment. +pub fn get_counter_value(registry: &MetricsRegistry, name: &str, labels: &[(&str, &str)]) -> u64 { + let Some(family) = registry.gather().into_iter().find(|m| m.name() == name) else { + return 0; + }; + family + .get_metric() + .iter() + .filter(|metric| { + labels.iter().all(|(name, value)| { + metric + .get_label() + .iter() + .any(|label| label.name() == *name && label.value() == *value) + }) + }) + .map(|metric| u64::cast_lossy(metric.get_counter().value())) + .sum() +} diff --git a/src/environmentd/tests/read_then_write.rs b/src/environmentd/tests/read_then_write.rs new file mode 100644 index 0000000000000..a762ed30bc65d --- /dev/null +++ b/src/environmentd/tests/read_then_write.rs @@ -0,0 +1,1616 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Integration tests for read-then-write statements: `DELETE`, `UPDATE` and +//! `INSERT ... SELECT`, plus the constant `INSERT` that shares their planning +//! path. +//! +//! Most tests here enable `enable_adapter_frontend_occ_read_then_write` and so +//! cover the frontend OCC path. `test_counts_query_total` runs with the flag +//! both off and on, because the property it checks must hold whichever path +//! sequenced the statement. `test_cancel_read_then_write` covers the +//! coordinator path only, and is the other half of the cancellation behavior +//! its OCC counterpart pins. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Barrier, Mutex}; +use std::thread; +use std::time::{Duration, Instant}; + +use mz_environmentd::test_util; +use mz_ore::assert_contains; +use mz_ore::error::ErrorExt; +use mz_ore::retry::Retry; +use tokio_postgres::error::SqlState; + +/// A harness with frontend OCC read-then-write enabled and nothing else. +/// +/// Callers add what they need on top, notably `unsafe_mode` for the tests that +/// hold a statement open with `mz_unsafe` functions. +fn frontend_occ_harness() -> test_util::TestHarness { + test_util::TestHarness::default().with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + "true".to_string(), + ) +} + +/// The server's message for a client error, or the client-side rendering when +/// the error never reached the server. `postgres::Error::to_string` is only "db +/// error" for a server error, so matching on it tells us nothing. +fn server_error_message(err: &postgres::Error) -> String { + match err.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => err.to_string(), + } +} + +// `mz_query_total` feeds product telemetry, and DML the session task sequences +// itself must be counted there exactly once, like DML the coordinator +// sequences. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_counts_query_total() { + for frontend_occ in [false, true] { + let server = test_util::TestHarness::default() + .with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + frontend_occ.to_string(), + ) + .start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE query_total_t (x INT)") + .unwrap(); + + for (statement_type, sql) in [ + ("insert", "INSERT INTO query_total_t SELECT 1"), + ("update", "UPDATE query_total_t SET x = 2"), + ("delete", "DELETE FROM query_total_t"), + ] { + let labels = [("session_type", "user"), ("statement_type", statement_type)]; + let before = + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &labels); + client.batch_execute(sql).unwrap(); + let after = + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &labels); + assert_eq!( + after, + before + 1, + "mz_query_total{{statement_type={statement_type}}} moved from {before} to {after} \ + across `{sql}`, with frontend OCC read-then-write {frontend_occ}" + ); + } + } +} + +// 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_long_running_write() { + let server = frontend_occ_harness().unsafe_mode().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(); + + // The last cancel request the thread sent is processed asynchronously, so + // it can still land on this read-back. Retry it in that case. + let rows = Retry::default() + .max_tries(5) + .clamp_backoff(Duration::from_millis(100)) + .retry(|_| 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 = Retry::default() + .max_tries(5) + .clamp_backoff(Duration::from_millis(100)) + .retry(|_| client.query_one("SELECT count(*) FROM t", &[])) + .unwrap() + .get::<_, i64>(0); + assert_eq!( + rows, 1, + "cancelled statement should not have committed writes" + ); + + // The read-then-write ran its selection through an internal subscribe, so a + // `SubscribeHandle` whose drop never reached the coordinator would leave + // that dataflow installed. + wait_for_no_dataflows(&mut client, "after cancelling a read-then-write"); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_constant_insert_prepares_unmaterializable_functions() { + let server = frontend_occ_harness().unsafe_mode().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_rejected_in_multi_statement_batch() { + let server = frontend_occ_harness().unsafe_mode().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"); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_constant_insert_respects_max_result_size() { + let server = frontend_occ_harness() + .unsafe_mode() + .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_constant_insert_rejects_mz_now() { + let server = frontend_occ_harness().unsafe_mode().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_returning_error_does_not_commit_write() { + let server = frontend_occ_harness().unsafe_mode().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_empty_snapshot_returns_zero() { + let server = frontend_occ_harness().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_concurrent_updates_retry() { + const NUM_WORKERS: usize = 4; + const UPDATES_PER_WORKER: usize = 25; + + let server = frontend_occ_harness().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(), + ); +} + +// A frontend OCC read-then-write that times out must not commit its writes. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_timeout_does_not_commit_write() { + let server = frontend_occ_harness().unsafe_mode().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"); +} + +/// 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. +/// +/// The sum oracle alone cannot distinguish "no over-deletion under concurrency" +/// from "there was no concurrency": workers running one after another satisfy it +/// too. So the clients are connected before any worker starts, released +/// together by a barrier, and the OCC retry histogram must show at least one +/// attempt that saw its read timestamp move. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_concurrent_delete_does_not_over_delete() { + const MULTIPLICITY: i32 = 7; + const NUM_WORKERS: usize = 6; + const ROUNDS: usize = 10; + + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + let mut setup = server.connect(postgres::NoTls).unwrap(); + setup.batch_execute("CREATE TABLE t (id INT)").unwrap(); + + for _round in 0..ROUNDS { + setup.batch_execute("DELETE FROM t").unwrap(); + setup + .execute( + "INSERT INTO t SELECT 1 FROM generate_series(1, $1)", + &[&MULTIPLICITY], + ) + .unwrap(); + + // Connecting inside the spawn loop would let the first worker finish + // its DELETE before the last client even exists. + let clients: Vec<_> = (0..NUM_WORKERS) + .map(|_| server.connect(postgres::NoTls).unwrap()) + .collect(); + + let barrier = Arc::new(Barrier::new(NUM_WORKERS)); + let total_deleted = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::new(); + for mut client in clients { + let barrier = Arc::clone(&barrier); + let total_deleted = Arc::clone(&total_deleted); + handles.push(thread::spawn(move || { + barrier.wait(); + 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)", + ); + } + + // Proof that the deletes actually raced: with all workers deleting the same + // rows at once, at least one attempt must have found its read timestamp + // passed and retried against fresh state. + 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 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 DELETE to retry under contention. \ + all {} observations landed in the 0-retry bucket, so the workers \ + did not actually run concurrently", + histogram.get_sample_count(), + ); +} + +/// Multiset multiplicity must be reflected in affected-row counts for +/// DELETE / UPDATE / INSERT...SELECT. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_duplicate_row_multiplicity_counts() { + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + 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 test_not_null_constraint_enforced() { + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + 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 report the inserted rows, in both the constant and the +/// read-dependent `INSERT ... SELECT` shape. Only `INSERT` accepts it, so +/// `UPDATE` and `DELETE` have nothing to check here. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_insert_returning_values() { + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + 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 test_update_moves_overlapping_rows() { + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + 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 test_insert_select_from_materialized_view() { + let server = frontend_occ_harness().unsafe_mode().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 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 +/// conserve exactly the writes it reported, and may only fail with errors a +/// correct implementation is allowed to return. +/// +/// Only one of the four statement arms mutates anything, which is what makes an +/// exact oracle available: `v` never goes negative so `WHERE v < 0` matches +/// nothing, no row is ever deleted so the `NOT EXISTS` guard never fires, and +/// `SET v = v` consolidates to no diffs. So `sum(v)` must equal the affected-row +/// counts the `v = v + 1` arm reported, and `count(*)` must not move. A write +/// that committed while reporting an error, reported an affected row without +/// committing, or applied its diffs twice across a retry all break that +/// equality. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_concurrent_mixed_dml_conserves_writes() { + const NUM_WORKERS: usize = 8; + const ITERS: usize = 30; + const NUM_ROWS: usize = 20; + // The only failure a correct implementation may return here. An internal + // error, a concurrently modified write target, or an indeterminate write + // are all bugs, so the error set is checked against this list instead of + // against a handful of bad-news substrings. + const ALLOWED_ERRORS: &[&str] = + &["read-then-write exceeded maximum retry attempts under contention"]; + + let server = frontend_occ_harness().unsafe_mode().start_blocking(); + let mut setup = server.connect(postgres::NoTls).unwrap(); + setup + .batch_execute("CREATE TABLE t (id INT, v INT)") + .unwrap(); + setup + .execute( + &format!("INSERT INTO t SELECT generate_series(1, {NUM_ROWS}), 0"), + &[], + ) + .unwrap(); + + let reported_increments = Arc::new(AtomicUsize::new(0)); + let 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 reported_increments = Arc::clone(&reported_increments); + let errors = Arc::clone(&errors); + handles.push(thread::spawn(move || { + for i in 0..ITERS { + let id = (w * 7 + i) % NUM_ROWS + 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}"), + }; + match client.execute(stmt.as_str(), &[]) { + Ok(affected) if i % 4 == 0 => { + reported_increments + .fetch_add(usize::try_from(affected).unwrap(), Ordering::SeqCst); + } + Ok(_) => {} + Err(e) => errors + .lock() + .unwrap() + .push(format!("`{stmt}`: {}", server_error_message(&e))), + } + } + })); + } + for h in handles { + h.join().expect("worker panicked"); + } + + let errors = errors.lock().unwrap(); + let unexpected: Vec<&String> = errors + .iter() + .filter(|error| !ALLOWED_ERRORS.iter().any(|allowed| error.contains(allowed))) + .collect(); + assert!( + unexpected.is_empty(), + "concurrent mixed DML returned errors outside the allow-list: {unexpected:#?}" + ); + + let sum: i64 = setup + .query_one("SELECT coalesce(sum(v), 0)::bigint FROM t", &[]) + .unwrap() + .get(0); + let expected = i64::try_from(reported_increments.load(Ordering::SeqCst)).unwrap(); + assert_eq!( + sum, expected, + "sum(v) must equal the number of rows the incrementing UPDATEs reported" + ); + + let count: i64 = setup + .query_one("SELECT count(*) FROM t", &[]) + .unwrap() + .get(0); + assert_eq!( + count, + i64::try_from(NUM_ROWS).unwrap(), + "no arm of this workload may add or remove a row" + ); +} + +/// A read-then-write whose read resolves to a far-future timestamp (here, a +/// `REFRESH AT ` materialized view) parks in +/// `ensure_read_linearized`'s sleep loop. `statement_timeout` has to end that +/// park, which it does because +/// `SessionClient::try_frontend_read_then_write_with_cancel` bounds the +/// *entire* operation, not just the OCC loop. +/// +/// NOTE: This holds for the OCC path only. The coordinator path also bounds the +/// scenario, but through the `statement_timeout` it arms around the row stream +/// it reads the selection from, and it blocks while holding the target table's +/// write lock rather than an OCC permit. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_far_future_refresh_mv_respects_statement_timeout() { + let server = frontend_occ_harness() + .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(); + // 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", so preserve + // the SqlState code and server message for the assertion to inspect. + 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)) => { + // `statement_timeout` bounds the whole operation, so the far-future + // op must error out rather than silently succeed, and do so 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})." + ); + } + } +} + +/// A far-future RTW parked in `ensure_read_linearized` holds its OCC semaphore +/// permit for as long as it is parked. With a bounded permit pool, one such op +/// therefore starves every other read-then-write in the process, including ones +/// on unrelated tables, because a victim blocks on permit acquisition *before* +/// the OCC loop. +/// +/// The victim must still honor its own `statement_timeout`, which it does +/// because `try_frontend_read_then_write_with_cancel`'s `select!` bounds the +/// *whole* operation, permit-acquisition wait included. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_far_future_read_then_write_starves_permit_pool() { + let server = frontend_occ_harness() + .unsafe_mode() + .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 grabs the single OCC permit and parks + // in ensure_read_linearized. + // + // `mz_query_total` is bumped when the session task takes the statement + // over, which is the last observable point before it acquires the permit. + // Polling for that instead of sleeping for a fixed span keeps a loaded + // machine from starting the victim while the parked op is still connecting + // or planning, which would fail the test without the permit pool ever being + // starved. + let insert_labels = [("session_type", "user"), ("statement_type", "insert")]; + let inserts_before = + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &insert_labels); + 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"); + }); + Retry::default() + .max_duration(Duration::from_secs(60)) + .clamp_backoff(Duration::from_millis(100)) + .retry(|_| { + let inserts_now = test_util::get_counter_value( + server.metrics_registry(), + "mz_query_total", + &insert_labels, + ); + if inserts_now > inserts_before { + Ok(()) + } else { + Err("far-future INSERT has not started executing") + } + }) + .expect("far-future INSERT never started executing"); + // The counter moves a few planning steps before the permit is taken. + thread::sleep(Duration::from_secs(1)); + + // 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. Its own + // `statement_timeout = '3s'` bounds the permit-acquisition wait, so + // it returns a timeout error rather than hanging. + 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. Because \ + statement_timeout bounds the permit wait, the victim should time out on \ + permit acquisition within ~3s instead of hanging ({recv_err})." + ); + } + } +} + +/// A cancelled or timed-out read-then-write must give back its OCC permit. +/// +/// The permit pool is sized to one here, so a permit that leaks wedges every +/// later read-then-write in the process, and the follow-up UPDATE is the oracle: +/// it needs the permit the abandoned statement held, and its own +/// `statement_timeout` turns a wedge into a failure instead of a hang. The cycle +/// repeats so that leaking one permit per iteration is fatal rather than +/// tolerable. +/// +/// The abandoned statement parks in `ensure_read_linearized`, waiting on a +/// far-future REFRESH materialized view. That holds a permit while occupying no +/// cluster worker, which the oracle depends on: a statement that blocks by +/// sleeping inside its dataflow keeps the worker busy after it is cancelled, so +/// the follow-up would be measuring cluster occupancy rather than permit +/// availability. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_cancel_and_timeout_release_permit() { + const ITERATIONS: i32 = 5; + // Bounds the follow-up write, so a leaked permit surfaces as this error + // rather than as a hang. + const FOLLOW_UP_TIMEOUT: &str = "30s"; + + let server = frontend_occ_harness() + .unsafe_mode() + .with_system_parameter_default("enable_refresh_every_mvs".to_string(), "true".to_string()) + // One permit: a single leak starves every read-then-write. + .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 t (n INT)").unwrap(); + client.batch_execute("INSERT INTO t VALUES (0)").unwrap(); + client + .batch_execute(&format!("SET statement_timeout = '{FOLLOW_UP_TIMEOUT}'")) + .unwrap(); + + // The read cannot be linearized until the year 3000, so the statement parks + // holding the sole permit. + let parking_insert = "INSERT INTO dst SELECT a FROM mv"; + + for iteration in 1..=ITERATIONS { + let mut parked = server.connect(postgres::NoTls).unwrap(); + let cancel_token = parked.cancel_token(); + let baseline = user_insert_count(&server); + let parked_handle = thread::spawn(move || parked.batch_execute(parking_insert)); + wait_until_parked_holding_a_permit(&server, baseline, iteration, "the cancel half"); + cancel_token.cancel_query(postgres::NoTls).unwrap(); + let err = parked_handle + .join() + .unwrap() + .expect_err("the parked INSERT should have been cancelled"); + let message = server_error_message(&err); + assert_eq!( + err.code(), + Some(&SqlState::QUERY_CANCELED), + "iteration {iteration}: cancelled INSERT reported {message}" + ); + // A statement timeout reports this same SQLSTATE, so without the message + // check a cancel that never arrived would look like a pass: the victim + // would sit on the default 60s timeout and fail with QUERY_CANCELED too. + assert!( + message.contains("user request"), + "iteration {iteration}: expected a cancellation, got {message}" + ); + + follow_up_write_gets_a_permit(&mut client, iteration, "a cancelled"); + + // Long enough that the deadline cannot fire before the permit is taken, + // short enough to keep the test quick. + let mut parked = server.connect(postgres::NoTls).unwrap(); + parked + .batch_execute("SET statement_timeout = '5s'") + .unwrap(); + let baseline = user_insert_count(&server); + let parked_handle = thread::spawn(move || parked.batch_execute(parking_insert)); + wait_until_parked_holding_a_permit(&server, baseline, iteration, "the timeout half"); + let err = parked_handle + .join() + .unwrap() + .expect_err("the parked INSERT should have timed out"); + let message = server_error_message(&err); + assert_eq!( + err.code(), + Some(&SqlState::QUERY_CANCELED), + "iteration {iteration}: timed-out INSERT reported {message}" + ); + assert!( + message.contains("statement timeout"), + "iteration {iteration}: expected a statement-timeout error, got {message}" + ); + + follow_up_write_gets_a_permit(&mut client, iteration, "a timed-out"); + } + + // Two follow-up writes per iteration, each of them a single-row UPDATE. + let n: i32 = client.query_one("SELECT n FROM t", &[]).unwrap().get(0); + assert_eq!(n, 2 * ITERATIONS, "a follow-up write did not commit"); + let count: i64 = client + .query_one("SELECT count(*) FROM dst", &[]) + .unwrap() + .get(0); + assert_eq!( + count, 0, + "a cancelled or timed-out INSERT committed its rows" + ); +} + +/// Counts user INSERTs the process has executed, the observable that +/// [`wait_until_parked_holding_a_permit`] watches. +fn user_insert_count(server: &test_util::TestServerWithRuntime) -> u64 { + const LABELS: [(&str, &str); 2] = [("session_type", "user"), ("statement_type", "insert")]; + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &LABELS) +} + +/// Waits until the INSERT under test holds an OCC permit. +/// +/// Without this, a loaded machine can let the cancel or the deadline land while +/// the statement is still connecting or planning. That exercises an exit which +/// never held a permit, so it proves nothing about releasing one, and it passes +/// anyway. +/// +/// There is no metric for permit acquisition, so this polls the closest +/// observable, the `mz_query_total` bump in `ExecutionLogging::take_over`, and +/// then settles. The bump happens a few steps before `acquire_owned`, hence the +/// settle. +/// +/// `baseline` must be read with [`user_insert_count`] before the statement is +/// started. Every user INSERT in the process shares that counter, so a baseline +/// read afterwards can already include the statement we are waiting for, and +/// then no bump ever arrives. +fn wait_until_parked_holding_a_permit( + server: &test_util::TestServerWithRuntime, + baseline: u64, + iteration: i32, + half: &str, +) { + const SETTLE: Duration = Duration::from_secs(1); + + Retry::default() + .max_duration(Duration::from_secs(60)) + .clamp_backoff(Duration::from_millis(50)) + .retry(|_| { + let inserts = user_insert_count(server); + if inserts > baseline { + Ok(()) + } else { + Err(inserts) + } + }) + .unwrap_or_else(|inserts| { + panic!( + "iteration {iteration}, {half}: the parked INSERT never reached its session task, \ + mz_query_total stuck at {inserts}" + ) + }); + thread::sleep(SETTLE); +} + +/// For `test_cancel_and_timeout_release_permit`: runs a read-then-write +/// that can only make progress if the abandoned statement released its permit. +fn follow_up_write_gets_a_permit(client: &mut postgres::Client, iteration: i32, predecessor: &str) { + let affected = client + .execute("UPDATE t SET n = n + 1", &[]) + .unwrap_or_else(|err| { + panic!( + "iteration {iteration}: the read-then-write following {predecessor} one failed \ + with `{}`, which is what a leaked OCC permit looks like", + server_error_message(&err) + ) + }); + assert_eq!( + affected, 1, + "iteration {iteration}: follow-up UPDATE after {predecessor} one affected {affected} rows" + ); +} + +/// Waits until the cluster runs no dataflows other than the ones introspection +/// installs for itself. An OCC read-then-write's internal subscribe is a +/// dataflow, so this catches a `SubscribeHandle` whose drop never tore it down. +fn wait_for_no_dataflows(client: &mut postgres::Client, context: &str) { + // Storage operators have their IDs offset by STORAGE_ID_OFFSET (1 << 48), + // so they are excluded by id. + const DATAFLOW_QUERY: &str = "SELECT count(*) \ + FROM mz_introspection.mz_dataflows \ + WHERE name NOT LIKE '%introspection-subscribe%' \ + AND id < 281474976710656"; + + Retry::default() + .max_duration(Duration::from_secs(60)) + .clamp_backoff(Duration::from_millis(500)) + .retry(|_| { + let count: i64 = client.query_one(DATAFLOW_QUERY, &[]).unwrap().get(0); + if count == 0 { Ok(()) } else { Err(count) } + }) + .unwrap_or_else(|count| panic!("{count} dataflows still installed {context}")); +} + +/// A read-then-write computed against one generation of its write target must +/// not commit once a concurrent `ALTER TABLE ... ADD COLUMN` has given the +/// target a new one. Either the write wins the race and commits in full, or the +/// group committer rejects it as a concurrent dependency mutation, which is a +/// retryable serialization failure. Both outcomes leave the table consistent +/// and neither may panic the coordinator. +/// +/// Which of the two happens depends on where the ALTER lands relative to the +/// generation the write captured, so the test accepts either. The sleep in the +/// selection makes the write slow enough that the ALTER usually lands inside +/// the window. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_write_racing_alter_table_add_column() { + const SLEEP_SECS: i32 = 3; + const ROUNDS: usize = 3; + + let server = frontend_occ_harness() + .unsafe_mode() + .with_system_parameter_default( + "unsafe_enable_unsafe_functions".to_string(), + "true".to_string(), + ) + .with_system_parameter_default( + "enable_alter_table_add_column".to_string(), + "true".to_string(), + ) + .start_blocking(); + + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE t (n INT, ts INT)") + .unwrap(); + client + .batch_execute(&format!("INSERT INTO t VALUES (0, {SLEEP_SECS})")) + .unwrap(); + + let update_labels = [("session_type", "user"), ("statement_type", "update")]; + let mut committed = 0; + for round in 0..ROUNDS { + let updates_before = test_util::get_counter_value( + server.metrics_registry(), + "mz_query_total", + &update_labels, + ); + let mut writer = server.connect(postgres::NoTls).unwrap(); + let (tx, rx) = std::sync::mpsc::channel(); + let handle = thread::spawn(move || { + // The CASE always takes its ELSE branch (`mz_sleep` returns NULL), + // so every row matches and the sleep runs in the subscribe + // dataflow, which is what keeps the write in flight. + let result = writer.execute( + "UPDATE t SET n = n + 1 \ + WHERE ts >= CASE WHEN mz_unsafe.mz_sleep(ts) > 0 THEN 1 ELSE 0 END", + &[], + ); + // Preserve the SqlState and server message: `to_string()` on a + // server error is only "db error". + let result = result.map_err(|err| { + ( + err.code().cloned(), + err.as_db_error().map(|db| db.message().to_string()), + ) + }); + let _ = tx.send(result); + }); + + // Land the ALTER after the UPDATE started executing, so it has a chance + // to invalidate the generation the UPDATE is writing against. + Retry::default() + .max_duration(Duration::from_secs(60)) + .clamp_backoff(Duration::from_millis(100)) + .retry(|_| { + let updates_now = test_util::get_counter_value( + server.metrics_registry(), + "mz_query_total", + &update_labels, + ); + if updates_now > updates_before { + Ok(()) + } else { + Err("racing UPDATE has not started executing") + } + }) + .expect("racing UPDATE never started executing"); + client + .batch_execute(&format!("ALTER TABLE t ADD COLUMN c{round} INT")) + .unwrap(); + + let outcome = rx + .recv_timeout(Duration::from_secs(120)) + .expect("UPDATE racing ALTER TABLE never returned"); + handle.join().expect("writer thread panicked"); + match outcome { + Ok(affected) => { + assert_eq!( + affected, 1, + "round {round}: the UPDATE committed but reported {affected} rows" + ); + committed += 1; + } + Err((code, message)) => { + assert_eq!( + code.as_ref(), + Some(&SqlState::T_R_SERIALIZATION_FAILURE), + "round {round}: the UPDATE failed with SqlState {code:?} \ + (message: {message:?}) rather than as a retryable conflict" + ); + assert!( + message + .as_deref() + .is_some_and(|m| m.contains("was concurrently modified")), + "round {round}: unexpected serialization-failure message {message:?}" + ); + } + } + + // Whichever way the race went, the table holds one row whose counter + // matches the number of UPDATEs that committed. + let rows = client.query("SELECT n FROM t", &[]).unwrap(); + assert_eq!(rows.len(), 1, "round {round}: unexpected row count"); + assert_eq!( + rows[0].get::<_, i32>(0), + committed, + "round {round}: a rejected UPDATE left its write behind, or a \ + committed one applied twice" + ); + } +} + +// A read-then-write the frontend OCC path committed must be visible to a +// strict serializable read that starts afterwards, on any session. The write +// commits at a timestamp it chose itself, so if it reported success before that +// timestamp was reflected in the timeline's oracle, a later linearized read +// could pick an earlier timestamp and miss the write. +// +// The reader connects fresh every iteration, so it cannot inherit the writer +// session's timestamp bookkeeping: the only thing that can carry the write +// forward is the global oracle. Iterating is what gives the race a chance to +// show up. A single pass could pass by luck. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_frontend_occ_write_visible_to_linearizable_read() { + const ITERATIONS: i32 = 50; + + let server = frontend_occ_harness().start_blocking(); + + let mut writer = server.connect(postgres::NoTls).unwrap(); + writer + .batch_execute("CREATE TABLE t (id INT, v INT)") + .unwrap(); + writer.batch_execute("INSERT INTO t VALUES (1, 0)").unwrap(); + + for iteration in 1..=ITERATIONS { + let affected = writer + .execute("UPDATE t SET v = v + 1 WHERE id = 1", &[]) + .unwrap(); + assert_eq!(affected, 1, "iteration {iteration}: UPDATE lost its row"); + + let mut reader = server.connect(postgres::NoTls).unwrap(); + reader + .batch_execute("SET transaction_isolation = 'strict serializable'") + .unwrap(); + let v: i32 = reader + .query_one("SELECT v FROM t WHERE id = 1", &[]) + .unwrap() + .get(0); + assert_eq!( + v, iteration, + "a strict serializable read on a fresh session did not observe the \ + committed UPDATE" + ); + } +} + +// Test that the server properly handles cancellation requests of read-then-write queries. +// See database-issues#6134. +#[mz_ore::test] +#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_wait` on OS `linux` +#[allow(clippy::disallowed_methods)] +fn test_cancel_read_then_write() { + let server = test_util::TestHarness::default() + .unsafe_mode() + .start_blocking(); + server.enable_feature_flags(&["unsafe_enable_unsafe_functions"]); + + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE foo (a TEXT, ts INT)") + .unwrap(); + + // Lots of races here, so try this whole thing in a loop. + Retry::default() + .clamp_backoff(Duration::ZERO) + .retry(|_state| { + let mut client1 = server.connect(postgres::NoTls).unwrap(); + let mut client2 = server.connect(postgres::NoTls).unwrap(); + let cancel_token = client2.cancel_token(); + + client1.batch_execute("DELETE FROM foo").unwrap(); + client1.batch_execute("SET statement_timeout = '5s'").unwrap(); + client1 + .batch_execute("INSERT INTO foo VALUES ('hello', 10)") + .unwrap(); + + let handle1 = thread::spawn(move || { + let err = client1 + .batch_execute("insert into foo select a, case when mz_unsafe.mz_sleep(ts) > 0 then 0 end as ts from foo") + .unwrap_err(); + assert_contains!( + err.to_string_with_causes(), + "statement timeout" + ); + client1 + }); + std::thread::sleep(Duration::from_millis(100)); + let handle2 = thread::spawn(move || { + let err = client2 + .batch_execute("insert into foo values ('blah', 1);") + .unwrap_err(); + assert_contains!( + err.to_string_with_causes(), + "canceling statement" + ); + }); + std::thread::sleep(Duration::from_millis(100)); + cancel_token.cancel_query(postgres::NoTls)?; + let mut client1 = handle1.join().unwrap(); + handle2.join().unwrap(); + let rows:i64 = client1.query_one ("SELECT count(*) FROM foo", &[]).unwrap().get(0); + // We ran 3 inserts. First succeeded. Second timedout. Third cancelled. + if rows !=1 { + anyhow::bail!("unexpected row count: {rows}"); + } + Ok::<_, anyhow::Error>(()) + }) + .unwrap(); +} diff --git a/src/environmentd/tests/server.rs b/src/environmentd/tests/server.rs index 7e3c0dd58f669..865523a8e9b98 100644 --- a/src/environmentd/tests/server.rs +++ b/src/environmentd/tests/server.rs @@ -2592,70 +2592,6 @@ fn test_github_20262() { } } -// Test that the server properly handles cancellation requests of read-then-write queries. -// See database-issues#6134. -#[mz_ore::test] -#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `epoll_wait` on OS `linux` -#[allow(clippy::disallowed_methods)] -fn test_cancel_read_then_write() { - let server = test_util::TestHarness::default() - .unsafe_mode() - .start_blocking(); - server.enable_feature_flags(&["unsafe_enable_unsafe_functions"]); - - let mut client = server.connect(postgres::NoTls).unwrap(); - client - .batch_execute("CREATE TABLE foo (a TEXT, ts INT)") - .unwrap(); - - // Lots of races here, so try this whole thing in a loop. - Retry::default() - .clamp_backoff(Duration::ZERO) - .retry(|_state| { - let mut client1 = server.connect(postgres::NoTls).unwrap(); - let mut client2 = server.connect(postgres::NoTls).unwrap(); - let cancel_token = client2.cancel_token(); - - client1.batch_execute("DELETE FROM foo").unwrap(); - client1.batch_execute("SET statement_timeout = '5s'").unwrap(); - client1 - .batch_execute("INSERT INTO foo VALUES ('hello', 10)") - .unwrap(); - - let handle1 = thread::spawn(move || { - let err = client1 - .batch_execute("insert into foo select a, case when mz_unsafe.mz_sleep(ts) > 0 then 0 end as ts from foo") - .unwrap_err(); - assert_contains!( - err.to_string_with_causes(), - "statement timeout" - ); - client1 - }); - std::thread::sleep(Duration::from_millis(100)); - let handle2 = thread::spawn(move || { - let err = client2 - .batch_execute("insert into foo values ('blah', 1);") - .unwrap_err(); - assert_contains!( - err.to_string_with_causes(), - "canceling statement" - ); - }); - std::thread::sleep(Duration::from_millis(100)); - cancel_token.cancel_query(postgres::NoTls)?; - let mut client1 = handle1.join().unwrap(); - handle2.join().unwrap(); - let rows:i64 = client1.query_one ("SELECT count(*) FROM foo", &[]).unwrap().get(0); - // We ran 3 inserts. First succeeded. Second timedout. Third cancelled. - if rows !=1 { - anyhow::bail!("unexpected row count: {rows}"); - } - Ok::<_, anyhow::Error>(()) - }) - .unwrap(); -} - #[mz_ore::test(tokio::test(flavor = "multi_thread", worker_threads = 1))] #[cfg_attr(miri, ignore)] // too slow async fn test_http_metrics() { diff --git a/src/environmentd/tests/sql.rs b/src/environmentd/tests/sql.rs index 93d8c09c9c5b6..0fc7425b65de1 100644 --- a/src/environmentd/tests/sql.rs +++ b/src/environmentd/tests/sql.rs @@ -1761,8 +1761,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 @@ -1776,7 +1780,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/environmentd/tests/statement_logging.rs b/src/environmentd/tests/statement_logging.rs index 859ef1dd493a1..c98cf9edfed01 100644 --- a/src/environmentd/tests/statement_logging.rs +++ b/src/environmentd/tests/statement_logging.rs @@ -14,14 +14,17 @@ use std::time::Duration; use chrono::{DateTime, Utc}; use mz_environmentd::test_util; +use mz_ore::assert_contains; use mz_ore::assert_none; use mz_ore::cast::{CastFrom, CastLossy, TryCastFrom}; use mz_ore::collections::CollectionExt; +use mz_ore::error::ErrorExt; use mz_ore::metrics::MetricsRegistry; use mz_ore::now::to_datetime; use mz_ore::retry::Retry; use mz_pgrepr::UInt8; use mz_sql_parser::ast::display::AstDisplay; +use tokio_postgres::error::SqlState; use tungstenite::Message; /// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped. @@ -48,6 +51,12 @@ impl TestServerWithStatementLoggingChecks { pub fn metrics_registry(&self) -> &MetricsRegistry { self.server.metrics_registry() } + + /// Returns a config for connecting to the __public__ SQL port, so a test + /// can connect as a role other than the default one. + pub fn pg_config(&self) -> postgres::Config { + self.server.pg_config() + } } /// Helper to get statement logging record counts from the metrics registry. @@ -1017,7 +1026,6 @@ fn test_statement_logging_ws_subscribe_no_crash() { // Give the server time to crash, if it's going to. std::thread::sleep(Duration::from_secs(1)) } - /// `finished_at` must record when execution finished, not when the coordinator /// got around to the end event. A statement the session task retires itself /// reports its own end timestamp, so a busy coordinator cannot inflate it. @@ -1089,3 +1097,675 @@ fn test_statement_logging_finished_at_excludes_coordinator_queue() { finished_at.timestamp_millis() - finished_bound ); } + +#[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 coordinator path does +// this through `set_statement_execution_timestamp` during group commit, so the +// frontend path has 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" + ); +} + +/// One case in the DML statement-logging parity table. +struct DmlLoggingCase { + /// Statements run before the one under test, to set up transaction state. + /// These must succeed. + before: &'static [&'static str], + /// The statement under test. Its log row is found by the redacted form of + /// this text, so no two cases may share it. + sql: &'static str, + /// Statements run after the one under test, to make the session usable + /// again. These must succeed. + after: &'static [&'static str], + /// Whether only the frontend path records an `execution_timestamp` for this + /// statement. It holds for every read-then-write that commits: the frontend + /// emits the write timestamp when its write lands, while the coordinator + /// retires the log entry in `sequence_read_then_write` and only reaches the + /// group commit that would set the timestamp afterwards, by which time the + /// entry is closed. The difference is pinned per case rather than excluded + /// from the comparison, so a change on either path fails this test. + frontend_only_execution_timestamp: bool, +} + +/// DML whose statement-logging record must not depend on which path sequenced +/// it. Failures matter as much as successes: the two paths reject a statement +/// at different points, so their error exits are where they drift apart. +const DML_LOGGING_PARITY_CASES: &[DmlLoggingCase] = &[ + DmlLoggingCase { + before: &[], + sql: "UPDATE parity_t SET x = x + 1", + after: &[], + frontend_only_execution_timestamp: true, + }, + DmlLoggingCase { + before: &[], + sql: "DELETE FROM parity_t WHERE x = 2", + after: &[], + frontend_only_execution_timestamp: true, + }, + DmlLoggingCase { + before: &[], + sql: "INSERT INTO parity_t SELECT x + 10 FROM parity_t", + after: &[], + frontend_only_execution_timestamp: true, + }, + DmlLoggingCase { + before: &[], + sql: "INSERT INTO parity_t VALUES (100) RETURNING x", + after: &[], + frontend_only_execution_timestamp: true, + }, + // A RETURNING insert that matches no rows. Both paths report a row count + // rather than an empty result set, so this pins the response kind (and with + // it `rows_returned` and `result_size`) of the zero-row case. + DmlLoggingCase { + before: &[], + sql: "INSERT INTO parity_t SELECT 101 WHERE false RETURNING x", + after: &[], + frontend_only_execution_timestamp: false, + }, + // Rejected while describing the portal, before either path begins an + // execution, so neither logs one. + DmlLoggingCase { + before: &[], + sql: "UPDATE parity_t SET x = nonexistent_col", + after: &[], + frontend_only_execution_timestamp: false, + }, + // Bounded staleness forbids writes. Both paths reject the statement after + // they have planned it and recorded its cluster, so the error row carries a + // cluster on both. + DmlLoggingCase { + before: &["SET transaction_isolation = 'bounded staleness 5s'"], + sql: "DELETE FROM parity_t WHERE x > 1000", + after: &["SET transaction_isolation = 'strict serializable'"], + frontend_only_execution_timestamp: false, + }, + DmlLoggingCase { + before: &["BEGIN"], + sql: "DELETE FROM parity_t", + after: &["ROLLBACK"], + frontend_only_execution_timestamp: false, + }, +]; + +/// The part of a statement's log row that both sequencing paths must agree on. +#[derive(Debug, PartialEq, Eq)] +struct DmlLoggingRecord { + finished_status: String, + error_message: Option, + /// Compared as "is it recorded at all", since the byte count itself is not + /// a property of the path. + result_size_is_null: bool, + rows_returned: Option, + execution_strategy: Option, + has_cluster: bool, + has_execution_timestamp: bool, +} + +/// SQL of the statement whose log row marks the end of the parity run. Once it +/// is visible, every earlier statement's row is too: the log's end-execution +/// events are recorded in the order the statements finished, and a flush +/// appends everything pending at once. +const DML_LOGGING_PARITY_SENTINEL: &str = "SELECT count(*) FROM parity_t"; + +/// The form of `sql` that statement logging records, which is what identifies a +/// statement's rows. DML is stored redacted. +fn redacted_sql(sql: &str) -> String { + mz_sql::parse::parse(sql) + .unwrap() + .into_element() + .ast + .to_ast_string_redacted() +} + +/// Reads the log row of the statement whose redacted SQL is `redacted_sql`. +/// +/// `None` means the execution was not logged at all, which is a comparable +/// outcome: a statement rejected before execution begins has no row on either +/// path. Only meaningful once the sentinel row is visible. +fn read_dml_logging_record( + mz_client: &mut postgres::Client, + redacted_sql: &str, +) -> Option { + let rows = mz_client + .query( + "SELECT + mseh.finished_status, + mseh.error_message, + mseh.result_size IS NULL, + mseh.rows_returned, + mseh.execution_strategy, + mseh.cluster_name IS NOT NULL, + mseh.execution_timestamp IS NOT NULL +FROM mz_internal.mz_statement_execution_history AS mseh +JOIN mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id +JOIN (SELECT DISTINCT sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE mst.redacted_sql = $1 AND mseh.finished_at IS NOT NULL", + &[&redacted_sql], + ) + .unwrap(); + + assert!( + rows.len() <= 1, + "expected at most one log row for {redacted_sql}, got {}", + rows.len() + ); + rows.first().map(|row| DmlLoggingRecord { + finished_status: row.get(0), + error_message: row.get(1), + result_size_is_null: row.get(2), + rows_returned: row.get(3), + execution_strategy: row.get(4), + has_cluster: row.get(5), + has_execution_timestamp: row.get(6), + }) +} + +/// Runs [`DML_LOGGING_PARITY_CASES`] on a server configured with the given +/// value of the frontend OCC read-then-write flag, and returns each case's log +/// row in table order. +#[allow(clippy::disallowed_methods)] +fn collect_dml_logging_records( + frontend_occ: bool, +) -> Vec<(&'static str, String, Option)> { + let harness = test_util::TestHarness::default().with_system_parameter_default( + "enable_adapter_frontend_occ_read_then_write".to_string(), + frontend_occ.to_string(), + ); + let (server, mut client) = setup_statement_logging_core(1.0, 1.0, "", harness); + + client.batch_execute("SET CLUSTER TO quickstart").unwrap(); + client + .batch_execute("CREATE TABLE parity_t (x INT)") + .unwrap(); + client + .batch_execute("INSERT INTO parity_t VALUES (1), (2)") + .unwrap(); + + for case in DML_LOGGING_PARITY_CASES { + for sql in case.before { + client.batch_execute(sql).unwrap(); + } + // Several cases fail on purpose. The log row is what the test reads, so + // the client-visible outcome is deliberately ignored here. + let _ = client.batch_execute(case.sql); + for sql in case.after { + client.batch_execute(sql).unwrap(); + } + } + client.batch_execute(DML_LOGGING_PARITY_SENTINEL).unwrap(); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + let sentinel = redacted_sql(DML_LOGGING_PARITY_SENTINEL); + Retry::default() + .max_duration(Duration::from_secs(60)) + .retry( + |_| match read_dml_logging_record(&mut mz_client, &sentinel) { + Some(_) => Ok(()), + None => Err(()), + }, + ) + .expect("statement log should flush the sentinel row"); + + DML_LOGGING_PARITY_CASES + .iter() + .map(|case| { + let redacted_sql = redacted_sql(case.sql); + let record = read_dml_logging_record(&mut mz_client, &redacted_sql); + (case.sql, redacted_sql, record) + }) + .collect() +} + +// DELETE/UPDATE/INSERT..SELECT are sequenced either by the coordinator or by +// the session task, depending on a flag fixed at process startup. What they +// record in `mz_statement_execution_history` must not depend on that: the log +// is a user-visible product surface, and a customer querying it cannot tell +// which path ran. +#[mz_ore::test] +fn test_statement_logging_dml_path_parity() { + let coordinator = collect_dml_logging_records(false); + let frontend = collect_dml_logging_records(true); + + let mut mismatches = Vec::new(); + for (case, ((sql, redacted, mut coordinator), (_, _, frontend))) in std::iter::zip( + DML_LOGGING_PARITY_CASES, + std::iter::zip(coordinator, frontend), + ) { + if case.frontend_only_execution_timestamp { + let timestamps = ( + coordinator + .as_ref() + .map(|record| record.has_execution_timestamp), + frontend + .as_ref() + .map(|record| record.has_execution_timestamp), + ); + assert_eq!( + timestamps, + (Some(false), Some(true)), + "{sql} is marked as recording an execution_timestamp on the frontend path only, \ + but the paths report {timestamps:?}" + ); + coordinator + .as_mut() + .expect("checked above") + .has_execution_timestamp = true; + } + if coordinator != frontend { + mismatches.push(format!( + "{sql} (logged as {redacted}):\n coordinator: {coordinator:?}\n frontend: {frontend:?}" + )); + } + } + assert!( + mismatches.is_empty(), + "statement log rows differ between the coordinator and the frontend OCC path:\n{}", + mismatches.join("\n") + ); +} + +/// Statement-logging outcome of one execution, as recorded once the log +/// flushes. +#[derive(Debug)] +struct StatementOutcome { + finished_status: String, + error_message: Option, +} + +/// Reads the outcomes of every finished execution whose logged SQL matches the +/// `LIKE` pattern `sql_pattern`, waiting for at least one to show up. +fn read_statement_outcomes( + mz_client: &mut postgres::Client, + sql_pattern: &str, +) -> Vec { + Retry::default() + .max_duration(Duration::from_secs(60)) + .retry(|_| { + let rows = mz_client + .query( + "SELECT mseh.finished_status, mseh.error_message +FROM mz_internal.mz_statement_execution_history AS mseh +JOIN mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id +JOIN (SELECT DISTINCT sql_hash, sql FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE mst.sql LIKE $1 AND mseh.finished_at IS NOT NULL", + &[&sql_pattern], + ) + .unwrap(); + if rows.is_empty() { + return Err(()); + } + Ok(rows + .into_iter() + .map(|row| StatementOutcome { + finished_status: row.get(0), + error_message: row.get(1), + }) + .collect::>()) + }) + .unwrap_or_else(|_| panic!("no finished log row matching {sql_pattern}")) +} + +// A cancelled frontend-sequenced write must be logged as the error the user +// received. Recording `aborted` instead loses the reason: `aborted` is the +// status for an execution whose outcome we never learned, and it carries no +// error message. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_cancel_frontend_read_then_write() { + let harness = 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( + "unsafe_enable_unsafe_functions".to_string(), + "true".to_string(), + ); + let (server, mut client) = setup_statement_logging_core(1.0, 1.0, "", harness); + + client + .batch_execute("CREATE TABLE cancel_logging_t (a TEXT, ts INT)") + .unwrap(); + client + .batch_execute("INSERT INTO cancel_logging_t VALUES ('hello', 10)") + .unwrap(); + + let cancel_token = client.cancel_token(); + let (shutdown_tx, shutdown_rx) = std::sync::mpsc::channel(); + let cancel_thread = thread::spawn(move || { + // The write below sleeps for ten seconds, so the first cancel lands + // well after it registered its cancellation watch. Cancelling before + // that would exercise a different exit. + thread::sleep(Duration::from_secs(1)); + loop { + match shutdown_rx.try_recv() { + Ok(()) | Err(std::sync::mpsc::TryRecvError::Disconnected) => return, + Err(std::sync::mpsc::TryRecvError::Empty) => { + let _ = cancel_token.cancel_query(postgres::NoTls); + } + } + thread::sleep(Duration::from_millis(500)); + } + }); + + let err = client + .batch_execute( + "INSERT INTO cancel_logging_t SELECT a, CASE WHEN mz_unsafe.mz_sleep(ts) > 0 THEN 0 END AS ts FROM cancel_logging_t", + ) + .unwrap_err(); + assert_eq!(err.code(), Some(&SqlState::QUERY_CANCELED)); + + shutdown_tx.send(()).unwrap(); + cancel_thread.join().unwrap(); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + let outcomes = read_statement_outcomes(&mut mz_client, "INSERT INTO cancel_logging_t SELECT%"); + assert_eq!(outcomes.len(), 1, "unexpected log rows: {outcomes:?}"); + let outcome = &outcomes[0]; + assert_ne!( + outcome.finished_status, "aborted", + "cancellation must not be recorded as an unknown outcome: {outcome:?}" + ); + assert_eq!(outcome.finished_status, "error", "{outcome:?}"); + assert_eq!( + outcome.error_message.as_deref(), + Some("canceling statement due to user request"), + "{outcome:?}" + ); +} + +// Same as cancellation, for the statement timeout: the log must carry the +// error the user received, not `aborted` with no message. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_timeout_frontend_read_then_write() { + let harness = 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( + "unsafe_enable_unsafe_functions".to_string(), + "true".to_string(), + ); + let (server, mut client) = setup_statement_logging_core(1.0, 1.0, "", harness); + + client + .batch_execute("CREATE TABLE timeout_logging_t (a TEXT, ts INT)") + .unwrap(); + client + .batch_execute("INSERT INTO timeout_logging_t VALUES ('hello', 10)") + .unwrap(); + client + .batch_execute("SET statement_timeout = '5s'") + .unwrap(); + + let err = client + .batch_execute( + "INSERT INTO timeout_logging_t SELECT a, CASE WHEN mz_unsafe.mz_sleep(ts) > 0 THEN 0 END AS ts FROM timeout_logging_t", + ) + .unwrap_err(); + assert_contains!(err.to_string_with_causes(), "statement timeout"); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + let outcomes = read_statement_outcomes(&mut mz_client, "INSERT INTO timeout_logging_t SELECT%"); + assert_eq!(outcomes.len(), 1, "unexpected log rows: {outcomes:?}"); + let outcome = &outcomes[0]; + assert_ne!( + outcome.finished_status, "aborted", + "a statement timeout must not be recorded as an unknown outcome: {outcome:?}" + ); + assert_eq!(outcome.finished_status, "error", "{outcome:?}"); + assert_eq!( + outcome.error_message.as_deref(), + Some("canceling statement due to statement timeout"), + "{outcome:?}" + ); +} + +// A statement that fails before the frontend commits to executing it is still +// the frontend's to record: the coordinator never sees it, so nothing else +// would. DML in an explicit transaction block is such a statement, rejected +// right after the frontend takes it over. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_frontend_read_then_write_transaction_error() { + 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 + .batch_execute("CREATE TABLE txn_error_t (x INT)") + .unwrap(); + client.batch_execute("BEGIN").unwrap(); + let err = client.batch_execute("DELETE FROM txn_error_t").unwrap_err(); + assert_contains!( + err.to_string_with_causes(), + "cannot be run inside a transaction block" + ); + client.batch_execute("ROLLBACK").unwrap(); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + let outcomes = read_statement_outcomes(&mut mz_client, "DELETE FROM txn_error_t"); + assert_eq!(outcomes.len(), 1, "unexpected log rows: {outcomes:?}"); + let outcome = &outcomes[0]; + assert_eq!(outcome.finished_status, "error", "{outcome:?}"); + assert_contains!( + outcome.error_message.as_deref().unwrap_or_default(), + "DELETE FROM txn_error_t cannot be run inside a transaction block" + ); +} + +// An RBAC denial is another exit that happens after the frontend takes the +// statement over. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_frontend_read_then_write_rbac_error() { + 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.batch_execute("CREATE TABLE rbac_t (x INT)").unwrap(); + + // The grants go through the system user: the session that owns the table + // does not own the schema, database and cluster the role also needs. + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + mz_client + .batch_execute("CREATE ROLE rbac_role INHERIT") + .unwrap(); + // Everything the UPDATE needs except UPDATE on the table itself, so that + // the privilege check on the table is what rejects it. + for grant in [ + "GRANT SELECT ON TABLE rbac_t TO rbac_role", + "GRANT USAGE ON SCHEMA public TO rbac_role", + "GRANT USAGE ON DATABASE materialize TO rbac_role", + "GRANT USAGE ON CLUSTER quickstart TO rbac_role", + ] { + mz_client.batch_execute(grant).unwrap(); + } + + let mut rbac_client = server + .pg_config() + .user("rbac_role") + .connect(postgres::NoTls) + .unwrap(); + let err = rbac_client + .batch_execute("UPDATE rbac_t SET x = 1") + .unwrap_err(); + assert_contains!(err.to_string_with_causes(), "permission denied"); + + let outcomes = read_statement_outcomes(&mut mz_client, "UPDATE rbac_t%"); + assert_eq!(outcomes.len(), 1, "unexpected log rows: {outcomes:?}"); + let outcome = &outcomes[0]; + assert_eq!(outcome.finished_status, "error", "{outcome:?}"); + assert_contains!( + outcome.error_message.as_deref().unwrap_or_default(), + "permission denied" + ); +} + +// A prepared DML statement that no frontend path handles. `EXECUTE` is +// unrolled in the session task, which takes over the EXECUTE's log entry, and +// the inner statement then falls back to the coordinator. The coordinator has +// to receive that entry and finish it, and the two statements have to be +// counted once each: the EXECUTE by the session task, the inner statement by +// the coordinator. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_prepared_dml_coordinator_fallback() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + + client + .batch_execute("CREATE TABLE prepared_dml_t (x INT)") + .unwrap(); + client + .batch_execute("PREPARE p AS INSERT INTO prepared_dml_t VALUES (1)") + .unwrap(); + + let insert_labels = [("session_type", "user"), ("statement_type", "insert")]; + let execute_labels = [("session_type", "user"), ("statement_type", "execute")]; + let inserts_before = + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &insert_labels); + let executes_before = + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &execute_labels); + + client.batch_execute("EXECUTE p").unwrap(); + + let rows: i64 = client + .query_one("SELECT count(*) FROM prepared_dml_t", &[]) + .unwrap() + .get(0); + assert_eq!(rows, 1); + assert_eq!( + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &insert_labels), + inserts_before + 1 + ); + assert_eq!( + test_util::get_counter_value(server.metrics_registry(), "mz_query_total", &execute_labels), + executes_before + 1 + ); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + let outcomes = read_statement_outcomes(&mut mz_client, "EXECUTE p"); + assert_eq!(outcomes.len(), 1, "unexpected log rows: {outcomes:?}"); + assert_eq!(outcomes[0].finished_status, "success", "{:?}", outcomes[0]); +} diff --git a/src/sql/src/session/vars.rs b/src/sql/src/session/vars.rs index 1cdc5dd702056..e00d20d2307f6 100644 --- a/src/sql/src/session/vars.rs +++ b/src/sql/src/session/vars.rs @@ -1225,6 +1225,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, @@ -1780,6 +1782,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 09e49d767be7e..01dd08603cd92 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), diff --git a/test/0dt/mzcompose.py b/test/0dt/mzcompose.py index c15458d8213aa..d9f3f6634633b 100644 --- a/test/0dt/mzcompose.py +++ b/test/0dt/mzcompose.py @@ -356,6 +356,14 @@ def workflow_read_only(c: Composition) -> None: 1 ! INSERT INTO t VALUES (3, 4); contains: cannot write in read-only mode + # A read-then-write reaches the write through a different path than + # a constant INSERT, so each shape needs its own rejection here. + ! DELETE FROM t WHERE a = 1; + contains: cannot write in read-only mode + ! UPDATE t SET b = b + 1; + contains: cannot write in read-only mode + ! INSERT INTO t SELECT a, b FROM t; + contains: cannot write in read-only mode > SET TRANSACTION_ISOLATION TO 'SERIALIZABLE'; > SELECT * FROM mv; 1 @@ -686,6 +694,14 @@ def workflow_basic(c: Composition) -> None: 1 ! INSERT INTO t VALUES (3, 4); contains: cannot write in read-only mode + # A read-then-write reaches the write through a different path than + # a constant INSERT, so each shape needs its own rejection here. + ! DELETE FROM t WHERE a = 1; + contains: cannot write in read-only mode + ! UPDATE t SET b = b + 1; + contains: cannot write in read-only mode + ! INSERT INTO t SELECT a, b FROM t; + contains: cannot write in read-only mode > SET TRANSACTION_ISOLATION TO 'SERIALIZABLE'; > SELECT * FROM mv; 1 @@ -843,6 +859,14 @@ def workflow_basic(c: Composition) -> None: 1 ! INSERT INTO t VALUES (3, 4); contains: cannot write in read-only mode + # A read-then-write reaches the write through a different path than + # a constant INSERT, so each shape needs its own rejection here. + ! DELETE FROM t WHERE a = 1; + contains: cannot write in read-only mode + ! UPDATE t SET b = b + 1; + contains: cannot write in read-only mode + ! INSERT INTO t SELECT a, b FROM t; + contains: cannot write in read-only mode > SET TRANSACTION_ISOLATION TO 'SERIALIZABLE'; > SELECT * FROM mv; 9 @@ -2494,6 +2518,14 @@ def workflow_ddl(c: Composition) -> None: 1 ! INSERT INTO t VALUES (3, 4); contains: cannot write in read-only mode + # A read-then-write reaches the write through a different path than + # a constant INSERT, so each shape needs its own rejection here. + ! DELETE FROM t WHERE a = 1; + contains: cannot write in read-only mode + ! UPDATE t SET b = b + 1; + contains: cannot write in read-only mode + ! INSERT INTO t SELECT a, b FROM t; + contains: cannot write in read-only mode > SET TRANSACTION_ISOLATION TO 'SERIALIZABLE'; > SELECT * FROM mv; 1 @@ -2659,6 +2691,14 @@ def workflow_ddl(c: Composition) -> None: 1 ! INSERT INTO t VALUES (3, 4); contains: cannot write in read-only mode + # A read-then-write reaches the write through a different path than + # a constant INSERT, so each shape needs its own rejection here. + ! DELETE FROM t WHERE a = 1; + contains: cannot write in read-only mode + ! UPDATE t SET b = b + 1; + contains: cannot write in read-only mode + ! INSERT INTO t SELECT a, b FROM t; + contains: cannot write in read-only mode > SET TRANSACTION_ISOLATION TO 'SERIALIZABLE'; > SELECT * FROM mv; 9 diff --git a/test/cluster/mzcompose.py b/test/cluster/mzcompose.py index 68bd562e5a380..7200b228aa567 100644 --- a/test/cluster/mzcompose.py +++ b/test/cluster/mzcompose.py @@ -4638,6 +4638,416 @@ def subscriber() -> None: ), "statement execution was ended twice; end-of-execution ownership handoff regressed" +def workflow_test_occ_zero_row_write_linearization(c: Composition) -> None: + """A read-then-write that reports zero rows must not retire before the write + that emptied its selection is readable through the timestamp oracle. + + The OCC path linearizes only its initial `as_of`, and its internal subscribe + follows Persist visibility, which runs ahead of the oracle: the group + committer appends before it applies the write timestamp. A DELETE or UPDATE + can therefore consolidate its selection to empty against state no + oracle-timestamped read can reach yet, report zero rows through + `NoRowsMatched`, and return. A strict-serializable read issued after that + response then still sees the row the response said was not there, and no + serial order explains that history. + + The `group_commit_before_apply_write` failpoint holds the winning writer + inside that window, the same one a second `environmentd` process opens on + its own with no ordering against local Persist visibility. + """ + + # Every txns-shard write parks here while armed, including the keepalives + # that advance table uppers, so this has to be a bounded `sleep` and not a + # `pause`: a keepalive would take the `pause` first and the winning DELETE + # would never get to append. The window only has to outlast a peek and one + # subscribe dataflow installation. + failpoint = "group_commit_before_apply_write" + arm = f"SET failpoints = '{failpoint}=sleep(10000)'" + disarm = f"SET failpoints = '{failpoint}=off'" + + def occ_writes() -> tuple[int, int]: + """Read-then-writes the OCC path sequenced, and how many of their write + attempts lost the race for their write timestamp.""" + metric = "mz_occ_read_then_write_retry_count" + metrics = c.exec( + "materialized", "curl", "localhost:6878/metrics", capture=True + ).stdout + values = { + line.split()[0]: int(float(line.split()[1])) + for line in metrics.splitlines() + if line.startswith((f"{metric}_count ", f"{metric}_sum ")) + } + return values[f"{metric}_count"], values[f"{metric}_sum"] + + def count(cur: Cursor, key: int) -> int: + cur.execute(f"SELECT count(*) FROM t WHERE k = {key}".encode()) + row = cur.fetchone() + assert row is not None + return int(row[0]) + + with c.override( + Materialized( + # Sampled once at startup, so this cannot be an `ALTER SYSTEM SET`. + additional_system_parameter_defaults={ + "enable_adapter_frontend_occ_read_then_write": "true" + }, + ) + ): + c.up("materialized") + c.sql("CREATE TABLE t (k int, v int)") + + # Ask the process rather than the catalog which path it takes: the UPDATE + # only reaches the histogram if the frontend sequenced it. + sequenced = occ_writes()[0] + c.sql("UPDATE t SET v = v + 1 WHERE k = 0") + assert occ_writes()[0] > sequenced, ( + "the UPDATE did not go through the OCC path, so this would exercise the " + "coordinator's lock-based path instead" + ) + + # Connections are opened before the failpoint is armed: starting a session + # appends to `mz_sessions`, which parks like any other write. + with ( + c.sql_cursor() as control, + c.sql_cursor() as probe, + c.sql_cursor() as winner_cur, + c.sql_cursor() as session, + ): + # A serializable read may pick a timestamp past the oracle's read + # timestamp, so it sees the winner's append while a strict-serializable + # read cannot. + probe.execute("SET transaction_isolation = 'serializable'") + + # The UPDATE's subscribe can instead report progress at the table's + # pre-append upper, with the row still there, then submit a write, + # queue behind the parked committer, and conclude zero rows only once + # the oracle has caught up. That proves nothing either way, so such an + # attempt is retried; the write conflict count tells the two apart. + for attempt in range(1, 4): + key = attempt + c.sql(f"INSERT INTO t VALUES ({key}, 1)") + armed = Event() + + def delete_winner(key: int = key) -> None: + armed.wait() + # Lands its append and advances t's upper, then the committer + # parks before applying that timestamp to the oracle. + winner_cur.execute(f"DELETE FROM t WHERE k = {key}".encode()) + + winner = PropagatingThread(target=delete_winner, name="winner") + winner.start() + control.execute(arm) + armed.set() + + # The winner's append is visible in Persist from here on ... + deadline = time.time() + 120 + while count(probe, key) > 0: + assert ( + time.time() < deadline + ), "the winning DELETE never became visible in Persist" + time.sleep(0.1) + # ... and the oracle cannot serve reads at it yet, which is what + # puts us inside the window. This witness says nothing about the + # UPDATE, so it stays valid once the zero-row path waits. + before = count(session, key) + + conflicts = occ_writes()[1] + started = time.time() + session.execute(f"UPDATE t SET v = v + 1 WHERE k = {key}".encode()) + matched = session.rowcount + elapsed = time.time() - started + after = count(session, key) + conflicts = occ_writes()[1] - conflicts + + control.execute(disarm) + # `off` does not interrupt a `sleep` under way, so this waits out + # the rest of the window. + winner.join(timeout=120) + assert not winner.is_alive(), "the winning DELETE never finished" + + print( + f"attempt {attempt}: UPDATE matched {matched} row(s) in " + f"{elapsed:.1f}s with {conflicts} write conflict(s); " + f"strict-serializable reads saw {before} row(s) before it and " + f"{after} row(s) after" + ) + assert matched == 0, ( + f"the UPDATE matched {matched} row(s) with the winner's delete " + "already visible, so it never took the zero-row path" + ) + if before == 0 or conflicts > 0: + continue + + assert after == 0, ( + "the UPDATE reported zero rows matched from state the oracle had " + f"not applied yet, and a strict-serializable read after it saw " + f"{after} row(s): the zero-row response was not linearized against " + "the write that emptied the selection" + ) + break + else: + raise AssertionError( + "no attempt got the UPDATE to report zero rows from the newer " + "state without first submitting a write, so the window was never " + "observed" + ) + + +def occ_sequenced_writes(c: Composition) -> int: + """How many read-then-writes the OCC path has sequenced. + + Observed once per read-then-write the frontend sequenced, so this is what + tells the two sequencing paths apart from outside the process. The + histogram is registered unconditionally, so a missing line means the scrape + itself did not land.""" + metric = "mz_occ_read_then_write_retry_count_count" + metrics = c.exec( + "materialized", "curl", "--silent", "localhost:6878/metrics", capture=True + ).stdout + for line in metrics.splitlines(): + if line.startswith(f"{metric} "): + return int(float(line.split()[1])) + raise RuntimeError(f"{metric} not found in materialized metrics") + + +def workflow_test_occ_read_then_write_dependency_dropped(c: Composition) -> None: + """A read-then-write whose selection reads a collection that is dropped + while the statement runs must report that, write nothing, and leave neither + its OCC permit nor its subscribe behind. + + The OCC path does not peek its selection, it installs a subscribe and reads + that, so a dropped read dependency does not reach the statement as a + planning error. It arrives as the subscribe's own termination, which the + loop has to turn into the statement's error instead of waiting for a + frontier that will never advance again. + + The sleep in the selection holds the loop in that wait. It blocks the worker + rendering the dataflow, so the dataflow only goes away once the sleep + returns, but retiring the subscribe is the coordinator's own bookkeeping and + does not wait for the cluster. That is what makes the deadline below a + statement about the adapter rather than about the sleep. + """ + + # The selection cannot finish on its own inside the deadline, so a DELETE + # that ends within it ended because of the DROP. + sleep_seconds = 30 + deadline_seconds = 15 + + with c.override( + Materialized( + # Sampled once at startup, so these cannot be an `ALTER SYSTEM SET`. + # One permit makes the follow-up write below depend on the failed + # statement having released the one it held. + additional_system_parameter_defaults={ + "enable_adapter_frontend_occ_read_then_write": "true", + "max_concurrent_occ_writes": "1", + "unsafe_enable_unsafe_functions": "true", + }, + ) + ): + c.up("materialized") + # `dependency` holds a single row, so the sleep in the selection below + # runs once rather than once per row, and it takes its duration from + # that row rather than from a literal, or the optimizer folds the + # predicate away before the dataflow ever sleeps. + c.sql(dedent(f""" + CREATE TABLE target (k int); + INSERT INTO target SELECT generate_series(1, 10); + CREATE TABLE dependency (k int, delay double precision); + INSERT INTO dependency VALUES (1, {sleep_seconds}); + CREATE CLUSTER other SIZE 'scale=1,workers=1'; + """)) + + # Ask the process rather than the catalog which path it takes: the + # DELETE only reaches the histogram if the frontend sequenced it. + sequenced = occ_sequenced_writes(c) + c.sql("DELETE FROM target WHERE k IN (SELECT k FROM dependency WHERE k < 0)") + assert occ_sequenced_writes(c) > sequenced, ( + "the DELETE did not go through the OCC path, so this would exercise the " + "coordinator's lock-based path instead" + ) + + outcome: list[str] = [] + + def delete() -> None: + with c.sql_cursor() as cur: + try: + cur.execute( + b"DELETE FROM target WHERE k IN (SELECT k FROM dependency " + b"WHERE mz_unsafe.mz_sleep(delay) IS NULL)" + ) + outcome.append("committed") + except DatabaseError as e: + outcome.append(str(e)) + + deleter = PropagatingThread(target=delete, name="deleter") + deleter.start() + # The statement has to be past planning and inside its subscribe before + # the dependency goes away, otherwise this tests planning instead. + time.sleep(5) + dropped_at = time.time() + c.sql("DROP TABLE dependency") + + deleter.join(timeout=deadline_seconds) + assert not deleter.is_alive(), ( + f"the DELETE was still running {time.time() - dropped_at:.0f}s after its read " + "dependency was dropped, so the subscribe's termination never reached the loop" + ) + [result] = outcome + assert ( + "was dropped" in result + ), f"the DELETE ended with {result!r} instead of reporting the dropped dependency" + + # Both checks run on the other cluster: the sleep still holds the worker + # this statement's subscribe ran on, and waiting that out here would + # blur a leaked permit into a busy replica. + with c.sql_cursor() as cur: + cur.execute(b"SET cluster = other") + cur.execute(b"SELECT count(*) FROM target") + row = cur.fetchone() + assert ( + row is not None and row[0] == 10 + ), f"target holds {row}, expected the 10 rows a failed DELETE leaves behind" + + started = time.time() + cur.execute(b"DELETE FROM target WHERE k = 1") + assert ( + cur.rowcount == 1 + ), f"the follow-up DELETE affected {cur.rowcount} rows, expected 1" + elapsed = time.time() - started + assert elapsed < deadline_seconds, ( + f"the follow-up read-then-write took {elapsed:.0f}s, which is what the " + "single OCC permit looks like when the failed statement kept it" + ) + + # The subscribe's dataflow goes away once the sleep releases the worker, + # so this waits that out. A `SubscribeHandle` whose drop never sent the + # cleanup leaves it installed for good. + deadline = time.time() + 120 + with c.sql_cursor() as cur: + while True: + cur.execute( + b"SELECT count(*) FROM mz_introspection.mz_dataflows " + b"WHERE name LIKE '%frontend-read-then-write%'" + ) + row = cur.fetchone() + assert row is not None + if row[0] == 0: + break + assert time.time() < deadline, ( + f"{row[0]} read-then-write subscribe dataflow(s) still installed after " + "the statement that owned them failed" + ) + time.sleep(1) + + +def workflow_test_occ_cancel_of_submitted_write(c: Composition) -> None: + """A cancel that arrives once a read-then-write's write is durable must not + take the statement's answer away from it. + + Having handed its diffs to the group committer, the OCC loop no longer knows + whether they landed, so the cancellation path has to wait for the + committer's answer rather than synthesize `canceling statement` on the spot. + Reporting a cancellation for a write the user can already read back is a + lost acknowledgement, not a cancelled statement. + + The `timestamped_write_before_result` failpoint holds the answer inside that + window, which is otherwise microseconds wide. It sits past the append and + past the oracle advance, so the increment being readable is what says the + window is open. + """ + + failpoint = "timestamped_write_before_result" + arm = f"SET failpoints = '{failpoint}=sleep(10000)'" + disarm = f"SET failpoints = '{failpoint}=off'" + + with c.override( + Materialized( + # Sampled once at startup, so this cannot be an `ALTER SYSTEM SET`. + additional_system_parameter_defaults={ + "enable_adapter_frontend_occ_read_then_write": "true" + }, + ) + ): + c.up("materialized") + c.sql("CREATE TABLE t (v int)") + c.sql("INSERT INTO t VALUES (0)") + + # Ask the process rather than the catalog which path it takes: only a + # frontend-sequenced write is a timestamped write, and only a + # timestamped write reaches the failpoint. + sequenced = occ_sequenced_writes(c) + c.sql("UPDATE t SET v = v WHERE v < 0") + assert occ_sequenced_writes(c) > sequenced, ( + "the UPDATE did not go through the OCC path, so the failpoint below would " + "never be reached" + ) + + # Connections are opened before the failpoint is armed: starting a + # session appends to `mz_sessions`, which is a write of its own. + with ( + c.sql_cursor() as control, + c.sql_cursor() as writer, + c.sql_cursor() as reader, + ): + writer.execute(b"SELECT pg_backend_pid()") + row = writer.fetchone() + assert row is not None + pid = int(row[0]) + + outcome: list[str] = [] + done = Event() + + def update() -> None: + try: + writer.execute(b"UPDATE t SET v = v + 1") + outcome.append("committed") + except DatabaseError as e: + outcome.append(str(e)) + finally: + done.set() + + control.execute(arm.encode()) + updater = PropagatingThread(target=update, name="updater") + updater.start() + + # The increment is readable, so the write is durable and the + # committer is holding its answer. + deadline = time.time() + 60 + while True: + reader.execute(b"SELECT v FROM t") + row = reader.fetchone() + assert row is not None + if row[0] == 1: + break + assert time.time() < deadline, ( + "the UPDATE's write never became durable, so the failpoint window " + "was never entered" + ) + time.sleep(0.1) + assert not done.is_set(), ( + "the UPDATE returned before the failpoint held its result, so the cancel " + "below would not race a submitted write" + ) + + control.execute(f"SELECT pg_cancel_backend({pid})".encode()) + updater.join(timeout=120) + assert not updater.is_alive(), "the UPDATE never finished" + control.execute(disarm.encode()) + + [result] = outcome + assert result == "committed", ( + f"the UPDATE reported {result!r} for a write the table already held, so " + "the cancel was answered without waiting for the write's outcome" + ) + + reader.execute(b"SELECT v FROM t") + row = reader.fetchone() + assert ( + row is not None and row[0] == 1 + ), f"the table holds {row}, expected the single increment the UPDATE reported" + + def workflow_test_refresh_mv_warmup( c: Composition, parser: WorkflowArgumentParser ) -> None: diff --git a/test/launchdarkly-flag-consistency/mzcompose.py b/test/launchdarkly-flag-consistency/mzcompose.py index 27ab80604948e..d3163533dc856 100644 --- a/test/launchdarkly-flag-consistency/mzcompose.py +++ b/test/launchdarkly-flag-consistency/mzcompose.py @@ -234,6 +234,7 @@ enable_0dt_caught_up_replica_status_check enable_0dt_caught_up_stability_check enable_0dt_deployment_panic_after_timeout + enable_adapter_frontend_occ_read_then_write enable_alter_table_add_column enable_auto_scaling_strategy enable_background_alter_cluster @@ -302,8 +303,10 @@ keep_n_sink_status_history_entries keep_n_source_status_history_entries log_filter_defaults + max_concurrent_occ_writes max_copy_from_row_size max_network_policies + max_occ_retries max_rules_per_network_policy max_sql_server_connections max_timestamp_interval diff --git a/test/sqllogictest/recursion_limit.slt b/test/sqllogictest/recursion_limit.slt index 12780ececb5b8..1bb544e78e0fe 100644 --- a/test/sqllogictest/recursion_limit.slt +++ b/test/sqllogictest/recursion_limit.slt @@ -10,12 +10,12 @@ # Tests to exercise recursion limits and stack overflow protection in # the query planner and optimizer. -# 418 UNIONs +# 388 UNIONs # NOTE: If this test fails, shorten it by removing some `UNION SELECT 1`s # from the end. Call stacks above the critical recursion can grow as we # add code elsewhere in the system (database-issues#9996). query I -SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1; +SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1 UNION SELECT 1; ---- 1 diff --git a/test/txn-wal-fencing/mzcompose.py b/test/txn-wal-fencing/mzcompose.py index 3e2e254692f11..eb52aa9b4f857 100644 --- a/test/txn-wal-fencing/mzcompose.py +++ b/test/txn-wal-fencing/mzcompose.py @@ -14,6 +14,7 @@ import argparse import random +import threading import time from concurrent import futures from dataclasses import dataclass @@ -96,6 +97,17 @@ class SuccessfulCommit: Materialized(name="mz_second"), ] +# Selects how a process sequences DELETE/UPDATE/INSERT ... SELECT: `false` keeps +# them on the Coordinator behind in-process write locks, `true` sequences them +# from the session task under optimistic concurrency control. The value is +# sampled once at process startup, so two processes in one environment can hold +# different values. +OCC_FLAG = "enable_adapter_frontend_occ_read_then_write" + +# Observed once per read-then-write that the OCC path sequenced, so the +# histogram's sample count identifies which path a process took. +OCC_METRIC = "mz_occ_read_then_write_retry_count_count" + def workflow_default(c: Composition, parser: WorkflowArgumentParser) -> None: parser.add_argument( @@ -281,3 +293,264 @@ def run_workload(c: Composition, workload: Workload, args: argparse.Namespace) - ), f"Unexpected result {result}; commit: {commit}; target {target}" print("Verification complete.") + + +# Connections driving increments. +MIXED_MODE_CONCURRENCY = 16 + +# The statement never reached a server, or a server refused it, so it wrote +# nothing. Keeping these apart from the indeterminate ones below matters: a +# fenced instance produces plenty, and each one widens the band the counter is +# checked against. +NOT_APPLIED = [ + # Targeted a container that is not up, or one that died before the + # connection was established, or one that stopped responding mid-connect. + "running docker compose failed", + "Connection refused", + "Connection timed out", + "canceling statement due to statement timeout", + # How the OCC path reports sustained contention. + "read-then-write exceeded maximum retry attempts under contention", +] + +# The connection went away with the statement in flight, so the increment may or +# may not be durable. A fenced container still publishes its port, so the +# connection is accepted and reset. +INDETERMINATE = [ + "server closed the connection unexpectedly", + "Connection reset by peer", +] + + +class Increment(Enum): + ACKED = 0 + REJECTED = 1 + UNKNOWN = 2 + + +def occ_sequenced_writes(c: Composition, service: str) -> int: + """How many read-then-writes `service` has sequenced through the OCC path.""" + metrics = c.exec( + service, "curl", "--silent", "localhost:6878/metrics", capture=True + ).stdout + for line in metrics.splitlines(): + if line.startswith(f"{OCC_METRIC} "): + return int(float(line.split()[1])) + # The histogram is registered unconditionally, so a missing line means the + # scrape itself did not land. + raise RuntimeError(f"{OCC_METRIC} not found in {service} metrics") + + +def increment_counter(args: tuple[Composition, str, bool]) -> Increment: + """Increments the shared counter by one through a read-then-write. + + `slow` stretches the read phase, so that the operation can straddle the + moment the second instance comes up. The sleep takes its duration from the + `delay` column rather than from a literal, so that it is not folded away at + plan time and instead runs while the selection is read. + """ + c, mz_service, slow = args + sleep = " AND mz_unsafe.mz_sleep(delay) IS NULL" if slow else "" + try: + c.sql_cursor(service=mz_service).execute( + f"UPDATE counter SET v = v + 1 WHERE k = 1{sleep}".encode() + ) + except Exception as e: + if any(msg in str(e) for msg in NOT_APPLIED): + return Increment.REJECTED + if any(msg in str(e) for msg in INDETERMINATE): + return Increment.UNKNOWN + raise RuntimeError(f"unexpected exception: {e}") + return Increment.ACKED + + +def workflow_mixed_mode_read_then_write( + c: Composition, parser: WorkflowArgumentParser +) -> None: + """Two instances in one environment sequencing read-then-write differently. + + `OCC_FLAG` is sampled once per process, so a rolling restart or a newly added + serving process can leave one instance on the Coordinator's write-lock path + and another on the OCC path. The two do not synchronize: the lock path + excludes concurrent writers, the OCC path detects them afterwards from the + write timestamp. A blind write from the lock path landing on top of an OCC + write leaves the row with a negative copy of the stale value and two copies + of the new one. + + Both orderings run, because they put the lock path on opposite sides of the + handover. Each one goes red on a lost update or a broken multiplicity, and + also on the precondition for either: the two instances committing at the same + time. + """ + parser.add_argument( + "--azurite", action="store_true", help="Use Azurite as blob store instead of S3" + ) + args = parser.parse_args() + + # Every increment opens its own connection, so leave out the per-invocation + # Docker Compose echo. + c.silent = True + + for first_occ, second_occ in [("false", "true"), ("true", "false")]: + print( + f"+++ Running with {OCC_FLAG} {first_occ} on 'mz_first', {second_occ} on 'mz_second' ..." + ) + run_mixed_mode(c, args.azurite, first_occ, second_occ) + + +def run_mixed_mode( + c: Composition, azurite: bool, first_occ: str, second_occ: str +) -> None: + """Runs one ordering: 'mz_first' comes up first, then 'mz_second' displaces it.""" + c.down(destroy_volumes=True) + c.up(c.metadata_store()) + + with c.override( + *[ + Materialized( + name=mz_name, + external_metadata_store=True, + external_blob_store=True, + blob_store_is_azure=azurite, + sanity_restart=False, + support_external_clusterd=True, + additional_system_parameter_defaults={OCC_FLAG: occ}, + ) + for mz_name, occ in [("mz_first", first_occ), ("mz_second", second_occ)] + ] + ): + c.up("mz_first") + + # Idempotent, because a retried connection re-runs the whole batch after + # the statements it already applied. `mz_sleep` blocks the timely worker + # it runs on, so `delay` stays short enough that the slow increments do + # not starve the replica. + c.sql( + """ + CREATE TABLE IF NOT EXISTS counter (k int, v bigint, delay double precision); + DELETE FROM counter; + INSERT INTO counter VALUES (1, 0, 0.5); + """, + service="mz_first", + ) + + print("--- Confirming the instances disagree on how to sequence") + assert ( + increment_counter((c, "mz_first", False)) == Increment.ACKED + ), "baseline increment on 'mz_first' did not commit" + occ_writes = occ_sequenced_writes(c, "mz_first") + assert (occ_writes > 0) == (first_occ == "true"), ( + f"'mz_first' sequenced {occ_writes} read-then-writes through OCC, " + f"which does not match {OCC_FLAG}={first_occ}" + ) + + print("--- Driving increments across both instances") + stop = threading.Event() + + def drive(worker: int) -> list[tuple[str, Increment, float, float]]: + # Worker 0 keeps a slow read-then-write in flight on 'mz_first' for + # the whole run, so that 'mz_second' comes up while an operation + # there sits between its read and its write. The workers aimed at + # 'mz_second' start before it can serve and take their rejections + # until it can. + mz_service = "mz_first" if worker % 2 == 0 else "mz_second" + op = (c, mz_service, worker == 0) + outcomes = [] + while not stop.is_set(): + issued = time.time() + outcomes.append( + (mz_service, increment_counter(op), issued, time.time()) + ) + return outcomes + + with futures.ThreadPoolExecutor(MIXED_MODE_CONCURRENCY) as executor: + drivers = [ + executor.submit(drive, worker) + for worker in range(MIXED_MODE_CONCURRENCY) + ] + try: + time.sleep(2) + c.up("mz_second") + # The fence lands while 'mz_second' opens the catalog, well + # before it reports healthy, so keep driving past that point to + # cover the handover from both sides. + time.sleep(15) + finally: + stop.set() + outcomes = [outcome for driver in drivers for outcome in driver.result()] + + # The baseline increment on 'mz_first' counts too. + acked = 1 + sum( + 1 for _, outcome, _, _ in outcomes if outcome == Increment.ACKED + ) + unknown = sum( + 1 for _, outcome, _, _ in outcomes if outcome == Increment.UNKNOWN + ) + print(f"acked: {acked}, unknown: {unknown}") + + occ_writes = occ_sequenced_writes(c, "mz_second") + assert (occ_writes > 0) == (second_occ == "true"), ( + f"'mz_second' sequenced {occ_writes} read-then-writes through OCC, " + f"which does not match {OCC_FLAG}={second_occ}" + ) + + # Without a commit from each side the run exercised one instance only. + acks = { + mz_service: [ + (issued, at) + for service, outcome, issued, at in outcomes + if service == mz_service and outcome == Increment.ACKED + ] + for mz_service in ["mz_first", "mz_second"] + } + for mz_service, times in acks.items(): + assert times, f"'{mz_service}' committed no increment" + + # What keeps the two modes from corrupting the row today is that they + # never commit at the same time: both paths advance the catalog upper + # before every write, so an instance stops committing once the other has + # fenced it, and the fence lands before the other instance reads anything. + # Overlapping commit windows are the bug's precondition, because a + # lock-path write can then land on top of an OCC write and leave the row + # with a negative copy of the stale value and two copies of the new one. + # Sampling the flag once per process does not prevent that, so a red + # assertion here means the modes have to be fenced against each other + # rather than merely fixed per process. + # Comparing when 'mz_first' last *issued* a statement that went on to + # commit against when 'mz_second' first committed keeps a slow response + # from reading as an overlap: a statement issued after the other instance + # had already committed cannot have committed before it. + gap = min(at for _, at in acks["mz_second"]) - max( + issued for issued, _ in acks["mz_first"] + ) + print(f"'mz_first' stopped committing {gap:.1f}s before 'mz_second' started") + assert gap > 0, ( + f"'mz_first' committed a statement it issued {-gap:.1f}s after " + f"'mz_second' had committed, so both sequencing modes were live at once" + ) + + print("--- Verifying the counter") + cursor = c.sql_cursor(service="mz_second") + + # A negative copy of the stale value has no rendering, so a broken + # multiplicity surfaces here as a missing row, an extra row, or a + # retraction error. + cursor.execute("SELECT v FROM counter WHERE k = 1") + rows = cursor.fetchall() + assert len(rows) == 1, f"counter holds {rows}, expected exactly one row" + + # An acked increment is durable, one whose connection died may or may not + # be. Anything below `acked` is a lost update. + v = rows[0][0] + assert ( + acked <= v <= acked + unknown + ), f"counter is {v}, expected between {acked} and {acked + unknown}" + + # Checked last so that a lost update is reported as such rather than as a + # missing fence. The fence bounds the window in which the two modes + # overlap: both paths advance the catalog upper before every write, so an + # instance stops writing once the other one has fenced it. + log = c.invoke("logs", "mz_first", capture=True).stdout + assert ( + "unable to advance catalog upper" in log or "fenced by envd" in log + ), "'mz_first' was never fenced, so the two instances never overlapped"