Merge pull request #13251 from teskje/internal-sinks

Internal sinks, part 1

Author: teskje

Cargo.lock

@@ -18,7 +18,8 @@ use differential_dataflow::operators::arrange::arrangement::ArrangeByKey;
 use differential_dataflow::{Collection, Hashable};
 use timely::dataflow::Scope;
 
-use mz_dataflow_types::sinks::*;
+use mz_dataflow_types::client::controller::storage::CollectionMetadata;
+use mz_dataflow_types::sinks::{SinkConnection, SinkDesc, SinkEnvelope};
 use mz_expr::{permutation_for_arrangement, MapFilterProject};
 use mz_interchange::envelopes::{combine_at_timestamp, dbz_format, upsert_format};
 use mz_repr::{Datum, Diff, GlobalId, Row, Timestamp};
@@ -36,7 +37,7 @@ where
         tokens: &mut std::collections::BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
         import_ids: BTreeSet<GlobalId>,
         sink_id: GlobalId,
-        sink: &SinkDesc,
+        sink: &SinkDesc<CollectionMetadata>,
     ) {
         let sink_render = get_sink_render_for(&sink.connection);
 
@@ -92,7 +93,7 @@ where
 
 #[allow(clippy::borrowed_box)]
 fn apply_sink_envelope<G>(
-    sink: &SinkDesc,
+    sink: &SinkDesc<CollectionMetadata>,
     sink_render: &Box<dyn SinkRender<G>>,
     collection: Collection<G, Row, Diff>,
 ) -> Collection<G, (Option<Row>, Option<Row>), Diff>
@@ -215,15 +216,15 @@ where
     fn render_continuous_sink(
         &self,
         compute_state: &mut crate::compute_state::ComputeState,
-        sink: &SinkDesc,
+        sink: &SinkDesc<CollectionMetadata>,
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
     ) -> Option<Rc<dyn Any>>
     where
         G: Scope<Timestamp = Timestamp>;
 }
 
