From 0d65c027f6a7ebeec2447cff1fd3bb21248e9e29 Mon Sep 17 00:00:00 2001 From: lihangyu Date: Thu, 12 Nov 2026 22:17:01 -0800 Subject: [PATCH] [refactor](variant) Make Variant metadata key registration allocation-safe (#70487) (cherry picked from commit 0c78027c8688456a2ff6b02ab7336b3a53856316) --- .../value/variant/variant_batch_builder.cpp | 118 +++++--- .../variant/v2/variant_path_builder.cpp | 20 +- .../variant_column_writer_reader_test.cpp | 31 +++ .../variant/variant_batch_builder_test.cpp | 258 ++++++++++++++++++ 4 files changed, 391 insertions(+), 36 deletions(-) diff --git a/be/src/core/value/variant/variant_batch_builder.cpp b/be/src/core/value/variant/variant_batch_builder.cpp index 6b2e77484e2a6c..d3bc1e1c65fd2b 100644 --- a/be/src/core/value/variant/variant_batch_builder.cpp +++ b/be/src/core/value/variant/variant_batch_builder.cpp @@ -22,11 +22,11 @@ #include #include #include -#include #include #include #include #include +#include #include #include "common/exception.h" @@ -36,6 +36,10 @@ #include "core/value/variant/variant_parquet_encoding.h" #include "util/utf8_check.h" +#ifdef BE_TEST +#include "util/debug_points.h" +#endif + namespace doris { namespace { @@ -130,9 +134,30 @@ void append_unsigned(String& output, uint64_t value, uint8_t width) { } } -template -void release_batch_container(Container& container) noexcept { - Container().swap(container); +template +void release_batch_vector(DorisVector& buffer) noexcept { + static_assert(noexcept(DorisVector())); + static_assert(noexcept(std::declval&>().swap(buffer))); + DorisVector().swap(buffer); +} + +using MetadataKeyBytes = DorisVector; +using MetadataKeys = DorisVector; + +// A key map entry points at the bytes owned by one MetadataKeyBytes. Growing MetadataKeys moves +// the inner vectors, so that move must only transfer the byte buffer. For example, reallocating +// 1024 keys to 2048 entries must not copy any key bytes or leave a map entry pointing at freed +// storage. +static_assert(std::is_nothrow_move_constructible_v, + "Variant metadata key bytes must move without allocating"); + +// An empty key owns no bytes, but its map entry still needs a stable, non-null address for both +// hashing and comparison. One sentinel byte serves every empty key. +constexpr char EMPTY_METADATA_KEY = '\0'; + +StringRef metadata_key_ref(const MetadataKeyBytes& bytes) noexcept { + DCHECK(bytes.empty() || bytes.data() != nullptr); + return {bytes.empty() ? &EMPTY_METADATA_KEY : bytes.data(), bytes.size()}; } } // namespace @@ -141,9 +166,13 @@ struct VariantMetadataBuilder::Impl { using KeyMapValue = std::pair; using KeyMap = phmap::flat_hash_map, CustomStdAllocator>; - using Keys = std::deque, CustomStdAllocator>>; + // register_key() publishes a map entry before appending the key, so a failed map allocation + // must not leave a half-inserted entry behind. phmap only moves entries after allocating a + // larger table, and moving this value type cannot throw, so a rejected growth leaves the map + // unchanged. + static_assert(std::is_nothrow_move_constructible_v); - Keys keys; + MetadataKeys keys; KeyMap key_to_temporary_id; DorisVector temporary_to_final_id; PaddedPODArray encoded; @@ -177,16 +206,33 @@ uint32_t VariantMetadataBuilder::register_key(StringRef key) { } const auto temporary_id = static_cast(_impl->keys.size()); - _impl->keys.emplace_back(); + + // Prepare every fallible step first: copy the bytes into a local owner and make room in the + // outer key vector. Only then publish the map entry and append the owner, so a rejected + // allocation leaves both containers untouched instead of needing a catch block. + MetadataKeyBytes owned_key; if (key.size != 0) { - _impl->keys.back().insert(key.data, key.data + key.size); + owned_key.insert(owned_key.end(), key.data, key.data + key.size); + } + // reserve() allocates exactly the requested capacity, so reserving size+1 every time would + // reallocate and move all previous keys on every registration. Grow geometrically only when + // the vector is full; the append below then needs no allocation. + if (_impl->keys.size() == _impl->keys.capacity()) { + _impl->keys.reserve(_impl->keys.capacity() == 0 ? 1 : _impl->keys.capacity() * 2); } - const StringRef owned_key {_impl->keys.back().data(), _impl->keys.back().size()}; + const size_t old_capacity = _impl->key_to_temporary_id.capacity(); - const auto [unused, inserted] = _impl->key_to_temporary_id.emplace(owned_key, temporary_id); - _impl->key_capacity_growths += old_capacity != _impl->key_to_temporary_id.capacity(); + // The entry is visible while `owned_key` still owns the bytes; the move below transfers that + // same buffer into `keys`, so the stored StringRef stays valid. + const auto [unused, inserted] = + _impl->key_to_temporary_id.emplace(metadata_key_ref(owned_key), temporary_id); static_cast(unused); DCHECK(inserted); + + // Commit: the outer vector has capacity for one more key and the move only transfers the byte + // buffer, so nothing below this point can allocate or throw. + _impl->keys.emplace_back(std::move(owned_key)); + _impl->key_capacity_growths += old_capacity != _impl->key_to_temporary_id.capacity(); return temporary_id; } @@ -211,8 +257,9 @@ void VariantMetadataBuilder::_abort_row() noexcept { DCHECK(_impl->row_active); DCHECK_LE(_impl->row_start_key_count, _impl->keys.size()); while (_impl->keys.size() > _impl->row_start_key_count) { - const StringRef key {_impl->keys.back().data(), _impl->keys.back().size()}; - const size_t erased = _impl->key_to_temporary_id.erase(key); + // Erase the map entry before releasing the bytes it points at; neither step allocates. + const size_t erased = + _impl->key_to_temporary_id.erase(metadata_key_ref(_impl->keys.back())); DCHECK_EQ(erased, 1); _impl->keys.pop_back(); } @@ -224,12 +271,13 @@ void VariantMetadataBuilder::_reserve_keys(size_t count) { throw Exception(ErrorCode::INVALID_ARGUMENT, "Cannot reserve Variant metadata keys after seal"); } + _impl->keys.reserve(count); _impl->key_to_temporary_id.reserve(count); } StringRef VariantMetadataBuilder::_temporary_key(uint32_t temporary_id) const noexcept { DCHECK_LT(temporary_id, _impl->keys.size()); - return {_impl->keys[temporary_id].data(), _impl->keys[temporary_id].size()}; + return metadata_key_ref(_impl->keys[temporary_id]); } size_t VariantMetadataBuilder::_key_capacity() const noexcept { @@ -255,8 +303,8 @@ void VariantMetadataBuilder::seal() { sorted_temporary_ids.push_back(temporary_id); } std::ranges::sort(sorted_temporary_ids, [this](uint32_t left, uint32_t right) { - const StringRef left_key {_impl->keys[left].data(), _impl->keys[left].size()}; - const StringRef right_key {_impl->keys[right].data(), _impl->keys[right].size()}; + const StringRef left_key = metadata_key_ref(_impl->keys[left]); + const StringRef right_key = metadata_key_ref(_impl->keys[right]); return left_key.compare(right_key) < 0; }); @@ -297,8 +345,12 @@ void VariantMetadataBuilder::seal() { append_unsigned(_impl->encoded, offset, offset_width); } for (uint32_t temporary_id : sorted_temporary_ids) { - const auto& key = _impl->keys[temporary_id]; - _impl->encoded.insert(key.data(), key.data() + key.size()); + const StringRef key = metadata_key_ref(_impl->keys[temporary_id]); + // PaddedPODArray::insert must not receive a null source pointer, and an empty key has no + // bytes to append. + if (key.size != 0) { + _impl->encoded.insert(key.data, key.data + key.size); + } } DCHECK_EQ(_impl->encoded.size(), encoded_size); _impl->sealed_key_count = count; @@ -342,9 +394,10 @@ StringRef VariantMetadataBuilder::encoded_metadata() const { PaddedPODArray VariantMetadataBuilder::_take_encoded_metadata() noexcept { DCHECK(_impl->sealed); PaddedPODArray encoded = std::move(_impl->encoded); - release_batch_container(_impl->keys); - release_batch_container(_impl->key_to_temporary_id); - release_batch_container(_impl->temporary_to_final_id); + // Destroy the scratch state directly: the encoded bytes were moved out first, and destroying + // the key map and key vectors only releases memory. Rebuilding empty containers here would + // allocate inside a noexcept cleanup path. + _impl.reset(); return encoded; } @@ -1237,14 +1290,14 @@ class VariantCollectionCore { } void release_scratch() noexcept { - release_batch_container(scalar_bytes); - release_batch_container(nodes); - release_batch_container(containers); - release_batch_container(children); - release_batch_container(scope_stack); - release_batch_container(object_id_scratch); - release_batch_container(container_plans); - release_batch_container(planned_object_children); + release_batch_vector(scalar_bytes); + release_batch_vector(nodes); + release_batch_vector(containers); + release_batch_vector(children); + release_batch_vector(scope_stack); + release_batch_vector(object_id_scratch); + release_batch_vector(container_plans); + release_batch_vector(planned_object_children); previous_object_tokens = nullptr; pending_object_tokens = nullptr; active_container = INVALID_INDEX; @@ -1575,12 +1628,13 @@ VariantBatchBuilder VariantBatchBuilder::finish_batch() { _impl->refresh_counters(_impl->metadata._key_capacity(), _impl->metadata._key_capacity_growths(), _impl->metadata.num_keys()); _impl->counters_finalized = true; + DBUG_EXECUTE_IF("VariantBatchBuilder.before_release_scratch", DBUG_RUN_CALLBACK()); #endif PaddedPODArray metadata = _impl->metadata._take_encoded_metadata(); _impl->collection.release_scratch(); - release_batch_container(_impl->roots); - release_batch_container(_impl->previous_object_tokens); - release_batch_container(_impl->pending_object_tokens); + release_batch_vector(_impl->roots); + release_batch_vector(_impl->previous_object_tokens); + release_batch_vector(_impl->pending_object_tokens); return {std::move(metadata), std::move(values), std::move(offsets)}; } diff --git a/be/src/storage/segment/variant/v2/variant_path_builder.cpp b/be/src/storage/segment/variant/v2/variant_path_builder.cpp index ea609d3bdacd47..788d7c61c7c22e 100644 --- a/be/src/storage/segment/variant/v2/variant_path_builder.cpp +++ b/be/src/storage/segment/variant/v2/variant_path_builder.cpp @@ -52,6 +52,9 @@ #include "core/value/vdatetime_value.h" #include "exec/common/variant_util.h" #include "exprs/function/parse/variant_jsonb_parse.h" +#ifdef BE_TEST +#include "util/debug_points.h" +#endif #include "util/jsonb_document.h" #include "util/jsonb_writer.h" @@ -1125,18 +1128,27 @@ Status VariantPathBuilder::append(VariantRef value, size_t row) { } try { +#ifdef BE_TEST + DBUG_EXECUTE_IF("VariantPathBuilder.before_append_value", DBUG_RUN_CALLBACK()); +#endif auto& nullable = assert_cast(*_impl->column); append_value(value, _impl->type, &nullable.get_nested_column()); - nullable.get_null_map_data().push_back(0); - } catch (const Exception&) { - if (is_array || _impl->type->get_primitive_type() == TYPE_JSONB) { + } catch (const Exception& exception) { + if (exception.code() == ErrorCode::MEM_ALLOC_FAILED || + exception.code() == ErrorCode::MEM_LIMIT_EXCEEDED || is_array || + _impl->type->get_primitive_type() == TYPE_JSONB) { throw; } RETURN_IF_ERROR(_impl->promote(jsonb_type(), false)); auto& nullable = assert_cast(*_impl->column); append_value(value, _impl->type, &nullable.get_nested_column()); - nullable.get_null_map_data().push_back(0); } + +#ifdef BE_TEST + DBUG_EXECUTE_IF("VariantPathBuilder.before_append_null_map", DBUG_RUN_CALLBACK()); +#endif + auto& nullable = assert_cast(*_impl->column); + nullable.get_null_map_data().push_back(0); _impl->rowids.push_back(static_cast(row)); _impl->logical_rows = row + 1; #if defined(BE_TEST) && !defined(BE_BENCHMARK) diff --git a/be/test/storage/variant/variant_column_writer_reader_test.cpp b/be/test/storage/variant/variant_column_writer_reader_test.cpp index a38886c880f1a2..c425fb7680a484 100644 --- a/be/test/storage/variant/variant_column_writer_reader_test.cpp +++ b/be/test/storage/variant/variant_column_writer_reader_test.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include #include @@ -284,6 +285,36 @@ TEST(VariantPathBuilderTest, PromotesValuesAndMaterializesMissingRows) { } } +TEST(VariantPathBuilderTest, MemoryFailureBypassesJsonbFallback) { + VariantBatchBuilder value_builder; + auto row = value_builder.begin_row(); + row.add_int(42); + row.finish(); + auto values = value_builder.finish_batch(); + const auto value = values.value_at(0); + + for (const auto point : + {"VariantPathBuilder.before_append_value", "VariantPathBuilder.before_append_null_map"}) { + SCOPED_TRACE(point); + segment_v2::VariantPathBuilder builder(PathInData("metric")); + ASSERT_TRUE(builder.append(value, 0).ok()); + const bool old_enable_debug_points = config::enable_debug_points; + config::enable_debug_points = true; + DebugPoints::instance()->add_with_callback( + point, std::function([] { + throw Exception(ErrorCode::MEM_ALLOC_FAILED, "Injected allocation failure"); + })); + Defer cleanup {[&] { + DebugPoints::instance()->remove(point); + config::enable_debug_points = old_enable_debug_points; + }}; + const Status status = builder.append(value, 1); + EXPECT_TRUE(status.is()) << status; + EXPECT_EQ(builder.promotion_count(), 0); + EXPECT_EQ(remove_nullable(builder.type())->get_primitive_type(), TYPE_BIGINT); + } +} + TEST(VariantPathBuilderTest, UsesBigintForEncodedIntegersWithoutPromotion) { VariantBatchBuilder value_builder; auto tiny_row = value_builder.begin_row(); diff --git a/be/test/util/variant/variant_batch_builder_test.cpp b/be/test/util/variant/variant_batch_builder_test.cpp index 257767eef20e51..415fc8bb875d15 100644 --- a/be/test/util/variant/variant_batch_builder_test.cpp +++ b/be/test/util/variant/variant_batch_builder_test.cpp @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -36,6 +37,7 @@ #include "core/value/variant/variant_parquet_encoding.h" #include "runtime/memory/mem_tracker_limiter.h" #include "runtime/thread_context.h" +#include "util/debug_points.h" #include "variant_test_utils.h" namespace doris { @@ -685,6 +687,147 @@ TEST(VariantMetadataBuilderTest, RegisterKeyPreservesValidationOnDictionaryMiss) EXPECT_EQ(metadata.metadata_ref().key_at(0), string_ref("valid")); } +TEST(VariantMetadataBuilderTest, EmptyKeyRoundTripsThroughMetadataEncoding) { + VariantMetadataBuilder metadata; + EXPECT_EQ(metadata.register_key(string_ref("")), 0); + EXPECT_EQ(metadata.register_key(string_ref("")), 0); + EXPECT_EQ(metadata.register_key(string_ref("a")), 1); + + metadata.seal(); + ASSERT_EQ(metadata.metadata_ref().dict_size(), 2); + // The empty key has no bytes to append but still occupies the first dictionary slot. + EXPECT_EQ(metadata.metadata_ref().key_at(0), string_ref("")); + EXPECT_EQ(metadata.metadata_ref().key_at(1), string_ref("a")); + EXPECT_EQ(metadata.final_id(0), 0); + EXPECT_EQ(metadata.final_id(1), 1); +} + +TEST(VariantMetadataBuilderTest, KeyReferencesStayValidAcrossKeyStorageGrowth) { + constexpr uint32_t KEY_COUNT = 8'192; + VariantMetadataBuilder metadata; + std::vector keys; + keys.reserve(KEY_COUNT); + for (uint32_t index = 0; index < KEY_COUNT; ++index) { + keys.push_back(numbered_key(index)); + } + for (uint32_t index = 0; index < KEY_COUNT; ++index) { + ASSERT_EQ(metadata.register_key(StringRef(keys[index])), index); + } + // Re-registering after several outer vector growths must still find the original entry. A + // dangling StringRef in the map would either read freed bytes or return a different id. + for (uint32_t index = 0; index < KEY_COUNT; index += 97) { + EXPECT_EQ(metadata.register_key(StringRef(keys[index])), index); + } + + metadata.seal(); + ASSERT_EQ(metadata.metadata_ref().dict_size(), KEY_COUNT); + for (uint32_t index = 0; index < KEY_COUNT; ++index) { + EXPECT_EQ(metadata.metadata_ref().key_at(metadata.final_id(index)), StringRef(keys[index])); + } +} + +// Reject exactly one allocation inside register_key and verify that the failed registration +// publishes nothing and can be retried with the same temporary id. Sweeping the allowance covers +// the byte copy, the outer vector reservation and the first key-map insertion. +TEST(VariantMetadataBuilderTest, RegisterKeyFailureIsAtomicAcrossAllocationSizes) { + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantMetadataBuilderKeyFailure"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + ++enable_thread_catch_bad_alloc; + Defer restore_bad_alloc {[] { --enable_thread_catch_bad_alloc; }}; + + const std::string key = "a metadata key long enough to own a buffer"; + for (const int64_t allowance : + {int64_t {0}, int64_t {1}, int64_t {4}, int64_t {16}, int64_t {32}, int64_t {64}, + int64_t {128}, int64_t {1'024}, int64_t {8'192}}) { + VariantMetadataBuilder metadata; + thread_context()->thread_mem_tracker_mgr->flush_untracked_mem(); + const int64_t baseline = tracker->consumption(); + tracker->set_limit(baseline + allowance); + bool failed = false; + try { + metadata.register_key(StringRef(key)); + } catch (const Exception& exception) { + failed = true; + EXPECT_EQ(exception.code(), ErrorCode::MEM_ALLOC_FAILED); + } + EXPECT_EQ(metadata.num_keys(), failed ? 0 : 1); + tracker->set_limit(-1); + ASSERT_EQ(metadata.num_keys(), failed ? 0 : 1); + // A rejected registration must leave no key and no map entry behind. + EXPECT_EQ(metadata.register_key(StringRef(key)), 0); + EXPECT_EQ(metadata.num_keys(), 1); + } +} + +// Fill to the point where the outer key vector has no spare capacity, then reject the growth +// allocation. The key bytes have already been copied at that point, so this also proves the +// copy is discarded without leaving an unindexed key. +TEST(VariantMetadataBuilderTest, RegisterKeyFailureDuringKeyStorageGrowthIsAtomic) { + // The key vector grows geometrically and is full at every power of two. + constexpr uint32_t FILLED = 1'024; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantMetadataBuilderKeyGrowth"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + VariantMetadataBuilder metadata; + for (uint32_t index = 0; index < FILLED; ++index) { + ASSERT_EQ(metadata.register_key(StringRef(numbered_key(index))), index); + } + ASSERT_EQ(metadata.num_keys(), FILLED); + + const std::string next = numbered_key(FILLED); + thread_context()->thread_mem_tracker_mgr->flush_untracked_mem(); + const int64_t baseline = tracker->consumption(); + // The next key copies its bytes first and then grows the outer vector; allow only the copy. + tracker->set_limit(baseline + static_cast(next.size())); + ++enable_thread_catch_bad_alloc; + Defer restore_bad_alloc {[] { --enable_thread_catch_bad_alloc; }}; + expect_builder_exception_code(ErrorCode::MEM_ALLOC_FAILED, + [&] { metadata.register_key(StringRef(next)); }); + EXPECT_EQ(metadata.num_keys(), FILLED); + // Existing entries remain reachable through their map StringRefs. + for (uint32_t index = 0; index < FILLED; index += 61) { + EXPECT_EQ(metadata.register_key(StringRef(numbered_key(index))), index); + } + + tracker->set_limit(-1); + EXPECT_EQ(metadata.register_key(StringRef(next)), FILLED); + EXPECT_EQ(metadata.num_keys(), FILLED + 1); +} + +// Fill just below the flat_hash_map growth threshold with spare capacity in the outer key vector, +// so the rejected allocation is the map table growth rather than a key buffer or vector growth. +TEST(VariantMetadataBuilderTest, RegisterKeyFailureDuringKeyMapGrowthIsAtomic) { + // phmap keeps a 7/8 load factor: a 2047-slot table grows on its 1793rd entry, where the outer + // key vector still has room, so only the map table allocation can fail. + constexpr uint32_t FILLED = 1'792; + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantMetadataBuilderMapGrowth"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + VariantMetadataBuilder metadata; + for (uint32_t index = 0; index < FILLED; ++index) { + ASSERT_EQ(metadata.register_key(StringRef(numbered_key(index))), index); + } + ASSERT_EQ(metadata.num_keys(), FILLED); + + const std::string next = numbered_key(FILLED); + thread_context()->thread_mem_tracker_mgr->flush_untracked_mem(); + const int64_t baseline = tracker->consumption(); + tracker->set_limit(baseline + static_cast(next.size())); + ++enable_thread_catch_bad_alloc; + Defer restore_bad_alloc {[] { --enable_thread_catch_bad_alloc; }}; + expect_builder_exception_code(ErrorCode::MEM_ALLOC_FAILED, + [&] { metadata.register_key(StringRef(next)); }); + EXPECT_EQ(metadata.num_keys(), FILLED); + for (uint32_t index = 0; index < FILLED; index += 61) { + EXPECT_EQ(metadata.register_key(StringRef(numbered_key(index))), index); + } + + tracker->set_limit(-1); + EXPECT_EQ(metadata.register_key(StringRef(next)), FILLED); + EXPECT_EQ(metadata.num_keys(), FILLED + 1); +} + std::string build_array_with_nulls(uint32_t count, uint8_t* value_header_out) { const OwnedBuilderValue owned = build_owned_value([count](VariantBatchBuilder::Row& row) { auto array = row.start_array(); @@ -1579,6 +1722,36 @@ TEST(VariantBatchBuilderTest, AbortAndRowErrorsRollbackBeforeTheNextRow) { EXPECT_EQ(required_field(block.value_at(0), "retained").get_int(), 9); } +TEST(VariantBatchBuilderTest, AbortedRowReleasesItsKeysAndTemporaryIds) { + VariantBatchBuilder builder; + { + auto row = builder.begin_row(); + auto object = row.start_object(); + object.add_key(string_ref("first")); + row.add_int(1); + object.add_key(string_ref("second")); + row.add_int(2); + object.finish(); + row.abort(); + } + { + auto row = builder.begin_row(); + auto object = row.start_object(); + // Aborting the first row must remove its keys from both the dictionary and the map, so + // this key starts a fresh key vector at temporary id 0. + object.add_key(string_ref("third")); + row.add_int(3); + object.finish(); + row.finish(); + } + + VariantBatchBuilder block = builder.finish_batch(); + ASSERT_EQ(block.num_rows(), 1); + ASSERT_EQ(block.metadata_ref().dict_size(), 1); + EXPECT_EQ(block.metadata_ref().key_at(0), string_ref("third")); + EXPECT_EQ(required_field(block.value_at(0), "third").get_int(), 3); +} + TEST(VariantBatchBuilderTest, CopiesBorrowedKeysStringsAndBinaryBeforeRowReturns) { VariantBatchBuilder builder; std::string key = "borrowed"; @@ -1901,6 +2074,91 @@ TEST(VariantBatchBuilderTest, CapacityGrowthIsBatchBoundedRatherThanPerRow) { EXPECT_GE(large.pending_object_token_capacity, small.pending_object_token_capacity); } +TEST(VariantBatchBuilderTest, ReleasesScratchWhenAllocationsAreRejected) { + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantBatchBuilderScratchRelease"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + VariantBatchBuilder builder; + auto row = builder.begin_row(); + auto object = row.start_object(); + object.add_key(string_ref("key")); + row.add_int(42); + object.finish(); + row.finish(); + + const bool old_debug_points = config::enable_debug_points; + config::enable_debug_points = true; + constexpr auto point = "VariantBatchBuilder.before_release_scratch"; + bool injected = false; + DebugPoints::instance()->add_with_callback(point, std::function([&] { + tracker->set_limit(1); + injected = true; + })); + Defer cleanup {[&] { + DebugPoints::instance()->remove(point); + config::enable_debug_points = old_debug_points; + tracker->set_limit(-1); + }}; + ++enable_thread_catch_bad_alloc; + Defer restore_catch {[] { --enable_thread_catch_bad_alloc; }}; + VariantBatchBuilder block = builder.finish_batch(); + tracker->set_limit(-1); + ASSERT_TRUE(injected); + ASSERT_EQ(block.num_rows(), 1); + ASSERT_EQ(block.metadata_ref().dict_size(), 1); + EXPECT_EQ(required_field(block.value_at(0), "key").get_int(), 42); +} + +TEST(VariantBatchBuilderTest, FailedKeyRegistrationLeavesNoUnindexedKey) { + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantBatchBuilderKeyRegistration"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + // A nonempty key first allocates its bytes; an empty key still needs room in the key vector. + // Neither may publish a map entry before the key is committed. + for (const auto key : {string_ref("new key"), string_ref("")}) { + VariantMetadataBuilder metadata; + { + tracker->set_limit(1); + ++enable_thread_catch_bad_alloc; + Defer restore {[&] { + --enable_thread_catch_bad_alloc; + tracker->set_limit(-1); + }}; + expect_builder_exception_code(ErrorCode::MEM_ALLOC_FAILED, + [&] { metadata.register_key(key); }); + } + EXPECT_EQ(metadata.num_keys(), 0); + EXPECT_EQ(metadata.register_key(key), 0); + EXPECT_EQ(metadata.num_keys(), 1); + } +} + +TEST(VariantBatchBuilderTest, FailedSealLeavesTheBatchTerminal) { + auto tracker = MemTrackerLimiter::create_shared(MemTrackerLimiter::Type::QUERY, + "VariantBatchBuilderFailedSeal"); + auto scoped_tracker = SwitchThreadMemTrackerLimiter(tracker); + VariantBatchBuilder builder; + { + auto row = builder.begin_row(); + row.add_int(1); + row.finish(); + } + + tracker->set_limit(1); + ++enable_thread_catch_bad_alloc; + Defer restore {[&] { + --enable_thread_catch_bad_alloc; + tracker->set_limit(-1); + }}; + expect_builder_exception_code(ErrorCode::MEM_ALLOC_FAILED, [&] { builder.finish_batch(); }); + tracker->set_limit(-1); + + // finish_batch marks the encoding unit terminal before seal(), so a failed seal cannot be + // retried or extended with another row. + EXPECT_THROW(builder.begin_row(), Exception); + EXPECT_THROW(builder.finish_batch(), Exception); +} + TEST(VariantBatchBuilderTest, BatchLifetimeAllocationsAreMemTrackerVisible) { constexpr size_t ROWS = 2'048; constexpr size_t PAYLOAD_BYTES = 4'096;