storage: Make render sources take a proper struct

Signed-off-by: Moritz Hoffmann <mh@materialize.com>

Author: antiguru

src/dataflow-types/src/client.rs

@@ -153,23 +153,26 @@ pub struct CreateSourceCommand<T> {
     pub ts_bindings: Vec<(PartitionId, T, crate::sources::MzOffset)>,
 }
 
+/// A command to render a single source
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+pub struct RenderSourcesCommand<T> {
+    /// A human-readable name.
+    pub debug_name: String,
+    /// The dataflow's ID.
+    pub dataflow_id: uuid::Uuid,
+    /// An optional frontier to which the input should be advanced.
+    pub as_of: Option<Antichain<T>>,
+    /// Sources instantiations made available to the dataflow.
+    pub source_imports: BTreeMap<GlobalId, SourceInstanceDesc<T>>,
+}
+
 /// Commands related to the ingress and egress of collections.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub enum StorageCommand<T = mz_repr::Timestamp> {
     /// Create the enumerated sources, each associated with its identifier.
     CreateSources(Vec<CreateSourceCommand<T>>),
     /// Render the enumerated sources.
-    ///
-    /// Each source has a name for debugging purposes, an optional "as of" frontier and collection
-    /// of sources to import.
-    RenderSources(
-        Vec<(
-            /* debug_name */ String,
-            /* dataflow_id */ uuid::Uuid,
-            /* as_of */ Option<Antichain<T>>,
-            /* source_imports*/ BTreeMap<GlobalId, SourceInstanceDesc<T>>,
-        )>,
-    ),
+    RenderSources(Vec<RenderSourcesCommand<T>>),
     /// Enable compaction in storage-managed collections.
     ///
     /// Each entry in the vector names a collection and provides a frontier after which

src/storage/src/boundary/boundary.rs

@@ -114,10 +114,13 @@ pub use boundary_hook::BoundaryHook;
 mod boundary_hook {
     use std::collections::BTreeMap;
     use std::fmt;
+    use std::iter::once;
 
     use async_trait::async_trait;
 
-    use mz_dataflow_types::client::{GenericClient, StorageCommand, StorageResponse};
+    use mz_dataflow_types::client::{
+        GenericClient, RenderSourcesCommand, StorageCommand, StorageResponse,
+    };
     use mz_dataflow_types::sources::SourceDesc;
     use mz_dataflow_types::{SourceInstanceDesc, SourceInstanceId, SourceInstanceRequest};
     use mz_expr::GlobalId;
@@ -169,23 +172,22 @@ mod boundary_hook {
                 for source in sources.iter() {
                     if let Some(requests) = self.pending.remove(&source.id) {
                         render_requests.extend(requests.into_iter().map(|request| {
-                            (
-                                format!(
+                            RenderSourcesCommand {
+                                debug_name: format!(
                                     "SourceDataflow({:?}, {:?})",
                                     request.dataflow_id, request.source_id
                                 ),
-                                request.dataflow_id,
-                                Some(request.as_of.clone()),
-                                Some((
+                                dataflow_id: request.dataflow_id,
+                                as_of: Some(request.as_of.clone()),
+                                source_imports: once((
                                     request.source_id,
                                     SourceInstanceDesc {
                                         description: source.desc.clone(),
                                         arguments: request.arguments,
                                     },
                                 ))
-                                .into_iter()
                                 .collect(),
-                            )
+                            }
                         }));
                     }
                     self.sources.insert(source.id, source.desc.clone());
@@ -211,15 +213,15 @@ mod boundary_hook {
                             let unique_id = request.unique_id();
                             if !self.suppress.contains_key(&unique_id) {
                                 if let Some(source) = self.sources.get(&request.source_id) {
-                                    let command = StorageCommand::RenderSources(vec![(
-                                        format!("SourceDataflow({:?}, {:?})", request.dataflow_id, request.source_id),
-                                        request.dataflow_id,
-                                        Some(request.as_of.clone()),
-                                        Some((request.source_id, SourceInstanceDesc {
+                                    let command = StorageCommand::RenderSources(vec![RenderSourcesCommand {
+                                        debug_name: format!("SourceDataflow({:?}, {:?})", request.dataflow_id, request.source_id),
+                                        dataflow_id: request.dataflow_id,
+                                        as_of: Some(request.as_of.clone()),
+                                        source_imports: once((request.source_id, SourceInstanceDesc {
                                             description: source.clone(),
                                             arguments: request.arguments,
-                                        })).into_iter().collect(),
-                                    )]);
+                                        })).collect(),
+                                    }]);
                                     self.client.send(command).await.unwrap()
                                 } else {
                                     self.pending.entry(request.source_id).or_insert(Vec::new()).push(request);

src/storage/src/storage_state.rs

@@ -6,7 +6,7 @@
 //! Worker-local state for storage timely instances.
 
 use std::cell::RefCell;
-use std::collections::{BTreeMap, HashMap};
+use std::collections::HashMap;
 use std::rc::{Rc, Weak};
 use std::time::{Duration, Instant};
 
@@ -22,11 +22,11 @@ use tokio::sync::mpsc;
 use tracing::{debug, trace};
 
 use mz_dataflow_types::client::{
-    CreateSourceCommand, StorageCommand, StorageResponse, TimestampBindingFeedback,
+    CreateSourceCommand, RenderSourcesCommand, StorageCommand, StorageResponse,
+    TimestampBindingFeedback,
 };
 use mz_dataflow_types::sources::AwsExternalId;
 use mz_dataflow_types::sources::{ExternalSourceConnector, SourceConnector};
-use mz_dataflow_types::SourceInstanceDesc;
 use mz_expr::{GlobalId, PartitionId};
 use mz_ore::now::NowFn;
 use mz_persist::client::RuntimeClient;
@@ -117,7 +117,7 @@ pub struct ActiveStorageState<'a, A: Allocate, B: StorageCapture> {
 
 impl<'a, A: Allocate, B: StorageCapture> ActiveStorageState<'a, A, B> {
     /// Sets up the timestamp binding machinery if needed for this source
-    fn setup_timestamp_binding_state(&mut self, source: &CreateSourceCommand<u64>) {
+    fn setup_timestamp_binding_state(&mut self, source: &CreateSourceCommand<Timestamp>) {
         let ts_history = if let SourceConnector::External {
             connector,
             ts_frequency,
@@ -347,16 +347,14 @@ impl<'a, A: Allocate, B: StorageCapture> ActiveStorageState<'a, A, B> {
         }
     }
 
-    fn build_storage_dataflow(
-        &mut self,
-        dataflows: Vec<(
-            String,
-            uuid::Uuid,
-            Option<Antichain<Timestamp>>,
-            BTreeMap<GlobalId, SourceInstanceDesc>,
-        )>,
-    ) {
-        for (debug_name, dataflow_id, as_of, source_imports) in dataflows {
+    fn build_storage_dataflow(&mut self, dataflows: Vec<RenderSourcesCommand<Timestamp>>) {
+        for RenderSourcesCommand {
+            debug_name,
+            dataflow_id,
+            as_of,
+            source_imports,
+        } in dataflows
+        {
             crate::render::build_storage_dataflow(
                 self.timely_worker,
                 &mut self.storage_state,