LockFreeFixedSizeAllocator.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "AutoRTFM.h"
#include "Misc/AssertionMacros.h"
#include "HAL/UnrealMemory.h"
#include "Misc/NoopCounter.h"
#include "Containers/LockFreeList.h"
 
#ifndef USE_NAIVE_TLockFreeFixedSizeAllocator_TLSCacheBase
#define USE_NAIVE_TLockFreeFixedSizeAllocator_TLSCacheBase (0) // this is useful for find who really leaked
#endif
 
template<int32 SIZE, typename TBundleRecycler, typename TTrackingCounter = FNoopCounter>
class TLockFreeFixedSizeAllocator_TLSCacheBase : public FNoncopyable
{
    enum
    {
        SIZE_PER_BUNDLE = 65536,
        NUM_PER_BUNDLE = SIZE_PER_BUNDLE / SIZE
    };
 
public:
 
    TLockFreeFixedSizeAllocator_TLSCacheBase()
    {
        static_assert(SIZE >= sizeof(void*) && SIZE % sizeof(void*) == 0, "Blocks in TLockFreeFixedSizeAllocator must be at least the size of a pointer.");
        TlsSlot = FPlatformTLS::AllocTlsSlot();
        check(FPlatformTLS::IsValidTlsSlot(TlsSlot));
    }
    ~TLockFreeFixedSizeAllocator_TLSCacheBase()
    {
        FPlatformTLS::FreeTlsSlot(TlsSlot);
        TlsSlot = FPlatformTLS::InvalidTlsSlot;
    }
 
    inline void* Allocate()
    {
#if USE_NAIVE_TLockFreeFixedSizeAllocator_TLSCacheBase
        return FMemory::Malloc(SIZE);
#else
        FThreadLocalCache& TLS = GetTLS();
 
        if (!TLS.PartialBundle)
        {
            if (TLS.FullBundle)
            {
                TLS.PartialBundle = TLS.FullBundle;
                TLS.FullBundle = nullptr;
            }
            else
            {
                TLS.PartialBundle = GlobalFreeListBundles.Pop();
                if (!TLS.PartialBundle)
                {
                    TLS.PartialBundle = (void**)FMemory::Malloc(SIZE_PER_BUNDLE);
                    void **Next = TLS.PartialBundle;
                    for (int32 Index = 0; Index < NUM_PER_BUNDLE - 1; Index++)
                    {
                        void* NextNext = (void*)(((uint8*)Next) + SIZE);
                        *Next = NextNext;
                        Next = (void**)NextNext;
                    }
                    *Next = nullptr;
                    NumFree.Add(NUM_PER_BUNDLE);
                }
            }
            TLS.NumPartial = NUM_PER_BUNDLE;
        }
        NumUsed.Increment();
        NumFree.Decrement();
        void* Result = (void*)TLS.PartialBundle;
        TLS.PartialBundle = (void**)*TLS.PartialBundle;
        TLS.NumPartial--;
        check(TLS.NumPartial >= 0 && ((!!TLS.NumPartial) == (!!TLS.PartialBundle)));
        return Result;
#endif
    }
 
    inline void Free(void *Item)
    {
#if USE_NAIVE_TLockFreeFixedSizeAllocator_TLSCacheBase
        return FMemory::Free(Item);
#else
        NumUsed.Decrement();
        NumFree.Increment();
        FThreadLocalCache& TLS = GetTLS();
        if (TLS.NumPartial >= NUM_PER_BUNDLE)
        {
            if (TLS.FullBundle)
            {
                GlobalFreeListBundles.Push(TLS.FullBundle);
                //TLS.FullBundle = nullptr;
            }
            TLS.FullBundle = TLS.PartialBundle;
            TLS.PartialBundle = nullptr;
            TLS.NumPartial = 0;
        }
        *(void**)Item = (void*)TLS.PartialBundle;
        TLS.PartialBundle = (void**)Item;
        TLS.NumPartial++;
#endif
    }
 
    const TTrackingCounter& GetNumUsed() const
    {
        return NumUsed;
    }
 
    const TTrackingCounter& GetNumFree() const
    {
        return NumFree;
    }
 
private:
 
    struct FThreadLocalCache
    {
        void **FullBundle;
        void **PartialBundle;
        int32 NumPartial;
 
        FThreadLocalCache()
            : FullBundle(nullptr)
            , PartialBundle(nullptr)
            , NumPartial(0)
        {
        }
    };
 
    FThreadLocalCache& GetTLS()
    {
        checkSlow(FPlatformTLS::IsValidTlsSlot(TlsSlot));
        FThreadLocalCache* TLS = (FThreadLocalCache*)FPlatformTLS::GetTlsValue(TlsSlot);
        if (!TLS)
        {
            TLS = new FThreadLocalCache();
            FPlatformTLS::SetTlsValue(TlsSlot, TLS);
        }
        return *TLS;
    }
 
