Merge pull request #25184 from ggevay/precise_acquire_read_holds2

Acquire precise read holds for `REFRESH AT`s

Author: ggevay

src/adapter/src/coord.rs

@@ -1253,7 +1253,9 @@ pub struct Coordinator {
     ///
     /// Upon completing a transaction, this timestamp should be removed from the holds
     /// in `self.read_capability[id]`, using the `release_read_holds` method.
-    txn_read_holds: BTreeMap<ConnectionId, read_policy::ReadHolds<Timestamp>>,
+    ///
+    /// We use a Vec because `ReadHolds` doesn't have a way of tracking multiplicity.
+    txn_read_holds: BTreeMap<ConnectionId, Vec<read_policy::ReadHolds<Timestamp>>>,
 
     /// Access to the peek fields should be restricted to methods in the [`peek`] API.
     /// A map from pending peek ids to the queue into which responses are sent, and

src/adapter/src/coord/command_handler.rs

@@ -843,7 +843,8 @@ impl Coordinator {
             }
 
             if acquire_read_holds {
-                self.acquire_read_holds_auto_cleanup(session, timestamp, &ids);
+                self.acquire_read_holds_auto_cleanup(session, timestamp, &ids, false)
+                    .expect("precise==false, so acquiring read holds always succeeds");
             }
 
             Ok(Some(timestamp))

src/adapter/src/coord/id_bundle.rs

@@ -54,6 +54,49 @@ impl CollectionIdBundle {
         }
     }
 
