This is an automated email from the ASF dual-hosted git repository.
Mryange pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 11a8d6217cd [improvement](be) Share encoded protobuf blocks across
broadcast RPCs (#68961)
11a8d6217cd is described below
commit 11a8d6217cd912fe484586911df3e55458b1e75d
Author: HappenLee <[email protected]>
AuthorDate: Wed Oct 14 16:24:41 2026 +0800
[improvement](be) Share encoded protobuf blocks across broadcast RPCs
(#68961)
### What problem does this PR solve?
Related PR: #63258
Remote broadcast channels share the same `PBlock`, but brpc normally
serializes its large `column_values` into a separate request buffer for
every destination. Broadcasting an S-byte block to N destinations
therefore copies roughly N*S bytes while preparing the protobuf
requests.
Encode the `block` or repeated `blocks` field once in
`BroadcastPBlockHolder`, before publishing it to the destination queues.
Each RPC serializes its routing fields and appends references to the
cached encoding through `brpc::SerializedRequest`. The receiver
continues to parse the existing `PTransmitDataParams` wire format. Queue
order, packet sequences, EOS, and per-RPC callback ownership are
preserved.
The original PBlock remains available to the existing send paths. Queue
admission accounts for both it and the cache once per holder.
Single-destination broadcasts and oversized fragments retain the
original path; assembled requests retain required-field validation and
the protobuf size limit. The HTTP borrowed-data fix remains in #63258.
The earlier standalone prototype measured packing CPU at 1 MiB x 8
destinations dropping from 3370.81 us to 402.43 us. This is motivation
for the change, not an end-to-end measurement of this production patch.
Network bytes and receiver processing are unchanged.
### Release note
Reduce repeated protobuf serialization and payload copies when
broadcasting exchange blocks to multiple remote destinations.
### Check List (For Author)
- Test
- [x] Unit Test: `broadcast_protobuf_test.cpp` covers wire equivalence,
shared backing buffers, batching/EOS, the actual send branch, queue
accounting, size limits, and asynchronous RPC completion/cancellation.
- ASAN: `ENABLE_PCH=OFF ENABLE_UNITY_BUILD=ON ./run-be-ut.sh -j 48 --run
--filter='*BroadcastProtobufTest*:*ExchangeSinkOperatorTest*:*ExchangeSinkTest*'`
— 24 tests passed (14 new and 10 existing).
- `build-support/clang-format.sh` / `build-support/check-format.sh` with
clang-format 16.0.6; `build-support/check-build-hygiene.sh` and `git
diff --check` passed.
- `build-support/run-clang-tidy.sh`: all changed production lines and
the entire new test file passed. The local LLVM 20 package needed an
explicit resource-directory argument.
- Behavior changed:
- [x] No. SQL results and the exchange wire protocol are unchanged.
- Does this need documentation?
- [x] No.
### Check List (For Reviewer who merge this PR)
- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label
---
be/src/exec/operator/exchange_sink_buffer.cpp | 143 +++++--
be/src/exec/operator/exchange_sink_buffer.h | 19 +
be/src/exec/operator/exchange_sink_operator.cpp | 3 +
be/src/util/proto_util.h | 12 +
be/test/exec/operator/broadcast_protobuf_test.cpp | 464 ++++++++++++++++++++++
5 files changed, 618 insertions(+), 23 deletions(-)
diff --git a/be/src/exec/operator/exchange_sink_buffer.cpp
b/be/src/exec/operator/exchange_sink_buffer.cpp
index f6938a5eccd..4db029bfad5 100644
--- a/be/src/exec/operator/exchange_sink_buffer.cpp
+++ b/be/src/exec/operator/exchange_sink_buffer.cpp
@@ -18,6 +18,7 @@
#include "exec/operator/exchange_sink_buffer.h"
#include <brpc/controller.h>
+#include <brpc/serialized_request.h>
#include <butil/errno.h>
#include <butil/iobuf_inl.h>
#include <fmt/format.h>
@@ -28,6 +29,7 @@
#include <pdqsort.h>
#include <stddef.h>
+#include <algorithm>
#include <atomic>
#include <cstdint>
#include <exception>
@@ -49,6 +51,54 @@
namespace doris {
+namespace {
+
+void add_borrowed_broadcast_block(PTransmitDataParams& request, PBlock&
source) {
+ auto* block = request.add_blocks();
+ block->mutable_column_metas()->CopyFrom(source.column_metas());
+ block->set_be_exec_version(source.be_exec_version());
+ block->set_compressed(source.compressed());
+ block->set_compression_type(source.compression_type());
+ block->set_uncompressed_size(source.uncompressed_size());
+ block->set_allocated_column_values(source.mutable_column_values());
+}
+
+} // namespace
+
+Status BroadcastPBlockHolder::prepare_serialized_block(bool send_multi_blocks,
+ size_t num_receivers) {
+ DCHECK(_serialized_block.empty());
+ if (num_receivers <= 1 || _pblock->column_metas().empty()) {
+ return Status::OK();
+ }
+ PTransmitDataParams fragment;
+ if (send_multi_blocks) {
+ // Match the existing multi-block request's metadata and proto2 field
presence.
+ add_borrowed_broadcast_block(fragment, *_pblock);
+ } else {
+ fragment.set_allocated_block(_pblock.get());
+ }
+ Defer detach([&] {
+ if (send_multi_blocks) {
+
static_cast<void>(fragment.mutable_blocks(0)->release_column_values());
+ } else {
+ static_cast<void>(fragment.release_block());
+ }
+ });
+ // Protobuf cannot encode >= 2 GiB messages. Leave these on the existing
path,
+ // where the complete request is checked for HTTP attachment transport.
+ if (fragment.ByteSizeLong() >= MIN_HTTP_BRPC_SIZE) {
+ return Status::OK();
+ }
+ butil::IOBufAsZeroCopyOutputStream stream(&_serialized_block);
+ // Routing fields are intentionally absent from this reusable message
fragment.
+ // PBlock and its nested metadata contain no required fields.
+ if (!fragment.SerializePartialToZeroCopyStream(&stream)) {
+ return Status::InternalError("Failed to serialize broadcast block");
+ }
+ return Status::OK();
+}
+
BroadcastPBlockHolder::~BroadcastPBlockHolder() {
// lock the parent queue, if the queue could lock success, then return the
block
// to the queue, to reuse the block
@@ -64,8 +114,9 @@ void
BroadcastPBlockHolderMemLimiter::acquire(BroadcastPBlockHolder& holder) {
std::unique_lock l(_holders_lock);
DCHECK(_broadcast_dependency != nullptr);
holder.set_parent_creator(shared_from_this());
- auto size = holder._pblock->column_values().size();
- _total_queue_buffer_size += size;
+ holder._queue_charge_bytes =
+ holder._pblock->column_values().size() +
holder._serialized_block.size();
+ _total_queue_buffer_size += holder._queue_charge_bytes;
_total_queue_blocks_count++;
if (_total_queue_buffer_size >= _total_queue_buffer_size_limit ||
_total_queue_blocks_count >= _total_queue_blocks_count_limit) {
@@ -76,8 +127,7 @@ void
BroadcastPBlockHolderMemLimiter::acquire(BroadcastPBlockHolder& holder) {
void BroadcastPBlockHolderMemLimiter::release(const BroadcastPBlockHolder&
holder) {
std::unique_lock l(_holders_lock);
DCHECK(_broadcast_dependency != nullptr);
- auto size = holder._pblock->column_values().size();
- _total_queue_buffer_size -= size;
+ _total_queue_buffer_size -= holder._queue_charge_bytes;
_total_queue_blocks_count--;
if (_total_queue_buffer_size <= 0) {
_broadcast_dependency->set_ready();
@@ -418,34 +468,34 @@ Status ExchangeSinkBuffer::_send_rpc(RpcInstance&
instance_data) {
}
}
+ const bool use_shared_protobuf =
+ std::ranges::any_of(requests,
+ [](const auto& req) {
+ return
!req.block_holder->serialized_block().empty();
+ }) &&
+ std::ranges::all_of(requests, [](const auto& req) {
+ return !req.block_holder->serialized_block().empty() ||
+
req.block_holder->get_block()->column_metas().empty();
+ });
auto& request = requests[0];
auto& brpc_request = instance_data.request;
brpc_request->set_sender_id(channel->_parent->sender_id());
brpc_request->set_be_number(channel->_parent->be_number());
- if (_send_multi_blocks) {
- for (int i = 0; i < requests.size(); i++) {
- auto& req = requests[i];
+ if (_send_multi_blocks && !use_shared_protobuf) {
+ for (const auto& req : requests) {
if (auto block = req.block_holder->get_block();
block && !block->column_metas().empty()) {
- auto add_block = brpc_request->add_blocks();
- for (int j = 0; j < block->column_metas_size(); ++j) {
-
add_block->add_column_metas()->CopyFrom(block->column_metas(j));
- }
- add_block->set_be_exec_version(block->be_exec_version());
- add_block->set_compressed(block->compressed());
- add_block->set_compression_type(block->compression_type());
-
add_block->set_uncompressed_size(block->uncompressed_size());
-
add_block->set_allocated_column_values(block->mutable_column_values());
+ add_borrowed_broadcast_block(*brpc_request, *block);
}
}
- } else {
+ } else if (!_send_multi_blocks) {
if (request.block_holder->get_block() &&
!request.block_holder->get_block()->column_metas().empty()) {
brpc_request->set_allocated_block(request.block_holder->get_block());
}
}
- Defer release_block([&]() {
+ auto detach_blocks = [&]() {
if (!_send_multi_blocks && request.block_holder->get_block()) {
static_cast<void>(brpc_request->release_block());
} else {
@@ -454,7 +504,8 @@ Status ExchangeSinkBuffer::_send_rpc(RpcInstance&
instance_data) {
}
brpc_request->clear_blocks();
}
- });
+ };
+ Defer release_block(detach_blocks);
instance_data.seq += requests.size();
brpc_request->set_packet_seq(instance_data.seq);
brpc_request->set_eos(requests.back().eos);
@@ -510,15 +561,30 @@ Status ExchangeSinkBuffer::_send_rpc(RpcInstance&
instance_data) {
}
});
{
- auto send_remote_block_closure = AutoReleaseClosure<
- PTransmitDataParams,
-
ExchangeSendCallback<PTransmitDataResult>>::create_unique(brpc_request,
-
send_callback);
if (enable_http_send_block(*brpc_request)) {
+ auto send_remote_block_closure = AutoReleaseClosure<
+ PTransmitDataParams,
+
ExchangeSendCallback<PTransmitDataResult>>::create_unique(brpc_request,
+
send_callback);
RETURN_IF_ERROR(transmit_block_httpv2(_context->exec_env(),
std::move(send_remote_block_closure),
channel->_brpc_dest_addr));
+ } else if (use_shared_protobuf) {
+ // The HTTP decision above needs the complete singular
request. Detach its
+ // borrowed block before encoding only the per-RPC routing
fields.
+ detach_blocks();
+ auto serialized_request =
std::make_shared<brpc::SerializedRequest>();
+ RETURN_IF_ERROR(_serialize_broadcast_request(*brpc_request,
requests,
+
serialized_request.get()));
+ auto closure = AutoReleaseClosure<brpc::SerializedRequest,
+
ExchangeSendCallback<PTransmitDataResult>>::
+ create_unique(serialized_request, send_callback);
+ transmit_serialized_block(channel->_brpc_stub.get(),
std::move(closure));
} else {
+ auto send_remote_block_closure = AutoReleaseClosure<
+ PTransmitDataParams,
+
ExchangeSendCallback<PTransmitDataResult>>::create_unique(brpc_request,
+
send_callback);
transmit_blockv2(channel->_brpc_stub.get(),
std::move(send_remote_block_closure));
}
}
@@ -529,6 +595,37 @@ Status ExchangeSinkBuffer::_send_rpc(RpcInstance&
instance_data) {
return Status::OK();
}
+Status ExchangeSinkBuffer::_serialize_broadcast_request(
+ const PTransmitDataParams& header, const
std::vector<BroadcastTransmitInfo>& requests,
+ brpc::SerializedRequest* serialized_request) {
+ DCHECK(!header.has_block());
+ DCHECK_EQ(header.blocks_size(), 0);
+ if (!header.IsInitialized()) {
+ return Status::InternalError("Uninitialized broadcast request header:
{}",
+ header.InitializationErrorString());
+ }
+ size_t size = header.ByteSizeLong();
+ for (const auto& request : requests) {
+ size += request.block_holder->serialized_block().size();
+ }
+ // SerializedRequest bypasses protobuf's normal size check. Preserve that
limit
+ // even when a multi-block batch consists of individually encodable
fragments.
+ if (size >= MIN_HTTP_BRPC_SIZE) {
+ return Status::InternalError("Broadcast protobuf request exceeds 2
GiB: {} bytes", size);
+ }
+ auto& data = serialized_request->serialized_data();
+ {
+ butil::IOBufAsZeroCopyOutputStream stream(&data);
+ if (!header.SerializeToZeroCopyStream(&stream)) {
+ return Status::InternalError("Failed to serialize broadcast
request header");
+ }
+ }
+ for (const auto& request : requests) {
+ data.append(request.block_holder->serialized_block());
+ }
+ return Status::OK();
+}
+
void ExchangeSinkBuffer::_ended(RpcInstance& ins) {
std::unique_lock<std::mutex> lock(*ins.mutex);
ins.running_sink_count--;
diff --git a/be/src/exec/operator/exchange_sink_buffer.h
b/be/src/exec/operator/exchange_sink_buffer.h
index 7478c63f4e8..2942fb04f57 100644
--- a/be/src/exec/operator/exchange_sink_buffer.h
+++ b/be/src/exec/operator/exchange_sink_buffer.h
@@ -18,6 +18,7 @@
#pragma once
#include <brpc/controller.h>
+#include <butil/iobuf.h>
#include <gen_cpp/data.pb.h>
#include <gen_cpp/internal_service.pb.h>
#include <gen_cpp/types.pb.h>
@@ -38,6 +39,10 @@
#include "service/backend_options.h"
#include "util/brpc_closure.h"
+namespace brpc {
+class SerializedRequest;
+}
+
namespace doris {
class PTransmitDataParams;
class TUniqueId;
@@ -64,9 +69,18 @@ public:
void reset_block() { _pblock->Clear(); }
+ // Called once by the producer, before publishing the holder to any
channel.
+ Status prepare_serialized_block(bool send_multi_blocks, size_t
num_receivers);
+
+ const butil::IOBuf& serialized_block() const { return _serialized_block; }
+
private:
friend class BroadcastPBlockHolderMemLimiter;
std::unique_ptr<PBlock> _pblock;
+ // Includes the outer PTransmitDataParams.block/blocks field tag and
length.
+ // Keep _pblock for the existing protobuf and HTTP paths.
+ butil::IOBuf _serialized_block;
+ int64_t _queue_charge_bytes = 0;
std::weak_ptr<BroadcastPBlockHolderMemLimiter> _parent_creator;
void set_parent_creator(std::shared_ptr<BroadcastPBlockHolderMemLimiter>
parent_creator) {
_parent_creator = parent_creator;
@@ -271,6 +285,7 @@ public:
Status add_block(Channel* channel, TransmitInfo&& request);
Status add_block(Channel* channel, BroadcastTransmitInfo&& request);
void close();
+ bool send_multi_blocks() const { return _send_multi_blocks; }
void update_rpc_time(RpcInstance& ins, int64_t start_rpc_time, int64_t
receive_rpc_time);
void update_profile(RuntimeProfile* profile);
@@ -308,6 +323,10 @@ private:
Status _send_rpc(RpcInstance& ins);
+ static Status _serialize_broadcast_request(const PTransmitDataParams&
header,
+ const
std::vector<BroadcastTransmitInfo>& requests,
+ brpc::SerializedRequest*
serialized_request);
+
#ifndef BE_TEST
inline void _ended(RpcInstance& ins);
inline void _failed(InstanceLoId id, const std::string& err);
diff --git a/be/src/exec/operator/exchange_sink_operator.cpp
b/be/src/exec/operator/exchange_sink_operator.cpp
index a4b3e682ebd..c4c6b6be830 100644
--- a/be/src/exec/operator/exchange_sink_operator.cpp
+++ b/be/src/exec/operator/exchange_sink_operator.cpp
@@ -516,6 +516,9 @@ Status ExchangeSinkOperatorX::sink_impl(RuntimeState*
state, Block* block, bool
block_holder->reset_block();
}
+ RETURN_IF_ERROR(block_holder->prepare_serialized_block(
+ local_state._sink_buffer->send_multi_blocks(),
+ local_state._rpc_channels_num));
local_state._broadcast_pb_mem_limiter->acquire(*block_holder);
size_t idx = 0;
diff --git a/be/src/util/proto_util.h b/be/src/util/proto_util.h
index aec345688a9..a03024f4684 100644
--- a/be/src/util/proto_util.h
+++ b/be/src/util/proto_util.h
@@ -66,6 +66,18 @@ void transmit_blockv2(PBackendService_Stub* stub,
std::unique_ptr<Closure> closu
closure.release();
}
+template <typename Closure>
+void transmit_serialized_block(PBackendService_Stub* stub,
std::unique_ptr<Closure> closure) {
+ // SerializedRequest bypasses brpc's protobuf serializer, including
request compression.
+ // Exchange compresses column_values before preparing the shared protobuf
fragment.
+ DCHECK_EQ(closure->cntl_->request_compress_type(),
brpc::COMPRESS_TYPE_NONE);
+ static const auto* method =
PBackendService::descriptor()->FindMethodByName("transmit_block");
+ closure->cntl_->http_request().Clear();
+ stub->channel()->CallMethod(method, closure->cntl_.get(),
closure->request_.get(),
+ closure->response_.get(), closure.get());
+ closure.release();
+}
+
template <typename Closure>
Status transmit_block_httpv2(ExecEnv* exec_env, std::unique_ptr<Closure>
closure,
TNetworkAddress brpc_dest_addr) {
diff --git a/be/test/exec/operator/broadcast_protobuf_test.cpp
b/be/test/exec/operator/broadcast_protobuf_test.cpp
new file mode 100644
index 00000000000..da593d5ea38
--- /dev/null
+++ b/be/test/exec/operator/broadcast_protobuf_test.cpp
@@ -0,0 +1,464 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <brpc/channel.h>
+#include <brpc/closure_guard.h>
+#include <brpc/protocol.h>
+#include <brpc/serialized_request.h>
+#include <brpc/server.h>
+#include <bthread/countdown_event.h>
+#include <gtest/gtest.h>
+
+#include <atomic>
+#include <memory>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "exec/exchange/vdata_stream_sender.h"
+#include "exec/operator/exchange_sink_buffer.h"
+#include "exec/operator/exchange_sink_operator.h"
+#include "exec/pipeline/dependency.h"
+#include "runtime/thread_context.h"
+#include "testutil/mock/mock_runtime_state.h"
+#include "util/debug/leakcheck_disabler.h"
+#include "util/defer_op.h"
+#include "util/proto_util.h"
+
+namespace doris {
+namespace {
+
+PTransmitDataParams make_header(int64_t destination = 1, int64_t sequence = 1,
bool eos = false) {
+ PTransmitDataParams request;
+ request.mutable_finst_id()->set_hi(11);
+ request.mutable_finst_id()->set_lo(destination);
+ request.mutable_query_id()->set_hi(21);
+ request.mutable_query_id()->set_lo(22);
+ request.set_node_id(3);
+ request.set_sender_id(4);
+ request.set_be_number(5);
+ request.set_packet_seq(sequence);
+ request.set_eos(eos);
+ request.mutable_query_statistics()->set_scan_rows(42);
+ request.mutable_exec_status()->set_status_code(0);
+ request.set_transfer_by_attachment(false);
+ return request;
+}
+
+std::shared_ptr<BroadcastPBlockHolder> make_holder(char value = 'a') {
+ auto holder = BroadcastPBlockHolder::create_shared();
+ auto* block = holder->get_block();
+ auto* meta = block->add_column_metas();
+ meta->set_name("values");
+ meta->set_type(PGenericType::STRING);
+ meta->set_is_nullable(true);
+ block->set_column_values(std::string(64 * 1024, value));
+ block->set_compressed(true);
+ block->set_compression_type(segment_v2::LZ4);
+ block->set_uncompressed_size(128 * 1024);
+ block->set_be_exec_version(7);
+ return holder;
+}
+
+PTransmitDataParams parse(const butil::IOBuf& data) {
+ PTransmitDataParams parsed;
+ butil::IOBufAsZeroCopyInputStream stream(data);
+ EXPECT_TRUE(parsed.ParseFromZeroCopyStream(&stream));
+ return parsed;
+}
+
+class BroadcastProtobufTest : public testing::TestWithParam<bool> {};
+
+class CapturingRpcChannel : public google::protobuf::RpcChannel {
+public:
+ ~CapturingRpcChannel() override {
+ if (completion != nullptr) {
+ completion->Run();
+ }
+ }
+
+ butil::IOBuf wire;
+ google::protobuf::Closure* completion = nullptr;
+
+ void CallMethod(const google::protobuf::MethodDescriptor* method,
+ google::protobuf::RpcController* controller,
+ const google::protobuf::Message* request,
google::protobuf::Message* response,
+ google::protobuf::Closure* done) override {
+ EXPECT_EQ(method,
PBackendService::descriptor()->FindMethodByName("transmit_block"));
+ EXPECT_EQ(request->GetDescriptor(),
brpc::SerializedRequest::descriptor());
+ brpc::SerializeRequestDefault(&wire,
static_cast<brpc::Controller*>(controller), request);
+ completion = done;
+ }
+};
+
+void expect_broadcast_packet(const PTransmitDataParams& packet, const
RpcInstance& instance,
+ bool multi, int sent, bool eos) {
+ EXPECT_EQ(packet.packet_seq(), multi ? 3 : sent);
+ EXPECT_EQ(packet.eos(), eos);
+ EXPECT_EQ(packet.finst_id().lo(), instance.id);
+ ASSERT_EQ(packet.blocks_size(), multi ? 2 : 0);
+ const auto& block = multi ? packet.blocks(0) : packet.block();
+ EXPECT_EQ(block.column_values(), std::string(64 * 1024, multi ? 'a' : 'a'
+ sent - 1));
+ EXPECT_FALSE(instance.request->has_block());
+ EXPECT_EQ(instance.request->blocks_size(), 0);
+}
+
+TEST_P(BroadcastProtobufTest,
SendRpcPreservesBatchSequenceEosAndRequestTemplate) {
+ auto state = std::make_shared<MockRuntimeState>();
+ ExchangeSinkBuffer buffer(state.get(), 1);
+ buffer._query_id = make_header().query_id();
+ buffer._dest_node_id = 3;
+ buffer._send_multi_blocks = GetParam();
+ buffer._send_multi_blocks_byte_size = 1024 * 1024;
+ TUniqueId destination;
+ destination.__set_hi(11);
+ destination.__set_lo(1);
+ buffer.construct_request(destination);
+ auto sink = ExchangeSinkLocalState::create_shared(state.get());
+ CapturingRpcChannel transport;
+ Channel channel(sink.get(), TNetworkAddress {}, destination, 0);
+ channel._brpc_stub = std::make_shared<PBackendService_Stub>(&transport);
+ auto& instance = *buffer._rpc_instances.at(destination.lo);
+ auto& queue = instance.broadcast_package_queue[&channel];
+ for (int i = 0; i < 2; ++i) {
+ auto holder = make_holder('a' + i);
+ ASSERT_TRUE(holder->prepare_serialized_block(GetParam(), 2).ok());
+ queue.push({.block_holder = holder, .eos = !GetParam() && i == 1});
+ }
+ if (GetParam()) {
+ queue.push({.block_holder = BroadcastPBlockHolder::create_shared(),
.eos = true});
+ }
+ int sent = 0;
+ while (!queue.empty()) {
+ ASSERT_TRUE(buffer._send_rpc(instance).ok());
+ ASSERT_NE(transport.completion, nullptr);
+ ++sent;
+ expect_broadcast_packet(parse(transport.wire), instance, GetParam(),
sent, queue.empty());
+ // We drive the next send explicitly; expire the callback before
completing this RPC.
+ channel._send_callback.reset();
+ transport.completion->Run();
+ transport.completion = nullptr;
+ transport.wire.clear();
+ }
+}
+
+void expect_shared_backing_blocks(const PTransmitDataParams& header,
+ const brpc::SerializedRequest& request,
+ const BroadcastPBlockHolder& holder) {
+ brpc::Controller controller;
+ butil::IOBuf wire;
+ brpc::SerializeRequestDefault(&wire, &controller, &request);
+ ASSERT_FALSE(controller.Failed());
+ wire.pop_front(header.ByteSizeLong());
+ ASSERT_EQ(wire.backing_block_num(),
holder.serialized_block().backing_block_num());
+ for (size_t i = 0; i < wire.backing_block_num(); ++i) {
+ EXPECT_EQ(wire.backing_block(i).data(),
holder.serialized_block().backing_block(i).data());
+ EXPECT_EQ(wire.backing_block(i).size(),
holder.serialized_block().backing_block(i).size());
+ }
+}
+
+TEST_P(BroadcastProtobufTest,
MatchesExistingWireMessageAndSharesBackingBlocks) {
+ auto holder = make_holder();
+ const auto original = holder->get_block()->SerializeAsString();
+ ASSERT_TRUE(holder->prepare_serialized_block(GetParam(), 8).ok());
+ ASSERT_FALSE(holder->serialized_block().empty());
+
+ for (int destination = 1; destination <= 8; ++destination) {
+ auto header = make_header(destination, destination + 100, destination
== 8);
+ brpc::SerializedRequest request;
+ ASSERT_TRUE(ExchangeSinkBuffer::_serialize_broadcast_request(
+ header, {{.block_holder = holder, .eos = false}},
&request)
+ .ok());
+ auto expected = header;
+ auto* block = GetParam() ? expected.add_blocks() :
expected.mutable_block();
+ block->CopyFrom(*holder->get_block());
+ EXPECT_EQ(parse(request.serialized_data()).SerializeAsString(),
+ expected.SerializeAsString());
+
+ // Exercise the actual brpc serializer: it must retain the same
backing data,
+ // rather than encoding the payload again for this destination.
+ expect_shared_backing_blocks(header, request, *holder);
+ }
+ EXPECT_EQ(holder->get_block()->SerializeAsString(), original);
+}
+
+TEST_P(BroadcastProtobufTest, ConcurrentDestinationsAndHolderRelease) {
+ auto holder = make_holder();
+ ASSERT_TRUE(holder->prepare_serialized_block(GetParam(), 8).ok());
+ std::weak_ptr<BroadcastPBlockHolder> weak_holder = holder;
+ std::vector<brpc::SerializedRequest> requests(8);
+ std::vector<std::thread> senders;
+ for (int i = 0; i < 8; ++i) {
+ senders.emplace_back([&, i, shared_holder = holder] {
+ SCOPED_INIT_THREAD_CONTEXT();
+ EXPECT_TRUE(ExchangeSinkBuffer::_serialize_broadcast_request(
+ make_header(i + 1, i + 10),
+ {{.block_holder = shared_holder, .eos =
false}}, &requests[i])
+ .ok());
+ });
+ }
+ for (auto& sender : senders) {
+ sender.join();
+ }
+ holder.reset();
+ ASSERT_TRUE(weak_holder.expired());
+ for (int i = 0; i < 8; ++i) {
+ auto parsed = parse(requests[i].serialized_data());
+ EXPECT_EQ(parsed.finst_id().lo(), i + 1);
+ EXPECT_EQ(parsed.packet_seq(), i + 10);
+ const auto& block = GetParam() ? parsed.blocks(0) : parsed.block();
+ EXPECT_EQ(block.column_values(), std::string(64 * 1024, 'a'));
+ }
+}
+
+TEST(BroadcastProtobufTest, RepeatedBlockPreservesLegacyDefaultFieldPresence) {
+ auto holder = make_holder();
+ holder->get_block()->clear_compressed();
+ holder->get_block()->clear_uncompressed_size();
+ holder->get_block()->clear_compression_type();
+ ASSERT_TRUE(holder->prepare_serialized_block(true, 2).ok());
+ brpc::SerializedRequest request;
+ ASSERT_TRUE(ExchangeSinkBuffer::_serialize_broadcast_request(
+ make_header(), {{.block_holder = holder, .eos =
false}}, &request)
+ .ok());
+ const auto parsed = parse(request.serialized_data());
+ ASSERT_EQ(parsed.blocks_size(), 1);
+ EXPECT_TRUE(parsed.blocks(0).has_compressed());
+ EXPECT_FALSE(parsed.blocks(0).compressed());
+ EXPECT_TRUE(parsed.blocks(0).has_uncompressed_size());
+ EXPECT_EQ(parsed.blocks(0).uncompressed_size(),
holder->get_block()->uncompressed_size());
+ EXPECT_TRUE(parsed.blocks(0).has_compression_type());
+ EXPECT_EQ(parsed.blocks(0).compression_type(),
holder->get_block()->compression_type());
+ EXPECT_FALSE(holder->get_block()->has_compressed());
+ EXPECT_FALSE(holder->get_block()->has_uncompressed_size());
+ EXPECT_FALSE(holder->get_block()->has_compression_type());
+}
+
+TEST(BroadcastProtobufTest, DifferentBatchBoundariesAndEmptyEos) {
+ auto first = make_holder('a');
+ auto second = make_holder('b');
+ auto eos = BroadcastPBlockHolder::create_shared();
+ for (const auto& holder : {first, second, eos}) {
+ ASSERT_TRUE(holder->prepare_serialized_block(true, 2).ok());
+ }
+ ASSERT_TRUE(eos->serialized_block().empty());
+
+ const std::vector<std::vector<BroadcastTransmitInfo>> batches = {
+ {{.block_holder = first, .eos = false},
+ {.block_holder = second, .eos = false},
+ {.block_holder = eos, .eos = true}},
+ {{.block_holder = first, .eos = false}},
+ {{.block_holder = second, .eos = false}, {.block_holder = eos,
.eos = true}},
+ {{.block_holder = eos, .eos = true}}};
+ const int64_t sequences[] = {3, 1, 3, 4};
+ const int expected_blocks[] = {2, 1, 1, 0};
+ for (size_t i = 0; i < batches.size(); ++i) {
+ brpc::SerializedRequest request;
+ auto header = make_header(i + 1, sequences[i], batches[i].back().eos);
+ ASSERT_TRUE(ExchangeSinkBuffer::_serialize_broadcast_request(header,
batches[i], &request)
+ .ok());
+ auto parsed = parse(request.serialized_data());
+ EXPECT_EQ(parsed.packet_seq(), sequences[i]);
+ EXPECT_EQ(parsed.eos(), batches[i].back().eos);
+ ASSERT_EQ(parsed.blocks_size(), expected_blocks[i]);
+ for (int j = 0; j < parsed.blocks_size(); ++j) {
+ EXPECT_EQ(parsed.blocks(j).SerializeAsString(),
+
batches[i][j].block_holder->get_block()->SerializeAsString());
+ }
+ }
+}
+
+TEST(BroadcastProtobufTest, SingleReceiverKeepsOriginalPath) {
+ for (bool multi : {false, true}) {
+ auto holder = make_holder();
+ const auto original = holder->get_block()->SerializeAsString();
+ ASSERT_TRUE(holder->prepare_serialized_block(multi, 1).ok());
+ EXPECT_TRUE(holder->serialized_block().empty());
+ EXPECT_EQ(holder->get_block()->SerializeAsString(), original);
+ }
+}
+
+TEST(BroadcastProtobufTest,
RejectsMissingRequiredFieldsBeforeSerializedRequestBypassesValidation) {
+ auto valid_holder = make_holder();
+ ASSERT_TRUE(valid_holder->prepare_serialized_block(false, 2).ok());
+ auto header = make_header();
+ header.clear_packet_seq();
+ brpc::SerializedRequest request;
+ EXPECT_FALSE(ExchangeSinkBuffer::_serialize_broadcast_request(
+ header, {{.block_holder = valid_holder, .eos =
false}}, &request)
+ .ok());
+ EXPECT_TRUE(request.serialized_data().empty());
+}
+
+TEST(BroadcastProtobufTest,
QueueAccountsSharedCacheOnceAndReleasesAdmissionCharge) {
+ auto dependency = Dependency::create_shared(1, 1, "broadcast", true);
+ auto limiter = BroadcastPBlockHolderMemLimiter::create_shared(dependency);
+ auto holder = make_holder();
+ ASSERT_TRUE(holder->prepare_serialized_block(false, 8).ok());
+ const auto charge =
+ holder->get_block()->column_values().size() +
holder->serialized_block().size();
+ limiter->_total_queue_buffer_size_limit = charge;
+ limiter->acquire(*holder);
+ EXPECT_EQ(limiter->_total_queue_buffer_size, charge);
+ EXPECT_EQ(limiter->_total_queue_blocks_count, 1);
+ EXPECT_EQ(dependency->is_blocked_by(nullptr), dependency.get());
+
+ std::vector<std::shared_ptr<BroadcastPBlockHolder>> channels(8, holder);
+ holder.reset();
+ EXPECT_EQ(limiter->_total_queue_buffer_size, charge);
+ // Release must subtract the admission charge even if the original HTTP
path
+ // moves column_values out of the PBlock before the holder is destroyed.
+ channels.front()->get_block()->clear_column_values();
+ channels.clear();
+ EXPECT_EQ(limiter->_total_queue_buffer_size, 0);
+ EXPECT_EQ(limiter->_total_queue_blocks_count, 0);
+ EXPECT_EQ(dependency->is_blocked_by(nullptr), nullptr);
+}
+
+TEST(BroadcastProtobufTest,
RejectsOversizedBatchBeforeBrpcBypassesProtobufSizeCheck) {
+ auto holder = BroadcastPBlockHolder::create_shared();
+ ASSERT_EQ(holder->_serialized_block.append(std::string(1024 * 1024, 'x')),
0);
+ // Grow references, not a 2 GiB allocation. Only the size check reads this
buffer.
+ while (holder->_serialized_block.size() < MIN_HTTP_BRPC_SIZE) {
+ auto copy = holder->_serialized_block;
+ holder->_serialized_block.append(copy);
+ }
+ brpc::SerializedRequest request;
+ auto status = ExchangeSinkBuffer::_serialize_broadcast_request(
+ make_header(), {{.block_holder = holder, .eos = false}}, &request);
+ EXPECT_FALSE(status.ok());
+ EXPECT_TRUE(request.serialized_data().empty());
+}
+
+class BroadcastService : public PBackendService {
+public:
+ bthread::CountdownEvent release {1};
+ std::atomic<int> received {0};
+
+ void release_requests() {
+ if (!_released) {
+ release.signal();
+ _released = true;
+ }
+ }
+
+ void transmit_block(google::protobuf::RpcController* controller,
+ const PTransmitDataParams* request,
PTransmitDataResult* response,
+ google::protobuf::Closure* done) override {
+ SCOPED_INIT_THREAD_CONTEXT();
+ brpc::ClosureGuard guard(done);
+ release.wait();
+ const auto& block = request->has_block() ? request->block() :
request->blocks(0);
+ EXPECT_EQ(block.column_values(), std::string(64 * 1024, 'a'));
+ EXPECT_TRUE(block.compressed());
+ EXPECT_EQ(request->packet_seq(), request->finst_id().lo() + 10);
+ received.fetch_add(1);
+ if (request->finst_id().lo() == 1) {
+ controller->SetFailed("injected transport failure");
+ } else if (request->finst_id().lo() == 2) {
+ Status::EndOfFile("receiver
finished").to_protobuf(response->mutable_status());
+ } else {
+ Status::OK().to_protobuf(response->mutable_status());
+ }
+ }
+
+private:
+ bool _released = false;
+};
+
+auto make_counting_callback(std::atomic<int>& succeeded, std::atomic<int>&
failed,
+ std::atomic<int>& eof) {
+ auto callback = ExchangeSendCallback<PTransmitDataResult>::create_shared();
+ callback->init(nullptr, false);
+ callback->start_rpc_time = 0;
+ callback->addFailedHandler([&](auto*, const auto&) { failed.fetch_add(1);
});
+ callback->addSuccessHandler([&](auto*, const auto&, const auto& result,
const auto&) {
+ if (Status::create(result.status()).template
is<ErrorCode::END_OF_FILE>()) {
+ eof.fetch_add(1);
+ } else {
+ EXPECT_TRUE(Status::create(result.status()).ok());
+ succeeded.fetch_add(1);
+ }
+ });
+ return callback;
+}
+
+TEST_P(BroadcastProtobufTest,
AsyncRpcPreservesPayloadAndReleasesCancelledCallbackRequest) {
+ BroadcastService service;
+ brpc::Server server;
+ ASSERT_EQ(server.AddService(&service, brpc::SERVER_DOESNT_OWN_SERVICE), 0);
+ {
+ // brpc registers process-lifetime built-in services on first startup.
+ debug::ScopedLeakCheckDisabler disable_lsan;
+ ASSERT_EQ(server.Start("127.0.0.1:0", nullptr), 0);
+ }
+ Defer stop([&] {
+ server.Stop(0);
+ server.Join();
+ });
+ Defer unblock([&] { service.release_requests(); });
+ brpc::Channel channel;
+ brpc::ChannelOptions options;
+ options.max_retry = 0;
+ options.timeout_ms = 10000;
+ ASSERT_EQ(channel.Init(server.listen_address(), &options), 0);
+ PBackendService_Stub stub(&channel);
+ std::atomic<int> succeeded {0};
+ std::atomic<int> failed {0};
+ std::atomic<int> eof {0};
+ std::vector<std::shared_ptr<ExchangeSendCallback<PTransmitDataResult>>>
callbacks;
+ std::vector<brpc::CallId> ids;
+ std::vector<std::weak_ptr<brpc::SerializedRequest>> weak_requests;
+ auto holder = make_holder();
+ ASSERT_TRUE(holder->prepare_serialized_block(GetParam(), 4).ok());
+ for (int i = 1; i <= 4; ++i) {
+ auto request = std::make_shared<brpc::SerializedRequest>();
+ ASSERT_TRUE(ExchangeSinkBuffer::_serialize_broadcast_request(
+ make_header(i, i + 10), {{.block_holder = holder,
.eos = false}},
+ request.get())
+ .ok());
+ auto callback = make_counting_callback(succeeded, failed, eof);
+ ids.push_back(callback->cntl_->call_id());
+ weak_requests.push_back(request);
+ auto closure = AutoReleaseClosure<
+ brpc::SerializedRequest,
+
ExchangeSendCallback<PTransmitDataResult>>::create_unique(request, callback);
+ transmit_serialized_block(&stub, std::move(closure));
+ // Simulate query cancellation before the fourth destination replies.
+ if (i != 4) {
+ callbacks.push_back(std::move(callback));
+ }
+ }
+ // No queue/producer owns the source block while these RPCs are
outstanding.
+ holder.reset();
+ service.release_requests();
+ for (auto id : ids) {
+ brpc::Join(id);
+ }
+ EXPECT_EQ(service.received, 4);
+ EXPECT_EQ(succeeded, 1);
+ EXPECT_EQ(failed, 1);
+ EXPECT_EQ(eof, 1);
+ for (const auto& request : weak_requests) {
+ EXPECT_TRUE(request.expired());
+ }
+}
+
+INSTANTIATE_TEST_SUITE_P(BlockField, BroadcastProtobufTest, testing::Bool());
+
+} // namespace
+} // namespace doris
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
