Limit memory usage of exchange sender (#10387) (#10417)

close #10337

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>

Co-authored-by: xufei <xufeixw@mail.ustc.edu.cn>
Co-authored-by: xufei <xufei@pingcap.com>

Author: ti-chi-bot

dbms/src/Flash/Coprocessor/DAGDriver.cpp

@@ -20,6 +20,7 @@
 #include <DataStreams/copyData.h>
 #include <Flash/Coprocessor/DAGContext.h>
 #include <Flash/Coprocessor/DAGDriver.h>
+#include <Flash/Coprocessor/DAGUtils.h>
 #include <Flash/Coprocessor/StreamWriter.h>
 #include <Flash/Coprocessor/StreamingDAGResponseWriter.h>
 #include <Flash/Coprocessor/UnaryDAGResponseWriter.h>
@@ -171,6 +172,7 @@ try
             streaming_writer,
             context.getSettingsRef().dag_records_per_chunk,
             context.getSettingsRef().batch_send_min_limit,
+            getMaxBufferedBytesInResponseWriter(context.getSettingsRef().max_buffered_bytes_in_executor, 1),
             dag_context);
         response_writer->prepare(query_executor->getSampleBlock());
         query_executor
@@ -225,6 +227,7 @@ try
             streaming_writer,
             context.getSettingsRef().dag_records_per_chunk,
             context.getSettingsRef().batch_send_min_limit,
+            getMaxBufferedBytesInResponseWriter(context.getSettingsRef().max_buffered_bytes_in_executor, 1),
             dag_context);
         response_writer->prepare(query_executor->getSampleBlock());
         query_executor

dbms/src/Flash/Coprocessor/DAGResponseWriter.h

@@ -49,10 +49,17 @@ class DAGResponseWriter
 
     virtual ~DAGResponseWriter() = default;
 
+    bool needFlush() const { return buffered_bytes >= max_buffered_bytes || buffered_rows >= max_buffered_rows; }
+
 protected:
     Int64 records_per_chunk;
     DAGContext & dag_context;
     bool has_pending_flush = false;
+
+    UInt64 max_buffered_bytes = 0;
+    UInt64 max_buffered_rows = 0;
+    UInt64 buffered_bytes = 0;
+    UInt64 buffered_rows = 0;
 };
 
 } // namespace DB

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include <Common/FmtUtils.h>
+#include <Common/ThresholdUtils.h>
 #include <Common/TiFlashException.h>
 #include <Core/Types.h>
 #include <DataTypes/DataTypeNullable.h>
@@ -1570,4 +1571,19 @@ tipb::ScalarFuncSig reverseGetFuncSigByFuncName(const String & name)
         throw Exception(fmt::format("Unsupported function {}", name));
     return func_name_sig_map[name];
 }
+
+constexpr ssize_t MAX_BATCH_SEND_MIN_LIMIT_MEM_SIZE
+    = 1024 * 1024 * 32; // 32MB (8192 Rows * 256 Byte/row * 16 partitions)
+
+UInt64 getMaxBufferedBytesInResponseWriter(Int64 max_buffered_bytes_in_executor, size_t concurrency)
+{
+    // max buffered bytes in ExchangeSender, we don't want to send too many small chunks so there is a minimum limit of MAX_BATCH_SEND_MIN_LIMIT_MEM_SIZE
+    UInt64 ret = MAX_BATCH_SEND_MIN_LIMIT_MEM_SIZE;
+    if (max_buffered_bytes_in_executor > 0)
+    {
+        ret = std::max(ret, getAverageThreshold(max_buffered_bytes_in_executor, concurrency));
+    }
+    return ret;
+}
+
 } // namespace DB

dbms/src/Flash/Coprocessor/DAGUtils.h

