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