Zero-copy for networked processes (#642)

* Zero-copy for networked processes

Introduce a intra- and inter-process serializing configuration variant.
Previously, it wasn't possible to construct a networked configuration that
would use zero-copy for process-local communication.

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

* Oversight

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

---------

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

Author: antiguru

communication/src/allocator/generic.rs

@@ -25,6 +25,8 @@ pub enum Generic {
     ProcessBinary(ProcessAllocator),
     /// Inter-process allocator.
     ZeroCopy(TcpAllocator<Process>),
+    /// Inter-process allocator, intra-process serializing allocator.
+    ZeroCopyBinary(TcpAllocator<ProcessAllocator>),
 }
 
 impl Generic {
@@ -35,6 +37,7 @@ impl Generic {
             Generic::Process(p) => p.index(),
             Generic::ProcessBinary(pb) => pb.index(),
             Generic::ZeroCopy(z) => z.index(),
+            Generic::ZeroCopyBinary(z) => z.index(),
         }
     }
     /// The number of workers.
@@ -44,6 +47,7 @@ impl Generic {
             Generic::Process(p) => p.peers(),
             Generic::ProcessBinary(pb) => pb.peers(),
             Generic::ZeroCopy(z) => z.peers(),
+            Generic::ZeroCopyBinary(z) => z.peers(),
         }
     }
     /// Constructs several send endpoints and one receive endpoint.
@@ -53,6 +57,7 @@ impl Generic {
             Generic::Process(p) => p.allocate(identifier),
             Generic::ProcessBinary(pb) => pb.allocate(identifier),
             Generic::ZeroCopy(z) => z.allocate(identifier),
+            Generic::ZeroCopyBinary(z) => z.allocate(identifier),
         }
     }
     /// Constructs several send endpoints and one receive endpoint.
@@ -62,6 +67,7 @@ impl Generic {
             Generic::Process(p) => p.broadcast(identifier),
             Generic::ProcessBinary(pb) => pb.broadcast(identifier),
             Generic::ZeroCopy(z) => z.broadcast(identifier),
+            Generic::ZeroCopyBinary(z) => z.broadcast(identifier),
         }
     }
     /// Perform work before scheduling operators.
@@ -71,6 +77,7 @@ impl Generic {
             Generic::Process(p) => p.receive(),
             Generic::ProcessBinary(pb) => pb.receive(),
             Generic::ZeroCopy(z) => z.receive(),
+            Generic::ZeroCopyBinary(z) => z.receive(),
         }
     }
     /// Perform work after scheduling operators.
@@ -80,6 +87,7 @@ impl Generic {
             Generic::Process(p) => p.release(),
             Generic::ProcessBinary(pb) => pb.release(),
             Generic::ZeroCopy(z) => z.release(),
+            Generic::ZeroCopyBinary(z) => z.release(),
         }
     }
     fn events(&self) -> &Rc<RefCell<Vec<usize>>> {
@@ -88,6 +96,7 @@ impl Generic {
             Generic::Process(ref p) => p.events(),
             Generic::ProcessBinary(ref pb) => pb.events(),
             Generic::ZeroCopy(ref z) => z.events(),
+            Generic::ZeroCopyBinary(ref z) => z.events(),
         }
     }
 }
@@ -110,6 +119,7 @@ impl Allocate for Generic {
             Generic::Process(p) => p.await_events(_duration),
             Generic::ProcessBinary(pb) => pb.await_events(_duration),
             Generic::ZeroCopy(z) => z.await_events(_duration),
+            Generic::ZeroCopyBinary(z) => z.await_events(_duration),
         }
     }
 }
@@ -129,6 +139,8 @@ pub enum GenericBuilder {
     ProcessBinary(ProcessBuilder),
     /// Builder for `ZeroCopy` allocator.
     ZeroCopy(TcpBuilder<TypedProcessBuilder>),
+    /// Builder for `ZeroCopyBinary` allocator.
+    ZeroCopyBinary(TcpBuilder<ProcessBuilder>),
 }
 
 impl AllocateBuilder for GenericBuilder {
@@ -139,6 +151,7 @@ impl AllocateBuilder for GenericBuilder {
             GenericBuilder::Process(p) => Generic::Process(p.build()),
             GenericBuilder::ProcessBinary(pb) => Generic::ProcessBinary(pb.build()),
             GenericBuilder::ZeroCopy(z) => Generic::ZeroCopy(z.build()),
+            GenericBuilder::ZeroCopyBinary(z) => Generic::ZeroCopyBinary(z.build()),
         }
     }
 }

communication/src/allocator/mod.rs

@@ -125,3 +125,11 @@ impl<T: Clone> Push<T> for Broadcaster<T> {
         }
     }
 }