@@ -85,6 +85,8 @@ void assertBlockSchema(const DataTypes & expected_types, const Block & block, co
 
 void assertBlockSchema(const Block & header, const Block & block, const String & context_description);
 
+UInt64 getMaxBufferedBytesInResponseWriter(Int64 max_buffered_bytes_in_executor, size_t concurrency);
+
 class UniqueNameGenerator
 {
 private:

dbms/src/Flash/Coprocessor/StreamingDAGResponseWriter.cpp

@@ -35,12 +35,11 @@ StreamingDAGResponseWriter<StreamWriterPtr>::StreamingDAGResponseWriter(
     StreamWriterPtr writer_,
     Int64 records_per_chunk_,
     Int64 batch_send_min_limit_,
+    UInt64 max_buffered_bytes_,
     DAGContext & dag_context_)
     : DAGResponseWriter(records_per_chunk_, dag_context_)
-    , batch_send_min_limit(batch_send_min_limit_)
     , writer(writer_)
 {
-    rows_in_blocks = 0;
     switch (dag_context.encode_type)
     {
     case tipb::EncodeType::TypeDefault:
@@ -56,15 +55,20 @@ StreamingDAGResponseWriter<StreamWriterPtr>::StreamingDAGResponseWriter(
         throw TiFlashException("Unsupported EncodeType", Errors::Coprocessor::Internal);
     }
     /// For other encode types, we will use records_per_chunk to control the batch size sent.
-    batch_send_min_limit
-        = dag_context.encode_type == tipb::EncodeType::TypeCHBlock ? batch_send_min_limit : (records_per_chunk - 1);
+    auto batch_send_min_limit
+        = dag_context.encode_type == tipb::EncodeType::TypeCHBlock ? batch_send_min_limit_ : (records_per_chunk - 1);
+    if (batch_send_min_limit <= 0)
+        max_buffered_rows = 1;
+    else
+        max_buffered_rows = static_cast<UInt64>(batch_send_min_limit);
+    max_buffered_bytes = max_buffered_bytes_;
 }
 
 template <class StreamWriterPtr>
 WriteResult StreamingDAGResponseWriter<StreamWriterPtr>::flush()
 {
     has_pending_flush = false;
-    if (rows_in_blocks > 0)
+    if (buffered_rows > 0)
     {
         auto wait_res = waitForWritable();
         if (wait_res == WaitResult::Ready)
@@ -96,11 +100,12 @@ WriteResult StreamingDAGResponseWriter<StreamWriterPtr>::write(const Block & blo
     size_t rows = block.rows();
     if (rows > 0)
     {
-        rows_in_blocks += rows;
+        buffered_rows += rows;
+        buffered_bytes += block.allocatedBytes();
         blocks.push_back(block);
     }
 
-    if (static_cast<Int64>(rows_in_blocks) > batch_send_min_limit)
+    if (needFlush())
     {
         return flush();
     }
@@ -157,7 +162,8 @@ void StreamingDAGResponseWriter<StreamWriterPtr>::encodeThenWriteBlocks()
     }
 
     assert(blocks.empty());
-    rows_in_blocks = 0;
+    buffered_rows = 0;
+    buffered_bytes = 0;
     writer->write(response.getResponse());
 }
 

dbms/src/Flash/Coprocessor/StreamingDAGResponseWriter.h

@@ -35,6 +35,7 @@ class StreamingDAGResponseWriter : public DAGResponseWriter
         StreamWriterPtr writer_,
         Int64 records_per_chunk_,
         Int64 batch_send_min_limit_,
+        UInt64 max_buffered_bytes_,
         DAGContext & dag_context_);
     WaitResult waitForWritable() const override;
 
@@ -45,10 +46,8 @@ class StreamingDAGResponseWriter : public DAGResponseWriter
     void encodeThenWriteBlocks();
 
 private:
-    Int64 batch_send_min_limit;
     StreamWriterPtr writer;
     std::vector<Block> blocks;
-    size_t rows_in_blocks;
     std::unique_ptr<ChunkCodecStream> chunk_codec_stream;
 };
 

dbms/src/Flash/Coprocessor/tests/gtest_streaming_writer.cpp

@@ -24,6 +24,7 @@
 #include <gtest/gtest.h>
 
 #include <Flash/Coprocessor/StreamingDAGResponseWriter.cpp>
+#include <cstddef>
 
 namespace DB
 {
@@ -113,6 +114,7 @@ try
         const size_t block_rows = 64;
         const size_t block_num = 64;
         const size_t batch_send_min_limit = 108;
+        const size_t max_buffered_bytes = 1024 * 1024 * 16;
 
         // 1. Build Blocks.
         std::vector<Block> blocks;
@@ -135,6 +137,7 @@ try
             mock_writer,
             batch_send_min_limit,
             batch_send_min_limit,
+            max_buffered_bytes,
             *dag_context_ptr);
         for (const auto & block : blocks)
             dag_writer->write(block);
@@ -169,5 +172,57 @@ try
 }
 CATCH
 
+TEST_F(TestStreamingWriter, testMaxBufferedBytes)
+try
+{
+    std::vector<tipb::EncodeType> encode_types{
+        tipb::EncodeType::TypeDefault,
+        tipb::EncodeType::TypeChunk,
+        tipb::EncodeType::TypeCHBlock};
+    for (const auto & encode_type : encode_types)
+    {
+        dag_context_ptr->encode_type = encode_type;
+
+        const size_t block_rows = 64;
+        const size_t block_num = 64;
+        const size_t batch_send_min_limit = block_rows * block_num / 2;
+        const std::vector<UInt64> max_buffered_bytes_vec{1, 1024 * 1024 * 1024};
+        const std::vector<UInt64> expected_response_num{block_num, 2};
+
+        for (size_t i = 0; i < max_buffered_bytes_vec.size(); ++i)
+        {
+            // 1. Build Blocks.
+            std::vector<Block> blocks;
+            for (size_t i = 0; i < block_num; ++i)
+            {
+                blocks.emplace_back(prepareBlock(block_rows));
+                blocks.emplace_back(prepareBlock(0));
+            }
+            Block header = blocks.back();
+
+            // 2. Build MockStreamWriter.
+            std::vector<tipb::SelectResponse> write_report;
+            auto checker = [&write_report](tipb::SelectResponse & response) {
+                write_report.emplace_back(std::move(response));
+            };
+            auto mock_writer = std::make_shared<MockStreamWriter>(checker);
+
+            // 3. Write blocks
+            auto dag_writer = std::make_shared<StreamingDAGResponseWriter<std::shared_ptr<MockStreamWriter>>>(
+                mock_writer,
+                batch_send_min_limit,
+                batch_send_min_limit,
+                max_buffered_bytes_vec[i],
+                *dag_context_ptr);
+            for (const auto & block : blocks)
+                dag_writer->write(block);
+            dag_writer->flush();
+
+            ASSERT_EQ(write_report.size(), expected_response_num[i]);
+        }
+    }
+}
+CATCH
+
 } // namespace tests
 } // namespace DB

