Пример #1
0
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");
    }
}
Пример #2
0
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);
}