decommit unusable page(s) in medium heap blocks (bug 5822302)

in medium heap blocks, if the object size is bigger than 1 page, it can cause the whole last 1~3 page(s) never been allocated. when such case is hit, decommit those pages to save memory as well as capture corruption. when returning the pages back to page allocator, we should commit those pages again for reuse, in case of OOM here, just decommit all pages in the heap block and let the page allocator to manage the decommitted pages.
in rescan code, asserting that we never scan those unallocatable pages

Author: leirocks

lib/Common/Core/FaultInjection.h

@@ -156,6 +156,10 @@ namespace Js
     if(Js::FaultInjection::Global.ShouldInjectFault(Js::FaultInjection::Global.NoThrow, name, size)) \
         return NULL;
 
+#define FAULTINJECT_MEMORY_NOTHROW_RET(name, size, ret) \
+    if(Js::FaultInjection::Global.ShouldInjectFault(Js::FaultInjection::Global.NoThrow, name, size)) \
+        return ret;
+
 #define FAULTINJECT_MEMORY_THROW(name, size) \
     if(Js::FaultInjection::Global.ShouldInjectFault(Js::FaultInjection::Global.Throw, name, size)) \
         Js::Throw::OutOfMemory();
@@ -200,6 +204,7 @@ namespace Js
 #define IS_FAULTINJECT_NO_THROW_ON false
 
 #define FAULTINJECT_MEMORY_NOTHROW(name, size)
+#define FAULTINJECT_MEMORY_NOTHROW_RET(name, size, ret)
 #define FAULTINJECT_MEMORY_THROW(name, size)
 #define FAULTINJECT_MEMORY_MARK_THROW(name, size)
 #define FAULTINJECT_MEMORY_MARK_NOTHROW(name, size)

lib/Common/Memory/HeapBlock.cpp

@@ -313,6 +313,8 @@ SmallHeapBlockT<TBlockAttributes>::SetPage(__in_ecount_pagesize char * baseAddre
     MemoryBarrier();
 #endif
 
+    this->DecommitUnusablePages();
+
     return TRUE;
 }
 
@@ -334,14 +336,21 @@ SmallHeapBlockT<TBlockAttributes>::ReleasePages(Recycler * recycler)
     char* address = this->address;
 
 #ifdef RECYCLER_FREE_MEM_FILL
-    memset(address, DbgMemFill, AutoSystemInfo::PageSize * this->GetPageCount());
+    memset(address, DbgMemFill, AutoSystemInfo::PageSize * (this->GetPageCount()-this->GetUnusablePageCount()));
 #endif
 
-    this->GetPageAllocator(recycler)->ReleasePages(address, this->GetPageCount(), this->GetPageSegment());
+    if (!this->RecommitUnusablePages())
+    {
+        this->GetPageAllocator(recycler)->PartialDecommitPages(address, this->GetPageCount(), address,
+            this->GetPageCount() - this->GetUnusablePageCount(), this->segment);
+    }
+    else
+    {
+        this->GetPageAllocator(recycler)->ReleasePages(address, this->GetPageCount(), this->GetPageSegment());
+    }    
 
     this->segment = nullptr;
     this->address = nullptr;
-
 }
 
 template <class TBlockAttributes>
