Merge pull request #668 from antiguru/exchange_flush

Optionally flush container builders

Author: frankmcsherry

container/src/lib.rs

@@ -125,6 +125,14 @@ pub trait ContainerBuilder: Default + 'static {
         }
         container.clear();
     }
+
+    /// Indicates a good moment to release resources.
+    ///
+    /// By default, does nothing. Callers first needs to drain the contents using [`Self::finish`]
+    /// before calling this function. The implementation should not change the contents of the
+    /// builder.
+    #[inline]
+    fn relax(&mut self) { }
 }
 
 /// A wrapper trait indicating that the container building will preserve the number of records.

timely/examples/columnar.rs

@@ -300,6 +300,7 @@ mod builder {
     }
 
     impl<C: Columnar> Default for ColumnBuilder<C> {
+        #[inline]
         fn default() -> Self {
             ColumnBuilder {
                 current: Default::default(),
@@ -331,6 +332,15 @@ mod builder {
             self.empty = self.pending.pop_front();
             self.empty.as_mut()
         }
+
+        #[inline]
+        fn relax(&mut self) {
+            // The caller is responsible for draining all contents; assert that we are empty.
+            // The assertion is not strictly necessary, but it helps catch bugs.
+            assert!(self.current.is_empty());
+            assert!(self.pending.is_empty());
+            *self = Self::default();
+        }
     }
 
     impl<C: Columnar> LengthPreservingContainerBuilder for ColumnBuilder<C> where C::Container: Clone { }

timely/src/dataflow/channels/pushers/buffer.rs

@@ -81,6 +81,7 @@ impl<T, CB: ContainerBuilder, P: Push<Message<T, CB::Container>>> Buffer<T, CB,
     /// Flushes all data and pushes a `None` to `self.pusher`, indicating a flush.
     pub fn cease(&mut self) {
         self.flush();
+        self.builder.relax();
         self.pusher.push(&mut None);
     }
 

timely/src/dataflow/channels/pushers/exchange.rs

@@ -14,7 +14,7 @@ where
     for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
 {
     pushers: Vec<P>,
-    buffers: Vec<CB>,
+    builders: Vec<CB>,
     current: Option<T>,
     hash_func: H,
 }
@@ -27,20 +27,20 @@ where
 {
     /// Allocates a new `Exchange` from a supplied set of pushers and a distribution function.
     pub fn new(pushers: Vec<P>, key: H) -> Exchange<T, CB, P, H> {
-        let mut buffers = vec![];
+        let mut builders = vec![];
         for _ in 0..pushers.len() {
-            buffers.push(Default::default());
+            builders.push(Default::default());
         }
         Exchange {
             pushers,
             hash_func: key,
-            buffers,
+            builders,
             current: None,
         }
     }
     #[inline]
     fn flush(&mut self, index: usize) {
-        while let Some(container) = self.buffers[index].finish() {
+        while let Some(container) = self.builders[index].finish() {
             if let Some(ref time) = self.current {
                 Message::push_at(container, time.clone(), &mut self.pushers[index]);
             }
@@ -79,14 +79,14 @@ where
             // if the number of pushers is a power of two, use a mask
             if self.pushers.len().is_power_of_two() {
                 let mask = (self.pushers.len() - 1) as u64;
-                CB::partition(data, &mut self.buffers, |datum| ((hash_func)(datum) & mask) as usize);
+                CB::partition(data, &mut self.builders, |datum| ((hash_func)(datum) & mask) as usize);
             }
             // as a last resort, use mod (%)
             else {
                 let num_pushers = self.pushers.len() as u64;
-                CB::partition(data, &mut self.buffers, |datum| ((hash_func)(datum) % num_pushers) as usize);
+                CB::partition(data, &mut self.builders, |datum| ((hash_func)(datum) % num_pushers) as usize);
             }
-            for (buffer, pusher) in self.buffers.iter_mut().zip(self.pushers.iter_mut()) {
+            for (buffer, pusher) in self.builders.iter_mut().zip(self.pushers.iter_mut()) {
                 while let Some(container) = buffer.extract() {
                     Message::push_at(container, time.clone(), pusher);
                 }
@@ -96,6 +96,7 @@ where
             // flush
             for index in 0..self.pushers.len() {
                 self.flush(index);
+                self.builders[index].relax();
                 self.pushers[index].push(&mut None);
             }
         }