use std::rc::Rc;
use timely::communication::{Data, Pull, Push};
use timely::dataflow::channels::pact::{LogPuller, LogPusher, ParallelizationContract};
use timely::dataflow::channels::{Bundle, Message};
use timely::logging::TimelyLogger;
use timely::progress::Timestamp;
use timely::worker::AsWorker;
use timely::Container;
#[derive(Debug)]
pub struct Distribute;
impl<T: Timestamp, C: Container + Data> ParallelizationContract<T, C> for Distribute {
type Pusher = DistributePusher<LogPusher<T, C, Box<dyn Push<Bundle<T, C>>>>>;
type Puller = LogPuller<T, C, Box<dyn Pull<Bundle<T, C>>>>;
fn connect<A: AsWorker>(
self,
allocator: &mut A,
identifier: usize,
address: Rc<[usize]>,
logging: Option<TimelyLogger>,
) -> (Self::Pusher, Self::Puller) {
let (senders, receiver) = allocator.allocate::<Message<T, C>>(identifier, address);
let senders = senders
.into_iter()
.enumerate()
.map(|(i, x)| LogPusher::new(x, allocator.index(), i, identifier, logging.clone()))
.collect::<Vec<_>>();
(
DistributePusher::new(senders),
LogPuller::new(receiver, allocator.index(), identifier, logging.clone()),
)
}
}
pub struct DistributePusher<P> {
pushers: Vec<P>,
next: usize,
}
impl<P> DistributePusher<P> {
pub fn new(pushers: Vec<P>) -> DistributePusher<P> {
Self { pushers, next: 0 }
}
}
impl<T: Eq + Data, C: Container, P: Push<Bundle<T, C>>> Push<Bundle<T, C>> for DistributePusher<P> {
fn push(&mut self, message: &mut Option<Bundle<T, C>>) {
let worker_idx = self.next;
self.next = (self.next + 1) % self.pushers.len();
self.pushers[worker_idx].push(message);
}
}