    uint32 TlsSlot;
 
    TBundleRecycler GlobalFreeListBundles;
 
    TTrackingCounter NumUsed;
 
    TTrackingCounter NumFree;
};
 
template<int32 SIZE, int TPaddingForCacheContention, typename TTrackingCounter = FNoopCounter>
class TLockFreeFixedSizeAllocator
{
public:
    TLockFreeFixedSizeAllocator()
    {
        checkf(!AutoRTFM::IsClosed() || !AutoRTFM::IsOnCurrentTransactionStack(this), TEXT("Not allowed to construct a stack local within a transaction."));
    }
 
    ~TLockFreeFixedSizeAllocator()
    {
        AutoRTFM::PopAllOnAbortHandlers(this);
        UE_AUTORTFM_ONCOMMIT(this)
        {
            check(GetNumUsed() == 0);
            Trim();
        };
    }
 
    void* Allocate(int32 Alignment = MIN_ALIGNMENT)
    {
        void* Memory = nullptr;
 
        // We need the allocation to happen (so we get a valid `Memory`)
        // pointer to use in the transaction. So we do the actual meat
        // of the allocation in the open.
        UE_AUTORTFM_OPEN
        {
            NumUsed.Increment();
            if (Alignment <= 4096)
            {
                // Pop from a free list only if Alignment is not large than a memory page size
                Memory = FreeList.Pop();
                if (Memory)
                {
                    NumFree.Decrement();
                }
            }
            if (!Memory)
            {
                Memory = FMemory::Malloc(SIZE, Alignment);
            }
        };
 
        // But if we abort the transaction, we need to return the memory
        // region to the allocator (otherwise we'll leak).
        UE_AUTORTFM_ONABORT(this, Memory)
        {
            this->Free(Memory);
        };
        
        return Memory;
    }
 
    void Free(void* Item)
    {
        // We defer actually returning `Item` until on-commit time.
        // If we didn't do this an aborting transaction would not
        // be able to be undone.
        UE_AUTORTFM_ONCOMMIT(this, Item)
        {
            NumUsed.Decrement();
            FreeList.Push(Item);
            NumFree.Increment();
        };
    }
 
    void Trim()
    {
        // Similar to `Free`, we need to only do the destructive
        // trim operation at commit time - otherwise we'll wouldn't
        // be able to undo it.
        UE_AUTORTFM_ONCOMMIT(this)
        {
            while (void* Mem = FreeList.Pop())
            {
                FMemory::Free(Mem);
                NumFree.Decrement();
            }
        };
    }
 
    typename TTrackingCounter::IntegerType GetNumUsed() const
    {
        return AutoRTFM::Open([&] { return NumUsed.GetValue(); });
    }
 
    typename TTrackingCounter::IntegerType GetNumFree() const
    {
        return AutoRTFM::Open([&] { return NumFree.GetValue(); });
    }
 
private:
    UE_NONCOPYABLE(TLockFreeFixedSizeAllocator)
 
    
    TLockFreePointerListUnordered<void, TPaddingForCacheContention> FreeList;
 
    TTrackingCounter NumUsed;
 
    TTrackingCounter NumFree;
};
 
template<int32 SIZE, int TPaddingForCacheContention, typename TTrackingCounter = FNoopCounter>
class TLockFreeFixedSizeAllocator_TLSCache : public TLockFreeFixedSizeAllocator_TLSCacheBase<SIZE, TLockFreePointerListUnordered<void*, TPaddingForCacheContention>, TTrackingCounter>
{
};
 
template<class T, int TPaddingForCacheContention>
class TLockFreeClassAllocator : private TLockFreeFixedSizeAllocator<sizeof(T), TPaddingForCacheContention, FNoopCounter>
{
public:
    void* Allocate()
    {
        return TLockFreeFixedSizeAllocator<sizeof(T), TPaddingForCacheContention>::Allocate();
    }
 
    T* New()
    {
        return ::new (Allocate()) T();
    }
 
    void Free(T *Item)
    {
        Item->~T();
        TLockFreeFixedSizeAllocator<sizeof(T), TPaddingForCacheContention>::Free(Item);
    }
};
 
template<class T, int TPaddingForCacheContention>
class TLockFreeClassAllocator_TLSCache : private TLockFreeFixedSizeAllocator_TLSCache<sizeof(T), TPaddingForCacheContention, FNoopCounter>
{
public:
    void* Allocate()
    {
        return TLockFreeFixedSizeAllocator_TLSCache<sizeof(T), TPaddingForCacheContention>::Allocate();
    }
 
    T* New()
    {
        return ::new (Allocate()) T();
    }
 
    void Free(T *Item)
    {
        Item->~T();
        TLockFreeFixedSizeAllocator_TLSCache<sizeof(T), TPaddingForCacheContention>::Free(Item);
    }
};