[copy_from] Proper cancelation via CancelOneshotIngestion message (#31136)

This PR fixes the `TODO(cf1)` related to canceling oneshot ingestions.
It adds a `StorageCommand::CancelOneshotIngestion` that
`reduce`s/compacts away a corresponding
`StorageCommand::RunOneshotIngestion`, much like `ComputeCommand::Peek`
and `ComputeCommand::CancelPeek`.

We send a `StorageCommand::CancelOneshotIngestion` whenever a user has
canceled a `COPY FROM` statement, but also the storage controller will
send one whenever a `RunOneshotIngestion` command completes.

### Motivation

Fix `TODO(cf1)` related to cancelation

### Checklist

- [x] This PR has adequate test coverage / QA involvement has been duly
considered. ([trigger-ci for additional test/nightly
runs](https://trigger-ci.dev.materialize.com/))
- [x] This PR has an associated up-to-date [design
doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md),
is a design doc
([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)),
or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [x] If this PR evolves [an existing `$T ⇔ Proto$T`
mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md)
(possibly in a backwards-incompatible way), then it is tagged with a
`T-proto` label.
- [x] If this PR will require changes to cloud orchestration or tests,
there is a companion cloud PR to account for those changes that is
tagged with the release-blocker label
([example](https://github.com/MaterializeInc/cloud/pull/5021)).
<!-- Ask in #team-cloud on Slack if you need help preparing the cloud
PR. -->
- [x] If this PR includes major [user-facing behavior
changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note),
I have pinged the relevant PM to schedule a changelog post.

Author: ParkMyCar

src/adapter/src/active_compute_sink.rs

@@ -22,14 +22,15 @@ use mz_expr::compare_columns;
 use mz_ore::cast::CastFrom;
 use mz_ore::now::EpochMillis;
 use mz_repr::adt::numeric;
-use mz_repr::{Datum, GlobalId, IntoRowIterator, Row, Timestamp};
+use mz_repr::{CatalogItemId, Datum, GlobalId, IntoRowIterator, Row, Timestamp};
 use mz_sql::plan::SubscribeOutput;
+use mz_storage_types::instances::StorageInstanceId;
 use timely::progress::Antichain;
 use tokio::sync::{mpsc, oneshot};
 use uuid::Uuid;
 
 use crate::coord::peek::PeekResponseUnary;
-use crate::{AdapterError, ExecuteResponse};
+use crate::{AdapterError, ExecuteContext, ExecuteResponse};
 
 #[derive(Debug)]
 /// A description of an active compute sink from the coordinator's perspective.
@@ -435,3 +436,16 @@ impl ActiveCopyTo {
         let _ = self.tx.send(message);
     }
 }
+
+/// State we keep in the `Coordinator` to track active `COPY FROM` statements.
+#[derive(Debug)]
+pub(crate) struct ActiveCopyFrom {
+    /// ID of the ingestion running in clusterd.
+    pub ingestion_id: uuid::Uuid,
+    /// The cluster this is currently running on.
+    pub cluster_id: StorageInstanceId,
+    /// The table we're currently copying into.
+    pub table_id: CatalogItemId,
+    /// Context of the SQL session that ran the statement.
+    pub ctx: ExecuteContext,
+}

src/adapter/src/coord.rs

@@ -171,7 +171,7 @@ use tracing::{debug, info, info_span, span, warn, Instrument, Level, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use uuid::Uuid;
 
-use crate::active_compute_sink::ActiveComputeSink;
+use crate::active_compute_sink::{ActiveComputeSink, ActiveCopyFrom};
 use crate::catalog::{BuiltinTableUpdate, Catalog, OpenCatalogResult};
 use crate::client::{Client, Handle};
 use crate::command::{Command, ExecuteResponse};
@@ -1668,7 +1668,7 @@ pub struct Coordinator {
     active_webhooks: BTreeMap<CatalogItemId, WebhookAppenderInvalidator>,
     /// A map of active `COPY FROM` statements. The Coordinator waits for `clusterd`
     /// to stage Batches in Persist that we will then link into the shard.
-    active_copies: BTreeMap<ConnectionId, ExecuteContext>,
+    active_copies: BTreeMap<ConnectionId, ActiveCopyFrom>,
 
     /// A map from connection ids to a watch channel that is set to `true` if the connection
     /// received a cancel request.

src/adapter/src/coord/ddl.rs

@@ -207,6 +207,7 @@ impl Coordinator {
         let mut cluster_replicas_to_drop = vec![];
         let mut compute_sinks_to_drop = BTreeMap::new();
         let mut peeks_to_drop = vec![];
+        let mut copies_to_drop = vec![];
         let mut clusters_to_create = vec![];
         let mut cluster_replicas_to_create = vec![];
         let mut update_tracing_config = false;
@@ -454,6 +455,18 @@ impl Coordinator {
             }
         }
 
+        // Clean up any pending `COPY` statements that rely on dropped relations or clusters.
+        for (conn_id, pending_copy) in &self.active_copies {
+            let dropping_table = table_gids_to_drop
+                .iter()
+                .any(|(item_id, _gid)| pending_copy.table_id == *item_id);
+            let dropping_cluster = clusters_to_drop.contains(&pending_copy.cluster_id);
+
+            if dropping_table || dropping_cluster {
+                copies_to_drop.push(conn_id.clone());
+            }
+        }
+
         let storage_ids_to_drop = sources_to_drop
             .iter()
             .map(|(_, gid)| *gid)
@@ -666,6 +679,11 @@ impl Coordinator {
                     }
                 }
             }
+            if !copies_to_drop.is_empty() {
+                for conn_id in copies_to_drop {
+                    self.cancel_pending_copy(&conn_id);
+                }
+            }
             if !indexes_to_drop.is_empty() {
                 self.drop_indexes(indexes_to_drop);
             }

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

@@ -23,7 +23,7 @@ use url::Url;
 use uuid::Uuid;
 
 use crate::coord::sequencer::inner::return_if_err;
-use crate::coord::{Coordinator, TargetCluster};
+use crate::coord::{ActiveCopyFrom, Coordinator, TargetCluster};
 use crate::optimize::dataflows::{prep_scalar_expr, EvalTime, ExprPrepStyle};
 use crate::session::{TransactionOps, WriteOp};
 use crate::{AdapterError, ExecuteContext, ExecuteResponse};
@@ -177,8 +177,16 @@ impl Coordinator {
             });
         });
         // Stash the execute context so we can cancel the COPY.
-        self.active_copies
-            .insert(ctx.session().conn_id().clone(), ctx);
+        let conn_id = ctx.session().conn_id().clone();
+        self.active_copies.insert(
+            conn_id,
+            ActiveCopyFrom {
+                ingestion_id,
+                cluster_id,
+                table_id: id,
+                ctx,
+            },
+        );
 
         let _result = self
             .controller
@@ -193,11 +201,20 @@ impl Coordinator {
         table_id: CatalogItemId,
         batches: Vec<Result<ProtoBatch, String>>,
     ) {
-        let Some(mut ctx) = self.active_copies.remove(&conn_id) else {
-            tracing::warn!(?conn_id, "got response for canceled COPY FROM");
+        let Some(active_copy) = self.active_copies.remove(&conn_id) else {
+            // Getting a successful response for a cancel COPY FROM is unexpected.
+            tracing::warn!(%conn_id, ?batches, "got response for canceled COPY FROM");
             return;
         };
 
+        let ActiveCopyFrom {
+            ingestion_id,
+            cluster_id: _,
+            table_id: _,
+            mut ctx,
+        } = active_copy;
+        tracing::info!(%ingestion_id, num_batches = ?batches.len(), "received batches to append");
+
         let mut all_batches = SmallVec::with_capacity(batches.len());
         let mut all_errors = SmallVec::<[String; 1]>::with_capacity(batches.len());
         let mut row_count = 0u64;
@@ -248,8 +265,21 @@ impl Coordinator {
     /// Cancel any active `COPY FROM` statements/oneshot ingestions.
     #[mz_ore::instrument(level = "debug")]
     pub(crate) fn cancel_pending_copy(&mut self, conn_id: &ConnectionId) {
-        // TODO(cf1): Also cancel the dataflow running on clusterd.
-        if let Some(ctx) = self.active_copies.remove(conn_id) {
+        if let Some(ActiveCopyFrom {
+            ingestion_id,
+            cluster_id: _,
+            table_id: _,
+            ctx,
+        }) = self.active_copies.remove(conn_id)
+        {
+            let cancel_result = self
+                .controller
+                .storage
+                .cancel_oneshot_ingestion(ingestion_id);
+            if let Err(err) = cancel_result {
+                tracing::error!(?err, "failed to cancel OneshotIngestion");
+            }
+
             ctx.retire(Err(AdapterError::Canceled));
         }
     }

src/storage-client/src/client.proto

@@ -48,13 +48,21 @@ message ProtoRunIngestionCommand {
   mz_storage_types.sources.ProtoIngestionDescription description = 2;
 }
 
-message ProtoRunOneshotIngestionCommand {
+message ProtoRunOneshotIngestion {
   mz_proto.ProtoU128 ingestion_id = 1;
   mz_repr.global_id.ProtoGlobalId collection_id = 2;
   mz_storage_types.controller.ProtoCollectionMetadata storage_metadata = 3;
   mz_storage_types.oneshot_sources.ProtoOneshotIngestionRequest request = 4;
 }
 
+message ProtoRunOneshotIngestionsCommand {
+  repeated ProtoRunOneshotIngestion ingestions = 1;
+}
+
+message ProtoCancelOneshotIngestionsCommand {
+  repeated mz_proto.ProtoU128 ingestions = 1;
+}
+
 message ProtoCreateSources {
   repeated ProtoRunIngestionCommand sources = 1;
 }
@@ -94,7 +102,8 @@ message ProtoStorageCommand {
     google.protobuf.Empty allow_writes = 7;
     ProtoRunSinks run_sinks = 4;
     mz_storage_types.parameters.ProtoStorageParameters update_configuration = 5;
-    ProtoRunOneshotIngestionCommand oneshot_ingestion = 10;
+    ProtoRunOneshotIngestionsCommand run_oneshot_ingestions = 10;
+    ProtoCancelOneshotIngestionsCommand cancel_oneshot_ingestions = 11;
   }
 }
 

src/storage-client/src/client.rs

@@ -40,6 +40,7 @@ use smallvec::SmallVec;
 use timely::progress::frontier::{Antichain, MutableAntichain};
 use timely::PartialOrder;
 use tonic::{Request, Status as TonicStatus, Streaming};
+use uuid::Uuid;
 
 use crate::client::proto_storage_server::ProtoStorage;
 use crate::metrics::ReplicaMetrics;
@@ -123,18 +124,29 @@ pub enum StorageCommand<T = mz_repr::Timestamp> {
     UpdateConfiguration(StorageParameters),
     /// Run the enumerated sources, each associated with its identifier.
     RunIngestions(Vec<RunIngestionCommand>),
-    /// Run a dataflow which will ingest data from an external source and only __stage__ it in
-    /// Persist.
-    ///
-    /// Unlike regular ingestions/sources, some other component (e.g. `environmentd`) is
-    /// responsible for linking the staged data into a shard.
-    RunOneshotIngestion(RunOneshotIngestionCommand),
     /// Enable compaction in storage-managed collections.
     ///
     /// Each entry in the vector names a collection and provides a frontier after which
     /// accumulations must be correct.
     AllowCompaction(Vec<(GlobalId, Antichain<T>)>),
     RunSinks(Vec<RunSinkCommand<T>>),
+    /// Run a dataflow which will ingest data from an external source and only __stage__ it in
+    /// Persist.
+    ///
+    /// Unlike regular ingestions/sources, some other component (e.g. `environmentd`) is
+    /// responsible for linking the staged data into a shard.
+    RunOneshotIngestion(Vec<RunOneshotIngestion>),
+    /// `CancelOneshotIngestion` instructs the replica to cancel the identified oneshot ingestions.
+    ///
+    /// It is invalid to send a [`CancelOneshotIngestion`] command that references a oneshot
+    /// ingestion that was not created by a corresponding [`RunOneshotIngestion`] command before.
+    /// Doing so may cause the replica to exhibit undefined behavior.
+    ///
+    /// [`CancelOneshotIngestion`]: crate::client::StorageCommand::CancelOneshotIngestion
+    /// [`RunOneshotIngestion`]: crate::client::StorageCommand::RunOneshotIngestion
+    CancelOneshotIngestion {
+        ingestions: Vec<Uuid>,
+    },
 }
 
 impl<T> StorageCommand<T> {
@@ -146,7 +158,8 @@ impl<T> StorageCommand<T> {
             | InitializationComplete
             | AllowWrites
             | UpdateConfiguration(_)
-            | AllowCompaction(_) => false,
+            | AllowCompaction(_)
+            | CancelOneshotIngestion { .. } => false,
             // TODO(cf2): multi-replica oneshot ingestions. At the moment returning
             // true here means we can't run `COPY FROM` on multi-replica clusters, this
             // should be easy enough to support though.
@@ -199,7 +212,7 @@ impl RustType<ProtoRunIngestionCommand> for RunIngestionCommand {
 
 /// A command that starts ingesting the given ingestion description
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
-pub struct RunOneshotIngestionCommand {
+pub struct RunOneshotIngestion {
     /// The ID of the ingestion dataflow.
     pub ingestion_id: uuid::Uuid,
     /// The ID of collection we'll stage batches for.
@@ -210,30 +223,30 @@ pub struct RunOneshotIngestionCommand {
     pub request: OneshotIngestionRequest,
 }
 
-impl RustType<ProtoRunOneshotIngestionCommand> for RunOneshotIngestionCommand {
-    fn into_proto(&self) -> ProtoRunOneshotIngestionCommand {
-        ProtoRunOneshotIngestionCommand {
+impl RustType<ProtoRunOneshotIngestion> for RunOneshotIngestion {
+    fn into_proto(&self) -> ProtoRunOneshotIngestion {
+        ProtoRunOneshotIngestion {
             ingestion_id: Some(self.ingestion_id.into_proto()),
             collection_id: Some(self.collection_id.into_proto()),
             storage_metadata: Some(self.collection_meta.into_proto()),
             request: Some(self.request.into_proto()),
         }
     }
 
-    fn from_proto(proto: ProtoRunOneshotIngestionCommand) -> Result<Self, TryFromProtoError> {
-        Ok(RunOneshotIngestionCommand {
+    fn from_proto(proto: ProtoRunOneshotIngestion) -> Result<Self, TryFromProtoError> {
+        Ok(RunOneshotIngestion {
             ingestion_id: proto
                 .ingestion_id
-                .into_rust_if_some("ProtoRunOneshotIngestionCommand::ingestion_id")?,
+                .into_rust_if_some("ProtoRunOneshotIngestion::ingestion_id")?,
             collection_id: proto
                 .collection_id
-                .into_rust_if_some("ProtoRunOneshotIngestionCommand::collection_id")?,
+                .into_rust_if_some("ProtoRunOneshotIngestion::collection_id")?,
             collection_meta: proto
                 .storage_metadata
-                .into_rust_if_some("ProtoRunOneshotIngestionCommand::storage_metadata")?,
+                .into_rust_if_some("ProtoRunOneshotIngestion::storage_metadata")?,
             request: proto
                 .request
-                .into_rust_if_some("ProtoRunOneshotIngestionCommand::request")?,
+                .into_rust_if_some("ProtoRunOneshotIngestion::request")?,
         })
     }
 }
@@ -300,12 +313,19 @@ impl RustType<ProtoStorageCommand> for StorageCommand<mz_repr::Timestamp> {
                 StorageCommand::RunIngestions(sources) => CreateSources(ProtoCreateSources {
                     sources: sources.into_proto(),
                 }),
-                StorageCommand::RunOneshotIngestion(oneshot) => {
-                    OneshotIngestion(oneshot.into_proto())
-                }
                 StorageCommand::RunSinks(sinks) => RunSinks(ProtoRunSinks {
                     sinks: sinks.into_proto(),
                 }),
+                StorageCommand::RunOneshotIngestion(ingestions) => {
+                    RunOneshotIngestions(ProtoRunOneshotIngestionsCommand {
+                        ingestions: ingestions.iter().map(|cmd| cmd.into_proto()).collect(),
+                    })
+                }
+                StorageCommand::CancelOneshotIngestion { ingestions } => {
+                    CancelOneshotIngestions(ProtoCancelOneshotIngestionsCommand {
+                        ingestions: ingestions.iter().map(|uuid| uuid.into_proto()).collect(),
+                    })
+                }
             }),
         }
     }
@@ -334,8 +354,21 @@ impl RustType<ProtoStorageCommand> for StorageCommand<mz_repr::Timestamp> {
             Some(RunSinks(ProtoRunSinks { sinks })) => {
                 Ok(StorageCommand::RunSinks(sinks.into_rust()?))
             }
-            Some(OneshotIngestion(oneshot)) => {
-                Ok(StorageCommand::RunOneshotIngestion(oneshot.into_rust()?))
+            Some(RunOneshotIngestions(oneshot)) => {
+                let ingestions = oneshot
+                    .ingestions
+                    .into_iter()
+                    .map(|cmd| cmd.into_rust())
+                    .collect::<Result<_, _>>()?;
+                Ok(StorageCommand::RunOneshotIngestion(ingestions))
+            }
+            Some(CancelOneshotIngestions(oneshot)) => {
+                let ingestions = oneshot
+                    .ingestions
+                    .into_iter()
+                    .map(|uuid| uuid.into_rust())
+                    .collect::<Result<_, _>>()?;
+                Ok(StorageCommand::CancelOneshotIngestion { ingestions })
             }
             None => Err(TryFromProtoError::missing_field(
                 "ProtoStorageCommand::kind",
@@ -802,7 +835,8 @@ where
             | StorageCommand::AllowWrites
             | StorageCommand::UpdateConfiguration(_)
             | StorageCommand::AllowCompaction(_)
-            | StorageCommand::RunOneshotIngestion(_) => {}
+            | StorageCommand::RunOneshotIngestion(_)
+            | StorageCommand::CancelOneshotIngestion { .. } => {}
         };
     }
 

src/storage-client/src/controller.rs

@@ -498,6 +498,12 @@ pub trait StorageController: Debug {
         result_tx: OneshotResultCallback<ProtoBatch>,
     ) -> Result<(), StorageError<Self::Timestamp>>;
 
+    /// Cancel a oneshot ingestion.
+    fn cancel_oneshot_ingestion(
+        &mut self,
+        ingestion_id: uuid::Uuid,
+    ) -> Result<(), StorageError<Self::Timestamp>>;
+
     /// Alter the sink identified by the given id to match the provided `ExportDescription`.
     async fn alter_export(
         &mut self,