1use 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
49pub 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 let mut needed_tokens = Vec::new();
86
87 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 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 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 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
162fn 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 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 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 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 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 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 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 needed_tokens.push(upsert_token);
381
382 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 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 let health = decode_health.into_iter().chain(envelope_health).collect();
442 (envelope_ok, needed_tokens, health)
443}
444
445fn 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 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 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
472fn 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#[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
515fn 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 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 let key_row = match upsert_envelope.style {
547 UpsertStyle::Debezium { .. }
549 | UpsertStyle::Default(KeyEnvelope::Flattened)
550 | UpsertStyle::ValueErrInline {
551 key_envelope: KeyEnvelope::Flattened,
552 error_column: _,
553 } => key,
554 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 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 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 packer.push_list(iter::once(Datum::String(&err_string)));
620 count += 1;
621 let metadata_len = metadata.as_row_ref().iter().count();
622 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
643fn 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 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
693fn 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 (_, Some(Err(e))) => Some(Err(e.into())),
702 (Some(Err(e)), _) => Some(Err(e.into())),
703 (None, None) => None,
704 _ => 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 assert_eq!(backpressure_inflight_bytes_limit, Some(1000));
760 }
761}