WriteLog.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#if (defined(__AUTORTFM) && __AUTORTFM)
 
#include "AutoRTFM.h"
#include "BuildMacros.h"
#include "Utils.h"
 
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdint.h>
 
namespace AutoRTFM
{
    struct FWriteLogEntry final
    {
        // The address of the write.
        std::byte* LogicalAddress = nullptr;
        
        // A pointer to the original data before the write occurred.
        std::byte* Data = nullptr;
        
        // The size of the write in bytes.
        size_t Size = 0;
        
        // If true, then this write will not be considered by the AutoRTFM
        // memory validator.
        bool bNoMemoryValidation = false;
    };
 
    // FWriteLog holds an ordered list of write records which can be iterated
    // forwards and backwards.
    // Ensure changes to this class are kept in sync with Unreal.natvis.
    class FWriteLog final
    {
    public:
        // Number of bits used by the FRecord to represent a write's size.
        static constexpr size_t RecordSizeBits = 15;
 
        // The maximum size for a single write log entry.
        // Writes will be split into multiple entries if the write is too large.
        static constexpr size_t RecordMaxSize = (1u << RecordSizeBits) - 1;
 
    private:
        struct FRecord
        {
            void Set(std::byte* Address, size_t Size, bool bNoMemoryValidation)
            {
                AUTORTFM_ASSERT_DEBUG((reinterpret_cast<uintptr_t>(Address) & 0xFFFF0000'00000000) == 0);
                AUTORTFM_ASSERT_DEBUG(Size <= 0x7FFF);
        
                Bits = reinterpret_cast<uint64_t>(Address);
                Bits <<= 1;
                Bits |= bNoMemoryValidation;
                Bits <<= 15;
                Bits |= Size;
            }
 
            void Grow(size_t Amount)
            {
                AUTORTFM_ASSERT_DEBUG(Size() + Amount <= 0x7FFF);
                Bits += Amount;
            }
 
            std::byte* Address() const
            {
                return reinterpret_cast<std::byte*>(Bits >> 16);
            }
 
            bool NoMemoryValidation() const
            {
                return Bits & 0x8000;
            }
 
            size_t Size() const
            {
                return Bits & 0x7FFF;
            }
 
            uint64_t Bits;
        };
 
        static_assert(sizeof(uintptr_t) == 8, "assumption: a pointer is 8 bytes");
        static_assert(sizeof(FRecord) == 8);
 
        // Ensure changes to this structure are kept in sync with Unreal.natvis.
        struct FBlock final
        {
            // ┌────────┬────┬────┬────┬────┬────────────────┬────┬────┬────┬────┐
            // │ FBlock │ D₀ │ D₁ │ D₂ │ D₃ │->            <-│ R₃ │ R₂ │ R₁ │ R₀ │
            // └────────┴────┴────┴────┴────┴────────────────┴────┴────┴────┴────┘
            //          ^                   ^                ^              ^
            //      DataStart()          DataEnd         LastRecord    FirstRecord
            // Where:
            //   Dₙ = Data n, Rₙ = Record n
 
            // Starting size of a heap-allocated block, including the FBlock struct header.
            static constexpr size_t DefaultSize = 2048;
 
            // Constructor
            // TotalSize is the total size of the allocated memory for the block including
            // the FBlock header.
            explicit FBlock(size_t TotalSize)
            {
                AUTORTFM_ENSURE((TotalSize & (alignof(FRecord) - 1)) == 0);
                std::byte* End = reinterpret_cast<std::byte*>(this) + TotalSize;
                DataEnd = DataStart();
                // Note: The initial empty state has LastRecord pointing one
                // FRecord beyond the immutable FirstRecord.
                LastRecord = reinterpret_cast<FRecord*>(End);
                FirstRecord = LastRecord - 1;
            }
 
            // Allocate performs a heap allocation of a new block.
            // TotalSize is the total size of the allocated memory for the block including
            // the FBlock header.
            static FBlock* Allocate(size_t TotalSize)
            {
                AUTORTFM_ASSERT(TotalSize > (sizeof(FBlock) + sizeof(FRecord)));
                void* Memory = AutoRTFM::Allocate(TotalSize, alignof(FBlock));
                // Disable false-positive warning C6386: Buffer overrun while writing to 'Memory'
                CA_SUPPRESS(6386)
                return new (Memory) FBlock(TotalSize);
            }
 
            // Free releases the heap-allocated memory for this block.
            // Note: This block must have been allocated with a call to Allocate().
            void Free()
            {
                AutoRTFM::Free(this);
            }
 
            // Returns a pointer to the data for the first entry
            std::byte* DataStart()
            {
                return reinterpret_cast<std::byte*>(this) + sizeof(FBlock);
            }
 
            // Returns a pointer to the data for the last entry
            std::byte* LastData()
            {
                return DataEnd - LastRecord->Size();
            }
 
            // Returns true if the block holds no entries.
            bool IsEmpty() const
            {
                return LastRecord > FirstRecord;
            }
 
            // This will return 0 if the passed-in entry can be folded into the previous write, and 1 if a new
            // write record will be required.
            UE_AUTORTFM_FORCEINLINE int NumRecordsNeededForWrite(std::byte* LogicalAddress, bool bNoMemoryValidation)
            {
                return static_cast<int>(
                    IsEmpty() ||
                    LogicalAddress != LastRecord->Address() + LastRecord->Size() ||
                    LastRecord->NoMemoryValidation() != bNoMemoryValidation);
            }
 
            // This determines the number of bytes that can safely be passed to `Push` below, given the number of
            // bytes in the write entry (`EntrySize`) and whether or not those bytes can be folded into an existing
            // record (which can be determined by calling `NumRecordsNeededForWrite` above). This algorithm does
            // _not_ clamp the input size to RecordMaxSize. Instead, the block allocation logic in `Push` and
            // `PushSmall` avoids creating an FBlock larger than MaxSize at all. (When very large pushes occur, the
            // `Push` algorithm _will_ create a large FBlock to satisfy it, but this logic does not call
            // `CalculatePushBytes`, and will create a new tail block afterwards.)
            // 
            // Returns zero if the block is entirely full, `EntrySize` if there's enough space, or a value
            // somewhere in-between if the entry needs to be split.
            UE_AUTORTFM_FORCEINLINE size_t CalculatePushBytes(size_t EntrySize, int NumRecordsNeeded)
            {
                std::byte* BlockEdge = reinterpret_cast<std::byte*>(LastRecord - NumRecordsNeeded);
                constexpr ptrdiff_t MinimumWorthwhileSplitSize = 8;
 
                if (DataEnd + EntrySize <= BlockEdge)
                {
                    // We will fit the entire entry in the block with room to spare.
                    AUTORTFM_ASSERT_DEBUG(EntrySize <= RecordMaxSize);
                    AUTORTFM_ASSERT_DEBUG((NumRecordsNeeded == 1) || (LastRecord->Size() + EntrySize <= ptrdiff_t(RecordMaxSize)));
                    return EntrySize;
                }
                else if (ptrdiff_t BlockCapacity = BlockEdge - DataEnd; BlockCapacity >= MinimumWorthwhileSplitSize)
                {
                    // We have enough data to fill up the entire block. This path will split the data across the end of this
                    // block and the start of the next block. We avoid this path when the block has fewer than eight bytes
                    // left, just as an efficiency measure, since it takes extra time to assemble two FRecords instead of one.
                    AUTORTFM_ASSERT_DEBUG(BlockCapacity <= ptrdiff_t(RecordMaxSize));
                    AUTORTFM_ASSERT_DEBUG((NumRecordsNeeded == 1) || (LastRecord->Size() + BlockCapacity <= ptrdiff_t(RecordMaxSize)));
                    return BlockCapacity;
                }
                else
                {
                    // The block is completely full.
                    return 0;
                }
            }
 
            // Grows this block by copying `NumBytes` bytes from `DataIn`, representing data originally from `LogicalAddress`. 
            // Asserts if the block does not have enough capacity to hold `NumBytes`. The caller should determine the available
            // capacity ahead of time by calling `NumRecordsNeededForWrite` and `CalculatePushBytes`.
            UE_AUTORTFM_FORCEINLINE void Push(std::byte* LogicalAddress, std::byte* DataIn, size_t NumBytes, bool bNoMemoryValidation, int NumRecordsNeeded)
            {
                AUTORTFM_ASSERT_DEBUG(DataEnd + NumBytes <= reinterpret_cast<std::byte*>(LastRecord - NumRecordsNeeded));
                AUTORTFM_ASSERT_DEBUG(NumBytes <= RecordMaxSize);
 
                if (NumRecordsNeeded == 1)
                {
                    LastRecord--;
                    LastRecord->Set(LogicalAddress, NumBytes, bNoMemoryValidation);
                }
                else
                {
                    AUTORTFM_ASSERT_DEBUG(NumRecordsNeeded == 0);
                    LastRecord->Grow(NumBytes);
                }
 
                memcpy(DataEnd, DataIn, NumBytes);
 
                DataEnd += NumBytes;
 
                AUTORTFM_ASSERT_DEBUG(DataEnd <= reinterpret_cast<std::byte*>(LastRecord));
            }
 
            // The next block in the linked list.
            FBlock* NextBlock = nullptr;
            // The previous block in the linked list.
            FBlock* PrevBlock = nullptr;
            // The pointer to the first entry's record
            FRecord* FirstRecord = nullptr;
            // The pointer to the last entry's record
            FRecord* LastRecord = nullptr;
            // One byte beyond the end of the last entry's data
            std::byte* DataEnd = nullptr;
        private:
            ~FBlock() = delete;
        };
 
    public:
        // Constructor
        FWriteLog()
        {
            new(HeadBlockMemory) FBlock(HeadBlockSize);
        }
 
        // Destructor
        ~FWriteLog()
        {
            Reset();
        }
 
        // Adds the write log entry to the log.
        // The log will make a copy of the FWriteLogEntry's data.
        void Push(FWriteLogEntry Entry)
        {
            AUTORTFM_ASSERT_DEBUG((reinterpret_cast<uintptr_t>(Entry.LogicalAddress) & 0xffff0000'00000000) == 0);
 
            {
                TotalSizeBytes += Entry.Size;
                int NumRecordsNeeded = TailBlock->NumRecordsNeededForWrite(Entry.LogicalAddress, Entry.bNoMemoryValidation);
                size_t NumBytes = TailBlock->CalculatePushBytes(Entry.Size, NumRecordsNeeded);
 
                if (NumBytes == Entry.Size)
                {
                    // This push fits into our existing block.
                    TailBlock->Push(Entry.LogicalAddress, Entry.Data, NumBytes, Entry.bNoMemoryValidation, NumRecordsNeeded);
                    NumEntries += NumRecordsNeeded;
                    return;
                }
 
                // The push doesn't fit into the existing block...
                if (NumBytes > 0)
                {
                    // ... but we can still use up the remainder of the block.
                    TailBlock->Push(Entry.LogicalAddress, Entry.Data, NumBytes, Entry.bNoMemoryValidation, NumRecordsNeeded);
                    NumEntries += NumRecordsNeeded;
 
                    // Adjust the entry to point to the remaining unlogged bytes.
                    Entry.Size -= NumBytes;
                    Entry.LogicalAddress += NumBytes;
                    Entry.Data += NumBytes;
                }
            }
 
            // Calculate how many maxed-out RecordMaxSize entries we can make.
            const size_t NumFullRecords = Entry.Size / RecordMaxSize;
            // Calculate how many bytes will be remaining once we have emitted the full-size entries.
            const size_t RemainingBytes = Entry.Size - (NumFullRecords * RecordMaxSize);
            // Calculate the exact required size of the block.
            const size_t RequiredSize =
                sizeof(FBlock) +                                                   // FBlock header
                (NumFullRecords * (RecordMaxSize + sizeof(FRecord))) +             // Bytes needed for full records
                (LIKELY(RemainingBytes) ? (RemainingBytes + sizeof(FRecord)) : 0); // Bytes needed for trailing partial record, if any
            // Add padding to the block to account for alignment.
            const size_t AlignedSize = AlignUp(RequiredSize, sizeof(FRecord));
 
            // Create a new empty tail block, large enough to hold the remainder of this push in its entirety, and
            // never smaller than the upcoming block size.
            const size_t BlockSize = std::max(AlignedSize, NextBlockSize);
            AllocateNewBlock(BlockSize);
 
            // Push all of the full records.
            for (size_t Index = NumFullRecords; Index--; )
            {
                constexpr int NumRecordsNeeded = 1;
                constexpr size_t NumBytes = RecordMaxSize;
 
                TailBlock->Push(Entry.LogicalAddress, Entry.Data, NumBytes, Entry.bNoMemoryValidation, NumRecordsNeeded);
                NumEntries += NumRecordsNeeded;
 
                Entry.Size -= NumBytes;
                Entry.LogicalAddress += NumBytes;
                Entry.Data += NumBytes;
            }
 
            // Push the final, partial record.
            if (LIKELY(RemainingBytes))
            {
                constexpr int NumRecordsNeeded = 1;
 
                TailBlock->Push(Entry.LogicalAddress, Entry.Data, RemainingBytes, Entry.bNoMemoryValidation, NumRecordsNeeded);
                NumEntries += NumRecordsNeeded;
            }
 
            // If we've just created an extra-large block, and alignment has caused it to be less than 100% full,
            // we preemptively allocate a new block here. This avoids a scenario where a future PushSmall could
            // accidentally overflow the record size. Normally this would be impossible because BlockMaxSize
            // isn't large enough to hold more than RecordMaxSize bytes, but in this case we are making an FBlock
            // which is potentially much larger than normal.
            if (AlignedSize > BlockMaxSize && AlignedSize > RequiredSize)
            {
                AllocateNewBlock(NextBlockSize);
            }
        }
 
        // Adds the write log entry to the log; assumes a payload small enough that splitting is not beneficial.
        // If you have large sizes, you should use `Push` instead so that splitting can occur.
        template <unsigned int SIZE>
        void PushSmall(std::byte* LogicalAddress)
        {
            static_assert(SIZE <= RecordMaxSize);  // This is a hard upper limit.
            AUTORTFM_ASSERT_DEBUG((reinterpret_cast<uintptr_t>(LogicalAddress) & 0xffff0000'00000000) == 0);
 
            TotalSizeBytes += SIZE;
 
            int NumRecordsNeeded = TailBlock->NumRecordsNeededForWrite(LogicalAddress, /*bNoMemoryValidation=*/false);
            size_t NumBytes = TailBlock->CalculatePushBytes(SIZE, NumRecordsNeeded);
 
            if (UNLIKELY(NumBytes != SIZE))
            {
                AllocateNewBlock(NextBlockSize);
                NumRecordsNeeded = 1;
            }
 
            TailBlock->Push(LogicalAddress, LogicalAddress, SIZE, /*bNoMemoryValidation=*/false, NumRecordsNeeded);
            NumEntries += NumRecordsNeeded;
        }
 
        // Iterator for enumerating the writes of the log.
        template<bool IS_FORWARD>
        struct TIterator final
        {
            TIterator() = default;
 
            TIterator(FBlock* StartBlock) : Block(StartBlock)
            {
                if (UNLIKELY(Block->IsEmpty()))
                {
                    if (UNLIKELY(!AdvanceBlock()))
                    {
                        // The write log is entirely empty.
                        return;
                    }
                }
 
                Data = IS_FORWARD ? Block->DataStart() : Block->LastData();
                Record = IS_FORWARD ? Block->FirstRecord : Block->LastRecord;
            }
 
            // Returns the entry at the current iterator's position.
            FWriteLogEntry operator*() const
            {
                return FWriteLogEntry
                {
                    .LogicalAddress = reinterpret_cast<std::byte*>(Record->Address()),
                    .Data = Data,
                    .Size = Record->Size(),
                    .bNoMemoryValidation = Record->NoMemoryValidation(),
                };
            }
 
            // Progresses the iterator to the next entry
            void operator++()
            {
                if constexpr (IS_FORWARD)
                {
                    if (Record == Block->LastRecord)
                    {
                        if (LIKELY(AdvanceBlock()))
                        {
                            Data = Block->DataStart();
                            Record = Block->FirstRecord;
                        }
                    }
                    else
                    {
                        Data += Record->Size();
                        Record--;
                    }
                }
                else
                {
                    if (Record == Block->FirstRecord)
                    {
                        if (LIKELY(AdvanceBlock()))
                        {
                            Data = Block->LastData();
                            Record = Block->LastRecord;
                        }
                    }
                    else
                    {
                        Record++;
                        Data -= Record->Size();
                    }
                }
            }
 
            // Inequality operator
            bool operator!=(const TIterator& Other) const
            {
                return (Other.Block != Block) || (Other.Record != Record);
            }
 
        private:
            // Resets the iterator (compares equal to the write log's end())
            UE_AUTORTFM_FORCEINLINE void Reset()
            {
                Block = nullptr;
                Data = nullptr;
                Record = nullptr;
            }
 
            // Moves from this block to the next (if IS_FORWARD) or previous (if not IS_FORWARD),
            // skipping any empty blocks. Returns true on success. If no more blocks exist, resets 
            // the iterator and returns false.
            UE_AUTORTFM_FORCEINLINE bool AdvanceBlock()
            {
                do
                {
                    Block = IS_FORWARD ? Block->NextBlock : Block->PrevBlock;
                }
                while (Block && Block->IsEmpty());
 
                if (!Block)
                {
                    Reset();
                    return false;
                }
 
                return true;
            }
 
            FBlock* Block = nullptr;
            std::byte* Data = nullptr;
            FRecord* Record = nullptr;
        };
 
        using Iterator = TIterator</* IS_FORWARD */ true>;
        using ReverseIterator = TIterator</* IS_FORWARD */ false>;
 
        Iterator begin() const
        {
            return (NumEntries > 0) ? Iterator(HeadBlock) : Iterator{};
        }
        ReverseIterator rbegin() const
        {
            return (NumEntries > 0) ? ReverseIterator(TailBlock) : ReverseIterator{};
        }
        Iterator end() const { return Iterator{}; }
        ReverseIterator rend() const { return ReverseIterator{}; }
 
        // Resets the write log to its initial state, freeing any allocated memory.
        void Reset()
        {
            // Skip HeadBlock, which is held as part of this structure.
            FBlock* Block = HeadBlock->NextBlock;
            while (nullptr != Block)
            {
                FBlock* const Next = Block->NextBlock;
                Block->Free();
                Block = Next;
            }
            new (HeadBlockMemory) FBlock(HeadBlockSize);
            HeadBlock = reinterpret_cast<FBlock*>(HeadBlockMemory);
            TailBlock = reinterpret_cast<FBlock*>(HeadBlockMemory);
            NumEntries = 0;
            TotalSizeBytes = 0;
            NextBlockSize = FBlock::DefaultSize;
        }
 
        // Returns true if the log holds no entries.
        UE_AUTORTFM_FORCEINLINE bool IsEmpty() const { return 0 == NumEntries; }
 
        // Return the number of entries in the log.
        UE_AUTORTFM_FORCEINLINE size_t Num() const { return NumEntries; }
 
        // Return the total size in bytes for all entries in the log.
        UE_AUTORTFM_FORCEINLINE size_t TotalSize() const { return TotalSizeBytes; }
 
        // Returns a hash of the first NumWriteEntries entries' logical memory
        // tracked by the write log. This is the memory post-write, not the
        // original memory that would be restored on abort.
        using FHash = uint64_t;
        FHash Hash(size_t NumWriteEntries) const;
 
    private:
        UE_AUTORTFM_FORCEINLINE void AllocateNewBlock(size_t Size)
        {
            FBlock* NewBlock = FBlock::Allocate(Size);
            NewBlock->PrevBlock = TailBlock;
            TailBlock->NextBlock = NewBlock;
            TailBlock = NewBlock;
 
            // Increase block sizes by 50% each time, capped at BlockMaxSize.
            NextBlockSize = std::min<size_t>((NextBlockSize * 3 / 2), BlockMaxSize);
        }
 
        template <size_t SIZE>
        static constexpr bool IsAlignedForFRecord = (SIZE & (alignof(FRecord) - 1)) == 0;
 
        // The size of the inline block, which is declared as a byte array (`HeadBlockMemory`)
        // inside the write log.
        static constexpr size_t HeadBlockSize = 256;
 
        // The upper bound on heap-allocated block size. We avoid making blocks larger 
        // than the maximum record size, so we don't need to insert size overflow 
        // checks throughout the push logic.
        static constexpr size_t BlockMaxSize = AlignDown(sizeof(FBlock) + sizeof(FRecord) + RecordMaxSize, alignof(FRecord));
 
        static_assert(IsAlignedForFRecord<HeadBlockSize>);
        static_assert(IsAlignedForFRecord<FBlock::DefaultSize>);
        static_assert(IsAlignedForFRecord<BlockMaxSize>);
 
        FHash HashAVX2(size_t NumWriteEntries) const;
 
        FBlock* HeadBlock = reinterpret_cast<FBlock*>(HeadBlockMemory);
        FBlock* TailBlock = reinterpret_cast<FBlock*>(HeadBlockMemory);
        size_t NumEntries = 0;
        size_t TotalSizeBytes = 0;
        size_t NextBlockSize = FBlock::DefaultSize;
        alignas(alignof(FBlock)) std::byte HeadBlockMemory[HeadBlockSize];
    };
}
 
#endif // (defined(__AUTORTFM) && __AUTORTFM)