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 futures::future::{BoxFuture, FutureExt};
43use mz_adapter_types::connection::ConnectionId;
44use mz_adapter_types::dyncfgs::GROUP_COMMIT_MAX_ATTEMPTS;
45use mz_catalog::builtin::{BuiltinTable, MZ_SESSIONS};
46use mz_dyncfg::{ConfigSet, ConfigValHandle};
47use mz_expr::CollectionPlan;
48use mz_ore::halt;
49use mz_ore::instrument;
50use mz_ore::now::NowFn;
51use mz_ore::task;
52use mz_repr::{CatalogItemId, GlobalId, Timestamp};
53use mz_sql::plan::{ExplainPlanPlan, ExplainTimestampPlan, Explainee, ExplaineeStatement, Plan};
54use mz_storage_client::client::TableData;
55use mz_storage_client::controller::{TableRegistration, TableWriteHandle};
56use mz_storage_types::controller::StorageError;
57use mz_timestamp_oracle::{TimestampOracle, WriteTimestamp};
58use smallvec::SmallVec;
59use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore, mpsc, oneshot};
60use tracing::{Instrument, Span, info, warn};
61
62use crate::AdapterError;
63use crate::catalog::{BuiltinTableUpdate, Catalog, CatalogUpperHandle};
64use crate::coord::timeline::write_ts_upper_bound;
65use crate::coord::{Coordinator, Message, PendingTxn};
66use crate::metrics::Metrics;
67use crate::session::{EndTransactionAction, Session};
68use crate::statement_logging::StatementLoggingId;
69use crate::util::{CompletedClientTransmitter, ResultExt};
70
71pub(crate) static REQUIRED_BUILTIN_TABLES: &[&LazyLock<BuiltinTable>] = &[&MZ_SESSIONS];
73
74#[derive(Debug)]
76pub(crate) enum BuiltinTableUpdateSource {
77 Internal(oneshot::Sender<()>),
79 Background(oneshot::Sender<()>),
81}
82
83#[derive(Debug, Clone)]
85pub enum WriteResult {
86 Success { timestamp: Timestamp },
88 TimestampPassed {
90 target_timestamp: Timestamp,
91 next_eligible_timestamp: Timestamp,
92 },
93 TimestampTooFarAhead {
96 target_timestamp: Timestamp,
97 limit: Timestamp,
98 },
99 Canceled,
101 ReadOnly,
103 TargetChanged,
105 Indeterminate,
107}
108
109#[derive(Debug)]
116pub struct InternalWriteResponder {
117 tx: Option<oneshot::Sender<WriteResult>>,
118}
119
120impl InternalWriteResponder {
121 pub(crate) fn new(tx: oneshot::Sender<WriteResult>) -> Self {
122 Self { tx: Some(tx) }
123 }
124
125 pub(crate) fn send(mut self, result: WriteResult) {
126 if let Some(tx) = self.tx.take() {
127 let _ = tx.send(result);
128 }
129 }
130}
131
132impl Drop for InternalWriteResponder {
133 fn drop(&mut self) {
134 if let Some(tx) = self.tx.take() {
135 let _ = tx.send(WriteResult::Indeterminate);
136 }
137 }
138}
139
140#[derive(Debug)]
142pub(crate) enum UserWriteResponder {
143 Session(PendingTxn),
146 Internal {
148 conn_id: ConnectionId,
149 target: WriteTarget,
153 result: InternalWriteResponder,
154 },
155}
156
157#[derive(Debug, Clone, Copy)]
159pub(crate) struct WriteTarget {
160 pub(crate) item_id: CatalogItemId,
161 pub(crate) global_id: GlobalId,
162}
163
164#[derive(Debug)]
166pub(crate) enum PendingWriteTxn {
167 User {
170 span: Span,
171 writes: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
173 responder: UserWriteResponder,
175 },
176 System {
178 updates: Vec<BuiltinTableUpdate>,
179 source: BuiltinTableUpdateSource,
180 },
181}
182
183impl PendingWriteTxn {
184 fn is_internal_system(&self) -> bool {
185 match self {
186 PendingWriteTxn::System {
187 source: BuiltinTableUpdateSource::Internal(_),
188 ..
189 } => true,
190 _ => false,
191 }
192 }
193}
194
195pub(crate) enum TableWriteCmd {
196 GroupCommit(GroupCommitRequest),
197 TimestampedWrite(TimestampedWriteRequest),
198 Register {
199 tables: Vec<TableRegistration>,
200 result: oneshot::Sender<Timestamp>,
201 },
202 Forget {
203 ids: Vec<GlobalId>,
204 result: oneshot::Sender<Timestamp>,
205 },
206}
207
208pub(crate) struct TimestampedWriteRequest {
217 pub(crate) appends: Vec<(GlobalId, Vec<TableData>)>,
218 pub(crate) target_timestamp: Timestamp,
219 pub(crate) result: InternalWriteResponder,
220 pub(crate) span: Span,
221}
222
223pub(crate) struct GroupCommitRequest {
225 appends: Vec<(GlobalId, Vec<TableData>)>,
227 responses: Vec<CompletedClientTransmitter>,
228 statement_logging_ids: Vec<StatementLoggingId>,
229 notifies: Vec<oneshot::Sender<()>>,
230 internal_results: Vec<InternalWriteResponder>,
231 permits: Vec<GroupCommitPermit>,
233 contains_internal_system_write: bool,
234 span: Span,
235}
236
237impl GroupCommitRequest {
238 fn merge(&mut self, other: GroupCommitRequest) {
239 let GroupCommitRequest {
240 appends,
241 responses,
242 statement_logging_ids,
243 notifies,
244 internal_results,
245 permits,
246 contains_internal_system_write,
247 span: _,
248 } = other;
249 self.appends.extend(appends);
250 self.responses.extend(responses);
251 self.statement_logging_ids.extend(statement_logging_ids);
252 self.notifies.extend(notifies);
253 self.internal_results.extend(internal_results);
254 self.permits.extend(permits);
255 self.contains_internal_system_write |= contains_internal_system_write;
256 }
257}
258
259pub(crate) struct GroupCommitter {
264 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
265 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
266 table_write_handle: Arc<dyn TableWriteHandle>,
267 catalog_upper: CatalogUpperHandle,
268 internal_cmd_tx: mpsc::UnboundedSender<Message>,
269 now: NowFn,
270 metrics: Metrics,
271 max_attempts: ConfigValHandle<usize>,
272}
273
274enum TxnsWriteAttempt {
276 Applied,
278 UpperConflict,
281 WorkerGone,
283}
284
285impl GroupCommitter {
286 async fn run(mut self) {
287 while let Some(cmd) = self.rx.recv().await {
288 let mut next = Some(cmd);
291 while let Some(cmd) = next.take() {
292 match cmd {
293 TableWriteCmd::GroupCommit(request) => {
294 let span = request.span.clone();
295 match self.commit(request).instrument(span).await {
296 ControlFlow::Continue(deferred) => next = deferred,
297 ControlFlow::Break(()) => return,
298 }
299 }
300 TableWriteCmd::TimestampedWrite(request) => {
301 let span = request.span.clone();
302 if self
303 .commit_timestamped(request)
304 .instrument(span)
305 .await
306 .is_break()
307 {
308 return;
309 }
310 }
311 TableWriteCmd::Register { tables, result } => {
312 let Some(write_ts) = self
313 .write_to_txns(None, |ts, _advance_to| {
314 self.table_write_handle.register(ts, tables.clone())
315 })
316 .await
317 else {
318 return;
319 };
320 let _ = result.send(write_ts.timestamp);
321 }
322 TableWriteCmd::Forget { ids, result } => {
323 let Some(write_ts) = self
324 .write_to_txns(None, |ts, _advance_to| {
325 self.table_write_handle.forget(ts, ids.clone())
326 })
327 .await
328 else {
329 return;
330 };
331 let _ = result.send(write_ts.timestamp);
332 }
333 }
334 }
335 }
336 }
337
338 async fn commit_timestamped(&self, request: TimestampedWriteRequest) -> ControlFlow<(), ()> {
359 let TimestampedWriteRequest {
360 appends,
361 target_timestamp,
362 result,
363 span: _,
364 } = request;
365
366 let oracle_write_ts = self.oracle.peek_write_ts().await;
367 if target_timestamp <= oracle_write_ts {
368 result.send(WriteResult::TimestampPassed {
369 target_timestamp,
370 next_eligible_timestamp: oracle_write_ts.step_forward(),
371 });
372 return ControlFlow::Continue(());
373 }
374
375 let now: Timestamp = (self.now)().into();
378 let limit = write_ts_upper_bound(&now);
379 if target_timestamp > limit {
380 result.send(WriteResult::TimestampTooFarAhead {
381 target_timestamp,
382 limit,
383 });
384 return ControlFlow::Continue(());
385 }
386
387 let write_ts = WriteTimestamp {
388 timestamp: target_timestamp,
389 advance_to: target_timestamp.step_forward(),
390 };
391 match self
392 .attempt_write_to_txns(
393 &write_ts,
394 Some(&self.metrics.append_table_duration_seconds),
395 |ts, advance_to| self.table_write_handle.append(ts, advance_to, appends),
396 )
397 .await
398 {
399 TxnsWriteAttempt::Applied => {}
400 TxnsWriteAttempt::UpperConflict => {
401 result.send(WriteResult::TimestampPassed {
402 target_timestamp,
403 next_eligible_timestamp: write_ts.advance_to,
404 });
405 return ControlFlow::Continue(());
406 }
407 TxnsWriteAttempt::WorkerGone => {
408 warn!("table write worker gone with a timestamped write outstanding");
409 return ControlFlow::Break(());
410 }
411 }
412
413 if self
414 .internal_cmd_tx
415 .send(Message::GroupCommitApplied {
416 responses: Vec::new(),
417 statement_logging_ids: Vec::new(),
418 internal_results: vec![result],
419 write_ts: target_timestamp,
420 })
421 .is_err()
422 {
423 warn!("coordinator shut down before a timestamped write could be finalized");
424 }
425 ControlFlow::Continue(())
426 }
427
428 async fn write_to_txns(
433 &self,
434 op_duration_metric: Option<&prometheus::Histogram>,
435 mut op: impl FnMut(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
436 ) -> Option<WriteTimestamp> {
437 let mut attempt = 0;
439 loop {
440 let max_attempts = self.max_attempts.get().max(1);
441 if attempt >= max_attempts {
442 halt!(
443 "txns-shard write reached attempt limit {max_attempts} after {attempt} conflicts, rebuilding"
444 );
445 }
446 attempt += 1;
447 let write_ts = self.oracle.write_ts().await;
448
449 match self
452 .attempt_write_to_txns(&write_ts, op_duration_metric, |ts, advance_to| {
453 op(ts, advance_to)
454 })
455 .await
456 {
457 TxnsWriteAttempt::Applied => return Some(write_ts),
458 TxnsWriteAttempt::UpperConflict => {
459 warn!(
460 write_ts = %write_ts.timestamp,
461 attempt,
462 "txns-shard write conflicted with another writer, retrying at a fresh timestamp"
463 );
464 continue;
465 }
466 TxnsWriteAttempt::WorkerGone => {
467 warn!("table write worker gone (process shutting down), winding down");
469 return None;
470 }
471 }
472 }
473 }
474
475 async fn attempt_write_to_txns(
481 &self,
482 write_ts: &WriteTimestamp,
483 op_duration_metric: Option<&prometheus::Histogram>,
484 op: impl FnOnce(Timestamp, Timestamp) -> oneshot::Receiver<Result<(), StorageError>>,
485 ) -> TxnsWriteAttempt {
486 let catalog_upper_start = Instant::now();
487 self.catalog_upper
488 .advance_upper(write_ts.advance_to)
489 .await
490 .unwrap_or_terminate("unable to advance catalog upper");
491 self.metrics
492 .group_commit_catalog_upper_seconds
493 .observe(catalog_upper_start.elapsed().as_secs_f64());
494
495 let op_start = Instant::now();
496 let op_res = op(write_ts.timestamp, write_ts.advance_to).await;
497 if let Some(metric) = op_duration_metric {
498 metric.observe(op_start.elapsed().as_secs_f64());
499 }
500
501 match op_res {
502 Ok(Ok(())) => {}
503 Ok(Err(StorageError::InvalidUppers(_))) => return TxnsWriteAttempt::UpperConflict,
504 Ok(Err(other)) => {
505 Err::<(), _>(other).unwrap_or_terminate("cannot fail to write to txns shard");
506 unreachable!("unwrap_or_terminate does not return on Err");
507 }
508 Err(_recv) => return TxnsWriteAttempt::WorkerGone,
509 }
510
511 let now: Timestamp = (self.now)().into();
512 crate::coord::timeline::check_runaway_write_ts(&now, write_ts.timestamp);
513
514 fail::fail_point!("group_commit_before_apply_write");
522
523 self.oracle.apply_write(write_ts.timestamp).await;
524
525 TxnsWriteAttempt::Applied
526 }
527
528 async fn commit(
533 &mut self,
534 mut request: GroupCommitRequest,
535 ) -> ControlFlow<(), Option<TableWriteCmd>> {
536 let mut deferred_cmd = None;
537 let mut rx_closed = false;
540
541 loop {
545 while deferred_cmd.is_none() {
546 match self.rx.try_recv() {
547 Ok(TableWriteCmd::GroupCommit(other)) => request.merge(other),
548 Ok(other) => deferred_cmd = Some(other),
549 Err(_) => break,
550 }
551 }
552
553 if request.contains_internal_system_write || deferred_cmd.is_some() {
556 break;
557 }
558 let ts = self.oracle.peek_write_ts().await;
559 let now: Timestamp = (self.now)().into();
560 if ts <= now {
561 break;
562 }
563 let remaining_ms = std::cmp::min(ts.saturating_sub(now), Timestamp::from(1_000u64));
566 let sleep = tokio::time::sleep(Duration::from_millis(remaining_ms.into()));
567 tokio::pin!(sleep);
568 loop {
569 tokio::select! {
570 _ = &mut sleep => break,
571 cmd = self.rx.recv(), if deferred_cmd.is_none() && !rx_closed => match cmd {
572 Some(TableWriteCmd::GroupCommit(other)) => request.merge(other),
573 Some(other) => deferred_cmd = Some(other),
574 None => rx_closed = true,
575 },
576 }
577 if request.contains_internal_system_write || deferred_cmd.is_some() {
578 break;
579 }
580 }
581 }
582
583 let GroupCommitRequest {
584 appends,
585 responses,
586 statement_logging_ids,
587 notifies,
588 internal_results,
589 permits,
590 contains_internal_system_write: _,
591 span: _,
592 } = request;
593
594 let append_metric = self.metrics.append_table_duration_seconds.clone();
595 let Some(write_ts) = self
596 .write_to_txns(Some(&append_metric), |ts, advance_to| {
597 self.table_write_handle
598 .append(ts, advance_to, appends.clone())
599 })
600 .await
601 else {
602 return ControlFlow::Break(());
604 };
605 let timestamp = write_ts.timestamp;
606
607 let modified_tables: Vec<_> = appends
608 .iter()
609 .filter_map(|(id, updates)| {
610 (id.is_user() && !updates.iter().all(|u| u.is_empty())).then_some(id)
611 })
612 .collect();
613 if !modified_tables.is_empty() {
614 info!(
615 "Appending to tables, {modified_tables:?}, at {timestamp}, advancing to {}",
616 write_ts.advance_to
617 );
618 }
619
620 drop(permits);
622
623 for notify in notifies {
624 let _ = notify.send(());
625 }
626
627 if self
630 .internal_cmd_tx
631 .send(Message::GroupCommitApplied {
632 responses,
633 statement_logging_ids,
634 internal_results,
635 write_ts: timestamp,
636 })
637 .is_err()
638 {
639 warn!("coordinator shut down before a group commit could be finalized");
640 }
641
642 ControlFlow::Continue(deferred_cmd)
643 }
644}
645
646pub(crate) fn spawn_group_committer(
647 rx: mpsc::UnboundedReceiver<TableWriteCmd>,
648 oracle: Arc<dyn TimestampOracle<Timestamp> + Send + Sync>,
649 table_write_handle: Arc<dyn TableWriteHandle>,
650 catalog_upper: CatalogUpperHandle,
651 internal_cmd_tx: mpsc::UnboundedSender<Message>,
652 now: NowFn,
653 metrics: Metrics,
654 dyncfgs: &ConfigSet,
655) {
656 let committer = GroupCommitter {
657 rx,
658 oracle,
659 table_write_handle,
660 catalog_upper,
661 internal_cmd_tx,
662 now,
663 metrics,
664 max_attempts: GROUP_COMMIT_MAX_ATTEMPTS.handle(dyncfgs),
665 };
666 task::spawn(|| "group_committer", committer.run());
667}
668
669impl Coordinator {
670 pub(crate) fn trigger_group_commit(&mut self) {
672 self.group_commit_tx.notify();
673 self.advance_timelines_interval.reset();
677 }
678
679 #[instrument(name = "coord::stage_group_commit")]
683 pub(crate) fn stage_group_commit(&mut self, permit: Option<GroupCommitPermit>) {
684 let pending_writes = std::mem::take(&mut self.pending_writes);
685 let contains_internal_system_write = pending_writes
686 .iter()
687 .any(|write| write.is_internal_system());
688
689 let mut appends: BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>> = BTreeMap::new();
690 let mut responses = Vec::with_capacity(pending_writes.len());
691 let mut statement_logging_ids = Vec::new();
692 let mut notifies = Vec::new();
693 let mut internal_results = Vec::new();
694
695 for pending_write in pending_writes {
696 match pending_write {
697 PendingWriteTxn::User {
698 span: _,
699 writes,
700 responder:
701 UserWriteResponder::Session(PendingTxn {
702 ctx,
703 response,
704 action,
705 }),
706 } => {
707 if let Some(id) = Self::stale_write_target(self.catalog(), &writes) {
714 let err = AdapterError::ConcurrentDependencyMutation {
715 dependency_id: id.to_string(),
716 };
717 let (ctx, result) = CompletedClientTransmitter::new(
718 ctx,
719 Err(err),
720 EndTransactionAction::Rollback,
721 )
722 .finalize();
723 ctx.retire(result);
724 continue;
725 }
726
727 for (id, table_data) in writes {
728 if self.catalog().try_get_entry(&id).is_some() {
734 appends.entry(id).or_default().extend(table_data);
735 }
736 }
737 if let Some(id) = ctx.extra().contents() {
738 statement_logging_ids.push(id);
739 }
740
741 responses.push(CompletedClientTransmitter::new(ctx, response, action));
742 }
743 PendingWriteTxn::User {
744 span: _,
745 writes,
746 responder: UserWriteResponder::Internal { target, result, .. },
747 } => {
748 let current_global_id = self
749 .catalog()
750 .try_get_entry(&target.item_id)
751 .map(|entry| entry.latest_global_id());
752 if current_global_id != Some(target.global_id) {
753 result.send(WriteResult::TargetChanged);
754 continue;
755 }
756 assert!(
762 writes.keys().all(|id| *id == target.item_id),
763 "frontend write for {:?} carries other tables: {:?}",
764 target.item_id,
765 writes.keys().collect::<Vec<_>>(),
766 );
767 appends
768 .entry(target.item_id)
769 .or_default()
770 .extend(writes.into_values().flatten());
771 internal_results.push(result);
772 }
773 PendingWriteTxn::System { updates, source } => {
774 for update in updates {
775 appends.entry(update.id).or_default().push(update.data);
776 }
777 match source {
779 BuiltinTableUpdateSource::Internal(tx)
780 | BuiltinTableUpdateSource::Background(tx) => notifies.push(tx),
781 }
782 }
783 }
784 }
785
786 let mut all_appends = Vec::with_capacity(appends.len());
789 for (item_id, table_data) in appends.into_iter() {
790 let mut all_rows = Vec::new();
791 let mut all_data = Vec::new();
792 for data in table_data {
793 match data {
794 TableData::Rows(rows) => all_rows.extend(rows),
795 TableData::Batches(_) => all_data.push(data),
796 }
797 }
798 differential_dataflow::consolidation::consolidate(&mut all_rows);
799 all_data.push(TableData::Rows(all_rows));
800
801 all_appends.push((item_id, all_data));
803 }
804
805 let appends: Vec<_> = all_appends
806 .into_iter()
807 .map(|(id, updates)| {
808 let gid = self.catalog().get_entry(&id).latest_global_id();
809 (gid, updates)
810 })
811 .collect();
812
813 let request = GroupCommitRequest {
815 appends,
816 responses,
817 statement_logging_ids,
818 notifies,
819 internal_results,
820 permits: permit.into_iter().collect(),
821 contains_internal_system_write,
822 span: Span::current(),
823 };
824 if self
825 .group_committer_tx
826 .send(TableWriteCmd::GroupCommit(request))
827 .is_err()
828 {
829 warn!("group committer task gone, dropping staged group commit");
831 }
832 }
833
834 fn stale_write_target(
844 catalog: &Catalog,
845 writes: &BTreeMap<CatalogItemId, SmallVec<[TableData; 1]>>,
846 ) -> Option<CatalogItemId> {
847 writes.iter().find_map(|(id, table_data)| {
848 let entry = catalog.try_get_entry(id)?;
850 let arity = entry
851 .relation_desc_latest()
852 .expect("write target is a table")
853 .arity();
854 let stale = table_data.iter().any(|data| match data {
855 TableData::Rows(rows) => rows
856 .first()
857 .is_some_and(|(row, _)| row.iter().count() != arity),
858 TableData::Batches(_) => false,
859 });
860 stale.then_some(*id)
861 })
862 }
863
864 pub(crate) async fn register_tables_via_committer(
866 &self,
867 tables: Vec<TableRegistration>,
868 ) -> Timestamp {
869 let (tx, rx) = oneshot::channel();
870 if self
871 .group_committer_tx
872 .send(TableWriteCmd::Register { tables, result: tx })
873 .is_err()
874 {
875 halt!("group committer terminated before a table registration could be submitted");
876 }
877 match rx.await {
878 Ok(ts) => ts,
879 Err(_) => halt!("group committer terminated with a table registration outstanding"),
880 }
881 }
882
883 pub(crate) async fn forget_tables_via_committer(&self, ids: Vec<GlobalId>) -> Timestamp {
885 let (tx, rx) = oneshot::channel();
886 if self
887 .group_committer_tx
888 .send(TableWriteCmd::Forget { ids, result: tx })
889 .is_err()
890 {
891 halt!("group committer terminated before a table forget could be submitted");
892 }
893 match rx.await {
894 Ok(ts) => ts,
895 Err(_) => halt!("group committer terminated with a table forget outstanding"),
896 }
897 }
898
899 pub(crate) fn submit_write(&mut self, pending_write_txn: PendingWriteTxn) {
901 if self.controller.read_only() {
902 panic!(
903 "attempting table write in read-only mode: {:?}",
904 pending_write_txn
905 );
906 }
907 self.pending_writes.push(pending_write_txn);
908 self.trigger_group_commit();
909 }
910
911 pub(crate) fn builtin_table_update<'a>(&'a mut self) -> BuiltinTableAppend<'a> {
913 BuiltinTableAppend { coord: self }
914 }
915}
916
917pub struct BuiltinTableAppend<'a> {
919 coord: &'a mut Coordinator,
920}
921
922pub type BuiltinTableAppendNotify = Pin<Box<dyn Future<Output = ()> + Send + Sync + 'static>>;
933
934pub struct BuiltinTableAppendCompletion {
936 notify: Option<BuiltinTableAppendNotify>,
937}
938
939impl BuiltinTableAppendCompletion {
940 pub fn new(notify: BuiltinTableAppendNotify) -> Self {
941 Self {
942 notify: Some(notify),
943 }
944 }
945
946 pub fn completed() -> Self {
948 Self { notify: None }
949 }
950
951 pub fn into_notify(self) -> Option<BuiltinTableAppendNotify> {
952 self.notify
953 }
954}
955
956impl<'a> BuiltinTableAppend<'a> {
957 pub fn background(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
968 if self.coord.controller.read_only() {
969 self.coord
970 .buffered_builtin_table_updates
971 .as_mut()
972 .expect("in read-only mode")
973 .append(&mut updates);
974
975 return Box::pin(futures::future::ready(()));
976 }
977
978 let (tx, rx) = oneshot::channel();
979 self.coord.pending_writes.push(PendingWriteTxn::System {
980 updates,
981 source: BuiltinTableUpdateSource::Background(tx),
982 });
983
984 Box::pin(rx.map(|_| ()))
985 }
986
987 pub fn defer(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
996 if self.coord.controller.read_only() {
997 self.coord
998 .buffered_builtin_table_updates
999 .as_mut()
1000 .expect("in read-only mode")
1001 .append(&mut updates);
1002
1003 return Box::pin(futures::future::ready(()));
1004 }
1005
1006 let (tx, rx) = oneshot::channel();
1007 self.coord.pending_writes.push(PendingWriteTxn::System {
1008 updates,
1009 source: BuiltinTableUpdateSource::Internal(tx),
1010 });
1011 self.coord.trigger_group_commit();
1012
1013 Box::pin(rx.map(|_| ()))
1014 }
1015
1016 pub fn execute(self, mut updates: Vec<BuiltinTableUpdate>) -> BuiltinTableAppendNotify {
1020 if self.coord.controller.read_only() {
1021 self.coord
1022 .buffered_builtin_table_updates
1023 .as_mut()
1024 .expect("in read-only mode")
1025 .append(&mut updates);
1026
1027 return Box::pin(futures::future::ready(()));
1028 }
1029
1030 let (tx, rx) = oneshot::channel();
1031
1032 self.coord.pending_writes.push(PendingWriteTxn::System {
1035 updates,
1036 source: BuiltinTableUpdateSource::Internal(tx),
1037 });
1038 self.coord.stage_group_commit(None);
1039
1040 self.coord.advance_timelines_interval.reset();
1042
1043 Box::pin(rx.map(|_| ()))
1044 }
1045}
1046
1047pub fn notifier() -> (GroupCommitNotifier, GroupCommitWaiter) {
1050 let notify = Arc::new(Notify::new());
1051 let in_progress = Arc::new(Semaphore::new(1));
1052
1053 let notifier = GroupCommitNotifier {
1054 notify: Arc::clone(¬ify),
1055 };
1056 let waiter = GroupCommitWaiter {
1057 notify,
1058 in_progress,
1059 };
1060
1061 (notifier, waiter)
1062}
1063
1064#[derive(Debug, Clone)]
1067pub struct GroupCommitNotifier {
1068 notify: Arc<Notify>,
1070}
1071
1072impl GroupCommitNotifier {
1073 pub fn notify(&self) {
1075 self.notify.notify_one()
1076 }
1077}
1078
1079#[derive(Debug)]
1082pub struct GroupCommitWaiter {
1083 notify: Arc<Notify>,
1085 in_progress: Arc<Semaphore>,
1087}
1088static_assertions::assert_not_impl_all!(GroupCommitWaiter: Clone);
1089
1090impl GroupCommitWaiter {
1091 pub async fn ready(&self) -> GroupCommitPermit {
1102 let permit = Semaphore::acquire_owned(Arc::clone(&self.in_progress))
1103 .await
1104 .expect("semaphore should not close");
1105
1106 self.notify.notified().await;
1109
1110 GroupCommitPermit {
1111 _permit: Some(permit),
1112 }
1113 }
1114}
1115
1116#[derive(Debug)]
1121pub struct GroupCommitPermit {
1122 _permit: Option<OwnedSemaphorePermit>,
1126}
1127
1128pub(crate) fn waiting_on_startup_appends(
1141 catalog: &Catalog,
1142 session: &mut Session,
1143 plan: &Plan,
1144) -> Option<BoxFuture<'static, ()>> {
1145 let depends_on = match plan {
1149 Plan::Select(plan) => plan.source.depends_on(),
1150 Plan::ReadThenWrite(plan) => plan.selection.depends_on(),
1151 Plan::ShowColumns(plan) => plan.select_plan.source.depends_on(),
1152 Plan::Subscribe(plan) => plan.from.depends_on(),
1153 Plan::ExplainPlan(ExplainPlanPlan {
1154 explainee: Explainee::Statement(ExplaineeStatement::Select { plan, .. }),
1155 ..
1156 }) => plan.source.depends_on(),
1157 Plan::ExplainTimestamp(ExplainTimestampPlan { raw_plan, .. }) => raw_plan.depends_on(),
1158 Plan::CreateConnection(_)
1159 | Plan::CreateDatabase(_)
1160 | Plan::CreateSchema(_)
1161 | Plan::CreateRole(_)
1162 | Plan::CreateNetworkPolicy(_)
1163 | Plan::CreateCluster(_)
1164 | Plan::CreateClusterReplica(_)
1165 | Plan::CreateSource(_)
1166 | Plan::CreateSources(_)
1167 | Plan::CreateSecret(_)
1168 | Plan::CreateSink(_)
1169 | Plan::CreateTable(_)
1170 | Plan::CreateView(_)
1171 | Plan::CreateMaterializedView(_)
1172 | Plan::CreateIndex(_)
1173 | Plan::CreateMetricSink(_)
1174 | Plan::CreateType(_)
1175 | Plan::Comment(_)
1176 | Plan::DiscardTemp
1177 | Plan::DiscardAll
1178 | Plan::DropObjects(_)
1179 | Plan::DropOwned(_)
1180 | Plan::EmptyQuery
1181 | Plan::ShowAllVariables
1182 | Plan::ShowCreate(_)
1183 | Plan::ShowVariable(_)
1184 | Plan::InspectShard(_)
1185 | Plan::SetVariable(_)
1186 | Plan::ResetVariable(_)
1187 | Plan::SetTransaction(_)
1188 | Plan::StartTransaction(_)
1189 | Plan::CommitTransaction(_)
1190 | Plan::AbortTransaction(_)
1191 | Plan::CopyFrom(_)
1192 | Plan::CopyTo(_)
1193 | Plan::ExplainPlan(_)
1194 | Plan::ExplainPushdown(_)
1195 | Plan::ExplainSinkSchema(_)
1196 | Plan::Insert(_)
1197 | Plan::AlterNetworkPolicy(_)
1198 | Plan::AlterNoop(_)
1199 | Plan::AlterClusterRename(_)
1200 | Plan::AlterClusterSwap(_)
1201 | Plan::AlterClusterReplicaRename(_)
1202 | Plan::AlterCluster(_)
1203 | Plan::AlterConnection(_)
1204 | Plan::AlterSource(_)
1205 | Plan::AlterSetCluster(_)
1206 | Plan::AlterItemRename(_)
1207 | Plan::AlterRetainHistory(_)
1208 | Plan::AlterSourceTimestampInterval(_)
1209 | Plan::AlterSchemaRename(_)
1210 | Plan::AlterSchemaSwap(_)
1211 | Plan::AlterSecret(_)
1212 | Plan::AlterSink(_)
1213 | Plan::AlterSystemSet(_)
1214 | Plan::AlterSystemReset(_)
1215 | Plan::AlterSystemResetAll(_)
1216 | Plan::AlterRole(_)
1217 | Plan::AlterOwner(_)
1218 | Plan::AlterTableAddColumn(_)
1219 | Plan::AlterMaterializedViewApplyReplacement(_)
1220 | Plan::Declare(_)
1221 | Plan::Fetch(_)
1222 | Plan::Close(_)
1223 | Plan::Prepare(_)
1224 | Plan::Execute(_)
1225 | Plan::Deallocate(_)
1226 | Plan::Raise(_)
1227 | Plan::GrantRole(_)
1228 | Plan::RevokeRole(_)
1229 | Plan::GrantPrivileges(_)
1230 | Plan::RevokePrivileges(_)
1231 | Plan::AlterDefaultPrivileges(_)
1232 | Plan::ReassignOwned(_)
1233 | Plan::ValidateConnection(_)
1234 | Plan::SideEffectingFunc(_) => BTreeSet::default(),
1235 };
1236 let depends_on_required_id = REQUIRED_BUILTIN_TABLES
1237 .iter()
1238 .map(|table| catalog.resolve_builtin_table(&**table))
1239 .any(|id| {
1240 catalog
1241 .get_global_ids(&id)
1242 .any(|gid| depends_on.contains(&gid))
1243 });
1244
1245 if !depends_on_required_id {
1248 return None;
1249 }
1250
1251 match session.clear_builtin_table_updates() {
1258 Some(wait_future) => Some(wait_future.boxed()),
1259 None => None,
1260 }
1261}
1262
1263#[cfg(test)]
1264mod tests {
1265 use std::sync::Mutex;
1266 use std::sync::atomic::{AtomicUsize, Ordering};
1267
1268 use async_trait::async_trait;
1269 use mz_ore::metrics::MetricsRegistry;
1270 use mz_ore::now::SYSTEM_TIME;
1271 use timely::progress::Antichain;
1272
1273 use super::*;
1274 use crate::catalog::Catalog;
1275
1276 #[mz_ore::test(tokio::test)]
1277 async fn internal_write_responder_reports_indeterminate_on_drop() {
1278 let (tx, rx) = oneshot::channel();
1279 drop(InternalWriteResponder::new(tx));
1280 assert!(matches!(rx.await, Ok(WriteResult::Indeterminate)));
1281 }
1282
1283 #[derive(Debug, Default)]
1284 struct MemTimestampOracle {
1285 read_write_ts: Mutex<(Timestamp, Timestamp)>,
1286 apply_writes: AtomicUsize,
1287 }
1288
1289 impl MemTimestampOracle {
1290 fn starting_at(ts: Timestamp) -> Self {
1291 Self {
1292 read_write_ts: Mutex::new((ts, ts)),
1293 apply_writes: AtomicUsize::new(0),
1294 }
1295 }
1296 }
1297
1298 #[async_trait]
1299 impl TimestampOracle<Timestamp> for MemTimestampOracle {
1300 async fn write_ts(&self) -> WriteTimestamp {
1301 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1302 let new_write_ts = std::cmp::max(*read_ts, *write_ts).step_forward();
1303 *write_ts = new_write_ts;
1304 WriteTimestamp {
1305 timestamp: new_write_ts,
1306 advance_to: new_write_ts.step_forward(),
1307 }
1308 }
1309
1310 async fn peek_write_ts(&self) -> Timestamp {
1311 let (_, write_ts) = &*self.read_write_ts.lock().expect("lock poisoned");
1312 *write_ts
1313 }
1314
1315 async fn read_ts(&self) -> Timestamp {
1316 let (read_ts, _) = &*self.read_write_ts.lock().expect("lock poisoned");
1317 *read_ts
1318 }
1319
1320 async fn apply_write(&self, lower_bound: Timestamp) {
1321 self.apply_writes.fetch_add(1, Ordering::SeqCst);
1322 let (read_ts, write_ts) = &mut *self.read_write_ts.lock().expect("lock poisoned");
1323 *read_ts = std::cmp::max(*read_ts, lower_bound);
1324 *write_ts = std::cmp::max(*read_ts, *write_ts);
1325 }
1326 }
1327
1328 #[derive(Debug)]
1329 struct ConflictingTableWriteHandle {
1330 conflicts: usize,
1331 calls: AtomicUsize,
1332 write_timestamps: Mutex<Vec<Timestamp>>,
1333 }
1334
1335 impl ConflictingTableWriteHandle {
1336 fn new(conflicts: usize) -> Self {
1337 Self {
1338 conflicts,
1339 calls: AtomicUsize::new(0),
1340 write_timestamps: Mutex::new(Vec::new()),
1341 }
1342 }
1343
1344 fn respond(&self, write_ts: Timestamp) -> oneshot::Receiver<Result<(), StorageError>> {
1345 let call = self.calls.fetch_add(1, Ordering::SeqCst);
1346 self.write_timestamps
1347 .lock()
1348 .expect("lock poisoned")
1349 .push(write_ts);
1350 let (tx, rx) = oneshot::channel();
1351 let result = if call < self.conflicts {
1352 Err(StorageError::InvalidUppers(vec![
1353 mz_storage_types::controller::InvalidUpper {
1354 id: GlobalId::User(1),
1355 current_upper: Antichain::from_elem(write_ts.step_forward()),
1356 },
1357 ]))
1358 } else {
1359 Ok(())
1360 };
1361 tx.send(result).expect("receiver still in scope");
1362 rx
1363 }
1364 }
1365
1366 impl TableWriteHandle for ConflictingTableWriteHandle {
1367 fn append(
1368 &self,
1369 write_ts: Timestamp,
1370 _advance_to: Timestamp,
1371 _commands: Vec<(GlobalId, Vec<TableData>)>,
1372 ) -> oneshot::Receiver<Result<(), StorageError>> {
1373 self.respond(write_ts)
1374 }
1375
1376 fn register(
1377 &self,
1378 register_ts: Timestamp,
1379 _tables: Vec<TableRegistration>,
1380 ) -> oneshot::Receiver<Result<(), StorageError>> {
1381 self.respond(register_ts)
1382 }
1383
1384 fn forget(
1385 &self,
1386 forget_ts: Timestamp,
1387 _ids: Vec<GlobalId>,
1388 ) -> oneshot::Receiver<Result<(), StorageError>> {
1389 self.respond(forget_ts)
1390 }
1391 }
1392
1393 #[mz_ore::test(tokio::test)]
1394 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_conflict_retry() {
1396 Catalog::with_debug(|catalog| async move {
1397 let initial_upper = catalog.current_upper().await;
1399 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1400 let handle = Arc::new(ConflictingTableWriteHandle::new(1));
1401 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1402 Arc::<MemTimestampOracle>::clone(&oracle);
1403 let handle_dyn: Arc<dyn TableWriteHandle> =
1404 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1405 let (_tx, rx) = mpsc::unbounded_channel();
1406 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1407 let committer = GroupCommitter {
1408 rx,
1409 oracle: oracle_dyn,
1410 table_write_handle: handle_dyn,
1411 catalog_upper: catalog.upper_handle(),
1412 internal_cmd_tx,
1413 now: SYSTEM_TIME.clone(),
1414 metrics: Metrics::register_into(&MetricsRegistry::new()),
1415 max_attempts: ConfigValHandle::disconnected(2),
1416 };
1417
1418 let write_ts = committer
1419 .write_to_txns(None, |ts, advance_to| {
1420 assert_eq!(advance_to, ts.step_forward());
1421 handle.append(ts, advance_to, Vec::new())
1422 })
1423 .await
1424 .expect("worker stays alive in this test");
1425
1426 let attempts = handle
1427 .write_timestamps
1428 .lock()
1429 .expect("lock poisoned")
1430 .clone();
1431 assert_eq!(attempts.len(), 2, "one conflict, one success");
1432 assert!(
1433 attempts[0] > initial_upper,
1434 "attempts start past the initial catalog upper: {attempts:?} vs {initial_upper}"
1435 );
1436 assert!(
1437 attempts[0] < attempts[1],
1438 "retry must use a fresh, larger timestamp: {attempts:?}"
1439 );
1440 assert_eq!(write_ts.timestamp, attempts[1]);
1441
1442 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1443 assert_eq!(oracle.read_ts().await, write_ts.timestamp);
1444
1445 let catalog_upper = catalog.current_upper().await;
1446 assert!(
1447 catalog_upper >= write_ts.advance_to,
1448 "catalog upper {catalog_upper} must cover the write's advance_to {}",
1449 write_ts.advance_to
1450 );
1451
1452 catalog.expire().await;
1453 })
1454 .await;
1455 }
1456
1457 #[mz_ore::test(tokio::test)]
1458 #[cfg_attr(miri, ignore)] async fn test_write_to_txns_zero_limit_allows_one_attempt() {
1460 Catalog::with_debug(|catalog| async move {
1461 let initial_upper = catalog.current_upper().await;
1462 let oracle = Arc::new(MemTimestampOracle::starting_at(initial_upper));
1463 let handle = Arc::new(ConflictingTableWriteHandle::new(0));
1464 let oracle_dyn: Arc<dyn TimestampOracle<Timestamp> + Send + Sync> =
1465 Arc::<MemTimestampOracle>::clone(&oracle);
1466 let handle_dyn: Arc<dyn TableWriteHandle> =
1467 Arc::<ConflictingTableWriteHandle>::clone(&handle);
1468 let (_tx, rx) = mpsc::unbounded_channel();
1469 let (internal_cmd_tx, _internal_cmd_rx) = mpsc::unbounded_channel();
1470 let committer = GroupCommitter {
1471 rx,
1472 oracle: oracle_dyn,
1473 table_write_handle: handle_dyn,
1474 catalog_upper: catalog.upper_handle(),
1475 internal_cmd_tx,
1476 now: SYSTEM_TIME.clone(),
1477 metrics: Metrics::register_into(&MetricsRegistry::new()),
1478 max_attempts: ConfigValHandle::disconnected(0),
1479 };
1480
1481 let write_ts = committer
1482 .write_to_txns(None, |ts, advance_to| {
1483 handle.append(ts, advance_to, Vec::new())
1484 })
1485 .await
1486 .expect("zero limit permits one successful attempt");
1487
1488 let attempts = handle
1489 .write_timestamps
1490 .lock()
1491 .expect("lock poisoned")
1492 .clone();
1493 assert_eq!(attempts, vec![write_ts.timestamp]);
1494 assert_eq!(oracle.apply_writes.load(Ordering::SeqCst), 1);
1495
1496 catalog.expire().await;
1497 })
1498 .await;
1499 }
1500}