diff --git a/include/lavik/storage/scan_hash_map.h b/include/lavik/storage/scan_hash_map.h index 35eac75f9..a93238152 100644 --- a/include/lavik/storage/scan_hash_map.h +++ b/include/lavik/storage/scan_hash_map.h @@ -49,14 +49,15 @@ namespace lavik::storage { // Single-owner backing store for compact hash-table entry handles. Small pages -// are carved from 1 MiB spans so mimalloc pays the 64 KiB alignment overhead -// once per 16 logical pages instead of once per page. Large entries still own -// dedicated 64 KiB-aligned allocations. Mimalloc remains responsible for OS -// memory, accounting, and purge behavior; only slots and span pages are managed -// here. A page belongs to one size class, and freed slots carry their own -// intrusive free-list link. Empty logical pages are reusable immediately, but -// their backing memory is returned only when all sixteen pages in the span are -// empty; this is the reclaim-granularity cost of amortizing alignment waste. +// are carved from 1 MiB spans. A smaller aligned request can supply the whole +// span from allocator slack; configurations that cannot expand it in place use +// a full-size request. Large entries still own dedicated 64 KiB-aligned +// allocations. Mimalloc remains responsible for OS memory, accounting, and +// purge behavior; only slots and span pages are managed here. A page belongs +// to one size class, and freed slots carry their own intrusive free-list link. +// Empty logical pages are reusable immediately, but their backing memory is +// returned only when all sixteen pages in the span are empty; this is the +// reclaim-granularity cost of sharing span storage. // // Handles are runtime-only. Zero is invalid; the high 21 bits select a page // directory entry and the low 11 bits select a slot in that page. Page IDs may @@ -77,9 +78,9 @@ class ScanHashMapEntryArena { static constexpr std::size_t kPageBytes = 64 * 1024; static constexpr std::size_t kPagesPerSpan = 16; static constexpr std::size_t kSpanBytes = kPagesPerSpan * kPageBytes; - // A 64 KiB-aligned allocation may consume one extra alignment unit in the - // allocator size class. External admission users must cover this whole - // first-span cost because the arena deliberately skips its own reservation. + // Admission must cover the full-size fallback and its alignment allowance, + // even when the smaller request normally supplies the entire span. + // Externally admitted arenas deliberately skip their own reservation. static constexpr std::size_t kSmallSpanAdmissionBytes = kSpanBytes + kPageBytes; static constexpr unsigned kSlotBits = 11; @@ -536,8 +537,22 @@ ScanHashMapEntryArena::AllocateSmallSpan() { auto span_owner = TryMakeRetainedUnique(allocation_domain_); if (span_owner == nullptr) return nullptr; auto* span = span_owner.get(); - span->storage_ = - TryAllocateRetainedBytes(allocation_domain_, kSpanBytes, kPageBytes); + // The pinned mimalloc's aligned path adds alignment - 1 before rounding. + // Requesting 960 KiB normally yields a full 1 MiB instead of 1088 KiB. + // mi_expand neither moves nor initializes storage; its result is the only + // capacity check needed. Debug padding disables expansion, and direct OS + // allocations can lose usable bytes to alignment, so neither may publish + // the sixteenth page without a successful expansion or a full-size request. + span->storage_ = TryAllocateRetainedBytes( + allocation_domain_, kSpanBytes - kPageBytes, kPageBytes); + if (span->storage_ == nullptr) return nullptr; + if (mi_expand(span->storage_, kSpanBytes) == nullptr) [[unlikely]] { + // Release the smaller allocation and its retained charge before admitting + // the fallback. A rejected fallback leaves no partially published span. + DeallocateRetainedBytes(allocation_domain_, span->storage_, kPageBytes); + span->storage_ = + TryAllocateRetainedBytes(allocation_domain_, kSpanBytes, kPageBytes); + } if (span->storage_ == nullptr) return nullptr; span->all_previous_ = nullptr; span->all_next_ = small_spans_head_; diff --git a/tests/scan_hash_map_test.cpp b/tests/scan_hash_map_test.cpp index ad92422ff..8c8662f2f 100644 --- a/tests/scan_hash_map_test.cpp +++ b/tests/scan_hash_map_test.cpp @@ -19,7 +19,9 @@ #include #include +#include #include +#include #include #include #include @@ -789,29 +791,170 @@ TEST(ScanHashMapTest, AccountsForOverflowPoolVectorGrowth) { 8ULL * 1024); } +TEST(ScanHashMapTest, SmallSpanUsesCompactAllocationWhenExpandable) { + constexpr std::size_t kSpanBytes = ScanHashMapEntryArena::kSpanBytes; + constexpr std::size_t kPageBytes = ScanHashMapEntryArena::kPageBytes; + using Allocation = std::unique_ptr; + Allocation compact(mi_malloc_aligned(kSpanBytes - kPageBytes, kPageBytes), + &mi_free); + Allocation full(mi_malloc_aligned(kSpanBytes, kPageBytes), &mi_free); + ASSERT_NE(compact, nullptr); + ASSERT_NE(full, nullptr); + // Compare arena retention with the allocator's actual capabilities. This + // exercises expansion with release mimalloc and fallback with Debug padding + // or MIMALLOC_DISALLOW_ARENA_ALLOC=1, without assuming a portable size class. + const bool can_expand = mi_expand(compact.get(), kSpanBytes) != nullptr; + const std::size_t expected_bytes = + mi_usable_size(can_expand ? compact.get() : full.get()); + ASSERT_GE(expected_bytes, kSpanBytes); + + 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); + { + ScanHashMapEntryArena arena(16); + ASSERT_TRUE(arena.PreallocateDirectory()); + const auto directory_bytes = lavik::WorkerMemoryAccountingBytes(0); + const auto allocation = arena.Allocate(32); + ASSERT_NE(allocation.pointer_, nullptr); + // The first page starts at the allocation base; later logical pages must + // never be passed to allocator introspection or free. + void* span = reinterpret_cast( + reinterpret_cast(allocation.pointer_) & + ~(kPageBytes - 1)); + EXPECT_EQ(mi_usable_size(span), expected_bytes); + const auto retained = + lavik::WorkerMemoryAccountingBytes(0) - directory_bytes; + EXPECT_GE(retained, static_cast(expected_bytes)); + EXPECT_LT(retained, static_cast(expected_bytes + kPageBytes)); + arena.Deallocate(allocation.handle_); + EXPECT_EQ(arena.allocated_spans(), 0); + EXPECT_EQ(lavik::WorkerMemoryAccountingBytes(0), directory_bytes); + } + EXPECT_EQ(lavik::WorkerMemoryAccountingBytes(0), before); +} + +TEST(ScanHashMapTest, SmallSpanRejectionReleasesTemporaryStorageAndPageId) { + constexpr std::size_t kSpanBytes = ScanHashMapEntryArena::kSpanBytes; + constexpr std::size_t kPageBytes = ScanHashMapEntryArena::kPageBytes; + std::unique_ptr probe( + mi_malloc_aligned(kSpanBytes - kPageBytes, kPageBytes), &mi_free); + ASSERT_NE(probe, nullptr); + const bool can_expand = mi_expand(probe.get(), kSpanBytes) != nullptr; + + 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; + ScanHashMapEntryArena arena(1); + ASSERT_TRUE(arena.PreallocateDirectory()); + const auto memory = lavik::GetWorkerMemoryStats(0); + constexpr std::size_t kMetadataHeadroom = 4096; + constexpr auto kMaximumBytes = std::numeric_limits::max(); + const auto compact_usable = + lavik::AllocatorUsableSizeForRequest(kSpanBytes - kPageBytes); + const auto full_usable = lavik::AllocatorUsableSizeForRequest(kSpanBytes); + ASSERT_LE(compact_usable, kMaximumBytes - kPageBytes - kMetadataHeadroom); + ASSERT_LE(full_usable, kMaximumBytes - kPageBytes); + const auto compact_admission = compact_usable + kPageBytes; + const auto full_admission = full_usable + kPageBytes; + const auto headroom = compact_admission + kMetadataHeadroom; + ASSERT_GT(full_admission, headroom); + ASSERT_GE(memory.retained_limit_bytes_, memory.retained_bytes_); + const auto available = memory.retained_limit_bytes_ - memory.retained_bytes_; + ASSERT_GT(available, headroom); + // Mirror allocator admission, including its alignment allowance. The + // compact request and span metadata fit, while the full-size fallback must + // be rejected. Neither a temporary charge nor a page ID may survive that + // rejection; restoring headroom must allow a retry in the same arena. + auto blocker = lavik::TryReserveMemory(available - headroom); + ASSERT_TRUE(blocker.has_value()); + const auto allocation = arena.Allocate(32); + if (can_expand) { + ASSERT_NE(allocation.pointer_, nullptr); + arena.Deallocate(allocation.handle_); + } else { + EXPECT_EQ(allocation.pointer_, nullptr); + EXPECT_EQ(allocation.handle_, 0); + } + EXPECT_EQ(arena.allocated_pages(), 0); + EXPECT_EQ(arena.allocated_spans(), 0); + EXPECT_EQ(lavik::GetWorkerMemoryStats(0).retained_bytes_, + memory.retained_bytes_); + EXPECT_EQ(lavik::GetWorkerMemoryStats(0).admission_pending_bytes_, + blocker->bytes()); + blocker.reset(); + + const auto retry = arena.Allocate(32); + ASSERT_NE(retry.pointer_, nullptr); + EXPECT_EQ(arena.Resolve(retry.handle_), retry.pointer_); + arena.Deallocate(retry.handle_); +} + TEST(ScanHashMapTest, AmortizesAlignmentAcrossSixteenLogicalPages) { ScanHashMapEntryArena arena; std::vector handles; + // This class fills every byte after the page header, including the final + // byte of the sixteenth page. A smaller request without successful expansion + // would overwrite Debug padding or run past a direct OS allocation. + constexpr std::size_t kSlotBytes = 32; constexpr std::size_t kSlotsPerPage = - (ScanHashMapEntryArena::kPageBytes - 64) / 4096; + (ScanHashMapEntryArena::kPageBytes - 64) / kSlotBytes; handles.reserve(ScanHashMapEntryArena::kPagesPerSpan * kSlotsPerPage); for (std::size_t page = 0; page < ScanHashMapEntryArena::kPagesPerSpan; ++page) { for (std::size_t slot = 0; slot < kSlotsPerPage; ++slot) { - handles.push_back(arena.Allocate(4096).handle_); + const auto allocation = arena.Allocate(kSlotBytes); + ASSERT_NE(allocation.pointer_, nullptr); + ASSERT_EQ(arena.HandleOf(allocation.pointer_), allocation.handle_); + std::memset(allocation.pointer_, static_cast(page + 1), kSlotBytes); + handles.push_back(allocation.handle_); } EXPECT_EQ(arena.allocated_spans(), 1); if (page + 1 < ScanHashMapEntryArena::kPagesPerSpan) { - EXPECT_LT(arena.AllocationBytesIfNewPage(4096), + EXPECT_LT(arena.AllocationBytesIfNewPage(kSlotBytes), ScanHashMapEntryArena::kSpanBytes); } } - EXPECT_GE(arena.AllocationBytesIfNewPage(4096), + EXPECT_GE(arena.AllocationBytesIfNewPage(kSlotBytes), ScanHashMapEntryArena::kSpanBytes); EXPECT_EQ(arena.allocated_pages(), ScanHashMapEntryArena::kPagesPerSpan); EXPECT_EQ(arena.allocated_spans(), 1); + for (std::size_t index = 0; index < handles.size(); ++index) { + const auto* data = + static_cast(arena.Resolve(handles[index])); + for (std::size_t byte = 0; byte < kSlotBytes; ++byte) { + ASSERT_EQ(data[byte], index / kSlotsPerPage + 1); + } + } + + // Recycle the final page while the other pages keep the span alive. + for (std::size_t slot = 0; slot < kSlotsPerPage; ++slot) { + arena.Deallocate(handles.back()); + handles.pop_back(); + } + EXPECT_EQ(arena.allocated_pages(), ScanHashMapEntryArena::kPagesPerSpan - 1); + EXPECT_EQ(arena.allocated_spans(), 1); + for (std::size_t slot = 0; slot < kSlotsPerPage; ++slot) { + const auto allocation = arena.Allocate(kSlotBytes); + ASSERT_NE(allocation.pointer_, nullptr); + std::memset(allocation.pointer_, 0xa5, kSlotBytes); + handles.push_back(allocation.handle_); + } + EXPECT_EQ(arena.allocated_spans(), 1); + for (auto handle = handles.rbegin(); handle != handles.rend(); ++handle) { arena.Deallocate(*handle); }