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