Partition storage and compute dataflows (reissue) (MaterializeInc#10531)

antiguru · frankmcsherry · web-flow · commit 087dba3584e8 · 2022-02-08T14:08:43.000Z
* Partition storage and compute dataflows (MaterializeInc#10519) * Partition storage and compute dataflows * Improve abstractions and comments * tidying comments * dataflow: Remove mut requirement for RcActivator The functions use inner mutability and don't require mutable access to the object. * dataflow: Event-based replay for sources Signed-off-by: Moritz Hoffmann <mh@materialize.com> * Wrap only push half with activator Co-authored-by: Frank McSherry <fmcsherry@me.com>
diff --git a/src/dataflow/src/activator.rs b/src/dataflow/src/activator.rs
@@ -47,20 +47,20 @@ impl RcActivator {
     }
 
     /// Register an additional [Activator] with this [RcActivator].
-    pub fn register(&mut self, activator: Activator) {
+    pub fn register(&self, activator: Activator) {
         self.inner.borrow_mut().register(activator)
     }
 
     /// Activate all contained activators.
     ///
     /// The implementation is free to ignore activations and only release them once a sufficient
     /// volume has been accumulated.
-    pub fn activate(&mut self) {
+    pub fn activate(&self) {
         self.inner.borrow_mut().activate()
     }
 
     /// Acknowledge the activation, which enables new activations to be scheduled.
-    pub fn ack(&mut self) {
+    pub fn ack(&self) {
         self.inner.borrow_mut().ack()
     }
 }
diff --git a/src/dataflow/src/event.rs b/src/dataflow/src/event.rs
@@ -0,0 +1,38 @@
+// Copyright Materialize, Inc. and contributors. All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+//! Traits and types for describing captured timely dataflow streams.
+//!
+//! This is roughly based on [timely::dataflow::operators::capture::event].
+
+use crate::activator::RcActivator;
+use timely::dataflow::operators::capture::{EventCore, EventPusherCore};
+
+/// An event pusher wrapper that activates targets on push.
+#[derive(Clone, Debug)]
+pub struct ActivatedEventPusher<E> {
+    pub inner: E,
+    pub activator: RcActivator,
+}
+
+impl<E> ActivatedEventPusher<E> {
+    /// Create a new activated event link wrapper.
+    ///
+    /// * inner: A wrapped event pusher/iterator.
+    pub fn new(inner: E, activator: RcActivator) -> Self {
+        Self { inner, activator }
+    }
+}
+
+impl<T, D, E: EventPusherCore<T, D>> EventPusherCore<T, D> for ActivatedEventPusher<E> {
+    fn push(&mut self, event: EventCore<T, D>) {
+        self.inner.push(event);
+        self.activator.activate();
+    }
+}
diff --git a/src/dataflow/src/lib.rs b/src/dataflow/src/lib.rs
@@ -14,6 +14,7 @@
 mod activator;
 mod arrangement;
 mod decode;
+mod event;
 mod metrics;
 mod operator;
 mod render;
diff --git a/src/dataflow/src/render/mod.rs b/src/dataflow/src/render/mod.rs
@@ -102,18 +102,23 @@
 
 use std::any::Any;
 use std::cell::RefCell;
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::rc::Rc;
 use std::rc::Weak;
+use std::time::Duration;
 
 use differential_dataflow::AsCollection;
 use timely::communication::Allocate;
+use timely::dataflow::operators::capture::EventLink;
 use timely::dataflow::operators::to_stream::ToStream;
+use timely::dataflow::operators::Capture;
 use timely::dataflow::scopes::Child;
 use timely::dataflow::Scope;
+
 use timely::progress::Antichain;
 use timely::worker::Worker as TimelyWorker;
 
+use crate::activator::RcActivator;
 use dataflow_types::*;
 use expr::{GlobalId, Id};
 use itertools::Itertools;
@@ -123,10 +128,12 @@ use persist::client::RuntimeClient;
 use repr::{Row, Timestamp};
 
 use crate::arrangement::manager::{TraceBundle, TraceManager};
+use crate::event::ActivatedEventPusher;
 use crate::metrics::Metrics;
 use crate::render::context::CollectionBundle;
 use crate::render::context::{ArrangementFlavor, Context};
 use crate::render::sources::PersistedSourceManager;
+use crate::replay::MzReplay;
 use crate::server::LocalInput;
 use crate::sink::SinkBaseMetrics;
 use crate::source::metrics::SourceBaseMetrics;
@@ -205,28 +212,47 @@ pub struct RelevantTokens {
     pub cdc_tokens: HashMap<GlobalId, Rc<dyn Any>>,
 }
 
-/// Build a dataflow from a description.
-pub fn build_dataflow<A: Allocate>(
+/// Information about each source that must be communicated between storage and compute layers.
+pub struct SourceBoundary {
+    /// Captured `row` updates representing a differential collection.
+    pub ok:
+        ActivatedEventPusher<Rc<EventLink<repr::Timestamp, (Row, repr::Timestamp, repr::Diff)>>>,
+    /// Captured error updates representing a differential collection.
+    pub err: ActivatedEventPusher<
+        Rc<EventLink<repr::Timestamp, (DataflowError, repr::Timestamp, repr::Diff)>>,
+    >,
+    /// A token that should be dropped to terminate the source.
+    pub token: Rc<Option<crate::source::SourceToken>>,
+    /// Additional tokens that should be dropped to terminate the source.
+    pub additional_tokens: Vec<Rc<dyn std::any::Any>>,
+}
+
+/// Assemble the "storage" side of a dataflow, i.e. the sources.
+///
+/// This method creates a new dataflow to host the implementations of sources for the `dataflow`
+/// argument, and returns assets for each source that can import the results into a new dataflow.
+pub fn build_storage_dataflow<A: Allocate>(
     timely_worker: &mut TimelyWorker<A>,
-    compute_state: &mut ComputeState,
     storage_state: &mut StorageState,
-    dataflow: DataflowDescription<plan::Plan>,
+    dataflow: &DataflowDescription<plan::Plan>,
     now: NowFn,
     source_metrics: &SourceBaseMetrics,
-    sink_metrics: &SinkBaseMetrics,
-) {
+) -> BTreeMap<GlobalId, SourceBoundary> {
     let worker_logging = timely_worker.log_register().get("timely");
     let name = format!("Dataflow: {}", &dataflow.debug_name);
     let materialized_logging = timely_worker.log_register().get("materialized");
 
+    let mut results = BTreeMap::new();
+
     timely_worker.dataflow_core(&name, worker_logging, Box::new(()), |_, scope| {
         // The scope.clone() occurs to allow import in the region.
         // We build a region here to establish a pattern of a scope inside the dataflow,
         // so that other similar uses (e.g. with iterative scopes) do not require weird
         // alternate type signatures.
         scope.clone().region_named(&name, |region| {
-            let mut context = Context::for_dataflow(&dataflow, scope.addr().into_element());
-            let mut tokens = RelevantTokens::default();
+            let as_of = dataflow.as_of.clone().unwrap();
+            let dataflow_id = scope.addr().into_element();
+            let debug_name = format!("{}-sources", dataflow.debug_name);
 
             assert!(
                 !dataflow
@@ -241,11 +267,11 @@ pub fn build_dataflow<A: Allocate>(
 
             // Import declared sources into the rendering context.
             for (src_id, source) in &dataflow.source_imports {
-                let (collection_bundle, (source_token, additional_tokens)) =
+                let ((ok, err), (source_token, additional_tokens)) =
                     crate::render::sources::import_source(
-                        &context.debug_name,
-                        context.dataflow_id,
-                        &context.as_of_frontier,
+                        &debug_name,
+                        dataflow_id,
+                        &as_of,
                         source.clone(),
                         storage_state,
                         region,
@@ -255,17 +281,89 @@ pub fn build_dataflow<A: Allocate>(
                         source_metrics,
                     );
 
+                let ok_activator = RcActivator::new(format!("{debug_name}-ok"), 1);
+                let err_activator = RcActivator::new(format!("{debug_name}-err"), 1);
+
+                let ok_handle =
+                    ActivatedEventPusher::new(Rc::new(EventLink::new()), ok_activator.clone());
+                let err_handle =
+                    ActivatedEventPusher::new(Rc::new(EventLink::new()), err_activator.clone());
+
+                results.insert(
+                    *src_id,
+                    SourceBoundary {
+                        ok: ActivatedEventPusher::<_>::clone(&ok_handle),
+                        err: ActivatedEventPusher::<_>::clone(&err_handle),
+                        token: source_token,
+                        additional_tokens,
+                    },
+                );
+
+                ok.inner.capture_into(ok_handle);
+                err.inner.capture_into(err_handle);
+            }
+        })
+    });
+
+    results
+}
+
+/// Assemble the "compute"  side of a dataflow, i.e. all but the sources.
+///
+/// This method imports sources from provided assets, and then builds the remaining
+/// dataflow using "compute-local" assets like shared arrangements, and producing
+/// both arrangements and sinks.
+pub fn build_compute_dataflow<A: Allocate>(
+    timely_worker: &mut TimelyWorker<A>,
+    compute_state: &mut ComputeState,
+    sources: BTreeMap<GlobalId, SourceBoundary>,
+    dataflow: DataflowDescription<plan::Plan>,
+    sink_metrics: &SinkBaseMetrics,
+) {
+    let worker_logging = timely_worker.log_register().get("timely");
+    let name = format!("Dataflow: {}", &dataflow.debug_name);
+
+    timely_worker.dataflow_core(&name, worker_logging, Box::new(()), |_, scope| {
+        // The scope.clone() occurs to allow import in the region.
+        // We build a region here to establish a pattern of a scope inside the dataflow,
+        // so that other similar uses (e.g. with iterative scopes) do not require weird
+        // alternate type signatures.
+        scope.clone().region_named(&name, |region| {
+            let mut context = Context::for_dataflow(&dataflow, scope.addr().into_element());
+            let mut tokens = RelevantTokens::default();
+
+            // Import declared sources into the rendering context.
+            for (source_id, source) in sources.into_iter() {
                 // Associate collection bundle with the source identifier.
-                context.insert_id(Id::Global(*src_id), collection_bundle);
+                let ok = Some(source.ok.inner)
+                    .mz_replay(
+                        region,
+                        &format!("{name}-{source_id}"),
+                        Duration::MAX,
+                        source.ok.activator,
+                    )
+                    .as_collection();
+                let err = Some(source.err.inner)
+                    .mz_replay(
+                        region,
+                        &format!("{name}-{source_id}-err"),
+                        Duration::MAX,
+                        source.err.activator,
+                    )
+                    .as_collection();
+                context.insert_id(
+                    Id::Global(source_id),
+                    CollectionBundle::from_collections(ok, err),
+                );
 
                 // Associate returned tokens with the source identifier.
-                let prior = tokens.source_tokens.insert(*src_id, source_token);
+                let prior = tokens.source_tokens.insert(source_id, source.token);
                 assert!(prior.is_none());
                 tokens
                     .additional_tokens
-                    .entry(*src_id)
+                    .entry(source_id)
                     .or_insert_with(Vec::new)
-                    .extend(additional_tokens);
+                    .extend(source.additional_tokens);
             }
 
             // Import declared indexes into the rendering context.
diff --git a/src/dataflow/src/render/sources.rs b/src/dataflow/src/render/sources.rs
@@ -69,7 +69,7 @@ pub(crate) fn import_table<G>(
     persisted_name: Option<String>,
 ) -> (
     LocalInput,
-    crate::render::CollectionBundle<G, Row, Timestamp>,
+    (Collection<G, Row>, Collection<G, DataflowError>),
 )
 where
     G: Scope<Timestamp = Timestamp>,
@@ -124,10 +124,7 @@ where
         })
         .as_collection();
 
-    let collection_bundle =
-        crate::render::CollectionBundle::from_collections(ok_collection, err_collection);
-
-    (local_input, collection_bundle)
+    (local_input, (ok_collection, err_collection))
 }
 
 /// Constructs a `CollectionBundle` and tokens from source arguments.
@@ -150,7 +147,7 @@ pub(crate) fn import_source<G>(
     now: NowFn,
     base_metrics: &SourceBaseMetrics,
 ) -> (
-    crate::render::CollectionBundle<G, Row, Timestamp>,
+    (Collection<G, Row>, Collection<G, DataflowError>),
     (
         Rc<Option<crate::source::SourceToken>>,
         Vec<Rc<dyn std::any::Any>>,
@@ -180,12 +177,12 @@ where
         // Create a new local input (exposed as TABLEs to users). Data is inserted
         // via Command::Insert commands.
         SourceConnector::Local { persisted_name, .. } => {
-            let (local_input, collection_bundle) =
+            let (local_input, (ok, err)) =
                 import_table(as_of_frontier, storage_state, scope, persisted_name);
             storage_state.local_inputs.insert(src_id, local_input);
 
             // TODO(mcsherry): Local tables are a special non-source we should relocate.
-            (collection_bundle, (Rc::new(None), Vec::new()))
+            ((ok, err), (Rc::new(None), Vec::new()))
         }
 
         SourceConnector::External {
@@ -632,10 +629,6 @@ where
             use differential_dataflow::operators::consolidate::ConsolidateStream;
             collection = collection.consolidate_stream();
 
-            // Introduce the stream by name, as an unarranged collection.
-            let collection_bundle =
-                crate::render::CollectionBundle::from_collections(collection, err_collection);
-
             let source_token = Rc::new(capability);
 
             // We also need to keep track of this mapping globally to activate sources
@@ -647,7 +640,10 @@ where
                 .push(Rc::downgrade(&source_token));
 
             // Return the source token for capability manipulation, and any additional tokens.
-            (collection_bundle, (source_token, additional_tokens))
+            (
+                (collection, err_collection),
+                (source_token, additional_tokens),
+            )
         }
     }
 }
diff --git a/src/dataflow/src/replay.rs b/src/dataflow/src/replay.rs
@@ -43,12 +43,19 @@ where
     I: IntoIterator,
     <I as IntoIterator>::Item: EventIterator<T, D> + 'static,
 {
+    /// Replay a collection of [EventIterator]s into a Timely stream.
+    ///
+    /// * `scope`: The [Scope] to replay into.
+    /// * `name`: Human-readable debug name of the Timely operator.
+    /// * `period`: Reschedule the operator once the period has elapsed.
+    ///    Provide [Duration::MAX] to disable periodic scheduling.
+    /// * `rc_activator`: An activator to trigger the operator.
     fn mz_replay<S: Scope<Timestamp = T>>(
         self,
         scope: &mut S,
         name: &str,
         period: Duration,
-        mut rc_activator: RcActivator,
+        rc_activator: RcActivator,
     ) -> Stream<S, D> {
         let name = format!("Replay {}", name);
         let mut builder = OperatorBuilder::new(name, scope.clone());
@@ -68,9 +75,15 @@ where
 
         builder.build(move |progress| {
             rc_activator.ack();
-            if last_active + period <= Instant::now() || !started {
+            if last_active
+                .checked_add(period)
+                .map_or(false, |next_active| next_active <= Instant::now())
+                || !started
+            {
                 last_active = Instant::now();
-                activator.activate_after(period);
+                if period < Duration::MAX {
+                    activator.activate_after(period);
+                }
             }
 
             if !started {
diff --git a/src/dataflow/src/server.rs b/src/dataflow/src/server.rs

Original file line number	Diff line number	Diff line change
`@@ -47,20 +47,20 @@ impl RcActivator {`
`47`	`47`	`}`
`48`	`48`
`49`	`49`	`/// Register an additional [Activator] with this [RcActivator].`
`50`		`- pub fn register(&mut self, activator: Activator) {`
	`50`	`+ pub fn register(&self, activator: Activator) {`
`51`	`51`	`self.inner.borrow_mut().register(activator)`
`52`	`52`	`}`
`53`	`53`
`54`	`54`	`/// Activate all contained activators.`
`55`	`55`	`///`
`56`	`56`	`/// The implementation is free to ignore activations and only release them once a sufficient`
`57`	`57`	`/// volume has been accumulated.`
`58`		`- pub fn activate(&mut self) {`
	`58`	`+ pub fn activate(&self) {`
`59`	`59`	`self.inner.borrow_mut().activate()`
`60`	`60`	`}`
`61`	`61`
`62`	`62`	`/// Acknowledge the activation, which enables new activations to be scheduled.`
`63`		`- pub fn ack(&mut self) {`
	`63`	`+ pub fn ack(&self) {`
`64`	`64`	`self.inner.borrow_mut().ack()`
`65`	`65`	`}`
`66`	`66`	`}`