IoDispatcherFileBackend.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "FileIoDispatcherBackend.h"
 
#include "IoDispatcherFileBackendTypes.h"
#include "IO/IoDispatcher.h"
#include "IO/IoStore.h"
#include "Containers/Array.h"
#include "Containers/Map.h"
#include "Containers/SpscQueue.h"
#include "Stats/Stats.h"
#include "Tasks/Task.h"
#include "HAL/Runnable.h"
#include "Misc/AES.h"
#include "GenericPlatform/GenericPlatformFile.h"
 
class IMappedFileHandle;
 
struct FFileIoStoreCompressionContext
{
    FFileIoStoreCompressionContext* Next = nullptr;
    uint64 UncompressedBufferSize = 0;
    uint8* UncompressedBuffer = nullptr;
};
 
class FFileIoStoreReader
{
public:
    FFileIoStoreReader(IPlatformFileIoStore& InPlatformImpl, FFileIoStoreStats& InStats);
    ~FFileIoStoreReader();
    FIoStatus Initialize(const TCHAR* InTocFilePath, int32 Order);
    uint32 GetContainerInstanceId() const
    {
        return ContainerFile.ContainerInstanceId;
    }
    FIoStatus Close();
    bool DoesChunkExist(const FIoChunkId& ChunkId) const;
    TIoStatusOr<uint64> GetSizeForChunk(const FIoChunkId& ChunkId) const;
    const FIoOffsetAndLength* Resolve(const FIoChunkId& ChunkId) const;
    FFileIoStoreContainerFile* GetContainerFile() { return &ContainerFile; }
    const FFileIoStoreContainerFile* GetContainerFile() const { return &ContainerFile; }
    IMappedFileHandle* GetMappedContainerFileHandle(uint64 TocOffset);
    const FIoContainerId& GetContainerId() const { return ContainerId; }
    int32 GetOrder() const { return Order; }
    bool IsEncrypted() const { return EnumHasAnyFlags(ContainerFile.ContainerFlags, EIoContainerFlags::Encrypted); }
    bool IsSigned() const { return EnumHasAnyFlags(ContainerFile.ContainerFlags, EIoContainerFlags::Signed); }
    const FGuid& GetEncryptionKeyGuid() const { return ContainerFile.EncryptionKeyGuid; }
    void SetEncryptionKey(const FAES::FAESKey& Key) { ContainerFile.EncryptionKey = Key; }
    const FAES::FAESKey& GetEncryptionKey() const { return ContainerFile.EncryptionKey; }
    TIoStatusOr<FIoContainerHeader> ReadContainerHeader(bool bReadSoftRefs) const;
    void ReopenAllFileHandles();
 
private:
    const FIoOffsetAndLength* FindChunkInternal(const FIoChunkId& ChunkId) const;
    uint64 GetTocAllocatedSize() const;
    
    IPlatformFileIoStore& PlatformImpl;
    FFileIoStoreStats& Stats;
 
    struct FPerfectHashMap
    {
        TConstArrayView<int32> TocChunkHashSeeds;
        TConstArrayView<FIoChunkId> TocChunkIds;
        TConstArrayView<FIoOffsetAndLength> TocOffsetAndLengths;
    };
    FPerfectHashMap PerfectHashMap;
    TMap<FIoChunkId, FIoOffsetAndLength> TocImperfectHashMapFallback;
    FFileIoStoreContainerFile ContainerFile;
    FIoContainerId ContainerId;
    // Owns the data for ArrayViews in ContainerFile and in PerfectHash map
    FIoStoreTocResourceStorage DataContainer;
    int32 Order;
    bool bClosed = false;
    bool bHasPerfectHashMap = false;
 
    static TAtomic<uint32> GlobalPartitionIndex;
    static TAtomic<uint32> GlobalContainerInstanceId;
};
 
class FFileIoStoreRequestTracker
{
public:
    FFileIoStoreRequestTracker(FFileIoStoreRequestAllocator& RequestAllocator, FFileIoStoreRequestQueue& RequestQueue);
    ~FFileIoStoreRequestTracker();
 
    FFileIoStoreCompressedBlock* FindOrAddCompressedBlock(FFileIoStoreBlockKey Key, bool& bOutWasAdded);
    void RemoveCompressedBlock(const FFileIoStoreCompressedBlock* CompressedBlock, bool bRemoveFromCancel = false);
    FFileIoStoreReadRequest* FindOrAddRawBlock(FFileIoStoreBlockKey Key, bool& bOutWasAdded);
    void RemoveRawBlock(const FFileIoStoreReadRequest* RawBlock, bool bRemoveFromCancel = false);
    void AddReadRequestsToResolvedRequest(FFileIoStoreCompressedBlock* CompressedBlock, FFileIoStoreResolvedRequest& ResolvedRequest);
    void AddReadRequestsToResolvedRequest(const FFileIoStoreReadRequestList& Requests, FFileIoStoreResolvedRequest& ResolvedRequest);
    bool CancelIoRequest(FFileIoStoreResolvedRequest& ResolvedRequest);
    void UpdatePriorityForIoRequest(FFileIoStoreResolvedRequest& ResolvedRequest);
    void ReleaseIoRequestReferences(FFileIoStoreResolvedRequest& ResolvedRequest);
    int64 GetLiveReadRequestsCount() const;
 
private:
    FFileIoStoreRequestAllocator& RequestAllocator;
    FFileIoStoreRequestQueue& RequestQueue;
    TMap<FFileIoStoreBlockKey, FFileIoStoreCompressedBlock*> CompressedBlocksMap;
    TMap<FFileIoStoreBlockKey, FFileIoStoreReadRequest*> RawBlocksMap;
};
 
