mz_compute/render/
continual_task.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! A continual task presents as something like a `TRIGGER`: it watches some
11//! _input_ and whenever it changes at time `T`, executes a SQL txn, writing to
12//! some _output_ at the same time `T`. It can also read anything in materialize
13//! as a _reference_, most notably including the output.
14//!
15//! Only reacting to new inputs (and not the full history) makes a CT's
16//! rehydration time independent of the size of the inputs (NB this is not true
17//! for references), enabling things like writing UPSERT on top of an
18//! append-only shard in SQL (ignore the obvious bug with my upsert impl):
19//!
20//! ```sql
21//! CREATE CONTINUAL TASK upsert (key INT, val INT) ON INPUT append_only AS (
22//!     DELETE FROM upsert WHERE key IN (SELECT key FROM append_only);
23//!     INSERT INTO upsert SELECT key, max(val) FROM append_only GROUP BY key;
24//! )
25//! ```
26//!
27//! Unlike a materialized view, the continual task does not update outputs if
28//! references later change. This enables things like auditing:
29//!
30//! ```sql
31//! CREATE CONTINUAL TASK audit_log (count INT8) ON INPUT anomalies AS (
32//!     INSERT INTO audit_log SELECT * FROM anomalies;
33//! )
34//! ```
35//!
36//! Rough implementation overview:
37//! - A CT is created and starts at some `start_ts` optionally later dropped and
38//!   stopped at some `end_ts`.
39//! - A CT takes one or more _input_s. These must be persist shards (i.e. TABLE,
40//!   SOURCE, MV, but not VIEW).
41//! - A CT has one or more _output_s. The outputs are (initially) owned by the
42//!   task and cannot be written to by other parts of the system.
43//! - The task is run for each time one of the inputs changes starting at
44//!   `start_ts`.
45//! - It is given the changes in its inputs at time `T` as diffs.
46//!   - These are presented as two SQL relations with just the inserts/deletes.
47//!   - NB: A full collection for the input can always be recovered by also
48//!     using the input as a "reference" (see below) and applying the diffs.
49//! - The task logic is expressed as a SQL transaction that does all reads at
50//!   commits all writes at `T`
51//!   - The notable exception to this is self-referential reads of the CT
52//!     output. See below for how that works.
53//! - This logic can _reference_ any nameable object in the system, not just the
54//!   inputs.
55//!   - However, the logic/transaction can mutate only the outputs.
56//! - Summary of differences between inputs and references:
57//!   - The task receives snapshot + changes for references (like regular
58//!     dataflow inputs today) but only changes for inputs.
59//!   - The task only produces output in response to changes in the inputs but
60//!     not in response to changes in the references.
61//! - Instead of re-evaluating the task logic from scratch for each input time,
62//!   we maintain the collection representing desired writes to the output(s) as
63//!   a dataflow.
64//! - The task dataflow is tied to a `CLUSTER` and runs on each `REPLICA`.
65//!   - HA strategy: multi-replica clusters race to commit and the losers throw
66//!     away the result.
67//!
68//! ## Self-References
69//!
70//! Self-references must be handled differently from other reads. When computing
71//! the proposed write to some output at `T`, we can only know the contents of
72//! it through `T-1` (the exclusive upper is `T`).
73//!
74//! We address this by initially assuming that the output contains no changes at
75//! `T`, then evaluating each of the statements in order, allowing them to see
76//! the proposed output changes made by the previous statements. By default,
77//! this is stopped after one iteration and proposed output diffs are committed
78//! if possible. (We could also add options for iterating to a fixpoint,
79//! stop/error after N iters, etc.) Then to compute the changes at `T+1`, we
80//! read in what was actually written to the output at `T` (maybe some other
81//! replica wrote something different) and begin again.
82//!
83//! The above is very similar to how timely/differential dataflow iteration
84//! works, except that our feedback loop goes through persist and the loop
85//! timestamp is already `mz_repr::Timestamp`.
86//!
87//! This is implemented as follows:
88//! - `let I = persist_source(self-reference)`
89//! - Transform `I` such that the contents at `T-1` are presented at `T` (i.e.
90//!   initially assume `T` is unchanged from `T-1`).
91//! - TODO(ct3): Actually implement the following.
92//! - In an iteration sub-scope:
93//!   - Bring `I` into the sub-scope and `let proposed = Variable`.
94//!   - We need a collection that at `(T, 0)` is always the contents of `I` at
95//!     `T`, but at `(T, 1...)` contains the proposed diffs by the CT logic. We
96//!     can construct it by concatenating `I` with `proposed` except that we
97//!     also need to retract everything in `proposed` for the next `(T+1, 0)`
98//!     (because `I` is the source of truth for what actually committed).
99//!  - `let R = retract_at_next_outer_ts(proposed)`
100//!  - `let result = logic(concat(I, proposed, R))`
101//!  - `proposed.set(result)`
102//! - Then we return `proposed.leave()` for attempted write to persist.
103//!
104//! ## As Ofs and Output Uppers
105//!
106//! - A continual task is first created with an initial as_of `I`. It is
107//!   initially rendered at as_of `I==A` but as it makes progress, it may be
108//!   rendered at later as_ofs `I<A`.
109//! - It is required that the output collection springs into existence at `I`
110//!   (i.e. receives the initial contents at `I`).
111//!   - For a snapshot CT, the full contents of the input at `I` are run through
112//!     the CT logic and written at `I`.
113//!   - For a non-snapshot CT, the collection is defined to be empty at `I`
114//!     (i.e. if the input happened to be written exactly at `I`, we'd ignore
115//!     it) and then start writing at `I+1`.
116//! - As documented in [DataflowDescription::as_of], `A` is the time we render
117//!   the inputs.
118//!   - An MV with an as_of of `A` will both have inputs rendered at `A` and
119//!     also the first time it could write is also `A`.
120//!   - A CT is the same on the initial render (`I==A`), but on renders after it
121//!     has made progress (`I<A`) the first time that  it could potentially
122//!     write is `A+1`. This is because a persist_source started with
123//!     SnapshotMode::Exclude can only start emitting diffs at `as_of+1`.
124//!   - As a result, we hold back the since on inputs to be strictly less than
125//!     the upper of the output. (This is only necessary for CTs, but we also do
126//!     it for MVs to avoid the special case.)
127//!   - For CT "inputs" (which are disallowed from being the output), we render
128//!     the persist_source with as_of `A`.
129//!     - When `I==A` we include the snapshot iff the snapshot option is used.
130//!     - When `I<A` we always exclude the snapshot. It would be unnecessary and
131//!       this is an absolutely critical performance optimization to make CT
132//!       rehydration times independent of input size.
133//!   - For CT "references", we render the persist_source with as_of `A` and
134//!     always include the snapshot.
135//!     - There is one subtlety: self-references on the initial render. We need
136//!       the contents to be available at `A-1`, so that we can do the
137//!       step_forward described above to get it at `A`. However, the collection
138//!       springs into existence at `I`, so we when `I==A`, we're not allowed to
139//!       read it as_of `A-1` (the since of the shard may have advanced past
140//!       that). We address this by rendering the persist_source as normal at
141//!       `A`. On startup, persist_source immediately downgrades its frontier to
142//!       `A` (making `A-1` readable). Combined with step_forward, this is
143//!       enough to unblock the CT self-reference. We do however have to tweak
144//!       the `suppress_early_progress` operator to use `A-1` instead of `A` for
145//!       this case.
146//!     - On subsequent renders, self-references work as normal.
147
148use std::any::Any;
149use std::cell::RefCell;
150use std::collections::BTreeSet;
151use std::rc::Rc;
152use std::sync::Arc;
153
154use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
155use differential_dataflow::difference::Semigroup;
156use differential_dataflow::lattice::Lattice;
157use differential_dataflow::{AsCollection, Collection, Hashable};
158use futures::{Future, FutureExt, StreamExt};
159use mz_compute_types::dataflows::DataflowDescription;
160use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc, ContinualTaskConnection};
161use mz_ore::cast::CastFrom;
162use mz_ore::collections::HashMap;
163use mz_persist_client::Diagnostics;
164use mz_persist_client::error::UpperMismatch;
165use mz_persist_client::operators::shard_source::SnapshotMode;
166use mz_persist_client::write::WriteHandle;
167use mz_persist_types::codec_impls::UnitSchema;
168use mz_repr::{Diff, GlobalId, Row, Timestamp};
169use mz_storage_types::StorageDiff;
170use mz_storage_types::controller::CollectionMetadata;
171use mz_storage_types::errors::DataflowError;
172use mz_storage_types::sources::SourceData;
173use mz_timely_util::builder_async::{Button, Event, OperatorBuilder as AsyncOperatorBuilder};
174use mz_timely_util::operator::CollectionExt;
175use mz_timely_util::probe;
176use mz_timely_util::probe::ProbeNotify;
177use timely::dataflow::channels::pact::{Exchange, Pipeline};
178use timely::dataflow::operators::generic::builder_rc::OperatorBuilder;
179use timely::dataflow::operators::{Filter, FrontierNotificator, Map, Operator};
180use timely::dataflow::{ProbeHandle, Scope};
181use timely::progress::frontier::AntichainRef;
182use timely::progress::{Antichain, Timestamp as _};
183use timely::{Data, PartialOrder};
184use tracing::debug;
185
186use crate::compute_state::ComputeState;
187use crate::render::StartSignal;
188use crate::render::sinks::SinkRender;
189use crate::sink::ConsolidatingVec;
190
191pub(crate) struct ContinualTaskCtx<G: Scope<Timestamp = Timestamp>> {
192    name: Option<String>,
193    dataflow_as_of: Option<Antichain<Timestamp>>,
194    inputs_with_snapshot: Option<bool>,
195    ct_inputs: BTreeSet<GlobalId>,
196    ct_outputs: BTreeSet<GlobalId>,
197    pub ct_times: Vec<Collection<G, (), Diff>>,
198}
199
200/// An encapsulation of the transformation logic necessary on data coming into a
201/// continual task.
202///
203/// NB: In continual task jargon, an "input" contains diffs and a "reference" is
204/// a normal source/collection.
205pub(crate) enum ContinualTaskSourceTransformer {
206    /// A collection containing, at each time T, exactly the inserts at time T
207    /// in the transformed collection.
208    ///
209    /// For example:
210    /// - Input: {} at 0, {1} at 1, {1} at 2, ...
211    /// - Output: {} at 0, {1} at 1, {} at 2, ...
212    ///
213    /// We'll presumably have the same for deletes eventually, but it's not
214    /// exposed in the SQL frontend yet.
215    InsertsInput {
216        source_id: GlobalId,
217        with_snapshot: bool,
218    },
219    /// A self-reference to the continual task's output. This is essentially a
220    /// timely feedback loop via the persist shard. See module rustdoc for how
221    /// this works.
222    SelfReference { source_id: GlobalId },
223    /// A normal collection (no-op transformation).
224    NormalReference,
225}
226
227impl ContinualTaskSourceTransformer {
228    /// The persist_source `SnapshotMode` to use when reading this source.
229    pub fn snapshot_mode(&self) -> SnapshotMode {
230        use ContinualTaskSourceTransformer::*;
231        match self {
232            InsertsInput {
233                with_snapshot: false,
234                ..
235            } => SnapshotMode::Exclude,
236            InsertsInput {
237                with_snapshot: true,
238                ..
239            }
240            | SelfReference { .. }
241            | NormalReference => SnapshotMode::Include,
242        }
243    }
244
245    /// Returns the as_of to use with the suppress_early_progress operator for
246    /// this source. See the module rustdoc for context.
247    pub fn suppress_early_progress_as_of(
248        &self,
249        as_of: Antichain<Timestamp>,
250    ) -> Antichain<Timestamp> {
251        use ContinualTaskSourceTransformer::*;
252        match self {
253            InsertsInput { .. } => as_of,
254            SelfReference { .. } => as_of
255                .iter()
256                .map(|x| x.step_back().unwrap_or_else(Timestamp::minimum))
257                .collect(),
258            NormalReference => as_of,
259        }
260    }
261
262    /// Performs the necessary transformation on the source collection.
263    ///
264    /// Returns the transformed "oks" and "errs" collections. Also returns the
265    /// appropriate `ct_times` collection used to inform the sink which times
266    /// were changed in the inputs.
267    pub fn transform<S: Scope<Timestamp = Timestamp>>(
268        &self,
269        oks: Collection<S, Row, Diff>,
270        errs: Collection<S, DataflowError, Diff>,
271    ) -> (
272        Collection<S, Row, Diff>,
273        Collection<S, DataflowError, Diff>,
274        Collection<S, (), Diff>,
275    ) {
276        use ContinualTaskSourceTransformer::*;
277        match self {
278            // Make a collection s.t, for each time T in the input, the output
279            // contains the inserts at T.
280            InsertsInput { source_id, .. } => {
281                let name = source_id.to_string();
282                // Keep only the inserts.
283                let oks = oks.inner.filter(|(_, _, diff)| diff.is_positive());
284                // Grab the original times for use in the sink operator.
285                let (oks, times) = oks.as_collection().times_extract(&name);
286                // Then retract everything at the next timestamp.
287                let oks = oks.inner.flat_map(|(row, ts, diff)| {
288                    let retract_ts = ts.step_forward();
289                    let negation = -diff;
290                    [(row.clone(), ts, diff), (row, retract_ts, negation)]
291                });
292                (oks.as_collection(), errs, times)
293            }
294            NormalReference => {
295                let times = Collection::empty(&oks.scope());
296                (oks, errs, times)
297            }
298            // When computing an self-referential output at `T`, start by
299            // assuming there are no changes from the contents at `T-1`. See the
300            // module rustdoc for how this fits into the larger picture.
301            SelfReference { source_id } => {
302                let name = source_id.to_string();
303                let times = Collection::empty(&oks.scope());
304                // step_forward will panic at runtime if it receives a data or
305                // capability with a time that cannot be stepped forward (i.e.
306                // because it is already the max). We're safe here because this
307                // is stepping `T-1` forward to `T`.
308                let oks = oks.step_forward(&name);
309                let errs = errs.step_forward(&name);
310                (oks, errs, times)
311            }
312        }
313    }
314}
315
316impl<G: Scope<Timestamp = Timestamp>> ContinualTaskCtx<G> {
317    pub fn new<P, S>(dataflow: &DataflowDescription<P, S, Timestamp>) -> Self {
318        let mut name = None;
319        let mut ct_inputs = BTreeSet::new();
320        let mut ct_outputs = BTreeSet::new();
321        let mut inputs_with_snapshot = None;
322        for (sink_id, sink) in &dataflow.sink_exports {
323            match &sink.connection {
324                ComputeSinkConnection::ContinualTask(ContinualTaskConnection {
325                    input_id, ..
326                }) => {
327                    ct_outputs.insert(*sink_id);
328                    ct_inputs.insert(*input_id);
329                    // There's only one CT sink per dataflow at this point.
330                    assert_eq!(name, None);
331                    name = Some(sink_id.to_string());
332                    assert_eq!(inputs_with_snapshot, None);
333                    match (
334                        sink.with_snapshot,
335                        dataflow.as_of.as_ref(),
336                        dataflow.initial_storage_as_of.as_ref(),
337                    ) {
338                        // User specified no snapshot when creating the CT.
339                        (false, _, _) => inputs_with_snapshot = Some(false),
340                        // User specified a snapshot but we're past the initial
341                        // as_of.
342                        (true, Some(as_of), Some(initial_as_of))
343                            if PartialOrder::less_than(initial_as_of, as_of) =>
344                        {
345                            inputs_with_snapshot = Some(false)
346                        }
347                        // User specified a snapshot and we're either at the
348                        // initial creation, or we don't know (builtin CTs). If
349                        // we don't know, it's always safe to fall back to
350                        // snapshotting, at worst it's wasted work and will get
351                        // filtered.
352                        (true, _, _) => inputs_with_snapshot = Some(true),
353                    }
354                }
355                _ => continue,
356            }
357        }
358        let mut ret = ContinualTaskCtx {
359            name,
360            dataflow_as_of: None,
361            inputs_with_snapshot,
362            ct_inputs,
363            ct_outputs,
364            ct_times: Vec::new(),
365        };
366        // Only clone the as_of if we're in a CT dataflow.
367        if ret.is_ct_dataflow() {
368            ret.dataflow_as_of = dataflow.as_of.clone();
369            // Sanity check that we have a name if we're in a CT dataflow.
370            assert!(ret.name.is_some());
371        }
372        ret
373    }
374
375    pub fn is_ct_dataflow(&self) -> bool {
376        // Inputs are non-empty iff outputs are non-empty.
377        assert_eq!(self.ct_inputs.is_empty(), self.ct_outputs.is_empty());
378        !self.ct_outputs.is_empty()
379    }
380
381    pub fn get_ct_source_transformer(
382        &self,
383        source_id: GlobalId,
384    ) -> Option<ContinualTaskSourceTransformer> {
385        let Some(inputs_with_snapshot) = self.inputs_with_snapshot else {
386            return None;
387        };
388        let transformer = match (
389            self.ct_inputs.contains(&source_id),
390            self.ct_outputs.contains(&source_id),
391        ) {
392            (false, false) => ContinualTaskSourceTransformer::NormalReference,
393            (false, true) => ContinualTaskSourceTransformer::SelfReference { source_id },
394            (true, false) => ContinualTaskSourceTransformer::InsertsInput {
395                source_id,
396                with_snapshot: inputs_with_snapshot,
397            },
398            (true, true) => panic!("ct output is not allowed to be an input"),
399        };
400        Some(transformer)
401    }
402
403    pub fn input_times(&self, scope: &G) -> Option<Collection<G, (), Diff>> {
404        // We have a name iff this is a CT dataflow.
405        assert_eq!(self.is_ct_dataflow(), self.name.is_some());
406        let Some(name) = self.name.as_ref() else {
407            return None;
408        };
409        // Note that self.ct_times might be empty (if the user didn't reference
410        // the input), but this still does the correct, though maybe useless,
411        // thing: no diffs coming into the input means no times to write at.
412        let ct_times = differential_dataflow::collection::concatenate(
413            &mut scope.clone(),
414            self.ct_times.iter().cloned(),
415        );
416        // Reduce this down to one update per-time-per-worker before exchanging
417        // it, so we don't waste work on unnecessarily high data volumes.
418        let ct_times = ct_times.times_reduce(name);
419        Some(ct_times)
420    }
421}
422
423impl<G> SinkRender<G> for ContinualTaskConnection<CollectionMetadata>
424where
425    G: Scope<Timestamp = Timestamp>,
426{
427    fn render_sink(
428        &self,
429        compute_state: &mut ComputeState,
430        _sink: &ComputeSinkDesc<CollectionMetadata>,
431        sink_id: GlobalId,
432        as_of: Antichain<Timestamp>,
433        start_signal: StartSignal,
434        oks: Collection<G, Row, Diff>,
435        errs: Collection<G, DataflowError, Diff>,
436        append_times: Option<Collection<G, (), Diff>>,
437        flow_control_probe: &probe::Handle<Timestamp>,
438    ) -> Option<Rc<dyn Any>> {
439        let name = sink_id.to_string();
440
441        let to_append = oks
442            .map(|x| SourceData(Ok(x)))
443            .concat(&errs.map(|x| SourceData(Err(x))));
444        let append_times = append_times.expect("should be provided by ContinualTaskCtx");
445
446        let write_handle = {
447            let clients = Arc::clone(&compute_state.persist_clients);
448            let metadata = self.storage_metadata.clone();
449            let handle_purpose = format!("ct_sink({})", name);
450            async move {
451                let client = clients
452                    .open(metadata.persist_location)
453                    .await
454                    .expect("valid location");
455                client
456                    .open_writer(
457                        metadata.data_shard,
458                        metadata.relation_desc.into(),
459                        UnitSchema.into(),
460                        Diagnostics {
461                            shard_name: sink_id.to_string(),
462                            handle_purpose,
463                        },
464                    )
465                    .await
466                    .expect("codecs should match")
467            }
468        };
469
470        let collection = compute_state.expect_collection_mut(sink_id);
471        let probe = ProbeHandle::default();
472        let to_append = to_append
473            .probe_with(&probe)
474            .inner
475            .probe_notify_with(vec![flow_control_probe.clone()])
476            .as_collection();
477        collection.compute_probe = Some(probe);
478        let sink_write_frontier = Rc::new(RefCell::new(Antichain::from_elem(Timestamp::minimum())));
479        collection.sink_write_frontier = Some(Rc::clone(&sink_write_frontier));
480
481        // TODO(ct1): Obey `compute_state.read_only_rx`
482        //
483        // Seemingly, the read-only env needs to tail the output shard and keep
484        // historical updates around until it sees that the output frontier
485        // advances beyond their times.
486        let sink_button = continual_task_sink(
487            &name,
488            to_append,
489            append_times,
490            as_of,
491            write_handle,
492            start_signal,
493            sink_write_frontier,
494        );
495        Some(Rc::new(sink_button.press_on_drop()))
496    }
497}
498
499fn continual_task_sink<G: Scope<Timestamp = Timestamp>>(
500    name: &str,
501    to_append: Collection<G, SourceData, Diff>,
502    append_times: Collection<G, (), Diff>,
503    as_of: Antichain<Timestamp>,
504    write_handle: impl Future<Output = WriteHandle<SourceData, (), Timestamp, StorageDiff>>
505    + Send
506    + 'static,
507    start_signal: StartSignal,
508    output_frontier: Rc<RefCell<Antichain<Timestamp>>>,
509) -> Button {
510    let scope = to_append.scope();
511    let mut op = AsyncOperatorBuilder::new(format!("ct_sink({})", name), scope.clone());
512
513    // TODO(ct2): This all works perfectly well data parallel (assuming we
514    // broadcast the append_times). We just need to hook it up to the
515    // multi-worker persist-sink, but that requires some refactoring. This would
516    // also remove the need for this to be an async timely operator.
517    let active_worker = name.hashed();
518    let to_append_input =
519        op.new_input_for_many(&to_append.inner, Exchange::new(move |_| active_worker), []);
520    let append_times_input = op.new_input_for_many(
521        &append_times.inner,
522        Exchange::new(move |_| active_worker),
523        [],
524    );
525
526    let active_worker = usize::cast_from(active_worker) % scope.peers() == scope.index();
527    let button = op.build(move |_capabilities| async move {
528        if !active_worker {
529            output_frontier.borrow_mut().clear();
530            return;
531        }
532
533        // SUBTLE: The start_signal below may not be unblocked by the compute
534        // controller until it thinks the inputs are "ready" (i.e. readable at
535        // the as_of), but if the CT is self-referential, one of the inputs will
536        // be the output (which starts at `T::minimum()`, not the as_of). To
537        // break this cycle, before we even get the start signal, go ahead and
538        // advance the output's (exclusive) upper to the first time that this CT
539        // might write: `as_of+1`. Because we don't want this to happen on
540        // restarts, only do it if the upper is `T::minimum()`.
541        let mut write_handle = write_handle.await;
542        {
543            let res = write_handle
544                .compare_and_append_batch(
545                    &mut [],
546                    Antichain::from_elem(Timestamp::minimum()),
547                    as_of.clone(),
548                    true,
549                )
550                .await
551                .expect("usage was valid");
552            match res {
553                // We advanced the upper.
554                Ok(()) => {}
555                // Someone else advanced the upper.
556                Err(UpperMismatch { .. }) => {}
557            }
558        }
559
560        let () = start_signal.await;
561
562        #[derive(Debug)]
563        enum OpEvent<C> {
564            ToAppend(Event<Timestamp, C, Vec<(SourceData, Timestamp, Diff)>>),
565            AppendTimes(Event<Timestamp, C, Vec<((), Timestamp, Diff)>>),
566        }
567
568        impl<C: std::fmt::Debug> OpEvent<C> {
569            fn apply(self, state: &mut SinkState<SourceData, Timestamp>) {
570                debug!("ct_sink event {:?}", self);
571                match self {
572                    OpEvent::ToAppend(Event::Data(_cap, x)) => {
573                        for (k, t, d) in x {
574                            state.to_append.push(((k, t), d));
575                        }
576                    }
577                    OpEvent::ToAppend(Event::Progress(x)) => state.to_append_progress = x,
578                    OpEvent::AppendTimes(Event::Data(_cap, x)) => state
579                        .append_times
580                        .extend(x.into_iter().map(|((), t, _d)| t)),
581                    OpEvent::AppendTimes(Event::Progress(x)) => state.append_times_progress = x,
582                }
583            }
584        }
585
586        let to_insert_input = to_append_input.map(OpEvent::ToAppend);
587        let append_times_input = append_times_input.map(OpEvent::AppendTimes);
588        let mut op_inputs = futures::stream::select(to_insert_input, append_times_input);
589
590        let mut state = SinkState::new();
591        loop {
592            // Loop until we've processed all the work we can.
593            loop {
594                if PartialOrder::less_than(&*output_frontier.borrow(), &state.output_progress) {
595                    output_frontier.borrow_mut().clear();
596                    output_frontier
597                        .borrow_mut()
598                        .extend(state.output_progress.iter().cloned());
599                }
600
601                debug!("ct_sink about to process {:?}", state);
602                let Some((new_upper, to_append)) = state.process() else {
603                    break;
604                };
605                debug!("ct_sink got write {:?}: {:?}", new_upper, to_append);
606                state.output_progress =
607                    truncating_compare_and_append(&mut write_handle, to_append, new_upper).await;
608            }
609
610            // Then try to generate some more work by reading inputs.
611            let Some(event) = op_inputs.next().await else {
612                // Inputs exhausted, shutting down.
613                output_frontier.borrow_mut().clear();
614                return;
615            };
616            event.apply(&mut state);
617            // Also drain any other events that may be ready.
618            while let Some(Some(event)) = op_inputs.next().now_or_never() {
619                event.apply(&mut state);
620            }
621        }
622    });
623
624    button
625}
626
627/// Writes the given data to the shard, truncating it as necessary.
628///
629/// Returns the latest known upper for the shard.
630async fn truncating_compare_and_append(
631    write_handle: &mut WriteHandle<SourceData, (), Timestamp, StorageDiff>,
632    to_append: Vec<((&SourceData, &()), &Timestamp, StorageDiff)>,
633    new_upper: Antichain<Timestamp>,
634) -> Antichain<Timestamp> {
635    let mut expected_upper = write_handle.shared_upper();
636    loop {
637        if !PartialOrder::less_than(&expected_upper, &new_upper) {
638            debug!("ct_sink skipping {:?}", new_upper.elements());
639            return expected_upper;
640        }
641        let res = write_handle
642            .compare_and_append(&to_append, expected_upper.clone(), new_upper.clone())
643            .await
644            .expect("usage was valid");
645        debug!(
646            "ct_sink write res {:?}-{:?}: {:?}",
647            expected_upper.elements(),
648            new_upper.elements(),
649            res
650        );
651        match res {
652            Ok(()) => return new_upper,
653            Err(err) => {
654                expected_upper = err.current;
655                continue;
656            }
657        }
658    }
659}
660
661#[derive(Debug)]
662struct SinkState<D, T> {
663    /// The known times at which we're going to write data to the output. This
664    /// is guaranteed to include all times < append_times_progress, except that
665    /// ones < output_progress may have been truncated.
666    append_times: BTreeSet<T>,
667    append_times_progress: Antichain<T>,
668
669    /// The data we've collected to append to the output. This is often
670    /// compacted to advancing times and is expected to be ~empty in the steady
671    /// state.
672    to_append: ConsolidatingVec<(D, T)>,
673    to_append_progress: Antichain<T>,
674
675    /// A lower bound on the upper of the output.
676    output_progress: Antichain<T>,
677}
678
679impl<D: Ord> SinkState<D, Timestamp> {
680    fn new() -> Self {
681        SinkState {
682            append_times: BTreeSet::new(),
683            append_times_progress: Antichain::from_elem(Timestamp::minimum()),
684            to_append: ConsolidatingVec::new(128, 0),
685            to_append_progress: Antichain::from_elem(Timestamp::minimum()),
686            output_progress: Antichain::from_elem(Timestamp::minimum()),
687        }
688    }
689
690    /// Returns data to write to the output, if any, and the new upper to use.
691    fn process(
692        &mut self,
693    ) -> Option<(
694        Antichain<Timestamp>,
695        Vec<((&D, &()), &Timestamp, StorageDiff)>,
696    )> {
697        // We can only append at times >= the output_progress, so pop off
698        // anything unnecessary.
699        while let Some(x) = self.append_times.first() {
700            if self.output_progress.less_equal(x) {
701                break;
702            }
703            self.append_times.pop_first();
704        }
705
706        // Find the smallest append_time before append_time_progress. This is
707        // the next time we might need to write data at. Note that we can only
708        // act on append_times once the progress has passed them, because they
709        // could come out of order.
710        let write_ts = match self.append_times.first() {
711            Some(x) if !self.append_times_progress.less_equal(x) => x,
712            Some(_) | None => {
713                // The CT sink's contract is that it only writes data at times
714                // we received an input diff. There are none in
715                // `[output_progress, append_times_progress)`, so we can go
716                // ahead and advance the upper of the output, if it's not
717                // already.
718                //
719                // We could instead ensure liveness by basing this off of
720                // to_append, but for any CTs reading the output (expected to be
721                // a common case) we'd end up looping each timestamp through
722                // persist one-by-one.
723                if PartialOrder::less_than(&self.output_progress, &self.append_times_progress) {
724                    return Some((self.append_times_progress.clone(), Vec::new()));
725                }
726                // Otherwise, nothing to do!
727                return None;
728            }
729        };
730
731        if self.to_append_progress.less_equal(write_ts) {
732            // Don't have all the necessary data yet.
733            if self.output_progress.less_than(write_ts) {
734                // We can advance the output upper up to the write_ts. For
735                // self-referential CTs this might be necessary to ensure
736                // dataflow progress.
737                return Some((Antichain::from_elem(write_ts.clone()), Vec::new()));
738            }
739            return None;
740        }
741
742        // Time to write some data! Produce the collection as of write_ts by
743        // advancing timestamps, consolidating, and filtering out anything at
744        // future timestamps.
745        let as_of = std::slice::from_ref(write_ts);
746        for ((_, t), _) in self.to_append.iter_mut() {
747            t.advance_by(AntichainRef::new(as_of))
748        }
749        // TODO(ct2): Metrics for vec len and cap.
750        self.to_append.consolidate();
751
752        let append_data = self
753            .to_append
754            .iter()
755            .filter_map(|((k, t), d)| (t <= write_ts).then_some(((k, &()), t, d.into_inner())))
756            .collect();
757        Some((Antichain::from_elem(write_ts.step_forward()), append_data))
758    }
759}
760
761trait StepForward<G: Scope, D, R> {
762    /// Translates a collection one timestamp "forward" (i.e. `T` -> `T+1` as
763    /// defined by `TimestampManipulation::step_forward`).
764    ///
765    /// This includes:
766    /// - The differential timestamps in each data.
767    /// - The capabilities paired with that data.
768    /// - (As a consequence of the previous) the output frontier is one step forward
769    ///   of the input frontier.
770    ///
771    /// The caller is responsible for ensuring that all data and capabilities given
772    /// to this operator can be stepped forward without panicking, otherwise the
773    /// operator will panic at runtime.
774    fn step_forward(&self, name: &str) -> Collection<G, D, R>;
775}
776
777impl<G, D, R> StepForward<G, D, R> for Collection<G, D, R>
778where
779    G: Scope<Timestamp = Timestamp>,
780    D: Data,
781    R: Semigroup + 'static,
782{
783    fn step_forward(&self, name: &str) -> Collection<G, D, R> {
784        let name = format!("ct_step_forward({})", name);
785        let mut builder = OperatorBuilder::new(name, self.scope());
786        let (mut output, output_stream) = builder.new_output();
787        // We step forward (by one) each data timestamp and capability. As a
788        // result the output's frontier is guaranteed to be one past the input
789        // frontier, so make this promise to timely.
790        let step_forward_summary = Timestamp::from(1);
791        let mut input = builder.new_input_connection(
792            &self.inner,
793            Pipeline,
794            [(0, Antichain::from_elem(step_forward_summary))],
795        );
796        builder.set_notify(false);
797        builder.build(move |_caps| {
798            move |_frontiers| {
799                let mut output = output.activate();
800                while let Some((cap, data)) = input.next() {
801                    for (_, ts, _) in data.iter_mut() {
802                        *ts = ts.step_forward();
803                    }
804                    let cap = cap.delayed(&cap.time().step_forward());
805                    output.session(&cap).give_container(data);
806                }
807            }
808        });
809
810        output_stream.as_collection()
811    }
812}
813
814trait TimesExtract<G: Scope, D, R> {
815    /// Returns a collection with the times changed in the input collection.
816    ///
817    /// This works by mapping the data piece of the differential tuple to `()`.
818    /// It is essentially the same as the following, but without cloning
819    /// everything in the input.
820    ///
821    /// ```ignore
822    /// input.map(|(_data, ts, diff)| ((), ts, diff))
823    /// ```
824    ///
825    /// The output may be partially consolidated, but no consolidation
826    /// guarantees are made.
827    fn times_extract(&self, name: &str) -> (Collection<G, D, R>, Collection<G, (), R>);
828}
829
830impl<G, D, R> TimesExtract<G, D, R> for Collection<G, D, R>
831where
832    G: Scope<Timestamp = Timestamp>,
833    D: Clone + 'static,
834    R: Semigroup + 'static + std::fmt::Debug,
835{
836    fn times_extract(&self, name: &str) -> (Collection<G, D, R>, Collection<G, (), R>) {
837        let name = format!("ct_times_extract({})", name);
838        let mut builder = OperatorBuilder::new(name, self.scope());
839        let (mut passthrough, passthrough_stream) = builder.new_output();
840        let (mut times, times_stream) = builder.new_output::<ConsolidatingContainerBuilder<_>>();
841        let mut input = builder.new_input(&self.inner, Pipeline);
842        builder.set_notify(false);
843        builder.build(|_caps| {
844            move |_frontiers| {
845                let mut passthrough = passthrough.activate();
846                let mut times = times.activate();
847                while let Some((cap, data)) = input.next() {
848                    let times_iter = data.iter().map(|(_data, ts, diff)| ((), *ts, diff.clone()));
849                    times.session_with_builder(&cap).give_iterator(times_iter);
850                    passthrough.session(&cap).give_container(data);
851                }
852            }
853        });
854        (
855            passthrough_stream.as_collection(),
856            times_stream.as_collection(),
857        )
858    }
859}
860
861trait TimesReduce<G: Scope, R> {
862    /// This is essentially a specialized impl of consolidate, with a HashMap
863    /// instead of the Trace.
864    fn times_reduce(&self, name: &str) -> Collection<G, (), R>;
865}
866
867impl<G, R> TimesReduce<G, R> for Collection<G, (), R>
868where
869    G: Scope<Timestamp = Timestamp>,
870    R: Semigroup + 'static + std::fmt::Debug,
871{
872    fn times_reduce(&self, name: &str) -> Collection<G, (), R> {
873        let name = format!("ct_times_reduce({})", name);
874        self.inner
875            .unary_frontier(Pipeline, &name, |_caps, _info| {
876                let mut notificator = FrontierNotificator::default();
877                let mut stash = HashMap::<_, R>::new();
878                move |input, output| {
879                    while let Some((cap, data)) = input.next() {
880                        for ((), ts, diff) in data.drain(..) {
881                            notificator.notify_at(cap.delayed(&ts));
882                            if let Some(sum) = stash.get_mut(&ts) {
883                                sum.plus_equals(&diff);
884                            } else {
885                                stash.insert(ts, diff);
886                            }
887                        }
888                    }
889                    notificator.for_each(&[input.frontier()], |cap, _not| {
890                        if let Some(diff) = stash.remove(cap.time()) {
891                            output.session(&cap).give(((), cap.time().clone(), diff));
892                        }
893                    });
894                }
895            })
896            .as_collection()
897    }
898}
899
900#[cfg(test)]
901mod tests {
902    use differential_dataflow::AsCollection;
903    use mz_repr::Timestamp;
904    use timely::Config;
905    use timely::dataflow::ProbeHandle;
906    use timely::dataflow::operators::capture::Extract;
907    use timely::dataflow::operators::{Capture, Input, ToStream};
908    use timely::progress::Antichain;
909
910    use super::*;
911
912    #[mz_ore::test]
913    fn step_forward() {
914        timely::execute(Config::thread(), |worker| {
915            let (mut input, probe, output) = worker.dataflow(|scope| {
916                let (handle, input) = scope.new_input();
917                let probe = ProbeHandle::<Timestamp>::new();
918                let output = input
919                    .as_collection()
920                    .step_forward("test")
921                    .probe_with(&probe)
922                    .inner
923                    .capture();
924                (handle, probe, output)
925            });
926
927            let mut expected = Vec::new();
928            for i in 0u64..10 {
929                let in_ts = Timestamp::new(i);
930                let out_ts = in_ts.step_forward();
931                input.send((i, in_ts, 1));
932                input.advance_to(in_ts.step_forward());
933
934                // We should get the data out advanced by `step_forward` and
935                // also, crucially, the output frontier should do the same (i.e.
936                // this is why we can't simply use `Collection::delay`).
937                worker.step_while(|| probe.less_than(&out_ts.step_forward()));
938                expected.push((i, out_ts, 1));
939            }
940            // Closing the input should allow the output to advance and the
941            // dataflow to shut down.
942            input.close();
943            while worker.step() {}
944
945            let actual = output
946                .extract()
947                .into_iter()
948                .flat_map(|x| x.1)
949                .collect::<Vec<_>>();
950            assert_eq!(actual, expected);
951        })
952        .unwrap();
953    }
954
955    #[mz_ore::test]
956    fn times_extract() {
957        struct PanicOnClone;
958
959        impl Clone for PanicOnClone {
960            fn clone(&self) -> Self {
961                panic!("boom")
962            }
963        }
964
965        let output = timely::execute_directly(|worker| {
966            worker.dataflow(|scope| {
967                let input = [
968                    (PanicOnClone, Timestamp::new(0), 0),
969                    (PanicOnClone, Timestamp::new(1), 1),
970                    (PanicOnClone, Timestamp::new(1), 1),
971                    (PanicOnClone, Timestamp::new(2), -2),
972                    (PanicOnClone, Timestamp::new(2), 1),
973                ]
974                .to_stream(scope)
975                .as_collection();
976                let (_passthrough, times) = input.times_extract("test");
977                times.inner.capture()
978            })
979        });
980        let expected = vec![((), Timestamp::new(1), 2), ((), Timestamp::new(2), -1)];
981        let actual = output
982            .extract()
983            .into_iter()
984            .flat_map(|x| x.1)
985            .collect::<Vec<_>>();
986        assert_eq!(actual, expected);
987    }
988
989    #[mz_ore::test]
990    fn times_reduce() {
991        let output = timely::execute_directly(|worker| {
992            worker.dataflow(|scope| {
993                let input = [
994                    ((), Timestamp::new(3), 1),
995                    ((), Timestamp::new(2), 1),
996                    ((), Timestamp::new(1), 1),
997                    ((), Timestamp::new(2), 1),
998                    ((), Timestamp::new(3), 1),
999                    ((), Timestamp::new(3), 1),
1000                ]
1001                .to_stream(scope)
1002                .as_collection();
1003                input.times_reduce("test").inner.capture()
1004            })
1005        });
1006        let expected = vec![
1007            ((), Timestamp::new(1), 1),
1008            ((), Timestamp::new(2), 2),
1009            ((), Timestamp::new(3), 3),
1010        ];
1011        let actual = output
1012            .extract()
1013            .into_iter()
1014            .flat_map(|x| x.1)
1015            .collect::<Vec<_>>();
1016        assert_eq!(actual, expected);
1017    }
1018
1019    #[mz_ore::test]
1020    fn ct_sink_state() {
1021        #[track_caller]
1022        fn assert_noop(state: &mut super::SinkState<&'static str, Timestamp>) {
1023            if let Some(ret) = state.process() {
1024                panic!("should be nothing to write: {:?}", ret);
1025            }
1026        }
1027
1028        #[track_caller]
1029        fn assert_write(
1030            state: &mut super::SinkState<&'static str, Timestamp>,
1031            expected_upper: u64,
1032            expected_append: &[&str],
1033        ) {
1034            let (new_upper, to_append) = state.process().expect("should be something to write");
1035            assert_eq!(
1036                new_upper,
1037                Antichain::from_elem(Timestamp::new(expected_upper))
1038            );
1039            let to_append = to_append
1040                .into_iter()
1041                .map(|((k, ()), _ts, _diff)| *k)
1042                .collect::<Vec<_>>();
1043            assert_eq!(to_append, expected_append);
1044        }
1045
1046        let mut s = super::SinkState::new();
1047
1048        // Nothing to do at the initial state.
1049        assert_noop(&mut s);
1050
1051        // Getting data to append is not enough to do anything yet.
1052        s.to_append.push((("a", 1.into()), Diff::ONE));
1053        s.to_append.push((("b", 1.into()), Diff::ONE));
1054        assert_noop(&mut s);
1055
1056        // Knowing that this is the only data we'll get for that timestamp is
1057        // still not enough.
1058        s.to_append_progress = Antichain::from_elem(2.into());
1059        assert_noop(&mut s);
1060
1061        // Even knowing that we got input at that time is not quite enough yet
1062        // (we could be getting these out of order).
1063        s.append_times.insert(1.into());
1064        assert_noop(&mut s);
1065
1066        // Indeed, it did come out of order. Also note that this checks the ==
1067        // case for time vs progress.
1068        s.append_times.insert(0.into());
1069        assert_noop(&mut s);
1070
1071        // Okay, now we know that we've seen all the times we got input up to 3.
1072        // This is enough to allow the empty write of `[0,1)`.
1073        s.append_times_progress = Antichain::from_elem(3.into());
1074        assert_write(&mut s, 1, &[]);
1075
1076        // That succeeded, now we can write the data at 1.
1077        s.output_progress = Antichain::from_elem(1.into());
1078        assert_write(&mut s, 2, &["a", "b"]);
1079
1080        // That succeeded, now we know about some empty time.
1081        s.output_progress = Antichain::from_elem(2.into());
1082        assert_write(&mut s, 3, &[]);
1083
1084        // That succeeded, now nothing to do.
1085        s.output_progress = Antichain::from_elem(3.into());
1086        assert_noop(&mut s);
1087
1088        // Find out about a new time to write at. Even without the data, we can
1089        // do an empty write up to that time.
1090        s.append_times.insert(5.into());
1091        s.append_times_progress = Antichain::from_elem(6.into());
1092        assert_write(&mut s, 5, &[]);
1093
1094        // That succeeded, now nothing to do again.
1095        s.output_progress = Antichain::from_elem(5.into());
1096
1097        // Retract one of the things currently in the collection and add a new
1098        // thing, to verify the consolidate.
1099        s.to_append.push((("a", 5.into()), Diff::MINUS_ONE));
1100        s.to_append.push((("c", 5.into()), Diff::ONE));
1101        s.to_append_progress = Antichain::from_elem(6.into());
1102        assert_write(&mut s, 6, &["b", "c"]);
1103    }
1104}