github-actions[bot] commented on code in PR #66820: URL: https://github.com/apache/doris/pull/66820#discussion_r3795664791
########## 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 <thrift/TToString.h> +#include <thrift/protocol/TProtocol.h> + +#include <algorithm> +#include <cstddef> +#include <cstdint> +#include <memory> +#include <ostream> +#include <utility> +#include <vector> + +namespace doris { + +class ThriftContainerMemoryCharge { +public: + virtual ~ThriftContainerMemoryCharge() = default; +}; + +class ThriftContainerMemoryChecker { +public: + virtual ~ThriftContainerMemoryChecker() = default; + virtual std::shared_ptr<ThriftContainerMemoryCharge> reserve_container_memory( + uint32_t count, size_t element_size) = 0; + virtual void retain_temporary_container_charge( + std::shared_ptr<ThriftContainerMemoryCharge> charge) = 0; +}; + +template <typename T> +class ThriftMemoryTrackedVector : public std::vector<T> { + using Base = std::vector<T>; + +public: + using Base::Base; + using Base::operator=; + + ThriftMemoryTrackedVector() = default; + ThriftMemoryTrackedVector(const ThriftMemoryTrackedVector&) = default; Review Comment: [P1] Do not share one reservation across deep copies. These default copy operations deep-copy the `std::vector` storage but only copy `_memory_charge`. The v1 page-index path reaches this directly: after decoding a charged `OffsetIndex`, `vparquet_reader.cpp:1233` lvalue-copies it into `_col_offsets`, so two large buffers coexist while only one is admitted. Lines 1322-1323 also copy charged `min_values`/`max_values` into ordinary vectors whose storage outlives the source charge. Please make copied storage acquire its own admission/charge (or make these transfers move-only and preserve ownership), and add a limit test that copies one admissible decoded index under headroom below the two-copy peak. ########## 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 <thrift/TToString.h> +#include <thrift/protocol/TProtocol.h> + +#include <algorithm> +#include <cstddef> +#include <cstdint> +#include <memory> +#include <ostream> +#include <utility> +#include <vector> + +namespace doris { + +class ThriftContainerMemoryCharge { +public: + virtual ~ThriftContainerMemoryCharge() = default; +}; + +class ThriftContainerMemoryChecker { +public: + virtual ~ThriftContainerMemoryChecker() = default; + virtual std::shared_ptr<ThriftContainerMemoryCharge> reserve_container_memory( + uint32_t count, size_t element_size) = 0; + virtual void retain_temporary_container_charge( + std::shared_ptr<ThriftContainerMemoryCharge> charge) = 0; +}; + +template <typename T> +class ThriftMemoryTrackedVector : public std::vector<T> { + using Base = std::vector<T>; + +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<ThriftContainerMemoryCharge> 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<const Base&>(lhs) == static_cast<const Base&>(rhs); + } + + friend bool operator<(const ThriftMemoryTrackedVector& lhs, + const ThriftMemoryTrackedVector& rhs) { + return static_cast<const Base&>(lhs) < static_cast<const Base&>(rhs); + } + + friend std::ostream& operator<<(std::ostream& out, const ThriftMemoryTrackedVector& values) { + return out << apache::thrift::to_string(static_cast<const Base&>(values)); + } + +private: + std::shared_ptr<ThriftContainerMemoryCharge> _memory_charge; +}; + +template <typename Container> +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<ThriftContainerMemoryChecker*>(protocol); checker != nullptr) { + const size_t elements = std::max<size_t>(count, container->capacity()); Review Comment: [P1] Budget the allocation that `resize` will actually make, including reallocation overlap. Thrift accepts repeated field IDs, so an external page index can present a tracked list first with N entries and then N+1. This hook replaces the N-element charge with only `max(N+1, capacity)==N+1`; on current libstdc++, the following resize grows capacity to 2N while the old N buffer is still live. The result is 2N steady / 3N transient storage admitted as N+1. Please tie the charge to actual allocator capacity/peak (without releasing the old charge early) and add repeated-list compact/binary cases. ########## 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 Review Comment: [P2] Avoid leaving no-op generated targets permanently stale. After this checker changes in an incremental checkout, this prerequisite reruns every target; Thrift 0.16 preserves an existing `_types.cpp` mtime when its content is unchanged, and the postprocessor also does not write when it finds no resize. IDLs such as `Metrics.thrift`, `QueryCache.thrift`, and `QueryPlanExtra.thrift` therefore remain older than the checker and recompile on every later `generated-source.sh noclean`. Please complete the rule with a stamp/mtime strategy for no-op outputs and test that a second noclean pass does no work. ########## be/src/util/thrift_util.cpp: ########## @@ -56,6 +61,82 @@ class TProtocol; #include <thread> 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<apache::thrift::protocol::TProtocol> 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<ThriftContainerMemoryCharge> reserve_container_memory( + uint32_t count, size_t element_size) override { + if (count > std::numeric_limits<size_t>::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<size_t>(count) * element_size; + if (bytes > static_cast<size_t>(std::numeric_limits<int64_t>::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<int64_t>(bytes)); + if (!status.ok()) { + throw Exception(status); + } + return std::make_shared<ReservedMemoryCharge>(_memory_manager->take_reserved_memory()); Review Comment: [P1] Avoid one uncharged heap token per nested container. `reserve_container_memory(0, ...)` still reaches this `make_shared`, and ordinary generated vectors then append another `shared_ptr` to `_temporary_charges`. For a reachable `list<list<...>>` such as `TRepeatNode.grouping_list`, N empty children admit only the outer `N * sizeof(vector)` storage but allocate N control blocks/tokens plus N temporary-vector entries outside that admission, defeating the boundary by another O(N) amount. Please aggregate temporary reservations (and special-case zero bytes), or include this bookkeeping in admission; cover many empty nested lists under constrained headroom. ########## 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 Review Comment: [P1] Make this source/header transformation failure-atomic. The recipe tracks only `_types.cpp`, while the checker writes that source first and its sibling header second. If header processing exits nonzero after the source write, Make keeps the newly dated `.cpp`; after the transient cause is repaired, the normal `noclean` retry sees this sole target newer than the IDL/checker and skips both transformations. The ordinary-vector header is still compile-compatible with the source hook, so retained page-index fields silently fall back to protocol-lifetime charging. Please use an atomic/grouped completion stamp (or delete/invalidate the target on failure) and test a failure between the two writes followed by a noclean retry. ########## 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<indent>\s*)(?P<container>.+)\.resize\((?P<size>_size\d*)\);$") Review Comment: [P1] Cover generated maps and sets as well as list resizes. This is the only source pattern the postprocessor instruments, but Thrift 0.16 fills maps with `operator[]` and sets with `insert`, so those node allocations never call `reserve_thrift_container_memory()`. This is reachable in the contextless plan-fragment deserialize path (`TPipelineFragmentParamsList` contains several scalar maps), where a compact wire map can expand into much larger `std::map` nodes without either task or fallback-process admission. Please instrument these generated allocation loops using the declared target type and add compact/binary map and set budget regressions. ########## 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): Review Comment: [P2] Register this test with a normal repository test target. No gensrc Make target, `build.sh`/`run-*-ut.sh` path, build-support test, or checked-in workflow invokes this file; the only entry point is the direct `unittest.main()` block. As added, both the postprocessor-shape assertion and the Make prerequisite assertion can regress while standard CI remains green. Please wire it into an executed target, then extend it with the real two-pass noclean case. ########## be/test/format_v2/parquet/parquet_statistics_test.cpp: ########## @@ -1276,7 +1293,65 @@ TEST(ParquetBloomFilterPruningTest, NativeBloomReportsConservativeReadOutcomes) auto truncated = make_valid_bloom(); Review Comment: [P2] Make this fixture reach the v2 successful-short-read guard. `StatisticsMemoryFileReader::size()` equals `_bytes.size()` and its read either fills the whole request or returns `IOError`; after this shrink, `validate_native_bloom_filter_layout()` rejects `header_size + numBytes > file_size` before the payload read. Thus deleting the new `bytes_read != bloom_filter->size()` check still leaves the test green. Let the fake advertise the validated logical range while returning a shorter successful payload read, assert that branch is reached, and keep this range-truncation case separately. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
