1#![allow(missing_docs)]
24#![allow(clippy::needless_borrow)]
25
26use std::cell::RefCell;
27use std::collections::{BTreeMap, VecDeque};
28use std::hash::{Hash, Hasher};
29use std::rc::Rc;
30use std::sync::Arc;
31use std::time::Duration;
32
33use differential_dataflow::lattice::Lattice;
34use differential_dataflow::{AsCollection, Hashable, VecCollection};
35use futures::stream::StreamExt;
36use mz_ore::cast::CastFrom;
37use mz_ore::collections::CollectionExt;
38use mz_ore::now::NowFn;
39use mz_persist_client::cache::PersistClientCache;
40use mz_repr::{Diff, GlobalId, RelationDesc, Row};
41use mz_storage_types::configuration::StorageConfiguration;
42use mz_storage_types::controller::CollectionMetadata;
43use mz_storage_types::errors::DataflowError;
44use mz_storage_types::sources::{SourceConnection, SourceExport, SourceTimestamp};
45use mz_timely_util::antichain::AntichainExt;
46use mz_timely_util::builder_async::{OperatorBuilder as AsyncOperatorBuilder, PressOnDropButton};
47use mz_timely_util::capture::PusherCapture;
48use mz_timely_util::operator::ConcatenateFlatten;
49use mz_timely_util::reclock::reclock;
50use timely::PartialOrder;
51use timely::container::CapacityContainerBuilder;
52use timely::dataflow::channels::pact::Pipeline;
53use timely::dataflow::operators::capture::capture::Capture;
54use timely::dataflow::operators::core::Map as _;
55use timely::dataflow::operators::generic::OutputBuilder;
56use timely::dataflow::operators::generic::builder_rc::OperatorBuilder as OperatorBuilderRc;
57use timely::dataflow::operators::vec::Broadcast;
58use timely::dataflow::operators::{CapabilitySet, InspectCore, Leave};
59use timely::dataflow::{Scope, StreamVec};
60use timely::order::TotalOrder;
61use timely::progress::frontier::MutableAntichain;
62use timely::progress::{Antichain, Timestamp};
63use tokio::sync::{Semaphore, watch};
64use tokio_stream::wrappers::WatchStream;
65use tracing::trace;
66
67use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate};
68use crate::metrics::StorageMetrics;
69use crate::metrics::source::SourceMetrics;
70use crate::source::reclock::ReclockOperator;
71use crate::source::types::{Probe, SourceMessage, SourceOutput, SourceRender, StackedCollection};
72use crate::statistics::SourceStatistics;
73
74#[derive(Clone)]
77pub struct RawSourceCreationConfig {
78 pub name: String,
80 pub id: GlobalId,
82 pub source_exports: BTreeMap<GlobalId, SourceExport<CollectionMetadata>>,
84 pub worker_id: usize,
86 pub worker_count: usize,
88 pub timestamp_interval: Duration,
91 pub now_fn: NowFn,
93 pub metrics: StorageMetrics,
95 pub as_of: Antichain<mz_repr::Timestamp>,
97 pub resume_uppers: BTreeMap<GlobalId, Antichain<mz_repr::Timestamp>>,
100 pub source_resume_uppers: BTreeMap<GlobalId, Vec<Row>>,
107 pub persist_clients: Arc<PersistClientCache>,
109 pub statistics: BTreeMap<GlobalId, SourceStatistics>,
111 pub shared_remap_upper: Rc<RefCell<Antichain<mz_repr::Timestamp>>>,
113 pub config: StorageConfiguration,
115 pub remap_collection_id: GlobalId,
117 pub remap_metadata: CollectionMetadata,
119 pub busy_signal: Arc<Semaphore>,
122}
123
124#[derive(Clone)]
127pub struct SourceExportCreationConfig {
128 pub id: GlobalId,
130 pub worker_id: usize,
132 pub metrics: StorageMetrics,
134 pub source_statistics: SourceStatistics,
136}
137
138impl RawSourceCreationConfig {
139 pub fn responsible_worker<P: Hash>(&self, partition: P) -> usize {
141 let mut h = std::hash::DefaultHasher::default();
142 (self.id, partition).hash(&mut h);
143 let key = usize::cast_from(h.finish());
144 key % self.worker_count
145 }
146
147 pub fn responsible_for<P: Hash>(&self, partition: P) -> bool {
149 self.responsible_worker(partition) == self.worker_id
150 }
151}
152
153pub fn create_raw_source<'scope, 'root, C>(
166 scope: Scope<'scope, mz_repr::Timestamp>,
167 root_scope: Scope<'root, ()>,
168 storage_state: &crate::storage_state::StorageState,
169 committed_upper: StreamVec<'scope, mz_repr::Timestamp, ()>,
170 config: &RawSourceCreationConfig,
171 source_connection: C,
172 start_signal: impl std::future::Future<Output = ()> + 'static,
173) -> (
174 BTreeMap<
175 GlobalId,
176 VecCollection<
177 'scope,
178 mz_repr::Timestamp,
179 Result<SourceOutput<C::Time>, DataflowError>,
180 Diff,
181 >,
182 >,
183 StreamVec<'root, (), HealthStatusMessage>,
184 Vec<PressOnDropButton>,
185)
186where
187 C: SourceConnection + SourceRender + Clone + 'static,
188{
189 let worker_id = config.worker_id;
190 let id = config.id;
191
192 let mut tokens = vec![];
193
194 let (probed_upper_tx, probed_upper_rx) = watch::channel(None);
195
196 let source_metrics = Arc::new(config.metrics.get_source_metrics(id, worker_id));
197
198 let timestamp_desc = source_connection.timestamp_desc();
199
200 let (remap_collection, remap_token) = remap_operator(
201 scope,
202 storage_state,
203 config.clone(),
204 probed_upper_rx,
205 timestamp_desc,
206 );
207 let remap_collection = remap_collection.inner.broadcast().as_collection();
209 tokens.push(remap_token);
210
211 let committed_upper = reclock_committed_upper(
212 remap_collection.clone(),
213 config.as_of.clone(),
214 committed_upper,
215 id,
216 Arc::clone(&source_metrics),
217 );
218
219 let mut reclocked_exports = BTreeMap::new();
220
221 let reclocked_exports2 = &mut reclocked_exports;
222 let (health, source_tokens) = root_scope.scoped("SourceTimeDomain", move |scope| {
223 let (exports, health_stream, source_tokens) = source_render_operator(
224 scope,
225 config,
226 source_connection,
227 probed_upper_tx,
228 committed_upper,
229 start_signal,
230 );
231
232 for (id, export) in exports {
233 let (reclock_pusher, reclocked) =
234 reclock(remap_collection.clone(), config.as_of.clone());
235 export
236 .inner
237 .map(move |(result, from_time, diff)| {
238 let result = match result {
239 Ok(msg) => Ok(SourceOutput {
240 key: msg.key,
241 value: msg.value,
242 metadata: msg.metadata,
243 from_time: from_time.clone(),
244 }),
245 Err(err) => Err(err),
246 };
247 (result, from_time, diff)
248 })
249 .capture_into(PusherCapture(reclock_pusher));
250 reclocked_exports2.insert(id, reclocked);
251 }
252
253 (health_stream.leave(root_scope), source_tokens)
254 });
255
256 tokens.extend(source_tokens);
257
258 (reclocked_exports, health, tokens)
259}
260
261fn source_render_operator<'scope, C>(
264 scope: Scope<'scope, C::Time>,
265 config: &RawSourceCreationConfig,
266 source_connection: C,
267 probed_upper_tx: watch::Sender<Option<Probe<C::Time>>>,
268 resume_uppers: impl futures::Stream<Item = Antichain<C::Time>> + 'static,
269 start_signal: impl std::future::Future<Output = ()> + 'static,
270) -> (
271 BTreeMap<GlobalId, StackedCollection<'scope, C::Time, Result<SourceMessage, DataflowError>>>,
272 StreamVec<'scope, C::Time, HealthStatusMessage>,
273 Vec<PressOnDropButton>,
274)
275where
276 C: SourceRender + 'static,
277{
278 let source_id = config.id;
279 let worker_id = config.worker_id;
280
281 let resume_uppers = resume_uppers.inspect(move |upper| {
282 let upper = upper.pretty();
283 trace!(%upper, "timely-{worker_id} source({source_id}) received resume upper");
284 });
285
286 let (exports, health, probe_stream, tokens) =
287 source_connection.render(scope, config, resume_uppers, start_signal);
288
289 let mut export_collections = BTreeMap::new();
290
291 let source_metrics = config.metrics.get_source_metrics(config.id, worker_id);
292
293 let resume_upper = Antichain::from_iter(
295 config
296 .resume_uppers
297 .values()
298 .flat_map(|f| f.iter().cloned()),
299 );
300 source_metrics
301 .resume_upper
302 .set(mz_persist_client::metrics::encode_ts_metric(&resume_upper));
303
304 let mut health_streams = vec![];
305
306 for (id, export) in exports {
307 let name = format!("SourceGenericStats({})", id);
308 let mut builder = OperatorBuilderRc::new(name, scope.clone());
309
310 let (health_output, derived_health) = builder.new_output();
311 let mut health_output =
312 OutputBuilder::<_, CapacityContainerBuilder<_>>::from(health_output);
313 health_streams.push(derived_health);
314
315 let (output, new_export) = builder.new_output();
316 let mut output = OutputBuilder::<_, CapacityContainerBuilder<_>>::from(output);
317
318 let mut input = builder.new_input(export.inner, Pipeline);
319 export_collections.insert(id, new_export.as_collection());
320
321 let bytes_read_counter = config.metrics.source_defs.bytes_read.clone();
322 let source_statistics = config
323 .statistics
324 .get(&id)
325 .expect("statistics initialized")
326 .clone();
327
328 builder.build(move |mut caps| {
329 let mut health_cap = Some(caps.remove(0));
330
331 move |frontiers| {
332 let mut last_status = None;
333 let mut health_output = health_output.activate();
334
335 if frontiers[0].is_empty() {
336 health_cap = None;
337 return;
338 }
339 let health_cap = health_cap.as_mut().unwrap();
340
341 input.for_each(|cap, data| {
342 for (message, _, _) in data.iter() {
343 match message {
344 Ok(message) => {
345 source_statistics.inc_messages_received_by(1);
346 let key_len = u64::cast_from(message.key.byte_len());
347 let value_len = u64::cast_from(message.value.byte_len());
348 bytes_read_counter.inc_by(key_len + value_len);
349 source_statistics.inc_bytes_received_by(key_len + value_len);
350 }
351 Err(error) => {
352 let update = HealthStatusUpdate::stalled(
356 error.to_string(),
357 Some(
358 "retracting the errored value may resume the source"
359 .to_string(),
360 ),
361 );
362 let status = HealthStatusMessage {
363 id: Some(id),
364 namespace: C::STATUS_NAMESPACE.clone(),
365 update,
366 };
367 if last_status.as_ref() != Some(&status) {
368 last_status = Some(status.clone());
369 health_output.session(&health_cap).give(status);
370 }
371 }
372 }
373 }
374 let mut output = output.activate();
375 output.session(&cap).give_container(data);
376 });
377 }
378 });
379 }
380
381 probe_stream.broadcast().inspect_container(move |event| {
392 if let Ok((_, data)) = event {
393 for probe in data {
394 let _ = probed_upper_tx.send(Some(probe.clone()));
396 }
397 }
398 });
399
400 (
401 export_collections,
402 health.concatenate_flatten::<_, CapacityContainerBuilder<_>>(health_streams),
403 tokens,
404 )
405}
406
407fn remap_operator<'scope, FromTime>(
413 scope: Scope<'scope, mz_repr::Timestamp>,
414 storage_state: &crate::storage_state::StorageState,
415 config: RawSourceCreationConfig,
416 mut probed_upper: watch::Receiver<Option<Probe<FromTime>>>,
417 remap_relation_desc: RelationDesc,
418) -> (
419 VecCollection<'scope, mz_repr::Timestamp, FromTime, Diff>,
420 PressOnDropButton,
421)
422where
423 FromTime: SourceTimestamp,
424{
425 let RawSourceCreationConfig {
426 name,
427 id,
428 source_exports: _,
429 worker_id,
430 worker_count,
431 timestamp_interval: _,
432 remap_metadata,
433 as_of,
434 resume_uppers: _,
435 source_resume_uppers: _,
436 metrics: _,
437 now_fn,
438 persist_clients,
439 statistics: _,
440 shared_remap_upper,
441 config: _,
442 remap_collection_id,
443 busy_signal: _,
444 } = config;
445
446 let read_only_rx = storage_state.read_only_rx.clone();
447 let error_handler = storage_state.error_handler("remap_operator", id);
448
449 let chosen_worker = usize::cast_from(id.hashed() % u64::cast_from(worker_count));
450 let active_worker = chosen_worker == worker_id;
451
452 let operator_name = format!("remap({})", id);
453 let mut remap_op = AsyncOperatorBuilder::new(operator_name, scope.clone());
454 let (remap_output, remap_stream) = remap_op.new_output::<CapacityContainerBuilder<_>>();
455
456 let button = remap_op.build(move |capabilities| async move {
457 if !active_worker {
458 shared_remap_upper.borrow_mut().clear();
461 return;
462 }
463
464 let mut cap_set = CapabilitySet::from_elem(capabilities.into_element());
465
466 let remap_handle = crate::source::reclock::compat::PersistHandle::<FromTime, _>::new(
467 Arc::clone(&persist_clients),
468 read_only_rx,
469 remap_metadata.clone(),
470 as_of.clone(),
471 shared_remap_upper,
472 id,
473 "remap",
474 worker_id,
475 worker_count,
476 remap_relation_desc,
477 remap_collection_id,
478 )
479 .await;
480
481 let remap_handle = match remap_handle {
482 Ok(handle) => handle,
483 Err(e) => {
484 error_handler
485 .report_and_stop(
486 e.context(format!("Failed to create remap handle for source {name}")),
487 )
488 .await
489 }
490 };
491
492 let (mut timestamper, mut initial_batch) = ReclockOperator::new(remap_handle).await;
493
494 trace!(
497 "timely-{worker_id} remap({id}) emitting remap snapshot: trace_updates={:?}",
498 &initial_batch.updates
499 );
500
501 let cap = cap_set.delayed(cap_set.first().unwrap());
502 remap_output.give_container(&cap, &mut initial_batch.updates);
503 drop(cap);
504 cap_set.downgrade(initial_batch.upper);
505
506 let mut prev_probe_ts: Option<mz_repr::Timestamp> = None;
507
508 while !cap_set.is_empty() {
509 let new_probe = probed_upper
511 .wait_for(|new_probe| match (prev_probe_ts, new_probe) {
512 (None, Some(_)) => true,
513 (Some(prev_ts), Some(new)) => prev_ts < new.probe_ts,
514 _ => false,
515 })
516 .await
517 .map(|probe| (*probe).clone())
518 .unwrap_or_else(|_| {
519 Some(Probe {
520 probe_ts: now_fn().into(),
521 upstream_frontier: Antichain::new(),
522 })
523 });
524
525 let probe = new_probe.expect("known to be Some");
526 prev_probe_ts = Some(probe.probe_ts);
527
528 let binding_ts = probe.probe_ts;
529 let cur_source_upper = probe.upstream_frontier;
530
531 let new_into_upper = Antichain::from_elem(binding_ts.step_forward());
532
533 let mut remap_trace_batch = timestamper
534 .mint(binding_ts, new_into_upper, cur_source_upper.borrow())
535 .await;
536
537 trace!(
538 "timely-{worker_id} remap({id}) minted new bindings: \
539 updates={:?} \
540 source_upper={} \
541 trace_upper={}",
542 &remap_trace_batch.updates,
543 cur_source_upper.pretty(),
544 remap_trace_batch.upper.pretty()
545 );
546
547 let cap = cap_set.delayed(cap_set.first().unwrap());
548 remap_output.give_container(&cap, &mut remap_trace_batch.updates);
549 cap_set.downgrade(remap_trace_batch.upper);
550 }
551 });
552
553 (remap_stream.as_collection(), button.press_on_drop())
554}
555
556fn reclock_committed_upper<'scope, T, FromTime>(
560 bindings: VecCollection<'scope, T, FromTime, Diff>,
561 as_of: Antichain<T>,
562 committed_upper: StreamVec<'scope, T, ()>,
563 id: GlobalId,
564 metrics: Arc<SourceMetrics>,
565) -> impl futures::stream::Stream<Item = Antichain<FromTime>> + 'static
566where
567 T: Timestamp + Lattice + TotalOrder,
568 FromTime: SourceTimestamp,
569{
570 let (tx, rx) = watch::channel(Antichain::from_elem(FromTime::minimum()));
571 let scope = bindings.scope().clone();
572
573 let name = format!("ReclockCommitUpper({id})");
574 let mut builder = OperatorBuilderRc::new(name, scope);
575
576 let mut bindings = builder.new_input(bindings.inner.clone(), Pipeline);
577 let _ = builder.new_input(committed_upper.clone(), Pipeline);
578
579 builder.build(move |_| {
580 use timely::progress::ChangeBatch;
582 let mut accepted_times: ChangeBatch<(T, FromTime)> = ChangeBatch::new();
583 let mut upper = Antichain::from_elem(Timestamp::minimum());
585 let mut ready_times = VecDeque::new();
587 let mut source_upper = MutableAntichain::new();
588
589 move |frontiers| {
590 bindings.for_each(|_, data| {
592 accepted_times.extend(data.drain(..).map(|(from, mut into, diff)| {
593 into.advance_by(as_of.borrow());
594 ((into, from), diff.into_inner())
595 }));
596 });
597 let new_upper = frontiers[0].frontier();
599 if PartialOrder::less_than(&upper.borrow(), &new_upper) {
600 upper = new_upper.to_owned();
601 let mut pending_times = std::mem::take(&mut accepted_times).into_inner();
604 pending_times.sort_unstable_by(|a, b| a.0.cmp(&b.0));
606 for ((into, from), diff) in pending_times.drain(..) {
607 if !upper.less_equal(&into) {
608 ready_times.push_back((from, into, diff));
609 } else {
610 accepted_times.update((into, from), diff);
611 }
612 }
613 }
614
615 if as_of.iter().all(|t| !upper.less_equal(t)) {
617 let committed_upper = frontiers[1].frontier();
618 if as_of.iter().all(|t| !committed_upper.less_equal(t)) {
619 let reclocked_upper = match committed_upper.as_option() {
660 Some(t_next) => {
661 let idx = ready_times.partition_point(|(_, t, _)| t < t_next);
662 let updates = ready_times
663 .drain(0..idx)
664 .map(|(from_time, _, diff)| (from_time, diff));
665 source_upper.update_iter(updates);
666 source_upper.frontier().to_owned()
669 }
670 None => Antichain::new(),
671 };
672 tx.send_replace(reclocked_upper);
673 }
674 }
675
676 metrics
677 .commit_upper_accepted_times
678 .set(u64::cast_from(accepted_times.len()));
679 metrics
680 .commit_upper_ready_times
681 .set(u64::cast_from(ready_times.len()));
682 }
683 });
684
685 WatchStream::from_changes(rx)
686}