mz_storage/source/
generator.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
10use std::collections::BTreeMap;
11use std::convert::Infallible;
12use std::ops::Rem;
13use std::sync::Arc;
14use std::time::Duration;
15
16use differential_dataflow::AsCollection;
17use futures::StreamExt;
18use itertools::Itertools;
19use mz_ore::cast::CastFrom;
20use mz_ore::iter::IteratorExt;
21use mz_repr::{Diff, GlobalId, Row};
22use mz_storage_types::errors::DataflowError;
23use mz_storage_types::sources::load_generator::{
24    Event, Generator, KeyValueLoadGenerator, LoadGenerator, LoadGeneratorOutput,
25    LoadGeneratorSourceConnection,
26};
27use mz_storage_types::sources::{MzOffset, SourceExportDetails, SourceTimestamp};
28use mz_timely_util::builder_async::{OperatorBuilder as AsyncOperatorBuilder, PressOnDropButton};
29use mz_timely_util::containers::stack::AccountedStackBuilder;
30use timely::container::CapacityContainerBuilder;
31use timely::dataflow::operators::core::Partition;
32use timely::dataflow::{Scope, Stream};
33use timely::progress::Antichain;
34use tokio::time::{Instant, interval_at};
35
36use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace};
37use crate::source::types::{Probe, SignaledFuture, SourceRender, StackedCollection};
38use crate::source::{RawSourceCreationConfig, SourceMessage};
39
40mod auction;
41mod clock;
42mod counter;
43mod datums;
44mod key_value;
45mod marketing;
46mod tpch;
47
48pub use auction::Auction;
49pub use clock::Clock;
50pub use counter::Counter;
51pub use datums::Datums;
52pub use tpch::Tpch;
53
54use self::marketing::Marketing;
55
56enum GeneratorKind {
57    Simple {
58        generator: Box<dyn Generator>,
59        tick_micros: Option<u64>,
60        as_of: u64,
61        up_to: u64,
62    },
63    KeyValue(KeyValueLoadGenerator),
64}
65
66impl GeneratorKind {
67    fn new(g: &LoadGenerator, tick_micros: Option<u64>, as_of: u64, up_to: u64) -> Self {
68        match g {
69            LoadGenerator::Auction => GeneratorKind::Simple {
70                generator: Box::new(Auction {}),
71                tick_micros,
72                as_of,
73                up_to,
74            },
75            LoadGenerator::Clock => GeneratorKind::Simple {
76                generator: Box::new(Clock {
77                    tick_ms: tick_micros
78                        .map(Duration::from_micros)
79                        .unwrap_or(Duration::from_secs(1))
80                        .as_millis()
81                        .try_into()
82                        .expect("reasonable tick interval"),
83                    as_of_ms: as_of,
84                }),
85                tick_micros,
86                as_of,
87                up_to,
88            },
89            LoadGenerator::Counter { max_cardinality } => GeneratorKind::Simple {
90                generator: Box::new(Counter {
91                    max_cardinality: max_cardinality.clone(),
92                }),
93                tick_micros,
94                as_of,
95                up_to,
96            },
97            LoadGenerator::Datums => GeneratorKind::Simple {
98                generator: Box::new(Datums {}),
99                tick_micros,
100                as_of,
101                up_to,
102            },
103            LoadGenerator::Marketing => GeneratorKind::Simple {
104                generator: Box::new(Marketing {}),
105                tick_micros,
106                as_of,
107                up_to,
108            },
109            LoadGenerator::Tpch {
110                count_supplier,
111                count_part,
112                count_customer,
113                count_orders,
114                count_clerk,
115            } => GeneratorKind::Simple {
116                generator: Box::new(Tpch {
117                    count_supplier: *count_supplier,
118                    count_part: *count_part,
119                    count_customer: *count_customer,
120                    count_orders: *count_orders,
121                    count_clerk: *count_clerk,
122                    // The default tick behavior 1s. For tpch we want to disable ticking
123                    // completely.
124                    tick: Duration::from_micros(tick_micros.unwrap_or(0)),
125                }),
126                tick_micros,
127                as_of,
128                up_to,
129            },
130            LoadGenerator::KeyValue(kv) => GeneratorKind::KeyValue(kv.clone()),
131        }
132    }
133
134    fn render<G: Scope<Timestamp = MzOffset>>(
135        self,
136        scope: &mut G,
137        config: &RawSourceCreationConfig,
138        committed_uppers: impl futures::Stream<Item = Antichain<MzOffset>> + 'static,
139        start_signal: impl std::future::Future<Output = ()> + 'static,
140    ) -> (
141        BTreeMap<GlobalId, StackedCollection<G, Result<SourceMessage, DataflowError>>>,
142        Stream<G, Infallible>,
143        Stream<G, HealthStatusMessage>,
144        Vec<PressOnDropButton>,
145    ) {
146        // figure out which output types from the generator belong to which output indexes
147        let mut output_map = BTreeMap::new();
148        // Maps the output index to export_id for statistics.
149        let mut idx_to_exportid = BTreeMap::new();
150        // Make sure that there's an entry for the default output, even if there are no exports
151        // that need data output. Certain implementations rely on it (at the time of this comment
152        // that includes the key-value load gen source).
153        output_map.insert(LoadGeneratorOutput::Default, Vec::new());
154        for (idx, (export_id, export)) in config.source_exports.iter().enumerate() {
155            let output_type = match &export.details {
156                SourceExportDetails::LoadGenerator(details) => details.output,
157                // This is an export that doesn't need any data output to it.
158                SourceExportDetails::None => continue,
159                _ => panic!("unexpected source export details: {:?}", export.details),
160            };
161            output_map
162                .entry(output_type)
163                .or_insert_with(Vec::new)
164                .push(idx);
165            idx_to_exportid.insert(idx, export_id.clone());
166        }
167
168        match self {
169            GeneratorKind::Simple {
170                tick_micros,
171                as_of,
172                up_to,
173                generator,
174            } => render_simple_generator(
175                generator,
176                tick_micros,
177                as_of.into(),
178                up_to.into(),
179                scope,
180                config,
181                committed_uppers,
182                output_map,
183            ),
184            GeneratorKind::KeyValue(kv) => key_value::render(
185                kv,
186                scope,
187                config.clone(),
188                committed_uppers,
189                start_signal,
190                output_map,
191                idx_to_exportid,
192            ),
193        }
194    }
195}
196
197impl SourceRender for LoadGeneratorSourceConnection {
198    type Time = MzOffset;
199
200    const STATUS_NAMESPACE: StatusNamespace = StatusNamespace::Generator;
201
202    fn render<G: Scope<Timestamp = MzOffset>>(
203        self,
204        scope: &mut G,
205        config: &RawSourceCreationConfig,
206        committed_uppers: impl futures::Stream<Item = Antichain<MzOffset>> + 'static,
207        start_signal: impl std::future::Future<Output = ()> + 'static,
208    ) -> (
209        BTreeMap<GlobalId, StackedCollection<G, Result<SourceMessage, DataflowError>>>,
210        Stream<G, Infallible>,
211        Stream<G, HealthStatusMessage>,
212        Option<Stream<G, Probe<MzOffset>>>,
213        Vec<PressOnDropButton>,
214    ) {
215        let generator_kind = GeneratorKind::new(
216            &self.load_generator,
217            self.tick_micros,
218            self.as_of,
219            self.up_to,
220        );
221        let (updates, uppers, health, button) =
222            generator_kind.render(scope, config, committed_uppers, start_signal);
223
224        (updates, uppers, health, None, button)
225    }
226}
227
228fn render_simple_generator<G: Scope<Timestamp = MzOffset>>(
229    generator: Box<dyn Generator>,
230    tick_micros: Option<u64>,
231    as_of: MzOffset,
232    up_to: MzOffset,
233    scope: &G,
234    config: &RawSourceCreationConfig,
235    committed_uppers: impl futures::Stream<Item = Antichain<MzOffset>> + 'static,
236    output_map: BTreeMap<LoadGeneratorOutput, Vec<usize>>,
237) -> (
238    BTreeMap<GlobalId, StackedCollection<G, Result<SourceMessage, DataflowError>>>,
239    Stream<G, Infallible>,
240    Stream<G, HealthStatusMessage>,
241    Vec<PressOnDropButton>,
242) {
243    let mut builder = AsyncOperatorBuilder::new(config.name.clone(), scope.clone());
244
245    let (data_output, stream) = builder.new_output::<AccountedStackBuilder<_>>();
246    let partition_count = u64::cast_from(config.source_exports.len());
247    let data_streams: Vec<_> = stream.partition::<CapacityContainerBuilder<_>, _, _>(
248        partition_count,
249        |((output, data), time, diff): &(
250            (usize, Result<SourceMessage, DataflowError>),
251            MzOffset,
252            Diff,
253        )| {
254            let output = u64::cast_from(*output);
255            (output, (data.clone(), time.clone(), diff.clone()))
256        },
257    );
258    let mut data_collections = BTreeMap::new();
259    for (id, data_stream) in config.source_exports.keys().zip_eq(data_streams) {
260        data_collections.insert(*id, data_stream.as_collection());
261    }
262
263    let (_progress_output, progress_stream) = builder.new_output::<CapacityContainerBuilder<_>>();
264    let (health_output, health_stream) = builder.new_output();
265
266    let busy_signal = Arc::clone(&config.busy_signal);
267    let source_resume_uppers = config.source_resume_uppers.clone();
268    let is_active_worker = config.responsible_for(());
269    let source_statistics = config.statistics.clone();
270    let button = builder.build(move |caps| {
271        SignaledFuture::new(busy_signal, async move {
272            let [mut cap, mut progress_cap, health_cap] = caps.try_into().unwrap();
273
274            // We only need this until we reported ourselves as Running.
275            let mut health_cap = Some(health_cap);
276
277            if !is_active_worker {
278                // Emit 0, to mark this worker as having started up correctly.
279                for stats in source_statistics.values() {
280                    stats.set_offset_known(0);
281                    stats.set_offset_committed(0);
282                }
283                return;
284            }
285
286            let resume_upper = Antichain::from_iter(
287                source_resume_uppers
288                    .values()
289                    .flat_map(|f| f.iter().map(MzOffset::decode_row)),
290            );
291
292            let Some(resume_offset) = resume_upper.into_option() else {
293                return;
294            };
295
296            let now_fn = mz_ore::now::SYSTEM_TIME.clone();
297
298            let start_instant = {
299                // We want to have our interval start at a nice round number...
300                // for example, if our tick interval is one minute, to start at a minute boundary.
301                // However, the `Interval` type from tokio can't be "floored" in that way.
302                // Instead, figure out the amount we should step forward based on the wall clock,
303                // then apply that to our monotonic clock to make things start at approximately the
304                // right time.
305                let now_millis = now_fn();
306                let now_instant = Instant::now();
307                let delay_millis = tick_micros
308                    .map(|tick_micros| tick_micros / 1000)
309                    .filter(|tick_millis| *tick_millis > 0)
310                    .map(|tick_millis| tick_millis - now_millis.rem(tick_millis))
311                    .unwrap_or(0);
312                now_instant + Duration::from_millis(delay_millis)
313            };
314            let tick = Duration::from_micros(tick_micros.unwrap_or(1_000_000));
315            let mut tick_interval = interval_at(start_instant, tick);
316
317            let mut rows = generator.by_seed(now_fn, None, resume_offset);
318
319            let mut committed_uppers = std::pin::pin!(committed_uppers);
320
321            // If we are just starting up, report 0 as our `offset_committed`.
322            let mut offset_committed = if resume_offset.offset == 0 {
323                Some(0)
324            } else {
325                None
326            };
327
328            while let Some((output_type, event)) = rows.next() {
329                match event {
330                    Event::Message(mut offset, (value, diff)) => {
331                        // Fast forward any data before the requested as of.
332                        if offset <= as_of {
333                            offset = as_of;
334                        }
335
336                        // If the load generator produces data at or beyond the
337                        // requested `up_to`, drop it. We'll terminate the load
338                        // generator when the capability advances to the `up_to`,
339                        // but the load generator might produce data far in advance
340                        // of its capability.
341                        if offset >= up_to {
342                            continue;
343                        }
344
345                        // Once we see the load generator start producing data for some offset,
346                        // we report progress beyond that offset, to ensure that a binding can be
347                        // minted for the data and it doesn't accumulate in reclocking.
348                        let _ = progress_cap.try_downgrade(&(offset + 1));
349
350                        let outputs = match output_map.get(&output_type) {
351                            Some(outputs) => outputs,
352                            // We don't have an output index for this output type, so drop it
353                            None => continue,
354                        };
355
356                        let message: Result<SourceMessage, DataflowError> = Ok(SourceMessage {
357                            key: Row::default(),
358                            value,
359                            metadata: Row::default(),
360                        });
361
362                        // Some generators always reproduce their TVC from the beginning which can
363                        // generate a significant amount of data that will overwhelm the dataflow.
364                        // Since those are not required downstream we eagerly ignore them here.
365                        if resume_offset <= offset {
366                            for (&output, message) in outputs.iter().repeat_clone(message) {
367                                data_output
368                                    .give_fueled(&cap, ((output, message), offset, diff))
369                                    .await;
370                            }
371                        }
372                    }
373                    Event::Progress(Some(offset)) => {
374                        if resume_offset <= offset && health_cap.is_some() {
375                            let health_cap = health_cap.take().expect("known to exist");
376                            health_output.give(
377                                &health_cap,
378                                HealthStatusMessage {
379                                    id: None,
380                                    namespace: StatusNamespace::Generator,
381                                    update: HealthStatusUpdate::running(),
382                                },
383                            );
384                        }
385
386                        // If we've reached the requested maximum offset, cease.
387                        if offset >= up_to {
388                            break;
389                        }
390
391                        // If the offset is at or below the requested `as_of`, don't
392                        // downgrade the capability.
393                        if offset <= as_of {
394                            continue;
395                        }
396
397                        cap.downgrade(&offset);
398                        let _ = progress_cap.try_downgrade(&offset);
399
400                        // We only sleep if we have surpassed the resume offset so that we can
401                        // quickly go over any historical updates that a generator might choose to
402                        // emit.
403                        // TODO(petrosagg): Remove the sleep below and make generators return an
404                        // async stream so that they can drive the rate of production directly
405                        if resume_offset < offset {
406                            loop {
407                                tokio::select! {
408                                    _tick = tick_interval.tick() => {
409                                        break;
410                                    }
411                                    Some(frontier) = committed_uppers.next() => {
412                                        if let Some(offset) = frontier.as_option() {
413                                            // Offset N means we have committed N offsets (offsets are
414                                            // 0-indexed)
415                                            offset_committed = Some(offset.offset);
416                                        }
417                                    }
418                                }
419                            }
420
421                            // TODO(guswynn): generators have various definitions of "snapshot", so
422                            // we are not going to implement snapshot progress statistics for them
423                            // right now, but will come back to it.
424                            if let Some(offset_committed) = offset_committed {
425                                for stats in source_statistics.values() {
426                                    stats.set_offset_committed(offset_committed);
427                                    // technically we could have _known_ a larger offset
428                                    // than the one that has been committed, but we can
429                                    // never recover that known amount on restart, so we
430                                    // just advance these in lock step.
431                                    stats.set_offset_known(offset_committed);
432                                }
433                            }
434                        }
435                    }
436                    Event::Progress(None) => return,
437                }
438            }
439        })
440    });
441
442    (
443        data_collections,
444        progress_stream,
445        health_stream,
446        vec![button.press_on_drop()],
447    )
448}