Make watch set installation fallible; make RegisterFrontendPeek wait for a result

Author: ggevay

src/adapter/src/command.rs

@@ -240,6 +240,7 @@ pub enum Command {
         is_fast_path: bool,
         /// If statement logging is enabled, contains the info necessary for installing watch sets.
         ids_to_watch: Option<IdsToWatch>,
+        tx: oneshot::Sender<Result<(), AdapterError>>,
     },
 
     /// Statement logging event from frontend peek sequencing.

src/adapter/src/coord.rs

@@ -104,7 +104,7 @@ use mz_catalog::memory::objects::{
 };
 use mz_cloud_resources::{CloudResourceController, VpcEndpointConfig, VpcEndpointEvent};
 use mz_compute_client::as_of_selection;
-use mz_compute_client::controller::error::{DataflowCreationError, InstanceMissing};
+use mz_compute_client::controller::error::{CollectionLookupError, DataflowCreationError, InstanceMissing};
 use mz_compute_types::ComputeInstanceId;
 use mz_compute_types::dataflows::DataflowDescription;
 use mz_compute_types::plan::Plan;
@@ -3764,13 +3764,14 @@ impl Coordinator {
         objects: BTreeSet<GlobalId>,
         t: Timestamp,
         state: WatchSetResponse,
-    ) {
-        let ws_id = self.controller.install_compute_watch_set(objects, t);
+    ) -> Result<(), CollectionLookupError> {
+        let ws_id = self.controller.install_compute_watch_set(objects, t)?;
         self.connection_watch_sets
             .entry(conn_id.clone())
             .or_default()
             .insert(ws_id);
         self.installed_watch_sets.insert(ws_id, (conn_id, state));
+        Ok(())
     }
 
     /// Install a _watch set_ in the controller that is automatically associated with the given
@@ -3782,13 +3783,14 @@ impl Coordinator {
         objects: BTreeSet<GlobalId>,
         t: Timestamp,
         state: WatchSetResponse,
-    ) {
-        let ws_id = self.controller.install_storage_watch_set(objects, t);
+    ) -> Result<(), CollectionLookupError> {
+        let ws_id = self.controller.install_storage_watch_set(objects, t)?;
         self.connection_watch_sets
             .entry(conn_id.clone())
             .or_default()
             .insert(ws_id);
         self.installed_watch_sets.insert(ws_id, (conn_id, state));
+        Ok(())
     }
 
     /// Cancels pending watchsets associated with the provided connection id.

src/adapter/src/coord/command_handler.rs

@@ -347,10 +347,6 @@ impl Coordinator {
                     tx,
                 } => {
                     soft_assert_eq_or_log!(ctx_extra.contents().is_some(), ids_to_watch.is_some());
-                    if let (Some(logging_id), Some(ids_to_watch)) = (ctx_extra.contents(), ids_to_watch) {
-                        self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch);
-                    }
-
                     let result = self
                         .implement_slow_path_peek(
                             *dataflow_plan,
@@ -364,9 +360,9 @@ impl Coordinator {
                             max_result_size,
                             max_query_result_size,
                             ctx_extra,
+                            ids_to_watch,
                         )
                         .await;
-
                     let _ = tx.send(result);
                 }
 
@@ -381,15 +377,6 @@ impl Coordinator {
                     tx,
                 } => {
                     soft_assert_eq_or_log!(ctx_extra.contents().is_some(), ids_to_watch.is_some());
-                    if let (Some(logging_id), Some(ids_to_watch)) = (ctx_extra.contents(), ids_to_watch) {
-                        // (The old peek sequencing forgot to do this for COPY TO.)
-                        self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch);
-                    }
-
-                    // We retire the execution context immediately, as the coordinator doesn't need
-                    // to track the execution for statement logging purposes (the frontend handles that).
-                    let _ = ctx_extra.retire();
-
                     // implement_copy_to spawns a background task that sends the response
                     // through tx when the COPY TO completes (or immediately if setup fails).
                     // We just call it and let it handle all response sending.
@@ -399,6 +386,8 @@ impl Coordinator {
                         target_replica,
                         source_ids,
                         conn_id,
+                        ctx_extra,
+                        ids_to_watch,
                         tx,
                     )
                     .await;
