1use std::collections::{BTreeMap, BTreeSet};
36use std::future::Future;
37use std::ops::ControlFlow;
38use std::pin::Pin;
39use std::sync::{Arc, LazyLock};
40use std::time::{Duration, Instant};
41
42use derivative::Derivative;
43use futures::future::{BoxFuture, FutureExt};
44use mz_adapter_types::connection::ConnectionId;
45use mz_adapter_types::dyncfgs::GROUP_COMMIT_MAX_ATTEMPTS;
46use mz_catalog::builtin::{BuiltinTable, MZ_SESSIONS};
47use mz_dyncfg::{ConfigSet, ConfigValHandle};
48use mz_expr::CollectionPlan;
49use mz_ore::assert_none;
50use mz_ore::halt;
51use mz_ore::instrument;
52use mz_ore::now::NowFn;
53use mz_ore::task;
54use mz_repr::{CatalogItemId, GlobalId, Timestamp};
55use mz_sql::names::ResolvedIds;
56use mz_sql::plan::{ExplainPlanPlan, ExplainTimestampPlan, Explainee, ExplaineeStatement, Plan};
57use mz_sql::session::metadata::SessionMetadata;
58use mz_storage_client::client::TableData;
59use mz_storage_client::controller::{TableRegistration, TableWriteHandle};
60use mz_storage_types::controller::StorageError;
61use mz_timestamp_oracle::{TimestampOracle, WriteTimestamp};
62use smallvec::SmallVec;
63use tokio::sync::{Notify, OwnedMutexGuard, OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
64use tracing::{Instrument, Span, info, warn};
65
66use crate::catalog::{BuiltinTableUpdate, Catalog, CatalogUpperHandle};
67use crate::coord::timeline::write_ts_upper_bound;
68use crate::coord::{Coordinator, Message, PendingTxn, PlanValidity};
69use crate::metrics::Metrics;
70use crate::session::{EndTransactionAction, GroupCommitWriteLocks, Session, WriteLocks};
71use crate::statement_logging::StatementLoggingId;
72use crate::util::{CompletedClientTransmitter, ResultExt};
73use crate::{AdapterError, ExecuteContext};
74
75pub(crate) static REQUIRED_BUILTIN_TABLES: &[&LazyLock<BuiltinTable>] = &[&MZ_SESSIONS];
77
78#[derive(Debug)]
82pub enum DeferredOp {
83 Plan(DeferredPlan),
85 Write(DeferredWrite),
87}
88
89impl DeferredOp {
90 pub(crate) fn can_be_optimistically_retried(&self) -> bool {
102 match self {
103 DeferredOp::Plan(_) => false,
104 DeferredOp::Write(_) => true,
105 }
106 }
107
108 pub fn required_locks(&self) -> impl Iterator<Item = CatalogItemId> + '_ {
110 match self {
111 DeferredOp::Plan(plan) => {
112 let iter = plan.requires_locks.iter().copied();
113 itertools::Either::Left(iter)
114 }
115 DeferredOp::Write(write) => {
116 let iter = write.writes.keys().copied();
117 itertools::Either::Right(iter)
118 }
119 }
120 }
121
122 pub fn conn_id(&self) -> &ConnectionId {
124 match self {
125 DeferredOp::Plan(plan) => plan.ctx.session().conn_id(),
126 DeferredOp::Write(write) => write.pending_txn.ctx.session().conn_id(),
127 }
128 }
129
130 pub fn into_ctx(self) -> ExecuteContext {
132 match self {
133 DeferredOp::Plan(plan) => plan.ctx,
134 DeferredOp::Write(write) => write.pending_txn.ctx,
135 }
136 }
137}
138
139#[derive(Derivative)]
141#[derivative(Debug)]
142pub struct DeferredPlan {
143 #[derivative(Debug = "ignore")]
144 pub ctx: ExecuteContext,
145 pub plan: Plan,
146 pub validity: PlanValidity,
147 pub requires_locks: BTreeSet<CatalogItemId>,
148 pub resolved_ids: ResolvedIds,
149 pub sql_impl_resolved_ids: ResolvedIds,
150}
151
152#[derive(Debug)]
153pub struct DeferredWrite {
154 pub span: Span,
155 pub writes: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
156 pub pending_txn: PendingTxn,
157}
158
159#[derive(Debug)]
161pub(crate) enum BuiltinTableUpdateSource {
162 Internal(oneshot::Sender<()>),
164 Background(oneshot::Sender<()>),
166}
167
168#[derive(Debug, Clone)]
170pub enum WriteResult {
171 Success { timestamp: Timestamp },
173 TimestampPassed {
175 target_timestamp: Timestamp,
176 next_eligible_timestamp: Timestamp,
177 },
178 TimestampTooFarAhead {
181 target_timestamp: Timestamp,
182 limit: Timestamp,
183 },
184 Canceled,
186 ReadOnly,
188 TargetChanged,
190 Indeterminate,
192}
193
194#[derive(Debug)]
201pub struct InternalWriteResponder {
202 tx: Option<oneshot::Sender<WriteResult>>,
203}
204
205impl InternalWriteResponder {
206 pub(crate) fn new(tx: oneshot::Sender<WriteResult>) -> Self {
207 Self { tx: Some(tx) }
208 }
209
210 pub(crate) fn send(mut self, result: WriteResult) {
211 if let Some(tx) = self.tx.take() {
212 let _ = tx.send(result);
213 }
214 }
215}
216
217impl Drop for InternalWriteResponder {
218 fn drop(&mut self) {
219 if let Some(tx) = self.tx.take() {
220 let _ = tx.send(WriteResult::Indeterminate);
221 }
222 }
223}
224
225#[derive(Debug)]
227pub(crate) enum UserWriteResponder {
228 Session(PendingTxn),
231 Internal {
233 conn_id: ConnectionId,
234 target: WriteTarget,
238 result: InternalWriteResponder,
239 },
240}
241
242#[derive(Debug, Clone, Copy)]
244pub(crate) struct WriteTarget {
245 pub(crate) item_id: CatalogItemId,
246 pub(crate) global_id: GlobalId,
247}
248
249impl UserWriteResponder {
250 pub(crate) fn conn_id(&self) -> &ConnectionId {
251 match self {
252 UserWriteResponder::Session(pending) => pending.ctx.session().conn_id(),
253 UserWriteResponder::Internal { conn_id, .. } => conn_id,
254 }
255 }
256}
257
258#[derive(Debug)]
260pub(crate) enum PendingWriteTxn {
261 User {
266 span: Span,
267 writes: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
269 write_locks: Option<WriteLocks>,
271 responder: UserWriteResponder,
273 },
274 System {
276 updates: Vec<BuiltinTableUpdate>,
277 source: BuiltinTableUpdateSource,
278 },
279}
280
281impl PendingWriteTxn {
282 fn is_internal_system(&self) -> bool {
283 match self {
284 PendingWriteTxn::System {
285 source: BuiltinTableUpdateSource::Internal(_),
286 ..
287 } => true,
288 _ => false,
289 }
290 }
291}
292
293pub(crate) enum TableWriteCmd {
294 GroupCommit(GroupCommitRequest),
295 TimestampedWrite(TimestampedWriteRequest),
296 Register {
297 tables: Vec<TableRegistration>,
298 result: oneshot::Sender<Timestamp>,
299 },
300 Forget {
301 ids: Vec<GlobalId>,
302 result: oneshot::Sender<Timestamp>,
303 },
304}
305
306pub(crate) struct TimestampedWriteRequest {
315 pub(crate) appends: Vec<(GlobalId, Vec<TableData>)>,
316 pub(crate) target_timestamp: Timestamp,
317 pub(crate) result: InternalWriteResponder,
318 pub(crate) span: Span,
319}
320
321pub(crate) struct GroupCommitRequest {
323 appends: Vec<(GlobalId, Vec<TableData>)>,
325 responses: Vec<CompletedClientTransmitter>,
326 statement_logging_ids: Vec<StatementLoggingId>,
327 notifies: Vec<oneshot::Sender<()>>,
328 internal_results: Vec<InternalWriteResponder>,
329 write_locks: GroupCommitWriteLocks,
330 permits: Vec<GroupCommitPermit>,
332 contains_internal_system_write: bool,
333 span: Span,
334}
335
336impl GroupCommitRequest {
337 fn merge(&mut self, other: GroupCommitRequest) {
338 let GroupCommitRequest {
339 appends,
340 responses,
341 statement_logging_ids,
342 notifies,
343 internal_results,
344 write_locks,
345 permits,
346 contains_internal_system_write,
347 span: _,
348 } = other;
349 self.appends.extend(appends);
350 self.responses.extend(responses);
351 self.statement_logging_ids.extend(statement_logging_ids);
352 self.notifies.extend(notifies);
353 self.internal_results.extend(internal_results);
354 self.write_locks.extend(write_locks);
355 self.permits.extend(permits);
356 self.contains_internal_system_write |= contains_internal_system_write;
357 }
358}
359
360pub(crate) struct GroupCommitter {
365 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
366 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
367 table_write_handle: Arc<dyn TableWriteHandle>,
368 catalog_upper: CatalogUpperHandle,
369 internal_cmd_tx: mpsc::UnboundedSender<Message>,
370 now: NowFn,
371 metrics: Metrics,
372 max_attempts: ConfigValHandle<usize>,
373}
374
375enum TxnsWriteAttempt {
377 Applied,
379 UpperConflict,
382 WorkerGone,
384}
385
386impl GroupCommitter {
387 async fn run(mut self) {
388 while let Some(cmd) = self.rx.recv().await {
389 let mut next = Some(cmd);
392 while let Some(cmd) = next.take() {
393 match cmd {
394 TableWriteCmd::GroupCommit(request) => {
395 let span = request.span.clone();
396 match self.commit(request).instrument(span).await {
397 ControlFlow::Continue(deferred) => next = deferred,
398 ControlFlow::Break(()) => return,
399 }
400 }
401 TableWriteCmd::TimestampedWrite(request) => {
402 let span = request.span.clone();
403 if self
404 .commit_timestamped(request)
405 .instrument(span)
406 .await
407 .is_break()
408 {
409 return;
410 }
411 }
412 TableWriteCmd::Register { tables, result } => {
413 let Some(write_ts) = self
414 .write_to_txns(None, |ts, _advance_to| {
415 self.table_write_handle.register(ts, tables.clone())
416 })
417 .await
418 else {
419 return;
420 };
421 let _ = result.send(write_ts.timestamp);
422 }
423 TableWriteCmd::Forget { ids, result } => {
424 let Some(write_ts) = self
425 .write_to_txns(None, |ts, _advance_to| {
426 self.table_write_handle.forget(ts, ids.clone())
427 })
428 .await
429 else {
430 return;
431 };
432 let _ = result.send(write_ts.timestamp);
433 }
434 }
435 }
436 }
437 }
438
439 async fn commit_timestamped(&self, request: TimestampedWriteRequest) -> ControlFlow<(), ()> {
462 let TimestampedWriteRequest {
463 appends,
464 target_timestamp,
465 result,
466 span: _,
467 } = request;
468
469 let oracle_write_ts = self.oracle.peek_write_ts().await;
470 if target_timestamp <= oracle_write_ts {
471 result.send(WriteResult::TimestampPassed {
472 target_timestamp,
473 next_eligible_timestamp: oracle_write_ts.step_forward(),
474 });
475 return ControlFlow::Continue(());
476 }
477
478 let now: Timestamp = (self.now)().into();
481 let limit = write_ts_upper_bound(&now);
482 if target_timestamp > limit {
483 result.send(WriteResult::TimestampTooFarAhead {
484 target_timestamp,
485 limit,
486 });
487 return ControlFlow::Continue(());
488 }
489
490 let write_ts = WriteTimestamp {
491 timestamp: target_timestamp,
492 advance_to: target_timestamp.step_forward(),
493 };
494 match self
495 .attempt_write_to_txns(
496 &write_ts,
497 Some(&self.metrics.append_table_duration_seconds),
498 |ts, advance_to| self.table_write_handle.append(ts, advance_to, appends),
499 )
500 .await
501 {
502 TxnsWriteAttempt::Applied => {}
503 TxnsWriteAttempt::UpperConflict => {
504 result.send(WriteResult::TimestampPassed {
505 target_timestamp,
506 next_eligible_timestamp: write_ts.advance_to,
507 });
508 return ControlFlow::Continue(());
509 }
510 TxnsWriteAttempt::WorkerGone => {
511 warn!("table write worker gone with a timestamped write outstanding");
512 return ControlFlow::Break(());
513 }
514 }
515
516 if self
517 .internal_cmd_tx
518 .send(Message::GroupCommitApplied {
519 responses: Vec::new(),
520 statement_logging_ids: Vec::new(),
521 internal_results: vec![result],
522 write_ts: target_timestamp,
523 })
524 .is_err()
525 {
526 warn!("coordinator shut down before a timestamped write could be finalized");
527 }
528 ControlFlow::Continue(())
529 }
530
531 async fn write_to_txns(
536 &self,
537 op_duration_metric: Option<&prometheus::Histogram>,
538 mut op: impl FnMut(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
539 ) -> Option<WriteTimestamp> {
540 let mut attempt = 0;
542 loop {
543 let max_attempts = self.max_attempts.get().max(1);
544 if attempt >= max_attempts {
545 halt!(
546 "txns-shard write reached attempt limit {max_attempts} after {attempt} conflicts, rebuilding"
547 );
548 }
549 attempt += 1;
550 let write_ts = self.oracle.write_ts().await;
551
552 match self
555 .attempt_write_to_txns(&write_ts, op_duration_metric, |ts, advance_to| {
556 op(ts, advance_to)
557 })
558 .await
559 {
560 TxnsWriteAttempt::Applied => return Some(write_ts),
561 TxnsWriteAttempt::UpperConflict => {
562 warn!(
563 write_ts = %write_ts.timestamp,
564 attempt,
565 "txns-shard write conflicted with another writer, retrying at a fresh timestamp"
566 );
567 continue;
568 }
569 TxnsWriteAttempt::WorkerGone => {
570 warn!("table write worker gone (process shutting down), winding down");
572 return None;
573 }
574 }
575 }
576 }
577
578 async fn attempt_write_to_txns(
584 &self,
585 write_ts: &WriteTimestamp,
586 op_duration_metric: Option<&prometheus::Histogram>,
587 op: impl FnOnce(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
588 ) -> TxnsWriteAttempt {
589 let catalog_upper_start = Instant::now();
590 self.catalog_upper
591 .advance_upper(write_ts.advance_to)
592 .await
593 .unwrap_or_terminate("unable to advance catalog upper");
594 self.metrics
595 .group_commit_catalog_upper_seconds
596 .observe(catalog_upper_start.elapsed().as_secs_f64());
597
598 let op_start = Instant::now();
599 let op_res = op(write_ts.timestamp, write_ts.advance_to).await;
600 if let Some(metric) = op_duration_metric {
601 metric.observe(op_start.elapsed().as_secs_f64());
602 }
603
604 match op_res {
605 Ok(Ok(())) => {}
606 Ok(Err(StorageError::InvalidUppers(_))) => return TxnsWriteAttempt::UpperConflict,
607 Ok(Err(other)) => {
608 Err::<(), _>(other).unwrap_or_terminate("cannot fail to write to txns shard");
609 unreachable!("unwrap_or_terminate does not return on Err");
610 }
611 Err(_recv) => return TxnsWriteAttempt::WorkerGone,
612 }
613
614 let now: Timestamp = (self.now)().into();
615 crate::coord::timeline::check_runaway_write_ts(&now, write_ts.timestamp);
616
617 fail::fail_point!("group_commit_before_apply_write");
626
627 self.oracle.apply_write(write_ts.timestamp).await;
628
629 TxnsWriteAttempt::Applied
630 }
631
632 async fn commit(
637 &mut self,
638 mut request: GroupCommitRequest,
639 ) -> ControlFlow<(), Option<TableWriteCmd>> {
640 let mut deferred_cmd = None;
641 let mut rx_closed = false;
644
645 loop {
649 while deferred_cmd.is_none() {
650 match self.rx.try_recv() {
651 Ok(TableWriteCmd::GroupCommit(other)) => request.merge(other),
652 Ok(other) => deferred_cmd = Some(other),
653 Err(_) => break,
654 }
655 }
656
657 if request.contains_internal_system_write || deferred_cmd.is_some() {
660 break;
661 }
662 let ts = self.oracle.peek_write_ts().await;
663 let now: Timestamp = (self.now)().into();
664 if ts <= now {
665 break;
666 }
667 let remaining_ms = std::cmp::min(ts.saturating_sub(now), Timestamp::from(1_000u64));
670 let sleep = tokio::time::sleep(Duration::from_millis(remaining_ms.into()));
671 tokio::pin!(sleep);
672 loop {
673 tokio::select! {
674 _ = &mut sleep => break,
675 cmd = self.rx.recv(), if deferred_cmd.is_none() && !rx_closed => match cmd {
676 Some(TableWriteCmd::GroupCommit(other)) => request.merge(other),
677 Some(other) => deferred_cmd = Some(other),
678 None => rx_closed = true,
679 },
680 }
681 if request.contains_internal_system_write || deferred_cmd.is_some() {
682 break;
683 }
684 }
685 }
686
687 let GroupCommitRequest {
688 appends,
689 responses,
690 statement_logging_ids,
691 notifies,
692 internal_results,
693 write_locks,
694 permits,
695 contains_internal_system_write: _,
696 span: _,
697 } = request;
698
699 let append_metric = self.metrics.append_table_duration_seconds.clone();
700 let Some(write_ts) = self
701 .write_to_txns(Some(&append_metric), |ts, advance_to| {
702 self.table_write_handle
703 .append(ts, advance_to, appends.clone())
704 })
705 .await
706 else {
707 return ControlFlow::Break(());
709 };
710 let timestamp = write_ts.timestamp;
711
712 let modified_tables: Vec<_> = appends
713 .iter()
714 .filter_map(|(id, updates)| {
715 (id.is_user() && !updates.iter().all(|u| u.is_empty())).then_some(id)
716 })
717 .collect();
718 if !modified_tables.is_empty() {
719 info!(
720 "Appending to tables, {modified_tables:?}, at {timestamp}, advancing to {}",
721 write_ts.advance_to
722 );
723 }
724
725 drop(permits);
728 drop(write_locks);
729
730 for notify in notifies {
731 let _ = notify.send(());
732 }
733
734 if self
737 .internal_cmd_tx
738 .send(Message::GroupCommitApplied {
739 responses,
740 statement_logging_ids,
741 internal_results,
742 write_ts: timestamp,
743 })
744 .is_err()
745 {
746 warn!("coordinator shut down before a group commit could be finalized");
747 }
748
749 ControlFlow::Continue(deferred_cmd)
750 }
751}
752
753pub(crate) fn spawn_group_committer(
754 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
755 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
756 table_write_handle: Arc<dyn TableWriteHandle>,
757 catalog_upper: CatalogUpperHandle,
758 internal_cmd_tx: mpsc::UnboundedSender<Message>,
759 now: NowFn,
760 metrics: Metrics,
761 dyncfgs: &ConfigSet,
762) {
763 let committer = GroupCommitter {
764 rx,
765 oracle,
766 table_write_handle,
767 catalog_upper,
768 internal_cmd_tx,
769 now,
770 metrics,
771 max_attempts: GROUP_COMMIT_MAX_ATTEMPTS.handle(dyncfgs),
772 };
773 task::spawn(|| "group_committer", committer.run());
774}
775
776impl Coordinator {
777 pub(crate) fn trigger_group_commit(&mut self) {
779 self.group_commit_tx.notify();
780 self.advance_timelines_interval.reset();
784 }
785
786 pub(crate) async fn try_deferred(
791 &mut self,
792 conn_id: ConnectionId,
793 acquired_lock: Option<(CatalogItemId, tokio::sync::OwnedMutexGuard<()>)>,
794 ) {
795 let Some(op) = self.deferred_write_ops.remove(&conn_id) else {
797 tracing::warn!(%conn_id, "no deferred op found, it must have been canceled?");
798 return;
799 };
800 tracing::info!(%conn_id, "trying deferred plan");
801
802 let write_locks = match acquired_lock {
804 Some((acquired_gid, acquired_lock)) => {
805 let mut write_locks = WriteLocks::builder(op.required_locks());
806
807 write_locks.insert_lock(acquired_gid, acquired_lock);
809
810 for gid in op.required_locks().filter(|gid| *gid != acquired_gid) {
812 if let Some(lock) = self.try_grant_object_write_lock(gid) {
813 write_locks.insert_lock(gid, lock);
814 }
815 }
816
817 let locks = match write_locks.all_or_nothing(op.conn_id()) {
819 Ok(locks) => locks,
820 Err(failed_to_acquire) => {
821 let acquire_future = self
822 .grant_object_write_lock(failed_to_acquire)
823 .map(Option::Some);
824 self.defer_op(acquire_future, op);
825 return;
826 }
827 };
828
829 Some(locks)
830 }
831 None => None,
832 };
833
834 match op {
835 DeferredOp::Plan(mut deferred) => {
836 if let Err(e) = deferred.validity.check(self.catalog()) {
837 deferred.ctx.retire(Err(e))
838 } else {
839 if let Some(locks) = write_locks {
841 let conn_id = deferred.ctx.session().conn_id().clone();
842 if let Err(existing) =
843 deferred.ctx.session_mut().try_grant_write_locks(locks)
844 {
845 tracing::error!(
846 %conn_id,
847 ?existing,
848 "session already write locks granted?",
849 );
850 return deferred.ctx.retire(Err(AdapterError::WrongSetOfLocks));
851 }
852 };
853
854 self.sequence_plan(
856 deferred.ctx,
857 deferred.plan,
858 deferred.resolved_ids,
859 deferred.sql_impl_resolved_ids,
860 )
861 .await;
862 }
863 }
864 DeferredOp::Write(DeferredWrite {
865 span,
866 writes,
867 pending_txn,
868 }) => {
869 self.submit_write(PendingWriteTxn::User {
870 span,
871 writes,
872 write_locks,
873 responder: UserWriteResponder::Session(pending_txn),
874 });
875 }
876 }
877 }
878
879 #[instrument(name = "coord::stage_group_commit")]
883 pub(crate) fn stage_group_commit(&mut self, permit: Option<GroupCommitPermit>) {
884 let mut validated_writes = Vec::new();
885 let mut deferred_writes = Vec::new();
886 let mut group_write_locks = GroupCommitWriteLocks::default();
887
888 let pending_writes: Vec<_> = self.pending_writes.drain(..).collect();
891
892 for pending_write in pending_writes {
894 match pending_write {
895 PendingWriteTxn::System { .. } => validated_writes.push(pending_write),
896 PendingWriteTxn::User {
897 span,
898 write_locks: Some(write_locks),
899 writes,
900 responder,
901 } => match write_locks.validate(writes.keys().copied()) {
902 Ok(validated_locks) => {
903 group_write_locks.merge(validated_locks);
907 validated_writes.push(PendingWriteTxn::User {
908 span,
909 writes,
910 write_locks: None,
911 responder,
912 });
913 }
914 Err(missing) => {
918 let writes: Vec<_> = writes.keys().collect();
919 panic!(
920 "got to group commit with partial set of locks!\nmissing: {:?}, writes: {:?}, conn_id: {}",
921 missing,
922 writes,
923 responder.conn_id(),
924 );
925 }
926 },
927 PendingWriteTxn::User {
930 span,
931 writes,
932 write_locks: None,
933 responder,
934 } => {
935 let missing = group_write_locks.missing_locks(writes.keys().copied());
936 if missing.is_empty() {
937 validated_writes.push(PendingWriteTxn::User {
938 span,
939 writes,
940 write_locks: None,
941 responder,
942 });
943 continue;
944 }
945
946 match responder {
947 UserWriteResponder::Session(pending_txn) => {
948 let mut just_in_time_locks = WriteLocks::builder(missing.clone());
949 for collection in missing {
950 if let Some(lock) = self.try_grant_object_write_lock(collection) {
951 just_in_time_locks.insert_lock(collection, lock);
952 }
953 }
954 match just_in_time_locks
955 .all_or_nothing(pending_txn.ctx.session().conn_id())
956 {
957 Ok(locks) => {
958 group_write_locks.merge(locks);
959 validated_writes.push(PendingWriteTxn::User {
960 span,
961 writes,
962 write_locks: None,
963 responder: UserWriteResponder::Session(pending_txn),
964 });
965 }
966 Err(missing) => {
967 let acquire_future =
968 self.grant_object_write_lock(missing).map(Option::Some);
969 deferred_writes.push((
970 acquire_future,
971 DeferredWrite {
972 span,
973 writes,
974 pending_txn,
975 },
976 ));
977 }
978 }
979 }
980 UserWriteResponder::Internal {
981 conn_id,
982 target,
983 result,
984 } => {
985 let acquired = missing
992 .into_iter()
993 .map(|id| {
994 self.try_grant_object_write_lock(id).map(|lock| (id, lock))
995 })
996 .collect::<Option<Vec<_>>>();
997 if let Some(acquired) = acquired {
998 for (id, lock) in acquired {
999 group_write_locks.insert_lock(id, lock);
1000 }
1001 validated_writes.push(PendingWriteTxn::User {
1002 span,
1003 writes,
1004 write_locks: None,
1005 responder: UserWriteResponder::Internal {
1006 conn_id,
1007 target,
1008 result,
1009 },
1010 });
1011 } else {
1012 self.pending_writes.push(PendingWriteTxn::User {
1031 span,
1032 writes,
1033 write_locks: None,
1034 responder: UserWriteResponder::Internal {
1035 conn_id,
1036 target,
1037 result,
1038 },
1039 });
1040 }
1041 }
1042 }
1043 }
1044 }
1045 }
1046
1047 for (acquire_future, write) in deferred_writes {
1049 self.defer_op(acquire_future, DeferredOp::Write(write));
1050 }
1051
1052 let contains_internal_system_write = validated_writes
1053 .iter()
1054 .any(|write| write.is_internal_system());
1055
1056 let mut appends: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>> = BTreeMap::new();
1057 let mut responses = Vec::with_capacity(validated_writes.len());
1058 let mut statement_logging_ids = Vec::new();
1059 let mut notifies = Vec::new();
1060 let mut internal_results = Vec::new();
1061
1062 for validated_write_txn in validated_writes {
1063 match validated_write_txn {
1064 PendingWriteTxn::User {
1065 span: _,
1066 writes,
1067 write_locks,
1068 responder:
1069 UserWriteResponder::Session(PendingTxn {
1070 ctx,
1071 response,
1072 action,
1073 }),
1074 } => {
1075 assert_none!(write_locks, "should have merged together all locks above");
1076
1077 if let Some(id) = Self::stale_write_target(self.catalog(), &writes) {
1084 let err = AdapterError::ConcurrentDependencyMutation {
1085 dependency_id: id.to_string(),
1086 };
1087 let (ctx, result) = CompletedClientTransmitter::new(
1088 ctx,
1089 Err(err),
1090 EndTransactionAction::Rollback,
1091 )
1092 .finalize();
1093 ctx.retire(result);
1094 continue;
1095 }
1096
1097 for (id, table_data) in writes {
1098 if self.catalog().try_get_entry(&id).is_some() {
1104 appends.entry(id).or_default().extend(table_data);
1105 }
1106 }
1107 if let Some(id) = ctx.extra().contents() {
1108 statement_logging_ids.push(id);
1109 }
1110
1111 responses.push(CompletedClientTransmitter::new(ctx, response, action));
1112 }
1113 PendingWriteTxn::User {
1114 span: _,
1115 writes,
1116 write_locks,
1117 responder: UserWriteResponder::Internal { target, result, .. },
1118 } => {
1119 assert_none!(write_locks, "should have merged together all locks above");
1120 let current_global_id = self
1121 .catalog()
1122 .try_get_entry(&target.item_id)
1123 .map(|entry| entry.latest_global_id());
1124 if current_global_id != Some(target.global_id) {
1125 result.send(WriteResult::TargetChanged);
1126 continue;
1127 }
1128 assert!(
1134 writes.keys().all(|id| *id == target.item_id),
1135 "frontend write for {:?} carries other tables: {:?}",
1136 target.item_id,
1137 writes.keys().collect::<Vec<_>>(),
1138 );
1139 appends
1140 .entry(target.item_id)
1141 .or_default()
1142 .extend(writes.into_values().flatten());
1143 internal_results.push(result);
1144 }
1145 PendingWriteTxn::System { updates, source } => {
1146 for update in updates {
1147 appends.entry(update.id).or_default().push(update.data);
1148 }
1149 match source {
1151 BuiltinTableUpdateSource::Internal(tx)
1152 | BuiltinTableUpdateSource::Background(tx) => notifies.push(tx),
1153 }
1154 }
1155 }
1156 }
1157
1158 let mut all_appends = Vec::with_capacity(appends.len());
1161 for (item_id, table_data) in appends.into_iter() {
1162 let mut all_rows = Vec::new();
1163 let mut all_data = Vec::new();
1164 for data in table_data {
1165 match data {
1166 TableData::Rows(rows) => all_rows.extend(rows),
1167 TableData::Batches(_) => all_data.push(data),
1168 }
1169 }
1170 differential_dataflow::consolidation::consolidate(&mut all_rows);
1171 all_data.push(TableData::Rows(all_rows));
1172
1173 all_appends.push((item_id, all_data));
1175 }
1176
1177 let appends: Vec<_> = all_appends
1178 .into_iter()
1179 .map(|(id, updates)| {
1180 let gid = self.catalog().get_entry(&id).latest_global_id();
1181 (gid, updates)
1182 })
1183 .collect();
1184
1185 let request = GroupCommitRequest {
1187 appends,
1188 responses,
1189 statement_logging_ids,
1190 notifies,
1191 internal_results,
1192 write_locks: group_write_locks,
1193 permits: permit.into_iter().collect(),
1194 contains_internal_system_write,
1195 span: Span::current(),
1196 };
1197 if self
1198 .group_committer_tx
1199 .send(TableWriteCmd::GroupCommit(request))
1200 .is_err()
1201 {
1202 warn!("group committer task gone, dropping staged group commit");
1204 }
1205 }
1206
1207 fn stale_write_target(
1217 catalog: &Catalog,
1218 writes: &BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
1219 ) -> Option<CatalogItemId> {
1220 writes.iter().find_map(|(id, table_data)| {
1221 let entry = catalog.try_get_entry(id)?;
1223 let arity = entry
1224 .relation_desc_latest()
1225 .expect("write target is a table")
1226 .arity();
1227 let stale = table_data.iter().any(|data| match data {
1228 TableData::Rows(rows) => rows
1229 .first()
1230 .is_some_and(|(row, _)| row.iter().count() != arity),
1231 TableData::Batches(_) => false,
1232 });
1233 stale.then_some(*id)
1234 })
1235 }
1236
1237 pub(crate) async fn register_tables_via_committer(
1239 &self,
1240 tables: Vec<TableRegistration>,
1241 ) -> Timestamp {
1242 let (tx, rx) = oneshot::channel();
1243 if self
1244 .group_committer_tx
1245 .send(TableWriteCmd::Register { tables, result: tx })
1246 .is_err()
1247 {
1248 halt!("group committer terminated before a table registration could be submitted");
1249 }
1250 match rx.await {
1251 Ok(ts) => ts,
1252 Err(_) => halt!("group committer terminated with a table registration outstanding"),
1253 }
1254 }
1255
1256 pub(crate) async fn forget_tables_via_committer(&self, ids: Vec<GlobalId>) -> Timestamp {
1258 let (tx, rx) = oneshot::channel();
1259 if self
1260 .group_committer_tx
1261 .send(TableWriteCmd::Forget { ids, result: tx })
1262 .is_err()
1263 {
1264 halt!("group committer terminated before a table forget could be submitted");
1265 }
1266 match rx.await {
1267 Ok(ts) => ts,
1268 Err(_) => halt!("group committer terminated with a table forget outstanding"),
1269 }
1270 }
1271
1272 pub(crate) fn submit_write(&mut self, pending_write_txn: PendingWriteTxn) {
1274 if self.controller.read_only() {
1275 panic!(
1276 "attempting table write in read-only mode: {:?}",
1277 pending_write_txn
1278 );
1279 }
1280 self.pending_writes.push(pending_write_txn);
1281 self.trigger_group_commit();
1282 }
1283
1284 pub(crate) fn builtin_table_update<'a>(&'a mut self) -> BuiltinTableAppend<'a> {
1286 BuiltinTableAppend { coord: self }
1287 }
1288
1289 pub(crate) fn defer_op<F>(&mut self, acquire_future: F, op: DeferredOp)
1290 where
1291 F: Future<Output = Option<(CatalogItemId, tokio::sync::OwnedMutexGuard<()>)>>
1292 + Send
1293 + 'static,
1294 {
1295 let conn_id = op.conn_id().clone();
1296
1297 let is_optimistic = op.can_be_optimistically_retried();
1299 self.deferred_write_ops.insert(conn_id.clone(), op);
1300
1301 let internal_cmd_tx = self.internal_cmd_tx.clone();
1302 let conn_id_ = conn_id.clone();
1303 mz_ore::task::spawn(|| format!("defer op {conn_id_}"), async move {
1304 tracing::info!(%conn_id, "deferring plan");
1305 let acquired_lock = acquire_future.await;
1310
1311 let acquired_lock = match (acquired_lock, is_optimistic) {
1315 (Some(_lock), true) => None,
1316 (Some(lock), false) => Some(lock),
1317 (None, _) => None,
1318 };
1319
1320 let _ = internal_cmd_tx.send(Message::TryDeferred {
1322 conn_id,
1323 acquired_lock,
1324 });
1325 });
1326 }
1327
1328 pub(crate) fn grant_object_write_lock(
1330 &mut self,
1331 object_id: CatalogItemId,
1332 ) -> impl Future<Output = (CatalogItemId, OwnedMutexGuard<()>)> + 'static {
1333 let write_lock_handle = self
1334 .write_locks
1335 .entry(object_id)
1336 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())));
1337 let write_lock_handle = Arc::clone(write_lock_handle);
1338
1339 write_lock_handle
1340 .lock_owned()
1341 .map(move |guard| (object_id, guard))
1342 }
1343
1344 pub(crate) fn try_grant_object_write_lock(
1347 &mut self,
1348 object_id: CatalogItemId,
1349 ) -> Option<OwnedMutexGuard<()>> {
1350 let write_lock_handle = self
1351 .write_locks
1352 .entry(object_id)
1353 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())));
1354 let write_lock_handle = Arc::clone(write_lock_handle);
1355
1356 write_lock_handle.try_lock_owned().ok()
1357 }
1358}
1359
1360pub struct BuiltinTableAppend<'a> {
1362 coord: &'a mut Coordinator,
1363}
1364
1365pub type BuiltinTableAppendNotify = Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>;
1376
1377pub struct BuiltinTableAppendCompletion {
1379 notify: BuiltinTableAppendNotify,
1380}
1381
1382impl BuiltinTableAppendCompletion {
1383 pub fn new(notify: BuiltinTableAppendNotify) -> Self {
1384 Self { notify }
1385 }
1386
1387 pub fn into_notify(self) -> BuiltinTableAppendNotify {
1388 self.notify
1389 }
1390}
1391
1392impl<'a> BuiltinTableAppend<'a> {
1393 pub fn background(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1404 if self.coord.controller.read_only() {
1405 self.coord
1406 .buffered_builtin_table_updates
1407 .as_mut()
1408 .expect("in read-only mode")
1409 .append(&mut updates);
1410
1411 return Box::pin(futures::future::ready(()));
1412 }
1413
1414 let (tx, rx) = oneshot::channel();
1415 self.coord.pending_writes.push(PendingWriteTxn::System {
1416 updates,
1417 source: BuiltinTableUpdateSource::Background(tx),
1418 });
1419
1420 Box::pin(rx.map(|_| ()))
1421 }
1422
1423 pub fn defer(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1432 if self.coord.controller.read_only() {
1433 self.coord
1434 .buffered_builtin_table_updates
1435 .as_mut()
1436 .expect("in read-only mode")
1437 .append(&mut updates);
1438
1439 return Box::pin(futures::future::ready(()));
1440 }
1441
1442 let (tx, rx) = oneshot::channel();
1443 self.coord.pending_writes.push(PendingWriteTxn::System {
1444 updates,
1445 source: BuiltinTableUpdateSource::Internal(tx),
1446 });
1447 self.coord.trigger_group_commit();
1448
1449 Box::pin(rx.map(|_| ()))
1450 }
1451
1452 pub fn execute(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1456 if self.coord.controller.read_only() {
1457 self.coord
1458 .buffered_builtin_table_updates
1459 .as_mut()
1460 .expect("in read-only mode")
1461 .append(&mut updates);
1462
1463 return Box::pin(futures::future::ready(()));
1464 }
1465
1466 let (tx, rx) = oneshot::channel();
1467
1468 self.coord.pending_writes.push(PendingWriteTxn::System {
1471 updates,
1472 source: BuiltinTableUpdateSource::Internal(tx),
1473 });
1474 self.coord.stage_group_commit(None);
1475
1476 self.coord.advance_timelines_interval.reset();
1478
1479 Box::pin(rx.map(|_| ()))
1480 }
1481}
1482
1483pub fn notifier() -> (GroupCommitNotifier, GroupCommitWaiter) {
1486 let notify = Arc::new(Notify::new());
1487 let in_progress = Arc::new(Semaphore::new(1));
1488
1489 let notifier = GroupCommitNotifier {
1490 notify: Arc::clone(¬ify),
1491 };
1492 let waiter = GroupCommitWaiter {
1493 notify,
1494 in_progress,
1495 };
1496
1497 (notifier, waiter)
1498}
1499
1500#[derive(Debug, Clone)]
1503pub struct GroupCommitNotifier {
1504 notify: Arc<Notify>,
1506}
1507
1508impl GroupCommitNotifier {
1509 pub fn notify(&self) {
1511 self.notify.notify_one()
1512 }
1513}
1514
1515#[derive(Debug)]
1518pub struct GroupCommitWaiter {
1519 notify: Arc<Notify>,
1521 in_progress: Arc<Semaphore>,
1523}
1524static_assertions::assert_not_impl_all!(GroupCommitWaiter: Clone);
1525
1526impl GroupCommitWaiter {
1527 pub async fn ready(&self) -> GroupCommitPermit {
1538 let permit = Semaphore::acquire_owned(Arc::clone(&self.in_progress))
1539 .await
1540 .expect("semaphore should not close");
1541
1542 self.notify.notified().await;
1545
1546 GroupCommitPermit {
1547 _permit: Some(permit),
1548 }
1549 }
1550}
1551
1552#[derive(Debug)]
1557pub struct GroupCommitPermit {
1558 _permit: Option<OwnedSemaphorePermit>,
1562}
1563
1564pub(crate) fn waiting_on_startup_appends(
1580 catalog: &Catalog,
1581 session: &mut Session,
1582 plan: &Plan,
1583) -> Option<(BTreeSet<CatalogItemId>, BoxFuture<'static, ()>)> {
1584 let depends_on = match plan {
1588 Plan::Select(plan) => plan.source.depends_on(),
1589 Plan::ReadThenWrite(plan) => plan.selection.depends_on(),
1590 Plan::ShowColumns(plan) => plan.select_plan.source.depends_on(),
1591 Plan::Subscribe(plan) => plan.from.depends_on(),
1592 Plan::ExplainPlan(ExplainPlanPlan {
1593 explainee: Explainee::Statement(ExplaineeStatement::Select { plan, .. }),
1594 ..
1595 }) => plan.source.depends_on(),
1596 Plan::ExplainTimestamp(ExplainTimestampPlan { raw_plan, .. }) => raw_plan.depends_on(),
1597 Plan::CreateConnection(_)
1598 | Plan::CreateDatabase(_)
1599 | Plan::CreateSchema(_)
1600 | Plan::CreateRole(_)
1601 | Plan::CreateNetworkPolicy(_)
1602 | Plan::CreateCluster(_)
1603 | Plan::CreateClusterReplica(_)
1604 | Plan::CreateSource(_)
1605 | Plan::CreateSources(_)
1606 | Plan::CreateSecret(_)
1607 | Plan::CreateSink(_)
1608 | Plan::CreateTable(_)
1609 | Plan::CreateView(_)
1610 | Plan::CreateMaterializedView(_)
1611 | Plan::CreateIndex(_)
1612 | Plan::CreateMetricSink(_)
1613 | Plan::CreateType(_)
1614 | Plan::Comment(_)
1615 | Plan::DiscardTemp
1616 | Plan::DiscardAll
1617 | Plan::DropObjects(_)
1618 | Plan::DropOwned(_)
1619 | Plan::EmptyQuery
1620 | Plan::ShowAllVariables
1621 | Plan::ShowCreate(_)
1622 | Plan::ShowVariable(_)
1623 | Plan::InspectShard(_)
1624 | Plan::SetVariable(_)
1625 | Plan::ResetVariable(_)
1626 | Plan::SetTransaction(_)
1627 | Plan::StartTransaction(_)
1628 | Plan::CommitTransaction(_)
1629 | Plan::AbortTransaction(_)
1630 | Plan::CopyFrom(_)
1631 | Plan::CopyTo(_)
1632 | Plan::ExplainPlan(_)
1633 | Plan::ExplainPushdown(_)
1634 | Plan::ExplainSinkSchema(_)
1635 | Plan::Insert(_)
1636 | Plan::AlterNetworkPolicy(_)
1637 | Plan::AlterNoop(_)
1638 | Plan::AlterClusterRename(_)
1639 | Plan::AlterClusterSwap(_)
1640 | Plan::AlterClusterReplicaRename(_)
1641 | Plan::AlterCluster(_)
1642 | Plan::AlterConnection(_)
1643 | Plan::AlterSource(_)
1644 | Plan::AlterSetCluster(_)
1645 | Plan::AlterItemRename(_)
1646 | Plan::AlterRetainHistory(_)
1647 | Plan::AlterSourceTimestampInterval(_)
1648 | Plan::AlterSchemaRename(_)
1649 | Plan::AlterSchemaSwap(_)
1650 | Plan::AlterSecret(_)
1651 | Plan::AlterSink(_)
1652 | Plan::AlterSystemSet(_)
1653 | Plan::AlterSystemReset(_)
1654 | Plan::AlterSystemResetAll(_)
1655 | Plan::AlterRole(_)
1656 | Plan::AlterOwner(_)
1657 | Plan::AlterTableAddColumn(_)
1658 | Plan::AlterMaterializedViewApplyReplacement(_)
1659 | Plan::Declare(_)
1660 | Plan::Fetch(_)
1661 | Plan::Close(_)
1662 | Plan::Prepare(_)
1663 | Plan::Execute(_)
1664 | Plan::Deallocate(_)
1665 | Plan::Raise(_)
1666 | Plan::GrantRole(_)
1667 | Plan::RevokeRole(_)
1668 | Plan::GrantPrivileges(_)
1669 | Plan::RevokePrivileges(_)
1670 | Plan::AlterDefaultPrivileges(_)
1671 | Plan::ReassignOwned(_)
1672 | Plan::ValidateConnection(_)
1673 | Plan::SideEffectingFunc(_) => BTreeSet::default(),
1674 };
1675 let depends_on_required_id = REQUIRED_BUILTIN_TABLES
1676 .iter()
1677 .map(|table| catalog.resolve_builtin_table(&**table))
1678 .any(|id| {
1679 catalog
1680 .get_global_ids(&id)
1681 .any(|gid| depends_on.contains(&gid))
1682 });
1683
1684 if !depends_on_required_id {
1687 return None;
1688 }
1689
1690 match session.clear_builtin_table_updates() {
1697 Some(wait_future) => {
1698 let depends_on = depends_on
1699 .into_iter()
1700 .map(|gid| catalog.get_entry_by_global_id(&gid).id())
1701 .collect();
1702 Some((depends_on, wait_future.boxed()))
1703 }
1704 None => None,
1705 }
1706}
1707
1708#[cfg(test)]
1709mod tests {
1710 use std::sync::Mutex;
1711 use std::sync::atomic::{AtomicUsize, Ordering};
1712
1713 use async_trait::async_trait;
1714 use mz_ore::metrics::MetricsRegistry;
1715 use mz_ore::now::SYSTEM_TIME;
1716 use timely::progress::Antichain;
1717
1718 use super::*;
1719 use crate::catalog::Catalog;
1720
1721 #[mz_ore::test(tokio::test)]
1722 async fn internal_write_responder_reports_indeterminate_on_drop() {
1723 let (tx, rx) = oneshot::channel();
1724 drop(InternalWriteResponder::new(tx));
1725 assert!(matches!(rx.await, Ok(WriteResult::Indeterminate)));
1726 }
1727
1728 #[derive(Debug, Default)]
1729 struct MemTimestampOracle {
1730 read_write_ts: Mutex<(Timestamp, Timestamp)>,
1731 apply_writes: AtomicUsize,
1732 }
1733
1734 impl MemTimestampOracle {
1735 fn starting_at(ts: Timestamp) -> Self {
1736 Self {
1737 read_write_ts: Mutex::new((ts, ts)),
1738 apply_writes: AtomicUsize::new(0),
1739 }
1740 }
1741 }
1742
1743 #[async_trait]
1744 impl TimestampOracle<Timestamp> for MemTimestampOracle {
1745 async fn write_ts(&self) -> WriteTimestamp {
1746 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1747 let new_write_ts = std::cmp::max(*read_ts, *write_ts).step_forward();
1748 *write_ts = new_write_ts;
1749 WriteTimestamp {
1750 timestamp: new_write_ts,
1751 advance_to: new_write_ts.step_forward(),
1752 }
1753 }
1754
1755 async fn peek_write_ts(&self) -> Timestamp {
1756 let (_, write_ts) = &*self.read_write_ts.lock().expect("lock poisoned");
1757 *write_ts
1758 }
1759
1760 async fn read_ts(&self) -> Timestamp {
1761 let (read_ts, _) = &*self.read_write_ts.lock().expect("lock poisoned");
1762 *read_ts
1763 }
1764
1765 async fn apply_write(&self, lower_bound: Timestamp) {
1766 self.apply_writes.fetch_add(1, Ordering::SeqCst);
1767 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1768 *read_ts = std::cmp::max(*read_ts, lower_bound);
1769 *write_ts = std::cmp::max(*read_ts, *write_ts);
1770 }
1771 }
1772
1773 #[derive(Debug)]
1774 struct ConflictingTableWriteHandle {
1775 conflicts: usize,
1776 calls: AtomicUsize,
1777 write_timestamps: Mutex<Vec<Timestamp>>,
1778 }
1779
1780 impl ConflictingTableWriteHandle {
1781 fn new(conflicts: usize) -> Self {
1782 Self {
1783 conflicts,
1784 calls: AtomicUsize::new(0),
1785 write_timestamps: Mutex::new(Vec::new()),
1786 }
1787 }
1788
1789 fn respond(&self, write_ts: Timestamp) -> oneshot::Receiver<Result<(), StorageError>> {
1790 let call = self.calls.fetch_add(1, Ordering::SeqCst);
1791 self.write_timestamps
1792 .lock()
1793 .expect("lock poisoned")
1794 .push(write_ts);
1795 let (tx, rx) = oneshot::channel();
1796 let result = if call < self.conflicts {
1797 Err(StorageError::InvalidUppers(vec![
1798 mz_storage_types::controller::InvalidUpper {
1799 id: GlobalId::User(1),
1800 current_upper: Antichain::from_elem(write_ts.step_forward()),
1801 },
1802 ]))
1803 } else {
1804 Ok(())
1805 };
1806 tx.send(result).expect("receiver still in scope");
1807 rx
1808 }
1809 }
1810
1811 impl TableWriteHandle for ConflictingTableWriteHandle {
1812 fn append(
1813 &self,
1814 write_ts: Timestamp,
1815 _advance_to: Timestamp,
1816 _commands: Vec<(GlobalId, Vec<TableData>)>,
1817 ) -> oneshot::Receiver<Result<(), StorageError>> {
1818 self.respond(write_ts)
1819 }
1820
1821 fn register(
1822 &self,
1823 register_ts: Timestamp,
1824 _tables: Vec<TableRegistration>,
1825 ) -> oneshot::Receiver<Result<(), StorageError>> {
1826 self.respond(register_ts)
1827 }
1828
1829 fn forget(
1830 &self,
1831 forget_ts: Timestamp,
1832 _ids: Vec<GlobalId>,
1833 ) -> oneshot::Receiver<Result<(), StorageError>> {
1834 self.respond(forget_ts)
1835 }
1836 }
1837
1838 #[mz_ore::test(tokio::test)]
1839 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_conflict_retry() {
1841 Catalog::with_debug(|catalog| async move {
1842 let initial_upper = catalog.current_upper().await;
1844 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1845 let handle = Arc::new(ConflictingTableWriteHandle::new(1));
1846 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1847 Arc::<MemTimestampOracle>::clone(&oracle);
1848 let handle_dyn: Arc<dyn TableWriteHandle> =
1849 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1850 let (_tx, rx) = mpsc::unbounded_channel();
1851 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1852 let committer = GroupCommitter {
1853 rx,
1854 oracle: oracle_dyn,
1855 table_write_handle: handle_dyn,
1856 catalog_upper: catalog.upper_handle(),
1857 internal_cmd_tx,
1858 now: SYSTEM_TIME.clone(),
1859 metrics: Metrics::register_into(&MetricsRegistry::new()),
1860 max_attempts: ConfigValHandle::disconnected(2),
1861 };
1862
1863 let write_ts = committer
1864 .write_to_txns(None, |ts, advance_to| {
1865 assert_eq!(advance_to, ts.step_forward());
1866 handle.append(ts, advance_to, Vec::new())
1867 })
1868 .await
1869 .expect("worker stays alive in this test");
1870
1871 let attempts = handle
1872 .write_timestamps
1873 .lock()
1874 .expect("lock poisoned")
1875 .clone();
1876 assert_eq!(attempts.len(), 2, "one conflict, one success");
1877 assert!(
1878 attempts[0] > initial_upper,
1879 "attempts start past the initial catalog upper: {attempts:?} vs {initial_upper}"
1880 );
1881 assert!(
1882 attempts[0] < attempts[1],
1883 "retry must use a fresh, larger timestamp: {attempts:?}"
1884 );
1885 assert_eq!(write_ts.timestamp, attempts[1]);
1886
1887 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1888 assert_eq!(oracle.read_ts().await, write_ts.timestamp);
1889
1890 let catalog_upper = catalog.current_upper().await;
1891 assert!(
1892 catalog_upper >= write_ts.advance_to,
1893 "catalog upper {catalog_upper} must cover the write's advance_to {}",
1894 write_ts.advance_to
1895 );
1896
1897 catalog.expire().await;
1898 })
1899 .await;
1900 }
1901
1902 #[mz_ore::test(tokio::test)]
1903 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_zero_limit_allows_one_attempt() {
1905 Catalog::with_debug(|catalog| async move {
1906 let initial_upper = catalog.current_upper().await;
1907 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1908 let handle = Arc::new(ConflictingTableWriteHandle::new(0));
1909 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1910 Arc::<MemTimestampOracle>::clone(&oracle);
1911 let handle_dyn: Arc<dyn TableWriteHandle> =
1912 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1913 let (_tx, rx) = mpsc::unbounded_channel();
1914 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1915 let committer = GroupCommitter {
1916 rx,
1917 oracle: oracle_dyn,
1918 table_write_handle: handle_dyn,
1919 catalog_upper: catalog.upper_handle(),
1920 internal_cmd_tx,
1921 now: SYSTEM_TIME.clone(),
1922 metrics: Metrics::register_into(&MetricsRegistry::new()),
1923 max_attempts: ConfigValHandle::disconnected(0),
1924 };
1925
1926 let write_ts = committer
1927 .write_to_txns(None, |ts, advance_to| {
1928 handle.append(ts, advance_to, Vec::new())
1929 })
1930 .await
1931 .expect("zero limit permits one successful attempt");
1932
1933 let attempts = handle
1934 .write_timestamps
1935 .lock()
1936 .expect("lock poisoned")
1937 .clone();
1938 assert_eq!(attempts, vec![write_ts.timestamp]);
1939 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1940
1941 catalog.expire().await;
1942 })
1943 .await;
1944 }
1945}