@@ -351,7 +360,15 @@ SmallHeapBlockT<TBlockAttributes>::BackgroundReleasePagesSweep(Recycler* recycle
     recycler->heapBlockMap.ClearHeapBlock(address, this->GetPageCount());
     char* address = this->address;
 
-    this->GetPageAllocator(recycler)->BackgroundReleasePages(address, this->GetPageCount(), this->GetPageSegment());
+    if (!this->RecommitUnusablePages())
+    {
+        this->GetPageAllocator(recycler)->PartialDecommitPages(address, this->GetPageCount(), address, 
+            this->GetPageCount() - this->GetUnusablePageCount(), this->segment);
+    }
+    else
+    {
+        this->GetPageAllocator(recycler)->BackgroundReleasePages(address, this->GetPageCount(), this->GetPageSegment());
+    }
 
     this->address = nullptr;
     this->segment = nullptr;

lib/Common/Memory/HeapBlock.h

@@ -212,6 +212,7 @@ template <class TBlockAttributes> class SmallFinalizableWithBarrierHeapBlockT;
 class RecyclerHeapObjectInfo;
 class HeapBlock
 {
+    friend MediumAllocationBlockAttributes;
 public:
     enum HeapBlockType : byte
     {
@@ -444,6 +445,20 @@ class SmallHeapBlockT : public HeapBlock
 public:
     ~SmallHeapBlockT();
 
+    void DecommitUnusablePages() 
+    {
+        TBlockAttributes::DecommitUnusablePages(this);
+    }
+
+    BOOL RecommitUnusablePages()
+    {
+        return TBlockAttributes::RecommitUnusablePages(this);
+    }
+
+    uint GetUnusablePageCount()
+    {
+        return TBlockAttributes::GetUnusablePageCount(this->objectSize);
+    }
 
 #ifdef RECYCLER_WRITE_BARRIER
     bool IsWithBarrier() const;

lib/Common/Memory/HeapBlockMap.cpp

@@ -1067,7 +1067,8 @@ HeapBlockMap32::RescanHeapBlockOnOOM(TBlockType* heapBlock, char* pageAddress, H
     // The following assert makes sure that this method is called only once per heap block
     Assert(blockStartAddress == pageAddress);
 
-    for (int i = 0; i < TBlockType::HeapBlockAttributes::PageCount; i++)
+    int inUsePageCount = heapBlock->GetPageCount() - heapBlock->GetUnusablePageCount();
+    for (int i = 0; i < inUsePageCount; i++)
     {
         char* pageAddressToScan = blockStartAddress + (i * AutoSystemInfo::PageSize);
 

lib/Common/Memory/HeapConstants.h

@@ -26,6 +26,9 @@ class HeapConstants
     static const uint MediumBucketCount = (MaxMediumObjectSize - MaxSmallObjectSize) / MediumObjectGranularity;
 #endif
 };
+namespace Memory {
+    class HeapBlock;
+}
 
 ///
 /// BlockAttributes are used to determine the allocation characteristics of a heap block
@@ -61,6 +64,9 @@ class SmallAllocationBlockAttributes
     static const bool IsLargeBlock = false;
 
     static BOOL IsAlignedObjectSize(size_t sizeCat);
+    static uint GetUnusablePageCount(size_t sizeCat);
+    static void DecommitUnusablePages(Memory::HeapBlock* heapBlock);
+    static BOOL RecommitUnusablePages(Memory::HeapBlock* heapBlock);
 };
 
 class MediumAllocationBlockAttributes
@@ -83,6 +89,9 @@ class MediumAllocationBlockAttributes
     static const bool IsLargeBlock = false;
 
     static BOOL IsAlignedObjectSize(size_t sizeCat);
+    static uint GetUnusablePageCount(size_t sizeCat);
+    static void DecommitUnusablePages(Memory::HeapBlock* heapBlock);
+    static BOOL RecommitUnusablePages(Memory::HeapBlock* heapBlock);
 };
 
 class LargeAllocationBlockAttributes

lib/Common/Memory/HeapInfo.cpp

@@ -505,10 +505,11 @@ HeapInfo::Initialize(Recycler * recycler
     if (pageheapmode == PageHeapMode::PageHeapModeOff)
     {
 #ifdef ENABLE_DEBUG_CONFIG_OPTIONS
-        isPageHeapEnabled = recycler->GetRecyclerFlagsTable().PageHeap != PageHeapMode::PageHeapModeOff;
-        pageheapmode = (PageHeapMode)recycler->GetRecyclerFlagsTable().PageHeap;
-        blockTypeFilter = (PageHeapBlockTypeFilter)recycler->GetRecyclerFlagsTable().PageHeapBlockType;
-        pBucketNumberRange = &recycler->GetRecyclerFlagsTable().PageHeapBucketNumber;
+        auto& flags = recycler->GetRecyclerFlagsTable();
+        isPageHeapEnabled = flags.PageHeap != PageHeapMode::PageHeapModeOff;
+        pageheapmode = (PageHeapMode)flags.PageHeap;
+        blockTypeFilter = (PageHeapBlockTypeFilter)flags.PageHeapBlockType;
+        pBucketNumberRange = &flags.PageHeapBucketNumber;
 
 #else
         // @TODO in free build, use environment var or other way to enable page heap
@@ -1767,12 +1768,60 @@ BOOL SmallAllocationBlockAttributes::IsAlignedObjectSize(size_t sizeCat)
 {
     return HeapInfo::IsAlignedSmallObjectSize(sizeCat);
 }
+/* static */
+uint SmallAllocationBlockAttributes::GetUnusablePageCount(size_t sizeCat)
+{
+    UNREFERENCED_PARAMETER(sizeCat);
+    return 0;
+}
+/* static */
+void SmallAllocationBlockAttributes::DecommitUnusablePages(HeapBlock* heapBlock)
+{
+    UNREFERENCED_PARAMETER(heapBlock);
+}
+/* static */
+BOOL SmallAllocationBlockAttributes::RecommitUnusablePages(HeapBlock* heapBlock)
+{
+    UNREFERENCED_PARAMETER(heapBlock);
+    return TRUE;
+}
 
 /* static */
 BOOL MediumAllocationBlockAttributes::IsAlignedObjectSize(size_t sizeCat)
 {
     return HeapInfo::IsAlignedMediumObjectSize(sizeCat);
 }
+/* static */
+uint MediumAllocationBlockAttributes::GetUnusablePageCount(size_t sizeCat)
+{
+    return ((MediumAllocationBlockAttributes::PageCount*AutoSystemInfo::PageSize) % sizeCat) / AutoSystemInfo::PageSize;
+}
+/* static */
+void MediumAllocationBlockAttributes::DecommitUnusablePages(HeapBlock* heapBlock)
+{
+    size_t count = MediumAllocationBlockAttributes::GetUnusablePageCount(heapBlock->GetObjectSize(nullptr));
+    if (count > 0)
+    {
+        char* startPage = (char*)heapBlock->address + (MediumAllocationBlockAttributes::PageCount - count)*AutoSystemInfo::PageSize;
+#pragma warning(suppress: 6250)
+        ::VirtualFree(startPage, count*AutoSystemInfo::PageSize, MEM_DECOMMIT);
+        ::ResetWriteWatch(startPage, count*AutoSystemInfo::PageSize);
+    }
+}
+/* static */
+BOOL MediumAllocationBlockAttributes::RecommitUnusablePages(HeapBlock* heapBlock)
+{
+    size_t count = MediumAllocationBlockAttributes::GetUnusablePageCount(heapBlock->GetObjectSize(nullptr));
+    if (count > 0)
+    {
+        FAULTINJECT_MEMORY_NOTHROW_RET(_u("recommit unusable pages"), count*AutoSystemInfo::PageSize, FALSE);
+
+        char* startPage = (char*)heapBlock->address + (MediumAllocationBlockAttributes::PageCount - count)*AutoSystemInfo::PageSize;
+#pragma warning(suppress: 6250)
+        return startPage==::VirtualAlloc(startPage, count*AutoSystemInfo::PageSize, MEM_COMMIT, PAGE_READWRITE);
+    }
+    return TRUE;
+}
 
 template class HeapInfo::ValidPointersMap<SmallAllocationBlockAttributes>;
 template class ValidPointers<SmallAllocationBlockAttributes>;

lib/Common/Memory/SmallNormalHeapBucket.cpp

@@ -6,7 +6,7 @@
 
 
 template <typename TBlockType>
-SmallNormalHeapBucketBase<TBlockType>::SmallNormalHeapBucketBase() 
+SmallNormalHeapBucketBase<TBlockType>::SmallNormalHeapBucketBase()
 #if ENABLE_PARTIAL_GC
     : partialHeapBlockList(nullptr)
 #if ENABLE_CONCURRENT_GC
@@ -96,6 +96,7 @@ SmallNormalHeapBucketBase<TBlockType>::RescanObjectsOnPage(TBlockType * block, c
     const uint pageByteOffset = static_cast<uint>((char*)pageAddress - blockStartAddress);
     uint firstObjectOnPageIndex = pageByteOffset / localObjectSize;
 
+
     // This is not necessarily the address on the first object that starts on the page
     // If the last object on the previous page spans two pages, this is the address of that object
     // We do it this way so that we can figure out if we need to rescan the first few bytes of the page
@@ -114,6 +115,12 @@ SmallNormalHeapBucketBase<TBlockType>::RescanObjectsOnPage(TBlockType * block, c
     const uint pageObjectCount = blockInfoForPage.pageObjectCount;
     const uint localObjectCount = (TBlockAttributes::PageCount * AutoSystemInfo::PageSize) / localObjectSize;
 
+    // With decommitting unallocatable ending pages and reset writewatch, we should never be scanning on these pages.
+    if (firstObjectOnPageIndex >= localObjectCount)
+    {
+        ReportFatalException(NULL, E_FAIL, Fatal_Recycler_MemoryCorruption, 3);
+    }
+
     // If all objects are marked, rescan whole block at once
     if (TBlockType::CanRescanFullBlock() && rescanMarkCount == pageObjectCount)
     {
@@ -588,7 +595,7 @@ template class SmallNormalHeapBucketBase<MediumFinalizableHeapBlock>;
 template class SmallNormalHeapBucketBase<SmallFinalizableWithBarrierHeapBlock>;
 template class SmallNormalHeapBucketBase<MediumFinalizableWithBarrierHeapBlock>;
 #endif
-   
+
 template void SmallNormalHeapBucketBase<SmallNormalHeapBlock>::Sweep(RecyclerSweep& recyclerSweep);
 template void SmallNormalHeapBucketBase<MediumNormalHeapBlock>::Sweep(RecyclerSweep& recyclerSweep);