[Bug](shuffle) fix mem leak in data stream sender (#16685)

This commit is contained in:
Gabriel
2023-02-14 16:40:13 +08:00
committed by GitHub
parent ea78184551
commit 784c27deeb
3 changed files with 25 additions and 6 deletions

View File

@ -240,7 +240,7 @@ Status ExchangeSinkBuffer::_send_rpc(InstanceLoId id) {
if (!q.empty()) {
// If we have data to shuffle which is not broadcasted
DO_RPC(q, block, nullptr)
DO_RPC(q, block.get(), nullptr)
} else if (!broadcast_q.empty()) {
// If we have data to shuffle which is broadcasted
DO_RPC(broadcast_q, block_holder->get_block(), request.block_holder)

View File

@ -38,7 +38,7 @@ namespace pipeline {
using InstanceLoId = int64_t;
struct TransmitInfo {
vectorized::PipChannel* channel;
PBlock* block;
std::unique_ptr<PBlock> block;
bool eos;
};

View File

@ -293,7 +293,7 @@ public:
bool is_local() const { return _is_local; }
void ch_roll_pb_block();
virtual void ch_roll_pb_block();
bool can_write() {
if (!is_local()) {
@ -392,14 +392,32 @@ Status VDataStreamSender::channel_add_rows(Channels& channels, int num_channels,
return Status::OK();
}
class PipChannel : public Channel {
class PipChannel final : public Channel {
public:
PipChannel(VDataStreamSender* parent, const RowDescriptor& row_desc,
const TNetworkAddress& brpc_dest, const TUniqueId& fragment_instance_id,
PlanNodeId dest_node_id, int buffer_size, bool is_transfer_chain,
bool send_query_statistics_with_every_batch)
: Channel(parent, row_desc, brpc_dest, fragment_instance_id, dest_node_id, buffer_size,
is_transfer_chain, send_query_statistics_with_every_batch) {}
is_transfer_chain, send_query_statistics_with_every_batch) {
ch_roll_pb_block();
}
~PipChannel() override {
if (_ch_cur_pb_block) {
delete _ch_cur_pb_block;
}
}
void ch_roll_pb_block() override {
// We have two choices here.
// 1. Use a PBlock pool and fetch an available PBlock if we need one. In this way, we can
// reuse the memory, but we have to use a lock to synchronize.
// 2. Create a new PBlock every time. In this way we don't need a lock but have to allocate
// new memory.
// Now we use the second way.
_ch_cur_pb_block = new PBlock();
}
// Asynchronously sends a block
// Returns the status of the most recently finished transmit_data
@ -414,7 +432,8 @@ public:
}
}
if (eos || block->column_metas_size()) {
RETURN_IF_ERROR(_buffer->add_block({this, block, eos}));
RETURN_IF_ERROR(_buffer->add_block(
{this, block ? std::make_unique<PBlock>(*block) : nullptr, eos}));
}
return Status::OK();
}