Pass number of peers at construction time

Signed-off-by: Moritz Hoffmann <[email protected]>

Author: antiguru

timely/src/dataflow/channels/pact.rs

@@ -96,7 +96,7 @@ where
     fn connect<A: AsWorker>(self, allocator: &mut A, identifier: usize, address: Rc<[usize]>, logging: Option<Logger>) -> (Self::Pusher, Self::Puller) {
         let (senders, receiver) = allocator.allocate::<Message<T, CB::Container>>(identifier, address);
         let senders = senders.into_iter().enumerate().map(|(i,x)| LogPusher::new(x, allocator.index(), i, identifier, logging.clone())).collect::<Vec<_>>();
-        let distributor = DrainContainerDistributor::new(self.hash_func);
+        let distributor = DrainContainerDistributor::new(self.hash_func, allocator.peers());
         (ExchangePusher::new(senders, distributor), LogPuller::new(receiver, allocator.index(), identifier, logging.clone()))
     }
 }

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

@@ -31,10 +31,11 @@ pub struct DrainContainerDistributor<CB, H> {
 }
 
 impl<CB: Default, H> DrainContainerDistributor<CB, H> {
-    /// Constructs a new `DrainContainerDistributor` with the given hash function.
-    pub fn new(hash_func: H) -> Self {
+    /// Constructs a new `DrainContainerDistributor` with the given hash function for a number of
+    /// peers.
+    pub fn new(hash_func: H, peers: usize) -> Self {
         Self {
-            builders: Vec::new(),
+            builders: std::iter::repeat_with(Default::default).take(peers).collect(),
             hash_func,
         }
     }
@@ -46,12 +47,7 @@ where
     for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64,
 {
     fn partition<T: Clone, P: Push<Message<T, CB::Container>>>(&mut self, container: &mut CB::Container, time: &T, pushers: &mut [P]) {
-        if self.builders.len() <= pushers.len() {
-            self.builders.resize_with(pushers.len(), Default::default);
-        }
-        else {
-            debug_assert_eq!(self.builders.len(), pushers.len());
-        }
+        debug_assert_eq!(self.builders.len(), pushers.len());
         if pushers.len().is_power_of_two() {
             let mask = (pushers.len() - 1) as u64;
             for datum in container.drain() {