MallocStomp2.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
#include "CoreTypes.h"
#include "HAL/MemoryBase.h"
#include "HAL/PlatformTLS.h"
 
#ifndef WITH_MALLOC_STOMP2
#define WITH_MALLOC_STOMP2 WITH_EDITOR && PLATFORM_WINDOWS
#endif
#if WITH_MALLOC_STOMP2
#include <atomic>
#if PLATFORM_MICROSOFT
    #include "Microsoft/WindowsHWrapper.h"
#endif
 
class FMallocStomp2 final : public FMalloc
{
private:
    UE_FORCEINLINE_HINT SIZE_T GetAlignedSize(SIZE_T InSize, SIZE_T InAlignment)
    {
        return (((InSize)+(InAlignment - 1)) & ~(InAlignment - 1));
    }
 
 
#if PLATFORM_64BITS
    static constexpr SIZE_T SentinelExpectedValue = 0xdeadbeefdeadbeef;
#else
    static constexpr SIZE_T SentinelExpectedValue = 0xdeadbeef;
#endif
 
    const SIZE_T PageSize;
 
    struct FAllocationData
    {
        void    *FullAllocationPointer;
        SIZE_T  FullSize;
        SIZE_T  Size;
        SIZE_T  Sentinel;
    };
 
public:
    struct FAddressSpaceStompDataBlockRangeEntry
    {
        UPTRINT StartAddress;
        UPTRINT EndAddress;
        int     StartIndex;
        int     EndIndex;
    };
 
 
 
    struct FAddressSpaceStompDataBlock
    {
        static constexpr SIZE_T VirtualAddressSpaceToReserve = 1024LL * 1024LL * 1024LL * 64LL; // 64 GB blocks, we try to reserve that much, but will half it and retry
        std::atomic<UPTRINT>    CurrentAddress;
        UPTRINT StartAddress;
        UPTRINT EndAddress;
        SIZE_T  Size;
 
        bool Init()
        {
            SIZE_T SizeToTryAndReserve = VirtualAddressSpaceToReserve;
            // if we can't reserve the full size, half the size and try again as long as we are over 1GB
            while (SizeToTryAndReserve >= 1024LL * 1024LL * 1024LL)
            {
                StartAddress = (UPTRINT) ::VirtualAlloc(nullptr, SizeToTryAndReserve, MEM_RESERVE, PAGE_READWRITE);
                if (StartAddress)
                {
                    EndAddress = StartAddress + SizeToTryAndReserve;
                    Size = SizeToTryAndReserve;
                    CurrentAddress = StartAddress;
                    return true;
                }
                SizeToTryAndReserve /= 2;
            }
            return false;
        }
 
        void* Malloc(SIZE_T InSizeToAllocate, SIZE_T InPageSize, bool InUnderrun = false)
        {
            UPTRINT Ret = CurrentAddress.fetch_add(InSizeToAllocate);
            if (Ret + InSizeToAllocate > EndAddress)
            {
                return nullptr;
            }
            if (!InUnderrun)
            {
                // leave the last page without backing store
                Ret = (UPTRINT) ::VirtualAlloc((void*)Ret, InSizeToAllocate - InPageSize, MEM_COMMIT, PAGE_READWRITE);
            }
            else
            {
                // leave the first page without backing store
                Ret = (UPTRINT) ::VirtualAlloc((void*)(Ret+InPageSize), InSizeToAllocate, MEM_COMMIT, PAGE_READWRITE);
            }
            return (void*)Ret;
        }
        
        void Free(void* InPtr, SIZE_T InSize)
        {
#if PLATFORM_WINDOWS && USING_CODE_ANALYSIS
    MSVC_PRAGMA(warning(push))
    MSVC_PRAGMA(warning(disable : 6250)) // Suppress C6250, as virtual address space is "leaked" by design.
#endif
            // Unmap physical memory, but keep virtual address range reserved to catch use-after-free errors.
            ::VirtualFree(InPtr, InSize, MEM_DECOMMIT);
 
#if PLATFORM_WINDOWS && USING_CODE_ANALYSIS
    MSVC_PRAGMA(warning(pop))
#endif
        }
    };
 
 
private:
    static constexpr SIZE_T TargetAddressSpaceToReserve = 1024LL * 1024LL * 1024LL * 1024LL * 32LL; // how much address space we will try and reserve, 32 TB
 
 
    static constexpr SIZE_T MaxAddressSpaceStompDataBlocks = 4096; // this * VirtualAddressSpaceToReserve is how much address space we will try and allocate
 
    std::atomic<int64> CurrentAddressSpaceStompDataBlockIndex;
    FAddressSpaceStompDataBlock     AddressSpaceStompDataBlocks[MaxAddressSpaceStompDataBlocks];
 
    int NumberOfRangeEntries;
    FAddressSpaceStompDataBlockRangeEntry AddressSpaceStompDataBlocksRangeEntries[MaxAddressSpaceStompDataBlocks];
 
    FMalloc* UsedMalloc;
 
    int IsPartOf(void* InPtr);
    int FindAddressSpaceStompDataBlockIndex(void* InPtr, int Index);
 
 
    const bool bUseUnderrunMode;
 
 
public:
    // FMalloc interface.
    FMallocStomp2(FMalloc* InMalloc, const bool InUseUnderrunMode = false);
 
    virtual void* Malloc(SIZE_T Size, uint32 Alignment) override;
 
    virtual void* TryMalloc(SIZE_T Size, uint32 Alignment) override;
 
    virtual void* Realloc(void* InPtr, SIZE_T NewSize, uint32 Alignment) override;
 
    virtual void* TryRealloc(void* InPtr, SIZE_T NewSize, uint32 Alignment) override;
 
    virtual void Free(void* InPtr) override;
 
    virtual bool GetAllocationSize(void *Original, SIZE_T &SizeOut) override;
 
    virtual void DumpAllocatorStats( FOutputDevice& Ar ) override
    {
        // No meaningful stats to dump.
    }
 
    virtual bool ValidateHeap() override
    {
        // Nothing to do here since validation happens as data is accessed
        // through page protection, and on each free checking the sentinel.
        return true;
    }
 
    virtual const TCHAR* GetDescriptiveName() override
    {
        return TEXT( "Stomp2" );
    }
 
    virtual bool IsInternallyThreadSafe() const override
    {
        // Stomp allocator 2 is thread-safe
        return true;
    }
 
    static FMalloc* OverrideIfEnabled(FMalloc* InUsedAlloc);
 
    void Init();
 
 
    static uint32 DisableTlsSlot;
 
 
    static SIZE_T MinSize;
    static SIZE_T MaxSize;
};
 
 
class FScopeDisableMallocStomp2
{
public:
 
    FScopeDisableMallocStomp2()
    {
        uint64 DisableCounter = (uint64)FPlatformTLS::GetTlsValue(FMallocStomp2::DisableTlsSlot);
        ++DisableCounter;
        FPlatformTLS::SetTlsValue(FMallocStomp2::DisableTlsSlot, (void*)DisableCounter);
    }
 
    ~FScopeDisableMallocStomp2()
    {
        uint64 DisableCounter = (uint64)FPlatformTLS::GetTlsValue(FMallocStomp2::DisableTlsSlot);
        --DisableCounter;
        FPlatformTLS::SetTlsValue(FMallocStomp2::DisableTlsSlot, (void*)DisableCounter);
    }
};
 
extern CORE_API FMallocStomp2* GMallocStomp2;
extern CORE_API bool GMallocStomp2Enabled;
 
#endif  //WITH_MALLOC_STOMP2