Merge pull request #24930 from mjibson/refactor-peek-timestamp

adapter: refactor peek timestamps

src/adapter/src/coord.rs

@@ -365,9 +365,7 @@ pub enum RealTimeRecencyContext {
         timeline_context: TimelineContext,
         oracle_read_ts: Option<Timestamp>,
         source_ids: BTreeSet<GlobalId>,
-        in_immediate_multi_stmt_txn: bool,
         optimizer: optimize::peek::Optimizer,
-        global_mir_plan: optimize::peek::GlobalMirPlan,
         peek_ctx: PeekContext,
     },
 }
@@ -392,9 +390,10 @@ pub enum PeekContext {
 #[derive(Debug)]
 pub enum PeekStage {
     Validate(PeekStageValidate),
-    Timestamp(PeekStageTimestamp),
-    OptimizeMir(PeekStageOptimizeMir),
+    LinearizeTimestamp(PeekStageLinearizeTimestamp),
     RealTimeRecency(PeekStageRealTimeRecency),
+    TimestampReadHold(PeekStageTimestampReadHold),
+    OptimizeMir(PeekStageOptimizeMir),
     OptimizeLir(PeekStageOptimizeLir),
     Finish(PeekStageFinish),
     Explain(PeekStageExplain),
@@ -404,9 +403,10 @@ impl PeekStage {
     fn validity(&mut self) -> Option<&mut PlanValidity> {
         match self {
             PeekStage::Validate(_) => None,
-            PeekStage::Timestamp(PeekStageTimestamp { validity, .. })
-            | PeekStage::OptimizeMir(PeekStageOptimizeMir { validity, .. })
+            PeekStage::LinearizeTimestamp(PeekStageLinearizeTimestamp { validity, .. })
             | PeekStage::RealTimeRecency(PeekStageRealTimeRecency { validity, .. })
+            | PeekStage::TimestampReadHold(PeekStageTimestampReadHold { validity, .. })
+            | PeekStage::OptimizeMir(PeekStageOptimizeMir { validity, .. })
             | PeekStage::OptimizeLir(PeekStageOptimizeLir { validity, .. })
             | PeekStage::Finish(PeekStageFinish { validity, .. }) => Some(validity),
             PeekStage::Explain(PeekStageExplain { validity, .. }) => Some(validity),
@@ -418,66 +418,67 @@ impl PeekStage {
 pub struct PeekStageValidate {
     plan: mz_sql::plan::SelectPlan,
     target_cluster: TargetCluster,
-    /// Context from where this peek initiated.
     peek_ctx: PeekContext,
 }
 
 #[derive(Debug)]
-pub struct PeekStageTimestamp {
+pub struct PeekStageLinearizeTimestamp {
     validity: PlanValidity,
     plan: mz_sql::plan::SelectPlan,
     source_ids: BTreeSet<GlobalId>,
     target_replica: Option<ReplicaId>,
     timeline_context: TimelineContext,
-    in_immediate_multi_stmt_txn: bool,
     optimizer: optimize::peek::Optimizer,
-    /// Context from where this peek initiated.
     peek_ctx: PeekContext,
 }
 
 #[derive(Debug)]
-pub struct PeekStageOptimizeMir {
+pub struct PeekStageRealTimeRecency {
     validity: PlanValidity,
     plan: mz_sql::plan::SelectPlan,
     source_ids: BTreeSet<GlobalId>,
     target_replica: Option<ReplicaId>,
     timeline_context: TimelineContext,
     oracle_read_ts: Option<Timestamp>,
-    in_immediate_multi_stmt_txn: bool,
     optimizer: optimize::peek::Optimizer,
-    /// Context from where this peek initiated.
     peek_ctx: PeekContext,
 }
 
 #[derive(Debug)]
-pub struct PeekStageRealTimeRecency {
+pub struct PeekStageTimestampReadHold {
     validity: PlanValidity,
     plan: mz_sql::plan::SelectPlan,
     source_ids: BTreeSet<GlobalId>,
-    id_bundle: CollectionIdBundle,
     target_replica: Option<ReplicaId>,
     timeline_context: TimelineContext,
     oracle_read_ts: Option<Timestamp>,
-    in_immediate_multi_stmt_txn: bool,
+    real_time_recency_ts: Option<mz_repr::Timestamp>,
+    optimizer: optimize::peek::Optimizer,
+    peek_ctx: PeekContext,
+}
+
+#[derive(Debug)]
+pub struct PeekStageOptimizeMir {
+    validity: PlanValidity,
+    plan: mz_sql::plan::SelectPlan,
+    source_ids: BTreeSet<GlobalId>,
+    id_bundle: CollectionIdBundle,
+    target_replica: Option<ReplicaId>,
+    determination: TimestampDetermination<mz_repr::Timestamp>,
     optimizer: optimize::peek::Optimizer,
-    global_mir_plan: optimize::peek::GlobalMirPlan,
-    /// Context from where this peek initiated.
     peek_ctx: PeekContext,
 }
 
 #[derive(Debug)]
 pub struct PeekStageOptimizeLir {
     validity: PlanValidity,
     plan: mz_sql::plan::SelectPlan,
-    id_bundle: Option<CollectionIdBundle>,
+    id_bundle: CollectionIdBundle,
     target_replica: Option<ReplicaId>,
-    timeline_context: TimelineContext,
-    oracle_read_ts: Option<Timestamp>,
+    determination: TimestampDetermination<mz_repr::Timestamp>,
     source_ids: BTreeSet<GlobalId>,
-    real_time_recency_ts: Option<mz_repr::Timestamp>,
     optimizer: optimize::peek::Optimizer,
     global_mir_plan: optimize::peek::GlobalMirPlan,
-    /// Context from where this peek initiated.
     peek_ctx: PeekContext,
 }
 
@@ -489,7 +490,6 @@ pub struct PeekStageFinish {
     target_replica: Option<ReplicaId>,
     source_ids: BTreeSet<GlobalId>,
     determination: TimestampDetermination<mz_repr::Timestamp>,
-    timestamp_context: TimestampContext<mz_repr::Timestamp>,
     optimizer: optimize::peek::Optimizer,
     global_lir_plan: optimize::peek::GlobalLirPlan,
 }

src/adapter/src/coord/message_handler.rs

@@ -34,7 +34,7 @@ use crate::command::Command;
 use crate::coord::appends::Deferred;
 use crate::coord::statement_logging::StatementLoggingId;
 use crate::coord::{
-    Coordinator, CreateConnectionValidationReady, Message, PeekStage, PeekStageOptimizeLir,
+    Coordinator, CreateConnectionValidationReady, Message, PeekStage, PeekStageTimestampReadHold,
     PendingReadTxn, PlanValidity, PurifiedStatementReady, RealTimeRecencyContext,
 };
 use crate::session::Session;
@@ -854,25 +854,21 @@ impl Coordinator {
                 timeline_context,
                 oracle_read_ts,
                 source_ids,
-                in_immediate_multi_stmt_txn: _,
                 optimizer,
-                global_mir_plan,
                 peek_ctx,
             } => {
                 self.execute_peek_stage(
                     ctx,
                     root_otel_ctx,
-                    PeekStage::OptimizeLir(PeekStageOptimizeLir {
+                    PeekStage::TimestampReadHold(PeekStageTimestampReadHold {
                         validity,
                         plan,
-                        id_bundle: None,
                         target_replica,
                         timeline_context,
                         oracle_read_ts,
                         source_ids,
                         real_time_recency_ts: Some(real_time_recency_ts),
                         optimizer,
-                        global_mir_plan,
                         peek_ctx,
                     }),
                 )

src/adapter/src/coord/sequencer/inner.rs

@@ -4209,7 +4209,7 @@ impl Coordinator {
         &self,
         session: &Session,
         source_ids: &BTreeSet<GlobalId>,
-        query_as_of: Antichain<Timestamp>,
+        query_as_of: &Antichain<Timestamp>,
         is_oneshot: bool,
     ) -> Result<Box<dyn mz_transform::StatisticsOracle>, AdapterError> {
         if !session.vars().enable_session_cardinality_estimates() {
@@ -4227,7 +4227,7 @@ impl Coordinator {
 
         let cached_stats = mz_ore::future::timeout(
             timeout,
-            CachedStatisticsOracle::new(source_ids, &query_as_of, self.controller.storage.as_ref()),
+            CachedStatisticsOracle::new(source_ids, query_as_of, self.controller.storage.as_ref()),
         )
         .await;