Clean up token use in dataflow (MaterializeInc#10521)

frankmcsherry · web-flow · commit 942e5a36d150 · 2022-02-08T09:54:25.000-05:00
* Clean up token use in dataflow

* update comments
diff --git a/src/dataflow/src/render/mod.rs b/src/dataflow/src/render/mod.rs
@@ -197,21 +197,6 @@ pub struct StorageState {
     pub persist: Option<RuntimeClient>,
 }
 
-/// A container for "tokens" that are relevant to an in-construction dataflow.
-///
-/// Tokens are used by consumers of data to keep their sources of data running.
-/// Once all tokens referencing a source are dropped, the source can shut down,
-/// which will wind down (eventually) the dataflow containing it.
-#[derive(Default)]
-pub struct RelevantTokens {
-    /// The source tokens for all sources that have been built in this context.
-    pub source_tokens: HashMap<GlobalId, Rc<Option<SourceToken>>>,
-    /// Any other tokens that need to be dropped when an object is dropped.
-    pub additional_tokens: HashMap<GlobalId, Vec<Rc<dyn Any>>>,
-    /// Tokens for CDCv2 capture sources that have been built in this context.
-    pub cdc_tokens: HashMap<GlobalId, Rc<dyn Any>>,
-}
-
 /// Information about each source that must be communicated between storage and compute layers.
 pub struct SourceBoundary {
     /// Captured `row` updates representing a differential collection.
@@ -222,9 +207,7 @@ pub struct SourceBoundary {
         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>>,
+    pub token: Rc<dyn std::any::Any>,
 }
 
 /// Assemble the "storage" side of a dataflow, i.e. the sources.
@@ -267,19 +250,18 @@ pub fn build_storage_dataflow<A: Allocate>(
 
             // Import declared sources into the rendering context.
             for (src_id, source) in &dataflow.source_imports {
-                let ((ok, err), (source_token, additional_tokens)) =
-                    crate::render::sources::import_source(
-                        &debug_name,
-                        dataflow_id,
-                        &as_of,
-                        source.clone(),
-                        storage_state,
-                        region,
-                        materialized_logging.clone(),
-                        src_id.clone(),
-                        now.clone(),
-                        source_metrics,
-                    );
+                let ((ok, err), token) = crate::render::sources::import_source(
+                    &debug_name,
+                    dataflow_id,
+                    &as_of,
+                    source.clone(),
+                    storage_state,
+                    region,
+                    materialized_logging.clone(),
+                    src_id.clone(),
+                    now.clone(),
+                    source_metrics,
+                );
 
                 let ok_activator = RcActivator::new(format!("{debug_name}-ok"), 1);
                 let err_activator = RcActivator::new(format!("{debug_name}-err"), 1);
@@ -294,8 +276,7 @@ pub fn build_storage_dataflow<A: Allocate>(
                     SourceBoundary {
                         ok: ActivatedEventPusher::<_>::clone(&ok_handle),
                         err: ActivatedEventPusher::<_>::clone(&err_handle),
-                        token: source_token,
-                        additional_tokens,
+                        token,
                     },
                 );
 
@@ -330,7 +311,7 @@ pub fn build_compute_dataflow<A: Allocate>(
         // 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 mut tokens = BTreeMap::new();
 
             // Import declared sources into the rendering context.
             for (source_id, source) in sources.into_iter() {
@@ -357,13 +338,8 @@ pub fn build_compute_dataflow<A: Allocate>(
                 );
 
                 // Associate returned tokens with the source identifier.
-                let prior = tokens.source_tokens.insert(source_id, source.token);
+                let prior = tokens.insert(source_id, source.token);
                 assert!(prior.is_none());
-                tokens
-                    .additional_tokens
-                    .entry(source_id)
-                    .or_insert_with(Vec::new)
-                    .extend(source.additional_tokens);
             }
 
             // Import declared indexes into the rendering context.
@@ -423,7 +399,7 @@ where
     fn import_index(
         &mut self,
         compute_state: &mut ComputeState,
-        tokens: &mut RelevantTokens,
+        tokens: &mut BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
         scope: &mut G,
         region: &mut Child<'g, G, G::Timestamp>,
         idx_id: GlobalId,
@@ -450,15 +426,10 @@ where
                     ArrangementFlavor::Trace(idx_id, ok_arranged, err_arranged),
                 ),
             );
-            tokens
-                .additional_tokens
-                .entry(idx_id)
-                .or_insert_with(Vec::new)
-                .push(Rc::new((
-                    ok_button.press_on_drop(),
-                    err_button.press_on_drop(),
-                    token,
-                )));
+            tokens.insert(
+                idx_id,
+                Rc::new((ok_button.press_on_drop(), err_button.press_on_drop(), token)),
+            );
         } else {
             panic!(
                 "import of index {} failed while building dataflow {}",
@@ -480,23 +451,18 @@ where
     fn export_index(
         &mut self,
         compute_state: &mut ComputeState,
-        tokens: &mut RelevantTokens,
+        tokens: &mut BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
         import_ids: HashSet<GlobalId>,
         idx_id: GlobalId,
         idx: &IndexDesc,
     ) {
         // put together tokens that belong to the export
-        let mut needed_source_tokens = Vec::new();
-        let mut needed_additional_tokens = Vec::new();
+        let mut needed_tokens = Vec::new();
         for import_id in import_ids {
-            if let Some(addls) = tokens.additional_tokens.get(&import_id) {
-                needed_additional_tokens.extend_from_slice(addls);
-            }
-            if let Some(source_token) = tokens.source_tokens.get(&import_id) {
-                needed_source_tokens.push(Rc::clone(&source_token));
+            if let Some(token) = tokens.get(&import_id) {
+                needed_tokens.push(Rc::clone(&token));
             }
         }
-        let tokens = Rc::new((needed_source_tokens, needed_additional_tokens));
         let bundle = self.lookup_id(Id::Global(idx_id)).unwrap_or_else(|| {
             panic!(
                 "Arrangement alarmingly absent! id: {:?}",
@@ -507,7 +473,7 @@ where
             Some(ArrangementFlavor::Local(oks, errs)) => {
                 compute_state.traces.set(
                     idx_id,
-                    TraceBundle::new(oks.trace, errs.trace).with_drop(tokens),
+                    TraceBundle::new(oks.trace, errs.trace).with_drop(needed_tokens),
                 );
             }
             Some(ArrangementFlavor::Trace(gid, _, _)) => {
diff --git a/src/dataflow/src/render/sinks.rs b/src/dataflow/src/render/sinks.rs
@@ -24,7 +24,6 @@ use interchange::envelopes::{combine_at_timestamp, dbz_format, upsert_format};
 use repr::{Datum, Diff, Row, Timestamp};
 
 use crate::render::context::Context;
-use crate::render::RelevantTokens;
 use crate::sink::SinkBaseMetrics;
 
 impl<G> Context<G, Row, Timestamp>
@@ -35,7 +34,7 @@ where
     pub(crate) fn export_sink(
         &mut self,
         compute_state: &mut crate::render::ComputeState,
-        tokens: &mut RelevantTokens,
+        tokens: &mut std::collections::BTreeMap<GlobalId, Rc<dyn std::any::Any>>,
         import_ids: HashSet<GlobalId>,
         sink_id: GlobalId,
         sink: &SinkDesc,
@@ -44,15 +43,10 @@ where
         let sink_render = get_sink_render_for(&sink.connector);
 
         // put together tokens that belong to the export
-        let mut needed_source_tokens = Vec::new();
-        let mut needed_additional_tokens = Vec::new();
-        let mut needed_sink_tokens = Vec::new();
+        let mut needed_tokens = Vec::new();
         for import_id in import_ids {
-            if let Some(addls) = tokens.additional_tokens.get(&import_id) {
-                needed_additional_tokens.extend_from_slice(addls);
-            }
-            if let Some(source_token) = tokens.source_tokens.get(&import_id) {
-                needed_source_tokens.push(Rc::clone(&source_token));
+            if let Some(token) = tokens.get(&import_id) {
+                needed_tokens.push(Rc::clone(&token))
             }
         }
 
@@ -89,17 +83,12 @@ where
             sink_render.render_continuous_sink(compute_state, sink, sink_id, collection, metrics);
 
         if let Some(sink_token) = sink_token {
-            needed_sink_tokens.push(sink_token);
+            needed_tokens.push(sink_token);
         }
 
-        let tokens = Rc::new((
-            needed_sink_tokens,
-            needed_source_tokens,
-            needed_additional_tokens,
-        ));
         compute_state
             .dataflow_tokens
-            .insert(sink_id, Box::new(tokens));
+            .insert(sink_id, Box::new(needed_tokens));
     }
 }
 
@@ -231,7 +220,7 @@ where
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
         metrics: &SinkBaseMetrics,
-    ) -> Option<Box<dyn Any>>
+    ) -> Option<Rc<dyn Any>>
     where
         G: Scope<Timestamp = Timestamp>;
 }
diff --git a/src/dataflow/src/render/sources.rs b/src/dataflow/src/render/sources.rs
@@ -129,8 +129,8 @@ where
 
 /// Constructs a `CollectionBundle` and tokens from source arguments.
 ///
-/// The first return value is the collection bundle, and the second a pair of source and additional
-/// tokens, that are used to control the demolition of the source.
+/// The first returned pair are the row and error collections, and the
+/// second is a token that will keep the source alive as long as it is held.
 pub(crate) fn import_source<G>(
     dataflow_debug_name: &String,
     dataflow_id: usize,
@@ -148,10 +148,7 @@ pub(crate) fn import_source<G>(
     base_metrics: &SourceBaseMetrics,
 ) -> (
     (Collection<G, Row>, Collection<G, DataflowError>),
-    (
-        Rc<Option<crate::source::SourceToken>>,
-        Vec<Rc<dyn std::any::Any>>,
-    ),
+    Rc<dyn std::any::Any>,
 )
 where
     G: Scope<Timestamp = Timestamp>,
@@ -164,7 +161,7 @@ where
     }
 
     // Tokens that we should return from the method.
-    let mut additional_tokens: Vec<Rc<dyn std::any::Any>> = Vec::new();
+    let mut needed_tokens: Vec<Rc<dyn std::any::Any>> = Vec::new();
 
     // Before proceeding, we may need to remediate sources with non-trivial relational
     // expressions that post-process the bare source. If the expression is trivial, a
@@ -182,7 +179,7 @@ where
             storage_state.local_inputs.insert(src_id, local_input);
 
             // TODO(mcsherry): Local tables are a special non-source we should relocate.
-            ((ok, err), (Rc::new(None), Vec::new()))
+            ((ok, err), Rc::new(()))
         }
 
         SourceConnector::External {
@@ -377,7 +374,7 @@ where
                             schema_registry_config,
                             confluent_wire_format,
                         );
-                        additional_tokens.push(Rc::new(token));
+                        needed_tokens.push(Rc::new(token));
                         (oks, None)
                     } else {
                         let (results, extra_token) = match ok_source {
@@ -401,7 +398,7 @@ where
                             ),
                         };
                         if let Some(tok) = extra_token {
-                            additional_tokens.push(Rc::new(tok));
+                            needed_tokens.push(Rc::new(tok));
                         }
 
                         // render envelopes
@@ -467,7 +464,7 @@ where
                                         &src_id,
                                         &source_name,
                                         storage_state,
-                                        &mut additional_tokens,
+                                        &mut needed_tokens,
                                     );
 
                                     (sealed_upsert, upsert_err)
@@ -518,7 +515,7 @@ where
                                             &src_id,
                                             &source_name,
                                             storage_state,
-                                            &mut additional_tokens,
+                                            &mut needed_tokens,
                                         );
 
                                         // NOTE: Persistence errors don't go through the same
@@ -639,11 +636,10 @@ where
                 .or_insert_with(Vec::new)
                 .push(Rc::downgrade(&source_token));
 
-            // Return the source token for capability manipulation, and any additional tokens.
-            (
-                (collection, err_collection),
-                (source_token, additional_tokens),
-            )
+            needed_tokens.push(source_token);
+
+            // Return the collections and any needed tokens.
+            ((collection, err_collection), Rc::new(needed_tokens))
         }
     }
 }
@@ -799,7 +795,7 @@ fn seal_and_await<G, D1, D2, K1, V1, K2, V2>(
     source_id: &GlobalId,
     source_name: &str,
     storage_state: &mut StorageState,
-    additional_tokens: &mut Vec<Rc<dyn std::any::Any>>,
+    needed_tokens: &mut Vec<Rc<dyn std::any::Any>>,
 ) -> Stream<G, (D1, Timestamp, Diff)>
 where
     G: Scope<Timestamp = Timestamp>,
@@ -844,7 +840,7 @@ where
         compaction_handle,
     );
 
