timely/dataflow/channels/pushers/
buffer.rs

1//! Buffering and session mechanisms to provide the appearance of record-at-a-time sending,
2//! with the performance of batched sends.
3
4use crate::communication::Push;
5use crate::container::{ContainerBuilder, CapacityContainerBuilder, PushInto};
6use crate::dataflow::channels::Message;
7use crate::dataflow::operators::Capability;
8use crate::progress::Timestamp;
9use crate::{Container, Data};
10
11/// Buffers data sent at the same time, for efficient communication.
12///
13/// The `Buffer` type should be used by calling `session` with a time, which checks whether
14/// data must be flushed and creates a `Session` object which allows sending at the given time.
15#[derive(Debug)]
16pub struct Buffer<T, CB, P> {
17    /// The currently open time, if it is open.
18    time: Option<T>,
19    /// A builder for containers, to send at `self.time`.
20    builder: CB,
21    /// The pusher to send data downstream.
22    pusher: P,
23}
24
25impl<T, CB: Default, P> Buffer<T, CB, P> {
26    /// Creates a new `Buffer`.
27    pub fn new(pusher: P) -> Self {
28        Self {
29            time: None,
30            builder: Default::default(),
31            pusher,
32        }
33    }
34
35    /// Returns a reference to the inner `P: Push` type.
36    ///
37    /// This is currently used internally, and should not be used without some care.
38    pub fn inner(&mut self) -> &mut P { &mut self.pusher }
39
40    /// Access the builder. Immutable access to prevent races with flushing
41    /// the underlying buffer.
42    pub fn builder(&self) -> &CB {
43        &self.builder
44    }
45}
46
47impl<T, C: Container + Data, P: Push<Message<T, C>>> Buffer<T, CapacityContainerBuilder<C>, P> where T: Eq+Clone {
48    /// Returns a `Session`, which accepts data to send at the associated time
49    #[inline]
50    pub fn session(&mut self, time: &T) -> Session<T, CapacityContainerBuilder<C>, P> {
51        self.session_with_builder(time)
52    }
53
54    /// Allocates a new `AutoflushSession` which flushes itself on drop.
55    #[inline]
56    pub fn autoflush_session(&mut self, cap: Capability<T>) -> AutoflushSession<T, CapacityContainerBuilder<C>, P> where T: Timestamp {
57        self.autoflush_session_with_builder(cap)
58    }
59}
60
61impl<T, CB: ContainerBuilder, P: Push<Message<T, CB::Container>>> Buffer<T, CB, P> where T: Eq+Clone {
62    /// Returns a `Session`, which accepts data to send at the associated time
63    pub fn session_with_builder(&mut self, time: &T) -> Session<T, CB, P> {
64        if let Some(true) = self.time.as_ref().map(|x| x != time) { self.flush(); }
65        self.time = Some(time.clone());
66        Session { buffer: self }
67    }
68
69    /// Allocates a new `AutoflushSession` which flushes itself on drop.
70    pub fn autoflush_session_with_builder(&mut self, cap: Capability<T>) -> AutoflushSession<T, CB, P> where T: Timestamp {
71        if let Some(true) = self.time.as_ref().map(|x| x != cap.time()) { self.flush(); }
72        self.time = Some(cap.time().clone());
73        AutoflushSession {
74            buffer: self,
75            _capability: cap,
76        }
77    }
78}
79
80impl<T, CB: ContainerBuilder, P: Push<Message<T, CB::Container>>> Buffer<T, CB, P> where T: Eq+Clone {
81    /// Flushes all data and pushes a `None` to `self.pusher`, indicating a flush.
82    pub fn cease(&mut self) {
83        self.flush();
84        self.pusher.push(&mut None);
85    }
86
87    /// Extract pending data from the builder, but not forcing a flush.
88    #[inline]
89    fn extract_and_send(&mut self) {
90        while let Some(container) = self.builder.extract() {
91            let time = self.time.as_ref().unwrap().clone();
92            Message::push_at(container, time, &mut self.pusher);
93        }
94    }
95
96    /// Flush the builder, forcing all its contents to be written.
97    #[inline]
98    fn flush(&mut self) {
99        while let Some(container) = self.builder.finish() {
100            let time = self.time.as_ref().unwrap().clone();
101            Message::push_at(container, time, &mut self.pusher);
102        }
103    }
104
105    /// Gives an entire container at the current time. Maintains FIFO order with previously pushed
106    /// data. Only intended to be used through [`Session::give_container`].
107    // TODO: This method could exist without a container builder, but we can't express this as a
108    // buffer always requires a container builder. We could expose the buffer's underlying pusher
109    // directly, but this would bypass the buffer's time tracking.
110    fn give_container(&mut self, container: &mut CB::Container) {
111        if !container.is_empty() {
112            self.flush();
113            let time = self.time.as_ref().unwrap().clone();
114            Message::push_at(container, time, &mut self.pusher);
115        }
116    }
117    
118    /// An internal implementation of push that should only be called by sessions.
119    #[inline]
120    fn push_internal<D>(&mut self, item: D) where CB: PushInto<D> {
121        self.builder.push_into(item);
122        self.extract_and_send();
123    }
124}
125
126/// An output session for sending records at a specified time.
127///
128/// The `Session` struct provides the user-facing interface to an operator output, namely
129/// the `Buffer` type. A `Session` wraps a session of output at a specified time, and
130/// avoids what would otherwise be a constant cost of checking timestamp equality.
131pub struct Session<'a, T, CB, P> {
132    buffer: &'a mut Buffer<T, CB, P>,
133}
134
135impl<'a, T, CB: ContainerBuilder, P> Session<'a, T, CB, P>
136where
137    T: Eq + Clone + 'a,
138    P: Push<Message<T, CB::Container>> + 'a,
139{
140    /// Provide a container at the time specified by the [Session]. Maintains FIFO order with
141    /// previously pushed data.
142    pub fn give_container(&mut self, container: &mut CB::Container) {
143        self.buffer.give_container(container)
144    }
145}
146
147impl<'a, T, CB, P> Session<'a, T, CB, P>
148where
149    T: Eq + Clone + 'a,
150    CB: ContainerBuilder + 'a,
151    P: Push<Message<T, CB::Container>> + 'a
152{
153    /// Access the builder. Immutable access to prevent races with flushing
154    /// the underlying buffer.
155    pub fn builder(&self) -> &CB {
156        self.buffer.builder()
157    }
158
159    /// Provides one record at the time specified by the `Session`.
160    #[inline]
161    pub fn give<D>(&mut self, data: D) where CB: PushInto<D> {
162        self.push_into(data);
163    }
164
165    /// Provides an iterator of records at the time specified by the `Session`.
166    #[inline]
167    pub fn give_iterator<I>(&mut self, iter: I)
168    where
169        I: Iterator,
170        CB: PushInto<I::Item>,
171    {
172        for item in iter {
173            self.push_into(item);
174        }
175    }
176}
177
178impl<'a, T, CB, P, D> PushInto<D> for Session<'a, T, CB, P>
179where
180    T: Eq + Clone + 'a,
181    CB: ContainerBuilder + PushInto<D> + 'a,
182    P: Push<Message<T, CB::Container>> + 'a,
183{
184    #[inline]
185    fn push_into(&mut self, item: D) {
186        self.buffer.push_internal(item);
187    }
188}
189
190/// A session which will flush itself when dropped.
191pub struct AutoflushSession<'a, T, CB, P>
192where
193    T: Timestamp + 'a,
194    CB: ContainerBuilder + 'a,
195    P: Push<Message<T, CB::Container>> + 'a,
196{
197    /// A reference to the underlying buffer.
198    buffer: &'a mut Buffer<T, CB, P>,
199    /// The capability being used to send the data.
200    _capability: Capability<T>,
201}
202
203impl<'a, T, CB, P> AutoflushSession<'a, T, CB, P>
204where
205    T: Timestamp + 'a,
206    CB: ContainerBuilder + 'a,
207    P: Push<Message<T, CB::Container>> + 'a,
208{
209    /// Transmits a single record.
210    #[inline]
211    pub fn give<D>(&mut self, data: D)
212    where
213        CB: PushInto<D>,
214    {
215        self.push_into(data);
216    }
217
218    /// Transmits records produced by an iterator.
219    #[inline]
220    pub fn give_iterator<I, D>(&mut self, iter: I)
221    where
222        I: Iterator<Item=D>,
223        CB: PushInto<D>,
224    {
225        for item in iter {
226            self.push_into(item);
227        }
228    }
229}
230impl<'a, T, CB, P, D> PushInto<D> for AutoflushSession<'a, T, CB, P>
231where
232    T: Timestamp + 'a,
233    CB: ContainerBuilder + PushInto<D> + 'a,
234    P: Push<Message<T, CB::Container>> + 'a,
235{
236    #[inline]
237    fn push_into(&mut self, item: D) {
238        self.buffer.push_internal(item);
239    }
240}
241
242impl<'a, T, CB, P> Drop for AutoflushSession<'a, T, CB, P>
243where
244    T: Timestamp + 'a,
245    CB: ContainerBuilder + 'a,
246    P: Push<Message<T, CB::Container>> + 'a,
247{
248    fn drop(&mut self) {
249        self.buffer.cease();
250    }
251}
timely/dataflow/channels/pushers/buffer.rs

timely/dataflow/channels/pushers/
buffer.rs
