Skip to main content

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}