compute: controllable mz_join_core yielding

This commit provides `mz_join_core` with a yield functions, inspired by
`persist_source` and DD's half join operator, that allows the caller to
control the yield behavior by time or by amount of work performed.

Author: teskje

src/compute/src/render/join/linear_join.rs

@@ -67,7 +67,9 @@ impl LinearJoinImpl {
             Self::DifferentialDataflow => {
                 differential_dataflow::operators::JoinCore::join_core(arranged1, arranged2, result)
             }
-            Self::Materialize => mz_join_core(arranged1, arranged2, result),
+            Self::Materialize => mz_join_core(arranged1, arranged2, result, |_start, work| {
+                work >= 1_000_000
+            }),
         }
     }
 }

src/compute/src/render/join/mz_join_core.rs

@@ -38,6 +38,7 @@
 
 use std::cmp::Ordering;
 use std::collections::VecDeque;
+use std::time::Instant;
 
 use differential_dataflow::consolidation::consolidate_updates;
 use differential_dataflow::difference::Multiply;
@@ -61,10 +62,11 @@ use timely::PartialOrder;
 /// Each matching pair of records `(key, val1)` and `(key, val2)` are subjected to the `result` function,
 /// which produces something implementing `IntoIterator`, where the output collection will have an entry for
 /// every value returned by the iterator.
-pub(super) fn mz_join_core<G, Tr1, Tr2, L, I, K, V1, V2>(
+pub(super) fn mz_join_core<G, Tr1, Tr2, L, I, K, V1, V2, YFn>(
     arranged1: &Arranged<G, Tr1>,
     arranged2: &Arranged<G, Tr2>,
     mut result: L,
+    yield_fn: YFn,
 ) -> Collection<G, I::Item, Diff>
 where
     G: Scope,
@@ -77,6 +79,7 @@ where
     K: Data,
     V1: Data,
     V2: Data,
+    YFn: Fn(Instant, usize) -> bool + 'static,
 {
     let mut trace1 = arranged1.trace.clone();
     let mut trace2 = arranged2.trace.clone();
@@ -278,24 +281,30 @@ where
                     // input must scan all batches from the other input).
 
                     // Perform some amount of outstanding work.
-                    let mut fuel = 1_000_000;
-                    while !todo1.is_empty() && fuel > 0 {
-                        todo1
-                            .front_mut()
-                            .unwrap()
-                            .work(output, &mut result, &mut fuel);
+                    let start_time = Instant::now();
+                    let mut work = 0;
+                    while !todo1.is_empty() && !yield_fn(start_time, work) {
+                        todo1.front_mut().unwrap().work(
+                            output,
+                            &mut result,
+                            |w| yield_fn(start_time, w),
+                            &mut work,
+                        );
                         if !todo1.front().unwrap().work_remains() {
                             todo1.pop_front();
                         }
                     }
 
                     // Perform some amount of outstanding work.
-                    let mut fuel = 1_000_000;
-                    while !todo2.is_empty() && fuel > 0 {
-                        todo2
-                            .front_mut()
-                            .unwrap()
-                            .work(output, &mut result, &mut fuel);
+                    let start_time = Instant::now();
+                    let mut work = 0;
+                    while !todo2.is_empty() && !yield_fn(start_time, work) {
+                        todo2.front_mut().unwrap().work(
+                            output,
+                            &mut result,
+                            |w| yield_fn(start_time, w),
+                            &mut work,
+                        );
                         if !todo2.front().unwrap().work_remains() {
                             todo2.pop_front();
                         }
@@ -405,14 +414,16 @@ where
     }
 
     /// Process keys until at least `fuel` output tuples produced, or the work is exhausted.
-    fn work<L, I>(
+    fn work<L, I, YFn>(
         &mut self,
         output: &mut OutputHandle<T, (D, T, Diff), Tee<T, (D, T, Diff)>>,
         mut result: L,
-        fuel: &mut usize,
+        yield_fn: YFn,
+        work: &mut usize,
     ) where
         I: IntoIterator<Item = D>,
         L: FnMut(&C1::Key, &C1::Val, &C2::Val) -> I,
+        YFn: Fn(usize) -> bool,
     {
         let meet = self.capability.time();
 
@@ -443,7 +454,7 @@ where
                 Ordering::Less => cursor1.seek_key(storage1, cursor2.key(storage2)),
                 Ordering::Greater => cursor2.seek_key(storage2, cursor1.key(storage1)),
                 Ordering::Equal => {
-                    // Populate `temp` with the results, as long as fuel remains.
+                    // Populate `temp` with the results, until we should yield.
                     let key = cursor2.key(storage2);
                     while let Some(val1) = cursor1.get_val(storage1) {
                         while let Some(val2) = cursor2.get_val(storage2) {
@@ -463,14 +474,14 @@ where
                         cursor1.step_val(storage1);
                         cursor2.rewind_vals(storage2);
 
-                        *fuel = fuel.saturating_sub(temp.len());
+                        *work.saturating_add(temp.len());
 
-                        if *fuel == 0 {
-                            // The fuel is exhausted, so we should yield. Returning here is only
-                            // allowed because we leave the cursors in a state that will let us
-                            // pick up the work correctly on the next invocation.
-                            *fuel += flush(temp, &mut session);
-                            if *fuel == 0 {
+                        if yield_fn(*work) {
+                            // Returning here is only allowed because we leave the cursors in a
+                            // state that will let us pick up the work correctly on the next
+                            // invocation.
+                            *work -= flush(temp, &mut session);
+                            if yield_fn(*work) {
                                 return;
                             }
                         }
@@ -483,7 +494,7 @@ where
         }
 
         if !temp.is_empty() {
-            *fuel += flush(temp, &mut session);
+            *work -= flush(temp, &mut session);
         }
 
         // We only get here after having iterated through all keys.