void Multiplexer::OnMixStreamBlock(
Connection& s, const StreamBlockHeader& header,
const MixStreamPtr& stream, const ByteBlockPtr& bytes) {
size_t sender_worker_rank = header.sender_rank * num_workers_per_host_ + header.sender_local_worker_id;
sLOG << "got block on" << s << "in MixStream" << header.stream_id << "from worker" << sender_worker_rank;
stream->OnStreamBlock(
sender_worker_rank,
Block(bytes, 0, header.size, header.first_item, header.num_items));
AsyncReadBlockHeader(s);
}
explicit Multiplexer(mem::Manager& mem_manager,
data::BlockPool& block_pool,
size_t num_workers_per_host, net::Group& group)
: mem_manager_(mem_manager),
block_pool_(block_pool),
dispatcher_(
mem_manager, group,
"host " + mem::to_string(group.my_host_rank()) + " multiplexer"),
group_(group),
num_workers_per_host_(num_workers_per_host),
stream_sets_(num_workers_per_host) {
for (size_t id = 0; id < group_.num_hosts(); id++) {
if (id == group_.my_host_rank()) continue;
AsyncReadBlockHeader(group_.connection(id));
}
(void)mem_manager_; // silence unused variable warning.
}
void Multiplexer::OnBlockHeader(Connection& s, net::Buffer&& buffer) {
// received invalid Buffer: the connection has closed?
if (!buffer.IsValid()) return;
StreamBlockHeader header;
net::BufferReader br(buffer);
header.ParseHeader(br);
// received stream id
StreamId id = header.stream_id;
size_t local_worker = header.receiver_local_worker_id;
size_t sender_worker_rank =
header.sender_rank * num_workers_per_host_ + header.sender_local_worker_id;
if (header.magic == MagicByte::CAT_STREAM_BLOCK)
{
CatStreamPtr stream = GetOrCreateCatStream(id, local_worker);
if (header.IsEnd()) {
sLOG << "end of stream on" << s << "in CatStream" << id
<< "from worker" << sender_worker_rank;
stream->OnCloseStream(sender_worker_rank);
AsyncReadBlockHeader(s);
}
else {
sLOG << "stream header from" << s << "on CatStream" << id
<< "from worker" << sender_worker_rank;
ByteBlockPtr bytes = ByteBlock::Allocate(header.size, block_pool_);
dispatcher_.AsyncRead(
s, bytes,
[this, header, stream, bytes](Connection& s) {
OnCatStreamBlock(s, header, stream, bytes);
});
}
}
else if (header.magic == MagicByte::MIX_STREAM_BLOCK)
{
MixStreamPtr stream = GetOrCreateMixStream(id, local_worker);
if (header.IsEnd()) {
sLOG << "end of stream on" << s << "in MixStream" << id
<< "from worker" << sender_worker_rank;
stream->OnCloseStream(sender_worker_rank);
AsyncReadBlockHeader(s);
}
else {
sLOG << "stream header from" << s << "on MixStream" << id
<< "from worker" << sender_worker_rank;
ByteBlockPtr bytes = ByteBlock::Allocate(header.size, block_pool_);
dispatcher_.AsyncRead(
s, bytes,
[this, header, stream, bytes](Connection& s) {
OnMixStreamBlock(s, header, stream, bytes);
});
}
}
else {
die("Invalid magic byte in BlockHeader");
}
}