IMPALA-13475: Mem limit deferred RPC processing parallelism

Before this patch KrpcDataStreamRecvr::SenderQueue::GetBatch() didn't
consider the mem requirement of the batches when initiating the
deserialization of deferred RCPs (EnqueueDeserializeTask) and tried
to deserialize as many batches in parallel as possible, based
on FLAGS_datastream_service_num_deserialization_threads.

This meant that some of the awakened threads could be rejected due to
hitting mem limit. In this case the deferred RPC was put back in queue.
Besides wasting resources this could also lead to starving bigger RPCs
by rejecting them and letting smaller RPCs take over.

The patch introduces the concept of "pending deferred RPCs", which
are RPCs where processing was requested, but the deserialization
thread didn't pick up the RPCs yet. The deserialized size of pending
deferred RPCs is considered already taken and only N deserializations
are started if the first N RPC's payload would fit to memory (with
the exception of empty queue where at least 1 is started).

The state flow of deferred RPCs became the following where each
transition should be able to succeed:
1. no deserialization thread started (in deferred_rpcs_)
2. deserialization thread started (in pending_deferred_rpcs_)
3. currently deserialized (tracked in num_pending_enqueue_)
4. enqueued
Note that it is still possible for batches to take over in 3
if their deserialization is faster.

Merging exchanges need special handling as those had multiple
queues with shared mem limit before this patch. The patch splits
the mem limit evenly among the queues to allow them enforcing
the "respect mem limit, with the exception when empty" rule
individually, as at least one batch in queue per sender is
critical to allow the merge sort to progress. Mem limit of closed
queues is reused by counting the active queues in atomic int
'num_active_queues_' and using this to calculate the share of total
mem limit a queue can use. There is still some risk of perf regression
in skewed merging exchanges where one sender may "need" the queue
space more than others.

Some parts are generated by copilot + Claude Sonnet 4.6

Change-Id: I4d8e2b07a19e4d30210d7903c01e75fba632a1de
Reviewed-on: http://gerrit.cloudera.org:8080/22217
Reviewed-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>
Tested-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>

Author: csringhofer

be/src/runtime/krpc-data-stream-mgr.cc

@@ -60,7 +60,8 @@ DEFINE_int32(datastream_sender_timeout_ms, 120000, "(Advanced) The time, in ms,
     "elapse  before a plan fragment will time-out trying to send the initial row batch.");
 DEFINE_int32(datastream_service_num_deserialization_threads, 16,
     "Number of threads for deserializing RPC requests deferred due to the receiver "
-    "not ready or the soft limit of the receiver is reached.");
+    "not ready or the soft limit of the receiver is reached. "
+    "If 0, it will be set to number of CPU cores.");
 DEFINE_int32(datastream_service_deserialization_queue_size, 10000,
     "Number of deferred RPC requests that can be enqueued before being processed by a "
     "deserialization thread.");
