Skip to main content

mz_storage/render/
sources.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//! Logic related to the creation of dataflow sources.
11//!
12//! See [`render_source`] for more details.
13
14use std::collections::BTreeMap;
15use std::iter;
16use std::sync::Arc;
17
18use differential_dataflow::{AsCollection, VecCollection};
19use mz_ore::cast::CastLossy;
20use mz_persist_client::operators::shard_source::SnapshotMode;
21use mz_repr::{Datum, Diff, GlobalId, Row, RowPacker};
22use mz_storage_operators::persist_source;
23use mz_storage_operators::persist_source::Subtime;
24use mz_storage_types::controller::CollectionMetadata;
25use mz_storage_types::dyncfgs;
26use mz_storage_types::errors::{
27    DataflowError, DecodeError, EnvelopeError, UpsertError, UpsertNullKeyError, UpsertValueError,
28};
29use mz_storage_types::parameters::StorageMaxInflightBytesConfig;
30use mz_storage_types::sources::envelope::{KeyEnvelope, NoneEnvelope, UpsertEnvelope, UpsertStyle};
31use mz_storage_types::sources::*;
32use mz_timely_util::builder_async::PressOnDropButton;
33use mz_timely_util::operator::CollectionExt;
34use mz_timely_util::order::refine_antichain;
35use serde::{Deserialize, Serialize};
36use timely::container::CapacityContainerBuilder;
37use timely::dataflow::StreamVec;
38use timely::dataflow::operators::vec::Map;
39use timely::dataflow::operators::{ConnectLoop, Feedback, Leave, OkErr};
40use timely::dataflow::scope::Scope;
41use timely::progress::{Antichain, Timestamp};
42
43use crate::decode::{render_decode_cdcv2, render_decode_delimited};
44use crate::healthcheck::{HealthStatusMessage, StatusNamespace};
45use crate::source::types::{DecodeResult, SourceOutput, SourceRender};
46use crate::source::{self, RawSourceCreationConfig, SourceExportCreationConfig};
47use crate::upsert::{UpsertKey, UpsertSourceTime, UpsertValue};
48
49/// _Renders_ complete _differential_ collections
50/// that represent the final source and its errors
51/// as requested by the original `CREATE SOURCE` statement,
52/// encapsulated in the passed `SourceInstanceDesc`.
53///
54/// The first element in the returned tuple is the pair of Collections,
55/// the second is a type-erased token that will keep the source
56/// alive as long as it is not dropped.
57///
58/// This function is intended to implement the recipe described here:
59/// <https://github.com/MaterializeInc/materialize/blob/main/doc/developer/platform/architecture-storage.md#source-ingestion>
60pub fn render_source<'scope, 'root, C>(
61    scope: Scope<'scope, mz_repr::Timestamp>,
62    root_scope: Scope<'root, ()>,
63    dataflow_debug_name: &String,
64    connection: C,
65    description: IngestionDescription<CollectionMetadata>,
66    resume_stream: StreamVec<'scope, mz_repr::Timestamp, ()>,
67    storage_state: &crate::storage_state::StorageState,
68    base_source_config: RawSourceCreationConfig,
69) -> (
70    BTreeMap<
71        GlobalId,
72        (
73            VecCollection<'scope, mz_repr::Timestamp, Row, Diff>,
74            VecCollection<'scope, mz_repr::Timestamp, DataflowError, Diff>,
75        ),
76    >,
77    Vec<StreamVec<'root, (), HealthStatusMessage>>,
78    Vec<PressOnDropButton>,
79)
80where
81    C: SourceConnection + SourceRender + 'static,
82    C::Time: UpsertSourceTime,
83{
84    // Tokens that we should return from the method.
85    let mut needed_tokens = Vec::new();
86
87    // Note that this `render_source` attaches a single _instance_ of a source
88    // to the passed `Scope`, and this instance may be disabled if the
89    // source type does not support multiple instances. `render_source`
90    // is called on each timely worker as part of
91    // [`super::build_storage_dataflow`].
92
93    // A set of channels (1 per worker) used to signal rehydration being finished
94    // to raw sources. These are channels and not timely streams because they
95    // have to cross a scope boundary.
96    //
97    // Note that these will be entirely subsumed by full `hydration` backpressure,
98    // once that is implemented.
99    let (starter, mut start_signal) = tokio::sync::mpsc::channel::<()>(1);
100    let start_signal = async move {
101        let _ = start_signal.recv().await;
102    };
103
104    // Build the _raw_ ok and error sources using `create_raw_source` and the
105    // correct `SourceReader` implementations
106    let (exports, health, source_tokens) = source::create_raw_source(
107        scope,
108        root_scope,
109        storage_state,
110        resume_stream,
111        &base_source_config,
112        connection,
113        start_signal,
114    );
115
116    needed_tokens.extend(source_tokens);
117
118    let mut health_streams = Vec::with_capacity(exports.len() + 1);
119    health_streams.push(health);
120
121    let mut outputs = BTreeMap::new();
122    for (export_id, export) in exports {
123        type CB<C> = CapacityContainerBuilder<C>;
124        let (ok_stream, err_stream) =
125            export.map_fallible::<CB<_>, CB<_>, _, _, _>("export-demux-ok-err", |r| r);
126
127        // All sources should push their various error streams into this vector,
128        // whose contents will be concatenated and inserted along the collection.
129        // All subsources include the non-definite errors of the ingestion
130        let mut error_collections = Vec::new();
131
132        let data_config = base_source_config.source_exports[&export_id]
133            .data_config
134            .clone();
135        let (ok, extra_tokens, health_stream) = render_source_stream(
136            scope,
137            dataflow_debug_name,
138            export_id,
139            ok_stream,
140            data_config,
141            &description,
142            &mut error_collections,
143            storage_state,
144            &base_source_config,
145            starter.clone(),
146        );
147        needed_tokens.extend(extra_tokens);
148
149        // Flatten the error collections.
150        let err_collection = match error_collections.len() {
151            0 => err_stream,
152            _ => err_stream.concatenate(error_collections),
153        };
154
155        outputs.insert(export_id, (ok, err_collection));
156
157        health_streams.extend(health_stream.into_iter().map(|s| s.leave(root_scope)));
158    }
159    (outputs, health_streams, needed_tokens)
160}
161
162/// Completes the rendering of a particular source stream by applying decoding and envelope
163/// processing as necessary
164fn render_source_stream<'scope, FromTime>(
165    scope: Scope<'scope, mz_repr::Timestamp>,
166    dataflow_debug_name: &String,
167    export_id: GlobalId,
168    ok_source: VecCollection<'scope, mz_repr::Timestamp, SourceOutput<FromTime>, Diff>,
169    data_config: SourceExportDataConfig,
170    description: &IngestionDescription<CollectionMetadata>,
171    error_collections: &mut Vec<VecCollection<'scope, mz_repr::Timestamp, DataflowError, Diff>>,
172    storage_state: &crate::storage_state::StorageState,
173    base_source_config: &RawSourceCreationConfig,
174    rehydrated_token: impl std::any::Any + 'static,
175) -> (
176    VecCollection<'scope, mz_repr::Timestamp, Row, Diff>,
177    Vec<PressOnDropButton>,
178    Vec<StreamVec<'scope, mz_repr::Timestamp, HealthStatusMessage>>,
179)
180where
181    FromTime: Timestamp + Sync,
182    FromTime: UpsertSourceTime,
183{
184    let mut needed_tokens = vec![];
185
186    // Use the envelope and encoding configs for this particular source export
187    let SourceExportDataConfig { encoding, envelope } = data_config;
188
189    let SourceDesc {
190        connection: _,
191        timestamp_interval: _,
192    } = description.desc;
193
194    let (decoded_stream, decode_health) = match encoding {
195        None => (
196            ok_source.map(|r| DecodeResult {
197                // This is safe because the current set of sources produce
198                // either:
199                // 1. Non-nullable keys
200                // 2. No keys at all.
201                //
202                // Please see the comment on `key_envelope_no_encoding` in
203                // `mz_sql::plan::statement::ddl` for more details.
204                key: Some(Ok(r.key)),
205                value: Some(Ok(r.value)),
206                metadata: r.metadata,
207                from_time: r.from_time,
208            }),
209            None,
210        ),
211        Some(encoding) => {
212            let (decoded_stream, decode_health) = render_decode_delimited(
213                ok_source,
214                encoding.key,
215                encoding.value,
216                dataflow_debug_name.clone(),
217                storage_state.metrics.decode_defs.clone(),
218                storage_state.storage_configuration.clone(),
219            );
220            (decoded_stream, Some(decode_health))
221        }
222    };
223
224    // render envelopes
225    let (envelope_ok, envelope_health) = match &envelope {
226        SourceEnvelope::Upsert(upsert_envelope) => {
227            let upsert_input = upsert_commands(decoded_stream, upsert_envelope.clone());
228
229            let persist_clients = Arc::clone(&storage_state.persist_clients);
230            // TODO: Get this to work with the as_of.
231            let resume_upper = base_source_config.resume_uppers[&export_id].clone();
232
233            let upper_ts = resume_upper
234                .as_option()
235                .expect("resuming an already finished ingestion")
236                .clone();
237            let outer_mz_scope = scope.clone();
238            let (upsert, health_update) = scope.scoped(
239                &format!("upsert_rehydration_backpressure({})", export_id),
240                |scope| {
241                    let (previous, previous_token, feedback_handle, backpressure_metrics) = {
242                        let as_of = Antichain::from_elem(upper_ts.saturating_sub(1));
243
244                        let backpressure_max_inflight_bytes = get_backpressure_max_inflight_bytes(
245                            &storage_state
246                                .storage_configuration
247                                .parameters
248                                .storage_dataflow_max_inflight_bytes_config,
249                            &storage_state.instance_context.cluster_memory_limit,
250                        );
251
252                        let (feedback_handle, flow_control, backpressure_metrics) =
253                            if let Some(storage_dataflow_max_inflight_bytes) =
254                                backpressure_max_inflight_bytes
255                            {
256                                tracing::info!(
257                                    ?backpressure_max_inflight_bytes,
258                                    "timely-{} using backpressure in upsert for source {}",
259                                    base_source_config.worker_id,
260                                    export_id
261                                );
262                                if !storage_state
263                                    .storage_configuration
264                                    .parameters
265                                    .storage_dataflow_max_inflight_bytes_config
266                                    .disk_only
267                                    || storage_state.instance_context.scratch_directory.is_some()
268                                {
269                                    let (feedback_handle, feedback_data) =
270                                        scope.feedback(Default::default());
271
272                                    // TODO(guswynn): cleanup
273                                    let backpressure_metrics = Some(
274                                        base_source_config
275                                            .metrics
276                                            .get_backpressure_metrics(export_id, scope.index()),
277                                    );
278
279                                    (
280                                        Some(feedback_handle),
281                                        Some(persist_source::FlowControl {
282                                            progress_stream: feedback_data,
283                                            max_inflight_bytes: storage_dataflow_max_inflight_bytes,
284                                            summary: (Default::default(), Subtime::least_summary()),
285                                            metrics: backpressure_metrics
286                                                .as_ref()
287                                                .map(|m| m.operator_metrics()),
288                                        }),
289                                        backpressure_metrics,
290                                    )
291                                } else {
292                                    (None, None, None)
293                                }
294                            } else {
295                                (None, None, None)
296                            };
297
298                        let storage_metadata = description.source_exports[&export_id]
299                            .storage_metadata
300                            .clone();
301
302                        let error_handler =
303                            storage_state.error_handler("upsert_rehydration", export_id);
304
305                        let (stream, tok) = persist_source::persist_source_core(
306                            outer_mz_scope,
307                            scope,
308                            export_id,
309                            persist_clients,
310                            storage_metadata,
311                            None,
312                            Some(as_of),
313                            SnapshotMode::Include,
314                            Antichain::new(),
315                            None,
316                            flow_control,
317                            false.then_some(|| unreachable!()),
318                            async {},
319                            error_handler,
320                        );
321                        (
322                            stream.as_collection(),
323                            Some(tok),
324                            feedback_handle,
325                            backpressure_metrics,
326                        )
327                    };
328
329                    let export_statistics = storage_state
330                        .aggregated_statistics
331                        .get_source(&export_id)
332                        .expect("statistics initialized")
333                        .clone();
334                    let export_config = SourceExportCreationConfig {
335                        id: export_id,
336                        worker_id: base_source_config.worker_id,
337                        metrics: base_source_config.metrics.clone(),
338                        source_statistics: export_statistics,
339                    };
340                    let (upsert, health_update, snapshot_progress, upsert_token) =
341                        if dyncfgs::ENABLE_UPSERT_V2
342                            .get(storage_state.storage_configuration.config_set())
343                        {
344                            // Resolved here, at operator construction, so the
345                            // dataflow keeps one stash flavor for its whole
346                            // life even if the flag flips underneath it.
347                            let stash_flavor =
348                                crate::upsert_continual_feedback_v2::UpsertStashFlavor::from_config(
349                                    storage_state.storage_configuration.config_set(),
350                                );
351                            crate::upsert::upsert_v2(
352                                upsert_input.enter(scope),
353                                upsert_envelope.clone(),
354                                refine_antichain(&resume_upper),
355                                previous,
356                                previous_token,
357                                export_config,
358                                backpressure_metrics,
359                                stash_flavor,
360                            )
361                        } else {
362                            crate::upsert::upsert(
363                                upsert_input.enter(scope),
364                                upsert_envelope.clone(),
365                                refine_antichain(&resume_upper),
366                                previous,
367                                previous_token,
368                                export_config,
369                                &storage_state.instance_context,
370                                &storage_state.storage_configuration,
371                                &storage_state.dataflow_parameters,
372                                backpressure_metrics,
373                            )
374                        };
375
376                    // Even though we register the `persist_sink` token at a top-level,
377                    // which will stop any data from being committed, we also register
378                    // a token for the `upsert` operator which may be in the middle of
379                    // rehydration processing the `persist_source` input above.
380                    needed_tokens.push(upsert_token);
381
382                    // If configured, delay raw sources until we rehydrate the upsert
383                    // source. Otherwise, drop the token, unblocking the sources at the
384                    // end rendering.
385                    if dyncfgs::DELAY_SOURCES_PAST_REHYDRATION
386                        .get(storage_state.storage_configuration.config_set())
387                    {
388                        crate::upsert::rehydration_finished(
389                            scope.clone(),
390                            base_source_config,
391                            rehydrated_token,
392                            refine_antichain(&resume_upper),
393                            snapshot_progress.clone(),
394                        );
395                    } else {
396                        drop(rehydrated_token)
397                    };
398
399                    // If backpressure from persist is enabled, we connect the upsert operator's
400                    // snapshot progress to the persist source feedback handle.
401                    if let Some(feedback_handle) = feedback_handle {
402                        snapshot_progress.connect_loop(feedback_handle);
403                    }
404
405                    (
406                        upsert.leave(outer_mz_scope),
407                        health_update
408                            .map(|(id, update)| HealthStatusMessage {
409                                id,
410                                namespace: StatusNamespace::Upsert,
411                                update,
412                            })
413                            .leave(outer_mz_scope),
414                    )
415                },
416            );
417
418            let (upsert_ok, upsert_err) = upsert.inner.ok_err(split_ok_err);
419            error_collections.push(upsert_err.as_collection());
420
421            (upsert_ok.as_collection(), Some(health_update))
422        }
423        SourceEnvelope::None(none_envelope) => {
424            let results = append_metadata_to_value(decoded_stream);
425
426            let flattened_stream = flatten_results_prepend_keys(none_envelope, results);
427
428            let (stream, errors) = flattened_stream.inner.ok_err(split_ok_err);
429
430            error_collections.push(errors.as_collection());
431            (stream.as_collection(), None)
432        }
433        SourceEnvelope::CdcV2 => {
434            let (oks, token) = render_decode_cdcv2(&decoded_stream);
435            needed_tokens.push(token);
436            (oks, None)
437        }
438    };
439
440    // Return the collections and any needed tokens.
441    let health = decode_health.into_iter().chain(envelope_health).collect();
442    (envelope_ok, needed_tokens, health)
443}
444
445// Returns the maximum limit of inflight bytes for backpressure based on given config
446// and the current cluster size
447fn get_backpressure_max_inflight_bytes(
448    inflight_bytes_config: &StorageMaxInflightBytesConfig,
449    cluster_memory_limit: &Option<usize>,
450) -> Option<usize> {
451    let StorageMaxInflightBytesConfig {
452        max_inflight_bytes_default,
453        max_inflight_bytes_cluster_size_fraction,
454        disk_only: _,
455    } = inflight_bytes_config;
456
457    // Will use backpressure only if the default inflight value is provided
458    if max_inflight_bytes_default.is_some() {
459        let current_cluster_max_bytes_limit =
460            cluster_memory_limit.as_ref().and_then(|cluster_memory| {
461                max_inflight_bytes_cluster_size_fraction.map(|fraction| {
462                    // We just need close the correct % of bytes here, so we just use lossy casts.
463                    usize::cast_lossy(f64::cast_lossy(*cluster_memory) * fraction)
464                })
465            });
466        current_cluster_max_bytes_limit.or(*max_inflight_bytes_default)
467    } else {
468        None
469    }
470}
471
472// TODO: Maybe we should finally move this to some central place and re-use. There seem to be
473// enough instances of this by now.
474fn split_ok_err<O, E, T, D>(x: (Result<O, E>, T, D)) -> Result<(O, T, D), (E, T, D)> {
475    match x {
476        (Ok(ok), ts, diff) => Ok((ok, ts, diff)),
477        (Err(err), ts, diff) => Err((err, ts, diff)),
478    }
479}
480
481/// After handling metadata insertion, we split streams into key/value parts for convenience
482#[derive(
483    Debug,
484    Clone,
485    Hash,
486    PartialEq,
487    Eq,
488    Ord,
489    PartialOrd,
490    Serialize,
491    Deserialize
492)]
493struct KV {
494    key: Option<Result<Row, DecodeError>>,
495    val: Option<Result<Row, DecodeError>>,
496}
497
498fn append_metadata_to_value<'scope, T: Timestamp, FromTime: Timestamp>(
499    results: VecCollection<'scope, T, DecodeResult<FromTime>, Diff>,
500) -> VecCollection<'scope, T, KV, Diff> {
501    results.map(move |res| {
502        let val = res.value.map(|val_result| {
503            val_result.map(|mut val| {
504                if !res.metadata.is_empty() {
505                    RowPacker::for_existing_row(&mut val).extend_by_row(&res.metadata);
506                }
507                val
508            })
509        });
510
511        KV { val, key: res.key }
512    })
513}
514
515/// Convert from streams of [`DecodeResult`] to UpsertCommands, inserting the Key according to [`KeyEnvelope`]
516fn upsert_commands<'scope, T: Timestamp, FromTime: Timestamp>(
517    input: VecCollection<'scope, T, DecodeResult<FromTime>, Diff>,
518    upsert_envelope: UpsertEnvelope,
519) -> VecCollection<'scope, T, (UpsertKey, Option<UpsertValue>, FromTime), Diff> {
520    let mut row_buf = Row::default();
521    input.map(move |result| {
522        let from_time = result.from_time;
523
524        let key = match result.key {
525            Some(Ok(key)) => Ok(key),
526            None => Err(UpsertError::NullKey(UpsertNullKeyError)),
527            Some(Err(err)) => Err(UpsertError::KeyDecode(err)),
528        };
529
530        // If we have a well-formed key we can continue, otherwise we're upserting an error
531        let key = match key {
532            Ok(key) => key,
533            Err(err) => match result.value {
534                Some(_) => {
535                    return (
536                        UpsertKey::from_key(Err(&err)),
537                        Some(Err(Box::new(err))),
538                        from_time,
539                    );
540                }
541                None => return (UpsertKey::from_key(Err(&err)), None, from_time),
542            },
543        };
544
545        // We can now apply the key envelope
546        let key_row = match upsert_envelope.style {
547            // flattened or debezium
548            UpsertStyle::Debezium { .. }
549            | UpsertStyle::Default(KeyEnvelope::Flattened)
550            | UpsertStyle::ValueErrInline {
551                key_envelope: KeyEnvelope::Flattened,
552                error_column: _,
553            } => key,
554            // named
555            UpsertStyle::Default(KeyEnvelope::Named(_))
556            | UpsertStyle::ValueErrInline {
557                key_envelope: KeyEnvelope::Named(_),
558                error_column: _,
559            } => {
560                if key.iter().nth(1).is_none() {
561                    key
562                } else {
563                    row_buf.packer().push_list(key.iter());
564                    row_buf.clone()
565                }
566            }
567            UpsertStyle::Default(KeyEnvelope::None)
568            | UpsertStyle::ValueErrInline {
569                key_envelope: KeyEnvelope::None,
570                error_column: _,
571            } => unreachable!(),
572        };
573
574        let key = UpsertKey::from_key(Ok(&key_row));
575
576        let metadata = result.metadata;
577
578        let value = match result.value {
579            Some(Ok(ref row)) => match upsert_envelope.style {
580                UpsertStyle::Debezium { after_idx } => match row.iter().nth(after_idx).unwrap() {
581                    Datum::List(after) => {
582                        let mut packer = row_buf.packer();
583                        packer.extend(after.iter());
584                        packer.extend_by_row(&metadata);
585                        Some(Ok(row_buf.clone()))
586                    }
587                    Datum::Null => None,
588                    d => panic!("type error: expected record, found {:?}", d),
589                },
590                UpsertStyle::Default(_) => {
591                    let mut packer = row_buf.packer();
592                    packer.extend_by_row(&key_row);
593                    packer.extend_by_row(row);
594                    packer.extend_by_row(&metadata);
595                    Some(Ok(row_buf.clone()))
596                }
597                UpsertStyle::ValueErrInline { .. } => {
598                    let mut packer = row_buf.packer();
599                    packer.extend_by_row(&key_row);
600                    // The 'error' column is null
601                    packer.push(Datum::Null);
602                    packer.extend_by_row(row);
603                    packer.extend_by_row(&metadata);
604                    Some(Ok(row_buf.clone()))
605                }
606            },
607            Some(Err(inner)) => {
608                match upsert_envelope.style {
609                    UpsertStyle::ValueErrInline { .. } => {
610                        let mut count = 0;
611                        // inline the error in the data output
612                        let err_string = inner.to_string();
613                        let mut packer = row_buf.packer();
614                        for datum in key_row.iter() {
615                            packer.push(datum);
616                            count += 1;
617                        }
618                        // The 'error' column is a record with a 'description' column
619                        packer.push_list(iter::once(Datum::String(&err_string)));
620                        count += 1;
621                        let metadata_len = metadata.as_row_ref().iter().count();
622                        // push nulls for all value columns
623                        packer.extend(
624                            iter::repeat(Datum::Null)
625                                .take(upsert_envelope.source_arity - count - metadata_len),
626                        );
627                        packer.extend_by_row(&metadata);
628                        Some(Ok(row_buf.clone()))
629                    }
630                    _ => Some(Err(Box::new(UpsertError::Value(UpsertValueError {
631                        for_key: key_row,
632                        inner,
633                    })))),
634                }
635            }
636            None => None,
637        };
638
639        (key, value, from_time)
640    })
641}
642
643/// Convert from streams of [`DecodeResult`] to Rows, inserting the Key according to [`KeyEnvelope`]
644fn flatten_results_prepend_keys<'scope, T: Timestamp>(
645    none_envelope: &NoneEnvelope,
646    results: VecCollection<'scope, T, KV, Diff>,
647) -> VecCollection<'scope, T, Result<Row, DataflowError>, Diff> {
648    let NoneEnvelope {
649        key_envelope,
650        key_arity,
651    } = none_envelope;
652
653    let null_key_columns = Row::pack_slice(&vec![Datum::Null; *key_arity]);
654
655    match key_envelope {
656        KeyEnvelope::None => {
657            results.flat_map(|KV { val, .. }| val.map(|result| result.map_err(Into::into)))
658        }
659        KeyEnvelope::Flattened => results
660            .flat_map(raise_key_value_errors)
661            .map(move |maybe_kv| {
662                maybe_kv.map(|(key, value)| {
663                    let mut key = key.unwrap_or_else(|| null_key_columns.clone());
664                    RowPacker::for_existing_row(&mut key).extend_by_row(&value);
665                    key
666                })
667            }),
668        KeyEnvelope::Named(_) => {
669            results
670                .flat_map(raise_key_value_errors)
671                .map(move |maybe_kv| {
672                    maybe_kv.map(|(key, value)| {
673                        let mut key = key.unwrap_or_else(|| null_key_columns.clone());
674                        // Named semantics rename a key that is a single column, and encode a
675                        // multi-column field as a struct with that name
676                        let row = if key.iter().nth(1).is_none() {
677                            RowPacker::for_existing_row(&mut key).extend_by_row(&value);
678                            key
679                        } else {
680                            let mut new_row = Row::default();
681                            let mut packer = new_row.packer();
682                            packer.push_list(key.iter());
683                            packer.extend_by_row(&value);
684                            new_row
685                        };
686                        row
687                    })
688                })
689        }
690    }
691}
692
693/// Handle possibly missing key or value portions of messages
694fn raise_key_value_errors(
695    KV { key, val }: KV,
696) -> Option<Result<(Option<Row>, Row), DataflowError>> {
697    match (key, val) {
698        (Some(Ok(key)), Some(Ok(value))) => Some(Ok((Some(key), value))),
699        (None, Some(Ok(value))) => Some(Ok((None, value))),
700        // always prioritize the value error if either or both have an error
701        (_, Some(Err(e))) => Some(Err(e.into())),
702        (Some(Err(e)), _) => Some(Err(e.into())),
703        (None, None) => None,
704        // TODO(petrosagg): these errors would be better grouped under an EnvelopeError enum
705        _ => Some(Err(DataflowError::from(EnvelopeError::Flat(
706            "Value not present for message".into(),
707        )))),
708    }
709}
710
711#[cfg(test)]
712mod test {
713    use super::*;
714
715    #[mz_ore::test]
716    fn test_no_default() {
717        let config = StorageMaxInflightBytesConfig {
718            max_inflight_bytes_default: None,
719            max_inflight_bytes_cluster_size_fraction: Some(0.5),
720            disk_only: false,
721        };
722        let memory_limit = Some(1000);
723
724        let backpressure_inflight_bytes_limit =
725            get_backpressure_max_inflight_bytes(&config, &memory_limit);
726
727        assert_eq!(backpressure_inflight_bytes_limit, None)
728    }
729
730    #[mz_ore::test]
731    fn test_no_matching_size() {
732        let config = StorageMaxInflightBytesConfig {
733            max_inflight_bytes_default: Some(10000),
734            max_inflight_bytes_cluster_size_fraction: Some(0.5),
735            disk_only: false,
736        };
737
738        let backpressure_inflight_bytes_limit = get_backpressure_max_inflight_bytes(&config, &None);
739
740        assert_eq!(
741            backpressure_inflight_bytes_limit,
742            config.max_inflight_bytes_default
743        )
744    }
745
746    #[mz_ore::test]
747    fn test_calculated_cluster_limit() {
748        let config = StorageMaxInflightBytesConfig {
749            max_inflight_bytes_default: Some(10000),
750            max_inflight_bytes_cluster_size_fraction: Some(0.5),
751            disk_only: false,
752        };
753        let memory_limit = Some(2000);
754
755        let backpressure_inflight_bytes_limit =
756            get_backpressure_max_inflight_bytes(&config, &memory_limit);
757
758        // the limit should be 50% of 2000 i.e. 1000
759        assert_eq!(backpressure_inflight_bytes_limit, Some(1000));
760    }
761}