dbms/src/Flash/Coprocessor/tests/gtest_ti_remote_block_inputstream.cpp

@@ -345,6 +345,7 @@ class TestTiRemoteBlockInputStream : public testing::Test
         auto dag_writer = std::make_shared<BroadcastOrPassThroughWriter<MockWriterPtr>>(
             writer,
             batch_send_min_limit,
+            1024 * 1024 * 16,
             *dag_context_ptr,
             MPPDataPacketVersion::MPPDataPacketV1,
             tipb::CompressionMode::FAST,
@@ -374,6 +375,7 @@ class TestTiRemoteBlockInputStream : public testing::Test
             writer,
             0,
             batch_send_min_limit,
+            1024 * 1024 * 16,
             *dag_context_ptr);
 
         // 2. encode all blocks

dbms/src/Flash/Mpp/BroadcastOrPassThroughWriter.cpp

@@ -25,36 +25,35 @@ namespace DB
 template <class ExchangeWriterPtr>
 BroadcastOrPassThroughWriter<ExchangeWriterPtr>::BroadcastOrPassThroughWriter(
     ExchangeWriterPtr writer_,
-    Int64 batch_send_min_limit_,
+    Int64 max_buffered_rows_,
+    UInt64 max_buffered_bytes_,
     DAGContext & dag_context_,
     MPPDataPacketVersion data_codec_version_,
     tipb::CompressionMode compression_mode_,
     tipb::ExchangeType exchange_type_)
     : DAGResponseWriter(/*records_per_chunk=*/-1, dag_context_)
