ODSCThread.h

Go to the documentation of this file.

// Copyright Epic Games, Inc. All Rights Reserved.
 
#pragma once
 
#include "CoreMinimal.h"
#include "HAL/Runnable.h"
#include "Misc/SingleThreadRunnable.h"
#include "GenericPlatform/GenericPlatformFile.h"
#include "Containers/Queue.h"
#include "ShaderCompiler.h"
#include "RHIDefinitions.h"
 
class FEvent;
class FRunnableThread;
class FMaterialShaderMap;
class FMaterialShaderMapId;
class FPrimitiveSceneInfo;
 
namespace UE
{
    namespace Cook
    {
        class ICookOnTheFlyServerConnection;
        class FCookOnTheFlyMessage;
    }
}
 
enum class EODSCMetaDataType
{
    Default=0, // The material hasn't been seen by ODSCManager yet
    IsDependentOnMaterialName,
    IsNotDependentOnMaterialName,
};
 
class FODSCMessageHandler : public IPlatformFile::IFileServerMessageHandler
{
public:
    FODSCMessageHandler(EShaderPlatform InShaderPlatform, ERHIFeatureLevel::Type InFeatureLevel, EMaterialQualityLevel::Type InQualityLevel, ODSCRecompileCommand InRecompileCommandType);
    FODSCMessageHandler(
        const TArray<FString>& InMaterials,
        const FString& ShaderTypesToLoad,
        EShaderPlatform InShaderPlatform,
        ERHIFeatureLevel::Type InFeatureLevel,
        EMaterialQualityLevel::Type InQualityLevel,
        ODSCRecompileCommand InRecompileCommandType,
        const FShaderCompilerFlags& InExtraCompilerFlags = FShaderCompilerFlags()
    );
 
    virtual void FillPayload(FArchive& Payload) override;
 
    virtual void ProcessResponse(FArchive& Response) override;
 
    void AddPayload(const FODSCRequestPayload& Payload);
 
    const TArray<FString>& GetMaterialsToLoad() const;
    const TArray<uint8>& GetMeshMaterialMaps() const;
    const TArray<uint8>& GetGlobalShaderMap() const;
    bool ReloadGlobalShaders() const;
    ODSCRecompileCommand GetRecompileCommandType() const { return RecompileCommandType; };
    int32 NumPayloads() const { return RequestBatch.Num(); }
 
private:
    double RequestStartTime = 0.0;
 
    TArray<FString> MaterialsToLoad;
 
    FString ShaderTypesToLoad;
 
    EShaderPlatform ShaderPlatform;
 
    ERHIFeatureLevel::Type FeatureLevel;
 
    EMaterialQualityLevel::Type QualityLevel;
 
    ODSCRecompileCommand RecompileCommandType = ODSCRecompileCommand::None;
 
    FShaderCompilerFlags ExtraCompilerFlags = FShaderCompilerFlags();
 
    TArray<FODSCRequestPayload> RequestBatch;
 
    TArray<uint8> OutMeshMaterialMaps;
 
    TArray<uint8> OutGlobalShaderMap;
};
 
