Frontend peek: Split ExecuteCopyTo into 2 commands to not block the main task

Author: ggevay

src/adapter/src/client.rs

@@ -1032,6 +1032,7 @@ impl SessionClient {
                 | Command::GetTransactionReadHoldsBundle { .. }
                 | Command::StoreTransactionReadHolds { .. }
                 | Command::ExecuteSlowPathPeek { .. }
+                | Command::CopyToPreflight { .. }
                 | Command::ExecuteCopyTo { .. }
                 | Command::ExecuteSideEffectingFunc { .. }
                 | Command::RegisterFrontendPeek { .. }

src/adapter/src/command.rs

@@ -224,6 +224,18 @@ pub enum Command {
         tx: oneshot::Sender<Result<ExecuteResponse, AdapterError>>,
     },
 
+    /// Preflight check for COPY TO S3 operation. This runs the slow S3 operations
+    /// (loading SDK config, checking bucket path, verifying permissions, uploading sentinel)
+    /// in a background task to avoid blocking the coordinator.
+    CopyToPreflight {
+        /// The S3 connection info needed for preflight checks.
+        s3_sink_connection: mz_compute_types::sinks::CopyToS3OneshotSinkConnection,
+        /// The sink ID for logging and S3 key management.
+        sink_id: GlobalId,
+        /// Response channel for the preflight result.
+        tx: oneshot::Sender<Result<(), AdapterError>>,
+    },
+
     ExecuteCopyTo {
         df_desc: Box<DataflowDescription<mz_compute_types::plan::Plan>>,
         compute_instance: ComputeInstanceId,
@@ -309,6 +321,7 @@ impl Command {
             | Command::GetTransactionReadHoldsBundle { .. }
             | Command::StoreTransactionReadHolds { .. }
             | Command::ExecuteSlowPathPeek { .. }
+            | Command::CopyToPreflight { .. }
             | Command::ExecuteCopyTo { .. }
             | Command::ExecuteSideEffectingFunc { .. }
             | Command::RegisterFrontendPeek { .. }
@@ -341,6 +354,7 @@ impl Command {
             | Command::GetTransactionReadHoldsBundle { .. }
             | Command::StoreTransactionReadHolds { .. }
             | Command::ExecuteSlowPathPeek { .. }
+            | Command::CopyToPreflight { .. }
             | Command::ExecuteCopyTo { .. }
             | Command::ExecuteSideEffectingFunc { .. }
             | Command::RegisterFrontendPeek { .. }

src/adapter/src/coord.rs

@@ -376,6 +376,7 @@ impl Message {
                 }
                 Command::StoreTransactionReadHolds { .. } => "store-transaction-read-holds",
                 Command::ExecuteSlowPathPeek { .. } => "execute-slow-path-peek",
+                Command::CopyToPreflight { .. } => "copy-to-preflight",
                 Command::ExecuteCopyTo { .. } => "execute-copy-to",
                 Command::ExecuteSideEffectingFunc { .. } => "execute-side-effecting-func",
                 Command::RegisterFrontendPeek { .. } => "register-frontend-peek",

src/adapter/src/coord/command_handler.rs

@@ -365,6 +365,28 @@ impl Coordinator {
                     let _ = tx.send(result);
                 }
 
+                Command::CopyToPreflight {
+                    s3_sink_connection,
+                    sink_id,
+                    tx,
+                } => {
+                    // Spawn a background task to perform the slow S3 preflight operations.
+                    // This avoids blocking the coordinator's main task.
+                    let connection_context = self.connection_context().clone();
+                    task::spawn(|| "copy_to_preflight", async move {
+                        let result = mz_storage_types::sinks::s3_oneshot_sink::preflight(
+                            connection_context,
+                            &s3_sink_connection.aws_connection,
+                            &s3_sink_connection.upload_info,
+                            s3_sink_connection.connection_id,
+                            sink_id,
+                        )
+                        .await
+                        .map_err(AdapterError::from);
+                        let _ = tx.send(result);
+                    });
+                }
+
                 Command::ExecuteCopyTo {
                     df_desc,
                     compute_instance,

src/adapter/src/coord/peek.rs

@@ -28,7 +28,6 @@ use mz_compute_client::protocol::command::PeekTarget;
 use mz_compute_client::protocol::response::PeekResponse;
 use mz_compute_types::ComputeInstanceId;
 use mz_compute_types::dataflows::{DataflowDescription, IndexImport};
-use mz_compute_types::sinks::ComputeSinkConnection;
 use mz_controller_types::ClusterId;
 use mz_expr::explain::{HumanizedExplain, HumanizerMode, fmt_text_constant_rows};
 use mz_expr::row::RowCollection;
@@ -1255,12 +1254,14 @@ impl crate::coord::Coordinator {
         .await
     }
 
-    /// Implements a COPY TO command by installing peek watch sets, validating S3 connection,
+    /// Implements a `COPY TO` command by installing peek watch sets,
     /// 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
-    /// to avoid the complexity of the staging mechanism for the frontend peek sequencing path.
+    /// (The S3 preflight check must be completed successfully via the
+    /// `CopyToPreflight` command _before_ calling this method. The preflight is
+    /// handled separately to avoid blocking the coordinator's main task with
+    /// slow S3 network operations.)
     ///
     /// This method does NOT block waiting for completion. Instead, it spawns a background task that
     /// will send the response through the provided tx channel when the COPY TO completes.
@@ -1300,49 +1301,6 @@ impl crate::coord::Coordinator {
 
         let sink_id = df_desc.sink_id();
 
-        // # Inlined from peek_copy_to_preflight
-
-        let connection_context = self.connection_context().clone();
-        let sinks = &df_desc.sink_exports;
-
-        if sinks.len() != 1 {
-            send_err(
-                tx,
-                AdapterError::Internal("expected exactly one copy to s3 sink".into()),
-            );
-            return;
-        }
-        let (sink_id_check, sink_desc) = sinks
-            .first_key_value()
-            .expect("known to be exactly one copy to s3 sink");
-        assert_eq!(sink_id, *sink_id_check);
-
-        match &sink_desc.connection {
-            ComputeSinkConnection::CopyToS3Oneshot(conn) => {
-                if let Err(e) = mz_storage_types::sinks::s3_oneshot_sink::preflight(
-                    connection_context,
-                    &conn.aws_connection,
-                    &conn.upload_info,
-                    conn.connection_id,
-                    sink_id,
-                )
-                .await
-                {
-                    send_err(tx, e.into());
-                    return;
-                }
-            }
-            _ => {
-                send_err(
-                    tx,
-                    AdapterError::Internal("expected copy to s3 oneshot sink".into()),
-                );
-                return;
-            }
-        }
-
-        // # Inlined from peek_copy_to_dataflow
-
         // Create and register ActiveCopyTo.
         // Note: sink_tx/sink_rx is the channel for the compute sink to notify completion
         // This is different from the command's tx which sends the response to the client

src/adapter/src/frontend_peek.rs

@@ -1311,6 +1311,38 @@ impl PeekClient {
                     &df_meta.optimizer_notices,
                 );
 
+                // Extract S3 sink connection info for preflight check
+                let sink_id = df_desc.sink_id();
+                let sinks = &df_desc.sink_exports;
+                if sinks.len() != 1 {
+                    return Err(AdapterError::Internal(
+                        "expected exactly one copy to s3 sink".into(),
+                    ));
+                }
+                let (_, sink_desc) = sinks
+                    .first_key_value()
+                    .expect("known to be exactly one copy to s3 sink");
+                let s3_sink_connection = match &sink_desc.connection {
+                    mz_compute_types::sinks::ComputeSinkConnection::CopyToS3Oneshot(conn) => {
+                        conn.clone()
+                    }
+                    _ => {
+                        return Err(AdapterError::Internal(
+                            "expected copy to s3 oneshot sink".into(),
+                        ));
+                    }
+                };
+
+                // Perform S3 preflight check in background task (via coordinator).
+                // This runs slow S3 operations without blocking the coordinator's main task.
+                self.call_coordinator(|tx| Command::CopyToPreflight {
+                    s3_sink_connection,
+                    sink_id,
+                    tx,
+                })
+                .await?;
+
+                // Preflight succeeded, now execute the actual COPY TO dataflow
                 let watch_set = statement_logging_id.map(|logging_id| {
                     WatchSetCreation::new(
                         logging_id,