-fn get_sink_render_for<G>(connection: &SinkConnection) -> Box<dyn SinkRender<G>>
+fn get_sink_render_for<G>(connection: &SinkConnection<CollectionMetadata>) -> Box<dyn SinkRender<G>>
 where
     G: Scope<Timestamp = Timestamp>,
 {

src/compute/src/render/sinks.rs

@@ -22,7 +22,6 @@ use differential_dataflow::{AsCollection, Collection, Hashable};
 use futures::executor::block_on;
 use futures::{StreamExt, TryFutureExt};
 use itertools::Itertools;
-use mz_interchange::json::JsonEncoder;
 use prometheus::core::AtomicU64;
 use rdkafka::client::ClientContext;
 use rdkafka::config::ClientConfig;
@@ -45,6 +44,7 @@ use tokio::sync::Mutex;
 use tracing::{debug, error, info};
 
 use mz_avro::types::Value;
+use mz_dataflow_types::client::controller::storage::CollectionMetadata;
 use mz_dataflow_types::connections::ConnectionContext;
 use mz_dataflow_types::sinks::{
     KafkaSinkConnection, KafkaSinkConsistencyConnection, PublishedSchemaInfo, SinkAsOf, SinkDesc,
@@ -54,6 +54,7 @@ use mz_interchange::avro::{
     self, get_debezium_transaction_schema, AvroEncoder, AvroSchemaGenerator,
 };
 use mz_interchange::encode::Encode;
+use mz_interchange::json::JsonEncoder;
 use mz_kafka_util::client::{create_new_client_config, MzClientContext};
 use mz_ore::cast::CastFrom;
 use mz_ore::collections::CollectionExt;
@@ -88,7 +89,7 @@ where
     fn render_continuous_sink(
         &self,
         compute_state: &mut crate::compute_state::ComputeState,
-        sink: &SinkDesc,
+        sink: &SinkDesc<CollectionMetadata>,
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
     ) -> Option<Rc<dyn Any>>

src/compute/src/sink/kafka.rs

@@ -21,14 +21,15 @@ use timely::progress::Antichain;
 use timely::progress::Timestamp as TimelyTimestamp;
 use timely::PartialOrder;
 
+use mz_dataflow_types::client::controller::storage::CollectionMetadata;
 use mz_dataflow_types::sinks::{PersistSinkConnection, SinkDesc};
-use mz_persist_client::PersistLocation;
+use mz_dataflow_types::sources::SourceData;
 use mz_repr::{Diff, GlobalId, Row, Timestamp};
 use mz_timely_util::operators_async_ext::OperatorBuilderExt;
 
 use crate::render::sinks::SinkRender;
 
-impl<G> SinkRender<G> for PersistSinkConnection
+impl<G> SinkRender<G> for PersistSinkConnection<CollectionMetadata>
 where
     G: Scope<Timestamp = Timestamp>,
 {
@@ -47,15 +48,24 @@ where
     fn render_continuous_sink(
         &self,
         compute_state: &mut crate::compute_state::ComputeState,
-        _sink: &SinkDesc,
+        _sink: &SinkDesc<CollectionMetadata>,
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
     ) -> Option<Rc<dyn Any>>
     where
         G: Scope<Timestamp = Timestamp>,
     {
         let scope = sinked_collection.scope();
-        let operator_name = format!("persist_sink({})", self.shard_id);
+
+        let persist_clients = Arc::clone(&compute_state.persist_clients);
+        let persist_location = self.storage_metadata.persist_location.clone();
+        let shard_id = self.storage_metadata.persist_shard;
+
+        // Log the shard ID so we know which shard to read for testing.
+        // TODO(teskje): Remove once we have a built-in way for reading back sinked data.
+        tracing::info!("persist_sink shard ID: {shard_id}");
+
+        let operator_name = format!("persist_sink({})", shard_id);
         let mut persist_op = OperatorBuilder::new(operator_name, scope.clone());
 
         // We want exactly one worker (in the cluster) to send all the data to persist. It's fine
@@ -72,13 +82,6 @@ where
         let mut input =
             persist_op.new_input(&sinked_collection.inner, Exchange::new(move |_| hashed_id));
 
-        let persist_clients = Arc::clone(&compute_state.persist_clients);
-        let persist_location = PersistLocation {
-            consensus_uri: self.consensus_uri.clone(),
-            blob_uri: self.blob_uri.clone(),
-        };
-        let shard_id = self.shard_id.clone();
-
         let token = Rc::new(());
         let token_weak = Rc::downgrade(&token);
 
@@ -113,7 +116,7 @@ where
                     .open(persist_location)
                     .await
                     .expect("could not open persist client")
-                    .open_writer::<Row, Row, Timestamp, Diff>(shard_id)
+                    .open_writer::<SourceData, (), Timestamp, Diff>(shard_id)
                     .await
                     .expect("could not open persist shard");
 
@@ -131,9 +134,13 @@ where
                         data.swap(&mut buffer);
 
                         for ((key, value), ts, diff) in buffer.drain(..) {
-                            let key = key.unwrap_or_default();
-                            let value = value.unwrap_or_default();
-                            stash.entry(ts).or_default().push(((key, value), ts, diff));
+                            assert!(key.is_none(), "persist_source does not support keys");
+                            let row = value.expect("persist_source must have values");
+                            stash.entry(ts).or_default().push((
+                                (SourceData(Ok(row)), ()),
+                                ts,
+                                diff,
+                            ));
                         }
                     });
 

src/compute/src/sink/persist_sink.rs

@@ -21,6 +21,7 @@ use timely::progress::timestamp::Timestamp as TimelyTimestamp;
 use timely::progress::Antichain;
 
 use mz_dataflow_types::{
+    client::controller::storage::CollectionMetadata,
     sinks::{SinkAsOf, SinkDesc, TailSinkConnection},
     TailResponse,
 };
@@ -47,7 +48,7 @@ where
     fn render_continuous_sink(
         &self,
         compute_state: &mut crate::compute_state::ComputeState,
-        sink: &SinkDesc,
+        sink: &SinkDesc<CollectionMetadata>,
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
     ) -> Option<Rc<dyn Any>>

src/compute/src/sink/tail.rs

@@ -30,7 +30,10 @@ use mz_dataflow_types::client::{
     ComputeInstanceId, ConcreteComputeInstanceReplicaConfig, ProcessId, ReplicaId,
 };
 use mz_dataflow_types::logging::LoggingConfig as DataflowLoggingConfig;
-use mz_dataflow_types::sinks::{SinkConnection, SinkConnectionBuilder, SinkEnvelope};
+use mz_dataflow_types::sinks::{
+    PersistSinkConnection, PersistSinkConnectionBuilder, SinkConnection, SinkConnectionBuilder,
+    SinkEnvelope,
+};
 use mz_dataflow_types::sources::{ExternalSourceConnection, SourceConnection, Timeline};
 use mz_expr::{ExprHumanizer, MirScalarExpr, OptimizedMirRelationExpr};
 use mz_ore::collections::CollectionExt;
@@ -174,6 +177,26 @@ impl CatalogState {
                     desc: source.desc.clone(),
                 })
             }
