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

mz_timely_util/
builder_async.rs
