mz_timely_util/
builder_async.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Types to build async operators with general shapes.
17
18use std::cell::{Cell, RefCell};
19use std::collections::VecDeque;
20use std::future::Future;
21use std::pin::Pin;
22use std::rc::Rc;
23use std::sync::Arc;
24use std::sync::atomic::{AtomicBool, Ordering};
25use std::task::{Context, Poll, Waker, ready};
26
27use futures_util::Stream;
28use futures_util::task::ArcWake;
29use timely::communication::{Pull, Push};
30use timely::container::{CapacityContainerBuilder, PushInto};
31use timely::dataflow::channels::Message;
32use timely::dataflow::channels::pact::ParallelizationContract;
33use timely::dataflow::channels::pushers::Output;
34use timely::dataflow::operators::generic::builder_rc::OperatorBuilder as OperatorBuilderRc;
35use timely::dataflow::operators::generic::{InputHandleCore, OperatorInfo};
36use timely::dataflow::operators::{Capability, CapabilitySet, CapabilityTrait, InputCapability};
37use timely::dataflow::{Scope, Stream as TimelyStream, StreamVec};
38use timely::order::TotalOrder;
39use timely::progress::{Antichain, Stamp, Timestamp};
40use timely::scheduling::{Activator, SyncActivator};
41use timely::{Bincode, Container, ContainerBuilder, PartialOrder};
42
43use crate::containers::stack::FueledBuilder;
44
45/// Builds async operators with generic shape.
46pub struct OperatorBuilder<'scope, T: Timestamp> {
47    builder: OperatorBuilderRc<'scope, T>,
48    /// The activator for this operator
49    activator: Activator,
50    /// The waker set up to activate this timely operator when woken
51    operator_waker: Arc<TimelyWaker>,
52    /// The currently known upper frontier of each of the input handles.
53    input_frontiers: Vec<Antichain<T>>,
54    /// Input queues for each of the declared inputs of the operator.
55    input_queues: Vec<Box<dyn InputQueue<T>>>,
56    /// Holds type erased closures that flush an output handle when called. These handles will be
57    /// automatically drained when the operator is scheduled after the logic future has been polled
58    output_flushes: Vec<Box<dyn FnMut()>>,
59    /// A handle to check whether all workers have pressed the shutdown button.
60    shutdown_handle: ButtonHandle,
61    /// A button to coordinate shutdown of this operator among workers.
62    shutdown_button: Button,
63}
64
65/// A helper trait abstracting over an input handle. It facilitates keeping around type erased
66/// handles for each of the operator inputs.
67trait InputQueue<T: Timestamp> {
68    /// Accepts all available input into local queues.
69    fn accept_input(&mut self);
70
71    /// Drains all available input and empties the local queue.
72    fn drain_input(&mut self);
73
74    /// Registers a frontier notification to be delivered.
75    fn notify_progress(&mut self, upper: Antichain<T>);
76}
77
78impl<T, D, C, P> InputQueue<T> for InputHandleQueue<T, D, C, P>
79where
80    T: Timestamp,
81    D: Container,
82    C: InputConnection<T> + 'static,
83    P: Pull<Message<T, D>> + 'static,
84{
85    fn accept_input(&mut self) {
86        let mut queue = self.queue.borrow_mut();
87        let mut new_data = false;
88        self.handle.for_each(|cap, data| {
89            new_data = true;
90            let cap = self.connection.accept(cap);
91            queue.push_back(Event::Data(cap, std::mem::take(data)));
92        });
93        if new_data {
94            if let Some(waker) = self.waker.take() {
95                waker.wake();
96            }
97        }
98    }
99
100    fn drain_input(&mut self) {
101        self.queue.borrow_mut().clear();
102        self.handle.for_each(|_, _| {});
103    }
104
105    fn notify_progress(&mut self, upper: Antichain<T>) {
106        let mut queue = self.queue.borrow_mut();
107        // It's beneficial to consolidate two consecutive progress statements into one if the
108        // operator hasn't seen the previous progress yet. This also avoids accumulation of
109        // progress statements in the queue if the operator only conditionally checks this input.
110        match queue.back_mut() {
111            Some(&mut Event::Progress(ref mut prev_upper)) => *prev_upper = upper,
112            _ => queue.push_back(Event::Progress(upper)),
113        }
114        if let Some(waker) = self.waker.take() {
115            waker.wake();
116        }
117    }
118}
119
120struct InputHandleQueue<
121    T: Timestamp,
122    D: Container,
123    C: InputConnection<T>,
124    P: Pull<Message<T, D>> + 'static,
125> {
126    queue: Rc<RefCell<VecDeque<Event<T, C::Capability, D>>>>,
127    waker: Rc<Cell<Option<Waker>>>,
128    connection: C,
129    handle: InputHandleCore<T, D, P>,
130}
131
132/// An async Waker that activates a specific operator when woken and marks the task as ready
133struct TimelyWaker {
134    activator: SyncActivator,
135    active: AtomicBool,
136    task_ready: AtomicBool,
137}
138
139impl ArcWake for TimelyWaker {
140    fn wake_by_ref(arc_self: &Arc<Self>) {
141        arc_self.task_ready.store(true, Ordering::SeqCst);
142        // Only activate the timely operator if it's not already active to avoid an infinite loop
143        if !arc_self.active.load(Ordering::SeqCst) {
144            // We don't have any guarantees about how long the Waker will be held for and so we
145            // must be prepared for the receiving end to have hung up when we finally do get woken.
146            // This can happen if by the time the waker is called the receiving timely worker has
147            // been shutdown. For this reason we ignore the activation error.
148            let _ = arc_self.activator.activate();
149        }
150    }
151}
152
153/// Async handle to an operator's input stream
154pub struct AsyncInputHandle<T: Timestamp, D: Container, C: InputConnection<T>> {
155    queue: Rc<RefCell<VecDeque<Event<T, C::Capability, D>>>>,
156    waker: Rc<Cell<Option<Waker>>>,
157    /// Whether this handle has finished producing data
158    done: bool,
159}
160
161impl<T: Timestamp, D: Container, C: InputConnection<T>> AsyncInputHandle<T, D, C> {
162    pub fn next_sync(&mut self) -> Option<Event<T, C::Capability, D>> {
163        let mut queue = self.queue.borrow_mut();
164        match queue.pop_front()? {
165            Event::Data(cap, data) => Some(Event::Data(cap, data)),
166            Event::Progress(frontier) => {
167                self.done = frontier.is_empty();
168                Some(Event::Progress(frontier))
169            }
170        }
171    }
172
173    /// Waits for the handle to have data. After this function returns it is guaranteed that at
174    /// least one call to `next_sync` will be `Some(_)`.
175    pub async fn ready(&self) {
176        std::future::poll_fn(|cx| self.poll_ready(cx)).await
177    }
178
179    fn poll_ready(&self, cx: &Context<'_>) -> Poll<()> {
180        if self.queue.borrow().is_empty() {
181            self.waker.set(Some(cx.waker().clone()));
182            Poll::Pending
183        } else {
184            Poll::Ready(())
185        }
186    }
187}
188
189impl<T: Timestamp, D: Container, C: InputConnection<T>> Stream for AsyncInputHandle<T, D, C> {
190    type Item = Event<T, C::Capability, D>;
191
192    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
193        if self.done {
194            return Poll::Ready(None);
195        }
196        ready!(self.poll_ready(cx));
197        Poll::Ready(self.next_sync())
198    }
199
200    fn size_hint(&self) -> (usize, Option<usize>) {
201        (self.queue.borrow().len(), None)
202    }
203}
204
205/// An event of an input stream
206#[derive(Debug)]
207pub enum Event<T: Timestamp, C, D> {
208    /// A data event
209    Data(C, D),
210    /// A progress event
211    Progress(Antichain<T>),
212}
213
214/// Shared part of an async output handle
215struct AsyncOutputHandleInner<T: Timestamp, CB: ContainerBuilder> {
216    /// Handle to write to the output stream.
217    output: Output<T, CB::Container>,
218    /// Current capability held by this output handle.
219    capability: Option<CapabilitySet<T>>,
220    /// Container builder to accumulate data before sending at `capability`.
221    builder: CB,
222}
223
224impl<T: Timestamp, CB: ContainerBuilder> AsyncOutputHandleInner<T, CB> {
225    /// Write all pending data to the output stream.
226    fn flush(&mut self) {
227        while let Some(container) = self.builder.finish() {
228            self.output
229                .give(self.capability.as_ref().expect("must exist"), container);
230        }
231    }
232
233    /// Cease this output handle, flushing all pending data and releasing its capability.
234    fn cease(&mut self) {
235        self.flush();
236        let _ = self.output.activate();
237        self.capability = None;
238    }
239
240    /// Provides data under the given capability. Flushes automatically when the capability
241    /// changes, so each output message carries the capability its data was given under.
242    fn give<D, C>(&mut self, cap: &C, data: D)
243    where
244        CB: PushInto<D>,
245        C: ShipUnder<T>,
246    {
247        if let Some(capability) = &self.capability
248            && !cap.same_times(capability)
249        {
250            self.flush();
251            self.capability = None;
252        }
253        if self.capability.is_none() {
254            self.capability = Some(cap.to_set());
255        }
256
257        self.builder.push_into(data);
258        while let Some(container) = self.builder.extract() {
259            self.output
260                .give(self.capability.as_ref().expect("must exist"), container);
261        }
262    }
263}
264
265pub struct AsyncOutputHandle<T: Timestamp, CB: ContainerBuilder> {
266    inner: Rc<RefCell<AsyncOutputHandleInner<T, CB>>>,
267    index: usize,
268}
269
270impl<T, C> AsyncOutputHandle<T, CapacityContainerBuilder<C>>
271where
272    T: Timestamp,
273    C: Container + Clone + 'static,
274{
275    #[inline]
276    pub fn give_container<Cap: ShipUnder<T>>(&self, cap: &Cap, container: &mut C) {
277        let mut inner = self.inner.borrow_mut();
278        inner.flush();
279        inner.output.give(cap, container);
280    }
281}
282
283impl<T, CB> AsyncOutputHandle<T, CB>
284where
285    T: Timestamp,
286    CB: ContainerBuilder,
287{
288    fn new(output: Output<T, CB::Container>, index: usize) -> Self {
289        let inner = AsyncOutputHandleInner {
290            output,
291            capability: None,
292            builder: CB::default(),
293        };
294        Self {
295            inner: Rc::new(RefCell::new(inner)),
296            index,
297        }
298    }
299
300    fn cease(&self) {
301        self.inner.borrow_mut().cease();
302    }
303}
304
305impl<T, CB> AsyncOutputHandle<T, CB>
306where
307    T: Timestamp,
308    CB: ContainerBuilder,
309{
310    pub fn give<D, Cap: ShipUnder<T>>(&self, cap: &Cap, data: D)
311    where
312        CB: PushInto<D>,
313    {
314        self.inner.borrow_mut().give(cap, data);
315    }
316}
317
318/// Maximum bytes [`AsyncOutputHandle::give_fueled`] emits through one handle before yielding back
319/// to timely.
320pub const MAX_OUTSTANDING_BYTES: usize = 128 * 1024 * 1024;
321
322impl<T, D> AsyncOutputHandle<T, FueledBuilder<CapacityContainerBuilder<Vec<D>>>>
323where
324    D: Clone + 'static,
325    T: Timestamp,
326{
327    /// Provides one record at the time specified by the capability and
328    /// charges `size_bytes` against the builder's fuel counter. Once at least
329    /// [`MAX_OUTSTANDING_BYTES`] have been emitted since the last yield
330    /// this method yields back to timely and resets the counter.
331    pub async fn give_fueled<D2, Cap: ShipUnder<T>>(&self, cap: &Cap, data: D2, size_bytes: usize)
332    where
333        FueledBuilder<CapacityContainerBuilder<Vec<D>>>: PushInto<D2>,
334    {
335        let should_yield = {
336            let mut handle = self.inner.borrow_mut();
337            handle.give(cap, data);
338            let bytes = &handle.builder.bytes;
339            bytes.set(bytes.get() + size_bytes);
340            let should_yield = bytes.get() > MAX_OUTSTANDING_BYTES;
341            if should_yield {
342                bytes.set(0);
343            }
344            should_yield
345        };
346        if should_yield {
347            tokio::task::yield_now().await;
348        }
349    }
350}
351
352impl<T: Timestamp, CB: ContainerBuilder> Clone for AsyncOutputHandle<T, CB> {
353    fn clone(&self) -> Self {
354        Self {
355            inner: Rc::clone(&self.inner),
356            index: self.index,
357        }
358    }
359}
360
361/// A trait describing the connection behavior between an input of an operator and zero or more of
362/// its outputs.
363///
364/// The connected inputs ([`ConnectedToOne`], [`ConnectedToMany`]) hand the operator body the
365/// message's capabilities as a [`CapabilitySet`] per output, which works in any scope. An operator
366/// that wants a single point calls `CapabilitySet::delayed`, or [`sole_capability`] in a totally
367/// ordered scope. [`Disconnected`] holds no capability at all and so reports the message's stamp as
368/// plain times.
369pub trait InputConnection<T: Timestamp> {
370    /// The capability type associated with this connection behavior.
371    type Capability;
372
373    /// Generates a summary description of the connection behavior given the number of outputs.
374    fn describe(&self, outputs: usize) -> Vec<Antichain<T::Summary>>;
375
376    /// Accepts an input capability.
377    fn accept(&self, input_cap: InputCapability<T>) -> Self::Capability;
378}
379
380/// A marker type representing a disconnected input.
381pub struct Disconnected;
382
383impl<T: Timestamp> InputConnection<T> for Disconnected {
384    type Capability = Stamp<T>;
385
386    fn describe(&self, outputs: usize) -> Vec<Antichain<T::Summary>> {
387        vec![Antichain::new(); outputs]
388    }
389
390    fn accept(&self, input_cap: InputCapability<T>) -> Self::Capability {
391        input_cap.stamp().clone()
392    }
393}
394
395/// A marker type representing an input connected to exactly one output.
396pub struct ConnectedToOne(usize);
397
398impl<T: Timestamp> InputConnection<T> for ConnectedToOne {
399    type Capability = CapabilitySet<T>;
400
401    fn describe(&self, outputs: usize) -> Vec<Antichain<T::Summary>> {
402        let mut summary = vec![Antichain::new(); outputs];
403        summary[self.0] = Antichain::from_elem(T::Summary::default());
404        summary
405    }
406
407    fn accept(&self, input_cap: InputCapability<T>) -> Self::Capability {
408        input_cap.retain_stamp(self.0)
409    }
410}
411
412/// A marker type representing an input connected to many outputs.
413pub struct ConnectedToMany<const N: usize>([usize; N]);
414
415impl<const N: usize, T: Timestamp> InputConnection<T> for ConnectedToMany<N> {
416    type Capability = [CapabilitySet<T>; N];
417
418    fn describe(&self, outputs: usize) -> Vec<Antichain<T::Summary>> {
419        let mut summary = vec![Antichain::new(); outputs];
420        for output in self.0 {
421            summary[output] = Antichain::from_elem(T::Summary::default());
422        }
423        summary
424    }
425
426    fn accept(&self, input_cap: InputCapability<T>) -> Self::Capability {
427        self.0.map(|output| input_cap.retain_stamp(output))
428    }
429}
430
431/// The one capability of a message in a totally ordered scope.
432///
433/// A capability set over a totally ordered timestamp holds at most one capability. It holds none
434/// for a message sent under no capabilities, which makes no progress claims, and this function
435/// panics on such a set.
436pub fn sole_capability<T: Timestamp + TotalOrder>(
437    capabilities: &CapabilitySet<T>,
438) -> &Capability<T> {
439    match &capabilities[..] {
440        [capability] => capability,
441        [] => panic!("message stamped with no capabilities"),
442        _ => unreachable!("a capability set over a total order holds at most one capability"),
443    }
444}
445
446/// A capability an output handle can ship data under: one capability, or a whole set.
447///
448/// The set is the general form. A lone [`Capability`] is the one-element case, which is what an
449/// operator in a totally ordered scope holds.
450pub trait ShipUnder<T: Timestamp>: CapabilityTrait<T> {
451    /// The capabilities to hold while the data is buffered.
452    fn to_set(&self) -> CapabilitySet<T>;
453
454    /// Whether `set` holds capabilities for exactly the times of `self`.
455    ///
456    /// Called per record, so it compares borrowed times rather than building stamps.
457    fn same_times(&self, set: &CapabilitySet<T>) -> bool;
458}
459
460impl<T: Timestamp> ShipUnder<T> for Capability<T> {
461    fn to_set(&self) -> CapabilitySet<T> {
462        std::iter::once(self.clone()).collect()
463    }
464
465    fn same_times(&self, set: &CapabilitySet<T>) -> bool {
466        matches!(&set[..], [only] if only.time() == self.time())
467    }
468}
469
470impl<T: Timestamp> ShipUnder<T> for CapabilitySet<T> {
471    fn to_set(&self) -> CapabilitySet<T> {
472        self.clone()
473    }
474
475    fn same_times(&self, set: &CapabilitySet<T>) -> bool {
476        // Both sets are antichains, so equal lengths and containment make them equal.
477        self.len() == set.len()
478            && self
479                .iter()
480                .all(|cap| set.iter().any(|other| other.time() == cap.time()))
481    }
482}
483
484/// A helper trait abstracting over an output handle. It facilitates passing type erased
485/// output handles during operator construction.
486/// It is not meant to be implemented by users.
487pub trait OutputIndex {
488    /// The output index of this handle.
489    fn index(&self) -> usize;
490}
491
492impl<T: Timestamp, CB: ContainerBuilder> OutputIndex for AsyncOutputHandle<T, CB> {
493    fn index(&self) -> usize {
494        self.index
495    }
496}
497
498impl<'scope, T: Timestamp> OperatorBuilder<'scope, T> {
499    /// Allocates a new generic async operator builder from its containing scope.
500    pub fn new(name: String, scope: Scope<'scope, T>) -> Self {
501        let builder = OperatorBuilderRc::new(name, scope);
502        let info = builder.operator_info();
503        let activator = scope.activator_for(Rc::clone(&info.address));
504        let sync_activator = scope.worker().sync_activator_for(info.address.to_vec());
505        let operator_waker = TimelyWaker {
506            activator: sync_activator,
507            active: AtomicBool::new(false),
508            task_ready: AtomicBool::new(true),
509        };
510        let (shutdown_handle, shutdown_button) = button(scope, info.address);
511
512        OperatorBuilder {
513            builder,
514            activator,
515            operator_waker: Arc::new(operator_waker),
516            input_frontiers: Default::default(),
517            input_queues: Default::default(),
518            output_flushes: Default::default(),
519            shutdown_handle,
520            shutdown_button,
521        }
522    }
523
524    /// Adds a new input that is connected to the specified output, returning the async input handle to use.
525    pub fn new_input_for<D, P>(
526        &mut self,
527        stream: TimelyStream<'scope, T, D>,
528        pact: P,
529        output: &dyn OutputIndex,
530    ) -> AsyncInputHandle<T, D, ConnectedToOne>
531    where
532        D: Container + Clone + 'static,
533        P: ParallelizationContract<T, D>,
534    {
535        let index = output.index();
536        assert!(index < self.builder.shape().outputs());
537        self.new_input_connection(stream, pact, ConnectedToOne(index))
538    }
539
540    /// Adds a new input that is connected to the specified outputs, returning the async input handle to use.
541    pub fn new_input_for_many<const N: usize, D, P>(
542        &mut self,
543        stream: TimelyStream<'scope, T, D>,
544        pact: P,
545        outputs: [&dyn OutputIndex; N],
546    ) -> AsyncInputHandle<T, D, ConnectedToMany<N>>
547    where
548        D: Container + Clone + 'static,
549        P: ParallelizationContract<T, D>,
550    {
551        let indices = outputs.map(|output| output.index());
552        for index in indices {
553            assert!(index < self.builder.shape().outputs());
554        }
555        self.new_input_connection(stream, pact, ConnectedToMany(indices))
556    }
557
558    /// Adds a new input that is not connected to any output, returning the async input handle to use.
559    pub fn new_disconnected_input<D, P>(
560        &mut self,
561        stream: TimelyStream<'scope, T, D>,
562        pact: P,
563    ) -> AsyncInputHandle<T, D, Disconnected>
564    where
565        D: Container + Clone + 'static,
566        P: ParallelizationContract<T, D>,
567    {
568        self.new_input_connection(stream, pact, Disconnected)
569    }
570
571    /// Adds a new input with connection information, returning the async input handle to use.
572    pub fn new_input_connection<D, P, C>(
573        &mut self,
574        stream: TimelyStream<'scope, T, D>,
575        pact: P,
576        connection: C,
577    ) -> AsyncInputHandle<T, D, C>
578    where
579        D: Container + Clone + 'static,
580        P: ParallelizationContract<T, D>,
581        C: InputConnection<T> + 'static,
582    {
583        self.input_frontiers
584            .push(Antichain::from_elem(T::minimum()));
585
586        let outputs = self.builder.shape().outputs();
587        let handle = self.builder.new_input_connection(
588            stream,
589            pact,
590            connection.describe(outputs).into_iter().enumerate(),
591        );
592
593        let waker = Default::default();
594        let queue = Default::default();
595        let input_queue = InputHandleQueue {
596            queue: Rc::clone(&queue),
597            waker: Rc::clone(&waker),
598            connection,
599            handle,
600        };
601        self.input_queues.push(Box::new(input_queue));
602
603        AsyncInputHandle {
604            queue,
605            waker,
606            done: false,
607        }
608    }
609
610    /// Adds a new output, returning the output handle and stream.
611    pub fn new_output<CB: ContainerBuilder>(
612        &mut self,
613    ) -> (
614        AsyncOutputHandle<T, CB>,
615        TimelyStream<'scope, T, CB::Container>,
616    ) {
617        let index = self.builder.shape().outputs();
618
619        let (output, stream) = self.builder.new_output_connection([]);
620
621        let handle = AsyncOutputHandle::new(output, index);
622
623        let flush_handle = handle.clone();
624        self.output_flushes
625            .push(Box::new(move || flush_handle.cease()));
626
627        (handle, stream)
628    }
629
630    /// Creates an operator implementation from supplied logic constructor. It returns a shutdown
631    /// button that when pressed it will cause the logic future to be dropped and input handles to
632    /// be drained. The button can be converted into a token by using
633    /// [`Button::press_on_drop`]
634    pub fn build<B, L>(self, constructor: B) -> Button
635    where
636        B: FnOnce(Vec<Capability<T>>) -> L,
637        L: Future + 'static,
638    {
639        let operator_waker = self.operator_waker;
640        let mut input_frontiers = self.input_frontiers;
641        let mut input_queues = self.input_queues;
642        let mut output_flushes = self.output_flushes;
643        let mut shutdown_handle = self.shutdown_handle;
644        self.builder.build_reschedule(move |caps| {
645            let mut logic_fut = Some(Box::pin(constructor(caps)));
646            move |new_frontiers| {
647                operator_waker.active.store(true, Ordering::SeqCst);
648                for (i, queue) in input_queues.iter_mut().enumerate() {
649                    // First, discover if there are any frontier notifications
650                    let cur = &mut input_frontiers[i];
651                    let new = new_frontiers[i].frontier();
652                    if PartialOrder::less_than(&cur.borrow(), &new) {
653                        queue.notify_progress(new.to_owned());
654                        *cur = new.to_owned();
655                    }
656                    // Then accept all input into local queues. This step registers the received
657                    // messages with progress tracking.
658                    queue.accept_input();
659                }
660                operator_waker.active.store(false, Ordering::SeqCst);
661
662                // If our worker pressed the button we stop scheduling the logic future and/or
663                // draining the input handles to stop producing data and frontier updates
664                // downstream.
665                if shutdown_handle.local_pressed() {
666                    // When all workers press their buttons we drop the logic future and start
667                    // draining the input handles.
668                    if shutdown_handle.all_pressed() {
669                        logic_fut = None;
670                        for queue in input_queues.iter_mut() {
671                            queue.drain_input();
672                        }
673                        false
674                    } else {
675                        true
676                    }
677                } else {
678                    // Schedule the logic future if any of the wakers above marked the task as ready
679                    if let Some(fut) = logic_fut.as_mut() {
680                        if operator_waker.task_ready.load(Ordering::SeqCst) {
681                            let waker = futures_util::task::waker_ref(&operator_waker);
682                            let mut cx = Context::from_waker(&waker);
683                            operator_waker.task_ready.store(false, Ordering::SeqCst);
684                            if Pin::new(fut).poll(&mut cx).is_ready() {
685                                // We're done with logic so deallocate the task
686                                logic_fut = None;
687                            }
688                            // Flush all the outputs before exiting
689                            for flush in output_flushes.iter_mut() {
690                                (flush)();
691                            }
692                        }
693                    }
694
695                    // The timely operator needs to be kept alive if the task is pending
696                    if logic_fut.is_some() {
697                        true
698                    } else {
699                        // Othewise we should keep draining all inputs
700                        for queue in input_queues.iter_mut() {
701                            queue.drain_input();
702                        }
703                        false
704                    }
705                }
706            }
707        });
708
709        self.shutdown_button
710    }
711
712    /// Creates a fallible operator implementation from supplied logic constructor. If the `Future`
713    /// resolves to an error it will be emitted in the returned error stream and then the operator
714    /// will wait indefinitely until the shutdown button is pressed.
715    ///
716    /// # Capability handling
717    ///
718    /// Unlike [`OperatorBuilder::build`], this method does not give owned capabilities to the
719    /// constructor. All initial capabilities are wrapped in a `CapabilitySet` and a mutable
720    /// reference to them is given instead. This is done to avoid storing owned capabilities in the
721    /// state of the logic future which would make using the `?` operator unsafe, since the
722    /// frontiers would incorrectly advance, potentially causing incorrect actions downstream.
723    ///
724    /// ```ignore
725    /// builder.build_fallible(|caps| Box::pin(async move {
726    ///     // Assert that we have the number of capabilities we expect
727    ///     // `cap` will be a `&mut Option<Capability<T>>`:
728    ///     let [cap_set]: &mut [_; 1] = caps.try_into().unwrap();
729    ///
730    ///     // Using cap to send data:
731    ///     output.give(&cap_set[0], 42);
732    ///
733    ///     // Using cap_set to downgrade it:
734    ///     cap_set.downgrade([]);
735    ///
736    ///     // Explicitly dropping the capability:
737    ///     // Simply running `drop(cap_set)` will only drop the reference and not the capability set itself!
738    ///     *cap_set = CapabilitySet::new();
739    ///
740    ///     // !! BIG WARNING !!:
741    ///     // It is tempting to `take` the capability out of the set for convenience. This will
742    ///     // move the capability into the future state, tying its lifetime to it, which will get
743    ///     // dropped when an error is hit, causing incorrect progress statements.
744    ///     let cap = cap_set.delayed(&Timestamp::minimum());
745    ///     *cap_set = CapabilitySet::new(); // DO NOT DO THIS
746    /// }));
747    /// ```
748    pub fn build_fallible<E: 'static, F>(
749        mut self,
750        constructor: F,
751    ) -> (Button, StreamVec<'scope, T, Rc<E>>)
752    where
753        F: for<'a> FnOnce(
754                &'a mut [CapabilitySet<T>],
755            ) -> Pin<Box<dyn Future<Output = Result<(), E>> + 'a>>
756            + 'static,
757    {
758        // Create a new completely disconnected output
759        let (error_output, error_stream) = self.new_output::<CapacityContainerBuilder<_>>();
760        let button = self.build(|mut caps| async move {
761            let error_cap = caps.pop().unwrap();
762            let mut caps = caps
763                .into_iter()
764                .map(CapabilitySet::from_elem)
765                .collect::<Vec<_>>();
766            if let Err(err) = constructor(&mut *caps).await {
767                error_output.give(&error_cap, Rc::new(err));
768                drop(error_cap);
769                // IMPORTANT: wedge this operator until the button is pressed. Returning would drop
770                // the capabilities and could produce incorrect progress statements.
771                std::future::pending().await
772            }
773        });
774        (button, error_stream)
775    }
776
777    /// Creates operator info for the operator.
778    pub fn operator_info(&self) -> OperatorInfo {
779        self.builder.operator_info()
780    }
781
782    /// Returns the activator for the operator.
783    pub fn activator(&self) -> &Activator {
784        &self.activator
785    }
786}
787
788/// Creates a new coordinated button the worker configuration described by `scope`.
789pub fn button<'scope, T: Timestamp>(
790    scope: Scope<'scope, T>,
791    addr: Rc<[usize]>,
792) -> (ButtonHandle, Button) {
793    let index = scope.worker().new_identifier();
794    let (pushers, puller) = scope.worker().allocate(index, addr);
795
796    let local_pressed = Rc::new(Cell::new(false));
797
798    let handle = ButtonHandle {
799        buttons_remaining: scope.peers(),
800        local_pressed: Rc::clone(&local_pressed),
801        puller,
802    };
803
804    let token = Button {
805        pushers,
806        local_pressed,
807    };
808
809    (handle, token)
810}
811
812/// A button that can be used to coordinate an action after all workers have pressed it.
813pub struct ButtonHandle {
814    /// The number of buttons still unpressed among workers.
815    buttons_remaining: usize,
816    /// A flag indicating whether this worker has pressed its button.
817    local_pressed: Rc<Cell<bool>>,
818    puller: Box<dyn Pull<Bincode<bool>>>,
819}
820
821impl ButtonHandle {
822    /// Returns whether this worker has pressed its button.
823    pub fn local_pressed(&self) -> bool {
824        self.local_pressed.get()
825    }
826
827    /// Returns whether all workers have pressed their buttons.
828    pub fn all_pressed(&mut self) -> bool {
829        while self.puller.recv().is_some() {
830            self.buttons_remaining -= 1;
831        }
832        self.buttons_remaining == 0
833    }
834}
835
836pub struct Button {
837    pushers: Vec<Box<dyn Push<Bincode<bool>>>>,
838    local_pressed: Rc<Cell<bool>>,
839}
840
841impl Button {
842    /// Presses the button. It is safe to call this function multiple times.
843    pub fn press(&mut self) {
844        for mut pusher in self.pushers.drain(..) {
845            pusher.send(Bincode::from(true));
846            pusher.done();
847        }
848        self.local_pressed.set(true);
849    }
850
851    /// Converts this button into a deadman's switch that will automatically press the button when
852    /// dropped.
853    pub fn press_on_drop(self) -> PressOnDropButton {
854        PressOnDropButton(self)
855    }
856}
857
858pub struct PressOnDropButton(Button);
859
860impl Drop for PressOnDropButton {
861    fn drop(&mut self) {
862        self.0.press();
863    }
864}
865
866#[cfg(test)]
867mod test {
868    use futures_util::StreamExt;
869    use timely::WorkerConfig;
870    use timely::dataflow::channels::pact::Pipeline;
871    use timely::dataflow::operators::Capture;
872    use timely::dataflow::operators::capture::Extract;
873    use timely::dataflow::operators::vec::ToStream;
874
875    use super::*;
876
877    #[mz_ore::test]
878    fn async_operator() {
879        let capture = timely::example(|scope| {
880            let input = (0..10).to_stream(scope);
881
882            let mut op = OperatorBuilder::new("async_passthru".to_string(), input.scope());
883            let (output, output_stream) = op.new_output::<CapacityContainerBuilder<_>>();
884            let mut input_handle = op.new_input_for(input, Pipeline, &output);
885
886            op.build(move |_capabilities| async move {
887                tokio::task::yield_now().await;
888                while let Some(event) = input_handle.next().await {
889                    match event {
890                        Event::Data(cap, data) => {
891                            for item in data.iter().copied() {
892                                tokio::task::yield_now().await;
893                                output.give(&cap, item);
894                            }
895                        }
896                        Event::Progress(_frontier) => {}
897                    }
898                }
899            });
900
901            output_stream.capture()
902        });
903        let extracted = capture.extract();
904
905        assert_eq!(extracted, vec![(0, vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9])]);
906    }
907
908    #[mz_ore::test]
909    fn gh_18837() {
910        let (builders, other) = timely::CommunicationConfig::Process(2).try_build().unwrap();
911        timely::execute::execute_from(builders, other, WorkerConfig::default(), |worker| {
912            let index = worker.index();
913            let tokens = worker.dataflow::<u64, _, _>(move |scope| {
914                let mut producer = OperatorBuilder::new("producer".to_string(), scope.clone());
915                let (_output, output_stream) =
916                    producer.new_output::<CapacityContainerBuilder<Vec<usize>>>();
917                let producer_button = producer.build(move |mut capabilities| async move {
918                    let mut cap = capabilities.pop().unwrap();
919                    if index != 0 {
920                        return;
921                    }
922                    // Worker 0 downgrades to 1 and keeps the capability around forever
923                    cap.downgrade(&1);
924                    std::future::pending().await
925                });
926
927                let mut consumer = OperatorBuilder::new("consumer".to_string(), scope.clone());
928                let mut input_handle = consumer.new_disconnected_input(output_stream, Pipeline);
929                let consumer_button = consumer.build(move |_| async move {
930                    while let Some(event) = input_handle.next().await {
931                        if let Event::Progress(frontier) = event {
932                            // We should never observe a frontier greater than [1]
933                            assert!(frontier.less_equal(&1));
934                        }
935                    }
936                });
937
938                (
939                    producer_button.press_on_drop(),
940                    consumer_button.press_on_drop(),
941                )
942            });
943
944            // Run dataflow until only worker 0 holds the frontier to [1]
945            for _ in 0..100 {
946                worker.step();
947            }
948            // Then drop the tokens of worker 0
949            if index == 0 {
950                drop(tokens)
951            }
952            // And step the dataflow some more to ensure consumers don't observe frontiers advancing.
953            for _ in 0..100 {
954                worker.step();
955            }
956        })
957        .expect("timely panicked");
958    }
959}
mz_timely_util/builder_async.rs

mz_timely_util/
builder_async.rs
