From 0328b96eb70bfcf0dc2ca9b32b9cfc89f93609df Mon Sep 17 00:00:00 2001 From: Jerry Zhao Date: Tue, 20 Oct 2026 04:35:24 +0000 Subject: [PATCH] perf: lazily initialize ScanHashMap resize targets --- include/lavik/storage/scan_hash_map.h | 219 ++++++++++++++++++---- tests/record_location_test.cpp | 33 ++++ tests/scan_hash_map_test.cpp | 252 ++++++++++++++++++++++++++ 3 files changed, 465 insertions(+), 39 deletions(-) diff --git a/include/lavik/storage/scan_hash_map.h b/include/lavik/storage/scan_hash_map.h index 35eac75f9..0714bdc5c 100644 --- a/include/lavik/storage/scan_hash_map.h +++ b/include/lavik/storage/scan_hash_map.h @@ -1378,7 +1378,8 @@ class ScanHashMap { tables_[0].exponent_ < MaxBucketExponent) return true; const Table& table = tables_[0]; - const Bucket* bucket = &table.buckets_[Hash(digest) & BucketMask(table)]; + const Bucket* bucket = + table.template ReadBucket(Hash(digest) & BucketMask(table)); while (bucket != nullptr) { for (auto handle : bucket->entries_) if (handle == 0) return false; @@ -1452,13 +1453,22 @@ class ScanHashMap { std::numeric_limits::max() / sizeof(Bucket)) { return std::numeric_limits::max(); } - return add(required, add(next_buckets * sizeof(Bucket), - kSlowPathBookkeepingAllowance)); + std::size_t buckets_bytes = next_buckets * sizeof(Bucket); + if (next_buckets > kBucketsPerInitializationBlock) { + const std::size_t words = + (next_buckets / kBucketsPerInitializationBlock + 63) / 64; + buckets_bytes = add(buckets_bytes, sizeof(OverflowBuckets)); + buckets_bytes = + add(buckets_bytes, + AllocatorUsableSizeForRequest(words * sizeof(std::uint64_t))); + } + return add(required, add(buckets_bytes, kSlowPathBookkeepingAllowance)); } const Table& target = Rehashing() ? tables_[1] : tables_[0]; const std::uint64_t hash = Hash(digest); - const Bucket* bucket = &target.buckets_[hash & BucketMask(target)]; + const Bucket* bucket = target.ReadBucket(hash & BucketMask(target)); + if (bucket == nullptr) return required; while (bucket != nullptr) { if (std::any_of(bucket->entries_.begin(), bucket->entries_.end(), [](EntryHandle handle) { return handle == 0; })) { @@ -1535,6 +1545,19 @@ class ScanHashMap { // Exposes incremental-rehash state for capacity diagnostics and tests. bool rehashing() const noexcept { return Rehashing(); } + // Direct buckets reserved by a resize but not yet initialized. Read-only + // access treats these as empty; useful for capacity and resize diagnostics. + std::size_t uninitialized_bucket_count() const noexcept { + std::size_t count = 0; + for (const Table& table : tables_) { + if (table.lazy_buckets_) { + count += table.overflow_->uninitialized_blocks_ * + kBucketsPerInitializationBlock; + } + } + return count; + } + // Returns the low bucket-address bits callers must retain with an Entry // address across suspension. It is computed while the Entry is known live; // FindAddress can then validate membership without dereferencing stale @@ -1546,12 +1569,14 @@ class ScanHashMap { return static_cast(EntryHash(entry)); } - Entry* Find(const Digest& digest, std::string_view key) { + [[gnu::always_inline]] Entry* Find(const Digest& digest, + std::string_view key) { AdvanceRehashIfNeeded(); return FindWithoutStep(digest, key); } - const Entry* Find(const Digest& digest, std::string_view key) const { + [[gnu::always_inline]] const Entry* Find(const Digest& digest, + std::string_view key) const { return FindWithoutStep(digest, key); } @@ -1584,9 +1609,9 @@ class ScanHashMap { FindCandidateInTable(tables_[0], digest, key, hash, accept)) { return found; } - return Rehashing() - ? FindCandidateInTable(tables_[1], digest, key, hash, accept) - : nullptr; + return Rehashing() ? FindCandidateInTable(tables_[1], digest, key, + hash, accept) + : nullptr; } bool Contains(const Entry* entry, std::uint32_t hash) const noexcept { @@ -1619,7 +1644,7 @@ class ScanHashMap { for (int t = 0, count = Rehashing() ? 2 : 1; t < count; ++t) { const Table& table = tables_[t]; if (table.buckets_) - __builtin_prefetch(&table.buckets_[hash & BucketMask(table)], 0, 3); + __builtin_prefetch(table.ReadBucket(hash & BucketMask(table)), 0, 3); } } @@ -1671,7 +1696,7 @@ class ScanHashMap { if (!table.buckets_) { continue; } - Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + Bucket* bucket = table.ReadBucket(hash & BucketMask(table)); while (bucket != nullptr) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { if (Occupied(*bucket, slot) && @@ -1765,7 +1790,7 @@ class ScanHashMap { if (!table.buckets_) { continue; } - Bucket* top = &table.buckets_[hash & BucketMask(table)]; + Bucket* top = table.ReadBucket(hash & BucketMask(table)); for (Bucket* bucket = top; bucket != nullptr; bucket = Chained(*bucket) ? Child(table, bucket) : nullptr) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { @@ -1803,7 +1828,7 @@ class ScanHashMap { if (!table.buckets_) { continue; } - Bucket* top = &table.buckets_[hash & BucketMask(table)]; + Bucket* top = table.ReadBucket(hash & BucketMask(table)); for (Bucket* bucket = top; bucket != nullptr; bucket = Chained(*bucket) ? Child(table, bucket) : nullptr) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { @@ -1873,7 +1898,12 @@ class ScanHashMap { continue; } - const Bucket* bucket = &table.buckets_[cursor->bucket_]; + const Bucket* bucket = table.ReadBucket(cursor->bucket_); + if (bucket == nullptr) { + assert(cursor->chain_ == 0 && cursor->slot_ == 0); + ++cursor->bucket_; + continue; + } for (std::size_t chain = 0; chain < cursor->chain_; ++chain) { assert(Chained(*bucket)); bucket = Child(table, bucket); @@ -2026,6 +2056,12 @@ class ScanHashMap { static_assert(sizeof(Bucket) == 64); + // Bound each first-write initialization independently of the table size. + // Small tables and offline recovery remain eager. Only a resize target can + // contain uninitialized blocks, while existing resize exclusion keeps remote + // readers out; no atomics or worker-blocking synchronization is needed. + static constexpr std::size_t kBucketsPerInitializationBlock = 1024; + using BucketOwner = std::unique_ptr>; using BucketOwnerVector = std::vector>; @@ -2043,10 +2079,16 @@ class ScanHashMap { struct OverflowBuckets { explicit OverflowBuckets(RetainedAllocationDomain domain) : buckets_(RetainedAllocator(domain)), - free_ids_(RetainedAllocator(domain)) {} + free_ids_(RetainedAllocator(domain)), + initialized_blocks_(RetainedAllocator(domain)) {} BucketOwnerVector buckets_; OverflowFreeIds free_ids_; + // Resize-only state shares the optional overflow allocation so the many + // empty partition/DB maps do not each acquire another pointer or vector. + std::vector> + initialized_blocks_; + std::size_t uninitialized_blocks_ = 0; }; using OverflowOwner = @@ -2062,22 +2104,103 @@ class ScanHashMap { // Direct tables never grow in place. An explicit array allocation gives // optional shrinking a rejection channel without changing vector's fatal // allocation contract or bypassing admission for later overflow growth. - bool AllocateBuckets(std::size_t count, - bool independent_admission = false) { + bool AllocateBuckets(std::size_t count, bool independent_admission = false, + bool defer_initialization = false) { assert(!buckets_); auto domain = buckets_.get_deleter().domain_; if (independent_admission) domain.externally_admitted_ = false; + const bool lazy = + defer_initialization && count > kBucketsPerInitializationBlock; + OverflowOwner prepared(nullptr, overflow_.get_deleter()); + if (lazy) { + prepared = TryMakeRetainedUnique( + domain, buckets_.get_deleter().domain_); + if (prepared == nullptr) return false; + const std::size_t blocks = count / kBucketsPerInitializationBlock; + const std::size_t words = (blocks + 63) / 64; + if (!TryReserveOverflowVector(&prepared->initialized_blocks_, words, + domain)) { + return false; + } + prepared->initialized_blocks_.resize(words, 0); + prepared->uninitialized_blocks_ = blocks; + } auto* data = static_cast(TryAllocateRetainedBytes( domain, count * sizeof(Bucket), alignof(Bucket))); if (data == nullptr) return false; - std::uninitialized_value_construct_n(data, count); + if (!lazy) std::uninitialized_value_construct_n(data, count); buckets_.reset(data); + overflow_ = std::move(prepared); + lazy_buckets_ = lazy; return true; } + // Never form a readable Bucket in storage whose lifetime has not started. + // All observers, including scans, address checks, and admission estimates, + // must interpret an uninitialized block as empty without touching it. + template + const Bucket* ReadBucket(std::size_t index) const noexcept { + if (!buckets_) return nullptr; + if constexpr (CheckInitialization) { + if (lazy_buckets_) { + const std::size_t block = index / kBucketsPerInitializationBlock; + if ((overflow_->initialized_blocks_[block / 64] & + (std::uint64_t{1} << (block % 64))) == 0) { + return nullptr; + } + } + } else { + // Primary-table lookup keeps its original path: table zero is always + // fully constructed, including throughout the next resize. + assert(!lazy_buckets_); + } + return buckets_.get() + index; + } + + template + Bucket* ReadBucket(std::size_t index) noexcept { + return const_cast( + std::as_const(*this).template ReadBucket(index)); + } + + Bucket* WriteBucket(std::size_t index) { + assert(buckets_); + if (lazy_buckets_) { + const std::size_t block = index / kBucketsPerInitializationBlock; + auto& word = overflow_->initialized_blocks_[block / 64]; + const std::uint64_t bit = std::uint64_t{1} << (block % 64); + if ((word & bit) == 0) { + std::uninitialized_value_construct_n( + buckets_.get() + block * kBucketsPerInitializationBlock, + kBucketsPerInitializationBlock); + // A foreground insert may initialize a block before migration gets + // there. Publish only after construction and never clear it again. + word |= bit; + --overflow_->uninitialized_blocks_; + } + } + return buckets_.get() + index; + } + + void FinishBucketInitialization() { + if (!lazy_buckets_) return; + // Each source bucket prepares its destination(s), including empty ones, + // so completing migration leaves no deferred work on the settled table. + assert(overflow_->uninitialized_blocks_ == 0); + lazy_buckets_ = false; + if (overflow_->buckets_.empty()) { + overflow_.reset(); + } else { + decltype(overflow_->initialized_blocks_) empty( + overflow_->initialized_blocks_.get_allocator()); + overflow_->initialized_blocks_.swap(empty); + } + } + DirectBuckets buckets_{nullptr, BucketArrayDeleter{}}; OverflowOwner overflow_; std::uint8_t exponent_ = 0; + bool lazy_buckets_ = false; std::size_t used_ = 0; }; @@ -2482,7 +2605,9 @@ class ScanHashMap { assert(arena_ != nullptr); tables_[1] = Table(arena_->allocation_domain()); tables_[1].exponent_ = tables_[0].exponent_ + 1; - if (!tables_[1].AllocateBuckets(std::size_t{1} << tables_[1].exponent_)) { + if (!tables_[1].AllocateBuckets(std::size_t{1} << tables_[1].exponent_, + /*independent_admission=*/false, + /*defer_initialization=*/true)) { std::terminate(); } tables_[1].used_ = 0; @@ -2513,7 +2638,9 @@ class ScanHashMap { // Deletion has no enclosing growth permit, even for a primary index. // Reserve the temporary old+new footprint before publishing the target; // rejection must never turn a successful Erase into a failed command. - if (!prepared.AllocateBuckets(buckets / 2, true)) { + if (!prepared.AllocateBuckets(buckets / 2, + /*independent_admission=*/true, + /*defer_initialization=*/true)) { EndResize(); return false; } @@ -2543,6 +2670,7 @@ class ScanHashMap { ++rehash_index_; if (rehash_index_ == BucketCount(tables_[0])) { assert(tables_[0].used_ == 0); + tables_[1].FinishBucketInitialization(); tables_[0] = std::move(tables_[1]); tables_[1] = Table(arena_->allocation_domain()); rehash_index_ = kNotRehashing; @@ -2553,7 +2681,8 @@ class ScanHashMap { static std::size_t FreeSlots(const Table& table, std::size_t bucket_index) noexcept { - const Bucket* bucket = &table.buckets_[bucket_index]; + const Bucket* bucket = table.ReadBucket(bucket_index); + if (bucket == nullptr) return kEntriesPerBucket; std::size_t free = 0; while (bucket != nullptr) { free += static_cast(std::count( @@ -2574,7 +2703,7 @@ class ScanHashMap { static bool AppendEmptyBuckets(Table* table, std::size_t bucket_index, std::size_t count) { - Bucket* tail = &table->buckets_[bucket_index]; + Bucket* tail = table->WriteBucket(bucket_index); while (Chained(*tail)) tail = Child(*table, tail); while (count-- != 0) { const ChildAllocation child = AllocateChild(table); @@ -2593,8 +2722,8 @@ class ScanHashMap { std::size_t slot_ = 0; }; std::array cursors = { - Cursor{.bucket_ = &target->buckets_[target_indexes[0]]}, - Cursor{.bucket_ = &target->buckets_[target_indexes[1]]}}; + Cursor{.bucket_ = target->ReadBucket(target_indexes[0])}, + Cursor{.bucket_ = target->ReadBucket(target_indexes[1])}}; for (RehashMove& move : *plan) { Cursor& cursor = cursors[move.split_]; while (cursor.bucket_ != nullptr) { @@ -2718,10 +2847,13 @@ class ScanHashMap { return nullptr; } + template Entry* FindInTable(Table& table, const Digest& digest, std::string_view key, std::uint64_t hash) { if (!table.buckets_) return nullptr; - Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + Bucket* bucket = table.template ReadBucket( + hash & BucketMask(table)); + if (bucket == nullptr) return nullptr; const std::uint8_t tag = HashTag(hash); if (Entry* found = FindInBucket(*bucket, digest, key, tag)) return found; return Chained(*bucket) @@ -2757,10 +2889,13 @@ class ScanHashMap { return nullptr; } + template const Entry* FindInTable(const Table& table, const Digest& digest, std::string_view key, std::uint64_t hash) const { if (!table.buckets_) return nullptr; - const Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + const Bucket* bucket = table.template ReadBucket( + hash & BucketMask(table)); + if (bucket == nullptr) return nullptr; const std::uint8_t tag = HashTag(hash); if (const Entry* found = FindInBucket(*bucket, digest, key, tag)) return found; @@ -2781,7 +2916,7 @@ class ScanHashMap { if (!table.buckets_) { continue; } - const Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + const Bucket* bucket = table.ReadBucket(hash & BucketMask(table)); while (bucket != nullptr) { if constexpr (FilterTag) { // Flush receipts already retain the key's tag. Use the same SIMD @@ -2815,7 +2950,7 @@ class ScanHashMap { if (!table.buckets_) { return; } - Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + Bucket* bucket = table.ReadBucket(hash & BucketMask(table)); const std::uint8_t tag = HashTag(hash); while (bucket != nullptr) { for (std::uint32_t matches = MatchingTags(*bucket, tag); matches != 0; @@ -2864,13 +2999,15 @@ class ScanHashMap { return nullptr; } - template + template Entry* FindCandidateInTable(Table& table, const Digest& digest, std::string_view key, std::uint64_t hash, Predicate& accept) { if (!table.buckets_) return nullptr; - Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + Bucket* bucket = table.template ReadBucket( + hash & BucketMask(table)); const std::uint8_t tag = HashTag(hash); + if (bucket == nullptr) return nullptr; if (Entry* found = FindCandidateInBucket(*bucket, digest, key, tag, accept)) return found; return Chained(*bucket) @@ -2879,7 +3016,11 @@ class ScanHashMap { : nullptr; } - Entry* FindWithoutStep(const Digest& digest, std::string_view key) { + // Resize preparation is much larger than lookup. Keep it from changing the + // compiler's inlining decision for the primary-table read path at callers + // that also insert; only overflow and second-table lookup stay out of line. + [[gnu::always_inline]] Entry* FindWithoutStep(const Digest& digest, + std::string_view key) { const std::uint64_t hash = Hash(digest); if (Entry* found = FindInTable(tables_[0], digest, key, hash); found != nullptr) { @@ -2891,8 +3032,8 @@ class ScanHashMap { return nullptr; } - const Entry* FindWithoutStep(const Digest& digest, - std::string_view key) const { + [[gnu::always_inline]] const Entry* FindWithoutStep( + const Digest& digest, std::string_view key) const { const std::uint64_t hash = Hash(digest); if (const Entry* found = FindInTable(tables_[0], digest, key, hash); found != nullptr) { @@ -2911,18 +3052,18 @@ class ScanHashMap { [[gnu::noinline]] Entry* FindInRehashTable(const Digest& digest, std::string_view key, std::uint64_t hash) { - return FindInTable(tables_[1], digest, key, hash); + return FindInTable(tables_[1], digest, key, hash); } [[gnu::noinline]] const Entry* FindInRehashTable(const Digest& digest, std::string_view key, std::uint64_t hash) const { - return FindInTable(tables_[1], digest, key, hash); + return FindInTable(tables_[1], digest, key, hash); } bool AddToTable(Table& table, EntryHandle entry, std::uint64_t hash) { const std::uint8_t tag = HashTag(hash); - Bucket* bucket = &table.buckets_[hash & BucketMask(table)]; + Bucket* bucket = table.WriteBucket(hash & BucketMask(table)); while (true) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { if (!Occupied(*bucket, slot)) { @@ -3017,7 +3158,7 @@ class ScanHashMap { if (!table.buckets_) { return; } - const Bucket* bucket = &table.buckets_[index]; + const Bucket* bucket = table.ReadBucket(index); while (bucket != nullptr) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { if (Occupied(*bucket, slot)) { @@ -3039,7 +3180,7 @@ class ScanHashMap { template bool EmitBucketWhile(const Table& table, std::uint64_t index, Fn& fn) const { if (!table.buckets_) return true; - const Bucket* bucket = &table.buckets_[index]; + const Bucket* bucket = table.ReadBucket(index); while (bucket != nullptr) { for (std::size_t slot = 0; slot < kEntriesPerBucket; ++slot) { if (Occupied(*bucket, slot) && @@ -3064,7 +3205,7 @@ class ScanHashMap { void DestroyTable(Table& table, bool destroy_entries) noexcept { const std::size_t count = BucketCount(table); for (std::size_t index = 0; index < count; ++index) { - Bucket* top = &table.buckets_[index]; + Bucket* top = table.ReadBucket(index); Bucket* bucket = top; while (bucket != nullptr) { const bool chained = Chained(*bucket); diff --git a/tests/record_location_test.cpp b/tests/record_location_test.cpp index 3e102c39f..6e1b50dfa 100644 --- a/tests/record_location_test.cpp +++ b/tests/record_location_test.cpp @@ -133,6 +133,39 @@ TEST(IndexReadersTest, LastReaderNotifiesOnceAndIntentSurvivesResizeScope) { EXPECT_EQ(notifications, 1); } +TEST(IndexReadersTest, LazyResizeKeepsReadersClosedUntilEveryBlockIsReady) { + ConcurrentIndex map(ConcurrentArena()); + ASSERT_TRUE(map.PreallocateForExpectedSize(9 * 8192)); + const Digest keep{7168}; + auto* kept = map.InsertNew(keep, "external", 1, false); + auto* removed = map.InsertNew(Digest{5120}, "external", 2, false); + ASSERT_NE(kept, nullptr); + ASSERT_NE(removed, nullptr); + map.EnableConcurrentReads(); + { + auto reader = map.TryRead(keep); + ASSERT_TRUE(reader); + EXPECT_EQ(std::as_const(map).Find(keep, "external"), kept); + } + ASSERT_TRUE(map.Erase(removed)); + ASSERT_EQ(map.uninitialized_bucket_count(), 4096); + EXPECT_FALSE(map.TryRead(keep)); + for (std::size_t step = 0; step < 8192 && map.rehashing(); ++step) { + EXPECT_FALSE(map.TryRead(Digest{step})); + ASSERT_TRUE(map.Maintain()); + } + ASSERT_FALSE(map.rehashing()); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + std::thread remote([&] { + auto reader = map.TryRead(keep); + EXPECT_TRUE(reader); + if (!reader) return; + EXPECT_EQ(std::as_const(map).Find(keep, "external"), kept); + EXPECT_EQ(std::as_const(map).Find(Digest{1024}, "external"), nullptr); + }); + remote.join(); +} + TEST(IndexReadersTest, StripeClosureAndNestedGlobalClosureDrainExistingReaders) { ConcurrentIndex map(ConcurrentArena()); diff --git a/tests/scan_hash_map_test.cpp b/tests/scan_hash_map_test.cpp index ad92422ff..dd6661605 100644 --- a/tests/scan_hash_map_test.cpp +++ b/tests/scan_hash_map_test.cpp @@ -744,6 +744,258 @@ TEST(ScanHashMapTest, PreallocatesSettledTableForKnownPopulation) { EXPECT_EQ(map.size(), 91); } +TEST(ScanHashMapTest, LazyGrowthSkipsUntouchedBlocksAndPreservesEarlyWrites) { + using Map = ScanHashMap; + constexpr std::size_t kOldBuckets = 2048; + constexpr std::size_t kOldEntries = 9 * kOldBuckets; + Map map; + ASSERT_TRUE(map.PreallocateForExpectedSize(kOldEntries)); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + for (std::size_t i = 0; i < kOldEntries; ++i) { + ASSERT_NE(map.InsertNew(Digest{i}, "external", i, false), nullptr); + } + // This insertion lands in the last of four target blocks, long before + // sequential migration reaches that block. Only that block is constructed. + const Digest early{(1ULL << 32) | 3072}; + auto* entry = map.InsertNew(early, "external", kOldEntries, false); + ASSERT_NE(entry, nullptr); + ASSERT_TRUE(map.rehashing()); + EXPECT_EQ(map.uninitialized_bucket_count(), 3072); + const auto& reader = std::as_const(map); + const Digest missing{(1ULL << 40) | 1024}; + EXPECT_EQ(reader.Find(missing, "external"), nullptr); + EXPECT_EQ(reader.Find(early, "external"), entry); + const auto address = reinterpret_cast(entry); + EXPECT_EQ(reader.FindAddress(address, Map::AddressHash(missing)), nullptr); + EXPECT_EQ(reader.FindAddress(address, Map::AddressHash(missing), + Map::AddressTag(missing)), + nullptr); + reader.PrefetchAddress(Map::AddressHash(missing)); + EXPECT_EQ( + map.RequiredAllocationBytes(missing, "external", false, false, true), + map.RequiredAllocationBytes(Digest{3073}, "external", false, false, + true)); + + auto check_scan = [&] { + std::vector seen(map.size(), 0); + std::uint64_t cursor = 0; + do { + cursor = map.Scan(cursor, [&](auto& value) { ++seen.at(value.value()); }); + } while (cursor != 0); + for (unsigned count : seen) EXPECT_EQ(count, 1); + Map::StableScanCursor stable; + std::size_t visited = 0; + while (!map.ScanStableWhile(&stable, [&](const auto&) { + ++visited; + return visited % 37 != 0; + })) { + } + EXPECT_EQ(visited, map.size()); + visited = 0; + map.ForEach([&](const auto&) { ++visited; }); + EXPECT_EQ(visited, map.size()); + visited = 0; + EXPECT_TRUE(map.ForEachWhile([&](const auto&) { + ++visited; + return true; + })); + EXPECT_EQ(visited, map.size()); + }; + check_scan(); + EXPECT_EQ(map.uninitialized_bucket_count(), 3072); + + // These owner queries advance migration near bucket zero, initializing + // target blocks zero and two. The missing key's block one stays untouched. + EXPECT_EQ(map.Find(missing, "external"), nullptr); + EXPECT_TRUE(map.FindCandidates(missing, "external").empty()); + EXPECT_EQ(map.FindCandidateIf(missing, "external", + [](const auto&) { return true; }), + nullptr); + EXPECT_FALSE(map.Erase(missing, "external")); + EXPECT_EQ(map.uninitialized_bucket_count(), 1024); + check_scan(); + EXPECT_EQ(map.uninitialized_bucket_count(), 1024); + + // An early write into the last untouched block must not be erased when + // migration later visits another bucket in the same block. + auto* second = map.InsertNew(missing, "external", kOldEntries + 1, false); + ASSERT_NE(second, nullptr); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + // Keep an overflow chain in the target: finishing initialization must free + // the bitmap without freeing the shared overflow metadata or its entries. + for (std::uint64_t i = 0; i < 13; ++i) { + const Digest collision{((i + 2) << 40) | 1024}; + ASSERT_NE(map.InsertNew(collision, "external", kOldEntries + 2 + i, false), + nullptr); + } + for (std::size_t step = 0; step < kOldBuckets && map.rehashing(); ++step) { + ASSERT_TRUE(map.Maintain()); + } + ASSERT_FALSE(map.rehashing()); + EXPECT_EQ(map.allocated_bucket_count(), 2 * kOldBuckets); + EXPECT_EQ(reader.Find(early, "external"), entry); + EXPECT_EQ(reader.Find(missing, "external"), second); + for (std::uint64_t i = 0; i < 13; ++i) { + const Digest collision{((i + 2) << 40) | 1024}; + auto* found = reader.Find(collision, "external"); + ASSERT_NE(found, nullptr); + EXPECT_EQ(found->value(), kOldEntries + 2 + i); + } + for (std::size_t i = 0; i < kOldEntries; ++i) { + auto* found = reader.Find(Digest{i}, "external"); + ASSERT_NE(found, nullptr); + EXPECT_EQ(found->value(), i); + } + check_scan(); +} + +TEST(ScanHashMapTest, LazyShrinkInitializesEmptyDestinationsBeforeSettling) { + ScanHashMap map; + ASSERT_TRUE(map.PreallocateForExpectedSize(9 * 8192)); + const Digest keep{7168}; + auto* kept = map.InsertNew(keep, "external", 1, false); + auto* removed = map.InsertNew(Digest{5120}, "external", 2, false); + ASSERT_NE(kept, nullptr); + ASSERT_NE(removed, nullptr); + ASSERT_TRUE(map.Erase(removed)); + ASSERT_TRUE(map.rehashing()); + EXPECT_EQ(map.uninitialized_bucket_count(), 4096); + EXPECT_EQ(std::as_const(map).Find(keep, "external"), kept); + std::size_t seen = 0; + std::uint64_t cursor = 0; + do { + cursor = map.Scan(cursor, [&](const auto& entry) { + EXPECT_EQ(entry.value(), 1); + ++seen; + }); + } while (cursor != 0); + EXPECT_EQ(seen, 1); + EXPECT_EQ(map.uninitialized_bucket_count(), 4096); + for (std::size_t step = 0; step < 8192 && map.rehashing(); ++step) { + ASSERT_TRUE(map.Maintain()); + } + ASSERT_FALSE(map.rehashing()); + EXPECT_EQ(map.allocated_bucket_count(), 4096); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + EXPECT_EQ(std::as_const(map).Find(keep, "external"), kept); + // Even destinations of empty source buckets are now safe to inspect after + // the bitmap has been discarded. + for (std::uint64_t i = 0; i < 4096; ++i) { + EXPECT_EQ(std::as_const(map).Find(Digest{i}, "external"), nullptr); + } +} + +TEST(ScanHashMapTest, LazyResizeBitmapSpansWordsAndScanSurvivesMigration) { + ScanHashMap map; + constexpr std::size_t kOldBuckets = 1 << 18; + ASSERT_TRUE(map.PreallocateForExpectedSize(9 * kOldBuckets)); + ASSERT_NE(map.InsertNew(Digest{0}, "external", 0, false), nullptr); + auto* removed = map.InsertNew(Digest{1}, "external", 1, false); + ASSERT_NE(removed, nullptr); + ASSERT_TRUE(map.Erase(removed)); + ASSERT_EQ(map.uninitialized_bucket_count(), kOldBuckets / 2); + constexpr std::array blocks{0, 63, 64, 127}; + for (std::size_t i = 0; i < blocks.size(); ++i) { + const Digest digest{(1ULL << 32) | (blocks[i] * 1024)}; + ASSERT_NE(map.InsertNew(digest, "external", i + 1, false), nullptr); + EXPECT_EQ(map.uninitialized_bucket_count(), + kOldBuckets / 2 - (i + 1) * 1024); + } + std::array seen{}; + std::uint64_t cursor = 0; + do { + cursor = + map.Scan(cursor, [&](const auto& entry) { ++seen.at(entry.value()); }); + ASSERT_TRUE(map.Maintain()); + } while (cursor != 0); + for (unsigned count : seen) EXPECT_GE(count, 1); + for (std::size_t step = 0; step < kOldBuckets && map.rehashing(); ++step) { + ASSERT_TRUE(map.Maintain()); + } + ASSERT_FALSE(map.rehashing()); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + for (std::size_t i = 0; i < blocks.size(); ++i) { + const Digest digest{(1ULL << 32) | (blocks[i] * 1024)}; + auto* found = std::as_const(map).Find(digest, "external"); + ASSERT_NE(found, nullptr); + EXPECT_EQ(found->value(), i + 1); + } +} + +TEST(ScanHashMapTest, LazyResizeMoveDetachAndClearReleaseAllAccounting) { + ASSERT_TRUE(lavik::InitMemoryLimit(64 * 1024 * 1024, 1).ok()); + lavik::BindMemoryAccountingShard(0); + struct RestoreLimit { + ~RestoreLimit() { + (void)lavik::InitMemoryLimit(1024ULL * 1024 * 1024, 1); + lavik::BindMemoryAccountingShard(lavik::kMaxMemoryWorkers); + } + } restore; + const auto before = lavik::WorkerMemoryAccountingBytes(0); + for (bool initialize_one_block : {false, true}) { + { + ScanHashMap map; + ASSERT_TRUE(map.PreallocateForExpectedSize(9 * 8192)); + ASSERT_NE(map.InsertNew(Digest{7168}, "external", 1, false), nullptr); + auto* removed = map.InsertNew(Digest{5120}, "external", 2, false); + ASSERT_NE(removed, nullptr); + ASSERT_TRUE(map.Erase(removed)); + if (initialize_one_block) ASSERT_TRUE(map.Maintain()); + const std::size_t remaining = map.uninitialized_bucket_count(); + EXPECT_EQ(remaining, initialize_one_block ? 3072 : 4096); + auto detached = map.Detach(); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + ScanHashMap moved(std::move(detached)); + EXPECT_EQ(moved.uninitialized_bucket_count(), remaining); + EXPECT_NE(std::as_const(moved).Find(Digest{7168}, "external"), nullptr); + detached = std::move(moved); + EXPECT_EQ(detached.uninitialized_bucket_count(), remaining); + detached.Clear(); + EXPECT_EQ(detached.uninitialized_bucket_count(), 0); + EXPECT_FALSE(detached.has_allocated_storage()); + } + EXPECT_EQ(lavik::WorkerMemoryAccountingBytes(0), before); + } +} + +TEST(ScanHashMapTest, LazyShrinkRollsBackMetadataWhenBucketAdmissionFails) { + ASSERT_TRUE(lavik::InitMemoryLimit(4 * 1024 * 1024, 1).ok()); + lavik::BindMemoryAccountingShard(0); + struct RestoreLimit { + ~RestoreLimit() { + (void)lavik::InitMemoryLimit(1024ULL * 1024 * 1024, 1); + lavik::BindMemoryAccountingShard(lavik::kMaxMemoryWorkers); + } + } restore; + for (bool externally_admitted : {false, true}) { + auto arena = std::make_shared( + ScanHashMapEntryArena::kMaximumPageId, externally_admitted); + ScanHashMap map(arena); + ASSERT_TRUE(map.PreallocateForExpectedSize(9 * 8192)); + ASSERT_NE(map.InsertNew(Digest{7168}, "external", 1, false), nullptr); + auto* removed = map.InsertNew(Digest{5120}, "external", 2, false); + ASSERT_NE(removed, nullptr); + const auto memory = lavik::GetWorkerMemoryStats(0); + // Leave room for the optional metadata and bitmap, but not the target + // bucket array. Their temporary accounting must unwind on rejection. + auto blocker = lavik::TryReserveMemory(memory.retained_limit_bytes_ - + memory.retained_bytes_ - 512); + ASSERT_TRUE(blocker.has_value()); + EXPECT_TRUE(map.Erase(removed)); + EXPECT_EQ(map.size(), 1); + EXPECT_FALSE(map.rehashing()); + EXPECT_EQ(map.allocated_bucket_count(), 8192); + EXPECT_EQ(map.uninitialized_bucket_count(), 0); + EXPECT_EQ(lavik::WorkerMemoryAccountingBytes(0), memory.retained_bytes_); + EXPECT_FALSE(map.Maintain()); + EXPECT_EQ(lavik::WorkerMemoryAccountingBytes(0), memory.retained_bytes_); + blocker.reset(); + ASSERT_TRUE(map.Maintain()); + EXPECT_TRUE(map.rehashing()); + EXPECT_EQ(map.uninitialized_bucket_count(), 4096); + } +} + TEST(ScanHashMapTest, KnownEmptyPopulationAllocatesNothing) { ScanHashMap map; EXPECT_TRUE(map.PreallocateForExpectedSize(0));