From ea9b1dcdc87491f8b0f96a8b1cb04e650e15a936 Mon Sep 17 00:00:00 2001
From: assiduous <assiduous@diligentgraphics.com>
Date: Fri, 17 Jan 2025 17:42:19 -0800
Subject: [PATCH] HnGeometryPool: reworked space reserving

---
 Hydrogent/include/HnGeometryPool.hpp | 28 +++++++++++++++++--
 Hydrogent/src/HnGeometryPool.cpp     | 40 +++++++++++++++++++++++++---
 Hydrogent/src/HnMesh.cpp             | 26 +++---------------
 3 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/Hydrogent/include/HnGeometryPool.hpp b/Hydrogent/include/HnGeometryPool.hpp
index 5ae3318c..85b7955f 100644
--- a/Hydrogent/include/HnGeometryPool.hpp
+++ b/Hydrogent/include/HnGeometryPool.hpp
@@ -97,8 +97,32 @@ class HnGeometryPool final
     Int64 GetPendingVertexDataSize() const { return m_PendingVertexDataSize.load(); }
     Int64 GetPendingIndexDataSize() const { return m_PendingIndexDataSize.load(); }
 
-    Int64 ReserveDataSize(Uint64 Size);
-    void  ReleaseReservedDataSize(Uint64 Size);
+    class ReservedSpace
+    {
+    public:
+        ReservedSpace(ReservedSpace&& Other) noexcept;
+
+        ReservedSpace& operator=(const ReservedSpace&) = delete;
+        ReservedSpace& operator=(ReservedSpace&&) = delete;
+
+        ~ReservedSpace();
+        void Release();
+
+        Uint64 GetTotalPendingSize() const { return m_TotalPendingSize; }
+
+    private:
+        ReservedSpace(HnGeometryPool& Pool, Uint64 Size, Uint64 TotalPendingSize) noexcept;
+
+    private:
+        friend class HnGeometryPool;
+        HnGeometryPool& m_Pool;
+        Uint64          m_Size             = 0;
+        Uint64          m_TotalPendingSize = 0;
+    };
+    ReservedSpace ReserveSpace(Uint64 Size);
+
+private:
+    void ReleaseReservedSpace(Uint64 Size);
 
 private:
     RefCntAutoPtr<IRenderDevice> m_pDevice;
diff --git a/Hydrogent/src/HnGeometryPool.cpp b/Hydrogent/src/HnGeometryPool.cpp
index 5181d6e2..9a5688fd 100644
--- a/Hydrogent/src/HnGeometryPool.cpp
+++ b/Hydrogent/src/HnGeometryPool.cpp
@@ -1017,15 +1017,49 @@ void HnGeometryPool::Commit(IDeviceContext* pContext)
     }
 }
 
-Int64 HnGeometryPool::ReserveDataSize(Uint64 Size)
+HnGeometryPool::ReservedSpace::ReservedSpace(HnGeometryPool& Pool,
+                                             Uint64          Size,
+                                             Uint64          TotalPendingSize) noexcept :
+    m_Pool{Pool},
+    m_Size{Size},
+    m_TotalPendingSize{TotalPendingSize}
 {
-    return GetPendingVertexDataSize() +
+}
+
+HnGeometryPool::ReservedSpace::ReservedSpace(ReservedSpace&& Other) noexcept :
+    m_Pool{Other.m_Pool},
+    m_Size{Other.m_Size},
+    m_TotalPendingSize{Other.m_TotalPendingSize}
+{
+    Other.m_Size             = 0;
+    Other.m_TotalPendingSize = 0;
+}
+
+void HnGeometryPool::ReservedSpace::Release()
+{
+    if (m_Size != 0)
+    {
+        m_Pool.ReleaseReservedSpace(m_Size);
+        m_Size = 0;
+    }
+}
+
+HnGeometryPool::ReservedSpace::~ReservedSpace()
+{
+    Release();
+}
+
+HnGeometryPool::ReservedSpace HnGeometryPool::ReserveSpace(Uint64 Size)
+{
+    const Int64 TotalPendingSize =
+        GetPendingVertexDataSize() +
         GetPendingIndexDataSize() +
         m_ReservedDataSize.fetch_add(static_cast<Int64>(Size)) +
         static_cast<Int64>(Size);
+    return {*this, Size, static_cast<Uint64>(TotalPendingSize)};
 }
 
-void HnGeometryPool::ReleaseReservedDataSize(Uint64 Size)
+void HnGeometryPool::ReleaseReservedSpace(Uint64 Size)
 {
     VERIFY_EXPR(m_ReservedDataSize.load() >= static_cast<Int64>(Size));
     m_ReservedDataSize.fetch_add(-static_cast<Int64>(Size));
diff --git a/Hydrogent/src/HnMesh.cpp b/Hydrogent/src/HnMesh.cpp
index 5d2a4ce2..5d636f8f 100644
--- a/Hydrogent/src/HnMesh.cpp
+++ b/Hydrogent/src/HnMesh.cpp
@@ -412,30 +412,10 @@ bool HnMesh::UpdateRepr(pxr::HdSceneDelegate& SceneDelegate,
         ExpectedGeometryDataSize += MeshUtils.GetTotalIndexCount(IndexCountFlags) * sizeof(Uint32);
     }
 
-    class GeometryPoolReservedDataGuard
+    HnGeometryPool::ReservedSpace ReservedSpace = GeometryPool.ReserveSpace(ExpectedGeometryDataSize);
+    if (GeometryLoadBudget > 0 && ReservedSpace.GetTotalPendingSize() > GeometryLoadBudget)
     {
-    public:
-        GeometryPoolReservedDataGuard(HnGeometryPool& Pool, Uint64 Size) :
-            m_Pool{Pool},
-            m_Size{Size},
-            m_TotalSize{static_cast<Uint64>(Pool.ReserveDataSize(Size))}
-        {
-        }
-        ~GeometryPoolReservedDataGuard()
-        {
-            m_Pool.ReleaseReservedDataSize(m_Size);
-        }
-        Uint64 GetTotalSize() const { return m_TotalSize; }
-
-    private:
-        HnGeometryPool& m_Pool;
-        const Uint64    m_Size;
-        const Uint64    m_TotalSize;
-    };
-    GeometryPoolReservedDataGuard ReservedDataGuard{GeometryPool, ExpectedGeometryDataSize};
-    if (GeometryLoadBudget > 0 && ReservedDataGuard.GetTotalSize() > GeometryLoadBudget)
-    {
-        if (ReservedDataGuard.GetTotalSize() == ExpectedGeometryDataSize)
+        if (ReservedSpace.GetTotalPendingSize() == ExpectedGeometryDataSize)
         {
             LOG_WARNING_MESSAGE("Syncing rprim ", Id, " requires ", ExpectedGeometryDataSize,
                                 " bytes of geometry data, which exceeds the budget ", GeometryLoadBudget, ".");