diff --git a/crates/dbsp/src/circuit/circuit_builder.rs b/crates/dbsp/src/circuit/circuit_builder.rs index 91f59bccf5a..e07d24b8860 100644 --- a/crates/dbsp/src/circuit/circuit_builder.rs +++ b/crates/dbsp/src/circuit/circuit_builder.rs @@ -3939,31 +3939,12 @@ where SndOp: SinkOperator, RcvOp: SourceOperator, { - let sender_id = self.add_node(|id| { - self.log_circuit_event(&CircuitEvent::operator( - GlobalNodeId::child_of(self, id), - sender.name(), - sender.location(), - )); - - let node = SinkNode::new(sender, input_stream.clone(), self.clone(), id); - self.connect_stream(input_stream, id, input_preference); - (node, id) - }); - - let output_stream = self.add_node(|id| { - self.log_circuit_event(&CircuitEvent::operator( - GlobalNodeId::child_of(self, id), - receiver.name(), - receiver.location(), - )); - - let node = SourceNode::new(receiver, self.clone(), id); - let output_stream = node.output_stream(); - (node, output_stream) - }); - - self.add_dependency(sender_id, output_stream.local_node_id()); + let sender_id = self.add_sink_with_preference(sender, input_stream, input_preference); + let output_stream = self.add_source(receiver); + self.add_dependency( + sender_id.local_node_id().unwrap(), + output_stream.local_node_id(), + ); output_stream } diff --git a/crates/dbsp/src/circuit/runtime.rs b/crates/dbsp/src/circuit/runtime.rs index 8e24494cb96..9720ffe1dff 100644 --- a/crates/dbsp/src/circuit/runtime.rs +++ b/crates/dbsp/src/circuit/runtime.rs @@ -7,7 +7,7 @@ use crate::SchedulerError; use crate::circuit::checkpointer::Checkpointer; use crate::circuit::dbsp_handle::StepSize; use crate::error::Error as DbspError; -use crate::operator::communication::Exchange; +use crate::operator::communication::{Exchange, ExchangeActivity}; use crate::storage::backend::StorageBackend; use crate::storage::file::format::Compression; use crate::storage::file::writer::Parameters; @@ -1383,7 +1383,8 @@ where match Runtime::runtime() { Some(runtime) if Runtime::num_workers() > 1 => { let exchange_id = runtime.sequence_next().try_into().unwrap(); - let exchange = Exchange::with_runtime(&runtime, exchange_id); + let exchange = + Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps); let identifier = Arc::new(format!("broadcast {name} (exchange {exchange_id})")); Self::MultiThreaded { @@ -1426,8 +1427,9 @@ where .await; exchange - .receive_all(|data| serde_json::from_slice(&data).unwrap()) + .receive_all(|data| serde_json::from_slice(&data).unwrap(), None) .await + .0 }) .await .ok_or(SchedulerError::Killed) diff --git a/crates/dbsp/src/operator/communication.rs b/crates/dbsp/src/operator/communication.rs index ec1eb6b591a..73ab4ef33cc 100644 --- a/crates/dbsp/src/operator/communication.rs +++ b/crates/dbsp/src/operator/communication.rs @@ -3,6 +3,9 @@ mod gather; mod shard; pub(crate) use exchange::{ - Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, pop_flushed, + Exchange, ExchangeClients, ExchangeDelivery, ExchangeDirectory, ExchangeId, MessageType, + pop_flushed, +}; +pub use exchange::{ + ExchangeActivity, ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators, }; -pub use exchange::{ExchangeReceiver, ExchangeSender, Mailbox, new_exchange_operators}; diff --git a/crates/dbsp/src/operator/communication/exchange.rs b/crates/dbsp/src/operator/communication/exchange.rs index b6c2e6d289f..71f22513b01 100644 --- a/crates/dbsp/src/operator/communication/exchange.rs +++ b/crates/dbsp/src/operator/communication/exchange.rs @@ -23,6 +23,7 @@ use crate::{ }; use binrw::{BinRead, BinWrite}; use crossbeam_utils::CachePadded; +use enum_map::{Enum, EnumMap}; use feldera_samply::Span; use feldera_storage::fbuf::FBuf; use itertools::Itertools; @@ -239,6 +240,45 @@ struct ExchangeMessage { data: Vec, } +/// Distinguishes messages by size. +/// +/// We segregate big and small messages into separate queues so that simple +/// broadcast and consensus messages don't get delayed behind bigger messages. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)] +pub(crate) enum MessageSize { + /// A big message. + Big, + + /// A small message. + Small, +} + +impl MessageSize { + /// Constructs `MessageSize` from a count of `bytes`. + pub(crate) fn from_bytes(bytes: usize) -> Self { + if bytes <= 4096 { + Self::Small + } else { + Self::Big + } + } +} + +/// Categorizes an [ExchangeMessage]. +/// +/// We segregate messages in different categories into different queues so that +/// messages in one category do not delay those in other categories. +#[derive(Copy, Clone, Debug, PartialEq, Eq, Enum)] +pub(crate) enum MessageType { + /// Messages sent via synchronous exchange. + Synchronous( + /// Message size. + MessageSize, + ), + /// Messages sent via streaming exchange. + Streaming, +} + pub struct ExchangeClient { tx: ByteBoundedSender, } @@ -342,6 +382,10 @@ impl ExchangeClient { }) .expect("remote exchange failed") } + + pub async fn wait(&self) { + self.tx.bound.wait().await; + } } /// Uniquely identifies an `Exchange` or `ShardedAccumulator` within a circuit. @@ -484,9 +528,11 @@ pub struct ExchangeClients { /// tries to send data to one. listener: OnceCell>, - /// Maps from a range of worker IDs to the RPC client used to contact those + /// Maps from a range of worker IDs to the RPC clients used to contact those /// workers. Only worker IDs for remote workers appear in the map. - clients: Vec<(Host, OnceCell)>, + /// + /// We use one RPC client per [MessageType] per [Host]. + clients: Vec<(Host, EnumMap>)>, } impl ExchangeClients { @@ -509,14 +555,14 @@ impl ExchangeClients { clients: runtime .layout() .other_hosts() - .map(|host| (host.clone(), OnceCell::new())) + .map(|host| (host.clone(), Default::default())) .collect(), } } /// Returns a client for `worker`, which must be a remote worker ID, first /// establishing a connection if there isn't one yet. - pub async fn connect(&self, worker: usize) -> &ExchangeClient { + pub async fn connect(&self, worker: usize, message_type: MessageType) -> &ExchangeClient { self.listener .get_or_init(|| async { if let Some(runtime) = self.runtime.upgrade() @@ -540,9 +586,20 @@ impl ExchangeClients { .iter() .find(|(host, _client)| host.workers.contains(&worker)) .unwrap(); - cell.get_or_init(|| ExchangeClient::new(host.address, &host.workers)) + cell[message_type] + .get_or_init(|| ExchangeClient::new(host.address, &host.workers)) .await } + + pub async fn wait(&self) { + for (_, clients) in &self.clients { + for client in clients.values() { + if let Some(client) = client.get() { + client.wait().await; + } + } + } + } } struct CallbackInner { @@ -703,6 +760,9 @@ pub(crate) struct Exchange { /// The number of bytes serialized. deserialized_bytes: AtomicUsize, + + /// When the exchange is active. + activity: ExchangeActivity, } // Stop Rust from complaining about unused field. @@ -763,6 +823,7 @@ where clients: Arc, exchange_id: ExchangeId, directory: &ExchangeDirectory, + activity: ExchangeActivity, ) -> Arc { let npeers = Runtime::num_workers(); let mailboxes: Vec>>> = @@ -788,6 +849,7 @@ where mailboxes, deserialization_usecs: AtomicU64::new(0), deserialized_bytes: AtomicUsize::new(0), + activity, }); directory.insert(exchange_id, exchange.clone()); @@ -806,39 +868,6 @@ where sender * self.npeers + receiver } - /// True if all `receiver`'s incoming mailboxes contain data. - /// - /// Once this function returns true, a subsequent `try_receive_all` - /// operation is guaranteed for `receiver`. - fn ready_to_receive(&self, receiver: usize) -> bool { - debug_assert!(receiver < self.npeers); - self.receiver_counters[receiver].load(Ordering::Acquire) == self.npeers - } - - /// Register callback to be invoked whenever the `ready_to_receive` - /// condition becomes true. - /// - /// The callback can be setup at most once (e.g., when a scheduler attaches - /// to the circuit) and cannot be unregistered. Notifications delivered - /// before the callback is registered are lost. The client should call - /// `ready_to_receive` after installing the callback to check - /// the status. - /// - /// After the callback has been registered, notifications are delivered with - /// at-least-once semantics: a notification is generated whenever the - /// status changes from not ready to ready, but spurious notifications - /// can occur occasionally. The user must check the status explicitly - /// by calling `ready_to_receive` or be prepared that `receive_all` - /// can fail. - pub(crate) fn register_receiver_callback(&self, receiver: usize, cb: F) - where - F: Fn() + Send + Sync + 'static, - { - debug_assert!(receiver < self.npeers); - - self.receiver_callbacks[receiver].set_callback(cb); - } - /// Locks and returns the mailbox for the sender/receiver pair. fn mailbox(&self, sender: usize, receiver: usize) -> MutexGuard<'_, Option>> { self.mailboxes[self.mailbox_index(sender, receiver)] @@ -849,7 +878,11 @@ where /// Create a new `Exchange` instance if an instance with the same id /// (created by another thread) does not yet exist within `runtime`. /// The number of peers will be set to `runtime.num_workers()`. - pub(crate) fn with_runtime(runtime: &Runtime, exchange_id: ExchangeId) -> Arc { + pub(crate) fn with_runtime( + runtime: &Runtime, + exchange_id: ExchangeId, + activity: ExchangeActivity, + ) -> Arc { // It's tempting to move the following calls to create the // `ExchangeDirectory` and `ExchangeClients` into `Exchange::new`, but // don't do it: all three of these access `runtime.local_store` and @@ -859,7 +892,7 @@ where runtime .local_store() .entry(ExchangeCacheId::new(exchange_id)) - .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory)) + .or_insert_with(|| Exchange::new(runtime, clients, exchange_id, &directory, activity)) .value() .clone() } @@ -987,16 +1020,19 @@ where }) .collect_vec(); + let message_type = MessageType::Synchronous(MessageSize::from_bytes( + items.iter().map(|fbuf| fbuf.len()).sum(), + )); + // We discard the return value that could allow us to wait // for the channel tx buffer to drain, because exchange is // synchronous, meaning that it will drain before we send // the next message. - let _ = self.clients.connect(receivers.start).await.send( - global_node_id.clone(), - self.exchange_id, - sender, - items, - ); + let _ = self + .clients + .connect(receivers.start, message_type) + .await + .send(global_node_id.clone(), self.exchange_id, sender, items); } } } @@ -1004,7 +1040,16 @@ where /// Read all incoming messages for this worker, waiting for data to arrive /// as needed. - pub(crate) async fn receive_all(&self, deserialize: D) -> Vec + /// + /// When the data is ready, but before reading it, this method swaps + /// `start_wait_usecs` with 0 and returns the old value along with the data. + /// This allows the caller to obtain the waiting time incurred just after it + /// became ready. + pub(crate) async fn receive_all( + &self, + deserialize: D, + start_wait_usecs: Option<&AtomicU64>, + ) -> (Vec, Option) where D: Fn(AlignedVec) -> T, { @@ -1024,6 +1069,11 @@ where } } + let start_wait_usecs = start_wait_usecs.and_then(|v| { + let start_wait_usecs = v.swap(0, Ordering::Acquire); + (start_wait_usecs != 0).then_some(start_wait_usecs) + }); + let mut data = Vec::with_capacity(self.npeers); for sender in 0..self.npeers { let mailbox = self.mailbox(sender, receiver).take().unwrap(); @@ -1035,7 +1085,7 @@ where } } - data + (data, start_wait_usecs) } } @@ -1062,6 +1112,19 @@ where } } +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +enum Phase { + Active, + Flush, + Commit, +} + +impl Phase { + fn is_inactive(&self, activity: ExchangeActivity) -> bool { + *self == Phase::Commit && activity == ExchangeActivity::InputOnly + } +} + /// Operator that partitions incoming data across all workers. /// /// This operator works in tandem with [`ExchangeReceiver`], which reassembles @@ -1100,13 +1163,6 @@ where /// ExchangeSender ExchangeReceiver /// ``` /// -/// `ExchangeSender` is an asynchronous operator., i.e., -/// [`ExchangeSender::is_async`] returns `true`. It becomes schedulable -/// ([`ExchangeSender::ready`] returns `true`) once all peers have retrieved -/// values written by the operator in the previous clock cycle. The scheduler -/// should use [`ExchangeSender::register_ready_callback`] to get notified when -/// the operator becomes schedulable. -/// /// `ExchangeSender` doesn't have a public constructor and must be instantiated /// using the [`new_exchange_operators`] function, which creates an /// [`ExchangeSender`]/[`ExchangeReceiver`] pair of operators and connects them @@ -1182,7 +1238,7 @@ where /// use dbsp::{ /// operator::{communication::new_exchange_operators, Generator}, /// circuit::{WorkerLocation, WorkerLocations}, -/// operator::communication::Mailbox, +/// operator::communication::{ExchangeActivity, Mailbox}, /// Circuit, RootCircuit, Runtime, /// storage::file::to_bytes_dyn, /// trace::aligned_deserialize, @@ -1218,6 +1274,7 @@ where /// }, /// |data| aligned_deserialize(&data[..]),/// // Reassemble received values into a vector. /// |v: &mut Vec, n| v.push(n), +/// ExchangeActivity::AllSteps, /// ).unwrap(); /// /// // Add exchange operators to the circuit. @@ -1259,7 +1316,7 @@ where // Input batch sizes. input_batch_stats: BatchSizeStats, - flushed: bool, + phase: Phase, // The instant when the sender produced its outputs, and the // receiver starts waiting for all other workers to produce their @@ -1286,7 +1343,7 @@ where outputs: Vec::with_capacity(Runtime::num_workers()), exchange, input_batch_stats: BatchSizeStats::new(), - flushed: false, + phase: Phase::Active, start_wait_usecs, phantom: PhantomData, } @@ -1324,8 +1381,12 @@ where true } + fn start_transaction(&mut self) { + self.phase = Phase::Active; + } + fn flush(&mut self) { - self.flushed = true; + self.phase = Phase::Flush; } } @@ -1340,6 +1401,16 @@ where } async fn eval_owned(&mut self, input: D) { + if self.phase.is_inactive(self.exchange.activity) { + return; + }; + let flushed = if self.phase == Phase::Flush { + self.phase = Phase::Commit; + true + } else { + false + }; + self.input_batch_stats.add_batch(input.num_entries_deep()); debug_assert!(self.ready()); @@ -1350,16 +1421,14 @@ where let data = self.outputs.drain(..).map(|mailbox| match mailbox { Mailbox::Tx(mut data) => { - data.push(self.flushed as u8); + data.push(flushed as u8); Mailbox::Tx(data) } Mailbox::Rx(_) => unreachable!(), - Mailbox::Plain(item) => Mailbox::Plain((item, self.flushed)), + Mailbox::Plain(item) => Mailbox::Plain((item, flushed)), }); self.exchange.send_all(&self.global_node_id, data).await; - - self.flushed = false; } fn input_preference(&self) -> OwnershipPreference { @@ -1375,18 +1444,10 @@ where /// peer. /// /// See [`ExchangeSender`] documentation for details. -/// -/// `ExchangeReceiver` is an asynchronous operator., i.e., -/// [`ExchangeReceiver::is_async`] returns `true`. It becomes schedulable -/// ([`ExchangeReceiver::ready`] returns `true`) once all peers have sent values -/// for this worker in the current clock cycle. The scheduler should use -/// [`ExchangeReceiver::register_ready_callback`] to get notified when the -/// operator becomes schedulable. pub struct ExchangeReceiver where T: Send + 'static + Clone, { - worker_index: usize, location: OperatorLocation, init: IF, deserialize: D, @@ -1396,6 +1457,7 @@ where flush_complete: bool, start_wait_usecs: Arc, total_wait_time: Arc, + phase: Phase, // Output batch sizes. output_batch_stats: BatchSizeStats, @@ -1406,7 +1468,6 @@ where T: Send + 'static + Clone + Debug, { pub(crate) fn new( - worker_index: usize, location: OperatorLocation, exchange: Arc>, init: IF, @@ -1414,10 +1475,7 @@ where deserialize: D, combine: L, ) -> Self { - debug_assert!(worker_index < Runtime::num_workers()); - Self { - worker_index, location, init, combine, @@ -1428,12 +1486,9 @@ where output_batch_stats: BatchSizeStats::new(), start_wait_usecs, total_wait_time: Arc::new(AtomicU64::new(0)), + phase: Phase::Active, } } - - fn is_ready(&self) -> bool { - self.exchange.ready_to_receive(self.worker_index) - } } impl Operator for ExchangeReceiver @@ -1464,50 +1519,12 @@ where }); } - fn is_async(&self) -> bool { + fn fixedpoint(&self, _scope: Scope) -> bool { true } - fn register_ready_callback(&mut self, cb: F) - where - F: Fn() + Send + Sync + 'static, - { - let start_wait_usecs = self.start_wait_usecs.clone(); - let total_wait_time = self.total_wait_time.clone(); - let exchange = self.exchange.clone(); - let worker_index = self.worker_index; - - let cb = move || { - if exchange.ready_to_receive(worker_index) { - // The callback can be invoked multiple times per step. - // Reset start_wait_usecs to 0 to make sure we don't double-count. - let start = start_wait_usecs.swap(0, Ordering::Acquire); - if start != 0 { - let end = current_time_usecs(); - if end > start { - let wait_time_usecs = end - start; - // if worker_index == 0 { - // info!( - // "{worker_index}: {} +{wait_time_usecs}", - // exchange.exchange_id() - // ); - // } - total_wait_time.fetch_add(wait_time_usecs, Ordering::AcqRel); - } - } - } - cb() - }; - self.exchange - .register_receiver_callback(self.worker_index, cb) - } - - fn ready(&self) -> bool { - self.is_ready() - } - - fn fixedpoint(&self, _scope: Scope) -> bool { - true + fn start_transaction(&mut self) { + self.phase = Phase::Active; } fn flush(&mut self) { @@ -1542,7 +1559,10 @@ where D: Fn(AlignedVec) -> T + Send + Sync + 'static, { async fn eval(&mut self) -> O { - debug_assert!(self.ready()); + if self.phase.is_inactive(self.exchange.activity) { + return (self.init)(); + } + let deserialize = |mut vec: AlignedVec| { let flushed = pop_flushed(&mut vec); let value = (self.deserialize)(vec); @@ -1550,7 +1570,16 @@ where }; let mut combined = (self.init)(); - let res = self.exchange.receive_all(deserialize).await; + let (res, start_wait_usecs) = self + .exchange + .receive_all(deserialize, Some(&self.start_wait_usecs)) + .await; + if let Some(start_wait_usecs) = start_wait_usecs { + self.total_wait_time.fetch_add( + current_time_usecs().saturating_sub(start_wait_usecs), + Ordering::Release, + ); + } for (data, flushed) in res { if flushed { self.flush_count += 1; @@ -1566,6 +1595,7 @@ where self.flush_complete = true; self.flush_count = 0; + self.phase = Phase::Commit; } self.output_batch_stats @@ -1588,6 +1618,28 @@ impl TypedMapKey for DirectoryId { type Value = ExchangeDirectory; } +/// The microsteps during which an exchange is active. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum ExchangeActivity { + /// The exchange is active in every microstep. + /// + /// This includes pre-commit and commit microstep. + AllSteps, + + /// The exchange is active only during pre-commit microsteps. + /// + /// This allows for optimizations for exchanges used for sharding data from + /// input operators, which don't exchange any data during commit. + /// + /// # Limitation + /// + /// The current implementation only works for operators that flush in the + /// same (micro)step in every worker. This is true for input operators, + /// which flush as soon as the transaction starts committing, but it is not + /// necessarily true for other operators. + InputOnly, +} + /// Create an [`ExchangeSender`]/[`ExchangeReceiver`] operator pair. /// /// See [`ExchangeSender`] documentation for details and example usage. @@ -1619,6 +1671,7 @@ pub fn new_exchange_operators( partition: PL, deserialize: D, combine: CL, + activity: ExchangeActivity, ) -> Option<(ExchangeSender, ExchangeReceiver)> where TO: Clone, @@ -1632,11 +1685,10 @@ where return None; } let runtime = Runtime::runtime().unwrap(); - let worker_index = Runtime::worker_index(); let exchange_id = runtime.sequence_next().try_into().unwrap(); let start_wait_usecs = Arc::new(AtomicU64::new(0)); - let exchange = Exchange::with_runtime(&runtime, exchange_id); + let exchange = Exchange::with_runtime(&runtime, exchange_id, activity); let sender = ExchangeSender::new( location, exchange.clone(), @@ -1644,7 +1696,6 @@ where partition, ); let receiver = ExchangeReceiver::new( - worker_index, location, exchange, init, @@ -1670,7 +1721,7 @@ mod tests { }, operator::{ Generator, - communication::{Mailbox, new_exchange_operators}, + communication::{ExchangeActivity, Mailbox, new_exchange_operators}, }, storage::file::{to_bytes, to_bytes_dyn}, trace::aligned_deserialize, @@ -1691,7 +1742,8 @@ mod tests { // value `(sender, n)` to each receiver, where `sender` is the sender's // worker number in round `n`. fn circuit() { - let exchange = Exchange::with_runtime(&Runtime::runtime().unwrap(), 0); + let exchange = + Exchange::with_runtime(&Runtime::runtime().unwrap(), 0, ExchangeActivity::AllSteps); TOKIO.block_on(async { let sender = Runtime::worker_index(); let n_workers = Runtime::num_workers(); @@ -1703,8 +1755,8 @@ mod tests { }) .await; - let received = exchange - .receive_all(|data| aligned_deserialize(&data[..])) + let (received, _) = exchange + .receive_all(|data| aligned_deserialize(&data[..]), None) .await; let expected = (0..n_workers).map(|worker| (worker, round)).collect_vec(); @@ -1821,6 +1873,7 @@ mod tests { }, |data| aligned_deserialize(&data[..]), |v: &mut Vec, n| v.push(n), + ExchangeActivity::AllSteps, ) .unwrap(); diff --git a/crates/dbsp/src/operator/dynamic.rs b/crates/dbsp/src/operator/dynamic.rs index 5bf472b5f1e..f35b4db6108 100644 --- a/crates/dbsp/src/operator/dynamic.rs +++ b/crates/dbsp/src/operator/dynamic.rs @@ -34,6 +34,7 @@ pub mod time_series; pub mod trace; pub(crate) mod upsert; +pub use communication::shard::Sharder; pub(crate) use communication::shard::shard_batch; /// The "standard" indexed Z-set type used by monomorphic diff --git a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs index 06b87e61ab6..ec2ca8fb3a4 100644 --- a/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs +++ b/crates/dbsp/src/operator/dynamic/balance/accumulate_trace_balanced.rs @@ -20,7 +20,7 @@ use crate::{ operator::{ Z1, async_stream_operators::{StreamingTernarySinkOperator, StreamingTernarySinkWrapper}, - communication::{Exchange, ExchangeReceiver}, + communication::{Exchange, ExchangeActivity, ExchangeReceiver}, dynamic::{ accumulate_trace::{AccumulateBoundsId, AccumulateTraceAppend, AccumulateZ1Trace}, balance::{ @@ -145,13 +145,12 @@ where // during the current transaction. It will be set to `true` exactly once per transaction. // Once all senders are flushed, the receiver can report itself as flushed too. let exchange: Arc> = - Exchange::with_runtime(&runtime, exchange_id); + Exchange::with_runtime(&runtime, exchange_id, ExchangeActivity::AllSteps); // Exchange receiver let batch_factories_clone = batch_factories.clone(); let receiver = circuit.add_source(ExchangeReceiver::new( - worker_index, Some(location), exchange.clone(), || Vec::new(), diff --git a/crates/dbsp/src/operator/dynamic/communication/gather.rs b/crates/dbsp/src/operator/dynamic/communication/gather.rs index ce2c962ae82..40eb252178b 100644 --- a/crates/dbsp/src/operator/dynamic/communication/gather.rs +++ b/crates/dbsp/src/operator/dynamic/communication/gather.rs @@ -8,7 +8,7 @@ use crate::{ runtime::{WorkerLocation, WorkerLocations}, }, circuit_cache_key, - operator::communication::{Mailbox, new_exchange_operators}, + operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators}, trace::{Batch, deserialize_indexed_wset, merge_batches, serialize_indexed_wset}, }; use arc_swap::ArcSwap; @@ -159,6 +159,7 @@ where }, move |data| deserialize_indexed_wset(&factories_clone, &data), |batches: &mut Vec, batch: B| batches.push(batch), + ExchangeActivity::AllSteps, ) .unwrap(); diff --git a/crates/dbsp/src/operator/dynamic/communication/shard.rs b/crates/dbsp/src/operator/dynamic/communication/shard.rs index 44f43d59c54..f1e212d6f71 100644 --- a/crates/dbsp/src/operator/dynamic/communication/shard.rs +++ b/crates/dbsp/src/operator/dynamic/communication/shard.rs @@ -16,7 +16,7 @@ use crate::{ }, circuit_cache_key, dynamic::{Data, DataTrait, DynPair, DynPairs, Factory}, - operator::communication::{Mailbox, new_exchange_operators}, + operator::communication::{ExchangeActivity, Mailbox, new_exchange_operators}, storage::file::SerializerInner, trace::{ Batch, BatchReader, Builder, IndexedWSetSerializer, deserialize_indexed_wset, merge_batches, @@ -35,68 +35,59 @@ fn all_workers() -> Range { 0..Runtime::num_workers() } -impl Stream +/// Builder for sharding a batch or pairs. +pub struct Sharder<'a, C, B> { + stream: &'a Stream, + workers: Range, + activity: ExchangeActivity, +} + +impl<'a, C, B> Sharder<'a, C, B> where C: Circuit, - IB: BatchReader