@@ -69,10 +70,13 @@ using std::mutex;
 namespace impala {
 
 KrpcDataStreamMgr::KrpcDataStreamMgr(MetricGroup* metrics)
-  : deserialize_pool_("data-stream-mgr", "deserialize",
-      FLAGS_datastream_service_num_deserialization_threads,
-      FLAGS_datastream_service_deserialization_queue_size,
-      boost::bind(&KrpcDataStreamMgr::DeserializeThreadFn, this, _1, _2)) {
+  : num_deserialization_threads_(
+    FLAGS_datastream_service_num_deserialization_threads > 0
+        ? FLAGS_datastream_service_num_deserialization_threads : CpuInfo::num_cores()),
+    deserialize_pool_("data-stream-mgr", "deserialize",
+        num_deserialization_threads_,
+        FLAGS_datastream_service_deserialization_queue_size,
+        boost::bind(&KrpcDataStreamMgr::DeserializeThreadFn, this, _1, _2)) {
   MetricGroup* dsm_metrics = metrics->GetOrCreateChildGroup("datastream-manager");
   num_senders_waiting_ =
       dsm_metrics->AddGauge("senders-blocked-on-recvr-creation", 0L);
@@ -473,4 +477,11 @@ KrpcDataStreamMgr::~KrpcDataStreamMgr() {
   deserialize_pool_.Join();
 }
 
+TransmitDataCtx::TransmitDataCtx(
+    const TransmitDataRequestPB* request, TransmitDataResponsePB* response,
+    kudu::rpc::RpcContext* rpc_context)
+  : request(request), response(response), rpc_context(rpc_context),
+    deserialized_size(RowBatch::GetDeserializedSize(request->row_batch_header())) { }
+
+
 } // namespace impala

be/src/runtime/krpc-data-stream-mgr.h

@@ -177,9 +177,11 @@ struct TransmitDataCtx {
   /// has been responded to. Not owned.
   kudu::rpc::RpcContext* rpc_context;
 
+  /// Memory needed by the batch after deserialization.
+  int64_t deserialized_size;
+
   TransmitDataCtx(const TransmitDataRequestPB* request, TransmitDataResponsePB* response,
-      kudu::rpc::RpcContext* rpc_context)
-    : request(request), response(response), rpc_context(rpc_context) { }
+      kudu::rpc::RpcContext* rpc_context);
 };
 
 /// Context for an EndDataStream() RPC. This structure is constructed when the RPC is
@@ -288,6 +290,8 @@ class KrpcDataStreamMgr : public CacheLineAligned {
   /// Waits for maintenance thread and sender response thread pool to finish.
   ~KrpcDataStreamMgr();
 
+  int num_deserialization_threads() const { return num_deserialization_threads_; }
+
  private:
   friend class KrpcDataStreamRecvr;
   friend class DataStreamTest;
@@ -320,6 +324,9 @@ class KrpcDataStreamMgr : public CacheLineAligned {
     int64_t enqueue_time_ns;
   };
 
+  /// Sets the number of threads in deserialize_pool_.
+  const int num_deserialization_threads_;
+
   /// Set of threads which deserialize buffered row batches, and deliver them to their
   /// receivers. Used only if RPCs were deferred when their channel's batch queue was
   /// full or if the receiver was not yet prepared.

be/src/runtime/krpc-data-stream-recvr.cc

@@ -161,12 +161,6 @@ class KrpcDataStreamRecvr {
   /// sender queue. Called from KrpcDataStreamMgr.
   void RemoveSender(int sender_id);
 
-  /// Return true if the addition of a new batch of size 'batch_size' would exceed the
-  /// total buffer limit.
-  bool ExceedsLimit(int64_t batch_size) {
-    return num_buffered_bytes_.Load() + batch_size > total_buffer_limit_;
-  }
-
   /// Return the current number of deferred RPCs.
   int64_t num_deferred_rpcs() const { return num_deferred_rpcs_.Load(); }
 
@@ -180,11 +174,6 @@ class KrpcDataStreamRecvr {
   TUniqueId fragment_instance_id_;
   PlanNodeId dest_node_id_;
 
-  /// Soft upper limit on the total amount of buffering in bytes allowed for this stream
-  /// across all sender queues. We defer processing of incoming RPCs once the amount of
-  /// buffered data exceeds this value.
-  const int64_t total_buffer_limit_;
-
   /// Row schema. Not owned.
   const RowDescriptor* row_desc_;
 
@@ -196,12 +185,20 @@ class KrpcDataStreamRecvr {
   /// from the fragment execution thread.
   bool closed_;
 
-  /// Current number of bytes held across all sender queues.
-  AtomicInt32 num_buffered_bytes_;
-
   /// Current number of outstanding deferred RPCs across all sender queues.
   AtomicInt64 num_deferred_rpcs_;
 
+  /// Total buffer limit across all sender queues as passed to the constructor. Each
+  /// queue derives its effective per-queue limit at runtime by dividing this value by
+  /// the current value of 'num_active_queues_' below, so that finished queues' budget
+  /// is automatically absorbed by the remaining active queues.
+  const int64_t total_buffer_limit_;
+
+  /// Number of sender queues whose sender has not yet finished. Decremented atomically
+  /// in RemoveSender(). Queues read this to compute their effective buffer limit.
+  /// Only meaningful when is_merging_ is true.
+  AtomicInt32 num_active_queues_;
+
   /// Memtracker for payloads of deferred Rpcs in the sender queue(s). This must be
   /// accessed with a sender queue's lock held to avoid race with Close() of the queue.
   boost::scoped_ptr<MemTracker> deferred_rpc_tracker_;
@@ -296,6 +293,9 @@ class KrpcDataStreamRecvr {
   /// Total wall-clock time in which the 'deferred_rpcs_' queues are not empty.
   RuntimeProfile::Counter* total_has_deferred_rpcs_timer_;
 
+  /// Total wall-clock time in which the 'pending_deferred_rpcs_' queues are not empty.
+  RuntimeProfile::Counter* total_has_pending_deferred_rpcs_timer_;
+
   /// Summary stats of time which RPCs spent in KRPC service queue before
   /// being dispatched to the RPC handlers.
   RuntimeProfile::SummaryStatsCounter* dispatch_timer_;

be/src/runtime/krpc-data-stream-recvr.h

@@ -393,11 +393,9 @@ void RowBatch::SetMemTracker(MemTracker* new_tracker) {
   mem_tracker_ = new_tracker;
 }
 
-int64_t RowBatch::GetDeserializedSize(const RowBatchHeaderPB& header,
-    const kudu::Slice& tuple_offsets) {
-  DCHECK_EQ(tuple_offsets.size() % sizeof(int32_t), 0);
-  return header.uncompressed_size() +
-      (tuple_offsets.size() / sizeof(int32_t)) * sizeof(Tuple*);
+int64_t RowBatch::GetDeserializedSize(const RowBatchHeaderPB& header) {
+  return header.uncompressed_size()
+      + header.num_rows() * header.num_tuples_per_row() * sizeof(Tuple*);
 }
 
 int64_t RowBatch::GetDeserializedSize(const OutboundRowBatch& batch) {

be/src/runtime/row-batch.cc

@@ -366,8 +366,7 @@ class RowBatch {
   /// less than the deserialized size.
   static int64_t GetSerializedSize(const OutboundRowBatch& batch);
   static int64_t GetDeserializedSize(const OutboundRowBatch& batch);
-  static int64_t GetDeserializedSize(const RowBatchHeaderPB& header,
-      const kudu::Slice& tuple_offsets);
+  static int64_t GetDeserializedSize(const RowBatchHeaderPB& header);
 
   int ALWAYS_INLINE num_rows() const { return num_rows_; }
   int ALWAYS_INLINE capacity() const { return capacity_; }

be/src/runtime/row-batch.h

@@ -46,7 +46,7 @@ class SpinLock {
   }
 
   /// Verify that the lock is held.
-  void DCheckLocked() { DCHECK(l_.IsHeld()); }
+  void DCheckLocked() const { DCHECK(l_.IsHeld()); }
 
  private:
   /// The underlying SpinLock from gutil.

be/src/util/spinlock.h

@@ -17,50 +17,84 @@
 
 from __future__ import absolute_import, division, print_function
 from builtins import range
+import time
 import pytest
 from tests.common.custom_cluster_test_suite import CustomClusterTestSuite
 from tests.common.skip import SkipIfBuildType
+from tests.util.cancel_util import cancel_query_and_validate_state
 
 
 @SkipIfBuildType.not_dev_build
 class TestExchangeDeferredBatches(CustomClusterTestSuite):
 
-    @classmethod
-    def setup_class(cls):
-      if cls.exploration_strategy() != 'exhaustive':
-        pytest.skip('runs only in exhaustive')
-      super(TestExchangeDeferredBatches, cls).setup_class()
+  @classmethod
+  def setup_class(cls):
+    if cls.exploration_strategy() != 'exhaustive':
+      pytest.skip('runs only in exhaustive')
+    super(TestExchangeDeferredBatches, cls).setup_class()
 
-    @pytest.mark.execute_serially
-    @CustomClusterTestSuite.with_args(
-        "--stress_datastream_recvr_delay_ms=3000"
-        + " --exchg_node_buffer_size_bytes=1024"
-        + " --datastream_service_num_deserialization_threads=1"
-        + " --impala_slow_rpc_threshold_ms=500")
-    def test_exchange_small_buffer(self, vector):
-        """Exercise the code which handles deferred row batches. In particular,
-        the exchange buffer is set to a small value to cause incoming row batches
-        to be deferred at the receiver. Also, use a single deserialization thread
-        to limit the speed in which the deferred row batches are dequeued. These
-        settings help expose the race in IMPALA-8239 when there is any error
-        deserializing deferred row batches."""
+  @pytest.mark.execute_serially
+  @CustomClusterTestSuite.with_args(
+      "--stress_datastream_recvr_delay_ms=3000"
+      + " --exchg_node_buffer_size_bytes=1024"
+      + " --datastream_service_num_deserialization_threads=1"
+      + " --impala_slow_rpc_threshold_ms=500")
+  def test_exchange_small_buffer(self, vector):
+    """Exercise the code which handles deferred row batches. In particular,
+    the exchange buffer is set to a small value to cause incoming row batches
+    to be deferred at the receiver. Also, use a single deserialization thread
+    to limit the speed in which the deferred row batches are dequeued. These
+    settings help expose the race in IMPALA-8239 when there is any error
+    deserializing deferred row batches."""
 
-        TEST_QUERY = "select count(*) from tpch.lineitem t1, tpch.lineitem t2 " +\
-           "where t1.l_orderkey = t2.l_orderkey"
-        EXPECTED_RESULT = ['30012985']
+    TEST_QUERY = "select count(*) from tpch.lineitem t1, tpch.lineitem t2 " +\
+        "where t1.l_orderkey = t2.l_orderkey"
+    EXPECTED_RESULT = ['30012985']
 
-        for i in range(10):
-            # Simulate row batch insertion failure. This triggers IMPALA-8239.
-            debug_action = 'RECVR_ADD_BATCH:FAIL@0.8'
-            self.execute_query_expect_failure(self.client, TEST_QUERY,
-                query_options={'debug_action': debug_action})
+    for i in range(10):
+      # Simulate row batch insertion failure. This triggers IMPALA-8239.
+      debug_action = 'RECVR_ADD_BATCH:FAIL@0.8'
+      self.execute_query_expect_failure(self.client, TEST_QUERY,
+          query_options={'debug_action': debug_action})
 
-        for i in range(10):
-            # Simulate row batch insertion failure. This triggers IMPALA-8239.
-            debug_action = 'RECVR_UNPACK_PAYLOAD:FAIL@0.8'
-            self.execute_query_expect_failure(self.client, TEST_QUERY,
-                query_options={'debug_action': debug_action})
+    for i in range(10):
+      # Simulate row batch insertion failure. This triggers IMPALA-8239.
+      debug_action = 'RECVR_UNPACK_PAYLOAD:FAIL@0.8'
+      self.execute_query_expect_failure(self.client, TEST_QUERY,
+          query_options={'debug_action': debug_action})
 
-        # Do a run with no debug action to make sure things are sane.
-        result = self.execute_query(TEST_QUERY, vector.get_value('exec_option'))
-        assert result.data == EXPECTED_RESULT
+    # Do a run with no debug action to make sure things are sane.
+    result = self.execute_query(TEST_QUERY, vector.get_value('exec_option'))
+    assert result.data == EXPECTED_RESULT
+
+  @pytest.mark.execute_serially
+  @CustomClusterTestSuite.with_args(
+      " --exchg_node_buffer_size_bytes=32768"
+      + " --datastream_service_num_deserialization_threads=16")
+  def test_parallel_deserialization_buffer_limit(self, vector):
+    """The test reproduces the scenario in IMPALA-13475 - delay is added to
+    the receiver to cause most TransmitData RPCs to be deferred and buffer
+    size is set to a small enough value where the queue can only hold 2 RowBatches
+    which is not enough for a row batch from all 3 senders. This means that the
+    receiver must limit the number of parallel deserialization tasks started to
+    ensure that the queue has enough memory for the result of all started tasks."""
+
+    TEST_QUERY = "select count(*) from tpch.lineitem t1, tpch.lineitem t2 " +\
+        "where t1.l_orderkey = t2.l_orderkey"
+    EXPECTED_RESULT = ['30012985']
+
+    # There are ~1100 row batches per exchange node under the join, with 1ms
+    # sleep they combine to guaranteed >2s run time.
+    query_options = {'debug_action': 'RECVR_ADD_BATCH:SLEEP@1'}
+    start = time.time()
+    result = self.execute_query(TEST_QUERY, query_options=query_options)
+    elapsed = time.time() - start
+    assert result.data == EXPECTED_RESULT
+    assert elapsed > 2, \
+        "Query should take >2s due to RECVR_ADD_BATCH:SLEEP@1, got %.2fs" % elapsed
+
+    # Test that the query can be cancelled at several points. The query takes around
+    # 3sec due to the debug_action.
+    for cancel_delay in [0.5, 1.0, 1.5, 2.0, 2.5]:
+      cancel_query_and_validate_state(TEST_QUERY,
+          query_options, vector.get_value('table_format'), cancel_delay)