Deterministic simulation of multi-worker computations

communication/src/allocator/mod.rs

@@ -151,6 +151,33 @@ impl ProcessBuilder {
         }
     }
 
+    /// Sets whether the built allocator's receive gate holds messages until released.
+    ///
+    /// Only `Bytes` allocators have a gate; the setting is ignored for `Typed` allocators.
+    pub fn holding(self, holding: bool) -> Self {
+        match self {
+            ProcessBuilder::Typed(t) => ProcessBuilder::Typed(t),
+            ProcessBuilder::Bytes(b) => ProcessBuilder::Bytes(b.holding(holding)),
+        }
+    }
+
+    /// Builds a vector of peers on the calling thread.
+    ///
+    /// `Bytes` builders block in `build` until their peers are built, and so must be
+    /// built together when they share a thread; `Typed` builders build independently.
+    pub fn build_all(builders: Vec<Self>) -> Vec<Process> {
+        if builders.iter().all(|builder| matches!(builder, ProcessBuilder::Bytes(_))) {
+            let bytes = builders.into_iter().map(|builder| match builder {
+                ProcessBuilder::Bytes(b) => b,
+                ProcessBuilder::Typed(_) => unreachable!("checked above"),
+            }).collect();
+            BytesProcessBuilder::build_all(bytes).into_iter().map(Process::Bytes).collect()
+        }
+        else {
+            builders.into_iter().map(|builder| builder.build()).collect()
+        }
+    }
+
     /// Constructs a vector of regular (mpsc-based, "Typed") intra-process builders.
     pub fn new_typed_vector(peers: usize, refill: BytesRefill, spill: Option<SpillPolicyFn>) -> Vec<Self> {
         <TypedProcess as PeerBuilder>::new_vector(peers, refill, spill)

communication/src/allocator/zero_copy/allocator.rs

@@ -13,7 +13,8 @@ use crate::allocator::{Process, ProcessBuilder, Exchangeable};
 use crate::allocator::canary::Canary;
 use crate::allocator::zero_copy::bytes_slab::BytesRefill;
 use crate::allocator::zero_copy::spill::SpillPolicyFn;
-use super::bytes_exchange::{BytesPull, SendEndpoint, MergeQueue};
+use super::bytes_exchange::{SendEndpoint, MergeQueue};
+use super::gate::Gate;
 use super::push_pull::{Pusher, PullerInner};
 
 /// Builds an instance of a TcpAllocator.
@@ -124,7 +125,7 @@ impl TcpBuilder {
             channel_id_bound: None,
             staged: Vec::new(),
             sends,
-            recvs,
+            gate: Rc::new(RefCell::new(Gate::new(self.index, self.peers, recvs, false))),
             to_local: HashMap::new(),
         }
     }
@@ -145,10 +146,17 @@ pub struct TcpAllocator {
 
     // sending, receiving, and responding to binary buffers.
     sends:      Vec<Rc<RefCell<SendEndpoint<MergeQueue>>>>,     // sends[x] -> goes to process x.
-    recvs:      Vec<MergeQueue>,                                // recvs[x] <- from process x.
+    gate:       Rc<RefCell<Gate>>,                              // receives from all processes.
     to_local:   HashMap<usize, Rc<RefCell<VecDeque<Bytes>>>>,   // to worker-local typed pullers.
 }
 
