[BugFix] Revert 'Avoid brpc communication when using local pass through' (#61709)

Author: before-Sunrise

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

@@ -29,6 +29,7 @@
 #include "runtime/data_stream_mgr.h"
 #include "runtime/descriptors.h"
 #include "runtime/exec_env.h"
+#include "runtime/local_pass_through_buffer.h"
 #include "runtime/runtime_state.h"
 #include "serde/compress_strategy.h"
 #include "serde/protobuf_serde.h"
@@ -46,13 +47,15 @@ class ExchangeSinkOperator::Channel {
     // how much tuple data is getting accumulated before being sent; it only applies
     // when data is added via add_row() and not sent directly via send_batch().
     Channel(ExchangeSinkOperator* parent, const TNetworkAddress& brpc_dest, const TUniqueId& fragment_instance_id,
-            PlanNodeId dest_node_id, int32_t num_shuffles, bool enable_exchange_pass_through, bool enable_exchange_perf)
+            PlanNodeId dest_node_id, int32_t num_shuffles, bool enable_exchange_pass_through, bool enable_exchange_perf,
+            PassThroughChunkBuffer* pass_through_chunk_buffer)
             : _parent(parent),
               _brpc_dest_addr(brpc_dest),
               _fragment_instance_id(fragment_instance_id),
               _dest_node_id(dest_node_id),
               _enable_exchange_pass_through(enable_exchange_pass_through),
               _enable_exchange_perf(enable_exchange_perf),
+              _pass_through_context(pass_through_chunk_buffer, fragment_instance_id, dest_node_id),
               _chunks(num_shuffles) {}
 
     // Initialize channel.
@@ -114,6 +117,7 @@ class ExchangeSinkOperator::Channel {
     // enable it to profile exchange's performance, which ignores computing local data for exchange_speed/_bytes,
     // because local data isn't accessed by remote network.
     const bool _enable_exchange_perf;
+    PassThroughContext _pass_through_context;
 
     bool _is_first_chunk = true;
     std::shared_ptr<PInternalService_RecoverableStub> _brpc_stub = nullptr;
@@ -123,8 +127,6 @@ class ExchangeSinkOperator::Channel {
     // If pipeline level shuffle is disable, the size of _chunks
     // always be 1
     std::vector<std::unique_ptr<Chunk>> _chunks;
-    ChunkPassThroughVectorPtr _pass_through_chunks;
-    int64_t _pass_through_physical_bytes = 0;
     PTransmitChunkParamsPtr _chunk_request;
     size_t _current_request_bytes = 0;
 
@@ -152,6 +154,7 @@ bool ExchangeSinkOperator::Channel::_check_use_pass_through() {
 }
 
 void ExchangeSinkOperator::Channel::_prepare_pass_through() {
+    _pass_through_context.init();
     _use_pass_through = _check_use_pass_through();
 }
 
@@ -226,8 +229,6 @@ Status ExchangeSinkOperator::Channel::send_one_chunk(RuntimeState* state, const
         _chunk_request->set_node_id(_dest_node_id);
         _chunk_request->set_sender_id(_parent->_sender_id);
         _chunk_request->set_be_number(_parent->_be_number);
-        _pass_through_chunks = std::make_unique<ChunkPassThroughVector>();
-        _pass_through_physical_bytes = 0;
         if (_parent->_is_pipeline_level_shuffle) {
             _chunk_request->set_is_pipeline_level_shuffle(true);
         }
@@ -236,19 +237,18 @@ Status ExchangeSinkOperator::Channel::send_one_chunk(RuntimeState* state, const
     // If chunk is not null, append it to request
     if (chunk != nullptr) {
         if (_use_pass_through) {
-            int64_t before_bytes = CurrentThread::current().get_consumed_bytes();
-            auto clone = chunk->clone_unique();
-            int64_t physical_bytes = CurrentThread::current().get_consumed_bytes() - before_bytes;
-            _pass_through_physical_bytes += physical_bytes;
             size_t chunk_size = serde::ProtobufChunkSerde::max_serialized_size(*chunk);
-            _pass_through_chunks->emplace_back(std::move(clone), driver_sequence, chunk_size, physical_bytes);
-            COUNTER_UPDATE(_parent->_bytes_pass_through_counter, chunk_size);
+            // -1 means disable pipeline level shuffle
+            TRY_CATCH_BAD_ALLOC(
+                    _pass_through_context.append_chunk(_parent->_sender_id, chunk, chunk_size,
+                                                       _parent->_is_pipeline_level_shuffle ? driver_sequence : -1));
             _current_request_bytes += chunk_size;
+            COUNTER_UPDATE(_parent->_bytes_pass_through_counter, chunk_size);
+            COUNTER_SET(_parent->_pass_through_buffer_peak_mem_usage, _pass_through_context.total_bytes());
         } else {
             if (_parent->_is_pipeline_level_shuffle) {
                 _chunk_request->add_driver_sequences(driver_sequence);
             }
-
             auto pchunk = _chunk_request->add_chunks();
             TRY_CATCH_BAD_ALLOC(RETURN_IF_ERROR(_parent->serialize_chunk(chunk, pchunk, &_is_first_chunk)));
             _current_request_bytes += pchunk->data().size();
@@ -261,21 +261,12 @@ Status ExchangeSinkOperator::Channel::send_one_chunk(RuntimeState* state, const
         _chunk_request->set_eos(eos);
         _chunk_request->set_use_pass_through(_use_pass_through);
         butil::IOBuf attachment;
-        int64_t physical_bytes = _use_pass_through ? _pass_through_physical_bytes
-                                                   : _parent->construct_brpc_attachment(_chunk_request, attachment);
-        TransmitChunkInfo info = {this->_fragment_instance_id,
-                                  _brpc_stub,
-                                  std::move(_chunk_request),
-                                  std::move(_pass_through_chunks),
-                                  state->exec_env()->stream_mgr(),
-                                  attachment,
-                                  physical_bytes,
-                                  _brpc_dest_addr};
+        int64_t attachment_physical_bytes = _parent->construct_brpc_attachment(_chunk_request, attachment);
+        TransmitChunkInfo info = {this->_fragment_instance_id, _brpc_stub,     std::move(_chunk_request), attachment,
+                                  attachment_physical_bytes,   _brpc_dest_addr};
         RETURN_IF_ERROR(_parent->_buffer->add_request(info));
         _current_request_bytes = 0;
         _chunk_request.reset();
-        _pass_through_chunks = std::make_unique<ChunkPassThroughVector>();
-        _pass_through_physical_bytes = 0;
         *is_real_sent = true;
     }
 
@@ -293,8 +284,8 @@ Status ExchangeSinkOperator::Channel::send_chunk_request(RuntimeState* state, PT
     chunk_request->set_be_number(_parent->_be_number);
     chunk_request->set_eos(false);
     chunk_request->set_use_pass_through(_use_pass_through);
-    TransmitChunkInfo info = {this->_fragment_instance_id,     _brpc_stub, std::move(chunk_request),  nullptr,
-                              state->exec_env()->stream_mgr(), attachment, attachment_physical_bytes, _brpc_dest_addr};
+    TransmitChunkInfo info = {this->_fragment_instance_id, _brpc_stub,     std::move(chunk_request), attachment,
+                              attachment_physical_bytes,   _brpc_dest_addr};
     RETURN_IF_ERROR(_parent->_buffer->add_request(info));
 
     return Status::OK();
@@ -351,6 +342,10 @@ ExchangeSinkOperator::ExchangeSinkOperator(
           _output_columns(output_columns),
           _num_sinkers(num_sinkers) {
     std::map<int64_t, int64_t> fragment_id_to_channel_index;
+    RuntimeState* state = fragment_ctx->runtime_state();
+
+    PassThroughChunkBuffer* pass_through_chunk_buffer =
+            state->exec_env()->stream_mgr()->get_pass_through_chunk_buffer(state->query_id());
 
     _channels.reserve(destinations.size());
     std::vector<int> driver_sequence_per_channel(destinations.size(), 0);
@@ -364,7 +359,7 @@ ExchangeSinkOperator::ExchangeSinkOperator(
         } else {
             std::unique_ptr<Channel> channel = std::make_unique<Channel>(
                     this, destination.brpc_server, fragment_instance_id, dest_node_id, _num_shuffles_per_channel,
-                    enable_exchange_pass_through, enable_exchange_perf);
+                    enable_exchange_pass_through, enable_exchange_perf, pass_through_chunk_buffer);
             _channels.emplace_back(channel.get());
             _instance_id2channel.emplace(fragment_instance_id.lo, std::move(channel));
         }
@@ -460,6 +455,9 @@ Status ExchangeSinkOperator::prepare(RuntimeState* state) {
     _shuffle_chunk_append_counter = ADD_COUNTER(_unique_metrics, "ShuffleChunkAppendCounter", TUnit::UNIT);
     _shuffle_chunk_append_timer = ADD_TIMER(_unique_metrics, "ShuffleChunkAppendTime");
     _compress_timer = ADD_TIMER(_unique_metrics, "CompressTime");
+    _pass_through_buffer_peak_mem_usage = _unique_metrics->AddHighWaterMarkCounter(
+            "PassThroughBufferPeakMemoryUsage", TUnit::BYTES,
+            RuntimeProfile::Counter::create_strategy(TUnit::BYTES, TCounterMergeType::SKIP_FIRST_MERGE));
 
     for (auto& [_, channel] : _instance_id2channel) {
         RETURN_IF_ERROR(channel->init(state));
@@ -651,10 +649,8 @@ Status ExchangeSinkOperator::set_finishing(RuntimeState* state) {
         butil::IOBuf attachment;
         const int64_t attachment_physical_bytes = construct_brpc_attachment(_chunk_request, attachment);
         for (const auto& [_, channel] : _instance_id2channel) {
-            if (!channel->use_pass_through()) {
-                PTransmitChunkParamsPtr copy = std::make_shared<PTransmitChunkParams>(*_chunk_request);
-                RETURN_IF_ERROR(channel->send_chunk_request(state, copy, attachment, attachment_physical_bytes));
-            }
+            PTransmitChunkParamsPtr copy = std::make_shared<PTransmitChunkParams>(*_chunk_request);
+            RETURN_IF_ERROR(channel->send_chunk_request(state, copy, attachment, attachment_physical_bytes));
         }
         _current_request_bytes = 0;
         _chunk_request.reset();

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

@@ -186,6 +186,7 @@ class ExchangeSinkOperator final : public Operator {
     RuntimeProfile::Counter* _raw_input_bytes_counter = nullptr;
     RuntimeProfile::Counter* _serialized_bytes_counter = nullptr;
     RuntimeProfile::Counter* _compressed_bytes_counter = nullptr;
+    RuntimeProfile::HighWaterMarkCounter* _pass_through_buffer_peak_mem_usage = nullptr;
 
     std::atomic<bool> _is_finished = false;
     std::atomic<bool> _is_cancelled = false;

be/src/exec/pipeline/exchange/sink_buffer.cpp

@@ -20,11 +20,9 @@
 #include <memory>
 #include <mutex>
 #include <string_view>
-#include <thread>
 
 #include "exec/pipeline/schedule/utils.h"
 #include "fmt/core.h"
-#include "runtime/data_stream_mgr.h"
 #include "util/defer_op.h"
 #include "util/time.h"
 #include "util/uid_util.h"
@@ -56,7 +54,6 @@ SinkBuffer::SinkBuffer(FragmentContext* fragment_ctx, const std::vector<TPlanFra
             ctx.num_finished_rpcs = 0;
             ctx.num_in_flight_rpcs = 0;
             ctx.dest_addrs = dest.brpc_server;
-            ctx.pass_through_blocked = false;
 
             PUniqueId finst_id;
             finst_id.set_hi(instance_id.hi);
@@ -112,16 +109,7 @@ Status SinkBuffer::add_request(TransmitChunkInfo& request) {
         auto& instance_id = request.fragment_instance_id;
         auto& context = sink_ctx(instance_id.lo);
 
-        if (request.params->has_use_pass_through() && request.params->use_pass_through()) {
-            RETURN_IF_ERROR(_try_to_send_local(instance_id, [&, tracker = CurrentThread::current().mem_tracker()]() {
-                // Release allocated bytes in current MemTracker, since it would not be released at current MemTracker
-                tracker->release(request.physical_bytes);
-                GlobalEnv::GetInstance()->passthrough_mem_tracker()->consume(request.physical_bytes);
-                context.buffer.push(std::move(request));
-            }));
-        } else {
-            RETURN_IF_ERROR(_try_to_send_rpc(instance_id, [&]() { context.buffer.push(std::move(request)); }));
-        }
+        RETURN_IF_ERROR(_try_to_send_rpc(instance_id, [&]() { context.buffer.push(request); }));
     }
 
     return Status::OK();
@@ -163,7 +151,7 @@ bool SinkBuffer::is_finished() const {
         return false;
     }
 
-    return _num_sending == 0 && _total_in_flight_rpc == 0;
+    return _num_sending_rpc == 0 && _total_in_flight_rpc == 0;
 }
 
 void SinkBuffer::update_profile(RuntimeProfile* profile) {
@@ -233,9 +221,9 @@ void SinkBuffer::cancel_one_sinker(RuntimeState* const state) {
         // check how many cancel operations are issued, and show the state of that time.
         VLOG_OPERATOR << fmt::format(
                 "fragment_instance_id {}, _num_uncancelled_sinkers {}, _is_finishing {}, _num_remaining_eos {}, "
-                "_num_sending {}, chunk is full {}",
+                "_num_sending_rpc {}, chunk is full {}",
                 print_id(_fragment_ctx->fragment_instance_id()), _num_uncancelled_sinkers, _is_finishing,
-                _num_remaining_eos, _num_sending, is_full());
+                _num_remaining_eos, _num_sending_rpc, is_full());
     }
 }
 
@@ -268,114 +256,13 @@ void SinkBuffer::_process_send_window(const TUniqueId& instance_id, const int64_
     }
 }
 
-Status SinkBuffer::_try_to_send_local(const TUniqueId& instance_id, const std::function<void()>& pre_works) {
-    // switch to process tracker
-    SCOPED_THREAD_LOCAL_MEM_TRACKER_SETTER(nullptr);
-    auto& context = sink_ctx(instance_id.lo);
-    auto notify = this->finishing_defer();
-    std::lock_guard guard(context.mutex);
-    DeferOp decrease_defer([this]() { --_num_sending; });
-    ++_num_sending;
-
-    pre_works();
-
-    for (;;) {
-        DeferOp reset_owner_id([&context]() { context.owner_id = {}; });
-
-        if (_is_finishing) {
-            return Status::OK();
-        }
-
-        auto& buffer = context.buffer;
-        if (buffer.empty() || context.pass_through_blocked) {
-            return Status::OK();
-        }
-
-        TransmitChunkInfo& request = buffer.front();
-        DeferOp pop_defer([&buffer, mem_tracker = _mem_tracker]() {
-            // The request memory is acquired by ExchangeSinkOperator,
-            // so use the instance_mem_tracker passed from ExchangeSinkOperator to release memory.
-            // This must be invoked before decrease_defer desctructed to avoid sink_buffer and fragment_ctx released.
-            SCOPED_THREAD_LOCAL_MEM_TRACKER_SETTER(mem_tracker);
-            buffer.pop();
-        });
-
-        if (request.params->eos()) {
-            DeferOp eos_defer([this, &instance_id]() {
-                if (--_num_remaining_eos == 0) {
-                    _is_finishing = true;
-                }
-                sink_ctx(instance_id.lo).num_sinker--;
-            });
-            // Only the last eos is sent to ExchangeSourceOperator.
-            if (context.num_sinker > 1) {
-                if (request.pass_through_chunks->size() == 0) {
-                    continue;
-                }
-                request.params->set_eos(false);
-            } else {
-                // this is the last eos query, set query stats
-                if (auto final_stats = _fragment_ctx->runtime_state()->query_ctx()->intermediate_query_statistic(
-                            _delta_bytes_sent.exchange(0))) {
-                    final_stats->to_pb(request.params->mutable_query_statistics());
-                }
-            }
-        }
-
-        auto query_ctx = std::weak_ptr(_fragment_ctx->runtime_state()->query_ctx()->shared_from_this());
-        context.owner_id = std::this_thread::get_id();
-        auto* closure = new DisposablePassThroughClosure([query_ctx, this, instance_id]() noexcept {
-            auto guard = query_ctx.lock();
-            RETURN_IF(guard == nullptr, (void)0);
-            auto& context = sink_ctx(instance_id.lo);
-
-            // Avoid local-passthrough recursive calls.
-            if (context.owner_id == std::this_thread::get_id()) {
-                --_total_in_flight_rpc;
-                context.pass_through_blocked = false;
-                return;
-            }
-
-            auto notify = this->defer_notify();
-            auto defer = DeferOp([this]() { --_total_in_flight_rpc; });
-            context.pass_through_blocked = false;
-            static_cast<void>(_try_to_send_local(instance_id, []() {}));
-        });
-
-        context.pass_through_blocked = true;
-        if (_first_send_time == -1) {
-            _first_send_time = MonotonicNanos();
-        }
-
-        // Decrease memory from pass through before moving chunks to the reciever fragment.
-        GlobalEnv::GetInstance()->passthrough_mem_tracker()->release(request.physical_bytes);
-
-        DCHECK_EQ(request.params->use_pass_through(), true);
-        ::google::protobuf::Closure* done = closure;
-
-        _total_in_flight_rpc++;
-        auto defer = CancelableDefer([this]() { --_total_in_flight_rpc; });
-        Status st = request.stream_mgr->transmit_chunk(instance_id, *request.params,
-                                                       std::move(request.pass_through_chunks), &done);
-        if (st.ok()) defer.cancel();
-        if (st.ok() && done != nullptr) {
-            // if the closure was not removed delete it and continue transmitting chunks.
-            done->Run();
-            continue;
-        }
-        return st;
-    }
-    return Status::OK();
-}
-
 Status SinkBuffer::_try_to_send_rpc(const TUniqueId& instance_id, const std::function<void()>& pre_works) {
     auto& context = sink_ctx(instance_id.lo);
-    auto notify = this->finishing_defer();
     std::lock_guard guard(context.mutex);
     pre_works();
 
-    DeferOp decrease_defer([this]() { --_num_sending; });
-    ++_num_sending;
+    DeferOp decrease_defer([this]() { --_num_sending_rpc; });
+    ++_num_sending_rpc;
 
     for (;;) {
         if (_is_finishing) {
@@ -426,6 +313,7 @@ Status SinkBuffer::_try_to_send_rpc(const TUniqueId& instance_id, const std::fun
                     return;
                 }
                 if (--_num_remaining_eos == 0) {
+                    auto notify = this->defer_notify();
                     _is_finishing = true;
                 }
                 sink_ctx(instance_id.lo).num_sinker--;
@@ -527,8 +415,8 @@ Status SinkBuffer::_try_to_send_rpc(const TUniqueId& instance_id, const std::fun
 
         // Attachment will be released by process_mem_tracker in closure->Run() in bthread, when receiving the response,
         // so decrease the memory usage of attachment from instance_mem_tracker immediately before sending the request.
-        _mem_tracker->release(request.physical_bytes);
-        GlobalEnv::GetInstance()->process_mem_tracker()->consume(request.physical_bytes);
+        _mem_tracker->release(request.attachment_physical_bytes);
+        GlobalEnv::GetInstance()->process_mem_tracker()->consume(request.attachment_physical_bytes);
 
         closure->cntl.Reset();
         closure->cntl.set_timeout_ms(_brpc_timeout_ms);