mz_adapter/frontend_read_then_write.rs
1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Frontend sequencing for read-then-write operations.
11//!
12//! This module implements INSERT [...] SELECT FROM [...], DELETE and UPDATE
13//! operations using a subscribe with optimistic concurrency control (OCC),
14//! sequenced from the session task rather than the Coordinator.
15//!
16//! The motivation is correctness with concurrent writers, including writers in
17//! different `environmentd` processes, which the coordinator's in-process write
18//! locks cannot provide. The OCC path also fixes a serializability defect of the
19//! lock path, which reads at one timestamp and commits at a later one while
20//! locking only the selection's direct dependency items. See the design doc,
21//! `doc/developer/design/20260210_incremental_occ_read_then_write.md`, and the
22//! comment on the retry arm in `run_occ_loop`. Relieving the coordinator loop is
23//! a side benefit, and only sequencing moves off it. The subscribe's data path
24//! still runs through the coordinator.
25//!
26//! ## Whether the write reads persisted state
27//!
28//! Two predicates answer that one question, and they have to agree. Before
29//! anything runs, `SessionClient::try_frontend_read_then_write` decides it
30//! syntactically, from `depends_on()` on the planned selection, because inside
31//! a transaction a read-dependent write has to be refused while refusing is
32//! still possible. Once the dataflow runs, the subscribe answers it
33//! dynamically: the channel closes on its own only once the sink's output
34//! frontier reaches the empty antichain, which means the selection can never
35//! change again.
36//!
37//! Reading nothing persisted is the common case for a clean close, not the
38//! guarantee. An input whose frontier seals closes cleanly too, despite reading
39//! persisted state, for example a `REFRESH AT` materialized view past its last
40//! refresh, whose write frontier advances to the empty antichain. What holds in
41//! either case is the property the write side actually needs: past the close,
42//! the consolidated diffs are frontier-independent. The syntactic predicate is
43//! correspondingly stricter than the dynamic one, since it refuses a sealed-MV
44//! `INSERT ... SELECT` in a transaction that would technically be bufferable.
45//!
46//! The two answers are used for different things, and that separation matters.
47//!
48//! The syntactic answer decides whether the statement can belong to a
49//! transaction. A selection that reads nothing produces diffs that are valid at
50//! any timestamp, so they are staged as session write ops and land when the
51//! transaction commits, which is what makes the statement atomic with whatever
52//! surrounds it. A selection that reads persisted state cannot belong to a
53//! transaction, and we refuse it in an explicit one. An extended-protocol
54//! pipeline is an implicit transaction, so it must not quietly join one either:
55//! it ends its own transaction instead of spanning the rest of the pipeline,
56//! which is how PostgreSQL treats statements that cannot run in a transaction
57//! block. It is durable once it reports success, and a later failure in the
58//! pipeline does not undo it.
59//!
60//! The dynamic answer only decides how the write is submitted, a timestamped
61//! write from inside the loop or a blind submission after it. It cannot decide
62//! transaction membership, because it is a property of the inputs rather than
63//! of the statement. A sealed input closes the subscribe cleanly, so an
64//! `INSERT ... SELECT` over a `REFRESH AT` materialized view past its last
65//! refresh takes the blind exit while still reading persisted state. Its diffs
66//! really are frontier-independent and staging them would be safe, and we still
67//! do not stage them. Otherwise whether a statement's rows survive a failure
68//! later in the pipeline would depend on whether one of its inputs happened to
69//! pass its last refresh, which no one reading the statement could predict.
70//!
71//! Staging is also what earns the right to span a pipeline in the first place.
72//! `TransactionStatus::may_span_pipeline` lets an implicit transaction stay open
73//! only for writes, precisely because they are merely staged. A statement that
74//! committed on its own has no business claiming it.
75//!
76//! Disagreement is caught on both sides, and only one side can still refuse.
77//! `frontend_read_then_write` re-checks the syntactic predicate before running a
78//! dataflow, which catches a caller that skipped the gate. If the syntactic
79//! predicate were laxer than the dynamic one, that check would pass and a write
80//! meant for staging would commit on its own, so the loop asserts wherever the
81//! two answers can be compared. The `Committed` arm catches a write timestamp
82//! for a statement we meant to stage, and the zero-row arm catches a read
83//! timestamp for one. Neither can undo anything by then, the write is already
84//! durable in the first case and there was never anything to write in the
85//! second, so all they do is make the disagreement loud.
86//!
87//! ## The frontier certifies, the oracle chooses
88//!
89//! The target `T` comes from the oracle, and a progress message at `F` only
90//! certifies completeness below `F`, so it gates the write rather than choosing
91//! its timestamp. The design doc's "The OCC loop" says why. Three invariants
92//! hold for every write this path makes:
93//!
94//! * `F >= T` before it submits, so the payload is a complete view of `T - 1`.
95//! * The payload is every diff below `T`, strictly. A diff at `T` is concurrent
96//! with the write and waits for a later target.
97//! * `T > as_of`, so the snapshot, which arrives at `as_of`, is in the payload.
98//!
99//! NOTE: `F >= T` does not make the two equal, because `F` is a minimum over the
100//! selection's inputs. Where `F` runs above `T`, a selection that reads the
101//! target table makes the compare-and-append refuse, and retrying higher is the
102//! design rather than a failure.
103//!
104//! ## A zero-row answer
105//!
106//! A write linearizes itself, because group commit advances the oracle before it
107//! acknowledges. An answer of "no rows" performs no write, so the loop reports
108//! the timestamp its view is complete through and the caller waits for the
109//! oracle to reach it. A selection empty at `as_of` is complete through `as_of`,
110//! which the caller put behind the oracle before the subscribe started, so that
111//! answer waits for nothing, while reporting the frontier the subscribe observed
112//! would cost a group commit every time. The design doc's "Linearization" argues
113//! why the lower timestamp is not the weaker guarantee.
114//!
115//! ## Rollout note
116//!
117//! The `FRONTEND_READ_THEN_WRITE` dyncfg is read once at process startup and
118//! fixed for the lifetime of the `environmentd` process. This avoids a
119//! mixed-mode window where both the lock-based coordinator path and this OCC
120//! path are active concurrently. The coordinator path acquires write locks to
121//! prevent concurrent writes between its read and write phases, but this OCC
122//! path does not use write locks, so concurrent operation of both paths could
123//! allow an OCC write to slip between a coordinator-path reader's read and
124//! write.
125
126use std::collections::BTreeMap;
127use std::collections::BTreeSet;
128use std::num::{NonZeroI64, NonZeroUsize};
129use std::sync::atomic::{AtomicBool, Ordering};
130use std::sync::{Arc, Mutex};
131use std::time::Duration;
132
133use bytesize::ByteSize;
134use differential_dataflow::consolidation;
135use mz_catalog::memory::error::ErrorKind;
136use mz_catalog::memory::objects::CatalogItem;
137use mz_cluster_client::ReplicaId;
138use mz_compute_types::ComputeInstanceId;
139use mz_expr::Eval;
140use mz_expr::row::RowCollection;
141use mz_expr::{CollectionPlan, Id, LocalId, MirRelationExpr, MirScalarExpr, RowSetFinishing};
142use mz_ore::cast::CastFrom;
143use mz_ore::{soft_assert_or_log, soft_panic_or_log};
144use mz_repr::optimize::OverrideFrom;
145use mz_repr::{CatalogItemId, Diff, GlobalId, RelationDesc, Row, RowArena, Timestamp};
146use mz_sql::catalog::CatalogError;
147use mz_sql::plan::{self, MutationKind, QueryWhen};
148use mz_sql::session::metadata::SessionMetadata;
149use mz_sql::session::vars::IsolationLevel;
150use mz_storage_client::client::TableData;
151use mz_storage_types::sources::Timeline;
152use mz_timestamp_oracle::TimestampOracle;
153use prometheus::Histogram;
154use timely::progress::Antichain;
155use tokio::sync::mpsc;
156use uuid::Uuid;
157
158use crate::active_compute_sink::ActiveSubscribeOwner;
159use crate::catalog::Catalog;
160use crate::command::{Command, ExecuteResponse, WriteAttemptKind};
161use crate::coord::appends::WriteResult;
162use crate::coord::read_then_write::{DependencyPolicy, validate_read_then_write_dependencies};
163use crate::coord::timestamp_selection::TimestampProvider;
164use crate::coord::{Coordinator, TargetCluster};
165use crate::error::AdapterError;
166use crate::metrics::{OCC_CALLER_BACKGROUND, OCC_CALLER_SESSION};
167use crate::optimize::Optimize;
168use crate::optimize::dataflows::{ComputeInstanceSnapshot, EvalTime, ExprPrep, ExprPrepOneShot};
169use crate::peek_client::CoordinatorClient;
170use crate::session::{Session, TransactionOps, WriteOp};
171use crate::statement_logging::{StatementLifecycleEvent, StatementLoggingId};
172use crate::{PeekClient, PeekResponseUnary, TimelineContext, optimize};
173
174/// Reason a frontend write attempt is being torn down early.
175#[derive(Clone, Copy)]
176pub(crate) enum FrontendWriteCancellation {
177 Canceled,
178 StatementTimeout,
179}
180
181impl From<FrontendWriteCancellation> for AdapterError {
182 fn from(cancellation: FrontendWriteCancellation) -> Self {
183 match cancellation {
184 FrontendWriteCancellation::Canceled => AdapterError::Canceled,
185 FrontendWriteCancellation::StatementTimeout => AdapterError::StatementTimeout,
186 }
187 }
188}
189
190/// State shared between an in-flight frontend write attempt and its
191/// cancellation wrapper,
192/// `SessionClient::try_frontend_read_then_write_with_cancel`.
193///
194/// The contract: `write_submitted` is true from just before the
195/// `AttemptWrite` command is sent until the attempt resolves as definitively
196/// not committed. While it is true, cancellation and statement timeout must
197/// not synthesize an error but await the definitive write result instead,
198/// because the write may already be durable.
199///
200/// The wrapper and the attempt it wraps are polled by the same task, so the
201/// mutex and the atomic are here to satisfy `Send`, not to arbitrate between
202/// concurrent writers. There is one writer for each field.
203pub(crate) struct FrontendWriteAttemptState {
204 write_submitted: AtomicBool,
205 /// Set at most once, by the cancellation wrapper.
206 cancellation: Mutex<Option<FrontendWriteCancellation>>,
207}
208
209impl FrontendWriteAttemptState {
210 pub(crate) fn new() -> Self {
211 Self {
212 write_submitted: AtomicBool::new(false),
213 cancellation: Mutex::new(None),
214 }
215 }
216
217 pub(crate) fn mark_write_submitted(&self) {
218 self.write_submitted.store(true, Ordering::Release);
219 }
220
221 /// Marks the submitted write as definitively not committed.
222 ///
223 /// NOTE: This must only be called for outcomes where the write is known
224 /// to not have landed (`TimestampPassed`). Terminal outcomes leave
225 /// `write_submitted` set so a concurrent cancellation path can never
226 /// fabricate an error for a write that may have committed.
227 fn mark_write_resolved(&self) {
228 self.write_submitted.store(false, Ordering::Release);
229 }
230
231 pub(crate) fn write_submitted(&self) -> bool {
232 self.write_submitted.load(Ordering::Acquire)
233 }
234
235 /// Records why the attempt is being torn down. The first reason recorded
236 /// is the one the attempt reports.
237 pub(crate) fn request(&self, cancellation: FrontendWriteCancellation) {
238 self.cancellation
239 .lock()
240 .expect("cancellation lock poisoned")
241 .get_or_insert(cancellation);
242 }
243
244 fn requested_error(&self) -> Option<AdapterError> {
245 self.cancellation
246 .lock()
247 .expect("cancellation lock poisoned")
248 .map(AdapterError::from)
249 }
250}
251
252/// Which kind of caller is driving a read-then-write.
253///
254/// Dependency rules, replica selection, and write cancellation follow from the
255/// caller kind.
256#[derive(Clone, Copy)]
257enum RtwCaller {
258 /// A user statement. Cancelled with its connection, and restricted to
259 /// reading user tables.
260 Session,
261 /// Coordinator-owned maintenance, pinned to one replica.
262 ///
263 /// The caller must build the statement itself rather than accept one from a
264 /// user, and must tolerate reading a log relation that is sealed empty,
265 /// which is how a replica with introspection disabled presents one.
266 Background { replica_id: ReplicaId },
267}
268
269impl RtwCaller {
270 fn is_background(&self) -> bool {
271 matches!(self, RtwCaller::Background { .. })
272 }
273
274 /// Which relations the selection may read.
275 fn dependency_policy(&self) -> DependencyPolicy {
276 match self {
277 RtwCaller::Session => DependencyPolicy::UserDml,
278 RtwCaller::Background { .. } => DependencyPolicy::SystemReads,
279 }
280 }
281
282 /// The replica a background caller pins its subscribe to, overriding the
283 /// session's replica selection.
284 fn replica_override(&self) -> Option<ReplicaId> {
285 match self {
286 RtwCaller::Background { replica_id } => Some(*replica_id),
287 RtwCaller::Session => None,
288 }
289 }
290
291 /// Who owns the subscribe, which decides whether it is cancelled with a
292 /// connection and whether it counts against one.
293 fn subscribe_owner(
294 &self,
295 conn_id: &mz_adapter_types::connection::ConnectionId,
296 session_uuid: Uuid,
297 ) -> ActiveSubscribeOwner {
298 match self {
299 RtwCaller::Session => ActiveSubscribeOwner::Session {
300 conn_id: conn_id.clone(),
301 session_uuid,
302 },
303 RtwCaller::Background { .. } => ActiveSubscribeOwner::Background,
304 }
305 }
306
307 /// The connection a pending write is cancelled with, if any.
308 fn write_conn_id(
309 &self,
310 conn_id: &mz_adapter_types::connection::ConnectionId,
311 ) -> Option<mz_adapter_types::connection::ConnectionId> {
312 match self {
313 RtwCaller::Session => Some(conn_id.clone()),
314 RtwCaller::Background { .. } => None,
315 }
316 }
317}
318
319/// What the OCC loop produced.
320enum OccOutcome {
321 /// The write is durable at `write_ts`.
322 Committed {
323 response: ExecuteResponse,
324 write_ts: Timestamp,
325 },
326 /// The selection was empty, so there was nothing to write.
327 ///
328 /// `empty_as_of` is the timestamp the emptiness holds at, and the caller must
329 /// bring the oracle's read timestamp up to it before responding. `None` when
330 /// the subscribe ran to completion, where the emptiness holds at every
331 /// timestamp. See the module docs for why the choice of timestamp matters.
332 NoRowsMatched {
333 response: ExecuteResponse,
334 empty_as_of: Option<Timestamp>,
335 },
336 /// Diffs no frontier can change, from a subscribe that ran to completion.
337 /// The close says the selection can never change again, not that it reads
338 /// nothing persisted, and either way the caller chooses whether to submit
339 /// them now or buffer them into the transaction.
340 Blind {
341 response: ExecuteResponse,
342 diffs: Vec<(Row, Diff)>,
343 },
344}
345
346/// What the coordinator's answer to a submitted write means for the statement.
347enum WriteOutcome {
348 /// The write is durable at this timestamp.
349 Committed(Timestamp),
350 /// The write did not land, and resubmitting these diffs cannot change
351 /// that. This is the error to report.
352 Failed(AdapterError),
353 /// Another writer advanced the target's upper past the timestamp we asked
354 /// for. The diffs still describe the mutation, so the OCC loop can
355 /// resubmit them once the subscribe has caught up.
356 Conflict { next_eligible_timestamp: Timestamp },
357}
358
359/// Maps a [`WriteResult`] to the outcome the statement reports, or to the one
360/// conflict the OCC loop can retry.
361fn classify_write_result(
362 result: WriteResult,
363 target_id: CatalogItemId,
364 attempt_state: &FrontendWriteAttemptState,
365) -> WriteOutcome {
366 match result {
367 WriteResult::Success { timestamp } => WriteOutcome::Committed(timestamp),
368 WriteResult::TimestampPassed {
369 next_eligible_timestamp,
370 ..
371 } => WriteOutcome::Conflict {
372 next_eligible_timestamp,
373 },
374 WriteResult::Canceled => WriteOutcome::Failed(
375 attempt_state
376 .requested_error()
377 .unwrap_or(AdapterError::Canceled),
378 ),
379 WriteResult::TimestampTooFarAhead {
380 target_timestamp,
381 limit,
382 } => WriteOutcome::Failed(AdapterError::ReadThenWriteTimestampTooFarAhead {
383 target_timestamp,
384 limit,
385 }),
386 WriteResult::ReadOnly => WriteOutcome::Failed(AdapterError::ReadOnly),
387 WriteResult::TargetChanged => {
388 // A concurrent DDL gave the table a new generation after we
389 // computed these diffs against the old one. The same error the
390 // coordinator raises when a dependency changes underneath a
391 // statement, so clients see one retryable outcome for both.
392 WriteOutcome::Failed(AdapterError::ConcurrentDependencyMutation {
393 dependency_id: target_id.to_string(),
394 })
395 }
396 WriteResult::Indeterminate => WriteOutcome::Failed(AdapterError::Internal(
397 "write outcome is indeterminate because the group committer shut down".into(),
398 )),
399 }
400}
401
402/// Ends the implicit transaction that a statement which cannot run in a
403/// transaction block opened for itself.
404///
405/// A read-then-write that reads persisted state is refused inside a transaction,
406/// so it must not quietly become part of one. Clearing the ops it staged leaves
407/// [`crate::session::TransactionStatus::may_span_pipeline`] false, so pgwire
408/// commits the implicit transaction rather than letting the rest of an
409/// extended-protocol pipeline join it. The write is already durable at this
410/// point, and the caller has established there is nothing else staged to lose.
411fn end_own_transaction(session: &mut Session, stages_rows: bool) {
412 if !stages_rows {
413 session.clear_transaction_ops();
414 }
415}
416
417/// Checks that a read-then-write may read what its selection depends on.
418///
419/// An invalid selection is invalid wherever the statement runs, so a caller with
420/// something contextual to report, such as the transaction state, must ask this
421/// first. Reporting the transaction for a statement that can never work suggests
422/// it would work outside one.
423///
424/// `catalog` must be the snapshot the plan was built against. A missing item
425/// means the caller mixed snapshots, which is reported as a catalog error rather
426/// than treated as a dropped dependency.
427pub(crate) fn validate_selection_dependencies(
428 catalog: &Catalog,
429 depends_on: &BTreeSet<GlobalId>,
430 policy: DependencyPolicy,
431) -> Result<(), AdapterError> {
432 let dependency_ids = depends_on
433 .iter()
434 .copied()
435 .map(|gid| {
436 catalog.try_resolve_item_id(&gid).ok_or_else(|| {
437 AdapterError::Catalog(mz_catalog::memory::error::Error {
438 kind: ErrorKind::Sql(CatalogError::UnknownItem(gid.to_string())),
439 })
440 })
441 })
442 .collect::<Result<Vec<_>, _>>()?;
443 let max_rw_dependencies = mz_adapter_types::dyncfgs::READ_THEN_WRITE_MAX_DEPENDENCIES
444 .get(catalog.system_config().dyncfgs());
445 validate_read_then_write_dependencies(catalog, dependency_ids, max_rw_dependencies, policy)
446}
447
448/// Validates a read-then-write and resolves the context the rest of the
449/// pipeline runs against.
450///
451/// Rejects `mz_now()` in the selection, the assignments or the returning
452/// clause. `optimize_mir_read_then_write` relies on that rejection by name
453/// when it prepares unmaterializable functions one-shot. Also enforces the
454/// dependency cap, resolves the target cluster and requires it to have a
455/// live replica, honors the session's replica pin, computes the read side's
456/// `TimelineContext`, and fetches the target table's descriptor.
457///
458/// `catalog` must be the snapshot the plan was built against. One snapshot
459/// serves planning, validation and optimization, so items the plan names
460/// cannot disappear from it, and the missing-entry branches below are
461/// failsafes rather than a live concurrent-DDL path.
462///
463/// `dependency_policy` decides which relations the selection may read. Both
464/// policies reject `mz_now()` anywhere in the transitive dependencies.
465fn validate_read_then_write(
466 catalog: &Arc<Catalog>,
467 session: &Session,
468 plan: &plan::ReadThenWritePlan,
469 target_cluster: TargetCluster,
470 dependency_policy: DependencyPolicy,
471) -> Result<ValidationResult, AdapterError> {
472 if contains_mz_now(plan) {
473 return Err(AdapterError::Unsupported(
474 "calls to mz_now in write statements",
475 ));
476 }
477
478 // One walk of the selection serves both the dependency check and the
479 // timeline validation below.
480 let depends_on = plan.selection.depends_on();
481
482 validate_selection_dependencies(catalog, &depends_on, dependency_policy)?;
483
484 let cluster = catalog.resolve_target_cluster(target_cluster, session)?;
485 let cluster_id = cluster.id;
486
487 if cluster.replicas().next().is_none() {
488 return Err(AdapterError::NoClusterReplicasAvailable {
489 name: cluster.name.clone(),
490 is_managed: cluster.is_managed(),
491 });
492 }
493
494 let replica_id = session
495 .vars()
496 .cluster_replica()
497 .map(|name| {
498 cluster
499 .replica_id(name)
500 .ok_or(AdapterError::UnknownClusterReplica {
501 cluster_name: cluster.name.clone(),
502 replica_name: name.to_string(),
503 })
504 })
505 .transpose()?;
506
507 let timeline = catalog.validate_timeline_context(depends_on.iter().copied())?;
508
509 // The loop waits for the subscribe's frontier to reach a target timestamp
510 // taken from the `EpochMilliseconds` oracle, and it is the target table's
511 // upper the write then competes for. A selection in another timeline
512 // counts something else, transactions rather than milliseconds for a CDCv2
513 // source, so its frontier is not comparable with that target and the
514 // statement would burn until `statement_timeout`. Refuse it up front
515 // instead.
516 //
517 // Only `INSERT ... SELECT` reaches this. A DELETE or UPDATE selection
518 // includes the target table, so a foreign timeline already fails above
519 // as a mixed-timeline query.
520 if let TimelineContext::TimelineDependent(t) = &timeline {
521 if t != &Timeline::EpochMilliseconds {
522 return Err(AdapterError::Unsupported(
523 "read-then-write on a selection outside the EpochMilliseconds timeline",
524 ));
525 }
526 }
527
528 // Get the table descriptor for constraint validation. As above, a
529 // missing entry would mean the snapshot contract was broken.
530 let table_desc = match catalog.try_get_entry(&plan.id) {
531 Some(entry) => entry
532 .relation_desc_latest()
533 .expect("table has desc")
534 .into_owned(),
535 None => {
536 return Err(AdapterError::Catalog(mz_catalog::memory::error::Error {
537 kind: ErrorKind::Sql(CatalogError::UnknownItem(plan.id.to_string())),
538 }));
539 }
540 };
541
542 Ok(ValidationResult {
543 cluster_id,
544 replica_id,
545 timeline,
546 depends_on,
547 table_desc,
548 })
549}
550
551/// Builds the response for a write that is about to be submitted.
552///
553/// This runs before the write, so the result-size checks in here reject the
554/// statement without having written anything.
555fn build_success_response(
556 kind: &MutationKind,
557 returning: &[MirScalarExpr],
558 diffs: &[(Row, Diff)],
559 max_result_size: u64,
560 max_query_result_size: u64,
561 row_set_finishing_seconds: &Histogram,
562) -> Result<ExecuteResponse, AdapterError> {
563 if returning.is_empty() {
564 // For UPDATE each changed row produces a retraction (-1) and an
565 // insertion (+1), so we divide by 2 below.
566 let row_count = diffs
567 .iter()
568 .map(|(_, diff)| diff.into_inner().unsigned_abs())
569 .sum::<u64>();
570 let row_count = usize::try_from(row_count).expect("positive row count must fit in usize");
571
572 return Ok(match kind {
573 MutationKind::Delete => ExecuteResponse::Deleted(row_count),
574 MutationKind::Update => ExecuteResponse::Updated(row_count / 2),
575 MutationKind::Insert => ExecuteResponse::Inserted(row_count),
576 });
577 }
578
579 let mut returning_rows = Vec::new();
580 let arena = RowArena::new();
581 // RETURNING expressions are evaluated row-by-row in this loop, so an
582 // expression like `RETURNING repeat('x', 10_000_000)` will allocate
583 // unbounded data unless we bail mid-loop. The post-loop
584 // `RowSetFinishing::finish` below would also reject this, but only
585 // after we've materialized everything. The early-bail caps the
586 // temporary allocation. We pick the lower of the two configured caps,
587 // whichever fires first wins.
588 let mut projected_byte_size: u64 = 0;
589 let early_cap = std::cmp::min(max_result_size, max_query_result_size);
590
591 for (row, diff) in diffs {
592 let include = match kind {
593 MutationKind::Delete => diff.is_negative(),
594 MutationKind::Update | MutationKind::Insert => diff.is_positive(),
595 };
596
597 if !include {
598 continue;
599 }
600
601 let mut returning_row = Row::with_capacity(returning.len());
602 let mut packer = returning_row.packer();
603 let datums: Vec<_> = row.iter().collect();
604
605 for expr in returning {
606 match expr.eval(&datums, &arena) {
607 Ok(datum) => packer.push(datum),
608 Err(err) => return Err(err.into()),
609 }
610 }
611
612 let multiplicity = NonZeroUsize::try_from(
613 NonZeroI64::try_from(diff.into_inner().abs()).expect("diff is non-zero"),
614 )
615 .map_err(AdapterError::from)?;
616
617 let row_bytes = u64::cast_from(returning_row.byte_len())
618 .saturating_mul(u64::cast_from(multiplicity.get()));
619 projected_byte_size = projected_byte_size.saturating_add(row_bytes);
620 if projected_byte_size > early_cap {
621 return Err(AdapterError::ResultSize(format!(
622 "result exceeds max size of {}",
623 ByteSize::b(early_cap)
624 )));
625 }
626
627 returning_rows.push((returning_row, multiplicity));
628 }
629
630 // Run the canonical finish to enforce both caps with full precision
631 // (including the sorted-view memory overhead) and to register the
632 // row-set-finishing duration histogram, mirroring the legacy
633 // `send_diffs` path.
634 let finishing = RowSetFinishing {
635 order_by: Vec::new(),
636 limit: None,
637 offset: 0,
638 project: (0..returning.len()).collect(),
639 };
640 match finishing.finish(
641 RowCollection::new(returning_rows, &finishing.order_by),
642 max_result_size,
643 Some(max_query_result_size),
644 row_set_finishing_seconds,
645 ) {
646 Ok((rows, _size_bytes)) => Ok(ExecuteResponse::SendingRowsImmediate {
647 rows: Box::new(rows),
648 }),
649 Err(e) => Err(AdapterError::ResultSize(e)),
650 }
651}
652
653/// Whether a read-then-write mentions `mz_now()` anywhere.
654///
655/// Read time and write time differ on this path, so `mz_now()` has no single
656/// answer and is refused in every position.
657pub(crate) fn contains_mz_now(plan: &plan::ReadThenWritePlan) -> bool {
658 plan.selection.contains_temporal()
659 || plan.assignments.values().any(|e| e.contains_temporal())
660 || plan.returning.iter().any(|e| e.contains_temporal())
661}
662
663/// The timeline whose oracle governs a read-then-write with the given read-side
664/// [`TimelineContext`].
665///
666/// The write target is always a table on `EpochMilliseconds`, so a read side
667/// that pins no timeline (`TimestampDependent`) still maps to
668/// `EpochMilliseconds`. `None` only for timestamp-independent selections, which
669/// need no oracle at all.
670fn governing_timeline(timeline: &TimelineContext) -> Option<Timeline> {
671 <Coordinator as TimestampProvider>::get_timeline(timeline)
672}
673
674/// A handle to an internal subscribe, meaning one that writes no
675/// `mz_subscriptions` row. A `Drop` impl ensures the subscribe's dataflow is
676/// cleaned up when dropped.
677struct SubscribeHandle {
678 rx: mpsc::UnboundedReceiver<PeekResponseUnary>,
679 sink_id: GlobalId,
680 /// Wrapped in `Option` so we can move it out in `Drop`.
681 client: Option<CoordinatorClient>,
682}
683
684impl SubscribeHandle {
685 /// Receive the next message from the subscribe, waiting if necessary.
686 pub async fn recv(&mut self) -> Option<PeekResponseUnary> {
687 self.rx.recv().await
688 }
689
690 /// Try to receive a message without waiting.
691 pub fn try_recv(&mut self) -> Result<PeekResponseUnary, mpsc::error::TryRecvError> {
692 self.rx.try_recv()
693 }
694}
695
696impl Drop for SubscribeHandle {
697 fn drop(&mut self) {
698 if let Some(client) = self.client.take() {
699 // Fire-and-forget: if the coordinator is gone, the subscribe will
700 // be cleaned up when the process exits anyway.
701 client.try_send(Command::DropInternalSubscribe {
702 sink_id: self.sink_id,
703 });
704 }
705 }
706}
707
708impl PeekClient {
709 /// Execute a read-then-write operation using frontend sequencing.
710 ///
711 /// Called by session code when the frontend_read_then_write dyncfg is
712 /// enabled. The caller owns the end-of-execution logging for
713 /// `statement_logging_id` and verified and planned the portal against
714 /// `catalog`, which stays in force through optimization and write-target
715 /// generation capture.
716 pub(crate) async fn frontend_read_then_write(
717 &mut self,
718 session: &mut Session,
719 plan: plan::ReadThenWritePlan,
720 target_cluster: TargetCluster,
721 catalog: &Arc<Catalog>,
722 statement_logging_id: Option<StatementLoggingId>,
723 attempt_state: Arc<FrontendWriteAttemptState>,
724 ) -> Result<ExecuteResponse, AdapterError> {
725 self.read_then_write(
726 session,
727 plan,
728 target_cluster,
729 catalog,
730 statement_logging_id,
731 attempt_state,
732 RtwCaller::Session,
733 )
734 .await
735 }
736
737 /// Executes a coordinator-owned read-then-write against system relations,
738 /// pinned to `replica_id`.
739 ///
740 /// See `RtwCaller::Background` for what the caller takes on by using this.
741 pub(crate) async fn background_read_then_write(
742 &mut self,
743 session: &mut Session,
744 plan: plan::ReadThenWritePlan,
745 cluster_id: ComputeInstanceId,
746 replica_id: ReplicaId,
747 catalog: &Arc<Catalog>,
748 ) -> Result<ExecuteResponse, AdapterError> {
749 let is_system_table = matches!(plan.id, CatalogItemId::System(_))
750 && catalog
751 .try_get_entry(&plan.id)
752 .is_some_and(|entry| matches!(entry.item(), CatalogItem::Table(_)));
753 if !is_system_table {
754 soft_panic_or_log!(
755 "background read-then-write target {} is not a system table",
756 plan.id
757 );
758 return Err(AdapterError::Internal(
759 "background read-then-write target is not a system table".into(),
760 ));
761 }
762
763 self.read_then_write(
764 session,
765 plan,
766 TargetCluster::Transaction(cluster_id),
767 catalog,
768 None,
769 // Nothing cancels a background write, so this state only ever
770 // records that a write was submitted.
771 Arc::new(FrontendWriteAttemptState::new()),
772 RtwCaller::Background { replica_id },
773 )
774 .await
775 }
776
777 #[allow(clippy::too_many_arguments)]
778 async fn read_then_write(
779 &mut self,
780 session: &mut Session,
781 mut plan: plan::ReadThenWritePlan,
782 target_cluster: TargetCluster,
783 catalog: &Arc<Catalog>,
784 statement_logging_id: Option<StatementLoggingId>,
785 attempt_state: Arc<FrontendWriteAttemptState>,
786 caller: RtwCaller,
787 ) -> Result<ExecuteResponse, AdapterError> {
788 // The OCC dataflow emits raw diffs and does not apply top-level
789 // finishing. Silently dropping a LIMIT, OFFSET, projection or ordering
790 // can change the rows written, so this stage requires trivial finishing.
791 if !plan.finishing.is_trivial(plan.selection.arity()) {
792 soft_panic_or_log!("frontend read-then-write received nontrivial row-set finishing");
793 return Err(AdapterError::Internal(
794 "frontend read-then-write requires trivial row-set finishing".into(),
795 ));
796 }
797
798 // A transaction that has taken a timestamped read, was opened READ
799 // ONLY, or is committed to some other kind of operation cannot take a
800 // write. Check up front, mirroring `sequence_insert`: the marker op
801 // below rejects only some of those states, and only with its own
802 // errors, so without this check the reported error and SQLSTATE would
803 // depend on which path sequenced the statement.
804 //
805 // Both this and the marker op require an open transaction. The
806 // frontends start one before they execute anything, and a `Failed`
807 // transaction only ever admits COMMIT/ROLLBACK, so DML never arrives
808 // in a state where these panic.
809 if !session.transaction().allows_writes() {
810 return Err(AdapterError::ReadOnlyTransaction);
811 }
812
813 let validation_result = validate_read_then_write(
814 catalog,
815 session,
816 &plan,
817 target_cluster,
818 caller.dependency_policy(),
819 )?;
820
821 let ValidationResult {
822 cluster_id,
823 mut replica_id,
824 timeline,
825 depends_on,
826 table_desc,
827 } = validation_result;
828 if let Some(pinned) = caller.replica_override() {
829 replica_id = Some(pinned);
830 }
831
832 // A write that reads no persisted state may join a surrounding
833 // transaction. Its rows do not come from a snapshot, so staging them
834 // and letting the transaction flush them keeps them atomic with
835 // everything else the transaction does.
836 //
837 // A write that does read persisted state may not. We refuse it in an
838 // explicit transaction, and an extended-protocol pipeline is an
839 // implicit transaction, so it must not silently span one either. It
840 // runs as its own transaction instead, which is how PostgreSQL treats
841 // statements that cannot run in a transaction block.
842 let stages_rows = depends_on.is_empty();
843
844 // Snapshot this before the marker op below, which makes the predicate
845 // true unconditionally.
846 let in_transaction = session
847 .transaction()
848 .may_share_transaction_with_other_statements();
849
850 if !stages_rows && in_transaction {
851 // Defense in depth for the gate in
852 // `SessionClient::try_frontend_read_then_write`. Rejecting here,
853 // before we run a dataflow, is the last point where refusing is
854 // still possible: past the OCC loop the write may already be
855 // durable.
856 soft_panic_or_log!(
857 "read-dependent read-then-write reached the OCC path inside a transaction"
858 );
859 return Err(AdapterError::Internal(
860 "read-then-write cannot be run inside a transaction block".into(),
861 ));
862 }
863
864 // Mark this as a write transaction in the session state machine, so
865 // auto-commit treats the statement as a write. The rows follow once we
866 // know them.
867 session.add_transaction_ops(TransactionOps::Writes(vec![]))?;
868
869 // Prepare expressions (resolve unmaterializable functions like
870 // current_user())
871 let style = ExprPrepOneShot {
872 logical_time: EvalTime::NotAvailable, // We already errored out on mz_now above.
873 session,
874 catalog_state: catalog.state(),
875 };
876 for expr in plan
877 .assignments
878 .values_mut()
879 .chain(plan.returning.iter_mut())
880 {
881 style.prep_scalar_expr(expr)?;
882 }
883
884 let (mut optimizer, global_mir_plan) =
885 self.optimize_mir_read_then_write(catalog, session, &plan, cluster_id)?;
886
887 // Acquire the OCC semaphore permit *before* acquiring read holds in
888 // `frontend_determine_timestamp`. Under contention, waiters will
889 // otherwise sit on read holds on the RTW's read dependencies for the
890 // entire time they are queued, pinning compaction on those
891 // collections. Waiting on the permit first keeps queued operations
892 // hold-free. Once we have a permit we proceed to acquire the read holds
893 // needed for the rest of the operation.
894 //
895 // The cost of this ordering is that a permit held by a long-running
896 // operation stalls every read-then-write in the process, including ones
897 // on unrelated tables, where the coordinator's write lock would only
898 // stall writes to the target table. We accept that because the
899 // statement timeout in
900 // `SessionClient::try_frontend_read_then_write_with_cancel` covers the
901 // permit wait, so the stall is bounded for everyone but a session that
902 // disabled its own timeout.
903 //
904 // The semaphore is owned by the coordinator and outlives every
905 // session task, so `acquire_owned` cannot return `Err` in practice.
906 //
907 // Background maintenance skips the queue entirely. It is single-flight
908 // by construction, one sweep at a time and one mutation at a time, so it
909 // adds at most one concurrent read-then-write. Taking a permit instead
910 // would let it hold one for as long as a subscribe on a loaded user
911 // replica takes to hydrate, and with `max_concurrent_occ_writes` set low
912 // that stalls user DML behind a background sampler. The bound above does
913 // not apply to it either: it has no statement timeout, only its own much
914 // longer one.
915 let permit = if caller.is_background() {
916 None
917 } else {
918 Some(
919 Arc::clone(&self.occ_write_semaphore)
920 .acquire_owned()
921 .await
922 .expect("semaphore is never closed during coordinator lifetime"),
923 )
924 };
925
926 // Determine timestamp and acquire read holds.
927 let oracle_read_ts = self.oracle_read_ts(&timeline).await?;
928
929 // Real-time recency, on the same terms as the frontend peek path. The
930 // coordinator round trip polls the selection's upstream sources, so we
931 // only pay it when the session actually asked for recency.
932 let vars = session.vars();
933 let real_time_recency_ts: Option<Timestamp> = if vars.real_time_recency()
934 && vars.transaction_isolation() == &IsolationLevel::StrictSerializable
935 && !session.contains_read_timestamp()
936 {
937 let real_time_recency_timeout = *vars.real_time_recency_timeout();
938 self.call_coordinator(|tx| Command::DetermineRealTimeRecentTimestamp {
939 source_ids: depends_on.iter().copied().collect(),
940 real_time_recency_timeout,
941 tx,
942 })
943 .await??
944 } else {
945 None
946 };
947
948 let bundle = global_mir_plan.id_bundle(cluster_id);
949 let (determination, read_holds) = self
950 .frontend_determine_timestamp(
951 session,
952 &bundle,
953 &QueryWhen::FreshestTableWrite,
954 cluster_id,
955 &timeline,
956 oracle_read_ts,
957 real_time_recency_ts,
958 )
959 .await?;
960
961 let as_of = determination.timestamp_context.timestamp_or_default();
962
963 let global_mir_plan = global_mir_plan.resolve(Antichain::from_elem(as_of));
964 let global_lir_plan = optimizer.catch_unwind_optimize(global_mir_plan)?;
965
966 // Log optimization finished
967 if let Some(logging_id) = statement_logging_id {
968 self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished);
969 }
970
971 let sink_id = global_lir_plan.sink_id();
972 let target_id = plan.id;
973 let target_global_id = catalog.get_entry(&target_id).latest_global_id();
974 let kind = plan.kind.clone();
975 let returning = plan.returning.clone();
976
977 let (df_desc, df_meta) = global_lir_plan.unapply();
978
979 // The coordinator sequences this statement's read as a real peek, so the
980 // optimizer's notices and the timestamp notice reach the session there.
981 // Emit both here for the same statement to look the same on either path.
982 crate::coord::sequencer::emit_optimizer_notices(
983 &**catalog,
984 session,
985 &df_meta.optimizer_notices,
986 );
987 if session.vars().emit_timestamp_notice() {
988 let conn_id = session.conn_id().clone();
989 let session_wall_time = session.pcx().wall_time;
990 let explanation = self
991 .call_coordinator(|tx| Command::ExplainTimestamp {
992 conn_id,
993 session_wall_time,
994 cluster_id,
995 id_bundle: bundle,
996 determination,
997 tx,
998 })
999 .await?;
1000 session.add_notice(crate::AdapterNotice::QueryTimestamp { explanation });
1001 }
1002
1003 let arity = df_desc
1004 .sink_exports
1005 .values()
1006 .next()
1007 .expect("has sink")
1008 .from_desc
1009 .arity();
1010
1011 let conn_id = session.conn_id().clone();
1012 let session_uuid = session.uuid();
1013 let start_time = (self.statement_logging_frontend.now)();
1014 let max_result_size = catalog.system_config().max_result_size();
1015 let max_query_result_size = session.vars().max_query_result_size();
1016 let row_set_finishing_seconds = session.metrics().row_set_finishing_seconds().clone();
1017 let max_occ_retries = usize::cast_from(catalog.system_config().max_occ_retries());
1018
1019 // Linearize the read BEFORE subscribing or writing: block until
1020 // the oracle for this query's timeline has advanced to `as_of`.
1021 //
1022 // Ordering is load-bearing: this leaves the oracle at or above `as_of`,
1023 // which is what makes the loop's target clear `as_of` and so include the
1024 // snapshot. A far-future `as_of` parks here until the clock arrives,
1025 // bounded by `statement_timeout`.
1026 self.ensure_read_linearized(&timeline, as_of).await?;
1027
1028 // The loop takes its write target from this oracle, and reaching one takes
1029 // `&mut self`, which the loop does not have. `None` for a
1030 // timestamp-independent selection, which reads at `Timestamp::maximum()`
1031 // and so always leaves through the blind path rather than reaching a write.
1032 let write_oracle = match governing_timeline(&timeline) {
1033 Some(tl) => Some(Arc::clone(self.ensure_oracle(tl).await?)),
1034 None => None,
1035 };
1036
1037 let subscribe_handle = self
1038 .create_internal_subscribe(
1039 Box::new(df_desc),
1040 cluster_id,
1041 replica_id,
1042 depends_on.clone(),
1043 as_of,
1044 arity,
1045 sink_id,
1046 caller.subscribe_owner(&conn_id, session_uuid),
1047 start_time,
1048 read_holds,
1049 )
1050 .await?;
1051
1052 let (retry_count, result) = self
1053 .run_occ_loop(
1054 subscribe_handle,
1055 target_id,
1056 target_global_id,
1057 kind,
1058 returning,
1059 max_result_size,
1060 max_query_result_size,
1061 row_set_finishing_seconds,
1062 max_occ_retries,
1063 table_desc,
1064 caller.write_conn_id(&conn_id),
1065 statement_logging_id,
1066 as_of,
1067 write_oracle,
1068 &attempt_state,
1069 )
1070 .await;
1071
1072 let caller_label = match caller {
1073 RtwCaller::Session => OCC_CALLER_SESSION,
1074 RtwCaller::Background { .. } => OCC_CALLER_BACKGROUND,
1075 };
1076 self.coordinator_client()
1077 .metrics()
1078 .occ_retry_count
1079 .with_label_values(&[caller_label])
1080 .observe(f64::from(u32::try_from(retry_count).unwrap_or(u32::MAX)));
1081
1082 // Finish the operation, including a blind write's submission, before
1083 // releasing the OCC permit. Holding it for the entire operation is what
1084 // bounds concurrency. An early drop would let a waiter start its
1085 // subscribe while we are still consolidating diffs, retrying, or
1086 // waiting for our write to commit.
1087 //
1088 // The zero-row linearization wait below is the one exception, and hands
1089 // the permit back before it parks.
1090 let mut permit = permit;
1091 let response = match result {
1092 Ok(OccOutcome::Committed { response, write_ts }) => {
1093 // A committed write timestamp for a statement we meant to
1094 // stage means the two predicates disagreed: the syntactic one
1095 // said it reads nothing, the subscribe then read persisted
1096 // state. The write is already durable, so there is nothing to
1097 // refuse, and `apply_write` still has to run to keep the
1098 // session's read timestamps ahead of it.
1099 soft_assert_or_log!(
1100 !stages_rows,
1101 "read-then-write committed a write it meant to stage"
1102 );
1103 session.apply_write(write_ts);
1104 end_own_transaction(session, stages_rows);
1105 Ok(response)
1106 }
1107 Ok(OccOutcome::NoRowsMatched {
1108 response,
1109 empty_as_of,
1110 }) => {
1111 // An `empty_as_of` for a statement we meant to stage is the
1112 // same predicate disagreement the `Committed` arm guards
1113 // against: the syntactic answer said it reads nothing, the
1114 // subscribe then read persisted state. Nothing is durable here,
1115 // so there is nothing to undo, but the disagreement itself is
1116 // the bug and it would otherwise park silently.
1117 soft_assert_or_log!(
1118 !(stages_rows && empty_as_of.is_some()),
1119 "read-then-write observed a read timestamp for a statement \
1120 it meant to stage"
1121 );
1122 end_own_transaction(session, stages_rows);
1123 match empty_as_of {
1124 Some(empty_as_of) => {
1125 // The wait is a no-op where the oracle is already past
1126 // `empty_as_of`, which is the common `WHERE <no match>`,
1127 // and otherwise costs the group commit that
1128 // `ensure_read_linearized` asks for. Either way the
1129 // subscribe handle is gone, so the permit guards nothing
1130 // and holding it would throttle unrelated writes.
1131 drop(permit.take());
1132 self.ensure_read_linearized(&timeline, empty_as_of)
1133 .await
1134 .map(|()| response)
1135 }
1136 None => Ok(response),
1137 }
1138 }
1139 Ok(OccOutcome::Blind { response, diffs }) if stages_rows => {
1140 // Staging rather than writing here is what makes the statement
1141 // atomic with its transaction. An extended-protocol pipeline is
1142 // an implicit transaction, and it may still fail after us, so a
1143 // write of our own would survive a rollback that discards
1144 // everything around it.
1145 //
1146 // NOTE: A staged session write carries no target-generation
1147 // guard the way the immediate path's `target_global_id` does. A
1148 // `WriteOp` only names the `CatalogItemId`, and commit staging
1149 // resolves whatever global id is current then. What keeps it
1150 // safe is a check at the far end: group commit compares the
1151 // arity of the rows each staged write carries against the
1152 // target's latest `RelationDesc` and rolls the transaction back
1153 // with `ConcurrentDependencyMutation` instead of encoding old
1154 // rows against a new schema. So an `ALTER TABLE ... ADD COLUMN`
1155 // landing between here and the commit becomes the same
1156 // retryable failure the immediate path reports as
1157 // `TargetChanged`. The comparison reads one row per staged
1158 // write and looks only at arity, so it stands in for the
1159 // descriptor rather than pinning it.
1160 //
1161 // Missing the pin is true of every staged write. The arity
1162 // check is not: rows staged as a batch, which is how `COPY
1163 // FROM` arrives, carry their schema into persist instead.
1164 session
1165 .add_transaction_ops(TransactionOps::Writes(vec![WriteOp {
1166 id: target_id,
1167 rows: TableData::Rows(diffs),
1168 }]))
1169 .map(|()| response)
1170 }
1171 Ok(OccOutcome::Blind { response, diffs }) => {
1172 if caller.is_background() {
1173 return Err(AdapterError::Internal(
1174 "background read-then-write unexpectedly had no persisted dependency"
1175 .into(),
1176 ));
1177 }
1178 // The subscribe closed on its own even though the selection
1179 // reads persisted state, so the input is sealed and the diffs
1180 // are frontier-independent after all. Staging them would be
1181 // safe for that reason, and we still do not, because it would
1182 // make a statement's transaction semantics depend on whether an
1183 // input happens to be past its last refresh. What the statement
1184 // reads decides, so this commits as its own transaction like
1185 // any other read-dependent write.
1186 match self
1187 .submit_blind_write(
1188 conn_id,
1189 target_id,
1190 target_global_id,
1191 diffs,
1192 statement_logging_id,
1193 &attempt_state,
1194 )
1195 .await
1196 {
1197 Ok(write_ts) => {
1198 session.apply_write(write_ts);
1199 end_own_transaction(session, stages_rows);
1200 Ok(response)
1201 }
1202 Err(err) => Err(err),
1203 }
1204 }
1205 Err(err) => Err(err),
1206 };
1207
1208 drop(permit);
1209
1210 response
1211 }
1212
1213 /// Builds the subscribe optimizer and the unresolved global MIR plan for a
1214 /// read-then-write.
1215 ///
1216 /// The optimized expression is the selection with the mutation already
1217 /// applied, so the subscribe's sink emits ready-to-write table diffs rather
1218 /// than query results. `finishing` and `returning` are deliberately not part
1219 /// of the dataflow, and unmaterializable functions are prepared one-shot.
1220 fn optimize_mir_read_then_write(
1221 &self,
1222 catalog: &Arc<Catalog>,
1223 session: &dyn SessionMetadata,
1224 plan: &plan::ReadThenWritePlan,
1225 cluster_id: ComputeInstanceId,
1226 ) -> Result<
1227 (
1228 optimize::subscribe::Optimizer,
1229 optimize::subscribe::GlobalMirPlan<optimize::subscribe::Unresolved>,
1230 ),
1231 AdapterError,
1232 > {
1233 // `finishing` is unused: the OCC path emits raw diffs and
1234 // `apply_mutation_to_mir` handles update projection.
1235 let plan::ReadThenWritePlan {
1236 id: _,
1237 selection,
1238 finishing: _,
1239 assignments,
1240 kind,
1241 returning: _,
1242 } = plan;
1243
1244 let expr = selection.clone().lower(catalog.system_config(), None)?;
1245 let mut expr = apply_mutation_to_mir(expr, kind, assignments);
1246
1247 // Resolve unmaterializable functions (now(), current_user, ...) before
1248 // the subscribe optimizer sees them: it uses `ExprPrepMaintained`,
1249 // which rejects them, but our subscribe is a one-shot read so we can
1250 // resolve them to constants. `mz_now()` is rejected upstream by
1251 // `validate_read_then_write`.
1252 let style = ExprPrepOneShot {
1253 logical_time: EvalTime::NotAvailable,
1254 session,
1255 catalog_state: catalog.state(),
1256 };
1257 expr.try_visit_scalars_mut(&mut |s| style.prep_scalar_expr(s))?;
1258
1259 let compute_instance = ComputeInstanceSnapshot::new_without_collections(cluster_id);
1260 let (_, view_id) = self.transient_id_gen.allocate_id();
1261 let (_, sink_id) = self.transient_id_gen.allocate_id();
1262 let debug_name = format!("frontend-read-then-write-subscribe-{}", sink_id);
1263 let optimizer_config = optimize::OptimizerConfig::from(catalog.system_config())
1264 .override_from(&catalog.get_cluster(cluster_id).config.features())
1265 .override_from(
1266 &catalog
1267 .state()
1268 .cluster_scoped_optimizer_overrides(cluster_id),
1269 );
1270
1271 let mut optimizer = optimize::subscribe::Optimizer::new(
1272 Arc::<Catalog>::clone(catalog),
1273 compute_instance,
1274 view_id,
1275 sink_id,
1276 true, // with_snapshot
1277 None, // up_to
1278 debug_name,
1279 optimizer_config,
1280 self.optimizer_metrics.clone(),
1281 );
1282
1283 let expr_typ = expr.typ();
1284 let sql_typ = mz_repr::SqlRelationType::from_repr(&expr_typ);
1285 let column_names: Vec<String> = (0..sql_typ.column_types.len())
1286 .map(|i| format!("column{}", i))
1287 .collect();
1288 let relation_desc = RelationDesc::new(sql_typ, column_names.iter().map(|s| s.as_str()));
1289
1290 // MIR ⇒ MIR optimization (global). The mutation is already applied in
1291 // MIR, so we hand the expression to the subscribe optimizer directly
1292 // instead of going through the `SubscribePlan` path, which expects HIR.
1293 // An empty `output` makes the sink emit raw diffs.
1294 // `optimize_query` is not an `Optimize` impl, so wrap it by hand.
1295 let global_mir_plan = mz_transform::catch_unwind_optimize(|| {
1296 optimizer.optimize_query(expr, relation_desc, vec![])
1297 })?;
1298
1299 Ok((optimizer, global_mir_plan))
1300 }
1301
1302 /// The governing oracle's read timestamp, used as a lower bound for
1303 /// timestamp selection. `None` when the selection needs no oracle.
1304 async fn oracle_read_ts(
1305 &mut self,
1306 timeline: &TimelineContext,
1307 ) -> Result<Option<Timestamp>, AdapterError> {
1308 match governing_timeline(timeline) {
1309 Some(timeline) => {
1310 let oracle = self.ensure_oracle(timeline).await?;
1311 Ok(Some(oracle.read_ts().await))
1312 }
1313 None => Ok(None),
1314 }
1315 }
1316
1317 /// Block until the oracle for this query's timeline has advanced to
1318 /// `as_of`. Returns immediately if it already has.
1319 ///
1320 /// This implements the strict-serializable read guarantee for RTW:
1321 /// once this returns, any session observing the oracle sees a read
1322 /// timestamp at least as large as `as_of`, so reads at `as_of` (and
1323 /// writes derived from them) cannot appear to "go backwards" relative
1324 /// to subsequent queries.
1325 async fn ensure_read_linearized(
1326 &mut self,
1327 timeline: &TimelineContext,
1328 as_of: Timestamp,
1329 ) -> Result<(), AdapterError> {
1330 // Linearization must target the oracle future readers of the target
1331 // table will consult, which is why this uses `governing_timeline` and
1332 // not `TimelineContext::timeline()`. The latter answers "is there a
1333 // source-forced timeline?" and would skip linearization entirely for a
1334 // read side that pins none.
1335 let tl = match governing_timeline(timeline) {
1336 Some(tl) => tl,
1337 None => return Ok(()),
1338 };
1339
1340 // Cloned before `ensure_oracle` borrows `self` for the rest of this
1341 // function. The handle is an `Arc` internally, so this is cheap.
1342 let group_commit_notifier = self.group_commit_notifier.clone();
1343 let oracle = self.ensure_oracle(tl).await?;
1344
1345 // The oracle advances only when a group commit applies, and an empty
1346 // group commit is already the periodic keepalive. So when we have
1347 // nothing to write ourselves, waiting for the next tick costs up to a
1348 // full `default_timestamp_interval`. We ask for that commit instead of
1349 // waiting for it, which also spares the oracle the ~1ms poll below
1350 // running for the whole interval.
1351 //
1352 // Once per wait rather than once per poll. The committer never
1353 // allocates a write timestamp above wall clock, so a far-future `as_of`
1354 // cannot be reached by asking, and nudging per iteration would spin for
1355 // as long as such a statement legitimately parks. That case pays one
1356 // empty commit, which is what the keepalive would have done anyway.
1357 let mut nudged = false;
1358
1359 loop {
1360 let oracle_ts = oracle.read_ts().await;
1361 if as_of <= oracle_ts {
1362 return Ok(());
1363 }
1364
1365 if !nudged {
1366 group_commit_notifier.notify();
1367 nudged = true;
1368 }
1369
1370 // Sleep for roughly the difference between as_of and the current
1371 // oracle timestamp. Since timestamps are epoch milliseconds, the
1372 // difference is the approximate wall-clock time we need to wait.
1373 // Cap at 1s to avoid very long sleeps if clocks are skewed,
1374 // matching the cap in `message_linearize_reads`.
1375 let wait_ms = u64::from(as_of.saturating_sub(oracle_ts));
1376 let wait = Duration::from_millis(wait_ms).min(Duration::from_secs(1));
1377 tokio::time::sleep(wait).await;
1378 }
1379 }
1380
1381 /// Submits frontier-independent diffs to group commit, which picks the
1382 /// write timestamp, and returns the timestamp the write committed at.
1383 ///
1384 /// Only valid for diffs that do not depend on an observed read frontier:
1385 /// the write lands at a timestamp this caller does not choose.
1386 async fn submit_blind_write(
1387 &self,
1388 conn_id: mz_adapter_types::connection::ConnectionId,
1389 target_id: CatalogItemId,
1390 target_global_id: GlobalId,
1391 diffs: Vec<(Row, Diff)>,
1392 statement_logging_id: Option<StatementLoggingId>,
1393 attempt_state: &FrontendWriteAttemptState,
1394 ) -> Result<Timestamp, AdapterError> {
1395 attempt_state.mark_write_submitted();
1396 let result = self
1397 .call_coordinator(|tx| Command::AttemptWrite {
1398 attempt: WriteAttemptKind::Session {
1399 conn_id,
1400 write_ts: None,
1401 },
1402 target_id,
1403 target_global_id,
1404 diffs,
1405 tx,
1406 })
1407 .await?;
1408
1409 // Every outcome here terminates the attempt, so `write_submitted`
1410 // stays set per its contract.
1411 match classify_write_result(result, target_id, attempt_state) {
1412 WriteOutcome::Committed(timestamp) => {
1413 if let Some(id) = statement_logging_id {
1414 self.log_set_timestamp(id, timestamp);
1415 }
1416 Ok(timestamp)
1417 }
1418 WriteOutcome::Failed(err) => Err(err),
1419 WriteOutcome::Conflict { .. } => {
1420 // Unreachable: a write that requests no timestamp cannot have
1421 // one pass. Group commit resolves it through
1422 // `UserWriteResponder::Internal`, which only reports a conflict
1423 // to a write that asked for a specific timestamp.
1424 soft_panic_or_log!("blind read-then-write unexpectedly got TimestampPassed");
1425 Err(AdapterError::Internal(
1426 "blind write unexpectedly got TimestampPassed".into(),
1427 ))
1428 }
1429 }
1430 }
1431
1432 /// Creates an internal subscribe, meaning one that writes no
1433 /// `mz_subscriptions` row. Returns a [`SubscribeHandle`] that ensures
1434 /// cleanup on drop.
1435 async fn create_internal_subscribe(
1436 &self,
1437 df_desc: Box<optimize::LirDataflowDescription>,
1438 cluster_id: ComputeInstanceId,
1439 replica_id: Option<ReplicaId>,
1440 depends_on: BTreeSet<GlobalId>,
1441 as_of: Timestamp,
1442 arity: usize,
1443 sink_id: GlobalId,
1444 owner: ActiveSubscribeOwner,
1445 start_time: mz_ore::now::EpochMillis,
1446 read_holds: crate::ReadHolds,
1447 ) -> Result<SubscribeHandle, AdapterError> {
1448 let rx: mpsc::UnboundedReceiver<PeekResponseUnary> = self
1449 .call_coordinator(|tx| Command::CreateInternalSubscribe {
1450 df_desc,
1451 cluster_id,
1452 replica_id,
1453 depends_on,
1454 as_of,
1455 arity,
1456 sink_id,
1457 owner,
1458 start_time,
1459 read_holds,
1460 tx,
1461 })
1462 .await??;
1463
1464 Ok(SubscribeHandle {
1465 rx,
1466 sink_id,
1467 client: Some(self.coordinator_client().clone()),
1468 })
1469 }
1470
1471 /// Run the OCC loop: drain the subscribe at `as_of`, apply the
1472 /// mutation, and submit the resulting diffs as a write.
1473 ///
1474 /// Semantically a SELECT at `target - 1` followed by an INSERT at `target`.
1475 /// `write_oracle` chooses `target`, the subscribe's frontier certifies the
1476 /// payload is complete below it, and a target the target table has moved
1477 /// past comes back as `WriteResult::TimestampPassed`, whose next eligible
1478 /// timestamp the loop adopts. At most `max_occ_retries` attempts.
1479 ///
1480 /// A subscribe that ends on its own has diffs no frontier can change, and
1481 /// those are returned as [`OccOutcome::Blind`] rather than written.
1482 ///
1483 /// Contract on the caller, both ends of the read: the oracle's read
1484 /// timestamp must be at or above `as_of` on entry, and an
1485 /// [`OccOutcome::NoRowsMatched`] must be linearized against its
1486 /// `empty_as_of` before the response goes out.
1487 ///
1488 /// `write_oracle` is `None` only for a timestamp-independent selection. Such
1489 /// a statement reads at `Timestamp::maximum()`, so it observes no progress
1490 /// past its `as_of` and always leaves through the blind path.
1491 ///
1492 /// Returns `(retry_count, result)` so the caller can record OCC retry
1493 /// metrics regardless of whether the operation succeeded or failed.
1494 async fn run_occ_loop(
1495 &self,
1496 mut subscribe_handle: SubscribeHandle,
1497 target_id: CatalogItemId,
1498 target_global_id: GlobalId,
1499 kind: MutationKind,
1500 returning: Vec<MirScalarExpr>,
1501 max_result_size: u64,
1502 max_query_result_size: u64,
1503 row_set_finishing_seconds: Histogram,
1504 max_occ_retries: usize,
1505 table_desc: RelationDesc,
1506 write_conn_id: Option<mz_adapter_types::connection::ConnectionId>,
1507 statement_logging_id: Option<StatementLoggingId>,
1508 as_of: Timestamp,
1509 write_oracle: Option<Arc<dyn TimestampOracle<Timestamp> + Send + Sync>>,
1510 attempt_state: &FrontendWriteAttemptState,
1511 ) -> (usize, Result<OccOutcome, AdapterError>) {
1512 let mut state = OccState::new();
1513
1514 // The timestamp the next attempt writes at, chosen when we are first
1515 // ready to attempt one and replaced only by a conflict. `None` until
1516 // then.
1517 let mut write_target: Option<Timestamp> = None;
1518
1519 // The smallest timestamp an attempt may target. `as_of` itself is out,
1520 // since the payload has to contain the snapshot the subscribe emits
1521 // there.
1522 //
1523 // `as_of` is `Timestamp::MAX` for a selection with no timestamp at all,
1524 // which is the selection whose subscribe closes on its own and leaves
1525 // through the blind path rather than a write. There is no timestamp
1526 // above `MAX`, so saturating keeps this total instead of asserting a
1527 // property of a value the write path never uses.
1528 let min_target = as_of.try_step_forward().unwrap_or(as_of);
1529
1530 // Retry invariant: the payload is the selection consolidated at
1531 // `target - 1`, and the diffs at or above `target` are concurrent with
1532 // the write, so a retry folds them in only once it raises the target.
1533 let result = loop {
1534 if let Some(error) = attempt_state.requested_error() {
1535 break Err(error);
1536 }
1537
1538 // Before a target is chosen, `min_target` is a lower bound on it, so
1539 // folding there consolidates the snapshot without admitting a diff a
1540 // write would have to treat as concurrent.
1541 let fold_target = write_target.unwrap_or(min_target);
1542
1543 // Already certified for the target we hold? Write. Otherwise wait for
1544 // the next subscribe message. Waiting first would hang after a
1545 // conflict, since an input settled until its next refresh sends
1546 // nothing further and does not close the channel either.
1547 //
1548 // Termination: the write arm awaits a round trip and only a conflict
1549 // returns to this one, raising `retry_count` towards
1550 // `max_occ_retries`, so neither arm spins.
1551 let attempt_write = match write_target {
1552 Some(target) if state.current_upper.is_some_and(|upper| upper >= target) => true,
1553 _ => {
1554 let msg = match subscribe_handle.recv().await {
1555 Some(msg) => msg,
1556 None => {
1557 // The channel closed cleanly, which says the
1558 // selection can never change again, so these diffs
1559 // are frontier-independent. It does not say the
1560 // selection reads nothing persisted, a sealed
1561 // `REFRESH AT` MV closes cleanly too. Either way no
1562 // target separates them, and the caller decides
1563 // where they go.
1564 state.fold_all();
1565 if state.payload.is_empty() {
1566 break Ok(OccOutcome::NoRowsMatched {
1567 response: build_no_rows_response(&kind),
1568 empty_as_of: None,
1569 });
1570 }
1571 let success_response = match build_success_response(
1572 &kind,
1573 &returning,
1574 &state.payload,
1575 max_result_size,
1576 max_query_result_size,
1577 &row_set_finishing_seconds,
1578 ) {
1579 Ok(response) => response,
1580 Err(e) => break Err(e),
1581 };
1582
1583 break Ok(OccOutcome::Blind {
1584 response: success_response,
1585 diffs: std::mem::take(&mut state.payload),
1586 });
1587 }
1588 };
1589
1590 match process_message(
1591 msg,
1592 &mut state,
1593 as_of,
1594 fold_target,
1595 max_result_size,
1596 &table_desc,
1597 ) {
1598 ProcessResult::Continue { ready_to_write } => ready_to_write,
1599 ProcessResult::NoRowsMatched { empty_as_of } => {
1600 break Ok(OccOutcome::NoRowsMatched {
1601 response: build_no_rows_response(&kind),
1602 empty_as_of: Some(empty_as_of),
1603 });
1604 }
1605 ProcessResult::Error(e) => break Err(e),
1606 }
1607 }
1608 };
1609
1610 if !attempt_write {
1611 continue;
1612 }
1613
1614 // Drain buffered messages before attempting the write.
1615 let drain_err = loop {
1616 match subscribe_handle.try_recv() {
1617 Ok(msg) => {
1618 match process_message(
1619 msg,
1620 &mut state,
1621 as_of,
1622 fold_target,
1623 max_result_size,
1624 &table_desc,
1625 ) {
1626 ProcessResult::Continue { .. } => {}
1627 ProcessResult::NoRowsMatched { empty_as_of } => {
1628 break Some(Ok(OccOutcome::NoRowsMatched {
1629 response: build_no_rows_response(&kind),
1630 empty_as_of: Some(empty_as_of),
1631 }));
1632 }
1633 ProcessResult::Error(e) => {
1634 break Some(Err(e));
1635 }
1636 }
1637 }
1638 Err(mpsc::error::TryRecvError::Empty) => break None,
1639 // The subscribe can finish (coordinator drops the sender
1640 // after `process_response` returns true) between our last
1641 // recv() and this drain. This is benign, all buffered
1642 // messages have already been consumed via the Ok(msg) arm
1643 // above.
1644 Err(mpsc::error::TryRecvError::Disconnected) => break None,
1645 }
1646 };
1647 if let Some(result) = drain_err {
1648 break result;
1649 }
1650
1651 let upper = state
1652 .current_upper
1653 .expect("a write attempt requires an observed frontier");
1654
1655 let target = match write_target {
1656 Some(target) => target,
1657 None => {
1658 let Some(oracle) = &write_oracle else {
1659 // Invariant: a statement with no governing timeline
1660 // reads at `as_of == Timestamp::maximum()`, so it
1661 // observes no progress past its `as_of` and leaves
1662 // through the blind arm above rather than reaching a
1663 // write.
1664 soft_panic_or_log!(
1665 "read-then-write reached a write attempt with no governing timeline"
1666 );
1667 break Err(AdapterError::Internal(
1668 "read-then-write has no oracle to take a write timestamp from".into(),
1669 ));
1670 };
1671
1672 // One step above the oracle's write timestamp is the smallest
1673 // value `commit_timestamped` accepts.
1674 let peek_write_ts = oracle.peek_write_ts().await;
1675 let Some(chosen) = peek_write_ts.try_step_forward() else {
1676 // A timeline that reached `Timestamp::MAX` is a broken
1677 // environment, not anything this statement did.
1678 soft_panic_or_log!(
1679 "read-then-write cannot target a timestamp above the write \
1680 timeline's timestamp {peek_write_ts}"
1681 );
1682 break Err(AdapterError::Internal(format!(
1683 "write timeline exhausted at timestamp {peek_write_ts}"
1684 )));
1685 };
1686
1687 // Unreachable while the oracle's read timestamp is at or
1688 // above `as_of` on entry, and the clamp keeps the payload
1689 // rule rather than only reporting the violation.
1690 if chosen < min_target {
1691 soft_panic_or_log!(
1692 "read-then-write target {chosen} does not clear the as_of {as_of}, \
1693 so the payload would miss the snapshot"
1694 );
1695 }
1696 let chosen = std::cmp::max(chosen, min_target);
1697
1698 write_target = Some(chosen);
1699 chosen
1700 }
1701 };
1702
1703 // Fold in what the drain picked up, plus anything a target raised
1704 // by the last conflict now admits.
1705 state.fold_below(target);
1706
1707 // A write at `target` needs every diff below it, which is what
1708 // progress at or above `target` certifies. Waiting for the next
1709 // message is bounded by `statement_timeout`, like every wait here.
1710 if upper < target {
1711 continue;
1712 }
1713
1714 if state.payload.is_empty() {
1715 // Everything below `target` cancelled out, so there is nothing to
1716 // write and the answer holds as of `target - 1`. Diffs pending at
1717 // or above `target` are concurrent with the write this would have
1718 // been and do not enter it.
1719 break Ok(OccOutcome::NoRowsMatched {
1720 response: build_no_rows_response(&kind),
1721 empty_as_of: Some(empty_as_of(target)),
1722 });
1723 }
1724
1725 let success_response = match build_success_response(
1726 &kind,
1727 &returning,
1728 &state.payload,
1729 max_result_size,
1730 max_query_result_size,
1731 &row_set_finishing_seconds,
1732 ) {
1733 Ok(response) => response,
1734 Err(e) => break Err(e),
1735 };
1736
1737 // Submit write.
1738 //
1739 // TODO(aljoscha): Store `Arc<Row>` in the payload if this shows up
1740 // in profiles. Every attempt clones every row, and we retry up to
1741 // `max_occ_retries` times.
1742 attempt_state.mark_write_submitted();
1743 let result = match self
1744 .call_coordinator(|tx| Command::AttemptWrite {
1745 attempt: match write_conn_id.clone() {
1746 Some(conn_id) => WriteAttemptKind::Session {
1747 conn_id,
1748 write_ts: Some(target),
1749 },
1750 None => WriteAttemptKind::Background { write_ts: target },
1751 },
1752 target_id,
1753 target_global_id,
1754 diffs: state.payload.clone(),
1755 tx,
1756 })
1757 .await
1758 {
1759 Ok(result) => result,
1760 Err(error) => break Err(error),
1761 };
1762
1763 match classify_write_result(result, target_id, attempt_state) {
1764 WriteOutcome::Committed(timestamp) => {
1765 if let Some(id) = statement_logging_id {
1766 self.log_set_timestamp(id, timestamp);
1767 }
1768 // N.B. subscribe_handle is dropped here, which fires off
1769 // the cleanup message.
1770 break Ok(OccOutcome::Committed {
1771 response: success_response,
1772 write_ts: timestamp,
1773 });
1774 }
1775 WriteOutcome::Failed(err) => break Err(err),
1776 WriteOutcome::Conflict {
1777 next_eligible_timestamp,
1778 } => {
1779 // The write definitively did not land, so the attempt is
1780 // resolved. Clearing `write_submitted` lets a cancel or
1781 // statement timeout that fires during the upcoming
1782 // subscribe wait resolve promptly instead of awaiting a
1783 // write result.
1784 attempt_state.mark_write_resolved();
1785 // Adopt the timestamp the committer reported as next
1786 // eligible. The accumulated diffs say nothing about it
1787 // yet, and they do not have to: the readiness check above
1788 // holds the next attempt until the subscribe has certified
1789 // everything below the new target, and the fold then moves
1790 // the diffs in between into the payload.
1791 write_target = Some(next_eligible_timestamp);
1792 state.retry_count += 1;
1793 // Cancellation wins over the retry budget: if both apply,
1794 // the user asked us to stop and that is the more truthful
1795 // answer.
1796 if let Some(error) = attempt_state.requested_error() {
1797 break Err(error);
1798 }
1799 if state.retry_count > max_occ_retries {
1800 // Contention is a user-visible condition, not an
1801 // internal invariant violation, and every attempt was
1802 // refused before anything was appended, so the
1803 // statement is retryable.
1804 break Err(AdapterError::ReadThenWriteContention);
1805 }
1806 tracing::debug!(
1807 retry_count = state.retry_count,
1808 write_ts = %target,
1809 next_eligible_timestamp = %next_eligible_timestamp,
1810 "OCC write conflict, retrying"
1811 );
1812 continue;
1813 }
1814 }
1815 };
1816
1817 (state.retry_count, result)
1818 }
1819}
1820
1821/// Result of validating a read-then-write operation.
1822struct ValidationResult {
1823 cluster_id: ComputeInstanceId,
1824 replica_id: Option<ReplicaId>,
1825 timeline: TimelineContext,
1826 depends_on: BTreeSet<GlobalId>,
1827 /// The table descriptor, used for constraint validation.
1828 table_desc: RelationDesc,
1829}
1830
1831/// Accumulated state for the OCC loop in `run_occ_loop`.
1832///
1833/// Every diff the subscribe ever sent is kept, split at [`Self::split`]. The
1834/// split only rises, so each diff crosses it once.
1835struct OccState {
1836 /// Consolidated net diffs from strictly below [`Self::split`].
1837 payload: Vec<(Row, Diff)>,
1838 /// Diffs at or above [`Self::split`], consolidated by `(row, timestamp)`.
1839 pending: Vec<(Row, Timestamp, Diff)>,
1840 /// Where the last fold split the diffs, `None` before the first one.
1841 split: Option<Timestamp>,
1842 /// Timestamp of the last progress message, which certifies that no diff
1843 /// will arrive below it.
1844 current_upper: Option<Timestamp>,
1845 retry_count: usize,
1846 /// Row bytes held in `payload` and `pending` together, which is what the
1847 /// `max_result_size` check measures. `pending` is consolidated by
1848 /// `(row, timestamp)`, so a row touched at several timestamps occupies
1849 /// several entries until a rising split folds them together, and the count
1850 /// can exceed the size of the payload that eventually goes out.
1851 byte_size: u64,
1852}
1853
1854impl OccState {
1855 fn new() -> Self {
1856 Self {
1857 payload: Vec::new(),
1858 pending: Vec::new(),
1859 split: None,
1860 current_upper: None,
1861 retry_count: 0,
1862 byte_size: 0,
1863 }
1864 }
1865
1866 /// Raises the split to `split`, moving the diffs below it into the payload.
1867 ///
1868 /// Lowering it is a bug: the payload is consolidated and never re-split, so
1869 /// the diffs above a lowered split would stay in it. We clamp to the old
1870 /// split, which keeps the payload's contract intact.
1871 fn fold_below(&mut self, split: Timestamp) {
1872 let split = match self.split {
1873 Some(previous) if split < previous => {
1874 soft_panic_or_log!(
1875 "read-then-write folded at {split}, below its previous split {previous}"
1876 );
1877 previous
1878 }
1879 _ => split,
1880 };
1881 self.split = Some(split);
1882 self.fold(Some(split));
1883 }
1884
1885 /// Moves every accumulated diff into the payload, whatever its timestamp.
1886 ///
1887 /// Only valid once the subscribe has run to completion, where the diffs are
1888 /// frontier-independent and no split separates them.
1889 fn fold_all(&mut self) {
1890 self.fold(None);
1891 }
1892
1893 /// Moves the diffs below `split`, or all of them when it is `None`, into the
1894 /// payload, consolidates both halves, and recomputes `byte_size`.
1895 fn fold(&mut self, split: Option<Timestamp>) {
1896 for (row, ts, diff) in std::mem::take(&mut self.pending) {
1897 match split {
1898 Some(split) if ts >= split => self.pending.push((row, ts, diff)),
1899 _ => self.payload.push((row, diff)),
1900 }
1901 }
1902 consolidation::consolidate(&mut self.payload);
1903 consolidation::consolidate_updates(&mut self.pending);
1904 self.byte_size = self
1905 .payload
1906 .iter()
1907 .map(|(row, _)| u64::cast_from(row.byte_len()))
1908 .chain(
1909 self.pending
1910 .iter()
1911 .map(|(row, _, _)| u64::cast_from(row.byte_len())),
1912 )
1913 .sum();
1914 }
1915
1916 /// Whether nothing has been accumulated on either side of the split.
1917 fn is_empty(&self) -> bool {
1918 self.payload.is_empty() && self.pending.is_empty()
1919 }
1920}
1921
1922/// Result of processing a single subscribe message in the OCC loop.
1923enum ProcessResult {
1924 Continue {
1925 ready_to_write: bool,
1926 },
1927 /// The consolidated selection is empty, as of the timestamp reported. See
1928 /// [`OccOutcome::NoRowsMatched`].
1929 NoRowsMatched {
1930 empty_as_of: Timestamp,
1931 },
1932 Error(AdapterError),
1933}
1934
1935/// Process one subscribe message, updating `state` in place.
1936///
1937/// Data rows are accumulated into `state` (with per-row constraint and
1938/// max-result-size checks). Progress messages fold everything below
1939/// `fold_target` into the payload and can promote the accumulated diffs to
1940/// "ready to write".
1941///
1942/// `fold_target` must not exceed the timestamp the next write attempt uses, or
1943/// the payload takes in a diff that is concurrent with that write.
1944fn process_message(
1945 response: PeekResponseUnary,
1946 state: &mut OccState,
1947 as_of: Timestamp,
1948 fold_target: Timestamp,
1949 max_result_size: u64,
1950 table_desc: &RelationDesc,
1951) -> ProcessResult {
1952 match response {
1953 PeekResponseUnary::Rows(mut rows) => {
1954 let mut saw_progress = false;
1955
1956 while let Some(row) = rows.next() {
1957 let mut datums = row.iter();
1958
1959 // Extract mz_timestamp (SubscribeOutput::Diffs format:
1960 // mz_timestamp, mz_progressed, mz_diff, ...data columns...).
1961 //
1962 // Format drift would mean we'd silently commit an incorrect
1963 // write, so surface every shape mismatch as an internal
1964 // error rather than panicking the process.
1965 let Some(ts_datum) = datums.next() else {
1966 return ProcessResult::Error(AdapterError::Internal(
1967 "missing mz_timestamp in subscribe output".into(),
1968 ));
1969 };
1970 let ts = match ts_datum {
1971 mz_repr::Datum::Numeric(n) => match n.0.try_into() {
1972 Ok(ts_u64) => Timestamp::new(ts_u64),
1973 Err(_) => {
1974 return ProcessResult::Error(AdapterError::Internal(format!(
1975 "mz_timestamp in subscribe output is not a valid u64: {n}"
1976 )));
1977 }
1978 },
1979 other => {
1980 return ProcessResult::Error(AdapterError::Internal(format!(
1981 "unexpected mz_timestamp datum: {other:?}"
1982 )));
1983 }
1984 };
1985
1986 let Some(progressed_datum) = datums.next() else {
1987 return ProcessResult::Error(AdapterError::Internal(
1988 "missing mz_progressed in subscribe output".into(),
1989 ));
1990 };
1991 let is_progress = matches!(progressed_datum, mz_repr::Datum::True);
1992
1993 if is_progress {
1994 state.current_upper = Some(ts);
1995 saw_progress = true;
1996
1997 // Fold and consolidate incrementally on each progress
1998 // message. This keeps memory bounded by the consolidated
1999 // size and makes the byte_size check below accurate (except
2000 // for rows received between two progress messages, which is
2001 // a small window).
2002 state.fold_below(fold_target);
2003
2004 // NOTE: The first progress message is always at `as_of`,
2005 // emitted by `ActiveSubscribe::initialize` before any data
2006 // batch, so the accumulation is empty there whatever the
2007 // snapshot holds. Later progress is gated on `batch.upper >
2008 // as_of` (see `crate::active_compute_sink`), so `ts > as_of`
2009 // is what distinguishes a real answer from that first one.
2010 //
2011 // Nothing accumulated at all, so no write is coming and the
2012 // loop would otherwise wait for diffs that will not arrive.
2013 // Our view is complete below `ts` and the payload covers
2014 // below `fold_target`, so the emptiness holds as of one below
2015 // the earlier of the two.
2016 if ts > as_of && state.is_empty() {
2017 return ProcessResult::NoRowsMatched {
2018 empty_as_of: empty_as_of(std::cmp::min(ts, fold_target)),
2019 };
2020 }
2021 } else {
2022 let Some(diff_datum) = datums.next() else {
2023 return ProcessResult::Error(AdapterError::Internal(
2024 "missing mz_diff in subscribe output".into(),
2025 ));
2026 };
2027 let diff = match diff_datum {
2028 mz_repr::Datum::Int64(d) => Diff::from(d),
2029 other => {
2030 return ProcessResult::Error(AdapterError::Internal(format!(
2031 "unexpected mz_diff datum while processing read-then-write: {other:?}"
2032 )));
2033 }
2034 };
2035
2036 let data_row = Row::pack(datums);
2037
2038 // Validate constraints for rows being added (positive diff)
2039 if diff.is_positive() {
2040 for (idx, datum) in data_row.iter().enumerate() {
2041 if let Err(e) = table_desc.constraints_met(idx, &datum) {
2042 return ProcessResult::Error(e.into());
2043 }
2044 }
2045 }
2046
2047 state.byte_size = state
2048 .byte_size
2049 .saturating_add(u64::cast_from(data_row.byte_len()));
2050 if state.byte_size > max_result_size {
2051 return ProcessResult::Error(AdapterError::ResultSize(format!(
2052 "result exceeds max size of {}",
2053 ByteSize::b(max_result_size)
2054 )));
2055 }
2056 state.pending.push((data_row, ts, diff));
2057 }
2058 }
2059
2060 // The complement of the zero-row exit above: something accumulated
2061 // means a write is coming, nothing at all means there is none.
2062 let ready_to_write = saw_progress && !state.is_empty();
2063 ProcessResult::Continue { ready_to_write }
2064 }
2065 PeekResponseUnary::Error(e) => {
2066 ProcessResult::Error(AdapterError::Unstructured(anyhow::anyhow!(e)))
2067 }
2068 // Match the lock path's classification. `Unstructured` would render
2069 // this as an internal error (XX000) for what is an ordinary concurrent
2070 // DDL race.
2071 PeekResponseUnary::DependencyDropped(dep) => {
2072 ProcessResult::Error(dep.to_concurrent_dependency_drop())
2073 }
2074 PeekResponseUnary::Canceled => ProcessResult::Error(AdapterError::Canceled),
2075 }
2076}
2077
2078/// The timestamp an answer holds as of, given a view complete strictly below
2079/// `complete_below`.
2080///
2081/// Both callers derive `complete_below` from a timestamp strictly above the
2082/// subscribe's `as_of`, so it is never `Timestamp::MIN` and the saturating
2083/// fallback is unreachable.
2084fn empty_as_of(complete_below: Timestamp) -> Timestamp {
2085 complete_below.step_back().unwrap_or(complete_below)
2086}
2087
2088/// Build the response returned when no rows matched the selection.
2089///
2090/// Bug-compatible with the coordinator path, which evaluates RETURNING over the
2091/// diffs and so reports a plain row count when there are none. Postgres returns
2092/// an empty result set for a zero-row `INSERT ... RETURNING` instead, but
2093/// changing that is a change to the path that ships today, not to this one.
2094fn build_no_rows_response(kind: &MutationKind) -> ExecuteResponse {
2095 match kind {
2096 MutationKind::Delete => ExecuteResponse::Deleted(0),
2097 MutationKind::Update => ExecuteResponse::Updated(0),
2098 MutationKind::Insert => ExecuteResponse::Inserted(0),
2099 }
2100}
2101
2102/// Transform a MIR expression to produce the appropriate diffs for a mutation.
2103///
2104/// - DELETE: Negates the expression to produce `(row, -1)` diffs
2105/// - UPDATE: Unions negated old rows with mapped new rows to produce both
2106/// `(old_row, -1)` and `(new_row, +1)` diffs
2107fn apply_mutation_to_mir(
2108 expr: MirRelationExpr,
2109 kind: &MutationKind,
2110 assignments: &BTreeMap<usize, MirScalarExpr>,
2111) -> MirRelationExpr {
2112 match kind {
2113 MutationKind::Delete => MirRelationExpr::Negate {
2114 input: Box::new(expr),
2115 },
2116 MutationKind::Update => {
2117 let arity = expr.arity();
2118
2119 // Find a fresh LocalId that won't conflict with any in the expression.
2120 //
2121 // Invariant: `Let` and `LetRec` are the only MIR nodes that *bind*
2122 // LocalIds. `Get` references them but does not introduce new ones.
2123 // So scanning just those two node kinds and picking `max + 1` is
2124 // guaranteed to produce an id unused by the subtree.
2125 let mut max_id = 0_u64;
2126 expr.visit_pre(|e| match e {
2127 MirRelationExpr::Let { id, .. } => {
2128 max_id = std::cmp::max(max_id, id.into());
2129 }
2130 MirRelationExpr::LetRec { ids, .. } => {
2131 for id in ids {
2132 max_id = std::cmp::max(max_id, id.into());
2133 }
2134 }
2135 _ => {}
2136 });
2137 let binding_id = LocalId::new(max_id + 1);
2138
2139 let get_binding = MirRelationExpr::Get {
2140 id: Id::Local(binding_id),
2141 typ: expr.typ(),
2142 access_strategy: mz_expr::AccessStrategy::UnknownOrLocal,
2143 };
2144
2145 let map_scalars: Vec<MirScalarExpr> = (0..arity)
2146 .map(|i| {
2147 assignments
2148 .get(&i)
2149 .cloned()
2150 .unwrap_or_else(|| MirScalarExpr::column(i))
2151 })
2152 .collect();
2153
2154 let new_rows = get_binding
2155 .clone()
2156 .map(map_scalars)
2157 .project((arity..2 * arity).collect());
2158
2159 let old_rows = MirRelationExpr::Negate {
2160 input: Box::new(get_binding),
2161 };
2162
2163 let body = new_rows.union(old_rows);
2164
2165 MirRelationExpr::Let {
2166 id: binding_id,
2167 value: Box::new(expr),
2168 body: Box::new(body),
2169 }
2170 }
2171 // INSERT: rows pass through unchanged, the subscribe emits them with
2172 // diff +1.
2173 MutationKind::Insert => expr,
2174 }
2175}
2176
2177#[cfg(test)]
2178mod tests {
2179 use mz_repr::adt::numeric;
2180 use mz_repr::{Datum, IntoRowIterator};
2181
2182 use super::*;
2183
2184 fn row(value: i64) -> Row {
2185 Row::pack_slice(&[Datum::Int64(value)])
2186 }
2187
2188 /// A progress message in the subscribe's `SubscribeOutput::Diffs` shape:
2189 /// `mz_timestamp, mz_progressed, mz_diff, data...`.
2190 fn progress(ts: u64) -> PeekResponseUnary {
2191 let mut row = Row::default();
2192 let mut packer = row.packer();
2193 packer.push(Datum::from(numeric::Numeric::from(ts)));
2194 packer.push(Datum::True);
2195 packer.push(Datum::Null);
2196 PeekResponseUnary::Rows(Box::new(row.into_row_iter()))
2197 }
2198
2199 /// Accumulates `(row value, timestamp, diff)` triples the way
2200 /// `process_message` does, without going through a subscribe.
2201 fn accumulate(diffs: impl IntoIterator<Item = (i64, u64, i64)>) -> OccState {
2202 let mut state = OccState::new();
2203 for (value, ts, diff) in diffs {
2204 state
2205 .pending
2206 .push((row(value), Timestamp::new(ts), Diff::from(diff)));
2207 }
2208 state
2209 }
2210
2211 /// The target is the boundary the write turns on, so the off-by-one is the
2212 /// whole point: a diff at exactly the target is concurrent with the write
2213 /// and must not be in its payload.
2214 #[mz_ore::test]
2215 fn test_fold_below_splits_at_the_target() {
2216 let mut state = accumulate([(1, 9, 1), (2, 10, 1), (3, 11, 1)]);
2217 state.fold_below(Timestamp::new(10));
2218
2219 assert_eq!(state.payload, vec![(row(1), Diff::ONE)]);
2220 assert_eq!(
2221 state.pending,
2222 vec![
2223 (row(2), Timestamp::new(10), Diff::ONE),
2224 (row(3), Timestamp::new(11), Diff::ONE),
2225 ]
2226 );
2227 }
2228
2229 /// A retry raises the target, which is what admits the diffs that were
2230 /// concurrent with the attempt that lost.
2231 #[mz_ore::test]
2232 fn test_fold_below_moves_each_diff_once() {
2233 let mut state = accumulate([(1, 9, 1), (2, 10, 1), (3, 11, 1)]);
2234
2235 state.fold_below(Timestamp::new(10));
2236 assert_eq!(state.payload, vec![(row(1), Diff::ONE)]);
2237
2238 state.fold_below(Timestamp::new(11));
2239 assert_eq!(
2240 state.payload,
2241 vec![(row(1), Diff::ONE), (row(2), Diff::ONE)]
2242 );
2243 assert_eq!(state.pending, vec![(row(3), Timestamp::new(11), Diff::ONE)]);
2244
2245 state.fold_below(Timestamp::new(12));
2246 assert_eq!(
2247 state.payload,
2248 vec![
2249 (row(1), Diff::ONE),
2250 (row(2), Diff::ONE),
2251 (row(3), Diff::ONE),
2252 ]
2253 );
2254 assert!(state.pending.is_empty());
2255 }
2256
2257 /// A row inserted and retracted below the target leaves nothing behind,
2258 /// which is what makes the payload the net change rather than a log.
2259 #[mz_ore::test]
2260 fn test_fold_below_cancels_opposite_diffs() {
2261 let mut state = accumulate([(1, 9, 1), (1, 10, -1), (2, 9, 1)]);
2262 state.fold_below(Timestamp::new(11));
2263
2264 assert_eq!(state.payload, vec![(row(2), Diff::ONE)]);
2265 assert!(state.pending.is_empty());
2266 assert!(!state.is_empty());
2267 }
2268
2269 /// The same rows stay pending or move to the payload depending on the
2270 /// target, so a size check that saw only one half would let a statement
2271 /// past `max_result_size` by picking the other one.
2272 #[mz_ore::test]
2273 fn test_byte_size_counts_payload_and_pending() {
2274 let row_bytes = u64::cast_from(row(1).byte_len());
2275
2276 let mut state = accumulate([(1, 9, 1), (2, 10, 1), (3, 11, 1)]);
2277 state.fold_below(Timestamp::new(10));
2278 assert_eq!(state.payload.len(), 1);
2279 assert_eq!(state.pending.len(), 2);
2280 assert_eq!(state.byte_size, 3 * row_bytes);
2281
2282 state.fold_below(Timestamp::new(12));
2283 assert!(state.pending.is_empty());
2284 assert_eq!(state.byte_size, 3 * row_bytes);
2285 }
2286
2287 /// A subscribe that ran to completion has no target to split on, and
2288 /// cancellation still applies.
2289 #[mz_ore::test]
2290 fn test_fold_all_takes_every_timestamp() {
2291 let mut state = accumulate([(1, 9, 1), (1, 10, -1), (2, u64::MAX, 1)]);
2292 state.fold_all();
2293
2294 assert_eq!(state.payload, vec![(row(2), Diff::ONE)]);
2295 assert!(state.pending.is_empty());
2296 }
2297
2298 /// A zero-row answer holds as of the timestamp the answer was reached at,
2299 /// never an input's frontier. The caller waits for the oracle to reach
2300 /// whatever it gets, and an input settled until its next refresh reports a
2301 /// frontier days out, so reporting that would spend the statement's timeout
2302 /// on an answer of "0 rows".
2303 #[mz_ore::test]
2304 fn test_zero_rows_report_the_answer_not_the_frontier() {
2305 let as_of = Timestamp::new(10);
2306 let desc = RelationDesc::empty();
2307
2308 // First pass, where the fold target is `as_of + 1`. The answer holds at
2309 // `as_of`, which the caller linearized before the subscribe started, so
2310 // it costs no wait however far out the frontier is.
2311 let mut state = OccState::new();
2312 match process_message(
2313 progress(u64::MAX / 2),
2314 &mut state,
2315 as_of,
2316 as_of.step_forward(),
2317 u64::MAX,
2318 &desc,
2319 ) {
2320 ProcessResult::NoRowsMatched { empty_as_of } => assert_eq!(empty_as_of, as_of),
2321 _ => panic!("an empty selection past `as_of` must report no rows matched"),
2322 }
2323
2324 // A target raised by a conflict, with the frontier past it. The answer
2325 // holds at one below the target, the same timestamp a write there would
2326 // have been read at.
2327 let fold_target = Timestamp::new(20);
2328 let mut state = OccState::new();
2329 match process_message(
2330 progress(u64::MAX / 2),
2331 &mut state,
2332 as_of,
2333 fold_target,
2334 u64::MAX,
2335 &desc,
2336 ) {
2337 ProcessResult::NoRowsMatched { empty_as_of } => {
2338 assert_eq!(empty_as_of, Timestamp::new(19))
2339 }
2340 _ => panic!("an empty selection past `as_of` must report no rows matched"),
2341 }
2342
2343 // A frontier below the target certifies less, so the answer holds one
2344 // below the frontier instead.
2345 let mut state = OccState::new();
2346 match process_message(
2347 progress(15),
2348 &mut state,
2349 as_of,
2350 fold_target,
2351 u64::MAX,
2352 &desc,
2353 ) {
2354 ProcessResult::NoRowsMatched { empty_as_of } => {
2355 assert_eq!(empty_as_of, Timestamp::new(14))
2356 }
2357 _ => panic!("an empty selection past `as_of` must report no rows matched"),
2358 }
2359 }
2360
2361 /// Diffs waiting above the target mean a write is still coming, so the
2362 /// answer is not "no rows" yet even with an empty payload.
2363 #[mz_ore::test]
2364 fn test_pending_diffs_are_not_a_zero_row_answer() {
2365 let as_of = Timestamp::new(10);
2366 let desc = RelationDesc::empty();
2367
2368 let mut state = accumulate([(1, 30, 1)]);
2369 match process_message(
2370 progress(20),
2371 &mut state,
2372 as_of,
2373 as_of.step_forward(),
2374 u64::MAX,
2375 &desc,
2376 ) {
2377 ProcessResult::Continue { ready_to_write } => assert!(ready_to_write),
2378 _ => panic!("a selection with diffs above the target must not report no rows"),
2379 }
2380 }
2381}