From 2ed1b5957356ac62c357b5725b7922f1d9a2f456 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 17 Aug 2026 11:08:20 +0800 Subject: [PATCH 1/4] [fix](parquet) Bound Thrift and Bloom filter metadata allocations --- be/src/format/parquet/parquet_predicate.h | 66 +++++++++---- be/src/util/thrift_util.cpp | 15 +-- be/src/util/thrift_util.h | 18 +++- .../parquet/parquet_statistics_test.cpp | 95 +++++++++++++++++++ .../format/parquet/parquet_thrift_test.cpp | 14 +++ 5 files changed, 179 insertions(+), 29 deletions(-) diff --git a/be/src/format/parquet/parquet_predicate.h b/be/src/format/parquet/parquet_predicate.h index a0f6c75bdefb8b..e4730de7630d32 100644 --- a/be/src/format/parquet/parquet_predicate.h +++ b/be/src/format/parquet/parquet_predicate.h @@ -441,44 +441,76 @@ class ParquetPredicate { static Status read_bloom_filter(const tparquet::ColumnMetaData& column_meta_data, io::FileReaderSPtr file_reader, io::IOContext* io_ctx, ColumnStat* ans_stat) { - size_t size; if (!column_meta_data.__isset.bloom_filter_offset) { return Status::NotSupported("Can not use this parquet bloom filter."); } + if (column_meta_data.bloom_filter_offset < 0 || + (column_meta_data.__isset.bloom_filter_length && + column_meta_data.bloom_filter_length <= 0)) { + return Status::Corruption("Invalid Parquet bloom filter offset or declared length"); + } - if (column_meta_data.__isset.bloom_filter_length && - column_meta_data.bloom_filter_length > 0) { - size = column_meta_data.bloom_filter_length; - } else { - size = BLOOM_FILTER_MAX_HEADER_LENGTH; + const uint64_t bloom_offset = static_cast(column_meta_data.bloom_filter_offset); + if (bloom_offset >= file_reader->size()) { + return Status::Corruption("Parquet bloom filter offset exceeds file size"); } + const size_t available = file_reader->size() - bloom_offset; + const size_t declared_available = + column_meta_data.__isset.bloom_filter_length + ? std::min(column_meta_data.bloom_filter_length, available) + : available; + const size_t header_read_size = + std::min(declared_available, BLOOM_FILTER_MAX_HEADER_LENGTH); size_t bytes_read = 0; - std::vector header_buffer(size); + std::vector header_buffer(header_read_size); RETURN_IF_ERROR(file_reader->read_at(column_meta_data.bloom_filter_offset, - Slice(header_buffer.data(), size), &bytes_read, - io_ctx)); + Slice(header_buffer.data(), header_buffer.size()), + &bytes_read, io_ctx)); tparquet::BloomFilterHeader t_bloom_filter_header; uint32_t t_bloom_filter_header_size = static_cast(bytes_read); - RETURN_IF_ERROR(deserialize_thrift_msg(header_buffer.data(), &t_bloom_filter_header_size, - true, &t_bloom_filter_header)); + if (!deserialize_thrift_msg(header_buffer.data(), &t_bloom_filter_header_size, true, + &t_bloom_filter_header) + .ok()) { + return Status::Corruption("Malformed Parquet bloom filter header"); + } // TODO the bloom filter could be encrypted, too, so need to double check that this is NOT the case if (!t_bloom_filter_header.algorithm.__isset.BLOCK || !t_bloom_filter_header.compression.__isset.UNCOMPRESSED || - !t_bloom_filter_header.hash.__isset.XXHASH) { + !t_bloom_filter_header.hash.__isset.XXHASH || t_bloom_filter_header.numBytes <= 0) { return Status::NotSupported("Can not use this parquet bloom filter."); } - ans_stat->bloom_filter = std::make_unique(); + const int64_t payload_size = t_bloom_filter_header.numBytes; + if (payload_size < segment_v2::BloomFilter::MINIMUM_BYTES || + payload_size > segment_v2::BloomFilter::MAXIMUM_BYTES || payload_size % 32 != 0) { + return Status::Corruption("Invalid Parquet bloom filter payload size {}", payload_size); + } + const uint64_t total_size = + static_cast(t_bloom_filter_header_size) + payload_size; + if (total_size > available) { + return Status::Corruption("Parquet bloom filter range exceeds file size"); + } + if (column_meta_data.__isset.bloom_filter_length && + (static_cast(column_meta_data.bloom_filter_length) < total_size || + static_cast(column_meta_data.bloom_filter_length) > available)) { + return Status::Corruption("Invalid Parquet bloom filter declared length"); + } - std::vector data_buffer(t_bloom_filter_header.numBytes); + // Validate the full split-block layout before allocating or adding metadata-controlled + // offsets; the Bloom filter implementation assumes complete 32-byte blocks. + std::vector data_buffer(static_cast(payload_size)); RETURN_IF_ERROR(file_reader->read_at( - column_meta_data.bloom_filter_offset + t_bloom_filter_header_size, - Slice(data_buffer.data(), t_bloom_filter_header.numBytes), &bytes_read, io_ctx)); + static_cast(bloom_offset) + t_bloom_filter_header_size, + Slice(data_buffer.data(), data_buffer.size()), &bytes_read, io_ctx)); + if (bytes_read != data_buffer.size()) { + return Status::Corruption("Truncated Parquet bloom filter payload"); + } + ans_stat->bloom_filter = std::make_unique(); RETURN_IF_ERROR(ans_stat->bloom_filter->init( - reinterpret_cast(data_buffer.data()), t_bloom_filter_header.numBytes, + reinterpret_cast(data_buffer.data()), data_buffer.size(), segment_v2::HashStrategyPB::XX_HASH_64)); return Status::OK(); diff --git a/be/src/util/thrift_util.cpp b/be/src/util/thrift_util.cpp index b0d6d3b3093b45..90fd4450c7cc50 100644 --- a/be/src/util/thrift_util.cpp +++ b/be/src/util/thrift_util.cpp @@ -71,15 +71,16 @@ ThriftSerializer::ThriftSerializer(bool compact, int initial_buffer_size) } std::shared_ptr create_deserialize_protocol( - std::shared_ptr mem, bool compact) { + std::shared_ptr mem, bool compact, + int32_t size_limit) { if (compact) { - apache::thrift::protocol::TCompactProtocolFactoryT - tproto_factory; - return tproto_factory.getProtocol(mem); + return std::make_shared>(mem, size_limit, size_limit); } else { - apache::thrift::protocol::TBinaryProtocolFactoryT - tproto_factory; - return tproto_factory.getProtocol(mem); + return std::make_shared>(mem, size_limit, size_limit, + /*strict_read=*/false, + /*strict_write=*/true); } } diff --git a/be/src/util/thrift_util.h b/be/src/util/thrift_util.h index 73c753e2271c8f..2c1cced8db1837 100644 --- a/be/src/util/thrift_util.h +++ b/be/src/util/thrift_util.h @@ -22,9 +22,11 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -122,7 +124,8 @@ class ThriftDeserializer { // Utility to create a protocol (deserialization) object for 'mem'. std::shared_ptr create_deserialize_protocol( - std::shared_ptr mem, bool compact); + std::shared_ptr mem, bool compact, + int32_t size_limit); // Deserialize a thrift message from buf/len. buf/len must at least contain // all the bytes needed to store the thrift message. On return, len will be @@ -134,15 +137,20 @@ Status deserialize_thrift_msg(const uint8_t* buf, uint32_t* len, bool compact, // transport. TMemoryBuffer is not const-safe, although we use it in // a const-safe way, so we have to explicitly cast away the const. auto conf = std::make_shared(); - // On Thrift 0.14.0+, need use TConfiguration to raise the max message size. - // max message size is 100MB default, so make it unlimited. - conf->setMaxMessageSize(std::numeric_limits::max()); + const int32_t size_limit = + *len == 0 ? 1 + : static_cast(std::min( + *len, static_cast(std::numeric_limits::max()))); + // A serialized string or container cannot have more elements than the message has bytes. + // This bound preserves valid input while rejecting hostile lengths before generated readers + // allocate memory for them. + conf->setMaxMessageSize(size_limit); std::shared_ptr tmem_transport( new apache::thrift::transport::TMemoryBuffer( const_cast(buf), *len, apache::thrift::transport::TMemoryBuffer::OBSERVE, conf)); std::shared_ptr tproto = - create_deserialize_protocol(tmem_transport, compact); + create_deserialize_protocol(tmem_transport, compact, size_limit); try { deserialized_msg->read(tproto.get()); diff --git a/be/test/format/parquet/parquet_statistics_test.cpp b/be/test/format/parquet/parquet_statistics_test.cpp index d52320c51698e4..8748d5e8c64aca 100644 --- a/be/test/format/parquet/parquet_statistics_test.cpp +++ b/be/test/format/parquet/parquet_statistics_test.cpp @@ -17,16 +17,111 @@ #include +#include +#include +#include #include +#include #include "format/parquet/parquet_predicate.h" +#include "util/thrift_util.h" namespace doris { +namespace { + +class BloomFilterFileReader final : public io::FileReader { +public: + explicit BloomFilterFileReader(std::vector data) : _data(std::move(data)) {} + + Status close() override { + _closed = true; + return Status::OK(); + } + + const io::Path& path() const override { return _path; } + size_t size() const override { return _data.size(); } + bool closed() const override { return _closed; } + int64_t mtime() const override { return 0; } + +protected: + Status read_at_impl(size_t offset, Slice result, size_t* bytes_read, + const io::IOContext* io_ctx) override { + if (offset > _data.size()) { + return Status::IOError("Out of bounds"); + } + *bytes_read = std::min(result.size, _data.size() - offset); + memcpy(result.data, _data.data() + offset, *bytes_read); + return Status::OK(); + } + +private: + std::vector _data; + io::Path _path = "parquet_bloom_filter_test"; + bool _closed = false; +}; + +Status read_test_bloom_filter(int32_t header_payload_size, size_t actual_payload_size, + int32_t declared_length_adjustment = 0) { + tparquet::BloomFilterAlgorithm algorithm; + algorithm.__set_BLOCK(tparquet::SplitBlockAlgorithm()); + tparquet::BloomFilterHash hash; + hash.__set_XXHASH(tparquet::XxHash()); + tparquet::BloomFilterCompression compression; + compression.__set_UNCOMPRESSED(tparquet::Uncompressed()); + tparquet::BloomFilterHeader header; + header.__set_numBytes(header_payload_size); + header.__set_algorithm(algorithm); + header.__set_hash(hash); + header.__set_compression(compression); + + std::vector file_bytes; + ThriftSerializer serializer(/*compact=*/true, /*initial_buffer_size=*/64); + RETURN_IF_ERROR(serializer.serialize(&header, &file_bytes)); + const size_t header_size = file_bytes.size(); + file_bytes.resize(header_size + actual_payload_size); + + tparquet::ColumnMetaData metadata; + metadata.__set_bloom_filter_offset(0); + metadata.__set_bloom_filter_length(static_cast(file_bytes.size()) + + declared_length_adjustment); + auto reader = std::make_shared(std::move(file_bytes)); + ParquetPredicate::ColumnStat stat; + return ParquetPredicate::read_bloom_filter(metadata, reader, nullptr, &stat); +} + +} // namespace + class ParquetStatisticsTest : public testing::Test { public: ParquetStatisticsTest() = default; }; +TEST_F(ParquetStatisticsTest, reject_truncated_bloom_filter_payload) { + // The reader may legally return a short read at EOF, so accepting it would initialize a + // Bloom filter whose missing bytes came from zero-filled process memory. + EXPECT_FALSE( + read_test_bloom_filter(/*header_payload_size=*/64, /*actual_payload_size=*/32).ok()); +} + +TEST_F(ParquetStatisticsTest, reject_invalid_bloom_filter_block_sizes) { + EXPECT_FALSE( + read_test_bloom_filter(/*header_payload_size=*/16, /*actual_payload_size=*/16).ok()); + EXPECT_FALSE( + read_test_bloom_filter(/*header_payload_size=*/33, /*actual_payload_size=*/33).ok()); +} + +TEST_F(ParquetStatisticsTest, reject_nonpositive_bloom_filter_declared_length) { + const int32_t declared_length_adjustment = -1000; + EXPECT_FALSE(read_test_bloom_filter(/*header_payload_size=*/32, /*actual_payload_size=*/32, + declared_length_adjustment) + .ok()); +} + +TEST_F(ParquetStatisticsTest, accept_valid_bloom_filter_layout) { + EXPECT_TRUE( + read_test_bloom_filter(/*header_payload_size=*/32, /*actual_payload_size=*/32).ok()); +} + TEST_F(ParquetStatisticsTest, test_try_read_old_utf8_stats) { // [, bcé]: min is empty, max starts with ASCII { diff --git a/be/test/format/parquet/parquet_thrift_test.cpp b/be/test/format/parquet/parquet_thrift_test.cpp index 0fe101db598138..daa1c5135173a9 100644 --- a/be/test/format/parquet/parquet_thrift_test.cpp +++ b/be/test/format/parquet/parquet_thrift_test.cpp @@ -60,6 +60,7 @@ #include "io/fs/local_file_system.h" #include "runtime/descriptors.h" #include "util/slice.h" +#include "util/thrift_util.h" #include "util/timezone_utils.h" namespace doris { @@ -71,6 +72,19 @@ class ParquetThriftReaderTest : public testing::Test { void TearDown() override { TimezoneUtils::clear_timezone_caches(); } }; +TEST_F(ParquetThriftReaderTest, reject_compact_container_larger_than_input) { + // A tiny payload must not make generated Thrift code resize a container to an + // attacker-controlled length before the transport discovers that the bytes are absent. + std::vector compact_metadata {0x29, 0xfc, 0x88, 0x27, 0x00}; + uint32_t length = compact_metadata.size(); + tparquet::FileMetaData metadata; + + Status status = deserialize_thrift_msg(compact_metadata.data(), &length, true, &metadata); + + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(metadata.schema.empty()); +} + TEST_F(ParquetThriftReaderTest, normal) { auto local_fs = io::global_local_filesystem(); io::FileReaderSPtr reader; From 9a3671f98622d4e666698f5f2741a92723cbe67d Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 17 Aug 2026 13:04:59 +0800 Subject: [PATCH 2/4] [fix](parquet) Address metadata allocation review findings ### What problem does this PR solve? Issue Number: None Problem Summary: Parquet Bloom readers could accept contradictory declared lengths, allocate duplicate untracked payload buffers, and retain multiple large filters. Thrift generated readers could resize decoded containers before task memory admission. This change requires exact Bloom ranges, reads into tracked single-owner buffers, caps legacy cache retention, and reserves decoded Thrift container storage before resize. ### Release note Reject malformed Parquet Bloom metadata conservatively and enforce task memory admission for metadata decoding. ### Check List (For Author) - Test: Unit Test - Behavior changed: Yes, malformed metadata now falls back conservatively and excessive decoded metadata is rejected before allocation. - Does this need documentation: No --- .../parquet_block_split_bloom_filter.cpp | 39 ++++-- .../parquet_block_split_bloom_filter.h | 6 + be/src/format/parquet/parquet_predicate.h | 21 ++- be/src/format/parquet/vparquet_reader.cpp | 10 ++ .../format_v2/parquet/parquet_statistics.cpp | 16 ++- .../native/block_split_bloom_filter.cpp | 32 ++++- .../reader/native/block_split_bloom_filter.h | 6 + be/src/util/thrift_util.cpp | 122 +++++++++++++++++- be/src/util/thrift_util.h | 10 +- .../parquet/parquet_statistics_test.cpp | 52 +++++++- .../format/parquet/parquet_thrift_test.cpp | 115 +++++++++++++++++ .../parquet/parquet_statistics_test.cpp | 24 +++- 12 files changed, 409 insertions(+), 44 deletions(-) diff --git a/be/src/format/parquet/parquet_block_split_bloom_filter.cpp b/be/src/format/parquet/parquet_block_split_bloom_filter.cpp index 7f93681df03699..b8cbacbfa9be13 100644 --- a/be/src/format/parquet/parquet_block_split_bloom_filter.cpp +++ b/be/src/format/parquet/parquet_block_split_bloom_filter.cpp @@ -23,6 +23,22 @@ namespace doris { +ParquetBlockSplitBloomFilter::~ParquetBlockSplitBloomFilter() { + if (_data == nullptr) { + return; + } + if (_is_write) { + g_write_bloom_filter_total_bytes << -static_cast(_size); + g_write_bloom_filter_num << -1; + } else { + g_read_bloom_filter_total_bytes << -static_cast(_size); + g_read_bloom_filter_num << -1; + } + g_total_bloom_filter_total_bytes << -static_cast(_size); + // The derived owner uses Doris's tracked allocator, so the base must not delete this view. + _data = nullptr; +} + // for write Status ParquetBlockSplitBloomFilter::init(uint64_t filter_size, segment_v2::HashStrategyPB strategy) { @@ -37,7 +53,8 @@ Status ParquetBlockSplitBloomFilter::init(uint64_t filter_size, } _num_bytes = filter_size; _size = _num_bytes; - _data = new char[_size]; + _owned_data = make_unique_buffer(_size); + _data = _owned_data.get(); memset(_data, 0, _size); _has_null = nullptr; _is_write = true; @@ -51,10 +68,20 @@ Status ParquetBlockSplitBloomFilter::init(uint64_t filter_size, // use deep copy to acquire the data Status ParquetBlockSplitBloomFilter::init(const char* buf, size_t size, segment_v2::HashStrategyPB strategy) { - if (size <= 1) { + if (buf == nullptr || size <= 1) { return Status::InvalidArgument("invalid size:{}", size); } DCHECK(size > 1); + RETURN_IF_ERROR(init_for_read(size, strategy)); + memcpy(_data, buf, size); + return Status::OK(); +} + +Status ParquetBlockSplitBloomFilter::init_for_read(size_t size, + segment_v2::HashStrategyPB strategy) { + if (size <= 1) { + return Status::InvalidArgument("invalid size:{}", size); + } if (strategy == XX_HASH_64) { _hash_func = [](const void* buf, const int64_t len, const uint64_t seed, void* out) { auto h = @@ -64,12 +91,8 @@ Status ParquetBlockSplitBloomFilter::init(const char* buf, size_t size, } else { return Status::InvalidArgument("invalid strategy:{}", strategy); } - if (buf == nullptr) { - return Status::InvalidArgument("buf is nullptr"); - } - - _data = new char[size]; - memcpy(_data, buf, size); + _owned_data = make_unique_buffer(size); + _data = _owned_data.get(); _size = size; _num_bytes = _size; _has_null = nullptr; diff --git a/be/src/format/parquet/parquet_block_split_bloom_filter.h b/be/src/format/parquet/parquet_block_split_bloom_filter.h index b27e9bab8fe211..186ebb1931dbaa 100644 --- a/be/src/format/parquet/parquet_block_split_bloom_filter.h +++ b/be/src/format/parquet/parquet_block_split_bloom_filter.h @@ -19,6 +19,7 @@ #include +#include "core/custom_allocator.h" #include "storage/index/bloom_filter/bloom_filter.h" namespace doris { @@ -36,8 +37,11 @@ namespace doris { // https://parquet.apache.org/docs/file-format/bloomfilter/ class ParquetBlockSplitBloomFilter : public segment_v2::BloomFilter { public: + ~ParquetBlockSplitBloomFilter() override; Status init(uint64_t filter_size, segment_v2::HashStrategyPB strategy) override; Status init(const char* buf, size_t size, segment_v2::HashStrategyPB strategy) override; + Status init_for_read(size_t size, segment_v2::HashStrategyPB strategy); + char* mutable_data() { return _data; } void add_bytes(const char* buf, size_t size) override; bool test_bytes(const char* buf, size_t size) const override; void set_has_null(bool has_null) override; @@ -62,6 +66,8 @@ class ParquetBlockSplitBloomFilter : public segment_v2::BloomFilter { }; private: + DorisUniqueBufferPtr _owned_data; + void _set_masks(uint32_t key, BlockMask& block_mask) const { for (int i = 0; i < BITS_SET_PER_BLOCK; ++i) { block_mask.item[i] = key * SALT[i]; diff --git a/be/src/format/parquet/parquet_predicate.h b/be/src/format/parquet/parquet_predicate.h index e4730de7630d32..ee0f3b60d83059 100644 --- a/be/src/format/parquet/parquet_predicate.h +++ b/be/src/format/parquet/parquet_predicate.h @@ -492,26 +492,25 @@ class ParquetPredicate { if (total_size > available) { return Status::Corruption("Parquet bloom filter range exceeds file size"); } + const auto expected_declared_length = static_cast(total_size); if (column_meta_data.__isset.bloom_filter_length && - (static_cast(column_meta_data.bloom_filter_length) < total_size || - static_cast(column_meta_data.bloom_filter_length) > available)) { + column_meta_data.bloom_filter_length != expected_declared_length) { return Status::Corruption("Invalid Parquet bloom filter declared length"); } - // Validate the full split-block layout before allocating or adding metadata-controlled - // offsets; the Bloom filter implementation assumes complete 32-byte blocks. - std::vector data_buffer(static_cast(payload_size)); + auto bloom_filter = std::make_unique(); + // Read directly into one tracked allocation so a valid maximum-size filter is admitted + // against the task budget without doubling its peak memory during initialization. + RETURN_IF_ERROR(bloom_filter->init_for_read(static_cast(payload_size), + segment_v2::HashStrategyPB::XX_HASH_64)); RETURN_IF_ERROR(file_reader->read_at( static_cast(bloom_offset) + t_bloom_filter_header_size, - Slice(data_buffer.data(), data_buffer.size()), &bytes_read, io_ctx)); - if (bytes_read != data_buffer.size()) { + Slice(bloom_filter->mutable_data(), bloom_filter->size()), &bytes_read, io_ctx)); + if (bytes_read != bloom_filter->size()) { return Status::Corruption("Truncated Parquet bloom filter payload"); } - ans_stat->bloom_filter = std::make_unique(); - RETURN_IF_ERROR(ans_stat->bloom_filter->init( - reinterpret_cast(data_buffer.data()), data_buffer.size(), - segment_v2::HashStrategyPB::XX_HASH_64)); + ans_stat->bloom_filter = std::move(bloom_filter); return Status::OK(); } diff --git a/be/src/format/parquet/vparquet_reader.cpp b/be/src/format/parquet/vparquet_reader.cpp index b5c4c01be409bd..b1b233aa0d570c 100644 --- a/be/src/format/parquet/vparquet_reader.cpp +++ b/be/src/format/parquet/vparquet_reader.cpp @@ -1507,6 +1507,8 @@ Status ParquetReader::_process_column_stat_filter( // Cache bloom filters for each column to avoid reading the same bloom filter multiple times // when there are multiple predicates on the same column std::unordered_map> bloom_filter_cache; + constexpr size_t MAX_CACHED_BLOOM_FILTER_BYTES = 16 * 1024 * 1024; + size_t cached_bloom_filter_bytes = 0; // Initialize output parameters *filtered_by_min_max = false; @@ -1565,6 +1567,7 @@ Status ParquetReader::_process_column_stat_filter( auto cache_iter = bloom_filter_cache.find(parquet_col_id); if (cache_iter != bloom_filter_cache.end()) { // Bloom filter already loaded for this column, reuse it + cached_bloom_filter_bytes -= cache_iter->second->size(); stat->bloom_filter = std::move(cache_iter->second); bloom_filter_cache.erase(cache_iter); return stat->bloom_filter != nullptr; @@ -1607,6 +1610,12 @@ Status ParquetReader::_process_column_stat_filter( // After evaluating, if the bloom filter was used, cache it for subsequent predicates if (stat.bloom_filter) { + // Large filters remain tracked but are not retained across predicates; this bounds the + // row-group cache independently of the number of predicate columns. + if (stat.bloom_filter->size() > + MAX_CACHED_BLOOM_FILTER_BYTES - cached_bloom_filter_bytes) { + continue; + } // Find the column id for caching for (auto* slot : _tuple_descriptor->slots()) { if (_table_info_node_ptr->children_column_exists(slot->col_name())) { @@ -1616,6 +1625,7 @@ Status ParquetReader::_process_column_stat_filter( _file_metadata->schema().get_column(file_col_name); int parquet_col_id = col_schema->physical_column_index; if (stat.col_schema == col_schema) { + cached_bloom_filter_bytes += stat.bloom_filter->size(); bloom_filter_cache[parquet_col_id] = std::move(stat.bloom_filter); break; } diff --git a/be/src/format_v2/parquet/parquet_statistics.cpp b/be/src/format_v2/parquet/parquet_statistics.cpp index a894601a1607fa..ca3e5d621de912 100644 --- a/be/src/format_v2/parquet/parquet_statistics.cpp +++ b/be/src/format_v2/parquet/parquet_statistics.cpp @@ -75,7 +75,7 @@ Status validate_native_bloom_filter_layout(int64_t offset, uint32_t header_size, } if (declared_length >= 0) { const uint64_t unsigned_declared_length = static_cast(declared_length); - if (unsigned_declared_length < total_size || + if (unsigned_declared_length != total_size || unsigned_declared_length > file_size - unsigned_offset) { return Status::Corruption( "Parquet Bloom filter requires {} bytes, metadata declares {}, file has {}", @@ -190,15 +190,17 @@ Status read_native_bloom_filter(const tparquet::ColumnMetaData& metadata, metadata.__isset.bloom_filter_length ? metadata.bloom_filter_length : -1, file->size())); - std::vector data(cast_set(header.numBytes)); + auto bloom_filter = std::make_unique(); + // Read directly into one tracked allocation so maximum-size filters cannot bypass task + // admission or temporarily require a second payload-sized allocation. + RETURN_IF_ERROR(bloom_filter->init_for_read(cast_set(header.numBytes), + segment_v2::HashStrategyPB::XX_HASH_64)); RETURN_IF_ERROR(file->read_at(static_cast(metadata.bloom_filter_offset) + header_size, - Slice(data.data(), data.size()), &bytes_read, io_ctx)); - if (bytes_read != data.size()) { + Slice(bloom_filter->mutable_data(), bloom_filter->size()), + &bytes_read, io_ctx)); + if (bytes_read != bloom_filter->size()) { return Status::Corruption("Truncated Parquet Bloom filter payload"); } - auto bloom_filter = std::make_unique(); - RETURN_IF_ERROR(bloom_filter->init(reinterpret_cast(data.data()), data.size(), - segment_v2::HashStrategyPB::XX_HASH_64)); *result = std::move(bloom_filter); return Status::OK(); } diff --git a/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.cpp b/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.cpp index 3421ecdd296892..151e885093caa1 100644 --- a/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.cpp +++ b/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.cpp @@ -37,11 +37,28 @@ Status set_hash_strategy(segment_v2::HashStrategyPB strategy, } } // namespace +BlockSplitBloomFilter::~BlockSplitBloomFilter() { + if (_data == nullptr) { + return; + } + if (_is_write) { + segment_v2::g_write_bloom_filter_total_bytes << -static_cast(_size); + segment_v2::g_write_bloom_filter_num << -1; + } else { + segment_v2::g_read_bloom_filter_total_bytes << -static_cast(_size); + segment_v2::g_read_bloom_filter_num << -1; + } + segment_v2::g_total_bloom_filter_total_bytes << -static_cast(_size); + // The derived owner uses Doris's tracked allocator, so the base must not delete this view. + _data = nullptr; +} + Status BlockSplitBloomFilter::init(uint64_t filter_size, segment_v2::HashStrategyPB strategy) { RETURN_IF_ERROR(set_hash_strategy(strategy, &_hash_func)); _num_bytes = filter_size; _size = _num_bytes; - _data = new char[_size]; + _owned_data = make_unique_buffer(_size); + _data = _owned_data.get(); memset(_data, 0, _size); _has_null = nullptr; _is_write = true; @@ -56,9 +73,18 @@ Status BlockSplitBloomFilter::init(const char* buf, size_t size, if (buf == nullptr || size <= 1) { return Status::InvalidArgument("Invalid Parquet Bloom filter buffer of size {}", size); } - RETURN_IF_ERROR(set_hash_strategy(strategy, &_hash_func)); - _data = new char[size]; + RETURN_IF_ERROR(init_for_read(size, strategy)); memcpy(_data, buf, size); + return Status::OK(); +} + +Status BlockSplitBloomFilter::init_for_read(size_t size, segment_v2::HashStrategyPB strategy) { + if (size <= 1) { + return Status::InvalidArgument("Invalid Parquet Bloom filter buffer of size {}", size); + } + RETURN_IF_ERROR(set_hash_strategy(strategy, &_hash_func)); + _owned_data = make_unique_buffer(size); + _data = _owned_data.get(); _size = size; _num_bytes = size; _has_null = nullptr; diff --git a/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.h b/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.h index 38dc97712ac46c..cf64815ddee5a0 100644 --- a/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.h +++ b/be/src/format_v2/parquet/reader/native/block_split_bloom_filter.h @@ -20,6 +20,7 @@ #include #include +#include "core/custom_allocator.h" #include "storage/index/bloom_filter/bloom_filter.h" namespace doris::format::parquet::native { @@ -28,8 +29,11 @@ namespace doris::format::parquet::native { // legacy Parquet reader merely to evaluate footer Bloom filters. class BlockSplitBloomFilter final : public segment_v2::BloomFilter { public: + ~BlockSplitBloomFilter() override; Status init(uint64_t filter_size, segment_v2::HashStrategyPB strategy) override; Status init(const char* buf, size_t size, segment_v2::HashStrategyPB strategy) override; + Status init_for_read(size_t size, segment_v2::HashStrategyPB strategy); + char* mutable_data() { return _data; } void add_bytes(const char* buf, size_t size) override; bool test_bytes(const char* buf, size_t size) const override; void set_has_null(bool has_null) override; @@ -38,6 +42,8 @@ class BlockSplitBloomFilter final : public segment_v2::BloomFilter { bool test_hash(uint64_t hash) const override; private: + DorisUniqueBufferPtr _owned_data; + static constexpr int BYTES_PER_BLOCK = 32; static constexpr int BITS_SET_PER_BLOCK = 8; static constexpr uint32_t SALT[BITS_SET_PER_BLOCK] = {0x47b6137bU, 0x44974d91U, 0x8824ad5bU, diff --git a/be/src/util/thrift_util.cpp b/be/src/util/thrift_util.cpp index 90fd4450c7cc50..afb84ec190f3b7 100644 --- a/be/src/util/thrift_util.cpp +++ b/be/src/util/thrift_util.cpp @@ -20,14 +20,18 @@ #include #include #include +#include +#include #include #include // IWYU pragma: no_include #include // IWYU pragma: keep +#include #include #include "common/compiler_util.h" // IWYU pragma: keep #include "common/logging.h" +#include "runtime/thread_context.h" #include "storage/tablet/tablet_schema.h" #include "storage/tablet_info.h" #include "util/thrift_server.h" @@ -56,6 +60,118 @@ class TProtocol; #include namespace doris { +namespace { + +constexpr size_t DECODED_THRIFT_STRUCT_RESERVATION_BYTES = 1024; + +size_t decoded_thrift_value_reservation(apache::thrift::protocol::TType type) { + using apache::thrift::protocol::T_BOOL; + using apache::thrift::protocol::T_BYTE; + using apache::thrift::protocol::T_DOUBLE; + using apache::thrift::protocol::T_I16; + using apache::thrift::protocol::T_I32; + using apache::thrift::protocol::T_I64; + using apache::thrift::protocol::T_LIST; + using apache::thrift::protocol::T_MAP; + using apache::thrift::protocol::T_SET; + using apache::thrift::protocol::T_STRING; + using apache::thrift::protocol::T_STRUCT; + switch (type) { + case T_BOOL: + case T_BYTE: + return 1; + case T_I16: + return sizeof(int16_t); + case T_I32: + return sizeof(int32_t); + case T_I64: + case T_DOUBLE: + return sizeof(int64_t); + case T_STRING: + return sizeof(std::string); + case T_LIST: + case T_SET: + case T_MAP: + return sizeof(std::vector); + case T_STRUCT: + // Generated structs vary in size. Reserving a conservative inline object budget keeps + // their eager vector resize inside task admission; the serialized-size reservation covers + // their dynamic field payloads. + return DECODED_THRIFT_STRUCT_RESERVATION_BYTES; + default: + return 1; + } +} + +class MemoryBudgetProtocol final : public apache::thrift::protocol::TProtocolDecorator { +public: + MemoryBudgetProtocol(std::shared_ptr protocol, + int32_t serialized_size) + : TProtocolDecorator(std::move(protocol)), + _memory_manager(thread_context()->thread_mem_tracker_mgr.get()), + _prior_reservation(_memory_manager->take_reserved_memory()) { + reserve_or_throw(static_cast(serialized_size), /*restore_prior_on_failure=*/true); + } + + ~MemoryBudgetProtocol() override { + _memory_manager->shrink_reserved(); + _memory_manager->adopt_reserved_memory(std::move(_prior_reservation)); + } + + uint32_t readMapBegin_virt(apache::thrift::protocol::TType& key_type, + apache::thrift::protocol::TType& value_type, + uint32_t& size) override { + const uint32_t consumed = TProtocolDecorator::readMapBegin_virt(key_type, value_type, size); + const uint32_t count = size; + const size_t element_size = decoded_thrift_value_reservation(key_type) + + decoded_thrift_value_reservation(value_type) + + 4 * sizeof(void*); + reserve_container(count, element_size); + return consumed; + } + + uint32_t readListBegin_virt(apache::thrift::protocol::TType& element_type, + uint32_t& size) override { + const uint32_t consumed = TProtocolDecorator::readListBegin_virt(element_type, size); + reserve_container(size, decoded_thrift_value_reservation(element_type)); + return consumed; + } + + uint32_t readSetBegin_virt(apache::thrift::protocol::TType& element_type, + uint32_t& size) override { + const uint32_t consumed = TProtocolDecorator::readSetBegin_virt(element_type, size); + reserve_container(size, decoded_thrift_value_reservation(element_type) + 4 * sizeof(void*)); + return consumed; + } + +private: + void reserve_container(uint32_t count, size_t element_size) { + if (count > std::numeric_limits::max() / element_size) { + throw apache::thrift::protocol::TProtocolException( + apache::thrift::protocol::TProtocolException::SIZE_LIMIT, + "Decoded Thrift container size overflows"); + } + reserve_or_throw(static_cast(count) * element_size, + /*restore_prior_on_failure=*/false); + } + + void reserve_or_throw(size_t bytes, bool restore_prior_on_failure) { + const Status status = _memory_manager->try_reserve(static_cast(bytes)); + if (status.ok()) { + return; + } + if (restore_prior_on_failure) { + _memory_manager->adopt_reserved_memory(std::move(_prior_reservation)); + } + throw apache::thrift::protocol::TProtocolException( + apache::thrift::protocol::TProtocolException::SIZE_LIMIT, status.to_string()); + } + + ThreadMemTrackerMgr* _memory_manager; + ReservedMemoryToken _prior_reservation; +}; + +} // namespace ThriftSerializer::ThriftSerializer(bool compact, int initial_buffer_size) : _mem_buffer(new apache::thrift::transport::TMemoryBuffer(initial_buffer_size)) { @@ -73,15 +189,17 @@ ThriftSerializer::ThriftSerializer(bool compact, int initial_buffer_size) std::shared_ptr create_deserialize_protocol( std::shared_ptr mem, bool compact, int32_t size_limit) { + std::shared_ptr protocol; if (compact) { - return std::make_shared>(mem, size_limit, size_limit); } else { - return std::make_shared>(mem, size_limit, size_limit, /*strict_read=*/false, /*strict_write=*/true); } + return std::make_shared(std::move(protocol), size_limit); } // Comparator for THostPorts. Thrift declares this (in gen-cpp/Types_types.h) but diff --git a/be/src/util/thrift_util.h b/be/src/util/thrift_util.h index 2c1cced8db1837..208d9f317e7e4b 100644 --- a/be/src/util/thrift_util.h +++ b/be/src/util/thrift_util.h @@ -141,18 +141,16 @@ Status deserialize_thrift_msg(const uint8_t* buf, uint32_t* len, bool compact, *len == 0 ? 1 : static_cast(std::min( *len, static_cast(std::numeric_limits::max()))); - // A serialized string or container cannot have more elements than the message has bytes. - // This bound preserves valid input while rejecting hostile lengths before generated readers - // allocate memory for them. + // Wire-size limits reject impossible prefixes, while the protocol wrapper separately reserves + // decoded container storage before generated readers resize their C++ containers. conf->setMaxMessageSize(size_limit); std::shared_ptr tmem_transport( new apache::thrift::transport::TMemoryBuffer( const_cast(buf), *len, apache::thrift::transport::TMemoryBuffer::OBSERVE, conf)); - std::shared_ptr tproto = - create_deserialize_protocol(tmem_transport, compact, size_limit); - try { + std::shared_ptr tproto = + create_deserialize_protocol(tmem_transport, compact, size_limit); deserialized_msg->read(tproto.get()); } catch (std::exception& e) { return Status::InternalError("Couldn't deserialize thrift msg:\n{}", e.what()); diff --git a/be/test/format/parquet/parquet_statistics_test.cpp b/be/test/format/parquet/parquet_statistics_test.cpp index 8748d5e8c64aca..e542b029c2ec95 100644 --- a/be/test/format/parquet/parquet_statistics_test.cpp +++ b/be/test/format/parquet/parquet_statistics_test.cpp @@ -31,7 +31,9 @@ namespace { class BloomFilterFileReader final : public io::FileReader { public: - explicit BloomFilterFileReader(std::vector data) : _data(std::move(data)) {} + explicit BloomFilterFileReader(std::vector data, size_t logical_size = 0) + : _data(std::move(data)), + _logical_size(logical_size == 0 ? _data.size() : logical_size) {} Status close() override { _closed = true; @@ -39,9 +41,10 @@ class BloomFilterFileReader final : public io::FileReader { } const io::Path& path() const override { return _path; } - size_t size() const override { return _data.size(); } + size_t size() const override { return _logical_size; } bool closed() const override { return _closed; } int64_t mtime() const override { return 0; } + bool returned_short_nonzero_offset_read() const { return _returned_short_nonzero_offset_read; } protected: Status read_at_impl(size_t offset, Slice result, size_t* bytes_read, @@ -51,17 +54,22 @@ class BloomFilterFileReader final : public io::FileReader { } *bytes_read = std::min(result.size, _data.size() - offset); memcpy(result.data, _data.data() + offset, *bytes_read); + _returned_short_nonzero_offset_read |= offset > 0 && *bytes_read != result.size; return Status::OK(); } private: std::vector _data; + size_t _logical_size; io::Path _path = "parquet_bloom_filter_test"; bool _closed = false; + bool _returned_short_nonzero_offset_read = false; }; Status read_test_bloom_filter(int32_t header_payload_size, size_t actual_payload_size, - int32_t declared_length_adjustment = 0) { + int32_t declared_length_adjustment = 0, + size_t logical_payload_size = 0, bool* returned_short_read = nullptr, + bool* installed_bloom_filter = nullptr) { tparquet::BloomFilterAlgorithm algorithm; algorithm.__set_BLOCK(tparquet::SplitBlockAlgorithm()); tparquet::BloomFilterHash hash; @@ -84,9 +92,18 @@ Status read_test_bloom_filter(int32_t header_payload_size, size_t actual_payload metadata.__set_bloom_filter_offset(0); metadata.__set_bloom_filter_length(static_cast(file_bytes.size()) + declared_length_adjustment); - auto reader = std::make_shared(std::move(file_bytes)); + const size_t logical_size = + logical_payload_size == 0 ? file_bytes.size() : header_size + logical_payload_size; + auto reader = std::make_shared(std::move(file_bytes), logical_size); ParquetPredicate::ColumnStat stat; - return ParquetPredicate::read_bloom_filter(metadata, reader, nullptr, &stat); + Status status = ParquetPredicate::read_bloom_filter(metadata, reader, nullptr, &stat); + if (returned_short_read != nullptr) { + *returned_short_read = reader->returned_short_nonzero_offset_read(); + } + if (installed_bloom_filter != nullptr) { + *installed_bloom_filter = stat.bloom_filter != nullptr; + } + return status; } } // namespace @@ -99,8 +116,31 @@ class ParquetStatisticsTest : public testing::Test { TEST_F(ParquetStatisticsTest, reject_truncated_bloom_filter_payload) { // The reader may legally return a short read at EOF, so accepting it would initialize a // Bloom filter whose missing bytes came from zero-filled process memory. + bool returned_short_read = false; + bool installed_bloom_filter = true; + EXPECT_FALSE(read_test_bloom_filter(/*header_payload_size=*/64, /*actual_payload_size=*/32, + /*declared_length_adjustment=*/32, + /*logical_payload_size=*/64, &returned_short_read, + &installed_bloom_filter) + .ok()); + EXPECT_TRUE(returned_short_read); + EXPECT_FALSE(installed_bloom_filter); +} + +TEST_F(ParquetStatisticsTest, reject_bloom_filter_range_beyond_file) { + bool returned_short_read = false; + EXPECT_FALSE(read_test_bloom_filter(/*header_payload_size=*/64, /*actual_payload_size=*/32, + /*declared_length_adjustment=*/0, + /*logical_payload_size=*/0, &returned_short_read) + .ok()); + EXPECT_FALSE(returned_short_read); +} + +TEST_F(ParquetStatisticsTest, reject_declared_bloom_filter_length_mismatch) { + // A present length describes exactly one header and payload. Treating it as an upper bound can + // reinterpret a multi-block filter as a smaller filter and cause false-negative pruning. EXPECT_FALSE( - read_test_bloom_filter(/*header_payload_size=*/64, /*actual_payload_size=*/32).ok()); + read_test_bloom_filter(/*header_payload_size=*/32, /*actual_payload_size=*/64).ok()); } TEST_F(ParquetStatisticsTest, reject_invalid_bloom_filter_block_sizes) { diff --git a/be/test/format/parquet/parquet_thrift_test.cpp b/be/test/format/parquet/parquet_thrift_test.cpp index daa1c5135173a9..8d1fcb40085da5 100644 --- a/be/test/format/parquet/parquet_thrift_test.cpp +++ b/be/test/format/parquet/parquet_thrift_test.cpp @@ -25,6 +25,7 @@ #include #include +#include #include #include #include @@ -59,11 +60,66 @@ #include "io/fs/file_reader_writer_fwd.h" #include "io/fs/local_file_system.h" #include "runtime/descriptors.h" +#include "runtime/memory/mem_tracker_limiter.h" +#include "runtime/thread_context.h" #include "util/slice.h" #include "util/thrift_util.h" #include "util/timezone_utils.h" namespace doris { +namespace { + +struct LargeDecodedElement { + std::array bytes {}; +}; + +struct ThriftContainerProbe { + uint32_t read(apache::thrift::protocol::TProtocol* protocol) { + apache::thrift::protocol::TType element_type; + uint32_t size = 0; + uint32_t consumed = protocol->readListBegin(element_type, size); + elements.resize(size); + return consumed + protocol->readListEnd(); + } + + std::vector elements; +}; + +struct ThriftStringProbe { + uint32_t read(apache::thrift::protocol::TProtocol* protocol) { + return protocol->readString(value); + } + + std::string value; +}; + +std::vector thrift_list_bytes(bool compact, uint32_t count, size_t total_size) { + std::vector bytes; + if (compact) { + bytes = {0xfc, static_cast(count)}; + } else { + bytes = {static_cast(apache::thrift::protocol::T_STRUCT), + static_cast(count >> 24), static_cast(count >> 16), + static_cast(count >> 8), static_cast(count)}; + } + bytes.resize(std::max(bytes.size(), total_size)); + return bytes; +} + +std::vector thrift_string_bytes(bool compact, std::string_view value) { + std::vector bytes; + if (compact) { + bytes.push_back(static_cast(value.size())); + } else { + const uint32_t size = value.size(); + bytes = {static_cast(size >> 24), static_cast(size >> 16), + static_cast(size >> 8), static_cast(size)}; + } + bytes.insert(bytes.end(), value.begin(), value.end()); + return bytes; +} + +} // namespace class ParquetThriftReaderTest : public testing::Test { public: @@ -85,6 +141,65 @@ TEST_F(ParquetThriftReaderTest, reject_compact_container_larger_than_input) { EXPECT_TRUE(metadata.schema.empty()); } +TEST_F(ParquetThriftReaderTest, reject_decoded_container_that_exceeds_task_budget) { + constexpr uint32_t element_count = 64; + constexpr int64_t memory_limit = 4 * 1024; + for (const bool compact : {true, false}) { + auto bytes = thrift_list_bytes(compact, element_count, element_count + 8); + uint32_t length = bytes.size(); + ThriftContainerProbe probe; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "ThriftContainerProbe", memory_limit); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + + Status status = deserialize_thrift_msg(bytes.data(), &length, compact, &probe); + + EXPECT_FALSE(status.ok()) << "compact=" << compact; + EXPECT_TRUE(probe.elements.empty()) << "compact=" << compact; + } +} + +TEST_F(ParquetThriftReaderTest, bound_strings_and_accept_valid_controls_for_both_protocols) { + for (const bool compact : {true, false}) { + auto malformed_container = thrift_list_bytes(compact, /*count=*/64, /*total_size=*/8); + uint32_t malformed_container_length = malformed_container.size(); + ThriftContainerProbe malformed_container_probe; + EXPECT_FALSE(deserialize_thrift_msg(malformed_container.data(), &malformed_container_length, + compact, &malformed_container_probe) + .ok()) + << "compact=" << compact; + EXPECT_TRUE(malformed_container_probe.elements.empty()) << "compact=" << compact; + + std::vector malformed_string = + compact ? std::vector {0xff, 0xff, 0xff, 0xff, 0x07} + : std::vector {0x7f, 0xff, 0xff, 0xff}; + uint32_t malformed_length = malformed_string.size(); + ThriftStringProbe malformed_probe; + EXPECT_FALSE(deserialize_thrift_msg(malformed_string.data(), &malformed_length, compact, + &malformed_probe) + .ok()) + << "compact=" << compact; + + auto valid_string = thrift_string_bytes(compact, "valid"); + uint32_t valid_string_length = valid_string.size(); + ThriftStringProbe valid_string_probe; + EXPECT_TRUE(deserialize_thrift_msg(valid_string.data(), &valid_string_length, compact, + &valid_string_probe) + .ok()) + << "compact=" << compact; + EXPECT_EQ(valid_string_probe.value, "valid"); + + auto valid_container = thrift_list_bytes(compact, /*count=*/1, /*total_size=*/8); + uint32_t valid_container_length = valid_container.size(); + ThriftContainerProbe valid_container_probe; + EXPECT_TRUE(deserialize_thrift_msg(valid_container.data(), &valid_container_length, compact, + &valid_container_probe) + .ok()) + << "compact=" << compact; + EXPECT_EQ(valid_container_probe.elements.size(), 1); + } +} + TEST_F(ParquetThriftReaderTest, normal) { auto local_fs = io::global_local_filesystem(); io::FileReaderSPtr reader; diff --git a/be/test/format_v2/parquet/parquet_statistics_test.cpp b/be/test/format_v2/parquet/parquet_statistics_test.cpp index e95d00930782f5..3c3e2c97b3c00a 100644 --- a/be/test/format_v2/parquet/parquet_statistics_test.cpp +++ b/be/test/format_v2/parquet/parquet_statistics_test.cpp @@ -941,6 +941,23 @@ TEST(ParquetBloomFilterPruningTest, NativeUint32BloomUsesPhysicalInt32Hash) { bloom_filter)); } +TEST(ParquetBloomFilterPruningTest, NativeBloomRequiresExactDeclaredLength) { + constexpr int64_t offset = 10; + constexpr uint32_t header_size = 12; + constexpr int64_t payload_size = 32; + constexpr size_t file_size = 128; + + EXPECT_TRUE(format::parquet::detail::validate_native_bloom_filter_layout( + offset, header_size, payload_size, header_size + payload_size, file_size) + .ok()); + // A larger declared range can belong to a differently sized filter and must not be decoded + // using the smaller header payload size. + EXPECT_FALSE( + format::parquet::detail::validate_native_bloom_filter_layout( + offset, header_size, payload_size, header_size + payload_size + 32, file_size) + .ok()); +} + TEST(ParquetBloomFilterPruningTest, NativeFloatingBloomPreservesDorisEquality) { const auto check_type = []( tparquet::Type::type physical_type, @@ -1276,7 +1293,12 @@ TEST(ParquetBloomFilterPruningTest, NativeBloomReportsConservativeReadOutcomes) auto truncated = make_valid_bloom(); truncated.resize(truncated.size() - 16); run_case(std::move(truncated), true, false, 1); // truncated payload - run_case(make_valid_bloom(), true, true, 0); // I/O failure + auto contradictory = make_valid_bloom(); + contradictory.resize(contradictory.size() + segment_v2::BloomFilter::MINIMUM_BYTES); + // The declared range contains two blocks while the header declares one. Falling back keeps the + // row group; decoding the first block could falsely prune values mapped to the second block. + run_case(std::move(contradictory), true, false, 1); + run_case(make_valid_bloom(), true, true, 0); // I/O failure } TEST(ParquetBloomFilterPruningTest, NativeBloomPreservesFirstLogicalProbeOrder) { From d31359d639161d39dd7c2cb402e6cd72c0d303e3 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 17 Aug 2026 15:30:44 +0800 Subject: [PATCH 3/4] [fix](parquet) Complete metadata review follow-ups ### What problem does this PR solve? Issue Number: None Problem Summary: Thrift container admission used wire types instead of generated target allocations, contextless service workers lacked a valid limiter, readable windows were over-reserved, and memory-limit failures were retried as parse errors. This change instruments generated container resize sites, preserves task trackers and memory statuses, avoids retrying memory failures, and adds compatibility coverage for Bloom filters without declared lengths. ### Release note Make Parquet metadata deserialization memory admission accurate and preserve actionable memory-limit failures. ### Check List (For Author) - Test: Unit Test - Behavior changed: Yes, decoded Thrift containers are admitted at their actual allocation sites and memory failures are returned without parse retries. - Does this need documentation: No --- .../format/parquet/vparquet_page_reader.cpp | 25 ++++ be/src/util/thrift_container_size.h | 43 ++++++ be/src/util/thrift_util.cpp | 127 +++++----------- be/src/util/thrift_util.h | 8 + .../parquet_column_chunk_reader_test.cpp | 27 ++++ .../parquet/parquet_statistics_test.cpp | 41 +++++ .../format/parquet/parquet_thrift_test.cpp | 140 ++++++++++++++++++ .../parquet/parquet_statistics_test.cpp | 53 +++++++ gensrc/thrift/Makefile | 4 + gensrc/thrift/add_container_memory_check.py | 56 +++++++ 10 files changed, 435 insertions(+), 89 deletions(-) create mode 100644 be/src/util/thrift_container_size.h create mode 100644 gensrc/thrift/add_container_memory_check.py diff --git a/be/src/format/parquet/vparquet_page_reader.cpp b/be/src/format/parquet/vparquet_page_reader.cpp index 611451e0cb4e6f..ed68b10ceccf46 100644 --- a/be/src/format/parquet/vparquet_page_reader.cpp +++ b/be/src/format/parquet/vparquet_page_reader.cpp @@ -30,6 +30,7 @@ #include "io/fs/buffered_reader.h" #include "runtime/runtime_profile.h" #include "storage/cache/page_cache.h" +#include "util/debug_points.h" #include "util/slice.h" #include "util/thrift_util.h" @@ -42,6 +43,24 @@ struct IOContext; namespace doris { static constexpr size_t INIT_PAGE_HEADER_SIZE = 128; +namespace { + +void inject_page_header_memory_failure(Status* status) { + DBUG_EXECUTE_IF("ParquetPageReader.parse_page_header.memory_failure", { + *status = Status::Error( + "Injected page header memory admission failure"); + }); +} + +bool is_memory_failure(const Status& status) { + return status.is() || + status.is() || + status.is() || + status.is() || status.is(); +} + +} // namespace + void ParquetPageCacheKeyBuilder::init(const std::string& path, int64_t mtime) { _file_key_prefix = fmt::format("{}::{}", path, mtime); } @@ -167,9 +186,15 @@ Status PageReader::parse_page_header() { SCOPED_RAW_TIMER(&_page_statistics.decode_header_time); auto st = deserialize_thrift_msg(page_header_buf, &real_header_size, true, &_cur_page_header); + inject_page_header_memory_failure(&st); if (st.ok()) { break; } + if (is_memory_failure(st)) { + // A larger input window cannot recover a memory admission failure and would only turn + // the actionable status into a generic page-header I/O error after repeated reads. + return st; + } if (_offset + header_size >= _end_offset || real_header_size > MAX_PAGE_HEADER_SIZE) { return Status::IOError( "Failed to deserialize parquet page header. offset: {}, " diff --git a/be/src/util/thrift_container_size.h b/be/src/util/thrift_container_size.h new file mode 100644 index 00000000000000..6e170aad913280 --- /dev/null +++ b/be/src/util/thrift_container_size.h @@ -0,0 +1,43 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#pragma once + +#include + +#include +#include + +namespace doris { + +class ThriftContainerMemoryChecker { +public: + virtual ~ThriftContainerMemoryChecker() = default; + virtual void reserve_container_memory(uint32_t count, size_t element_size) = 0; +}; + +template +void reserve_thrift_container_memory(apache::thrift::protocol::TProtocol* protocol, + const Container*, uint32_t count) { + // The generated target type, rather than the untrusted wire tag, defines the allocation made + // by vector::resize. Unknown fields never reach this generated allocation hook. + if (auto* checker = dynamic_cast(protocol); checker != nullptr) { + checker->reserve_container_memory(count, sizeof(typename Container::value_type)); + } +} + +} // namespace doris diff --git a/be/src/util/thrift_util.cpp b/be/src/util/thrift_util.cpp index afb84ec190f3b7..ffcf7465ddf230 100644 --- a/be/src/util/thrift_util.cpp +++ b/be/src/util/thrift_util.cpp @@ -34,6 +34,7 @@ #include "runtime/thread_context.h" #include "storage/tablet/tablet_schema.h" #include "storage/tablet_info.h" +#include "util/thrift_container_size.h" #include "util/thrift_server.h" namespace apache::thrift::protocol { @@ -62,113 +63,61 @@ class TProtocol; namespace doris { namespace { -constexpr size_t DECODED_THRIFT_STRUCT_RESERVATION_BYTES = 1024; - -size_t decoded_thrift_value_reservation(apache::thrift::protocol::TType type) { - using apache::thrift::protocol::T_BOOL; - using apache::thrift::protocol::T_BYTE; - using apache::thrift::protocol::T_DOUBLE; - using apache::thrift::protocol::T_I16; - using apache::thrift::protocol::T_I32; - using apache::thrift::protocol::T_I64; - using apache::thrift::protocol::T_LIST; - using apache::thrift::protocol::T_MAP; - using apache::thrift::protocol::T_SET; - using apache::thrift::protocol::T_STRING; - using apache::thrift::protocol::T_STRUCT; - switch (type) { - case T_BOOL: - case T_BYTE: - return 1; - case T_I16: - return sizeof(int16_t); - case T_I32: - return sizeof(int32_t); - case T_I64: - case T_DOUBLE: - return sizeof(int64_t); - case T_STRING: - return sizeof(std::string); - case T_LIST: - case T_SET: - case T_MAP: - return sizeof(std::vector); - case T_STRUCT: - // Generated structs vary in size. Reserving a conservative inline object budget keeps - // their eager vector resize inside task admission; the serialized-size reservation covers - // their dynamic field payloads. - return DECODED_THRIFT_STRUCT_RESERVATION_BYTES; - default: - return 1; - } -} +class ScopedThreadContextHandle { +public: + ScopedThreadContextHandle() { ThreadLocalHandle::create_thread_local_if_not_exits(); } + ~ScopedThreadContextHandle() { ThreadLocalHandle::del_thread_local_if_count_is_zero(); } +}; -class MemoryBudgetProtocol final : public apache::thrift::protocol::TProtocolDecorator { +class MemoryBudgetProtocol final : public apache::thrift::protocol::TProtocolDecorator, + public ThriftContainerMemoryChecker { public: - MemoryBudgetProtocol(std::shared_ptr protocol, - int32_t serialized_size) - : TProtocolDecorator(std::move(protocol)), - _memory_manager(thread_context()->thread_mem_tracker_mgr.get()), - _prior_reservation(_memory_manager->take_reserved_memory()) { - reserve_or_throw(static_cast(serialized_size), /*restore_prior_on_failure=*/true); + explicit MemoryBudgetProtocol(std::shared_ptr protocol) + : TProtocolDecorator(std::move(protocol)) { + _memory_manager = thread_context()->thread_mem_tracker_mgr.get(); + if (_memory_manager->limiter_mem_tracker()->label() == "Orphan") { + // Apache Thrift worker threads have no Doris task context. Attach a process-accounted + // limiter so reservation checks never run against the forbidden orphan tracker. + _fallback_tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::OTHER, + "ThriftDeserialize"); + _memory_manager->attach_limiter_tracker(_fallback_tracker); + _switched_tracker = true; + } + _prior_reservation = _memory_manager->take_reserved_memory(); } ~MemoryBudgetProtocol() override { _memory_manager->shrink_reserved(); _memory_manager->adopt_reserved_memory(std::move(_prior_reservation)); + if (_switched_tracker) { + _memory_manager->detach_limiter_tracker(); + } } - uint32_t readMapBegin_virt(apache::thrift::protocol::TType& key_type, - apache::thrift::protocol::TType& value_type, - uint32_t& size) override { - const uint32_t consumed = TProtocolDecorator::readMapBegin_virt(key_type, value_type, size); - const uint32_t count = size; - const size_t element_size = decoded_thrift_value_reservation(key_type) + - decoded_thrift_value_reservation(value_type) + - 4 * sizeof(void*); - reserve_container(count, element_size); - return consumed; - } - - uint32_t readListBegin_virt(apache::thrift::protocol::TType& element_type, - uint32_t& size) override { - const uint32_t consumed = TProtocolDecorator::readListBegin_virt(element_type, size); - reserve_container(size, decoded_thrift_value_reservation(element_type)); - return consumed; - } - - uint32_t readSetBegin_virt(apache::thrift::protocol::TType& element_type, - uint32_t& size) override { - const uint32_t consumed = TProtocolDecorator::readSetBegin_virt(element_type, size); - reserve_container(size, decoded_thrift_value_reservation(element_type) + 4 * sizeof(void*)); - return consumed; - } - -private: - void reserve_container(uint32_t count, size_t element_size) { + void reserve_container_memory(uint32_t count, size_t element_size) override { if (count > std::numeric_limits::max() / element_size) { throw apache::thrift::protocol::TProtocolException( apache::thrift::protocol::TProtocolException::SIZE_LIMIT, "Decoded Thrift container size overflows"); } - reserve_or_throw(static_cast(count) * element_size, - /*restore_prior_on_failure=*/false); - } - - void reserve_or_throw(size_t bytes, bool restore_prior_on_failure) { - const Status status = _memory_manager->try_reserve(static_cast(bytes)); - if (status.ok()) { - return; + const size_t bytes = static_cast(count) * element_size; + if (bytes > static_cast(std::numeric_limits::max())) { + throw apache::thrift::protocol::TProtocolException( + apache::thrift::protocol::TProtocolException::SIZE_LIMIT, + "Decoded Thrift container size exceeds reservation range"); } - if (restore_prior_on_failure) { - _memory_manager->adopt_reserved_memory(std::move(_prior_reservation)); + const Status status = _memory_manager->try_reserve(static_cast(bytes)); + if (!status.ok()) { + throw Exception(status); } - throw apache::thrift::protocol::TProtocolException( - apache::thrift::protocol::TProtocolException::SIZE_LIMIT, status.to_string()); } - ThreadMemTrackerMgr* _memory_manager; +private: + ScopedThreadContextHandle _thread_context_handle; + ThreadMemTrackerMgr* _memory_manager = nullptr; + std::shared_ptr _fallback_tracker; ReservedMemoryToken _prior_reservation; + bool _switched_tracker = false; }; } // namespace @@ -199,7 +148,7 @@ std::shared_ptr create_deserialize_protocol /*strict_read=*/false, /*strict_write=*/true); } - return std::make_shared(std::move(protocol), size_limit); + return std::make_shared(std::move(protocol)); } // Comparator for THostPorts. Thrift declares this (in gen-cpp/Types_types.h) but diff --git a/be/src/util/thrift_util.h b/be/src/util/thrift_util.h index 208d9f317e7e4b..95f89cc492b19f 100644 --- a/be/src/util/thrift_util.h +++ b/be/src/util/thrift_util.h @@ -31,7 +31,9 @@ #include #include +#include "common/exception.h" #include "common/status.h" +#include "util/defer_op.h" namespace apache::thrift::protocol { class TProtocol; @@ -149,9 +151,15 @@ Status deserialize_thrift_msg(const uint8_t* buf, uint32_t* len, bool compact, const_cast(buf), *len, apache::thrift::transport::TMemoryBuffer::OBSERVE, conf)); try { + // Thrift-generated standard containers can throw Doris memory exceptions through the + // allocation hooks; preserve their original status for callers that must not retry OOM. + enable_thread_catch_bad_alloc++; + Defer defer_catch_bad_alloc {[&]() { enable_thread_catch_bad_alloc--; }}; std::shared_ptr tproto = create_deserialize_protocol(tmem_transport, compact, size_limit); deserialized_msg->read(tproto.get()); + } catch (const doris::Exception& e) { + return e.to_status(); } catch (std::exception& e) { return Status::InternalError("Couldn't deserialize thrift msg:\n{}", e.what()); } catch (...) { diff --git a/be/test/format/parquet/parquet_column_chunk_reader_test.cpp b/be/test/format/parquet/parquet_column_chunk_reader_test.cpp index be9616c523f638..c0bd548e91fa13 100644 --- a/be/test/format/parquet/parquet_column_chunk_reader_test.cpp +++ b/be/test/format/parquet/parquet_column_chunk_reader_test.cpp @@ -34,6 +34,7 @@ #include "io/fs/file_reader.h" #include "runtime/runtime_state.h" #include "util/coding.h" +#include "util/debug_points.h" #include "util/thrift_util.h" namespace doris { @@ -294,6 +295,32 @@ TEST(ParquetColumnChunkReaderTest, DictionaryProbeDoesNotParseDataPageHeader) { EXPECT_EQ(buffered_reader.read_count(), data_header_read_count); } +TEST(ParquetColumnChunkReaderTest, PageHeaderMemoryFailureIsNotRetried) { + tparquet::PageHeader header = make_data_page_header(tparquet::Encoding::PLAIN); + + std::vector data; + ThriftSerializer serializer(/*compact=*/true, /*initial_buffer_size=*/256); + ASSERT_TRUE(serializer.serialize(&header, &data).ok()); + data.push_back(0); + const size_t data_size = data.size(); + CountingBufferedReader buffered_reader(std::move(data)); + tparquet::ColumnMetaData metadata; + metadata.__set_codec(tparquet::CompressionCodec::UNCOMPRESSED); + ParquetPageReadContext page_read_ctx(false); + PageReader reader(&buffered_reader, nullptr, 0, data_size, 1, metadata, + page_read_ctx); + const bool debug_points_were_enabled = config::enable_debug_points; + config::enable_debug_points = true; + DebugPoints::instance()->add("ParquetPageReader.parse_page_header.memory_failure"); + + Status status = reader.parse_page_header(); + DebugPoints::instance()->remove("ParquetPageReader.parse_page_header.memory_failure"); + config::enable_debug_points = debug_points_were_enabled; + + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(buffered_reader.read_count(), 1); +} + TEST(ParquetColumnChunkReaderTest, ParsePageHeaderLoadsDictionaryOnFirstUse) { ColumnChunkFixture fixture; ASSERT_TRUE(make_dictionary_fixture(&fixture).ok()); diff --git a/be/test/format/parquet/parquet_statistics_test.cpp b/be/test/format/parquet/parquet_statistics_test.cpp index e542b029c2ec95..eeae08326cf0f1 100644 --- a/be/test/format/parquet/parquet_statistics_test.cpp +++ b/be/test/format/parquet/parquet_statistics_test.cpp @@ -162,6 +162,47 @@ TEST_F(ParquetStatisticsTest, accept_valid_bloom_filter_layout) { read_test_bloom_filter(/*header_payload_size=*/32, /*actual_payload_size=*/32).ok()); } +TEST_F(ParquetStatisticsTest, accept_bloom_filter_without_declared_length_before_trailing_bytes) { + constexpr int32_t present_value = 1; + ParquetBlockSplitBloomFilter source; + ASSERT_TRUE(source.init(segment_v2::BloomFilter::MINIMUM_BYTES, + segment_v2::HashStrategyPB::XX_HASH_64) + .ok()); + source.add_bytes(reinterpret_cast(&present_value), sizeof(present_value)); + int32_t absent_value = 2; + while (source.test_bytes(reinterpret_cast(&absent_value), sizeof(absent_value))) { + ++absent_value; + } + + tparquet::BloomFilterAlgorithm algorithm; + algorithm.__set_BLOCK(tparquet::SplitBlockAlgorithm()); + tparquet::BloomFilterHash hash; + hash.__set_XXHASH(tparquet::XxHash()); + tparquet::BloomFilterCompression compression; + compression.__set_UNCOMPRESSED(tparquet::Uncompressed()); + tparquet::BloomFilterHeader header; + header.__set_numBytes(static_cast(source.size())); + header.__set_algorithm(algorithm); + header.__set_hash(hash); + header.__set_compression(compression); + std::vector file_bytes; + ThriftSerializer serializer(/*compact=*/true, /*initial_buffer_size=*/64); + ASSERT_TRUE(serializer.serialize(&header, &file_bytes).ok()); + file_bytes.insert(file_bytes.end(), source.data(), source.data() + source.size()); + file_bytes.resize(file_bytes.size() + 64); + + tparquet::ColumnMetaData metadata; + metadata.__set_bloom_filter_offset(0); + auto reader = std::make_shared(std::move(file_bytes)); + ParquetPredicate::ColumnStat stat; + ASSERT_TRUE(ParquetPredicate::read_bloom_filter(metadata, reader, nullptr, &stat).ok()); + ASSERT_NE(stat.bloom_filter, nullptr); + EXPECT_TRUE(stat.bloom_filter->test_bytes(reinterpret_cast(&present_value), + sizeof(present_value))); + EXPECT_FALSE(stat.bloom_filter->test_bytes(reinterpret_cast(&absent_value), + sizeof(absent_value))); +} + TEST_F(ParquetStatisticsTest, test_try_read_old_utf8_stats) { // [, bcé]: min is empty, max starts with ASCII { diff --git a/be/test/format/parquet/parquet_thrift_test.cpp b/be/test/format/parquet/parquet_thrift_test.cpp index 8d1fcb40085da5..5580cd852fb3c4 100644 --- a/be/test/format/parquet/parquet_thrift_test.cpp +++ b/be/test/format/parquet/parquet_thrift_test.cpp @@ -15,17 +15,27 @@ // specific language governing permissions and limitations // under the License. +#include #include #include +#include #include #include #include #include #include +#include +#include #include +#include +#include +#include +#include +#include #include #include +#include #include #include #include @@ -63,6 +73,8 @@ #include "runtime/memory/mem_tracker_limiter.h" #include "runtime/thread_context.h" #include "util/slice.h" +#include "util/thrift_container_size.h" +#include "util/thrift_server.h" #include "util/thrift_util.h" #include "util/timezone_utils.h" @@ -78,6 +90,7 @@ struct ThriftContainerProbe { apache::thrift::protocol::TType element_type; uint32_t size = 0; uint32_t consumed = protocol->readListBegin(element_type, size); + reserve_thrift_container_memory(protocol, &elements, size); elements.resize(size); return consumed + protocol->readListEnd(); } @@ -93,6 +106,56 @@ struct ThriftStringProbe { std::string value; }; +struct ThriftSkipProbe { + uint32_t read(apache::thrift::protocol::TProtocol* protocol) { + return protocol->skip(apache::thrift::protocol::T_LIST); + } +}; + +std::vector thrift_string_bytes(bool compact, std::string_view value); + +class ContextlessDeserializeService final : public doristest::NetworkTestServiceIf { +public: + void Send(doristest::ThriftDataResult& result, + const doristest::ThriftDataParams& params) override { + contextless_before_deserialize = !pthread_context_ptr_init; + auto bytes = thrift_string_bytes(/*compact=*/true, params.data); + uint32_t length = bytes.size(); + ThriftStringProbe probe; + const Status status = deserialize_thrift_msg(bytes.data(), &length, true, &probe); + context_cleaned_after_deserialize = !pthread_context_ptr_init; + if (!status.ok()) { + throw apache::thrift::TException(status.to_string()); + } + result.__set_bytes_received(static_cast(probe.value.size())); + } + + std::atomic_bool contextless_before_deserialize = false; + std::atomic_bool context_cleaned_after_deserialize = false; +}; + +int find_available_port() { + const int socket_fd = socket(AF_INET, SOCK_STREAM, 0); + if (socket_fd < 0) { + return -1; + } + sockaddr_in address {}; + address.sin_family = AF_INET; + address.sin_addr.s_addr = htonl(INADDR_LOOPBACK); + address.sin_port = 0; + if (bind(socket_fd, reinterpret_cast(&address), sizeof(address)) != 0) { + close(socket_fd); + return -1; + } + socklen_t address_size = sizeof(address); + if (getsockname(socket_fd, reinterpret_cast(&address), &address_size) != 0) { + close(socket_fd); + return -1; + } + close(socket_fd); + return ntohs(address.sin_port); +} + std::vector thrift_list_bytes(bool compact, uint32_t count, size_t total_size) { std::vector bytes; if (compact) { @@ -159,6 +222,83 @@ TEST_F(ParquetThriftReaderTest, reject_decoded_container_that_exceeds_task_budge } } +TEST_F(ParquetThriftReaderTest, use_generated_target_type_for_container_budget) { + constexpr int64_t memory_limit = 4 * 1024; + // FileMetaData.schema is vector, but the forged compact wire tag advertises + // bool elements. Generated readers must budget the actual target vector before resizing it. + std::vector compact_metadata {0x29, 0xf1, 0x40, 0x00}; + compact_metadata.resize(72); + uint32_t length = compact_metadata.size(); + tparquet::FileMetaData metadata; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "GeneratedThriftTargetType", memory_limit); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + + Status status = deserialize_thrift_msg(compact_metadata.data(), &length, true, &metadata); + + EXPECT_TRUE(status.is()) << status; + EXPECT_TRUE(metadata.schema.empty()); +} + +TEST_F(ParquetThriftReaderTest, skip_unknown_container_without_phantom_reservation) { + constexpr uint32_t element_count = 64; + constexpr int64_t memory_limit = 4 * 1024; + for (const bool compact : {true, false}) { + auto bytes = thrift_list_bytes(compact, element_count, element_count + 8); + uint32_t length = bytes.size(); + ThriftSkipProbe probe; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "SkippedThriftContainer", memory_limit); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + + EXPECT_TRUE(deserialize_thrift_msg(bytes.data(), &length, compact, &probe).ok()) + << "compact=" << compact; + } +} + +TEST_F(ParquetThriftReaderTest, accept_small_message_from_large_readable_window) { + constexpr int64_t memory_limit = 16 * 1024; + for (const bool compact : {true, false}) { + auto bytes = thrift_string_bytes(compact, "valid"); + bytes.resize(64 * 1024); + uint32_t length = bytes.size(); + ThriftStringProbe probe; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "ThriftReadableWindow", memory_limit); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + + EXPECT_TRUE(deserialize_thrift_msg(bytes.data(), &length, compact, &probe).ok()) + << "compact=" << compact; + EXPECT_EQ(probe.value, "valid"); + EXPECT_LT(length, bytes.size()); + } +} + +TEST_F(ParquetThriftReaderTest, deserialize_on_contextless_threaded_service_worker) { + const int port = find_available_port(); + ASSERT_GT(port, 0); + auto handler = std::make_shared(); + auto processor = std::make_shared(handler); + ThriftServer server("contextless-deserialize-test", processor, port, + ThriftServer::DEFAULT_WORKER_THREADS, ThriftServer::THREADED); + ASSERT_TRUE(server.start().ok()); + + auto socket = std::make_shared("127.0.0.1", port); + auto transport = std::make_shared(socket); + auto protocol = std::make_shared(transport); + doristest::NetworkTestServiceClient client(protocol); + transport->open(); + doristest::ThriftDataParams params; + params.__set_data("valid"); + doristest::ThriftDataResult result; + client.Send(result, params); + transport->close(); + + EXPECT_EQ(result.bytes_received, params.data.size()); + EXPECT_TRUE(handler->contextless_before_deserialize.load()); + EXPECT_TRUE(handler->context_cleaned_after_deserialize.load()); +} + TEST_F(ParquetThriftReaderTest, bound_strings_and_accept_valid_controls_for_both_protocols) { for (const bool compact : {true, false}) { auto malformed_container = thrift_list_bytes(compact, /*count=*/64, /*total_size=*/8); diff --git a/be/test/format_v2/parquet/parquet_statistics_test.cpp b/be/test/format_v2/parquet/parquet_statistics_test.cpp index 3c3e2c97b3c00a..1c3f7418e404d7 100644 --- a/be/test/format_v2/parquet/parquet_statistics_test.cpp +++ b/be/test/format_v2/parquet/parquet_statistics_test.cpp @@ -1299,6 +1299,59 @@ TEST(ParquetBloomFilterPruningTest, NativeBloomReportsConservativeReadOutcomes) // row group; decoding the first block could falsely prune values mapped to the second block. run_case(std::move(contradictory), true, false, 1); run_case(make_valid_bloom(), true, true, 0); // I/O failure + + const auto run_without_declared_length = [&](int32_t predicate_value, bool expected_pruned) { + auto bytes = make_valid_bloom(); + bytes.resize(bytes.size() + segment_v2::BloomFilter::MINIMUM_BYTES); + auto type = std::make_shared(); + auto column_schema = std::make_unique(); + column_schema->kind = format::parquet::ParquetColumnSchemaKind::PRIMITIVE; + column_schema->local_id = 0; + column_schema->leaf_column_id = 0; + column_schema->type = type; + column_schema->type_descriptor.doris_type = type; + column_schema->type_descriptor.physical_type = tparquet::Type::INT32; + + tparquet::ColumnMetaData column_metadata; + column_metadata.__set_type(tparquet::Type::INT32); + column_metadata.__set_codec(tparquet::CompressionCodec::UNCOMPRESSED); + column_metadata.__set_num_values(1); + column_metadata.__set_total_compressed_size(0); + column_metadata.__set_data_page_offset(0); + column_metadata.__set_bloom_filter_offset(0); + tparquet::ColumnChunk chunk; + chunk.__set_meta_data(column_metadata); + tparquet::RowGroup row_group; + row_group.__set_columns({chunk}); + row_group.__set_total_byte_size(0); + row_group.__set_num_rows(1); + tparquet::FileMetaData metadata; + metadata.__set_version(1); + metadata.__set_num_rows(1); + metadata.__set_row_groups({row_group}); + + auto request = + request_with_bloom_conjunct(type, {Field::create_field(predicate_value)}); + std::vector> schema; + schema.push_back(std::move(column_schema)); + format::parquet::ParquetFileContext file_context; + file_context.native_file = std::make_shared(std::move(bytes)); + std::vector selected_row_groups; + format::parquet::ParquetPruningStats pruning_stats; + ASSERT_TRUE(format::parquet::select_row_groups_by_metadata( + metadata, schema, request, nullptr, &selected_row_groups, true, + &pruning_stats, nullptr, nullptr, &file_context) + .ok()); + EXPECT_EQ(selected_row_groups.empty(), expected_pruned); + EXPECT_EQ(pruning_stats.filtered_row_groups_by_bloom_filter, expected_pruned ? 1 : 0); + EXPECT_EQ(pruning_stats.bloom_filter_probe_successes, 1); + EXPECT_EQ(pruning_stats.bloom_filter_corrupt_rejections, 0); + }; + + // The header's numBytes bounds the filter when the optional metadata length is absent; bytes + // belonging to later file structures must not change present or absent probe outcomes. + run_without_declared_length(/*predicate_value=*/1, /*expected_pruned=*/false); + run_without_declared_length(/*predicate_value=*/2, /*expected_pruned=*/true); } TEST(ParquetBloomFilterPruningTest, NativeBloomPreservesFirstLogicalProbeOrder) { diff --git a/gensrc/thrift/Makefile b/gensrc/thrift/Makefile index f6a196390eecbc..950bacca7907ae 100644 --- a/gensrc/thrift/Makefile +++ b/gensrc/thrift/Makefile @@ -30,9 +30,13 @@ all: ${GEN_OBJECTS} ${OBJECTS} .PHONY: all THRIFT_CPP_ARGS = -I ${CURDIR} -I ${BUILD_DIR}/thrift/ --gen cpp:moveable_types,no_skeleton -out ${BUILD_DIR}/gen_cpp --allow-64bit-consts -strict +CONTAINER_MEMORY_CHECK = ${CURDIR}/add_container_memory_check.py ${BUILD_DIR}/gen_cpp: mkdir -p $@ # handwrite thrift +${BUILD_DIR}/gen_cpp/parquet_types.cpp: ${CONTAINER_MEMORY_CHECK} + ${BUILD_DIR}/gen_cpp/%_types.cpp: ${CURDIR}/%.thrift | ${BUILD_DIR}/gen_cpp ${THRIFT} ${THRIFT_CPP_ARGS} $< + python3 ${CONTAINER_MEMORY_CHECK} $@ diff --git a/gensrc/thrift/add_container_memory_check.py b/gensrc/thrift/add_container_memory_check.py new file mode 100644 index 00000000000000..efece114a44381 --- /dev/null +++ b/gensrc/thrift/add_container_memory_check.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pathlib +import re +import sys + + +INCLUDE = '#include "util/thrift_container_size.h"' +RESIZE = re.compile(r"^(?P\s*)(?P.+)\.resize\((?P_size\d*)\);$") + + +def main() -> int: + path = pathlib.Path(sys.argv[1]) + lines = path.read_text().splitlines() + output = [] + replacements = 0 + for line in lines: + match = RESIZE.match(line) + if match is not None: + check = ( + f"{match['indent']}::doris::reserve_thrift_container_memory(" + f"iprot, &{match['container']}, {match['size']});" + ) + if not output or output[-1] != check: + output.append(check) + replacements += 1 + output.append(line) + + if replacements == 0: + return 0 + if INCLUDE not in output: + own_header = next(index for index, line in enumerate(output) if line.startswith('#include "')) + output.insert(own_header + 1, INCLUDE) + path.write_text("\n".join(output) + "\n") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 4824fe849dbcf57c9494cf3d57a93d7b8e155d36 Mon Sep 17 00:00:00 2001 From: Gabriel Date: Mon, 17 Aug 2026 17:10:14 +0800 Subject: [PATCH 4/4] [fix](parquet) Retain decoded page index memory charges --- be/src/util/thrift_container_size.h | 86 ++++++++++++++++++- be/src/util/thrift_util.cpp | 18 +++- .../format/parquet/parquet_thrift_test.cpp | 40 +++++++++ gensrc/thrift/Makefile | 4 +- gensrc/thrift/add_container_memory_check.py | 60 +++++++++++-- .../thrift/test_add_container_memory_check.py | 75 ++++++++++++++++ 6 files changed, 268 insertions(+), 15 deletions(-) create mode 100644 gensrc/thrift/test_add_container_memory_check.py diff --git a/be/src/util/thrift_container_size.h b/be/src/util/thrift_container_size.h index 6e170aad913280..0cf813e65e7376 100644 --- a/be/src/util/thrift_container_size.h +++ b/be/src/util/thrift_container_size.h @@ -17,27 +17,107 @@ #pragma once +#include #include +#include #include #include +#include +#include +#include +#include namespace doris { +class ThriftContainerMemoryCharge { +public: + virtual ~ThriftContainerMemoryCharge() = default; +}; + class ThriftContainerMemoryChecker { public: virtual ~ThriftContainerMemoryChecker() = default; - virtual void reserve_container_memory(uint32_t count, size_t element_size) = 0; + virtual std::shared_ptr reserve_container_memory( + uint32_t count, size_t element_size) = 0; + virtual void retain_temporary_container_charge( + std::shared_ptr charge) = 0; +}; + +template +class ThriftMemoryTrackedVector : public std::vector { + using Base = std::vector; + +public: + using Base::Base; + using Base::operator=; + + ThriftMemoryTrackedVector() = default; + ThriftMemoryTrackedVector(const ThriftMemoryTrackedVector&) = default; + ThriftMemoryTrackedVector(ThriftMemoryTrackedVector&&) noexcept = default; + ThriftMemoryTrackedVector& operator=(const ThriftMemoryTrackedVector&) = default; + ThriftMemoryTrackedVector& operator=(ThriftMemoryTrackedVector&&) noexcept = default; + + ThriftMemoryTrackedVector& operator=(const Base& other) { + Base::operator=(other); + _memory_charge.reset(); + return *this; + } + + ThriftMemoryTrackedVector& operator=(Base&& other) noexcept { + Base::operator=(std::move(other)); + _memory_charge.reset(); + return *this; + } + + void set_thrift_memory_charge(std::shared_ptr charge) { + _memory_charge = std::move(charge); + } + + void swap(ThriftMemoryTrackedVector& other) noexcept { + Base::swap(other); + _memory_charge.swap(other._memory_charge); + } + + friend bool operator==(const ThriftMemoryTrackedVector& lhs, + const ThriftMemoryTrackedVector& rhs) { + return static_cast(lhs) == static_cast(rhs); + } + + friend bool operator<(const ThriftMemoryTrackedVector& lhs, + const ThriftMemoryTrackedVector& rhs) { + return static_cast(lhs) < static_cast(rhs); + } + + friend std::ostream& operator<<(std::ostream& out, const ThriftMemoryTrackedVector& values) { + return out << apache::thrift::to_string(static_cast(values)); + } + +private: + std::shared_ptr _memory_charge; }; template void reserve_thrift_container_memory(apache::thrift::protocol::TProtocol* protocol, - const Container*, uint32_t count) { + Container* container, uint32_t count) { // The generated target type, rather than the untrusted wire tag, defines the allocation made // by vector::resize. Unknown fields never reach this generated allocation hook. if (auto* checker = dynamic_cast(protocol); checker != nullptr) { - checker->reserve_container_memory(count, sizeof(typename Container::value_type)); + const size_t elements = std::max(count, container->capacity()); + auto charge = + checker->reserve_container_memory(elements, sizeof(typename Container::value_type)); + if constexpr (requires { container->set_thrift_memory_charge(charge); }) { + // Generated fields retain the admission charge until the decoded allocation dies. + container->set_thrift_memory_charge(std::move(charge)); + } else { + checker->retain_temporary_container_charge(std::move(charge)); + } } } +template +void swap(ThriftMemoryTrackedVector& lhs, ThriftMemoryTrackedVector& rhs) noexcept { + lhs.swap(rhs); +} + } // namespace doris diff --git a/be/src/util/thrift_util.cpp b/be/src/util/thrift_util.cpp index ffcf7465ddf230..fdd5ef13ad078d 100644 --- a/be/src/util/thrift_util.cpp +++ b/be/src/util/thrift_util.cpp @@ -94,7 +94,8 @@ class MemoryBudgetProtocol final : public apache::thrift::protocol::TProtocolDec } } - void reserve_container_memory(uint32_t count, size_t element_size) override { + std::shared_ptr reserve_container_memory( + uint32_t count, size_t element_size) override { if (count > std::numeric_limits::max() / element_size) { throw apache::thrift::protocol::TProtocolException( apache::thrift::protocol::TProtocolException::SIZE_LIMIT, @@ -110,13 +111,28 @@ class MemoryBudgetProtocol final : public apache::thrift::protocol::TProtocolDec if (!status.ok()) { throw Exception(status); } + return std::make_shared(_memory_manager->take_reserved_memory()); + } + + void retain_temporary_container_charge( + std::shared_ptr charge) override { + _temporary_charges.emplace_back(std::move(charge)); } private: + class ReservedMemoryCharge final : public ThriftContainerMemoryCharge { + public: + explicit ReservedMemoryCharge(ReservedMemoryToken token) : _token(std::move(token)) {} + + private: + ReservedMemoryToken _token; + }; + ScopedThreadContextHandle _thread_context_handle; ThreadMemTrackerMgr* _memory_manager = nullptr; std::shared_ptr _fallback_tracker; ReservedMemoryToken _prior_reservation; + std::vector> _temporary_charges; bool _switched_tracker = false; }; diff --git a/be/test/format/parquet/parquet_thrift_test.cpp b/be/test/format/parquet/parquet_thrift_test.cpp index 5580cd852fb3c4..35495702eedb4f 100644 --- a/be/test/format/parquet/parquet_thrift_test.cpp +++ b/be/test/format/parquet/parquet_thrift_test.cpp @@ -222,6 +222,46 @@ TEST_F(ParquetThriftReaderTest, reject_decoded_container_that_exceeds_task_budge } } +TEST_F(ParquetThriftReaderTest, retain_decoded_container_charge_until_output_is_destroyed) { + constexpr size_t element_count = 32; + tparquet::PageLocation page_location; + page_location.__set_offset(0); + page_location.__set_compressed_page_size(1); + page_location.__set_first_row_index(0); + tparquet::OffsetIndex source; + source.__set_page_locations(std::vector(element_count, page_location)); + std::vector bytes; + ThriftSerializer serializer(/*compact=*/true, 1024); + ASSERT_TRUE(serializer.serialize(&source, &bytes).ok()); + + const auto one_container_bytes = + static_cast(element_count * sizeof(tparquet::PageLocation)); + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "RetainedThriftContainers", + one_container_bytes * 3 / 2); + SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(tracker); + + { + tparquet::OffsetIndex first; + uint32_t first_length = bytes.size(); + ASSERT_TRUE(deserialize_thrift_msg(bytes.data(), &first_length, true, &first).ok()); + ASSERT_EQ(first.page_locations.size(), element_count); + + tparquet::OffsetIndex second; + uint32_t second_length = bytes.size(); + const Status second_status = + deserialize_thrift_msg(bytes.data(), &second_length, true, &second); + EXPECT_TRUE(second_status.is()) << second_status; + EXPECT_TRUE(second.page_locations.empty()); + } + + // Destroying the retained output must return its charge so the same task can deserialize again. + tparquet::OffsetIndex after_release; + uint32_t after_release_length = bytes.size(); + EXPECT_TRUE( + deserialize_thrift_msg(bytes.data(), &after_release_length, true, &after_release).ok()); +} + TEST_F(ParquetThriftReaderTest, use_generated_target_type_for_container_budget) { constexpr int64_t memory_limit = 4 * 1024; // FileMetaData.schema is vector, but the forged compact wire tag advertises diff --git a/gensrc/thrift/Makefile b/gensrc/thrift/Makefile index 950bacca7907ae..9f0984abc3f923 100644 --- a/gensrc/thrift/Makefile +++ b/gensrc/thrift/Makefile @@ -34,9 +34,7 @@ CONTAINER_MEMORY_CHECK = ${CURDIR}/add_container_memory_check.py ${BUILD_DIR}/gen_cpp: mkdir -p $@ -# handwrite thrift -${BUILD_DIR}/gen_cpp/parquet_types.cpp: ${CONTAINER_MEMORY_CHECK} -${BUILD_DIR}/gen_cpp/%_types.cpp: ${CURDIR}/%.thrift | ${BUILD_DIR}/gen_cpp +${BUILD_DIR}/gen_cpp/%_types.cpp: ${CURDIR}/%.thrift ${CONTAINER_MEMORY_CHECK} | ${BUILD_DIR}/gen_cpp ${THRIFT} ${THRIFT_CPP_ARGS} $< python3 ${CONTAINER_MEMORY_CHECK} $@ diff --git a/gensrc/thrift/add_container_memory_check.py b/gensrc/thrift/add_container_memory_check.py index efece114a44381..9a863ff484dd28 100644 --- a/gensrc/thrift/add_container_memory_check.py +++ b/gensrc/thrift/add_container_memory_check.py @@ -24,10 +24,32 @@ INCLUDE = '#include "util/thrift_container_size.h"' RESIZE = re.compile(r"^(?P\s*)(?P.+)\.resize\((?P_size\d*)\);$") +VECTOR_FIELD = re.compile( + r"^(?P\s+)std::vector<(?P.+)> (?P[A-Za-z_][A-Za-z0-9_]*);$" +) +LIFETIME_TRACKED_FIELDS = { + # Parquet metadata loaders can retain many page-index objects together, so these reservations + # must outlive the protocol. Other generated containers still need admission before resize but + # are not retained by the metadata-loading path that requires aggregate lifetime accounting. + "page_locations", + "unencoded_byte_array_data_bytes", + "null_pages", + "min_values", + "max_values", + "null_counts", + "repetition_level_histograms", + "definition_level_histograms", +} -def main() -> int: - path = pathlib.Path(sys.argv[1]) +def add_include(lines): + if INCLUDE in lines: + return + first_include = next(index for index, line in enumerate(lines) if line.startswith("#include ")) + lines.insert(first_include + 1, INCLUDE) + + +def instrument_source(path): lines = path.read_text().splitlines() output = [] replacements = 0 @@ -43,12 +65,34 @@ def main() -> int: replacements += 1 output.append(line) - if replacements == 0: - return 0 - if INCLUDE not in output: - own_header = next(index for index, line in enumerate(output) if line.startswith('#include "')) - output.insert(own_header + 1, INCLUDE) - path.write_text("\n".join(output) + "\n") + if replacements != 0: + add_include(output) + path.write_text("\n".join(output) + "\n") + + +def instrument_header(path): + lines = path.read_text().splitlines() + output = [] + replacements = 0 + for line in lines: + match = VECTOR_FIELD.match(line) + if match is not None and match["name"] in LIFETIME_TRACKED_FIELDS: + line = ( + f"{match['indent']}::doris::ThriftMemoryTrackedVector<{match['value']}> " + f"{match['name']};" + ) + replacements += 1 + output.append(line) + + if replacements != 0: + add_include(output) + path.write_text("\n".join(output) + "\n") + + +def main() -> int: + path = pathlib.Path(sys.argv[1]) + instrument_source(path) + instrument_header(path.with_suffix(".h")) return 0 diff --git a/gensrc/thrift/test_add_container_memory_check.py b/gensrc/thrift/test_add_container_memory_check.py new file mode 100644 index 00000000000000..2e6ecf1d4c0199 --- /dev/null +++ b/gensrc/thrift/test_add_container_memory_check.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pathlib +import subprocess +import sys +import tempfile +import unittest + + +SCRIPT_DIR = pathlib.Path(__file__).resolve().parent +CHECKER = SCRIPT_DIR / "add_container_memory_check.py" + + +class AddContainerMemoryCheckTest(unittest.TestCase): + def test_instruments_cpp_and_generated_vector_fields_idempotently(self): + with tempfile.TemporaryDirectory() as temp_dir: + source = pathlib.Path(temp_dir) / "Sample_types.cpp" + header = source.with_suffix(".h") + source.write_text( + '#include "Sample_types.h"\n' + "uint32_t Sample::read(TProtocol* iprot) {\n" + " this->page_locations.resize(_size0);\n" + "}\n" + ) + header.write_text( + '#include "thrift/TBase.h"\n' + "class Sample {\n" + " public:\n" + " std::vector page_locations;\n" + " void __set_page_locations(const std::vector & val);\n" + "};\n" + ) + + subprocess.run([sys.executable, str(CHECKER), str(source)], check=True) + first_source = source.read_text() + first_header = header.read_text() + subprocess.run([sys.executable, str(CHECKER), str(source)], check=True) + + self.assertEqual(source.read_text(), first_source) + self.assertEqual(header.read_text(), first_header) + self.assertIn("reserve_thrift_container_memory", first_source) + self.assertIn("ThriftMemoryTrackedVector page_locations", first_header) + self.assertIn( + "__set_page_locations(const std::vector & val)", first_header + ) + + def test_checker_is_prerequisite_of_every_generated_reader(self): + makefile = (SCRIPT_DIR / "Makefile").read_text() + pattern_rule = ( + "${BUILD_DIR}/gen_cpp/%_types.cpp: ${CURDIR}/%.thrift " + "${CONTAINER_MEMORY_CHECK} | ${BUILD_DIR}/gen_cpp" + ) + self.assertIn(pattern_rule, makefile) + self.assertNotIn("gen_cpp/parquet_types.cpp: ${CONTAINER_MEMORY_CHECK}", makefile) + + +if __name__ == "__main__": + unittest.main()