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::halt;
50use mz_ore::instrument;
51use mz_ore::now::NowFn;
52use mz_ore::task;
53use mz_repr::{CatalogItemId, GlobalId, Timestamp};
54use mz_sql::names::ResolvedIds;
55use mz_sql::plan::{ExplainPlanPlan, ExplainTimestampPlan, Explainee, ExplaineeStatement, Plan};
56use mz_sql::session::metadata::SessionMetadata;
57use mz_storage_client::client::TableData;
58use mz_storage_client::controller::{TableRegistration, TableWriteHandle};
59use mz_storage_types::controller::StorageError;
60use mz_timestamp_oracle::{TimestampOracle, WriteTimestamp};
61use smallvec::SmallVec;
62use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
63use tracing::{Instrument, Span, info, warn};
64
65use crate::catalog::{BuiltinTableUpdate, Catalog, CatalogUpperHandle};
66use crate::coord::timeline::write_ts_upper_bound;
67use crate::coord::{Coordinator, Message, PendingTxn, PlanValidity};
68use crate::metrics::Metrics;
69use crate::session::{EndTransactionAction, Session};
70use crate::statement_logging::StatementLoggingId;
71use crate::util::{CompletedClientTransmitter, ResultExt};
72use crate::{AdapterError, ExecuteContext};
73
74pub(crate) static REQUIRED_BUILTIN_TABLES: &[&LazyLock<BuiltinTable>] = &[&MZ_SESSIONS];
76
77#[derive(Derivative)]
79#[derivative(Debug)]
80pub struct DeferredPlan {
81 #[derivative(Debug = "ignore")]
82 pub ctx: ExecuteContext,
83 pub plan: Plan,
84 pub validity: PlanValidity,
85 pub resolved_ids: ResolvedIds,
86 pub sql_impl_resolved_ids: ResolvedIds,
87}
88
89#[derive(Debug)]
91pub(crate) enum BuiltinTableUpdateSource {
92 Internal(oneshot::Sender<()>),
94 Background(oneshot::Sender<()>),
96}
97
98#[derive(Debug, Clone)]
100pub enum WriteResult {
101 Success { timestamp: Timestamp },
103 TimestampPassed {
105 target_timestamp: Timestamp,
106 next_eligible_timestamp: Timestamp,
107 },
108 TimestampTooFarAhead {
111 target_timestamp: Timestamp,
112 limit: Timestamp,
113 },
114 Canceled,
116 ReadOnly,
118 TargetChanged,
120 Indeterminate,
122}
123
124#[derive(Debug)]
131pub struct InternalWriteResponder {
132 tx: Option<oneshot::Sender<WriteResult>>,
133}
134
135impl InternalWriteResponder {
136 pub(crate) fn new(tx: oneshot::Sender<WriteResult>) -> Self {
137 Self { tx: Some(tx) }
138 }
139
140 pub(crate) fn send(mut self, result: WriteResult) {
141 if let Some(tx) = self.tx.take() {
142 let _ = tx.send(result);
143 }
144 }
145}
146
147impl Drop for InternalWriteResponder {
148 fn drop(&mut self) {
149 if let Some(tx) = self.tx.take() {
150 let _ = tx.send(WriteResult::Indeterminate);
151 }
152 }
153}
154
155#[derive(Debug)]
157pub(crate) enum UserWriteResponder {
158 Session(PendingTxn),
161 Internal {
163 conn_id: ConnectionId,
164 target: WriteTarget,
168 result: InternalWriteResponder,
169 },
170}
171
172#[derive(Debug, Clone, Copy)]
174pub(crate) struct WriteTarget {
175 pub(crate) item_id: CatalogItemId,
176 pub(crate) global_id: GlobalId,
177}
178
179#[derive(Debug)]
181pub(crate) enum PendingWriteTxn {
182 User {
185 span: Span,
186 writes: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
188 responder: UserWriteResponder,
190 },
191 System {
193 updates: Vec<BuiltinTableUpdate>,
194 source: BuiltinTableUpdateSource,
195 },
196}
197
198impl PendingWriteTxn {
199 fn is_internal_system(&self) -> bool {
200 match self {
201 PendingWriteTxn::System {
202 source: BuiltinTableUpdateSource::Internal(_),
203 ..
204 } => true,
205 _ => false,
206 }
207 }
208}
209
210pub(crate) enum TableWriteCmd {
211 GroupCommit(GroupCommitRequest),
212 TimestampedWrite(TimestampedWriteRequest),
213 Register {
214 tables: Vec<TableRegistration>,
215 result: oneshot::Sender<Timestamp>,
216 },
217 Forget {
218 ids: Vec<GlobalId>,
219 result: oneshot::Sender<Timestamp>,
220 },
221}
222
223pub(crate) struct TimestampedWriteRequest {
232 pub(crate) appends: Vec<(GlobalId, Vec<TableData>)>,
233 pub(crate) target_timestamp: Timestamp,
234 pub(crate) result: InternalWriteResponder,
235 pub(crate) span: Span,
236}
237
238pub(crate) struct GroupCommitRequest {
240 appends: Vec<(GlobalId, Vec<TableData>)>,
242 responses: Vec<CompletedClientTransmitter>,
243 statement_logging_ids: Vec<StatementLoggingId>,
244 notifies: Vec<oneshot::Sender<()>>,
245 internal_results: Vec<InternalWriteResponder>,
246 permits: Vec<GroupCommitPermit>,
248 contains_internal_system_write: bool,
249 span: Span,
250}
251
252impl GroupCommitRequest {
253 fn merge(&mut self, other: GroupCommitRequest) {
254 let GroupCommitRequest {
255 appends,
256 responses,
257 statement_logging_ids,
258 notifies,
259 internal_results,
260 permits,
261 contains_internal_system_write,
262 span: _,
263 } = other;
264 self.appends.extend(appends);
265 self.responses.extend(responses);
266 self.statement_logging_ids.extend(statement_logging_ids);
267 self.notifies.extend(notifies);
268 self.internal_results.extend(internal_results);
269 self.permits.extend(permits);
270 self.contains_internal_system_write |= contains_internal_system_write;
271 }
272}
273
274pub(crate) struct GroupCommitter {
279 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
280 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
281 table_write_handle: Arc<dyn TableWriteHandle>,
282 catalog_upper: CatalogUpperHandle,
283 internal_cmd_tx: mpsc::UnboundedSender<Message>,
284 now: NowFn,
285 metrics: Metrics,
286 max_attempts: ConfigValHandle<usize>,
287}
288
289enum TxnsWriteAttempt {
291 Applied,
293 UpperConflict,
296 WorkerGone,
298}
299
300impl GroupCommitter {
301 async fn run(mut self) {
302 while let Some(cmd) = self.rx.recv().await {
303 let mut next = Some(cmd);
306 while let Some(cmd) = next.take() {
307 match cmd {
308 TableWriteCmd::GroupCommit(request) => {
309 let span = request.span.clone();
310 match self.commit(request).instrument(span).await {
311 ControlFlow::Continue(deferred) => next = deferred,
312 ControlFlow::Break(()) => return,
313 }
314 }
315 TableWriteCmd::TimestampedWrite(request) => {
316 let span = request.span.clone();
317 if self
318 .commit_timestamped(request)
319 .instrument(span)
320 .await
321 .is_break()
322 {
323 return;
324 }
325 }
326 TableWriteCmd::Register { tables, result } => {
327 let Some(write_ts) = self
328 .write_to_txns(None, |ts, _advance_to| {
329 self.table_write_handle.register(ts, tables.clone())
330 })
331 .await
332 else {
333 return;
334 };
335 let _ = result.send(write_ts.timestamp);
336 }
337 TableWriteCmd::Forget { ids, result } => {
338 let Some(write_ts) = self
339 .write_to_txns(None, |ts, _advance_to| {
340 self.table_write_handle.forget(ts, ids.clone())
341 })
342 .await
343 else {
344 return;
345 };
346 let _ = result.send(write_ts.timestamp);
347 }
348 }
349 }
350 }
351 }
352
353 async fn commit_timestamped(&self, request: TimestampedWriteRequest) -> ControlFlow<(), ()> {
374 let TimestampedWriteRequest {
375 appends,
376 target_timestamp,
377 result,
378 span: _,
379 } = request;
380
381 let oracle_write_ts = self.oracle.peek_write_ts().await;
382 if target_timestamp <= oracle_write_ts {
383 result.send(WriteResult::TimestampPassed {
384 target_timestamp,
385 next_eligible_timestamp: oracle_write_ts.step_forward(),
386 });
387 return ControlFlow::Continue(());
388 }
389
390 let now: Timestamp = (self.now)().into();
393 let limit = write_ts_upper_bound(&now);
394 if target_timestamp > limit {
395 result.send(WriteResult::TimestampTooFarAhead {
396 target_timestamp,
397 limit,
398 });
399 return ControlFlow::Continue(());
400 }
401
402 let write_ts = WriteTimestamp {
403 timestamp: target_timestamp,
404 advance_to: target_timestamp.step_forward(),
405 };
406 match self
407 .attempt_write_to_txns(
408 &write_ts,
409 Some(&self.metrics.append_table_duration_seconds),
410 |ts, advance_to| self.table_write_handle.append(ts, advance_to, appends),
411 )
412 .await
413 {
414 TxnsWriteAttempt::Applied => {}
415 TxnsWriteAttempt::UpperConflict => {
416 result.send(WriteResult::TimestampPassed {
417 target_timestamp,
418 next_eligible_timestamp: write_ts.advance_to,
419 });
420 return ControlFlow::Continue(());
421 }
422 TxnsWriteAttempt::WorkerGone => {
423 warn!("table write worker gone with a timestamped write outstanding");
424 return ControlFlow::Break(());
425 }
426 }
427
428 if self
429 .internal_cmd_tx
430 .send(Message::GroupCommitApplied {
431 responses: Vec::new(),
432 statement_logging_ids: Vec::new(),
433 internal_results: vec![result],
434 write_ts: target_timestamp,
435 })
436 .is_err()
437 {
438 warn!("coordinator shut down before a timestamped write could be finalized");
439 }
440 ControlFlow::Continue(())
441 }
442
443 async fn write_to_txns(
448 &self,
449 op_duration_metric: Option<&prometheus::Histogram>,
450 mut op: impl FnMut(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
451 ) -> Option<WriteTimestamp> {
452 let mut attempt = 0;
454 loop {
455 let max_attempts = self.max_attempts.get().max(1);
456 if attempt >= max_attempts {
457 halt!(
458 "txns-shard write reached attempt limit {max_attempts} after {attempt} conflicts, rebuilding"
459 );
460 }
461 attempt += 1;
462 let write_ts = self.oracle.write_ts().await;
463
464 match self
467 .attempt_write_to_txns(&write_ts, op_duration_metric, |ts, advance_to| {
468 op(ts, advance_to)
469 })
470 .await
471 {
472 TxnsWriteAttempt::Applied => return Some(write_ts),
473 TxnsWriteAttempt::UpperConflict => {
474 warn!(
475 write_ts = %write_ts.timestamp,
476 attempt,
477 "txns-shard write conflicted with another writer, retrying at a fresh timestamp"
478 );
479 continue;
480 }
481 TxnsWriteAttempt::WorkerGone => {
482 warn!("table write worker gone (process shutting down), winding down");
484 return None;
485 }
486 }
487 }
488 }
489
490 async fn attempt_write_to_txns(
496 &self,
497 write_ts: &WriteTimestamp,
498 op_duration_metric: Option<&prometheus::Histogram>,
499 op: impl FnOnce(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
500 ) -> TxnsWriteAttempt {
501 let catalog_upper_start = Instant::now();
502 self.catalog_upper
503 .advance_upper(write_ts.advance_to)
504 .await
505 .unwrap_or_terminate("unable to advance catalog upper");
506 self.metrics
507 .group_commit_catalog_upper_seconds
508 .observe(catalog_upper_start.elapsed().as_secs_f64());
509
510 let op_start = Instant::now();
511 let op_res = op(write_ts.timestamp, write_ts.advance_to).await;
512 if let Some(metric) = op_duration_metric {
513 metric.observe(op_start.elapsed().as_secs_f64());
514 }
515
516 match op_res {
517 Ok(Ok(())) => {}
518 Ok(Err(StorageError::InvalidUppers(_))) => return TxnsWriteAttempt::UpperConflict,
519 Ok(Err(other)) => {
520 Err::<(), _>(other).unwrap_or_terminate("cannot fail to write to txns shard");
521 unreachable!("unwrap_or_terminate does not return on Err");
522 }
523 Err(_recv) => return TxnsWriteAttempt::WorkerGone,
524 }
525
526 let now: Timestamp = (self.now)().into();
527 crate::coord::timeline::check_runaway_write_ts(&now, write_ts.timestamp);
528
529 fail::fail_point!("group_commit_before_apply_write");
538
539 self.oracle.apply_write(write_ts.timestamp).await;
540
541 TxnsWriteAttempt::Applied
542 }
543
544 async fn commit(
549 &mut self,
550 mut request: GroupCommitRequest,
551 ) -> ControlFlow<(), Option<TableWriteCmd>> {
552 let mut deferred_cmd = None;
553 let mut rx_closed = false;
556
557 loop {
561 while deferred_cmd.is_none() {
562 match self.rx.try_recv() {
563 Ok(TableWriteCmd::GroupCommit(other)) => request.merge(other),
564 Ok(other) => deferred_cmd = Some(other),
565 Err(_) => break,
566 }
567 }
568
569 if request.contains_internal_system_write || deferred_cmd.is_some() {
572 break;
573 }
574 let ts = self.oracle.peek_write_ts().await;
575 let now: Timestamp = (self.now)().into();
576 if ts <= now {
577 break;
578 }
579 let remaining_ms = std::cmp::min(ts.saturating_sub(now), Timestamp::from(1_000u64));
582 let sleep = tokio::time::sleep(Duration::from_millis(remaining_ms.into()));
583 tokio::pin!(sleep);
584 loop {
585 tokio::select! {
586 _ = &mut sleep => break,
587 cmd = self.rx.recv(), if deferred_cmd.is_none() && !rx_closed => match cmd {
588 Some(TableWriteCmd::GroupCommit(other)) => request.merge(other),
589 Some(other) => deferred_cmd = Some(other),
590 None => rx_closed = true,
591 },
592 }
593 if request.contains_internal_system_write || deferred_cmd.is_some() {
594 break;
595 }
596 }
597 }
598
599 let GroupCommitRequest {
600 appends,
601 responses,
602 statement_logging_ids,
603 notifies,
604 internal_results,
605 permits,
606 contains_internal_system_write: _,
607 span: _,
608 } = request;
609
610 let append_metric = self.metrics.append_table_duration_seconds.clone();
611 let Some(write_ts) = self
612 .write_to_txns(Some(&append_metric), |ts, advance_to| {
613 self.table_write_handle
614 .append(ts, advance_to, appends.clone())
615 })
616 .await
617 else {
618 return ControlFlow::Break(());
620 };
621 let timestamp = write_ts.timestamp;
622
623 let modified_tables: Vec<_> = appends
624 .iter()
625 .filter_map(|(id, updates)| {
626 (id.is_user() && !updates.iter().all(|u| u.is_empty())).then_some(id)
627 })
628 .collect();
629 if !modified_tables.is_empty() {
630 info!(
631 "Appending to tables, {modified_tables:?}, at {timestamp}, advancing to {}",
632 write_ts.advance_to
633 );
634 }
635
636 drop(permits);
638
639 for notify in notifies {
640 let _ = notify.send(());
641 }
642
643 if self
646 .internal_cmd_tx
647 .send(Message::GroupCommitApplied {
648 responses,
649 statement_logging_ids,
650 internal_results,
651 write_ts: timestamp,
652 })
653 .is_err()
654 {
655 warn!("coordinator shut down before a group commit could be finalized");
656 }
657
658 ControlFlow::Continue(deferred_cmd)
659 }
660}
661
662pub(crate) fn spawn_group_committer(
663 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
664 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
665 table_write_handle: Arc<dyn TableWriteHandle>,
666 catalog_upper: CatalogUpperHandle,
667 internal_cmd_tx: mpsc::UnboundedSender<Message>,
668 now: NowFn,
669 metrics: Metrics,
670 dyncfgs: &ConfigSet,
671) {
672 let committer = GroupCommitter {
673 rx,
674 oracle,
675 table_write_handle,
676 catalog_upper,
677 internal_cmd_tx,
678 now,
679 metrics,
680 max_attempts: GROUP_COMMIT_MAX_ATTEMPTS.handle(dyncfgs),
681 };
682 task::spawn(|| "group_committer", committer.run());
683}
684
685impl Coordinator {
686 pub(crate) fn trigger_group_commit(&mut self) {
688 self.group_commit_tx.notify();
689 self.advance_timelines_interval.reset();
693 }
694
695 pub(crate) async fn sequence_deferred_plan(&mut self, conn_id: ConnectionId) {
697 let Some(mut deferred) = self.deferred_plans.remove(&conn_id) else {
698 return;
700 };
701 if let Err(e) = deferred.validity.check(self.catalog()) {
702 deferred.ctx.retire(Err(e));
703 } else {
704 self.sequence_plan(
705 deferred.ctx,
706 deferred.plan,
707 deferred.resolved_ids,
708 deferred.sql_impl_resolved_ids,
709 )
710 .await;
711 }
712 }
713
714 #[instrument(name = "coord::stage_group_commit")]
718 pub(crate) fn stage_group_commit(&mut self, permit: Option<GroupCommitPermit>) {
719 let pending_writes = std::mem::take(&mut self.pending_writes);
720 let contains_internal_system_write = pending_writes
721 .iter()
722 .any(|write| write.is_internal_system());
723
724 let mut appends: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>> = BTreeMap::new();
725 let mut responses = Vec::with_capacity(pending_writes.len());
726 let mut statement_logging_ids = Vec::new();
727 let mut notifies = Vec::new();
728 let mut internal_results = Vec::new();
729
730 for pending_write in pending_writes {
731 match pending_write {
732 PendingWriteTxn::User {
733 span: _,
734 writes,
735 responder:
736 UserWriteResponder::Session(PendingTxn {
737 ctx,
738 response,
739 action,
740 }),
741 } => {
742 if let Some(id) = Self::stale_write_target(self.catalog(), &writes) {
749 let err = AdapterError::ConcurrentDependencyMutation {
750 dependency_id: id.to_string(),
751 };
752 let (ctx, result) = CompletedClientTransmitter::new(
753 ctx,
754 Err(err),
755 EndTransactionAction::Rollback,
756 )
757 .finalize();
758 ctx.retire(result);
759 continue;
760 }
761
762 for (id, table_data) in writes {
763 if self.catalog().try_get_entry(&id).is_some() {
769 appends.entry(id).or_default().extend(table_data);
770 }
771 }
772 if let Some(id) = ctx.extra().contents() {
773 statement_logging_ids.push(id);
774 }
775
776 responses.push(CompletedClientTransmitter::new(ctx, response, action));
777 }
778 PendingWriteTxn::User {
779 span: _,
780 writes,
781 responder: UserWriteResponder::Internal { target, result, .. },
782 } => {
783 let current_global_id = self
784 .catalog()
785 .try_get_entry(&target.item_id)
786 .map(|entry| entry.latest_global_id());
787 if current_global_id != Some(target.global_id) {
788 result.send(WriteResult::TargetChanged);
789 continue;
790 }
791 assert!(
797 writes.keys().all(|id| *id == target.item_id),
798 "frontend write for {:?} carries other tables: {:?}",
799 target.item_id,
800 writes.keys().collect::<Vec<_>>(),
801 );
802 appends
803 .entry(target.item_id)
804 .or_default()
805 .extend(writes.into_values().flatten());
806 internal_results.push(result);
807 }
808 PendingWriteTxn::System { updates, source } => {
809 for update in updates {
810 appends.entry(update.id).or_default().push(update.data);
811 }
812 match source {
814 BuiltinTableUpdateSource::Internal(tx)
815 | BuiltinTableUpdateSource::Background(tx) => notifies.push(tx),
816 }
817 }
818 }
819 }
820
821 let mut all_appends = Vec::with_capacity(appends.len());
824 for (item_id, table_data) in appends.into_iter() {
825 let mut all_rows = Vec::new();
826 let mut all_data = Vec::new();
827 for data in table_data {
828 match data {
829 TableData::Rows(rows) => all_rows.extend(rows),
830 TableData::Batches(_) => all_data.push(data),
831 }
832 }
833 differential_dataflow::consolidation::consolidate(&mut all_rows);
834 all_data.push(TableData::Rows(all_rows));
835
836 all_appends.push((item_id, all_data));
838 }
839
840 let appends: Vec<_> = all_appends
841 .into_iter()
842 .map(|(id, updates)| {
843 let gid = self.catalog().get_entry(&id).latest_global_id();
844 (gid, updates)
845 })
846 .collect();
847
848 let request = GroupCommitRequest {
850 appends,
851 responses,
852 statement_logging_ids,
853 notifies,
854 internal_results,
855 permits: permit.into_iter().collect(),
856 contains_internal_system_write,
857 span: Span::current(),
858 };
859 if self
860 .group_committer_tx
861 .send(TableWriteCmd::GroupCommit(request))
862 .is_err()
863 {
864 warn!("group committer task gone, dropping staged group commit");
866 }
867 }
868
869 fn stale_write_target(
879 catalog: &Catalog,
880 writes: &BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
881 ) -> Option<CatalogItemId> {
882 writes.iter().find_map(|(id, table_data)| {
883 let entry = catalog.try_get_entry(id)?;
885 let arity = entry
886 .relation_desc_latest()
887 .expect("write target is a table")
888 .arity();
889 let stale = table_data.iter().any(|data| match data {
890 TableData::Rows(rows) => rows
891 .first()
892 .is_some_and(|(row, _)| row.iter().count() != arity),
893 TableData::Batches(_) => false,
894 });
895 stale.then_some(*id)
896 })
897 }
898
899 pub(crate) async fn register_tables_via_committer(
901 &self,
902 tables: Vec<TableRegistration>,
903 ) -> Timestamp {
904 let (tx, rx) = oneshot::channel();
905 if self
906 .group_committer_tx
907 .send(TableWriteCmd::Register { tables, result: tx })
908 .is_err()
909 {
910 halt!("group committer terminated before a table registration could be submitted");
911 }
912 match rx.await {
913 Ok(ts) => ts,
914 Err(_) => halt!("group committer terminated with a table registration outstanding"),
915 }
916 }
917
918 pub(crate) async fn forget_tables_via_committer(&self, ids: Vec<GlobalId>) -> Timestamp {
920 let (tx, rx) = oneshot::channel();
921 if self
922 .group_committer_tx
923 .send(TableWriteCmd::Forget { ids, result: tx })
924 .is_err()
925 {
926 halt!("group committer terminated before a table forget could be submitted");
927 }
928 match rx.await {
929 Ok(ts) => ts,
930 Err(_) => halt!("group committer terminated with a table forget outstanding"),
931 }
932 }
933
934 pub(crate) fn submit_write(&mut self, pending_write_txn: PendingWriteTxn) {
936 if self.controller.read_only() {
937 panic!(
938 "attempting table write in read-only mode: {:?}",
939 pending_write_txn
940 );
941 }
942 self.pending_writes.push(pending_write_txn);
943 self.trigger_group_commit();
944 }
945
946 pub(crate) fn builtin_table_update<'a>(&'a mut self) -> BuiltinTableAppend<'a> {
948 BuiltinTableAppend { coord: self }
949 }
950
951 pub(crate) fn defer_plan<F>(&mut self, wait_future: F, plan: DeferredPlan)
952 where
953 F: Future<Output = ()> + Send + 'static,
954 {
955 let conn_id = plan.ctx.session().conn_id().clone();
956 self.deferred_plans.insert(conn_id.clone(), plan);
957
958 let internal_cmd_tx = self.internal_cmd_tx.clone();
959 let conn_id_ = conn_id.clone();
960 mz_ore::task::spawn(|| format!("defer plan {conn_id_}"), async move {
961 tracing::info!(%conn_id, "deferring plan");
962 wait_future.await;
963
964 let _ = internal_cmd_tx.send(Message::DeferredPlanReady { conn_id });
966 });
967 }
968}
969
970pub struct BuiltinTableAppend<'a> {
972 coord: &'a mut Coordinator,
973}
974
975pub type BuiltinTableAppendNotify = Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>;
986
987pub struct BuiltinTableAppendCompletion {
989 notify: BuiltinTableAppendNotify,
990}
991
992impl BuiltinTableAppendCompletion {
993 pub fn new(notify: BuiltinTableAppendNotify) -> Self {
994 Self { notify }
995 }
996
997 pub fn into_notify(self) -> BuiltinTableAppendNotify {
998 self.notify
999 }
1000}
1001
1002impl<'a> BuiltinTableAppend<'a> {
1003 pub fn background(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1014 if self.coord.controller.read_only() {
1015 self.coord
1016 .buffered_builtin_table_updates
1017 .as_mut()
1018 .expect("in read-only mode")
1019 .append(&mut updates);
1020
1021 return Box::pin(futures::future::ready(()));
1022 }
1023
1024 let (tx, rx) = oneshot::channel();
1025 self.coord.pending_writes.push(PendingWriteTxn::System {
1026 updates,
1027 source: BuiltinTableUpdateSource::Background(tx),
1028 });
1029
1030 Box::pin(rx.map(|_| ()))
1031 }
1032
1033 pub fn defer(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1042 if self.coord.controller.read_only() {
1043 self.coord
1044 .buffered_builtin_table_updates
1045 .as_mut()
1046 .expect("in read-only mode")
1047 .append(&mut updates);
1048
1049 return Box::pin(futures::future::ready(()));
1050 }
1051
1052 let (tx, rx) = oneshot::channel();
1053 self.coord.pending_writes.push(PendingWriteTxn::System {
1054 updates,
1055 source: BuiltinTableUpdateSource::Internal(tx),
1056 });
1057 self.coord.trigger_group_commit();
1058
1059 Box::pin(rx.map(|_| ()))
1060 }
1061
1062 pub fn execute(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1066 if self.coord.controller.read_only() {
1067 self.coord
1068 .buffered_builtin_table_updates
1069 .as_mut()
1070 .expect("in read-only mode")
1071 .append(&mut updates);
1072
1073 return Box::pin(futures::future::ready(()));
1074 }
1075
1076 let (tx, rx) = oneshot::channel();
1077
1078 self.coord.pending_writes.push(PendingWriteTxn::System {
1081 updates,
1082 source: BuiltinTableUpdateSource::Internal(tx),
1083 });
1084 self.coord.stage_group_commit(None);
1085
1086 self.coord.advance_timelines_interval.reset();
1088
1089 Box::pin(rx.map(|_| ()))
1090 }
1091}
1092
1093pub fn notifier() -> (GroupCommitNotifier, GroupCommitWaiter) {
1096 let notify = Arc::new(Notify::new());
1097 let in_progress = Arc::new(Semaphore::new(1));
1098
1099 let notifier = GroupCommitNotifier {
1100 notify: Arc::clone(¬ify),
1101 };
1102 let waiter = GroupCommitWaiter {
1103 notify,
1104 in_progress,
1105 };
1106
1107 (notifier, waiter)
1108}
1109
1110#[derive(Debug, Clone)]
1113pub struct GroupCommitNotifier {
1114 notify: Arc<Notify>,
1116}
1117
1118impl GroupCommitNotifier {
1119 pub fn notify(&self) {
1121 self.notify.notify_one()
1122 }
1123}
1124
1125#[derive(Debug)]
1128pub struct GroupCommitWaiter {
1129 notify: Arc<Notify>,
1131 in_progress: Arc<Semaphore>,
1133}
1134static_assertions::assert_not_impl_all!(GroupCommitWaiter: Clone);
1135
1136impl GroupCommitWaiter {
1137 pub async fn ready(&self) -> GroupCommitPermit {
1148 let permit = Semaphore::acquire_owned(Arc::clone(&self.in_progress))
1149 .await
1150 .expect("semaphore should not close");
1151
1152 self.notify.notified().await;
1155
1156 GroupCommitPermit {
1157 _permit: Some(permit),
1158 }
1159 }
1160}
1161
1162#[derive(Debug)]
1167pub struct GroupCommitPermit {
1168 _permit: Option<OwnedSemaphorePermit>,
1172}
1173
1174pub(crate) fn waiting_on_startup_appends(
1190 catalog: &Catalog,
1191 session: &mut Session,
1192 plan: &Plan,
1193) -> Option<(BTreeSet<CatalogItemId>, BoxFuture<'static, ()>)> {
1194 let depends_on = match plan {
1198 Plan::Select(plan) => plan.source.depends_on(),
1199 Plan::ReadThenWrite(plan) => plan.selection.depends_on(),
1200 Plan::ShowColumns(plan) => plan.select_plan.source.depends_on(),
1201 Plan::Subscribe(plan) => plan.from.depends_on(),
1202 Plan::ExplainPlan(ExplainPlanPlan {
1203 explainee: Explainee::Statement(ExplaineeStatement::Select { plan, .. }),
1204 ..
1205 }) => plan.source.depends_on(),
1206 Plan::ExplainTimestamp(ExplainTimestampPlan { raw_plan, .. }) => raw_plan.depends_on(),
1207 Plan::CreateConnection(_)
1208 | Plan::CreateDatabase(_)
1209 | Plan::CreateSchema(_)
1210 | Plan::CreateRole(_)
1211 | Plan::CreateNetworkPolicy(_)
1212 | Plan::CreateCluster(_)
1213 | Plan::CreateClusterReplica(_)
1214 | Plan::CreateSource(_)
1215 | Plan::CreateSources(_)
1216 | Plan::CreateSecret(_)
1217 | Plan::CreateSink(_)
1218 | Plan::CreateTable(_)
1219 | Plan::CreateView(_)
1220 | Plan::CreateMaterializedView(_)
1221 | Plan::CreateIndex(_)
1222 | Plan::CreateMetricSink(_)
1223 | Plan::CreateType(_)
1224 | Plan::Comment(_)
1225 | Plan::DiscardTemp
1226 | Plan::DiscardAll
1227 | Plan::DropObjects(_)
1228 | Plan::DropOwned(_)
1229 | Plan::EmptyQuery
1230 | Plan::ShowAllVariables
1231 | Plan::ShowCreate(_)
1232 | Plan::ShowVariable(_)
1233 | Plan::InspectShard(_)
1234 | Plan::SetVariable(_)
1235 | Plan::ResetVariable(_)
1236 | Plan::SetTransaction(_)
1237 | Plan::StartTransaction(_)
1238 | Plan::CommitTransaction(_)
1239 | Plan::AbortTransaction(_)
1240 | Plan::CopyFrom(_)
1241 | Plan::CopyTo(_)
1242 | Plan::ExplainPlan(_)
1243 | Plan::ExplainPushdown(_)
1244 | Plan::ExplainSinkSchema(_)
1245 | Plan::Insert(_)
1246 | Plan::AlterNetworkPolicy(_)
1247 | Plan::AlterNoop(_)
1248 | Plan::AlterClusterRename(_)
1249 | Plan::AlterClusterSwap(_)
1250 | Plan::AlterClusterReplicaRename(_)
1251 | Plan::AlterCluster(_)
1252 | Plan::AlterConnection(_)
1253 | Plan::AlterSource(_)
1254 | Plan::AlterSetCluster(_)
1255 | Plan::AlterItemRename(_)
1256 | Plan::AlterRetainHistory(_)
1257 | Plan::AlterSourceTimestampInterval(_)
1258 | Plan::AlterSchemaRename(_)
1259 | Plan::AlterSchemaSwap(_)
1260 | Plan::AlterSecret(_)
1261 | Plan::AlterSink(_)
1262 | Plan::AlterSystemSet(_)
1263 | Plan::AlterSystemReset(_)
1264 | Plan::AlterSystemResetAll(_)
1265 | Plan::AlterRole(_)
1266 | Plan::AlterOwner(_)
1267 | Plan::AlterTableAddColumn(_)
1268 | Plan::AlterMaterializedViewApplyReplacement(_)
1269 | Plan::Declare(_)
1270 | Plan::Fetch(_)
1271 | Plan::Close(_)
1272 | Plan::Prepare(_)
1273 | Plan::Execute(_)
1274 | Plan::Deallocate(_)
1275 | Plan::Raise(_)
1276 | Plan::GrantRole(_)
1277 | Plan::RevokeRole(_)
1278 | Plan::GrantPrivileges(_)
1279 | Plan::RevokePrivileges(_)
1280 | Plan::AlterDefaultPrivileges(_)
1281 | Plan::ReassignOwned(_)
1282 | Plan::ValidateConnection(_)
1283 | Plan::SideEffectingFunc(_) => BTreeSet::default(),
1284 };
1285 let depends_on_required_id = REQUIRED_BUILTIN_TABLES
1286 .iter()
1287 .map(|table| catalog.resolve_builtin_table(&**table))
1288 .any(|id| {
1289 catalog
1290 .get_global_ids(&id)
1291 .any(|gid| depends_on.contains(&gid))
1292 });
1293
1294 if !depends_on_required_id {
1297 return None;
1298 }
1299
1300 match session.clear_builtin_table_updates() {
1307 Some(wait_future) => {
1308 let depends_on = depends_on
1309 .into_iter()
1310 .map(|gid| catalog.get_entry_by_global_id(&gid).id())
1311 .collect();
1312 Some((depends_on, wait_future.boxed()))
1313 }
1314 None => None,
1315 }
1316}
1317
1318#[cfg(test)]
1319mod tests {
1320 use std::sync::Mutex;
1321 use std::sync::atomic::{AtomicUsize, Ordering};
1322
1323 use async_trait::async_trait;
1324 use mz_ore::metrics::MetricsRegistry;
1325 use mz_ore::now::SYSTEM_TIME;
1326 use timely::progress::Antichain;
1327
1328 use super::*;
1329 use crate::catalog::Catalog;
1330
1331 #[mz_ore::test(tokio::test)]
1332 async fn internal_write_responder_reports_indeterminate_on_drop() {
1333 let (tx, rx) = oneshot::channel();
1334 drop(InternalWriteResponder::new(tx));
1335 assert!(matches!(rx.await, Ok(WriteResult::Indeterminate)));
1336 }
1337
1338 #[derive(Debug, Default)]
1339 struct MemTimestampOracle {
1340 read_write_ts: Mutex<(Timestamp, Timestamp)>,
1341 apply_writes: AtomicUsize,
1342 }
1343
1344 impl MemTimestampOracle {
1345 fn starting_at(ts: Timestamp) -> Self {
1346 Self {
1347 read_write_ts: Mutex::new((ts, ts)),
1348 apply_writes: AtomicUsize::new(0),
1349 }
1350 }
1351 }
1352
1353 #[async_trait]
1354 impl TimestampOracle<Timestamp> for MemTimestampOracle {
1355 async fn write_ts(&self) -> WriteTimestamp {
1356 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1357 let new_write_ts = std::cmp::max(*read_ts, *write_ts).step_forward();
1358 *write_ts = new_write_ts;
1359 WriteTimestamp {
1360 timestamp: new_write_ts,
1361 advance_to: new_write_ts.step_forward(),
1362 }
1363 }
1364
1365 async fn peek_write_ts(&self) -> Timestamp {
1366 let (_, write_ts) = &*self.read_write_ts.lock().expect("lock poisoned");
1367 *write_ts
1368 }
1369
1370 async fn read_ts(&self) -> Timestamp {
1371 let (read_ts, _) = &*self.read_write_ts.lock().expect("lock poisoned");
1372 *read_ts
1373 }
1374
1375 async fn apply_write(&self, lower_bound: Timestamp) {
1376 self.apply_writes.fetch_add(1, Ordering::SeqCst);
1377 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1378 *read_ts = std::cmp::max(*read_ts, lower_bound);
1379 *write_ts = std::cmp::max(*read_ts, *write_ts);
1380 }
1381 }
1382
1383 #[derive(Debug)]
1384 struct ConflictingTableWriteHandle {
1385 conflicts: usize,
1386 calls: AtomicUsize,
1387 write_timestamps: Mutex<Vec<Timestamp>>,
1388 }
1389
1390 impl ConflictingTableWriteHandle {
1391 fn new(conflicts: usize) -> Self {
1392 Self {
1393 conflicts,
1394 calls: AtomicUsize::new(0),
1395 write_timestamps: Mutex::new(Vec::new()),
1396 }
1397 }
1398
1399 fn respond(&self, write_ts: Timestamp) -> oneshot::Receiver<Result<(), StorageError>> {
1400 let call = self.calls.fetch_add(1, Ordering::SeqCst);
1401 self.write_timestamps
1402 .lock()
1403 .expect("lock poisoned")
1404 .push(write_ts);
1405 let (tx, rx) = oneshot::channel();
1406 let result = if call < self.conflicts {
1407 Err(StorageError::InvalidUppers(vec![
1408 mz_storage_types::controller::InvalidUpper {
1409 id: GlobalId::User(1),
1410 current_upper: Antichain::from_elem(write_ts.step_forward()),
1411 },
1412 ]))
1413 } else {
1414 Ok(())
1415 };
1416 tx.send(result).expect("receiver still in scope");
1417 rx
1418 }
1419 }
1420
1421 impl TableWriteHandle for ConflictingTableWriteHandle {
1422 fn append(
1423 &self,
1424 write_ts: Timestamp,
1425 _advance_to: Timestamp,
1426 _commands: Vec<(GlobalId, Vec<TableData>)>,
1427 ) -> oneshot::Receiver<Result<(), StorageError>> {
1428 self.respond(write_ts)
1429 }
1430
1431 fn register(
1432 &self,
1433 register_ts: Timestamp,
1434 _tables: Vec<TableRegistration>,
1435 ) -> oneshot::Receiver<Result<(), StorageError>> {
1436 self.respond(register_ts)
1437 }
1438
1439 fn forget(
1440 &self,
1441 forget_ts: Timestamp,
1442 _ids: Vec<GlobalId>,
1443 ) -> oneshot::Receiver<Result<(), StorageError>> {
1444 self.respond(forget_ts)
1445 }
1446 }
1447
1448 #[mz_ore::test(tokio::test)]
1449 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_conflict_retry() {
1451 Catalog::with_debug(|catalog| async move {
1452 let initial_upper = catalog.current_upper().await;
1454 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1455 let handle = Arc::new(ConflictingTableWriteHandle::new(1));
1456 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1457 Arc::<MemTimestampOracle>::clone(&oracle);
1458 let handle_dyn: Arc<dyn TableWriteHandle> =
1459 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1460 let (_tx, rx) = mpsc::unbounded_channel();
1461 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1462 let committer = GroupCommitter {
1463 rx,
1464 oracle: oracle_dyn,
1465 table_write_handle: handle_dyn,
1466 catalog_upper: catalog.upper_handle(),
1467 internal_cmd_tx,
1468 now: SYSTEM_TIME.clone(),
1469 metrics: Metrics::register_into(&MetricsRegistry::new()),
1470 max_attempts: ConfigValHandle::disconnected(2),
1471 };
1472
1473 let write_ts = committer
1474 .write_to_txns(None, |ts, advance_to| {
1475 assert_eq!(advance_to, ts.step_forward());
1476 handle.append(ts, advance_to, Vec::new())
1477 })
1478 .await
1479 .expect("worker stays alive in this test");
1480
1481 let attempts = handle
1482 .write_timestamps
1483 .lock()
1484 .expect("lock poisoned")
1485 .clone();
1486 assert_eq!(attempts.len(), 2, "one conflict, one success");
1487 assert!(
1488 attempts[0] > initial_upper,
1489 "attempts start past the initial catalog upper: {attempts:?} vs {initial_upper}"
1490 );
1491 assert!(
1492 attempts[0] < attempts[1],
1493 "retry must use a fresh, larger timestamp: {attempts:?}"
1494 );
1495 assert_eq!(write_ts.timestamp, attempts[1]);
1496
1497 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1498 assert_eq!(oracle.read_ts().await, write_ts.timestamp);
1499
1500 let catalog_upper = catalog.current_upper().await;
1501 assert!(
1502 catalog_upper >= write_ts.advance_to,
1503 "catalog upper {catalog_upper} must cover the write's advance_to {}",
1504 write_ts.advance_to
1505 );
1506
1507 catalog.expire().await;
1508 })
1509 .await;
1510 }
1511
1512 #[mz_ore::test(tokio::test)]
1513 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_zero_limit_allows_one_attempt() {
1515 Catalog::with_debug(|catalog| async move {
1516 let initial_upper = catalog.current_upper().await;
1517 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1518 let handle = Arc::new(ConflictingTableWriteHandle::new(0));
1519 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1520 Arc::<MemTimestampOracle>::clone(&oracle);
1521 let handle_dyn: Arc<dyn TableWriteHandle> =
1522 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1523 let (_tx, rx) = mpsc::unbounded_channel();
1524 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1525 let committer = GroupCommitter {
1526 rx,
1527 oracle: oracle_dyn,
1528 table_write_handle: handle_dyn,
1529 catalog_upper: catalog.upper_handle(),
1530 internal_cmd_tx,
1531 now: SYSTEM_TIME.clone(),
1532 metrics: Metrics::register_into(&MetricsRegistry::new()),
1533 max_attempts: ConfigValHandle::disconnected(0),
1534 };
1535
1536 let write_ts = committer
1537 .write_to_txns(None, |ts, advance_to| {
1538 handle.append(ts, advance_to, Vec::new())
1539 })
1540 .await
1541 .expect("zero limit permits one successful attempt");
1542
1543 let attempts = handle
1544 .write_timestamps
1545 .lock()
1546 .expect("lock poisoned")
1547 .clone();
1548 assert_eq!(attempts, vec![write_ts.timestamp]);
1549 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1550
1551 catalog.expire().await;
1552 })
1553 .await;
1554 }
1555}