timely/dataflow/channels/pushers/
exchange.rs
1use crate::communication::Push;
4use crate::container::{ContainerBuilder, PushInto};
5use crate::dataflow::channels::Message;
6use crate::{Container, Data};
7
8pub 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 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 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 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 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 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 for index in 0..self.pushers.len() {
98 self.flush(index);
99 self.pushers[index].push(&mut None);
100 }
101 }
102 }
103}