+            // TODO(teskje): Replace once `CatalogItem::RecordedView` lands.
+            CatalogItem::Sink(Sink {
+                connection:
+                    SinkConnectionState::Ready(SinkConnection::Persist(PersistSinkConnection {
+                        value_desc,
+                        ..
+                    }))
+                    | SinkConnectionState::Pending(SinkConnectionBuilder::Persist(
+                        PersistSinkConnectionBuilder { value_desc },
+                    )),
+                ..
+            }) => {
+                let connection = SourceConnection::Local {
+                    timeline: Timeline::EpochMilliseconds,
+                };
+                Some(mz_dataflow_types::sources::SourceDesc {
+                    connection,
+                    desc: value_desc.clone(),
+                })
+            }
             _ => None,
         }
     }

src/coord/src/catalog.rs

@@ -794,6 +794,29 @@ impl<S: Append + 'static> Coordinator<S> {
                 }
                 CatalogItem::View(_) => (),
                 CatalogItem::Sink(sink) => {
+                    // Re-announce the source description.
+                    // TODO(teskje): Remove this once persist sinks are replaced by recorded views.
+                    if let Some(desc) = self.catalog.state().source_description_for(entry.id()) {
+                        let ingestion = IngestionDescription {
+                            id: entry.id(),
+                            desc,
+                            since: Antichain::from_elem(Timestamp::minimum()),
+                            source_imports: BTreeMap::new(),
+                            storage_metadata: (),
+                        };
+                        self.dataflow_client
+                            .storage_mut()
+                            .create_sources(vec![ingestion])
+                            .await
+                            .unwrap();
+                        self.initialize_storage_read_policies(
+                            vec![entry.id()],
+                            self.logical_compaction_window_ms,
+                        )
+                        .await;
+                    }
+
+                    // Re-create the sink on the compute instance.
                     let builder = match &sink.connection {
                         SinkConnectionState::Pending(builder) => builder,
                         SinkConnectionState::Ready(_) => {
@@ -2904,7 +2927,29 @@ impl<S: Append + 'static> Coordinator<S> {
             })
             .await;
         match transact_result {
-            Ok(()) => (),
+            Ok(()) => {
+                // Announce the creation of the sink's corresponding source.
+                // TODO(teskje): Remove this once persist sinks are replaced by recorded views.
+                if let Some(desc) = self.catalog.state().source_description_for(id) {
+                    let ingestion = IngestionDescription {
+                        id,
+                        desc,
+                        since: Antichain::from_elem(Timestamp::minimum()),
+                        source_imports: BTreeMap::new(),
+                        storage_metadata: (),
+                    };
+                    self.dataflow_client
+                        .storage_mut()
+                        .create_sources(vec![ingestion])
+                        .await
+                        .unwrap();
+                    self.initialize_storage_read_policies(
+                        vec![id],
+                        self.logical_compaction_window_ms,
+                    )
+                    .await;
+                }
+            }
             Err(CoordError::Catalog(catalog::Error {
                 kind: catalog::ErrorKind::ItemAlreadyExists(_),
                 ..
@@ -5077,6 +5122,17 @@ impl<S: Append + 'static> Coordinator<S> {
                 .entry(compute_instance)
                 .or_insert(vec![])
                 .push(id);
+
+            // Persist sinks write to storage collections, which need to be
+            // dropped when their sinks are dropped.
+            // TODO(teskje): Remove this once persist sinks are replaced by recorded views.
+            if self.dataflow_client.storage_mut().collection(id).is_ok() {
+                self.dataflow_client
+                    .storage_mut()
+                    .drop_sources(vec![id])
+                    .await
+                    .unwrap();
+            }
         }
         for (compute_instance, ids) in by_compute_instance {
             // A cluster could have been dropped, so verify it exists.

src/coord/src/coord.rs

@@ -342,9 +342,7 @@ fn build_persist_sink(
     _id: GlobalId,
 ) -> Result<SinkConnection, CoordError> {
     Ok(SinkConnection::Persist(PersistSinkConnection {
-        consensus_uri: builder.consensus_uri,
-        blob_uri: builder.blob_uri,
-        shard_id: builder.shard_id,
         value_desc: builder.value_desc,
+        storage_metadata: (),
     }))
 }

src/coord/src/sink_connection.rs

@@ -37,6 +37,7 @@ use crate::client::replicated::ActiveReplication;
 use crate::client::{ComputeClient, ComputeCommand, ComputeInstanceId, InstanceConfig, ReplicaId};
 use crate::client::{GenericClient, Peek};
 use crate::logging::LoggingConfig;
+use crate::sinks::{PersistSinkConnection, SinkConnection, SinkDesc};
 use crate::{DataflowDescription, SourceInstanceDesc};
 use mz_expr::RowSetFinishing;
 use mz_ore::tracing::OpenTelemetryContext;
@@ -338,8 +339,8 @@ where
             }
         }
 
-        // Here we augment all the imported sources with the appropriate storage metadata needed by
-        // the compute instance to read them
+        // Here we augment all imported sources and all exported sinks with with the appropriate
+        // storage metadata needed by the compute instance.
         let mut augmented_dataflows = Vec::with_capacity(dataflows.len());
         for d in dataflows {
             let mut source_imports = BTreeMap::new();
@@ -353,13 +354,37 @@ where
                 source_imports.insert(id, desc);
             }
 
+            let mut sink_exports = BTreeMap::new();
+            for (id, se) in d.sink_exports {
+                let connection = match se.connection {
+                    SinkConnection::Persist(conn) => {
+                        let metadata = self.storage_controller.collection_metadata(id)?;
+                        let conn = PersistSinkConnection {
+                            value_desc: conn.value_desc,
+                            storage_metadata: metadata,
+                        };
+                        SinkConnection::Persist(conn)
+                    }
+                    SinkConnection::Kafka(conn) => SinkConnection::Kafka(conn),
+                    SinkConnection::Tail(conn) => SinkConnection::Tail(conn),
+                };
+                let desc = SinkDesc {
+                    from: se.from,
+                    from_desc: se.from_desc,
+                    connection,
+                    envelope: se.envelope,
+                    as_of: se.as_of,
+                };
+                sink_exports.insert(id, desc);
+            }
+
             augmented_dataflows.push(DataflowDescription {
                 source_imports,
+                sink_exports,
                 // The rest of the fields are identical
                 index_imports: d.index_imports,
                 objects_to_build: d.objects_to_build,
                 index_exports: d.index_exports,
-                sink_exports: d.sink_exports,
                 as_of: d.as_of,
                 debug_name: d.debug_name,
                 id: d.id,
@@ -587,6 +612,21 @@ where
             self.update_read_capabilities(&mut read_capability_changes)
                 .await?;
         }
+
+        // Tell the storage controller about new write frontiers for storage
+        // collections that are advanced by compute sinks.
+        // TODO(teskje): The storage controller should have a task to directly
+        // keep track of the frontiers of storage collections, instead of
+        // relying on others for that information.
+        let storage_updates: Vec<_> = updates
+            .iter()
+            .filter(|(id, _)| self.storage_mut().collection(*id).is_ok())
+            .cloned()
+            .collect();
+        self.storage_mut()
+            .update_write_frontiers(&storage_updates)
+            .await?;
+
         Ok(())
     }