timely/dataflow/channels/pushers/
exchange.rs

1//! The exchange pattern distributes pushed data between many target pushees.
2
3use crate::ContainerBuilder;
4use crate::communication::Push;
5use crate::container::{DrainContainer, PushInto};
6use crate::dataflow::channels::Message;
7use crate::progress::Stamp;
8
9/// Distribute containers to several pushers.
10///
11/// A distributor sits behind an exchange pusher, and partitions containers at a given time
12/// into several pushers. It can use [`Message::push_at`] to push its outputs at the desired
13/// pusher.
14///
15/// It needs to uphold progress tracking requirements. The count of the input container
16/// must be preserved across the output containers, from the first call to `partition` until the
17/// call to `flush` for a specific time stamp.
18pub trait Distributor<C> {
19    /// Partition the contents of `container` at `stamp` into the `pushers`.
20    fn partition<T: Clone, P: Push<Message<T, C>>>(&mut self, container: &mut C, stamp: &Stamp<T>, pushers: &mut [P]);
21    /// Flush any remaining contents into the `pushers` at stamp `stamp`.
22    fn flush<T: Clone, P: Push<Message<T, C>>>(&mut self, stamp: &Stamp<T>, pushers: &mut [P]);
23    /// Optionally release resources, such as memory.
24    fn relax(&mut self) { }
25}
26
27/// A distributor creating containers from a drainable container based
28/// on a hash function of the container's item.
29pub struct DrainContainerDistributor<CB, H> {
30    builders: Vec<CB>,
31    hash_func: H,
32}
33
34impl<CB: Default, H> DrainContainerDistributor<CB, H> {
35    /// Constructs a new `DrainContainerDistributor` with the given hash function for a number of
36    /// peers.
37    pub fn new(hash_func: H, peers: usize) -> Self {
38        Self {
39            builders: std::iter::repeat_with(Default::default).take(peers).collect(),
40            hash_func,
41        }
42    }
43}
44
45impl<CB, H> Distributor<CB::Container> for DrainContainerDistributor<CB, H>
46where
47    CB: ContainerBuilder<Container: DrainContainer> + for<'a> PushInto<<CB::Container as DrainContainer>::Item<'a>>,
48    for<'a> H: FnMut(&<CB::Container as DrainContainer>::Item<'a>) -> u64,
49{
50    fn partition<T: Clone, P: Push<Message<T, CB::Container>>>(&mut self, container: &mut CB::Container, stamp: &Stamp<T>, pushers: &mut [P]) {
51        debug_assert_eq!(self.builders.len(), pushers.len());
52        if pushers.len().is_power_of_two() {
53            let mask = (pushers.len() - 1) as u64;
54            for datum in container.drain() {
55                let index = ((self.hash_func)(&datum) & mask) as usize;
56                self.builders[index].push_into(datum);
57                while let Some(produced) = self.builders[index].extract() {
58                    Message::push_at(produced, stamp.clone(), &mut pushers[index]);
59                }
60            }
61        }
62        else {
63            let num_pushers = pushers.len() as u64;
64            for datum in container.drain() {
65                let index = ((self.hash_func)(&datum) % num_pushers) as usize;
66                self.builders[index].push_into(datum);
67                while let Some(produced) = self.builders[index].extract() {
68                    Message::push_at(produced, stamp.clone(), &mut pushers[index]);
69                }
70            }
71        }
72    }
73
74    fn flush<T: Clone, P: Push<Message<T, CB::Container>>>(&mut self, stamp: &Stamp<T>, pushers: &mut [P]) {
75        for (builder, pusher) in self.builders.iter_mut().zip(pushers.iter_mut()) {
76            while let Some(container) = builder.finish() {
77                Message::push_at(container, stamp.clone(), pusher);
78            }
79        }
80    }
81
82    fn relax(&mut self) {
83        for builder in &mut self.builders {
84            builder.relax();
85        }
86    }
87}
88
89// TODO : Software write combining
90/// Distributes records among target pushees according to a distributor.
91pub struct Exchange<T, P, D> {
92    pushers: Vec<P>,
93    current: Option<Stamp<T>>,
94    distributor: D,
95}
96
97impl<T: Clone, P, D>  Exchange<T, P, D> {
98    /// Allocates a new `Exchange` from a supplied set of pushers and a distributor.
99    pub fn new(pushers: Vec<P>, distributor: D) -> Exchange<T, P, D> {
100        Exchange {
101            pushers,
102            current: None,
103            distributor,
104        }
105    }
106}
107
108impl<T: Eq+Clone, C, P, D> Push<Message<T, C>> for Exchange<T, P, D>
109where
110    P: Push<Message<T, C>>,
111    D: Distributor<C>,
112{
113    #[inline(never)]
114    fn push(&mut self, message: &mut Option<Message<T, C>>) {
115        // if only one pusher, no exchange
116        if self.pushers.len() == 1 {
117            self.pushers[0].push(message);
118        }
119        else if let Some(message) = message {
120
121            let stamp = &message.stamp;
122            let data = &mut message.data;
123
124            // if the stamp isn't right, flush everything.
125            match self.current.as_ref() {
126                // We have a current stamp, and it is different from the new stamp.
127                Some(current_stamp) if current_stamp != stamp => {
128                    self.distributor.flush(current_stamp, &mut self.pushers);
129                    self.current = Some(stamp.clone());
130                }
131                // We had no stamp before, or flushed.
132                None => self.current = Some(stamp.clone()),
133                // Stamp didn't change since last call.
134                _ => {}
135            }
136
137            self.distributor.partition(data, stamp, &mut self.pushers);
138        }
139        else {
140            // flush
141            if let Some(stamp) = self.current.take() {
142                self.distributor.flush(&stamp, &mut self.pushers);
143            }
144            self.distributor.relax();
145            for index in 0..self.pushers.len() {
146                self.pushers[index].done();
147            }
148        }
149    }
150}
timely/dataflow/channels/pushers/exchange.rs

timely/dataflow/channels/pushers/
exchange.rs