class FODSCThread
    : FRunnable, FSingleThreadRunnable
{
public:
 
    FODSCThread(const FString& HostIP);
    virtual ~FODSCThread();
 
    void StartThread();
 
    void StopThread();
 
    //~ Begin FSingleThreadRunnable Interface
    // Cannot be overriden to ensure identical behavior with the threaded tick
    virtual void Tick() override final;
    //~ End FSingleThreadRunnable Interface
 
    void AddRequest(
        const TArray<FString>& MaterialsToCompile,
        const FString& ShaderTypesToLoad,
        EShaderPlatform ShaderPlatform,
        ERHIFeatureLevel::Type FeatureLevel,
        EMaterialQualityLevel::Type QualityLevel,
        ODSCRecompileCommand RecompileCommandType,
        const FShaderCompilerFlags& ExtraCompilerFlags = FShaderCompilerFlags()
    );
 
    void AddShaderPipelineRequest(
        EShaderPlatform ShaderPlatform,
        ERHIFeatureLevel::Type FeatureLevel,
        EMaterialQualityLevel::Type QualityLevel,
        const FMaterial* Material,
        const FPrimitiveSceneInfo* PrimitiveSceneInfo,
        const FString& VertexFactoryName,
        const FString& PipelineName,
        const TArray<FString>& ShaderTypeNames,
        int32 PermutationId,
        const TArray<FShaderId>& RequestShaderIds
    );
 
    void GetCompletedRequests(TArray<FODSCMessageHandler*>& OutCompletedRequests);
 
    void Wakeup();
    
    void WaitUntilAllRequestsDone();
 
    bool GetPendingShaderData(bool& bOutIsConnectedToODSCServer, bool& bOutHasPendingGlobalShaders, uint32& OutNumPendingMaterialsRecompile, uint32& OutNumPendingMaterialsShaders) const;
 
    void ResetMaterialsODSCData(EShaderPlatform InShaderPlatform);
 
    bool CheckIfRequestAlreadySent(const TArray<FShaderId>& RequestShaderIds, const FMaterial* Material) const;
    const FString& GetODSCHostIP() const { return ODSCHostIP; };
 
    void UnregisterMaterialName(const FMaterial* Material);
    void RegisterMaterialShaderMaps(const FString& MaterialName, const TArray<TRefCountPtr<FMaterialShaderMap>>& LoadedShaderMaps);
    FMaterialShaderMap* FindMaterialShaderMap(const FString& MaterialName, const FMaterialShaderMapId& ShaderMapId) const;
 
    void RetrieveMissedMaterials(TArray<FString>& OutMaterialPaths) const;
 
protected:
 
    //~ Begin FRunnable Interface
    virtual bool Init() override;
    virtual uint32 Run() override final;
    virtual void Stop() override;
    virtual void Exit() override;
    //~ End FRunnable Interface
 
    FThreadSafeCounter ExitRequest;
 
private:
 
    void Process();
 
    bool ConnectToODSCHost();
    bool CheckODSCConnection();
 
    TQueue<FODSCMessageHandler*, EQueueMode::Mpsc> PendingMaterialThreadedRequests;
 
    TQueue<FODSCRequestPayload, EQueueMode::Mpsc> PendingMeshMaterialThreadedRequests;
 
    TQueue<FODSCMessageHandler*, EQueueMode::Spsc> CompletedThreadedRequests;
 
    mutable FRWLock RequestHashesRWLock;
    FCriticalSection RequestHashCriticalSection;
 
    struct FODSCShaderId
    {
    public:
        inline FODSCShaderId() {}
        FODSCShaderId(const FShaderId& ShaderId);
 
        FHashedName ShaderTypeHashedName = 0;
        FHashedName VFTypeHashedName = 0;
        FHashedName ShaderPipelineName = 0;
        int32 PermutationId = 0;
        uint32 Platform : SP_NumBits = SP_NumPlatforms;
 
        friend inline uint32 GetTypeHash( const FODSCShaderId& Id )
        {
            return HashCombine(
                GetTypeHash(Id.ShaderTypeHashedName),
                HashCombine(GetTypeHash(Id.VFTypeHashedName),
                            HashCombine(GetTypeHash(Id.ShaderPipelineName),
                                        HashCombine(GetTypeHash(Id.PermutationId), GetTypeHash(Id.Platform)))));
        }
 
        friend bool operator==(const FODSCShaderId& X, const FODSCShaderId& Y)
        {
            return X.ShaderTypeHashedName == Y.ShaderTypeHashedName
            && X.ShaderPipelineName == Y.ShaderPipelineName
            && X.VFTypeHashedName == Y.VFTypeHashedName
            && X.PermutationId == Y.PermutationId 
            && X.Platform == Y.Platform;
        }
 
        friend bool operator!=(const FODSCShaderId& X, const FODSCShaderId& Y)
        {
            return !(X == Y);
        }
    };
 
 
    struct FODSCShaderMapData
    {
        TArray<TRefCountPtr<FMaterialShaderMap>> MaterialShaderMaps;
        TSet<FODSCShaderId> CurrentRequests;
        
        FName ActorPath;
    };
 
    TMap<FName, FODSCShaderMapData> RequestHashes;
 
    TMap<UPTRINT, FName> ODSCPointerToNames; 
 
    FRunnableThread* Thread = nullptr;
 
    FEvent* WakeupEvent;
    
    FEvent* AllRequestsDoneEvent;
 
    bool SendMessageToServer(IPlatformFile::IFileServerMessageHandler* Handler);
 
    TArray<FODSCMessageHandler*> PendingRequestsMaterialAndGlobal;
    TArray<FODSCMessageHandler*> PendingRequestsPipeline;
 
    TUniquePtr<UE::Cook::ICookOnTheFlyServerConnection> CookOnTheFlyServerConnection;
 
    FString ODSCHostIP;
 
    int32 MaxPayloadPerRequests = INT32_MAX;
 
    std::atomic<bool> bIsConnectedToODSCServer = false;
    std::atomic<bool> bHasPendingGlobalShaders = false;
    std::atomic<uint32> NumPendingMaterialsRecompile = 0;
    std::atomic<uint32> NumPendingMaterialsShaders = 0;
    bool bHasDefaultConnection = false;
};