-    , batch_send_min_limit(batch_send_min_limit_)
     , writer(writer_)
     , exchange_type(exchange_type_)
     , data_codec_version(data_codec_version_)
     , compression_method(ToInternalCompressionMethod(compression_mode_))
 {
-    rows_in_blocks = 0;
     RUNTIME_CHECK(dag_context.encode_type == tipb::EncodeType::TypeCHBlock);
     RUNTIME_CHECK(exchange_type == tipb::ExchangeType::Broadcast || exchange_type == tipb::ExchangeType::PassThrough);
 
     switch (data_codec_version)
     {
     case MPPDataPacketV0:
-        if (batch_send_min_limit <= 0)
-            batch_send_min_limit = 1;
+        if (max_buffered_rows_ <= 0)
+            max_buffered_rows_ = 1;
         break;
     case MPPDataPacketV1:
     default:
     {
         // make `batch_send_min_limit` always GT 0
-        if (batch_send_min_limit <= 0)
+        if (max_buffered_rows_ <= 0)
         {
             // set upper limit if not specified
-            batch_send_min_limit = 8 * 1024 /* 8K */;
+            max_buffered_rows_ = 8 * 1024 /* 8K */;
         }
         for (const auto & field_type : dag_context.result_field_types)
         {
@@ -63,13 +62,15 @@ BroadcastOrPassThroughWriter<ExchangeWriterPtr>::BroadcastOrPassThroughWriter(
         break;
     }
     }
+    max_buffered_rows = static_cast<UInt64>(max_buffered_rows_);
+    max_buffered_bytes = max_buffered_bytes_;
 }
 
 template <class ExchangeWriterPtr>
 WriteResult BroadcastOrPassThroughWriter<ExchangeWriterPtr>::flush()
 {
     has_pending_flush = false;
-    if (rows_in_blocks > 0)
+    if (buffered_rows > 0)
     {
         auto wait_res = waitForWritable();
         if (wait_res == WaitResult::Ready)
@@ -102,11 +103,12 @@ WriteResult BroadcastOrPassThroughWriter<ExchangeWriterPtr>::write(const Block &
     size_t rows = block.rows();
     if (rows > 0)
     {
-        rows_in_blocks += rows;
+        buffered_rows += rows;
+        buffered_bytes += block.allocatedBytes();
         blocks.push_back(block);
     }
 
-    if (static_cast<Int64>(rows_in_blocks) >= batch_send_min_limit)
+    if (needFlush())
     {
         return flush();
     }
@@ -130,7 +132,8 @@ void BroadcastOrPassThroughWriter<ExchangeWriterPtr>::writeBlocks()
     else
         writer->passThroughWrite(blocks, data_codec_version, compression_method);
     blocks.clear();
-    rows_in_blocks = 0;
+    buffered_rows = 0;
+    buffered_bytes = 0;
 }
 
 template class BroadcastOrPassThroughWriter<SyncMPPTunnelSetWriterPtr>;

dbms/src/Flash/Mpp/BroadcastOrPassThroughWriter.h

@@ -31,7 +31,8 @@ class BroadcastOrPassThroughWriter : public DAGResponseWriter
 public:
     BroadcastOrPassThroughWriter(
         ExchangeWriterPtr writer_,
-        Int64 batch_send_min_limit_,
+        Int64 max_buffered_rows_,
+        UInt64 max_buffered_bytes_,
         DAGContext & dag_context_,
         MPPDataPacketVersion data_codec_version_,
         tipb::CompressionMode compression_mode_,
@@ -44,10 +45,8 @@ class BroadcastOrPassThroughWriter : public DAGResponseWriter
     void writeBlocks();
 
 private:
-    Int64 batch_send_min_limit;
     ExchangeWriterPtr writer;
     std::vector<Block> blocks;
-    size_t rows_in_blocks;
     const tipb::ExchangeType exchange_type;
 
     // support data compression