+impl TcpAllocator {
+    /// The gate through which this allocator receives messages from other processes.
+    ///
+    /// Messages from workers in the same process pass through the inner allocator instead.
+    pub fn gate(&self) -> Rc<RefCell<Gate>> { Rc::clone(&self.gate) }
+}
+
 impl Allocate for TcpAllocator {
     fn index(&self) -> usize { self.index }
     fn peers(&self) -> usize { self.peers }
@@ -264,9 +272,7 @@ impl Allocate for TcpAllocator {
 
         self.inner.receive();
 
-        for recv in self.recvs.iter_mut() {
-            recv.drain_into(&mut self.staged);
-        }
+        self.gate.borrow_mut().receive(&mut self.staged);
 
         let mut events = self.inner.events().borrow_mut();
 

communication/src/allocator/zero_copy/allocator_process.rs

@@ -14,7 +14,8 @@ use crate::allocator::{AllocateBuilder, Exchangeable, PeerBuilder};
 use crate::allocator::canary::Canary;
 use crate::allocator::zero_copy::bytes_slab::BytesRefill;
 use crate::allocator::zero_copy::spill::SpillPolicyFn;
-use super::bytes_exchange::{BytesPull, SendEndpoint, MergeQueue};
+use super::bytes_exchange::{SendEndpoint, MergeQueue};
+use super::gate::Gate;
 
 use super::push_pull::{Pusher, Puller};
 
@@ -31,6 +32,7 @@ pub struct ProcessBuilder {
     pullers: Vec<Sender<MergeQueue>>,   // for pulling bytes from other workers.
     refill: BytesRefill,
     spill: Option<SpillPolicyFn>,       // optional spill factory for recv queues.
+    holding: bool,                      // whether the receive gate holds messages.
 }
 
 impl PeerBuilder for ProcessBuilder {
@@ -55,19 +57,44 @@ impl PeerBuilder for ProcessBuilder {
                     pullers,
                     refill: refill.clone(),
                     spill: spill.clone(),
+                    holding: false,
                 }
             )
             .collect()
     }
 }
 
 impl ProcessBuilder {
+    /// Sets whether the built allocator's receive gate holds messages until released.
+    pub fn holding(mut self, holding: bool) -> Self {
+        self.holding = holding;
+        self
+    }
+
     /// Builds a `ProcessAllocator`, instantiating `Rc<RefCell<_>>` elements.
-    pub fn build(self) -> ProcessAllocator {
+    ///
+    /// Each builder must be built on its own thread, or all together with
+    /// [`build_all`](Self::build_all): `build` blocks until every peer has
+    /// fulfilled its obligations, which a peer does only when it is itself built.
+    pub fn build(mut self) -> ProcessAllocator {
+        let recvs = self.fulfill();
+        self.complete(recvs)
+    }
+
+    /// Builds a vector of peers on the calling thread.
+    ///
+    /// All obligations are fulfilled before any allocator is completed, so no
+    /// builder blocks waiting on a peer that has not yet been built.
+    pub fn build_all(mut builders: Vec<Self>) -> Vec<ProcessAllocator> {
+        let recvs: Vec<_> = builders.iter_mut().map(|builder| builder.fulfill()).collect();
+        builders.into_iter().zip(recvs).map(|(builder, recvs)| builder.complete(recvs)).collect()
+    }
 
-        // Fulfill puller obligations.
+    /// Fulfills puller obligations: creates this allocator's receive queues on the
+    /// calling thread, sending the writer halves to peers and returning the readers.
+    fn fulfill(&mut self) -> Vec<MergeQueue> {
         let mut recvs = Vec::with_capacity(self.peers);
-        for puller in self.pullers.into_iter() {
+        for puller in std::mem::take(&mut self.pullers) {
             let buzzer = crate::buzzer::Buzzer::default();
             let (writer, reader) = match self.spill.as_ref() {
                 Some(build_fn) => {
@@ -81,8 +108,11 @@ impl ProcessBuilder {
             recvs.push(reader);
             puller.send(writer).expect("Failed to send MergeQueue");
         }
+        recvs
+    }
 
-        // Extract pusher commitments.
+    /// Extracts pusher commitments, blocking until each peer has fulfilled them.
+    fn complete(self, recvs: Vec<MergeQueue>) -> ProcessAllocator {
         let mut sends = Vec::with_capacity(self.peers);
         for pusher in self.pushers.into_iter() {
             let queue = pusher.recv().expect("Failed to receive MergeQueue");
@@ -98,7 +128,7 @@ impl ProcessBuilder {
             channel_id_bound: None,
             staged: Vec::new(),
             sends,
-            recvs,
+            gate: Rc::new(RefCell::new(Gate::new(self.index, self.peers, recvs, self.holding))),
             to_local: HashMap::new(),
         }
     }
@@ -128,10 +158,17 @@ pub struct ProcessAllocator {
     // sending, receiving, and responding to binary buffers.
     staged:     Vec<Bytes>,
     sends:      Vec<Rc<RefCell<SendEndpoint<MergeQueue>>>>, // sends[x] -> goes to thread x.
-    recvs:      Vec<MergeQueue>,                            // recvs[x] <- from thread x.
+    gate:       Rc<RefCell<Gate>>,                          // receives from all threads.
     to_local:   HashMap<usize, Rc<RefCell<VecDeque<Bytes>>>>,          // to worker-local typed pullers.
 }
 
+impl ProcessAllocator {
+    /// The gate through which this allocator receives messages from its peers.
+    ///
+    /// A gate built holding lets a driver control message delivery, for deterministic simulation.
+    pub fn gate(&self) -> Rc<RefCell<Gate>> { Rc::clone(&self.gate) }
+}
+
 impl Allocate for ProcessAllocator {
     fn index(&self) -> usize { self.index }
     fn peers(&self) -> usize { self.peers }
@@ -191,9 +228,7 @@ impl Allocate for ProcessAllocator {
 
         let mut events = self.events.borrow_mut();
 
-        for recv in self.recvs.iter_mut() {
-            recv.drain_into(&mut self.staged);
-        }
+        self.gate.borrow_mut().receive(&mut self.staged);
 
         for mut bytes in self.staged.drain(..) {
 

communication/src/allocator/zero_copy/gate.rs

@@ -0,0 +1,108 @@
+//! Receive-side admission of framed messages.
+//!
+//! A [`Gate`] sits between the byte queues an allocator physically receives from and the
+//! typed channels it surfaces messages into. By default it is open: bytes pass through
+//! unchanged, and the gate costs one additional `Vec` move per `receive()`.
+//!
+//! A gate constructed holding splits received bytes into framed messages and files them by
+//! source worker. A held message is physically present but logically undelivered: it is
+//! not surfaced to its channel and raises no event, until the holder of the gate releases
+//! it. Per-source FIFO order is preserved; the holder chooses only the interleaving
+//! across sources, which is exactly the schedule space a real transport can produce.
+//! This is the mechanism behind deterministic simulation of multi-worker computations.
+//!
+//! Messages a worker sends to itself are never held. Real allocators surface them at the
+//! next `receive()`, and holding them would admit schedules no deployment can produce.
+
+use std::collections::VecDeque;
+
+use timely_bytes::arc::Bytes;
+
+use crate::networking::MessageHeader;
+use super::bytes_exchange::{BytesPull, MergeQueue};
+
+/// Admission control over bytes received from peers.
+pub struct Gate {
+    /// The owning worker's index, whose own messages are never held.
+    index: usize,
+    /// Physical sources of bytes. Each may carry messages from several source workers.
+    recvs: Vec<MergeQueue>,
+    /// Whether received messages are held until released.
+    holding: bool,
+    /// Held messages, one framed message per entry, indexed by source worker.
+    held: Vec<VecDeque<Bytes>>,
+    /// Bytes admitted for delivery, each containing whole framed messages.
+    admitted: Vec<Bytes>,
+}
+
+impl Gate {
+    /// Creates a gate over `recvs`, for worker `index` of `peers`, holding if `holding`.
+    pub fn new(index: usize, peers: usize, recvs: Vec<MergeQueue>, holding: bool) -> Self {
+        Gate {
+            index,
+            recvs,
+            holding,
+            held: (0 .. peers).map(|_| VecDeque::new()).collect(),
+            admitted: Vec::new(),
+        }
+    }
+
+    /// Drains the physical sources, admitting or holding what they contain.
+    ///
+    /// The allocator calls this from `receive()`. A driver may also call it, to see
+    /// what has arrived without stepping the worker.
+    pub fn fetch(&mut self) {
+        if !self.holding {
+            for recv in self.recvs.iter_mut() {
+                recv.drain_into(&mut self.admitted);
+            }
+        }
+        else {
+            let mut staged = Vec::new();
+            for recv in self.recvs.iter_mut() {
+                recv.drain_into(&mut staged);
+            }
+            for mut bytes in staged {
+                // Received bytes contain whole framed messages; no splitting across allocations.
+                while !bytes.is_empty() {
+                    let header = MessageHeader::try_read(&bytes[..]).expect("failed to read full header!");
+                    let message = bytes.extract_to(header.required_bytes());
+                    if header.source == self.index {
+                        self.admitted.push(message);
+                    }
+                    else {
+                        self.held[header.source].push_back(message);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Fetches, then moves all admitted bytes into `into`.
+    pub fn receive(&mut self, into: &mut Vec<Bytes>) {
+        self.fetch();
+        into.append(&mut self.admitted);
+    }
+
+    /// The number of held messages from `source`, as of the last fetch.
+    pub fn held(&self, source: usize) -> usize {
+        self.held[source].len()
+    }
+
+    /// Releases up to `count` held messages from `source`, in FIFO order.
+    ///
+    /// Returns the number released, which is less than `count` if fewer are held.
+    /// This clamping keeps any sequence of releases valid, which matters when
+    /// shrinking failing schedules.
+    pub fn release(&mut self, source: usize, count: usize) -> usize {
+        let held = &mut self.held[source];
+        let count = std::cmp::min(count, held.len());
+        self.admitted.extend(held.drain(.. count));
+        count
+    }
+
+    /// Releases all held messages, and returns the number released.
+    pub fn release_all(&mut self) -> usize {
+        (0 .. self.held.len()).map(|source| self.release(source, usize::MAX)).sum()
+    }
+}

communication/src/allocator/zero_copy/mod.rs

@@ -10,6 +10,7 @@
 
 pub mod bytes_slab;
 pub mod bytes_exchange;
+pub mod gate;
 pub mod spill;
 pub mod tcp;
 pub mod allocator;

timely/src/lib.rs

@@ -92,6 +92,7 @@ pub mod dataflow;
 pub mod synchronization;
 pub mod execute;
 pub mod order;
+pub mod simulate;
 
 pub mod logging;
 // pub mod log_events;