class FFileIoStore final
    : public FRunnable
    , public UE::IoStore::IFileIoDispatcherBackend
{
public:
    FFileIoStore(TUniquePtr<IPlatformFileIoStore>&& PlatformImpl);
    ~FFileIoStore();
    void Initialize(TSharedRef<const FIoDispatcherBackendContext> Context) override;
    void Shutdown() override;
    virtual TIoStatusOr<FIoContainerHeader> Mount(
        const TCHAR* InTocPath, 
        int32 Order, 
        const FGuid& EncryptionKeyGuid, 
        const FAES::FAESKey& EncryptionKey, 
        UE::IoStore::ETocMountOptions Options) override;
    virtual bool Unmount(const TCHAR* InTocPath) override;
    void ResolveIoRequests(FIoRequestList Requests, FIoRequestList& OutUnresolved) override;
    void CancelIoRequest(FIoRequestImpl* Request) override;
    void UpdatePriorityForIoRequest(FIoRequestImpl* Request) override;
    bool DoesChunkExist(const FIoChunkId& ChunkId) const override;
    TIoStatusOr<uint64> GetSizeForChunk(const FIoChunkId& ChunkId) const;
    FIoRequestImpl* GetCompletedIoRequests() override;
    TIoStatusOr<FIoMappedRegion> OpenMapped(const FIoChunkId& ChunkId, const FIoReadOptions& Options) override;
    const TCHAR* GetName() const override;
    virtual void ReopenAllFileHandles() override;
 
    virtual bool Init() override;
    virtual uint32 Run() override;
    virtual void Stop() override;
 
    uint32 GetThreadId() const;
 
private:
 
    bool Resolve(FIoRequestImpl* Request);
    void OnNewPendingRequestsAdded();
    void ReadBlocks(FFileIoStoreResolvedRequest& ResolvedRequest);
    void FreeBuffer(FFileIoStoreBuffer& Buffer);
    FFileIoStoreCompressionContext* AllocCompressionContext();
    void FreeCompressionContext(FFileIoStoreCompressionContext* CompressionContext);
    void ScatterBlock(FFileIoStoreCompressedBlock* CompressedBlock, bool bIsAsync);
    void CompleteDispatcherRequest(FFileIoStoreResolvedRequest* ResolvedRequest);
    void FinalizeCompressedBlock(FFileIoStoreCompressedBlock* CompressedBlock);
    void StopThread();
 
    uint64 ReadBufferSize = 0;
    TSharedPtr<const FIoDispatcherBackendContext> BackendContext;
    FFileIoStoreStats Stats;
    FFileIoStoreBlockCache BlockCache;
    FFileIoStoreBufferAllocator BufferAllocator;
    FFileIoStoreRequestAllocator RequestAllocator;
    FFileIoStoreRequestQueue RequestQueue;
    FFileIoStoreRequestTracker RequestTracker;
    TUniquePtr<IPlatformFileIoStore> PlatformImpl;
    FRunnableThread* Thread = nullptr;
    bool bIsMultithreaded;
    TAtomic<bool> bStopRequested{ false };
    mutable FRWLock IoStoreReadersLock;
    TArray<TUniquePtr<FFileIoStoreReader>> IoStoreReaders;
    TArray<TUniquePtr<FFileIoStoreCompressionContext>> CompressionContexts;
    TSpscQueue<UE::Tasks::FTask> DecompressionTasks;
    FFileIoStoreCompressionContext* FirstFreeCompressionContext = nullptr;
    FFileIoStoreCompressedBlock* ReadyForDecompressionHead = nullptr;
    FFileIoStoreCompressedBlock* ReadyForDecompressionTail = nullptr;
    FCriticalSection DecompressedBlocksCritical;
    FFileIoStoreCompressedBlock* FirstDecompressedBlock = nullptr;
    FIoRequestImpl* CompletedRequestsHead = nullptr;
    FIoRequestImpl* CompletedRequestsTail = nullptr;
    FDelegateHandle OversubscriptionLimitReached;
};
 
TSharedRef<FFileIoStore> CreateIoDispatcherFileBackend();