@@ -411,6 +400,7 @@ impl Coordinator {
                     ctx_extra,
                     is_fast_path,
                     ids_to_watch,
+                    tx,
                 } => {
                     self.handle_register_frontend_peek(
                         uuid,
@@ -420,6 +410,7 @@ impl Coordinator {
                         ctx_extra,
                         is_fast_path,
                         ids_to_watch,
+                        tx,
                     );
                 }
                 Command::FrontendStatementLogging(event) => {
@@ -1877,10 +1868,14 @@ impl Coordinator {
         ctx_extra: ExecuteContextExtra,
         is_fast_path: bool,
         ids_to_watch: Option<IdsToWatch>,
+        tx: oneshot::Sender<Result<(), AdapterError>>,
     ) {
         soft_assert_eq_or_log!(ctx_extra.contents().is_some(), ids_to_watch.is_some());
         if let (Some(logging_id), Some(ids_to_watch)) = (ctx_extra.contents(), ids_to_watch) {
-            self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch);
+            if let Err(e) = self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch) {
+                let _ = tx.send(Err(AdapterError::concurrent_dependency_drop_from_collection_lookup_error(e, cluster_id.into())));
+                return;
+            }
         }
 
         // Store the peek in pending_peeks for later retrieval when results arrive
@@ -1900,5 +1895,7 @@ impl Coordinator {
             .entry(conn_id)
             .or_default()
             .insert(uuid, cluster_id.into());
+
+        let _ = tx.send(Ok(()));
     }
 }

src/adapter/src/coord/peek.rs

@@ -56,6 +56,7 @@ use tracing::{Instrument, Span};
 use uuid::Uuid;
 
 use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyTo};
+use crate::coord::statement_logging::IdsToWatch;
 use crate::coord::timestamp_selection::TimestampDetermination;
 use crate::optimize::OptimizerError;
 use crate::statement_logging::{StatementEndedExecutionReason, StatementExecutionStrategy};
@@ -1211,7 +1212,19 @@ impl crate::coord::Coordinator {
         max_result_size: u64,
         max_query_result_size: Option<u64>,
         mut ctx_extra: ExecuteContextExtra,
+        ids_to_watch: Option<IdsToWatch>,
     ) -> Result<ExecuteResponse, AdapterError> {
+        // Install watch sets for statement lifecycle logging if enabled.
+        if let (Some(logging_id), Some(ids_to_watch)) = (ctx_extra.contents(), ids_to_watch) {
+            self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch)
+                .map_err(|e| {
+                    AdapterError::concurrent_dependency_drop_from_collection_lookup_error(
+                        e,
+                        compute_instance,
+                    )
+                })?;
+        }
+
         let source_arity = intermediate_result_type.arity();
 
         let planned_peek = PlannedPeek {
@@ -1239,8 +1252,8 @@ impl crate::coord::Coordinator {
         .await
     }
 
-    /// Implements a COPY TO command by validating S3 connection, shipping the dataflow,
-    /// and spawning a background task to wait for completion.
+    /// Implements a COPY TO command by installing peek watch sets, validating S3 connection,
+    /// shipping the dataflow, and spawning a background task to wait for completion.
     /// This is called from the command handler for ExecuteCopyTo.
     ///
     /// This method inlines the logic from peek_copy_to_preflight and peek_copy_to_dataflow
@@ -1256,6 +1269,8 @@ impl crate::coord::Coordinator {
         target_replica: Option<ReplicaId>,
         source_ids: BTreeSet<GlobalId>,
         conn_id: ConnectionId,
+        ctx_extra: ExecuteContextExtra,
+        ids_to_watch: Option<IdsToWatch>,
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     ) {
         // Helper to send error and return early
@@ -1264,6 +1279,22 @@ impl crate::coord::Coordinator {
             let _ = tx.send(Err(e));
         };
 
+        // Install watch sets for statement lifecycle logging if enabled.
+        if let (Some(logging_id), Some(ids_to_watch)) = (ctx_extra.contents(), ids_to_watch) {
+            if let Err(e) = self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch) {
+                let err = AdapterError::concurrent_dependency_drop_from_collection_lookup_error(
+                    e,
+                    compute_instance,
+                );
+                send_err(tx, err);
+                return;
+            }
+        }
+
+        // We retire the execution context immediately, as the coordinator doesn't need
+        // to track the execution for statement logging purposes (the frontend handles that).
+        let _ = ctx_extra.retire();
+
         let sink_id = df_desc.sink_id();
 
         // # Inlined from peek_copy_to_preflight

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

@@ -3301,7 +3301,7 @@ impl Coordinator {
                 plan_validity,
                 read_hold,
             }),
-        );
+        ).expect("plan validity verified above; we are on the coordinator main task, so they couldn't have gone away since then");
     }
 
     #[instrument]

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

@@ -851,7 +851,7 @@ impl Coordinator {
                 &id_bundle,
                 determination.timestamp_context.timestamp_or_default(),
             );
-            self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch);
+            self.install_peek_watch_sets(conn_id.clone(), logging_id, ids_to_watch).expect("the old peek sequencing re-verifies the dependencies' existence before installing the new watch sets");
         }
 
         let max_result_size = self.catalog().system_config().max_result_size();