-    additional_tokens.push(source_token);
+    needed_tokens.push(source_token);
 
     // Don't send data forward to "dataflow" until the frontier
     // tells us that we both persisted and sealed it.
diff --git a/src/dataflow/src/sink/avro_ocf.rs b/src/dataflow/src/sink/avro_ocf.rs
@@ -9,6 +9,7 @@
 
 use std::any::Any;
 use std::fs::OpenOptions;
+use std::rc::Rc;
 
 use differential_dataflow::{Collection, Hashable};
 
@@ -50,7 +51,7 @@ where
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
         _metrics: &SinkBaseMetrics,
-    ) -> Option<Box<dyn Any>>
+    ) -> Option<Rc<dyn Any>>
     where
         G: Scope<Timestamp = Timestamp>,
     {
diff --git a/src/dataflow/src/sink/kafka.rs b/src/dataflow/src/sink/kafka.rs
@@ -87,7 +87,7 @@ where
         sink_id: GlobalId,
         sinked_collection: Collection<G, (Option<Row>, Option<Row>), Diff>,
         metrics: &SinkBaseMetrics,
-    ) -> Option<Box<dyn Any>>
+    ) -> Option<Rc<dyn Any>>
     where
         G: Scope<Timestamp = Timestamp>,
     {
@@ -950,7 +950,7 @@ fn kafka<G>(
     as_of: SinkAsOf,
     write_frontier: Rc<RefCell<Antichain<Timestamp>>>,
     metrics: &KafkaBaseMetrics,
-) -> Box<dyn Any>
+) -> Rc<dyn Any>
 where
     G: Scope<Timestamp = Timestamp>,
 {
@@ -1027,7 +1027,7 @@ pub fn produce_to_kafka<G>(
     shared_gate_ts: Rc<Cell<Option<Timestamp>>>,
     write_frontier: Rc<RefCell<Antichain<Timestamp>>>,
     metrics: &KafkaBaseMetrics,
-) -> Box<dyn Any>
+) -> Rc<dyn Any>
 where
     G: Scope<Timestamp = Timestamp>,
 {
@@ -1287,7 +1287,7 @@ where
         }),
     );
 
-    Box::new(KafkaSinkToken { shutdown_flag })
+    Rc::new(KafkaSinkToken { shutdown_flag })
 }
 
 /// Encodes a stream of `(Option<Row>, Option<Row>)` updates using the specified encoder.
diff --git a/src/dataflow/src/sink/tail.rs b/src/dataflow/src/sink/tail.rs
diff --git a/src/dataflow/src/source/mod.rs b/src/dataflow/src/source/mod.rs
