Skip to content
Merged
31 changes: 6 additions & 25 deletions crates/dbsp/src/circuit/circuit_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3939,31 +3939,12 @@ where
SndOp: SinkOperator<I>,
RcvOp: SourceOperator<O>,
{
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
}

Expand Down
8 changes: 5 additions & 3 deletions crates/dbsp/src/circuit/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
Expand Down
7 changes: 5 additions & 2 deletions crates/dbsp/src/operator/communication.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Loading
Loading