src/adapter/src/coord/statement_logging.rs

@@ -32,6 +32,7 @@ use tokio::time::MissedTickBehavior;
 use tracing::debug;
 use uuid::Uuid;
 use mz_adapter_types::connection::ConnectionId;
+use mz_compute_client::controller::error::CollectionLookupError;
 use crate::coord::{ConnMeta, Coordinator, WatchSetResponse};
 use crate::session::{LifecycleTimestamps, Session};
 use crate::statement_logging::{
@@ -1461,7 +1462,7 @@ impl Coordinator {
         conn_id: ConnectionId,
         logging_id: StatementLoggingId,
         ids_to_watch: IdsToWatch,
-    ) {
+    ) -> Result<(), CollectionLookupError> {
         let IdsToWatch {
             timestamp,
             storage_ids,
@@ -1476,7 +1477,7 @@ impl Coordinator {
                 logging_id,
                 StatementLifecycleEvent::StorageDependenciesFinished,
             ),
-        );
+        )?;
         self.install_compute_watch_set(
             conn_id,
             compute_ids,
@@ -1485,6 +1486,7 @@ impl Coordinator {
                 logging_id,
                 StatementLifecycleEvent::ComputeDependenciesFinished,
             ),
-        );
+        )?;
+        Ok(())
     }
 }

src/adapter/src/peek_client.rs

@@ -343,15 +343,17 @@ impl PeekClient {
 
         // Register coordinator tracking of this peek. This has to complete before issuing the peek.
         let ctx_extra = crate::coord::ExecuteContextExtra::new(statement_logging_id);
-        self.coordinator_client.send(Command::RegisterFrontendPeek {
+        self.call_coordinator(|tx| Command::RegisterFrontendPeek {
             uuid,
             conn_id: conn_id.clone(),
             cluster_id: compute_instance.into(),
             depends_on,
             ctx_extra,
             is_fast_path: true,
             ids_to_watch,
-        });
+            tx,
+        })
+        .await?;
 
         // At this stage we don't know column names for the result because we
         // only know the peek's result type as a bare SqlRelationType.

src/controller/src/lib.rs

@@ -62,6 +62,7 @@ use mz_txn_wal::metrics::Metrics as TxnMetrics;
 use timely::progress::{Antichain, Timestamp};
 use tokio::sync::mpsc;
 use uuid::Uuid;
+use mz_compute_client::controller::error::CollectionLookupError;
 
 pub mod clusters;
 
@@ -372,15 +373,20 @@ where
         &mut self,
         mut objects: BTreeSet<GlobalId>,
         t: T,
-    ) -> WatchSetId {
+    ) -> Result<WatchSetId, CollectionLookupError> {
         let ws_id = self.watch_set_id_gen.allocate_id();
 
+        // Collect all frontiers first, returning any errors
+        let frontiers: BTreeMap<GlobalId, _> = objects
+            .iter()
+            .map(|id| {
+                self.compute
+                    .collection_frontiers(*id, None)
+                    .map(|f| (*id, f.write_frontier))
+            })
+            .collect::<Result<_, _>>()?;
         objects.retain(|id| {
-            let frontier = self
-                .compute
-                .collection_frontiers(*id, None)
-                .map(|f| f.write_frontier)
-                .expect("missing compute dependency");
+            let frontier = frontiers.get(id).expect("just collected");
             frontier.less_equal(&t)
         });
         if objects.is_empty() {
@@ -395,7 +401,7 @@ where
             self.unfulfilled_watch_sets.insert(ws_id, (objects, t));
         }
 
-        ws_id
+        Ok(ws_id)
     }
 
     /// Install a _watch set_ in the controller.
@@ -410,13 +416,12 @@ where
         &mut self,
         mut objects: BTreeSet<GlobalId>,
         t: T,
-    ) -> WatchSetId {
+    ) -> Result<WatchSetId, CollectionLookupError> {
         let ws_id = self.watch_set_id_gen.allocate_id();
 
         let uppers = self
             .storage
-            .collections_frontiers(objects.iter().cloned().collect())
-            .expect("missing storage dependencies")
+            .collections_frontiers(objects.iter().cloned().collect())?
             .into_iter()
             .map(|(id, _since, upper)| (id, upper))
             .collect::<BTreeMap<_, _>>();
@@ -436,7 +441,7 @@ where
             }
             self.unfulfilled_watch_sets.insert(ws_id, (objects, t));
         }
-        ws_id
+        Ok(ws_id)
     }
 
     /// Uninstalls a previously installed WatchSetId. The method is a no-op if the watch set has