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");
}
}
void Multiplexer::OnMultiplexerHeader(
size_t peer, uint32_t seq, Connection& s, net::Buffer&& buffer) {
die_unless(d_->ongoing_requests_[peer] > 0);
d_->ongoing_requests_[peer]--;
// received invalid Buffer: the connection has closed?
if (!buffer.IsValid()) return;
net::BufferReader br(buffer);
StreamMultiplexerHeader header = StreamMultiplexerHeader::Parse(br);
LOG << "OnMultiplexerHeader() header"
<< " magic=" << unsigned(header.magic)
<< " size=" << header.size
<< " num_items=" << header.num_items
<< " first_item=" << header.first_item
<< " typecode_verify=" << header.typecode_verify
<< " stream_id=" << header.stream_id;
// received stream id
StreamId id = header.stream_id;
size_t local_worker = header.receiver_local_worker;
// round of allocation size to next power of two
size_t alloc_size = header.size;
if (alloc_size < THRILL_DEFAULT_ALIGN) alloc_size = THRILL_DEFAULT_ALIGN;
alloc_size = tlx::round_up_to_power_of_two(alloc_size);
if (header.magic == MagicByte::CatStreamBlock)
{
CatStreamDataPtr stream = GetOrCreateCatStreamData(
id, local_worker, /* dia_id (unknown at this time) */ 0);
stream->rx_net_bytes_ += buffer.size();
if (header.IsEnd()) {
sLOG << "end of stream on" << s << "in CatStream" << id
<< "from worker" << header.sender_worker;
stream->OnStreamBlock(
header.sender_worker, header.seq, Block());
}
else {
sLOG << "stream header from" << s << "on CatStream" << id
<< "from worker" << header.sender_worker
<< "for local_worker" << local_worker
<< "seq" << header.seq
<< "size" << header.size;
PinnedByteBlockPtr bytes = block_pool_.AllocateByteBlock(
alloc_size, local_worker);
sLOG << "new PinnedByteBlockPtr bytes=" << *bytes;
d_->ongoing_requests_[peer]++;
dispatcher_.AsyncRead(
s, seq + 1, header.size, std::move(bytes),
[this, peer, header, stream](
Connection& s, PinnedByteBlockPtr&& bytes) {
OnCatStreamBlock(peer, s, header, stream, std::move(bytes));
});
}
}
else if (header.magic == MagicByte::MixStreamBlock)
{
MixStreamDataPtr stream = GetOrCreateMixStreamData(
id, local_worker, /* dia_id (unknown at this time) */ 0);
stream->rx_net_bytes_ += buffer.size();
if (header.IsEnd()) {
sLOG << "end of stream on" << s << "in MixStream" << id
<< "from worker" << header.sender_worker;
stream->OnStreamBlock(header.sender_worker, header.seq, Block());
}
else {
sLOG << "stream header from" << s << "on MixStream" << id
<< "from worker" << header.sender_worker
<< "for local_worker" << local_worker
<< "seq" << header.seq
<< "size" << header.size;
PinnedByteBlockPtr bytes = block_pool_.AllocateByteBlock(
alloc_size, local_worker);
d_->ongoing_requests_[peer]++;
dispatcher_.AsyncRead(
s, seq + 1, header.size, std::move(bytes),
[this, peer, header, stream](
Connection& s, PinnedByteBlockPtr&& bytes) mutable {
OnMixStreamBlock(peer, s, header, stream, std::move(bytes));
});
}
}
else {
die("Invalid magic byte in MultiplexerHeader");
}
AsyncReadMultiplexerHeader(peer, s);
}
