timely/dataflow/channels/pushers/
exchange.rs

1//! The exchange pattern distributes pushed data between many target pushees.
2
3use crate::communication::Push;
4use crate::container::{ContainerBuilder, PushInto};
5use crate::dataflow::channels::Message;
6use crate::{Container, Data};
7
8// TODO : Software write combining
9/// Distributes records among target pushees according to a distribution function.
10pub struct Exchange<T, CB, P, H>
11where
12    CB: ContainerBuilder,
13    P: Push<Message<T, CB::Container>>,
14    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
15{
16    pushers: Vec<P>,
17    buffers: Vec<CB>,
18    current: Option<T>,
19    hash_func: H,
20}
21
22impl<T: Clone, CB, P, H>  Exchange<T, CB, P, H>
23where
24    CB: ContainerBuilder,
25    P: Push<Message<T, CB::Container>>,
26    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
27{
28    /// Allocates a new `Exchange` from a supplied set of pushers and a distribution function.
29    pub fn new(pushers: Vec<P>, key: H) -> Exchange<T, CB, P, H> {
30        let mut buffers = vec![];
31        for _ in 0..pushers.len() {
32            buffers.push(Default::default());
33        }
34        Exchange {
35            pushers,
36            hash_func: key,
37            buffers,
38            current: None,
39        }
40    }
41    #[inline]
42    fn flush(&mut self, index: usize) {
43        while let Some(container) = self.buffers[index].finish() {
44            if let Some(ref time) = self.current {
45                Message::push_at(container, time.clone(), &mut self.pushers[index]);
46            }
47        }
48    }
49}
50
51impl<T: Eq+Data, CB, P, H> Push<Message<T, CB::Container>> for Exchange<T, CB, P, H>
52where
53    CB: ContainerBuilder,
54    CB: for<'a> PushInto<<CB::Container as Container>::Item<'a>>,
55    P: Push<Message<T, CB::Container>>,
56    for<'a> H: FnMut(&<CB::Container as Container>::Item<'a>) -> u64
57{
58    #[inline(never)]
59    fn push(&mut self, message: &mut Option<Message<T, CB::Container>>) {
60        // if only one pusher, no exchange
61        if self.pushers.len() == 1 {
62            self.pushers[0].push(message);
63        }
64        else if let Some(message) = message {
65
66            let time = &message.time;
67            let data = &mut message.data;
68
69            // if the time isn't right, flush everything.
70            if self.current.as_ref().is_some_and(|x| x != time) {
71                for index in 0..self.pushers.len() {
72                    self.flush(index);
73                }
74            }
75            self.current = Some(time.clone());
76
77            let hash_func = &mut self.hash_func;
78
79            // if the number of pushers is a power of two, use a mask
80            if self.pushers.len().is_power_of_two() {
81                let mask = (self.pushers.len() - 1) as u64;
82                CB::partition(data, &mut self.buffers, |datum| ((hash_func)(datum) & mask) as usize);
83            }
84            // as a last resort, use mod (%)
85            else {
86                let num_pushers = self.pushers.len() as u64;
87                CB::partition(data, &mut self.buffers, |datum| ((hash_func)(datum) % num_pushers) as usize);
88            }
89            for (buffer, pusher) in self.buffers.iter_mut().zip(self.pushers.iter_mut()) {
90                while let Some(container) = buffer.extract() {
91                    Message::push_at(container, time.clone(), pusher);
92                }
93            }
94        }
95        else {
96            // flush
97            for index in 0..self.pushers.len() {
98                self.flush(index);
99                self.pushers[index].push(&mut None);
100            }
101        }
102    }
103}