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 a0f6c75bdefb8b..ee0f3b60d83059 100644 --- a/be/src/format/parquet/parquet_predicate.h +++ b/be/src/format/parquet/parquet_predicate.h @@ -441,45 +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"); + } + const auto expected_declared_length = static_cast(total_size); + if (column_meta_data.__isset.bloom_filter_length && + column_meta_data.bloom_filter_length != expected_declared_length) { + return Status::Corruption("Invalid Parquet bloom filter declared length"); + } - std::vector data_buffer(t_bloom_filter_header.numBytes); + 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( - 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(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"); + } - RETURN_IF_ERROR(ans_stat->bloom_filter->init( - reinterpret_cast(data_buffer.data()), t_bloom_filter_header.numBytes, - 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_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/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_container_size.h b/be/src/util/thrift_container_size.h new file mode 100644 index 00000000000000..0cf813e65e7376 --- /dev/null +++ b/be/src/util/thrift_container_size.h @@ -0,0 +1,123 @@ +// 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 +#include +#include +#include +#include +#include +#include + +namespace doris { + +class ThriftContainerMemoryCharge { +public: + virtual ~ThriftContainerMemoryCharge() = default; +}; + +class ThriftContainerMemoryChecker { +public: + virtual ~ThriftContainerMemoryChecker() = default; + 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, + 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) { + 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 b0d6d3b3093b45..fdd5ef13ad078d 100644 --- a/be/src/util/thrift_util.cpp +++ b/be/src/util/thrift_util.cpp @@ -20,16 +20,21 @@ #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_container_size.h" #include "util/thrift_server.h" namespace apache::thrift::protocol { @@ -56,6 +61,82 @@ class TProtocol; #include namespace doris { +namespace { + +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, + public ThriftContainerMemoryChecker { +public: + 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(); + } + } + + 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, + "Decoded Thrift container size overflows"); + } + 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"); + } + const Status status = _memory_manager->try_reserve(static_cast(bytes)); + 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; +}; + +} // namespace ThriftSerializer::ThriftSerializer(bool compact, int initial_buffer_size) : _mem_buffer(new apache::thrift::transport::TMemoryBuffer(initial_buffer_size)) { @@ -71,16 +152,19 @@ 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) { + std::shared_ptr protocol; if (compact) { - apache::thrift::protocol::TCompactProtocolFactoryT - tproto_factory; - return tproto_factory.getProtocol(mem); + protocol = std::make_shared>(mem, size_limit, size_limit); } else { - apache::thrift::protocol::TBinaryProtocolFactoryT - tproto_factory; - return tproto_factory.getProtocol(mem); + protocol = std::make_shared>(mem, size_limit, size_limit, + /*strict_read=*/false, + /*strict_write=*/true); } + 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 73c753e2271c8f..95f89cc492b19f 100644 --- a/be/src/util/thrift_util.h +++ b/be/src/util/thrift_util.h @@ -22,14 +22,18 @@ #include #include +#include #include #include #include +#include #include #include #include +#include "common/exception.h" #include "common/status.h" +#include "util/defer_op.h" namespace apache::thrift::protocol { class TProtocol; @@ -122,7 +126,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,18 +139,27 @@ 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()))); + // 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); - 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 d52320c51698e4..eeae08326cf0f1 100644 --- a/be/test/format/parquet/parquet_statistics_test.cpp +++ b/be/test/format/parquet/parquet_statistics_test.cpp @@ -17,16 +17,192 @@ #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, size_t logical_size = 0) + : _data(std::move(data)), + _logical_size(logical_size == 0 ? _data.size() : logical_size) {} + + Status close() override { + _closed = true; + return Status::OK(); + } + + const io::Path& path() const override { return _path; } + 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, + 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); + _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, + 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; + 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); + 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; + 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 + 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. + 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=*/32, /*actual_payload_size=*/64).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, 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 0fe101db598138..35495702eedb4f 100644 --- a/be/test/format/parquet/parquet_thrift_test.cpp +++ b/be/test/format/parquet/parquet_thrift_test.cpp @@ -15,16 +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 @@ -59,10 +70,119 @@ #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_container_size.h" +#include "util/thrift_server.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); + reserve_thrift_container_memory(protocol, &elements, 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; +}; + +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) { + 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: @@ -71,6 +191,195 @@ 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, 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, 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 + // 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); + 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..1c3f7418e404d7 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,65 @@ 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 + + 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..9f0984abc3f923 100644 --- a/gensrc/thrift/Makefile +++ b/gensrc/thrift/Makefile @@ -30,9 +30,11 @@ 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/%_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 new file mode 100644 index 00000000000000..9a863ff484dd28 --- /dev/null +++ b/gensrc/thrift/add_container_memory_check.py @@ -0,0 +1,100 @@ +#!/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*)\);$") +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 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 + 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: + 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 + + +if __name__ == "__main__": + sys.exit(main()) 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()