+    /// Returns a new bundle with the identifiers that are present in both `self` and `other`.
+    pub fn intersection(&self, other: &CollectionIdBundle) -> CollectionIdBundle {
+        // Attend to storage ids.
+        let storage_ids = self
+            .storage_ids
+            .intersection(&other.storage_ids)
+            .cloned()
+            .collect();
+
+        // Intersect ComputeInstanceIds.
+        let self_compute_instances = self.compute_ids.keys().collect::<BTreeSet<_>>();
+        let other_compute_instances = other.compute_ids.keys().collect::<BTreeSet<_>>();
+        let compute_instances = self_compute_instances
+            .intersection(&other_compute_instances)
+            .cloned();
+
+        // For each ComputeInstanceId, intersect `self` with `other`.
+        let compute_ids = compute_instances
+            .map(|compute_instance_id| {
+                let self_compute_ids = self
+                    .compute_ids
+                    .get(compute_instance_id)
+                    .expect("id is in intersection, so should be found");
+                let other_compute_ids = other
+                    .compute_ids
+                    .get(compute_instance_id)
+                    .expect("id is in intersection, so should be found");
+                (
+                    compute_instance_id.clone(),
+                    self_compute_ids
+                        .intersection(other_compute_ids)
+                        .cloned()
+                        .collect(),
+                )
+            })
+            .collect();
+
+        CollectionIdBundle {
+            storage_ids,
+            compute_ids,
+        }
+    }
+
     /// Extends a `CollectionIdBundle` with the contents of another `CollectionIdBundle`.
     pub fn extend(&mut self, other: &CollectionIdBundle) {
         self.storage_ids.extend(&other.storage_ids);

src/adapter/src/coord/read_policy.rs

@@ -110,9 +110,15 @@ impl<T: Eq + Hash + Ord> ReadHolds<T> {
     }
 
     /// Extends a `ReadHolds` with the contents of another `ReadHolds`.
-    pub fn extend(&mut self, other: ReadHolds<T>) {
-        for (time, id_bundle) in other.holds {
-            self.holds.entry(time).or_default().extend(&id_bundle);
+    /// Asserts that the newly added read holds don't coincide with any of the existing read holds in self.
+    pub fn extend_with_new(&mut self, other: ReadHolds<T>) {
+        for (time, other_id_bundle) in other.holds {
+            let self_id_bundle = self.holds.entry(time).or_default();
+            assert!(
+                self_id_bundle.intersection(&other_id_bundle).is_empty(),
+                "extend_with_new encountered duplicate read holds",
+            );
+            self_id_bundle.extend(&other_id_bundle);
         }
     }
 
@@ -214,7 +220,7 @@ impl crate::coord::Coordinator {
                             .or_default()
                             .extend(id_bundle);
                     }
-                    read_holds.extend(new_read_holds);
+                    read_holds.extend_with_new(new_read_holds);
                 }
                 TimelineContext::TimestampIndependent | TimelineContext::TimestampDependent => {
                     id_bundles.entry(None).or_default().extend(&id_bundle);
@@ -382,14 +388,34 @@ impl crate::coord::Coordinator {
 
     /// Attempt to acquire read holds on the indicated collections at the indicated `time`.
     ///
-    /// If we are unable to acquire a read hold at the provided `time` for a specific id, then we
-    /// will acquire a read hold at the lowest possible time for that id.
+    /// If we are unable to acquire a read hold at the provided `time` for a specific id, then
+    /// depending on the `precise` argument, we either fall back to acquiring a read hold at
+    /// the lowest possible time for that id, or return an error. The returned error contains
+    /// those collection sinces that were later than the specified time.
     pub(crate) fn acquire_read_holds(
         &mut self,
-        time: mz_repr::Timestamp,
+        time: Timestamp,
         id_bundle: &CollectionIdBundle,
-    ) -> ReadHolds<mz_repr::Timestamp> {
+        precise: bool,
+    ) -> Result<ReadHolds<Timestamp>, Vec<(Antichain<Timestamp>, CollectionIdBundle)>> {
         let read_holds = self.initialize_read_holds(time, id_bundle);
+        if precise {
+            // If we are not able to acquire read holds precisely at the specified time (only later), then error out.
+            let too_late = read_holds
+                .holds
+                .iter()
+                .filter_map(|(antichain, ids)| {
+                    if antichain.iter().all(|hold_time| *hold_time == time) {
+                        None
+                    } else {
+                        Some((antichain.clone(), ids.clone()))
+                    }
+                })
+                .collect_vec();
+            if !too_late.is_empty() {
+                return Err(too_late);
+            }
+        }
         // Update STORAGE read policies.
         let mut policy_changes = Vec::new();
         for (time, id) in read_holds.storage_ids() {
@@ -418,26 +444,30 @@ impl crate::coord::Coordinator {
                 .unwrap_or_terminate("cannot fail to set read policy");
         }
 
-        read_holds
+        Ok(read_holds)
     }
 
     /// Attempt to acquire read holds on the indicated collections at the indicated `time`.
     /// This is similar to [Self::acquire_read_holds], but instead of returning the read holds,
     /// it arranges for them to be automatically released at the end of the transaction.
     ///
-    /// If we are unable to acquire a read hold at the provided `time` for a specific id, then we
-    /// will acquire a read hold at the lowest possible time for that id.
+    /// If we are unable to acquire a read hold at the provided `time` for a specific id, then
+    /// depending on the `precise` argument, we either fall back to acquiring a read hold at
+    /// the lowest possible time for that id, or return an error. The returned error contains
+    /// those collection sinces that were later than the specified time.
     pub(crate) fn acquire_read_holds_auto_cleanup(
         &mut self,
         session: &Session,
         time: Timestamp,
         id_bundle: &CollectionIdBundle,
-    ) {
-        let read_holds = self.acquire_read_holds(time, id_bundle);
+        precise: bool,
+    ) -> Result<(), Vec<(Antichain<Timestamp>, CollectionIdBundle)>> {
+        let read_holds = self.acquire_read_holds(time, id_bundle, precise)?;
         self.txn_read_holds
             .entry(session.conn_id().clone())
-            .or_insert_with(ReadHolds::new)
-            .extend(read_holds);
+            .or_insert_with(Vec::new)
+            .push(read_holds);
+        Ok(())
     }
 
     /// Attempt to update the timestamp of the read holds on the indicated collections from the
@@ -446,10 +476,10 @@ impl crate::coord::Coordinator {
     /// If we are unable to update a read hold at the provided `time` for a specific id, then we
     /// leave it unchanged.
     ///
-    /// This method relies on a previous call to `acquire_read_holds` with the same
-    /// `read_holds` argument or a previous call to `update_read_hold` that returned
+    /// This method relies on a previous call to
+    /// `initialize_read_holds`, `acquire_read_holds`, or `update_read_hold` that returned
     /// `read_holds`, and its behavior will be erratic if called on anything else.
-    pub(super) fn update_read_hold(
+    pub(super) fn update_read_holds(
         &mut self,
         read_holds: ReadHolds<mz_repr::Timestamp>,
         new_time: mz_repr::Timestamp,
@@ -547,34 +577,46 @@ impl crate::coord::Coordinator {
 
         new_read_holds
     }
-    /// Release read holds on the indicated collections at the indicated times.
+
+    /// Release the given read holds.
     ///
-    /// This method relies on a previous call to `acquire_read_holds` with the same
-    /// argument, or a previous call to `update_read_hold` that returned
-    /// `read_holds`, and its behavior will be erratic if called on anything else,
+    /// This method relies on a previous call to
+    /// `initialize_read_holds`, `acquire_read_holds`, or `update_read_hold` that returned
+    /// `ReadHolds`, and its behavior will be erratic if called on anything else,
     /// or if called more than once on the same bundle of read holds.
-    pub(super) fn release_read_hold(&mut self, read_holds: &ReadHolds<mz_repr::Timestamp>) {
+    pub(super) fn release_read_holds(&mut self, read_holdses: Vec<ReadHolds<Timestamp>>) {
         // Update STORAGE read policies.
-        let mut policy_changes = Vec::new();
-        for (time, id) in read_holds.storage_ids() {
-            // It's possible that a concurrent DDL statement has already dropped this GlobalId
-            if let Some(read_needs) = self.storage_read_capabilities.get_mut(id) {
-                read_needs.holds.update_iter(time.iter().map(|t| (*t, -1)));
-                policy_changes.push((*id, read_needs.policy()));
+        let mut storage_policy_changes = Vec::new();
+        for read_holds in read_holdses.iter() {
+            for (time, id) in read_holds.storage_ids() {
+                // It's possible that a concurrent DDL statement has already dropped this GlobalId
+                if let Some(read_needs) = self.storage_read_capabilities.get_mut(id) {
+                    read_needs.holds.update_iter(time.iter().map(|t| (*t, -1)));
+                    storage_policy_changes.push((*id, read_needs.policy()));
+                }
             }
         }
-        self.controller.storage.set_read_policy(policy_changes);
+        self.controller
+            .storage
+            .set_read_policy(storage_policy_changes);
         // Update COMPUTE read policies
         let mut compute = self.controller.active_compute();
-        for (compute_instance, compute_ids) in read_holds.compute_ids() {
-            let mut policy_changes = Vec::new();
-            for (time, id) in compute_ids {
-                // It's possible that a concurrent DDL statement has already dropped this GlobalId
-                if let Some(read_needs) = self.compute_read_capabilities.get_mut(id) {
-                    read_needs.holds.update_iter(time.iter().map(|t| (*t, -1)));
-                    policy_changes.push((*id, read_needs.policy()));
+        let mut policy_changes_per_instance = BTreeMap::new();
+        for read_holds in read_holdses.iter() {
+            for (compute_instance, compute_ids) in read_holds.compute_ids() {
+                let policy_changes = policy_changes_per_instance
+                    .entry(compute_instance)
+                    .or_insert_with(Vec::new);
+                for (time, id) in compute_ids {
+                    // It's possible that a concurrent DDL statement has already dropped this GlobalId
+                    if let Some(read_needs) = self.compute_read_capabilities.get_mut(id) {
+                        read_needs.holds.update_iter(time.iter().map(|t| (*t, -1)));
+                        policy_changes.push((*id, read_needs.policy()));
+                    }
                 }
             }
+        }
+        for (compute_instance, policy_changes) in policy_changes_per_instance {
             if compute.instance_exists(*compute_instance) {
                 compute
                     .set_read_policy(*compute_instance, policy_changes)

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

@@ -206,7 +206,10 @@ impl Coordinator {
         let plan::CreateMaterializedViewPlan {
             materialized_view:
                 plan::MaterializedView {
-                    expr, cluster_id, ..
+                    expr,
+                    cluster_id,
+                    refresh_schedule,
+                    ..
                 },
             ambiguous_columns,
             ..
@@ -243,6 +246,31 @@ impl Coordinator {
             role_metadata: session.role_metadata().clone(),
         };
 
+        // Acquire read holds at all the REFRESH AT times.
+        // Note that we already acquired a possibly non-precise read hold at mz_now() in the purification,
+        // if any of the REFRESH options involve mz_now(). But now we can acquire precise read holds, because by now
+        // the REFRESH AT expressions have been evaluated, so we can handle something like
+        // `mz_now()::text::int8 + 10000`;
+        if let Some(refresh_schedule) = refresh_schedule {
+            if !refresh_schedule.ats.is_empty() {
+                let ids = self
+                    .index_oracle(*cluster_id)
+                    .sufficient_collections(resolved_ids.0.iter());
+                for refresh_at_ts in &refresh_schedule.ats {
+                    match self.acquire_read_holds_auto_cleanup(session, *refresh_at_ts, &ids, true)
+                    {
+                        Ok(()) => {}
+                        Err(earliest_possible) => {
+                            return Err(AdapterError::InputNotReadableAtRefreshAtTime(
+                                *refresh_at_ts,
+                                earliest_possible,
+                            ));
+                        }
+                    };
+                }
+            }
+        }
+
         Ok(CreateMaterializedViewStage::Optimize(
             CreateMaterializedViewOptimize {
                 validity,

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

@@ -984,6 +984,9 @@ impl Coordinator {
         // we must acquire read holds here so they are held until the off-thread work
         // returns to the coordinator.
         if let Some(txn_reads) = self.txn_read_holds.get(session.conn_id()) {
+            // Transactions involving peeks will acquire read holds at most once.
+            assert_eq!(txn_reads.len(), 1);
+            let txn_reads = &txn_reads[0];
             // Find referenced ids not in the read hold. A reference could be caused by a
             // user specifying an object in a different schema than the first query. An
             // index could be caused by a CREATE INDEX after the transaction started.
@@ -1000,7 +1003,8 @@ impl Coordinator {
                 });
             }
         } else if let Some((timestamp, bundle)) = potential_read_holds {
-            self.acquire_read_holds_auto_cleanup(session, timestamp, bundle);
+            self.acquire_read_holds_auto_cleanup(session, timestamp, bundle, true)
+                .expect("able to acquire read holds at the time that we just got from `determine_timestamp`");
         }
 
         // TODO: Checking for only `InTransaction` and not `Implied` (also `Started`?) seems

src/adapter/src/coord/sql.rs

@@ -205,7 +205,7 @@ impl Coordinator {
 
         // Release this transaction's compaction hold on collections.
         if let Some(txn_reads) = self.txn_read_holds.remove(conn_id) {
-            self.release_read_hold(&txn_reads);
+            self.release_read_holds(txn_reads);
         }
     }
 

src/adapter/src/coord/timeline.rs

@@ -723,7 +723,7 @@ impl Coordinator {
             };
             let read_ts = oracle.read_ts().await;
             if read_holds.times().any(|time| time.less_than(&read_ts)) {
-                read_holds = self.update_read_hold(read_holds, read_ts);
+                read_holds = self.update_read_holds(read_holds, read_ts);
             }
             self.global_timelines
                 .insert(timeline, TimelineState { oracle, read_holds });