Skip to content

Commit e326ebb

Browse files
authored
[feature](pipelineX) control exchange sink by memory usage (#28814)
1 parent d42fd68 commit e326ebb

6 files changed

Lines changed: 38 additions & 7 deletions

File tree

‎be/src/pipeline/exec/exchange_sink_operator.cpp‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,17 +199,22 @@ Status ExchangeSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& inf
199199
_wait_channel_timer.resize(local_size);
200200
auto deps_for_channels = AndDependency::create_shared(
201201
_parent->operator_id(), _parent->node_id(), state->get_query_ctx());
202-
for (auto channel : channels) {
202+
auto deps_for_channels_mem_limit = AndDependency::create_shared(
203+
_parent->operator_id(), _parent->node_id(), state->get_query_ctx());
204+
for (auto* channel : channels) {
203205
if (channel->is_local()) {
204206
_local_channels_dependency[dep_id] = channel->get_local_channel_dependency();
205207
DCHECK(_local_channels_dependency[dep_id] != nullptr);
206208
deps_for_channels->add_child(_local_channels_dependency[dep_id]);
207209
_wait_channel_timer[dep_id] = ADD_CHILD_TIMER(
208210
_profile, fmt::format("WaitForLocalExchangeBuffer{}", dep_id), timer_name);
211+
auto local_recvr = channel->local_recvr();
212+
deps_for_channels_mem_limit->add_child(local_recvr->get_mem_limit_dependency());
209213
dep_id++;
210214
}
211215
}
212216
_exchange_sink_dependency->add_child(deps_for_channels);
217+
_exchange_sink_dependency->add_child(deps_for_channels_mem_limit);
213218
}
214219
if (p._part_type == TPartitionType::HASH_PARTITIONED) {
215220
_partition_count = channels.size();

‎be/src/pipeline/exec/exchange_sink_operator.h‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,9 +96,14 @@ class LocalExchangeChannelDependency final : public Dependency {
9696
LocalExchangeChannelDependency(int id, int node_id, QueryContext* query_ctx)
9797
: Dependency(id, node_id, "LocalExchangeChannelDependency", true, query_ctx) {}
9898
~LocalExchangeChannelDependency() override = default;
99-
// TODO(gabriel): blocked by memory
10099
};
101100

101+
class LocalExchangeMemLimitDependency final : public Dependency {
102+
ENABLE_FACTORY_CREATOR(LocalExchangeMemLimitDependency);
103+
LocalExchangeMemLimitDependency(int id, int node_id, QueryContext* query_ctx)
104+
: Dependency(id, node_id, "LocalExchangeMemLimitDependency", true, query_ctx) {}
105+
~LocalExchangeMemLimitDependency() override = default;
106+
};
102107
class ExchangeSinkLocalState final : public PipelineXSinkLocalState<AndDependency> {
103108
ENABLE_FACTORY_CREATOR(ExchangeSinkLocalState);
104109
using Base = PipelineXSinkLocalState<AndDependency>;

‎be/src/pipeline/exec/exchange_source_operator.cpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ Status ExchangeLocalState::init(RuntimeState* state, LocalStateInfo& info) {
7474
stream_recvr = state->exec_env()->vstream_mgr()->create_recvr(
7575
state, p.input_row_desc(), state->fragment_instance_id(), p.node_id(), p.num_senders(),
7676
profile(), p.is_merging(), p.sub_plan_query_statistics_recvr());
77+
stream_recvr->create_mem_limit_dependency(p.operator_id(), p.node_id(), state->get_query_ctx());
7778
auto* source_dependency = _dependency;
7879
const auto& queues = stream_recvr->sender_queues();
7980
deps.resize(queues.size());

‎be/src/vec/runtime/vdata_stream_recvr.cpp‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -354,8 +354,7 @@ VDataStreamRecvr::VDataStreamRecvr(
354354
_profile(profile),
355355
_peak_memory_usage_counter(nullptr),
356356
_sub_plan_query_statistics_recvr(sub_plan_query_statistics_recvr),
357-
_enable_pipeline(state->enable_pipeline_exec()),
358-
_mem_available(std::make_shared<bool>(true)) {
357+
_enable_pipeline(state->enable_pipeline_exec()) {
359358
// DataStreamRecvr may be destructed after the instance execution thread ends.
360359
_mem_tracker =
361360
std::make_unique<MemTracker>("VDataStreamRecvr:" + print_id(_fragment_instance_id));
@@ -506,12 +505,21 @@ void VDataStreamRecvr::update_blocks_memory_usage(int64_t size) {
506505
_blocks_memory_usage->add(size);
507506
auto val = _blocks_memory_usage_current_value.fetch_add(size);
508507
if (val + size > config::exchg_node_buffer_size_bytes) {
509-
*_mem_available = false;
508+
if (_exchange_sink_mem_limit_dependency) {
509+
_exchange_sink_mem_limit_dependency->block();
510+
}
510511
} else {
511-
*_mem_available = true;
512+
if (_exchange_sink_mem_limit_dependency) {
513+
_exchange_sink_mem_limit_dependency->set_ready();
514+
}
512515
}
513516
}
514517

518+
void VDataStreamRecvr::create_mem_limit_dependency(int id, int node_id, QueryContext* query_ctx) {
519+
_exchange_sink_mem_limit_dependency =
520+
pipeline::LocalExchangeMemLimitDependency::create_shared(id, node_id, query_ctx);
521+
}
522+
515523
void VDataStreamRecvr::close() {
516524
if (_is_closed) {
517525
return;

‎be/src/vec/runtime/vdata_stream_recvr.h‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
#include "common/object_pool.h"
4343
#include "common/status.h"
4444
#include "runtime/descriptors.h"
45+
#include "runtime/query_context.h"
4546
#include "runtime/query_statistics.h"
4647
#include "util/runtime_profile.h"
4748
#include "util/stopwatch.hpp"
@@ -61,6 +62,7 @@ class RuntimeState;
6162
namespace pipeline {
6263
struct ExchangeDataDependency;
6364
class LocalExchangeChannelDependency;
65+
class LocalExchangeMemLimitDependency;
6466
class ExchangeLocalState;
6567
} // namespace pipeline
6668

@@ -130,6 +132,10 @@ class VDataStreamRecvr {
130132
std::shared_ptr<pipeline::LocalExchangeChannelDependency> get_local_channel_dependency(
131133
int sender_id);
132134

135+
void create_mem_limit_dependency(int id, int node_id, QueryContext* query_ctx);
136+
137+
auto get_mem_limit_dependency() { return _exchange_sink_mem_limit_dependency; }
138+
133139
private:
134140
void update_blocks_memory_usage(int64_t size);
135141
class PipSenderQueue;
@@ -189,7 +195,8 @@ class VDataStreamRecvr {
189195
std::vector<std::shared_ptr<pipeline::LocalExchangeChannelDependency>>
190196
_sender_to_local_channel_dependency;
191197

192-
std::shared_ptr<bool> _mem_available;
198+
// use to limit sink write
199+
std::shared_ptr<pipeline::LocalExchangeMemLimitDependency> _exchange_sink_mem_limit_dependency;
193200
};
194201

195202
class ThreadClosure : public google::protobuf::Closure {

‎be/src/vec/sink/vdata_stream_sender.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,11 @@ class Channel {
321321

322322
void set_receiver_eof(Status st) { _receiver_status = st; }
323323

324+
auto local_recvr() {
325+
DCHECK(is_local());
326+
return _local_recvr;
327+
}
328+
324329
protected:
325330
bool _recvr_is_valid() {
326331
if (_local_recvr && !_local_recvr->is_closed()) {

0 commit comments

Comments
 (0)