+
+/// A builder for vectors of peers.
+pub trait PeerBuilder {
+    /// The peer type.
+    type Peer: AllocateBuilder + Sized;
+    /// Allocate a list of `Self::Peer` of length `peers`.
+    fn new_vector(peers: usize) -> Vec<Self::Peer>;
+}

communication/src/allocator/process.rs

@@ -9,7 +9,7 @@ use std::collections::{HashMap};
 use crossbeam_channel::{Sender, Receiver};
 
 use crate::allocator::thread::{ThreadBuilder};
-use crate::allocator::{Allocate, AllocateBuilder, Thread};
+use crate::allocator::{Allocate, AllocateBuilder, PeerBuilder, Thread};
 use crate::{Push, Pull};
 use crate::buzzer::Buzzer;
 
@@ -70,8 +70,12 @@ pub struct Process {
 impl Process {
     /// Access the wrapped inner allocator.
     pub fn inner(&mut self) -> &mut Thread { &mut self.inner }
+}
+
+impl PeerBuilder for Process {
+    type Peer = ProcessBuilder;
     /// Allocate a list of connected intra-process allocators.
-    pub fn new_vector(peers: usize) -> Vec<ProcessBuilder> {
+    fn new_vector(peers: usize) -> Vec<ProcessBuilder> {
 
         let mut counters_send = Vec::with_capacity(peers);
         let mut counters_recv = Vec::with_capacity(peers);

communication/src/allocator/zero_copy/allocator_process.rs

@@ -10,7 +10,7 @@ use timely_bytes::arc::Bytes;
 use crate::networking::MessageHeader;
 
 use crate::{Allocate, Push, Pull};
-use crate::allocator::{AllocateBuilder, Exchangeable};
+use crate::allocator::{AllocateBuilder, Exchangeable, PeerBuilder};
 use crate::allocator::canary::Canary;
 
 use super::bytes_exchange::{BytesPull, SendEndpoint, MergeQueue};
@@ -30,11 +30,12 @@ pub struct ProcessBuilder {
     pullers: Vec<Sender<MergeQueue>>,   // for pulling bytes from other workers.
 }
 
-impl ProcessBuilder {
+impl PeerBuilder for ProcessBuilder {
+    type Peer = ProcessBuilder;
     /// Creates a vector of builders, sharing appropriate state.
     ///
     /// This method requires access to a byte exchanger, from which it mints channels.
-    pub fn new_vector(count: usize) -> Vec<ProcessBuilder> {
+    fn new_vector(count: usize) -> Vec<ProcessBuilder> {
 
         // Channels for the exchange of `MergeQueue` endpoints.
         let (pullers_vec, pushers_vec) = crate::promise_futures(count, count);
@@ -53,7 +54,9 @@ impl ProcessBuilder {
             )
             .collect()
     }
+}
 
+impl ProcessBuilder {
     /// Builds a `ProcessAllocator`, instantiating `Rc<RefCell<_>>` elements.
     pub fn build(self) -> ProcessAllocator {
 

communication/src/allocator/zero_copy/initialize.rs

@@ -2,7 +2,7 @@
 
 use std::sync::Arc;
 use timely_logging::Logger;
-use crate::allocator::process::ProcessBuilder;
+use crate::allocator::PeerBuilder;
 use crate::logging::CommunicationEventBuilder;
 use crate::networking::create_sockets;
 use super::tcp::{send_loop, recv_loop};
@@ -34,17 +34,17 @@ impl Drop for CommsGuard {
 use crate::logging::CommunicationSetup;
 
 /// Initializes network connections
-pub fn initialize_networking(
+pub fn initialize_networking<P: PeerBuilder>(
     addresses: Vec<String>,
     my_index: usize,
     threads: usize,
     noisy: bool,
     log_sender: Arc<dyn Fn(CommunicationSetup)->Option<Logger<CommunicationEventBuilder>>+Send+Sync>,
 )
--> ::std::io::Result<(Vec<TcpBuilder<ProcessBuilder>>, CommsGuard)>
+-> ::std::io::Result<(Vec<TcpBuilder<P::Peer>>, CommsGuard)>
 {
     let sockets = create_sockets(addresses, my_index, noisy)?;
-    initialize_networking_from_sockets(sockets, my_index, threads, log_sender)
+    initialize_networking_from_sockets::<_, P>(sockets, my_index, threads, log_sender)
 }
 
 /// Initialize send and recv threads from sockets.
@@ -54,13 +54,13 @@ pub fn initialize_networking(
 ///
 /// It is important that the `sockets` argument contain sockets for each remote process, in order, and
 /// with position `my_index` set to `None`.
-pub fn initialize_networking_from_sockets<S: Stream + 'static>(
+pub fn initialize_networking_from_sockets<S: Stream + 'static, P: PeerBuilder>(
     mut sockets: Vec<Option<S>>,
     my_index: usize,
     threads: usize,
     log_sender: Arc<dyn Fn(CommunicationSetup)->Option<Logger<CommunicationEventBuilder>>+Send+Sync>,
 )
--> ::std::io::Result<(Vec<TcpBuilder<ProcessBuilder>>, CommsGuard)>
+-> ::std::io::Result<(Vec<TcpBuilder<P::Peer>>, CommsGuard)>
 {
     // Sockets are expected to be blocking,
     for socket in sockets.iter_mut().flatten() {
@@ -69,7 +69,7 @@ pub fn initialize_networking_from_sockets<S: Stream + 'static>(
 
     let processes = sockets.len();
 
-    let process_allocators = crate::allocator::process::Process::new_vector(threads);
+    let process_allocators = P::new_vector(threads);
     let (builders, promises, futures) = new_vector(process_allocators, my_index, processes);
 
     let mut promises_iter = promises.into_iter();

communication/src/initialize.rs

@@ -12,7 +12,7 @@ use getopts;
 use timely_logging::Logger;
 
 use crate::allocator::thread::ThreadBuilder;
-use crate::allocator::{AllocateBuilder, Process, Generic, GenericBuilder};
+use crate::allocator::{AllocateBuilder, Process, Generic, GenericBuilder, PeerBuilder};
 use crate::allocator::zero_copy::allocator_process::ProcessBuilder;
 use crate::allocator::zero_copy::initialize::initialize_networking;
 use crate::logging::{CommunicationEventBuilder, CommunicationSetup};
@@ -36,6 +36,8 @@ pub enum Config {
         addresses: Vec<String>,
         /// Verbosely report connection process
         report: bool,
+        /// Enable intra-process zero-copy
+        zerocopy: bool,
         /// Closure to create a new logger for a communication thread
         log_fn: Arc<dyn Fn(CommunicationSetup) -> Option<Logger<CommunicationEventBuilder>> + Send + Sync>,
     }
@@ -47,14 +49,14 @@ impl Debug for Config {
             Config::Thread => write!(f, "Config::Thread()"),
             Config::Process(n) => write!(f, "Config::Process({})", n),
             Config::ProcessBinary(n) => write!(f, "Config::ProcessBinary({})", n),
-            Config::Cluster { threads, process, addresses, report, .. } => f
+            Config::Cluster { threads, process, addresses, report, zerocopy, log_fn: _ } => f
                 .debug_struct("Config::Cluster")
                 .field("threads", threads)
                 .field("process", process)
                 .field("addresses", addresses)
                 .field("report", report)
-                // TODO: Use `.finish_non_exhaustive()` after rust/#67364 lands
-                .finish()
+                .field("zerocopy", zerocopy)
+                .finish_non_exhaustive()
         }
     }
 }
@@ -113,12 +115,13 @@ impl Config {
                 }
             }
 
-            assert!(processes == addresses.len());
+            assert_eq!(processes, addresses.len());
             Ok(Config::Cluster {
                 threads,
                 process,
                 addresses,
                 report,
+                zerocopy,
                 log_fn: Arc::new(|_| None),
             })
         } else if threads > 1 {
@@ -158,14 +161,22 @@ impl Config {
             Config::ProcessBinary(threads) => {
                 Ok((ProcessBuilder::new_vector(threads).into_iter().map(GenericBuilder::ProcessBinary).collect(), Box::new(())))
             },
-            Config::Cluster { threads, process, addresses, report, log_fn } => {
-                match initialize_networking(addresses, process, threads, report, log_fn) {
+            Config::Cluster { threads, process, addresses, report, zerocopy: false, log_fn } => {
+                match initialize_networking::<Process>(addresses, process, threads, report, log_fn) {
                     Ok((stuff, guard)) => {
                         Ok((stuff.into_iter().map(GenericBuilder::ZeroCopy).collect(), Box::new(guard)))
                     },
                     Err(err) => Err(format!("failed to initialize networking: {}", err))
                 }
             },
+            Config::Cluster { threads, process, addresses, report, zerocopy: true, log_fn } => {
+                match initialize_networking::<ProcessBuilder>(addresses, process, threads, report, log_fn) {
+                    Ok((stuff, guard)) => {
+                        Ok((stuff.into_iter().map(GenericBuilder::ZeroCopyBinary).collect(), Box::new(guard)))
+                    },
+                    Err(err) => Err(format!("failed to initialize networking: {}", err))
+                }
+            }
         }
     }
 }