Hydrogent: accurately compute texture loader memory usage

Hydrogent/include/HnTextureRegistry.hpp

@@ -226,7 +226,6 @@ class HnTextureRegistry final : public std::enable_shared_from_this<HnTextureReg
 
     std::atomic<Uint32> m_NextTextureId{0};
     std::atomic<Int32>  m_NumTexturesLoading{0};
-    std::atomic<Int64>  m_LoadingTexDataSize{0};
     std::atomic<Uint32> m_StorageVersion{0};
     std::atomic<Uint32> m_DataVersion{0};
     std::atomic<Int64>  m_AtlasDataSize{0};

Hydrogent/include/HnTextureUtils.hpp

@@ -38,6 +38,8 @@ namespace USD
 RefCntAutoPtr<ITextureLoader> CreateTextureLoaderFromSdfPath(const char*            SdfPath,
                                                              const TextureLoadInfo& LoadInfo);
 
+Int64 GetTextureLoaderMemoryUsage();
+
 } // namespace USD
 
 } // namespace Diligent

Hydrogent/src/HnTextureRegistry.cpp

@@ -82,9 +82,6 @@ void HnTextureRegistry::WaitForAsyncTasks()
 {
     if (m_pThreadPool)
     {
-        VERIFY(m_NumTexturesLoading != 0 || m_LoadingTexDataSize == 0,
-               "The number of loading textures is zero, but the loading data size is non-zero (", m_LoadingTexDataSize, ")");
-
         for (RefCntAutoPtr<IAsyncTask>& pTask : m_AsyncTasks)
         {
             pTask->Cancel();
@@ -223,9 +220,6 @@ void HnTextureRegistry::Commit(IDeviceContext* pContext)
             TextureHandle& Handle = *tex_it.second.Handle;
             VERIFY_EXPR(Handle.IsInitialized());
 
-            m_LoadingTexDataSize.fetch_add(-static_cast<Int64>(Handle.DataSize));
-            VERIFY_EXPR(m_LoadingTexDataSize.load() >= 0);
-
             m_DataVersion.fetch_add(1);
 
             if (Handle.GetTexture())
@@ -281,7 +275,6 @@ void HnTextureRegistry::LoadTexture(const pxr::TfToken
 
     const TextureDesc& TexDesc = pLoader->GetTextureDesc();
     TexHandle->DataSize        = GetStagingTextureDataSize(TexDesc);
-    m_LoadingTexDataSize.fetch_add(static_cast<Int64>(TexHandle->DataSize));
 
     // Try to allocate texture in the atlas first
     RefCntAutoPtr<ITextureAtlasSuballocation> pAtlasSuballocation;
@@ -363,7 +356,7 @@ HnTextureRegistry::TextureHandleSharedPtr HnTextureRegistry::Allocate(const pxr:
                 RefCntAutoPtr<IAsyncTask> pTask = EnqueueAsyncWork(
                     m_pThreadPool,
                     [=](Uint32 ThreadId) {
-                        if (m_LoadBudget != 0 && m_LoadingTexDataSize.load() > m_LoadBudget)
+                        if (m_LoadBudget != 0 && GetTextureLoaderMemoryUsage() > m_LoadBudget)
                         {
                             // Reschedule the task
                             return ASYNC_TASK_STATUS_NOT_STARTED;
@@ -453,7 +446,7 @@ HnTextureRegistry::UsageStats HnTextureRegistry::GetUsageStats() const
 {
     UsageStats Stats;
     Stats.NumTexturesLoading  = m_NumTexturesLoading.load();
-    Stats.LoadingTexDataSize  = m_LoadingTexDataSize.load();
+    Stats.LoadingTexDataSize  = GetTextureLoaderMemoryUsage();
     Stats.AtlasDataSize       = m_AtlasDataSize.load();
     Stats.SeparateTexDataSize = m_SeparateTexDataSize.load();
     return Stats;

Hydrogent/src/HnTextureUtils.cpp

@@ -28,6 +28,7 @@
 
 #include "ProxyDataBlob.hpp"
 #include "RefCntContainer.hpp"
+#include "DefaultRawMemoryAllocator.hpp"
 
 #include "pxr/usd/ar/asset.h"
 #include "pxr/usd/ar/resolver.h"
@@ -38,8 +39,132 @@ namespace Diligent
 namespace USD
 {
 
+namespace
+{
+
+class TextureMemoryAllocator final : public IMemoryAllocator
+{
+public:
+    virtual void* Allocate(size_t Size, const Char* dbgDescription, const char* dbgFileName, const Int32 dbgLineNumber) override final
+    {
+        void* Ptr = DefaultRawMemoryAllocator::GetAllocator().Allocate(Size, dbgDescription, dbgFileName, dbgLineNumber);
+        RegisterAllocation(Ptr, Size);
+        return Ptr;
+    }
+
+    virtual void Free(void* Ptr) override final
+    {
+        DefaultRawMemoryAllocator::GetAllocator().Free(Ptr);
+        UnregisterAllocation(Ptr);
+    }
+
+    virtual void* AllocateAligned(size_t Size, size_t Alignment, const Char* dbgDescription, const char* dbgFileName, const Int32 dbgLineNumber) override final
+    {
+        void* Ptr = DefaultRawMemoryAllocator::GetAllocator().AllocateAligned(Size, Alignment, dbgDescription, dbgFileName, dbgLineNumber);
+        RegisterAllocation(Ptr, Size);
+        return Ptr;
+    }
+
+    virtual void FreeAligned(void* Ptr) override final
+    {
+        DefaultRawMemoryAllocator::GetAllocator().FreeAligned(Ptr);
+        UnregisterAllocation(Ptr);
+    }
+
+    static TextureMemoryAllocator& Get()
+    {
+        static TextureMemoryAllocator Allocator;
+        return Allocator;
+    }
+
+    int64_t GetTotalAllocatedSize() const
+    {
+        return m_TotalAllocatedSize.load();
+    }
+
+    ~TextureMemoryAllocator()
+    {
+        DEV_CHECK_ERR(m_Allocations.empty(), "There are ", m_Allocations.size(), " outstanding allocations");
+    }
+
+private:
+    TextureMemoryAllocator()
+    {
+    }
+
+    void RegisterAllocation(const void* Ptr, size_t Size)
+    {
+        if (Ptr == nullptr)
+            return;
+
+        std::lock_guard<std::mutex> Guard{m_AllocationsMtx};
+        m_Allocations.emplace(Ptr, Size);
+        m_TotalAllocatedSize.fetch_add(Size);
+    }
+
+    void UnregisterAllocation(const void* Ptr)
+    {
+        if (Ptr == nullptr)
+            return;
+
+        std::lock_guard<std::mutex> Guard{m_AllocationsMtx};
+        auto                        it = m_Allocations.find(Ptr);
+        if (it != m_Allocations.end())
+        {
+            m_TotalAllocatedSize.fetch_add(-static_cast<int64_t>(it->second));
+            m_Allocations.erase(it);
+        }
+        else
+        {
+            UNEXPECTED("Failed to find allocation");
+        }
+    }
+
+private:
+    std::mutex                              m_AllocationsMtx;
+    std::unordered_map<const void*, size_t> m_Allocations;
+    std::atomic<int64_t>                    m_TotalAllocatedSize{0};
+};
+
+class AssetDataContainer : public RefCntContainer<std::shared_ptr<const char>>
+{
+public:
+    using TBase = RefCntContainer<std::shared_ptr<const char>>;
+
+    AssetDataContainer(IReferenceCounters* pRefCounters, std::shared_ptr<const char> _Data, size_t _Size) :
+        TBase{pRefCounters, std::move(_Data)},
+        Size{_Size}
+    {
+        s_AssetDataSize.fetch_add(static_cast<int64_t>(Size));
+    }
+
+    ~AssetDataContainer()
+    {
+        s_AssetDataSize.fetch_add(-static_cast<int64_t>(Size));
+    }
+
+    static RefCntAutoPtr<AssetDataContainer> Create(std::shared_ptr<const char> Data, size_t Size)
+    {
+        return RefCntAutoPtr<AssetDataContainer>{MakeNewRCObj<AssetDataContainer>()(std::move(Data), Size)};
+    }
+
+    static Int64 GetTotalAllocatedSize()
+    {
+        return s_AssetDataSize.load();
+    }
+
+private:
+    const size_t Size;
+
+    static std::atomic<int64_t> s_AssetDataSize;
+};
+
+std::atomic<int64_t> AssetDataContainer::s_AssetDataSize{0};
+
+} // namespace
+
 RefCntAutoPtr<ITextureLoader> CreateTextureLoaderFromSdfPath(const char*            SdfPath,
-                                                             const TextureLoadInfo& LoadInfo)
+                                                             const TextureLoadInfo& _LoadInfo)
 {
     pxr::ArResolvedPath ResolvedPath{SdfPath};
     if (ResolvedPath.empty())
@@ -55,15 +180,23 @@ RefCntAutoPtr<ITextureLoader> CreateTextureLoaderFromSdfPath(const char*
     if (!Buffer)
         return {};
 
-    RefCntAutoPtr<IObject>   pAssetData = RefCntContainer<std::shared_ptr<const char>>::Create(Buffer);
+    RefCntAutoPtr<IObject>   pAssetData = AssetDataContainer::Create(Buffer, Asset->GetSize());
     RefCntAutoPtr<IDataBlob> pDataBlob  = ProxyDataBlob::Create(Buffer.get(), Asset->GetSize(), pAssetData);
 
+    TextureLoadInfo LoadInfo = _LoadInfo;
+    LoadInfo.pAllocator      = &TextureMemoryAllocator::Get();
+
     RefCntAutoPtr<ITextureLoader> pLoader;
     CreateTextureLoaderFromDataBlob(pDataBlob, LoadInfo, &pLoader);
 
     return pLoader;
 }
 
+Int64 GetTextureLoaderMemoryUsage()
+{
+    return TextureMemoryAllocator::Get().GetTotalAllocatedSize() + AssetDataContainer::GetTotalAllocatedSize();
+}
+
 } // namespace USD
 
 } // namespace Diligent