diff --git a/src/storage/src/source/postgres.rs b/src/storage/src/source/postgres.rs index 48f4c9c1f76de..4f929f66f2505 100644 --- a/src/storage/src/source/postgres.rs +++ b/src/storage/src/source/postgres.rs @@ -76,23 +76,25 @@ //! │concat│ │ //! ╰──┬───╯ │ //! │ data │progress -//! │ output │output +//! │ outputs │output +//! │ (one per export) │ //! v v //! ``` +//! +//! Both the snapshot and the replication operators emit one output stream per source export, and +//! the streams of each export are concatenated pairwise. use std::collections::BTreeMap; use std::rc::Rc; use std::time::Duration; -use differential_dataflow::AsCollection; use itertools::Itertools as _; use mz_expr::EvalError; -use mz_ore::cast::CastFrom; use mz_ore::error::ErrorExt; use mz_postgres_util::desc::PostgresTableDesc; use mz_postgres_util::schema_change::SchemaChangeError; use mz_postgres_util::{Client, PostgresError, Sql, query_opt, simple_query_opt, sql}; -use mz_repr::{Datum, Diff, GlobalId, Row}; +use mz_repr::{Datum, GlobalId, Row}; use mz_storage_types::errors::{DataflowError, SourceError, SourceErrorDetails}; use mz_storage_types::sources::casts::StorageScalarExpr; use mz_storage_types::sources::postgres::CastType; @@ -101,9 +103,7 @@ use mz_storage_types::sources::{ }; use mz_timely_util::builder_async::PressOnDropButton; use serde::{Deserialize, Serialize}; -use timely::container::CapacityContainerBuilder; use timely::dataflow::operators::Concat; -use timely::dataflow::operators::core::Partition; use timely::dataflow::operators::vec::{Map, ToStream}; use timely::dataflow::{Scope, StreamVec}; use timely::progress::{Antichain, Timestamp}; @@ -199,24 +199,16 @@ impl SourceRender for PostgresSourceConnection { metrics, ); - let updates = snapshot_updates.concat(repl_updates); - let partition_count = u64::cast_from(config.source_exports.len()); - let data_streams: Vec<_> = updates - .inner - .partition::, _, _>( - partition_count, - |((output, data), time, diff): ( - (usize, Result), - MzOffset, - Diff, - )| { - let output = u64::cast_from(output); - (output, (data, time, diff)) - }, - ); + // Both operators produce one stream per export, in output index order, which matches the + // iteration order of `source_exports`. let mut data_collections = BTreeMap::new(); - for (id, data_stream) in config.source_exports.keys().zip_eq(data_streams) { - data_collections.insert(*id, data_stream.as_collection()); + for ((id, snapshot_updates), repl_updates) in config + .source_exports + .keys() + .zip_eq(snapshot_updates) + .zip_eq(repl_updates) + { + data_collections.insert(*id, snapshot_updates.concat(repl_updates)); } let export_ids = config.source_exports.keys().copied(); diff --git a/src/storage/src/source/postgres/replication.rs b/src/storage/src/source/postgres/replication.rs index 44911e9d26c5d..5a76625e715d6 100644 --- a/src/storage/src/source/postgres/replication.rs +++ b/src/storage/src/source/postgres/replication.rs @@ -69,6 +69,7 @@ //! [1]: https://www.postgresql.org/docs/15/protocol-replication.html#PROTOCOL-REPLICATION-START-REPLICATION //! [2]: https://www.postgresql.org/message-id/CAFPTHDZS9O9WG02EfayBd6oONzK%2BqfUxS6AbVLJ7W%2BKECza2gg%40mail.gmail.com +use std::cell::Cell; use std::collections::BTreeMap; use std::convert::Infallible; use std::pin::pin; @@ -117,7 +118,7 @@ use crate::source::postgres::verify_schema; use crate::source::postgres::{DefiniteError, ReplicationError, SourceOutputInfo, TransientError}; use crate::source::probe; use crate::source::types::{ - FuelSize, Probe, ResumeUppers, SignaledFuture, SourceMessage, StackedCollection, + Probe, ResumeUppers, SharedFuel, SignaledFuture, SourceMessage, StackedCollection, }; /// A logical replication message from the server. @@ -153,7 +154,7 @@ pub(crate) fn render<'scope>( committed_uppers: impl futures::Stream> + 'static, metrics: PgSourceMetrics, ) -> ( - StackedCollection<'scope, MzOffset, (usize, Result)>, + Vec>>, StreamVec<'scope, MzOffset, Probe>, StreamVec<'scope, MzOffset, ReplicationError>, PressOnDropButton, @@ -162,7 +163,16 @@ pub(crate) fn render<'scope>( let mut builder = AsyncOperatorBuilder::new(op_name, scope.clone()); let slot_reader = u64::cast_from(config.responsible_worker("slot")); - let (data_output, data_stream) = builder.new_output(); + // One data output port per source export, in output index order. All data port capabilities + // are managed in lockstep, so every export observes the same frontier. + let export_count = config.source_exports.len(); + let mut data_outputs = Vec::with_capacity(export_count); + let mut data_streams = Vec::with_capacity(export_count); + for _ in 0..export_count { + let (output, stream) = builder.new_output::>>(); + data_outputs.push(output); + data_streams.push(stream); + } let (definite_error_handle, definite_errors) = builder.new_output::>(); let (probe_output, probe_stream) = builder.new_output::>(); @@ -182,8 +192,9 @@ pub(crate) fn render<'scope>( let busy_signal = Arc::clone(&config.busy_signal); Box::pin(SignaledFuture::new(busy_signal, async move { let (id, worker_id) = (config.id, config.worker_id); - let [data_cap_set, definite_error_cap_set, probe_cap]: &mut [_; 3] = - caps.try_into().unwrap(); + let (data_cap_sets, caps) = caps.split_at_mut(export_count); + let [definite_error_cap_set, probe_cap]: &mut [_; 2] = caps.try_into().unwrap(); + let mut data_outputs = SharedFuel::new(data_outputs); if !config.responsible_for("slot") { // Emit 0, to mark this worker as having started up correctly. @@ -352,22 +363,21 @@ pub(crate) fn render<'scope>( ); // If the replication stream cannot be obtained from the resume point there is nothing // else to do. These errors are not retractable. - for (oid, outputs) in table_info.iter() { + for outputs in table_info.values() { for output_index in outputs.keys() { // We pick `u64::MAX` as the LSN which will (in practice) never conflict // any previously revealed portions of the TVC. let update = ( - ( - *oid, - *output_index, - Err(DataflowError::from(err.clone())), - ), + Err(DataflowError::from(err.clone())), MzOffset::from(u64::MAX), Diff::ONE, ); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + data_outputs + .give( + *output_index, + &data_cap_sets[*output_index][0], + update, + ) .await; } } @@ -405,18 +415,17 @@ pub(crate) fn render<'scope>( Err(err) => { // If the replication stream cannot be obtained in a definite way there is // nothing else to do. These errors are not retractable. - for (oid, outputs) in table_info.iter() { + for outputs in table_info.values() { for output_index in outputs.keys() { // We pick `u64::MAX` as the LSN which will (in practice) never conflict // any previously revealed portions of the TVC. let update = ( - (*oid, *output_index, Err(DataflowError::from(err.clone()))), + Err(DataflowError::from(err.clone())), MzOffset::from(u64::MAX), Diff::ONE, ); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + data_outputs + .give(*output_index, &data_cap_sets[*output_index][0], update) .await; } } @@ -476,23 +485,25 @@ pub(crate) fn render<'scope>( "new_upper={data_upper} tx_lsn={commit_lsn}", ); data_upper = commit_lsn + 1; - while let Some((oid, output_index, event, diff)) = tx.try_next().await? + while let Some((_oid, output_index, event, diff)) = + tx.try_next().await? { - let event = event.map_err(Into::into); - let data = (oid, output_index, event); + let event: Result<_, DataflowError> = event.map_err(Into::into); if let Some(req) = rewinds.get(&output_index) { if commit_lsn <= req.snapshot_lsn { - let update = (data.clone(), MzOffset::from(0), -diff); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + let update = (event.clone(), MzOffset::from(0), -diff); + data_outputs + .give( + output_index, + &data_cap_sets[output_index][0], + update, + ) .await; } } - let update = (data, commit_lsn, diff); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + let update = (event, commit_lsn, diff); + data_outputs + .give(output_index, &data_cap_sets[output_index][0], update) .await; } } @@ -510,20 +521,19 @@ pub(crate) fn render<'scope>( match error { Postgres(PostgresError::PublicationMissing(publication)) => { let err = DefiniteError::PublicationDropped(publication); - for (oid, outputs) in table_info.iter() { + for outputs in table_info.values() { for output_index in outputs.keys() { let update = ( - ( - *oid, - *output_index, - Err(DataflowError::from(err.clone())), - ), - data_cap_set[0].time().clone(), + Err(DataflowError::from(err.clone())), + data_cap_sets[*output_index][0].time().clone(), Diff::ONE, ); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + data_outputs + .give( + *output_index, + &data_cap_sets[*output_index][0], + update, + ) .await; } } @@ -542,21 +552,20 @@ pub(crate) fn render<'scope>( // definite, non-retryable error. let err = DefiniteError::InvalidPhysicalReplica { expected, actual }; - for (oid, outputs) in table_info.iter() { + for outputs in table_info.values() { for output_index in outputs.keys() { let update = ( - ( - *oid, - *output_index, - Err(DataflowError::from(err.clone())), - ), + Err(DataflowError::from(err.clone())), // We don't have a clean way to align on when the replica changed so jump straight to u64::MAX to avoid conflicts. MzOffset::from(u64::MAX), Diff::ONE, ); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + data_outputs + .give( + *output_index, + &data_cap_sets[*output_index][0], + update, + ) .await; } } @@ -577,13 +586,12 @@ pub(crate) fn render<'scope>( } let update = ( - (oid, output_index, Err(error.into())), - data_cap_set[0].time().clone(), + Err(error.into()), + data_cap_sets[output_index][0].time().clone(), Diff::ONE, ); - let size = update.fuel_size(); - data_output - .give_fueled(&data_cap_set[0], update, size) + data_outputs + .give(output_index, &data_cap_sets[output_index][0], update) .await; } } @@ -598,10 +606,12 @@ pub(crate) fn render<'scope>( if will_yield { trace!(%id, "timely-{worker_id} yielding at lsn={data_upper}"); rewinds.retain(|_, req| data_upper <= req.snapshot_lsn); - // As long as there are pending rewinds we can't downgrade our data capability - // since we must be able to produce data at offset 0. + // As long as there are pending rewinds we can't downgrade our data + // capabilities since we must be able to produce data at offset 0. if rewinds.is_empty() { - data_cap_set.downgrade([&data_upper]); + for cap_set in data_cap_sets.iter_mut() { + cap_set.downgrade([&data_upper]); + } } } } @@ -610,43 +620,61 @@ pub(crate) fn render<'scope>( })) }); - // We now process the slot updates and apply the cast expressions - let mut final_row = Row::default(); - let mut datum_vec = DatumVec::new(); - let mut next_worker = (0..u64::cast_from(scope.peers())) + // We now process the slot updates and apply the cast expressions, with one decode operator + // per export. + let mut output_info = BTreeMap::new(); + for outputs in table_info.into_values() { + for (output_index, info) in outputs { + output_info.insert(output_index, info); + } + } + let mut replication_updates = Vec::with_capacity(data_streams.len()); + let peers = u64::cast_from(scope.peers()); + let record = Rc::new(Cell::new(0_u64)); + for (output_index, data_stream) in data_streams.into_iter().enumerate() { + let info = output_info.get(&output_index).cloned(); + let mut final_row = Row::default(); + let mut datum_vec = DatumVec::new(); + // Round robin on 1000-records basis to avoid creating tiny containers when there are a - // small number of updates and a large number of workers. - .flat_map(|w| std::iter::repeat_n(w, 1000)) - .cycle(); - let round_robin = Exchange::new(move |_| next_worker.next().unwrap()); - let replication_updates = data_stream - .unary(round_robin, "PgCastReplicationRows", |_, _| { - move |input, output| { - input.for_each_time(|time, data| { - let mut session = output.session(&time); - for ((oid, output_index, event), time, diff) in - data.flat_map(|data| data.drain(..)) - { - let output = &table_info - .get(&oid) - .and_then(|outputs| outputs.get(&output_index)) - .expect("table_info contains all outputs"); - let event = event.and_then(|row| { - let datums = datum_vec.borrow_with(&row); - super::cast_row(&output.casts, &datums, &mut final_row)?; - Ok(SourceMessage { - key: Row::default(), - value: final_row.clone(), - metadata: Row::default(), - }) + // small number of updates and a large number of workers. The counter is shared by every + // export's exchange so rotation follows total volume. + let record = Rc::clone(&record); + let round_robin = Exchange::new(move |_| { + let n = record.get(); + record.set(n.wrapping_add(1)); + (n / 1000) % peers + }); + let updates = data_stream + .unary( + round_robin, + &format!("PgCastReplicationRows({output_index})"), + |_, _| { + move |input, output| { + input.for_each_time(|time, data| { + let mut session = output.session(&time); + for (event, time, diff) in data.flat_map(|data| data.drain(..)) { + let info = info + .as_ref() + .expect("only exports with output info receive data"); + let event = event.and_then(|row| { + let datums = datum_vec.borrow_with(&row); + super::cast_row(&info.casts, &datums, &mut final_row)?; + Ok(SourceMessage { + key: Row::default(), + value: final_row.clone(), + metadata: Row::default(), + }) + }); + session.give((event, time, diff)); + } }); - - session.give(((output_index, event), time, diff)); } - }); - } - }) - .as_collection(); + }, + ) + .as_collection(); + replication_updates.push(updates); + } let errors = definite_errors.concat(transient_errors.map(ReplicationError::from)); diff --git a/src/storage/src/source/postgres/snapshot.rs b/src/storage/src/source/postgres/snapshot.rs index 328444c3eda0b..6300591e97f5f 100644 --- a/src/storage/src/source/postgres/snapshot.rs +++ b/src/storage/src/source/postgres/snapshot.rs @@ -198,7 +198,7 @@ use crate::source::postgres::replication::RewindRequest; use crate::source::postgres::{ DefiniteError, ReplicationError, SourceOutputInfo, TransientError, verify_schema, }; -use crate::source::types::{FuelSize, SignaledFuture, SourceMessage, StackedCollection}; +use crate::source::types::{SharedFuel, SignaledFuture, SourceMessage, StackedCollection}; use crate::statistics::SourceStatistics; /// Information broadcasted from the snapshot leader to all workers. @@ -338,7 +338,7 @@ pub(crate) fn render<'scope>( table_info: BTreeMap>, metrics: PgSnapshotMetrics, ) -> ( - StackedCollection<'scope, MzOffset, (usize, Result)>, + Vec>>, StreamVec<'scope, MzOffset, RewindRequest>, StreamVec<'scope, MzOffset, Infallible>, StreamVec<'scope, MzOffset, ReplicationError>, @@ -349,7 +349,16 @@ pub(crate) fn render<'scope>( let (feedback_handle, feedback_data) = scope.feedback(Default::default()); - let (raw_handle, raw_data) = builder.new_output(); + // One data output port per source export, in output index order. All data port capabilities + // are managed in lockstep, so every export observes the same frontier. + let export_count = config.source_exports.len(); + let mut raw_handles = Vec::with_capacity(export_count); + let mut raw_streams = Vec::with_capacity(export_count); + for _ in 0..export_count { + let (handle, stream) = builder.new_output::>>(); + raw_handles.push(handle); + raw_streams.push(stream); + } let (rewinds_handle, rewinds) = builder.new_output::>(); // This output is used to signal to the replication operator that the replication slot has been // created. With the current state of execution serialization there isn't a lot of benefit @@ -404,13 +413,14 @@ pub(crate) fn render<'scope>( Box::pin(SignaledFuture::new(busy_signal, async move { let id = config.id; let worker_id = config.worker_id; + let (data_cap_sets, caps) = caps.split_at_mut(export_count); let [ - data_cap_set, rewind_cap_set, slot_ready_cap_set, snapshot_cap_set, definite_error_cap_set, - ]: &mut [_; 5] = caps.try_into().unwrap(); + ]: &mut [_; 4] = caps.try_into().unwrap(); + let mut raw_handles = SharedFuel::new(raw_handles); let connection_config = connection .connection @@ -530,7 +540,7 @@ pub(crate) fn render<'scope>( // nothing else to do. These errors are not retractable. Err(PostgresError::PublicationMissing(publication)) => { let err = DefiniteError::PublicationDropped(publication); - for (oid, outputs) in tables_to_snapshot.iter() { + for outputs in tables_to_snapshot.values() { // Produce a definite error here and then exit to ensure // a missing publication doesn't generate a transient // error and restart this dataflow indefinitely. @@ -540,13 +550,16 @@ pub(crate) fn render<'scope>( // portions of the TVC. for output_index in outputs.keys() { let update = ( - (*oid, *output_index, Err(err.clone().into())), + Err(err.clone().into()), MzOffset::from(u64::MAX), Diff::ONE, ); - let size = update.fuel_size(); - raw_handle - .give_fueled(&data_cap_set[0], update, size) + raw_handles + .give( + *output_index, + &data_cap_sets[*output_index][0], + update, + ) .await; } } @@ -635,14 +648,10 @@ pub(crate) fn render<'scope>( } // We pick `u64::MAX` as the LSN which will (in practice) never conflict // any previously revealed portions of the TVC. - let update = ( - (oid, output_index, Err(err.clone().into())), - MzOffset::from(u64::MAX), - Diff::ONE, - ); - let size = update.fuel_size(); - raw_handle - .give_fueled(&data_cap_set[0], update, size) + let update = + (Err(err.clone().into()), MzOffset::from(u64::MAX), Diff::ONE); + raw_handles + .give(output_index, &data_cap_sets[output_index][0], update) .await; } } @@ -664,14 +673,9 @@ pub(crate) fn render<'scope>( for (&oid, outputs) in tables_to_snapshot.iter() { for (&output_index, info) in outputs.iter() { if let Err(err) = verify_schema(oid, info, &upstream_info) { - let update = ( - (oid, output_index, Err(err.into())), - MzOffset::minimum(), - Diff::ONE, - ); - let size = update.fuel_size(); - raw_handle - .give_fueled(&data_cap_set[0], update, size) + let update = (Err(err.into()), MzOffset::minimum(), Diff::ONE); + raw_handles + .give(output_index, &data_cap_sets[output_index][0], update) .await; continue; } @@ -727,14 +731,9 @@ pub(crate) fn render<'scope>( let mut snapshot_staged = 0; while let Some(bytes) = stream.try_next().await? { - let update = ( - (oid, output_index, Ok(bytes)), - MzOffset::minimum(), - Diff::ONE, - ); - let size = update.fuel_size(); - raw_handle - .give_fueled(&data_cap_set[0], update, size) + let update = (Ok(bytes), MzOffset::minimum(), Diff::ONE); + raw_handles + .give(output_index, &data_cap_sets[output_index][0], update) .await; snapshot_staged += 1; if snapshot_staged % 1000 == 0 { @@ -797,43 +796,54 @@ pub(crate) fn render<'scope>( })) }); - // We now decode the COPY protocol and apply the cast expressions - let mut text_row = Row::default(); - let mut final_row = Row::default(); - let mut datum_vec = DatumVec::new(); - let snapshot_updates = raw_data - .unary(Pipeline, "PgCastSnapshotRows", |_, _| { - move |input, output| { - input.for_each_time(|time, data| { - let mut session = output.session(&time); - for ((oid, output_index, event), time, diff) in - data.flat_map(|data| data.drain(..)) - { - let output = &table_info - .get(&oid) - .and_then(|outputs| outputs.get(&output_index)) - .expect("table_info contains all outputs"); - - let event = event - .as_ref() - .map_err(|e: &DataflowError| e.clone()) - .and_then(|bytes| { - decode_copy_row(bytes, output.casts.len(), &mut text_row)?; - let datums = datum_vec.borrow_with(&text_row); - super::cast_row(&output.casts, &datums, &mut final_row)?; - Ok(SourceMessage { - key: Row::default(), - value: final_row.clone(), - metadata: Row::default(), - }) - }); - - session.give(((output_index, event), time, diff)); + // We now decode the COPY protocol and apply the cast expressions, with one decode operator + // per export. + let mut output_info = BTreeMap::new(); + for outputs in table_info.into_values() { + for (output_index, info) in outputs { + output_info.insert(output_index, info); + } + } + let mut snapshot_updates = Vec::with_capacity(raw_streams.len()); + for (output_index, raw_stream) in raw_streams.into_iter().enumerate() { + let info = output_info.get(&output_index).cloned(); + let mut text_row = Row::default(); + let mut final_row = Row::default(); + let mut datum_vec = DatumVec::new(); + let updates = raw_stream + .unary( + Pipeline, + &format!("PgCastSnapshotRows({output_index})"), + |_, _| { + move |input, output| { + input.for_each_time(|time, data| { + let mut session = output.session(&time); + for (event, time, diff) in data.flat_map(|data| data.drain(..)) { + let info = info + .as_ref() + .expect("only exports with output info receive data"); + let event = event + .as_ref() + .map_err(|e: &DataflowError| e.clone()) + .and_then(|bytes| { + decode_copy_row(bytes, info.casts.len(), &mut text_row)?; + let datums = datum_vec.borrow_with(&text_row); + super::cast_row(&info.casts, &datums, &mut final_row)?; + Ok(SourceMessage { + key: Row::default(), + value: final_row.clone(), + metadata: Row::default(), + }) + }); + session.give((event, time, diff)); + } + }); } - }); - } - }) - .as_collection(); + }, + ) + .as_collection(); + snapshot_updates.push(updates); + } let errors = definite_errors.concat(transient_errors.map(ReplicationError::from)); diff --git a/src/storage/src/source/types.rs b/src/storage/src/source/types.rs index d5f6294fc7fe6..dc9fa6f03c578 100644 --- a/src/storage/src/source/types.rs +++ b/src/storage/src/source/types.rs @@ -24,11 +24,14 @@ use mz_repr::{Diff, GlobalId, Row}; use mz_storage_types::errors::{DataflowError, DecodeError}; use mz_storage_types::sources::SourceTimestamp; use mz_timely_util::antichain::AntichainExt; -use mz_timely_util::builder_async::PressOnDropButton; +use mz_timely_util::builder_async::{AsyncOutputHandle, MAX_OUTSTANDING_BYTES, PressOnDropButton}; use pin_project::pin_project; use serde::{Deserialize, Serialize}; +use timely::ContainerBuilder; +use timely::container::PushInto; +use timely::dataflow::operators::Capability; use timely::dataflow::{Scope, StreamVec}; -use timely::progress::Antichain; +use timely::progress::{Antichain, Timestamp}; use tokio::sync::Semaphore; use tokio_util::sync::PollSemaphore; @@ -243,6 +246,55 @@ impl FuelSize for mz_timely_util::order::Partitioned { } } +/// Output handles of one operator that share a fuel budget, so that the operator yields back to +/// timely once per [`MAX_OUTSTANDING_BYTES`] of its total emission across them. +/// +/// [`AsyncOutputHandle::give_fueled`] counts bytes per handle, so an operator with many handles +/// could emit up to the bound on each of them before yielding. +pub struct SharedFuel { + handles: Vec>, + outstanding_bytes: usize, +} + +impl SharedFuel { + /// Takes ownership of `handles`, which [`Self::give`] addresses by position. + pub fn new(handles: Vec>) -> Self { + Self { + handles, + outstanding_bytes: 0, + } + } + + /// Gives `data` at `cap` to the handle at `index`, charging its [`FuelSize`] against the + /// shared budget. Yields back to timely and resets the budget once it exceeds + /// [`MAX_OUTSTANDING_BYTES`]. + /// + /// # Panics + /// + /// Panics if `index` is not the position of a handle passed to [`Self::new`]. + pub async fn give(&mut self, index: usize, cap: &Capability, data: D) + where + CB: PushInto, + D: FuelSize, + { + let size = data.fuel_size(); + self.handles[index].give(cap, data); + if self.charge(size) { + tokio::task::yield_now().await; + } + } + + /// Charges `size`, returning whether the budget was exceeded, in which case it resets. + fn charge(&mut self, size: usize) -> bool { + self.outstanding_bytes += size; + let exceeded = self.outstanding_bytes > MAX_OUTSTANDING_BYTES; + if exceeded { + self.outstanding_bytes = 0; + } + exceeded + } +} + /// The result of probing an upstream system for its write frontier. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Probe { diff --git a/src/timely-util/src/builder_async.rs b/src/timely-util/src/builder_async.rs index 51eb84b39d8b1..fd3d11fdfe809 100644 --- a/src/timely-util/src/builder_async.rs +++ b/src/timely-util/src/builder_async.rs @@ -313,16 +313,18 @@ where } } +/// Maximum bytes [`AsyncOutputHandle::give_fueled`] emits through one handle before yielding back +/// to timely. +pub const MAX_OUTSTANDING_BYTES: usize = 128 * 1024 * 1024; + impl AsyncOutputHandle>>> where D: Clone + 'static, T: Timestamp, { - pub const MAX_OUTSTANDING_BYTES: usize = 128 * 1024 * 1024; - /// Provides one record at the time specified by the capability and /// charges `size_bytes` against the builder's fuel counter. Once at least - /// [`Self::MAX_OUTSTANDING_BYTES`] have been emitted since the last yield + /// [`MAX_OUTSTANDING_BYTES`] have been emitted since the last yield /// this method yields back to timely and resets the counter. pub async fn give_fueled(&self, cap: &Capability, data: D2, size_bytes: usize) where @@ -333,7 +335,7 @@ where handle.give(cap, data); let bytes = &handle.builder.bytes; bytes.set(bytes.get() + size_bytes); - let should_yield = bytes.get() > Self::MAX_OUTSTANDING_BYTES; + let should_yield = bytes.get() > MAX_OUTSTANDING_BYTES; if should_yield { bytes.set(0); }