GPUUploadManager: fix potential issue with stale pages in ScheduleBufferUpdate

Author: TheMostDiligent

Graphics/GraphicsTools/include/GPUUploadManagerImpl.hpp

@@ -148,12 +148,12 @@ class GPUUploadManagerImpl final : public ObjectBase<IGPUUploadManager>
         // Returns the number of pending operations.
         size_t GetNumPendingOps() const { return m_NumPendingOps.load(std::memory_order_acquire); }
 
+        // Returns true if the page is sealed for new writes.
+        bool IsSealed() const { return (m_State.load(std::memory_order_acquire) & SEALED_BIT) != 0; }
+
 #ifdef DILIGENT_DEBUG
         // Returns the number of active writers. This is used for testing and debugging purposes.
         Uint32 DbgGetWriterCount() const { return m_State.load(std::memory_order_acquire) & WRITER_MASK; }
-
-        // Returns true if the page is sealed for new writes. This is used for testing and debugging purposes.
-        bool DbgIsSealed() const { return (m_State.load(std::memory_order_acquire) & SEALED_BIT) != 0; }
 #endif
 
         void ReleaseStagingBuffer(IDeviceContext* pContext);

Graphics/GraphicsTools/interface/GPUUploadManager.h

@@ -135,6 +135,8 @@ DILIGENT_BEGIN_INTERFACE(IGPUUploadManager, IObject)
                                               void*                         pCallbackData DEFAULT_VALUE(nullptr)) PURE;
 
     /// Retrieves GPU upload manager statistics.
+    ///
+    /// The method must not be called concurrently with RenderThreadUpdate().
     VIRTUAL void METHOD(GetStats)(THIS_
                                   GPUUploadManagerStats REF Stats) CONST PURE;
 };

Graphics/GraphicsTools/src/GPUUploadManagerImpl.cpp

@@ -234,7 +234,7 @@ bool GPUUploadManagerImpl::Page::ScheduleBufferUpdate(IBuffer*
 
 void GPUUploadManagerImpl::Page::ExecutePendingOps(IDeviceContext* pContext, Uint64 FenceValue)
 {
-    VERIFY(DbgIsSealed(), "Page must be sealed before executing pending operations");
+    VERIFY(IsSealed(), "Page must be sealed before executing pending operations");
     VERIFY(DbgGetWriterCount() == 0, "All writers must finish before executing pending operations");
 
     if (m_pData != nullptr && !m_PersistentMapped)
@@ -284,7 +284,7 @@ void GPUUploadManagerImpl::Page::Reset(IDeviceContext* pContext)
 
 bool GPUUploadManagerImpl::Page::TryEnqueue()
 {
-    VERIFY(DbgIsSealed(), "Page must be sealed before it can be enqueued for execution");
+    VERIFY(IsSealed(), "Page must be sealed before it can be enqueued for execution");
     VERIFY(DbgGetWriterCount() == 0, "All writers must finish before the page can be enqueued");
 
     bool Expected = false;
@@ -412,14 +412,38 @@ void GPUUploadManagerImpl::ScheduleBufferUpdate(IDeviceContext*               pC
             continue;
         }
 
-        const bool UpdateScheduled = Writer.ScheduleBufferUpdate(pDstBuffer, DstOffset, NumBytes, pSrcData, Callback, pCallbackData);
-        if (Writer.EndWriting() == Page::WritingStatus::LastWriterSealed)
+        auto EndWriting = [&]() {
+            if (Writer.EndWriting() == Page::WritingStatus::LastWriterSealed)
+            {
+                // We were the last writer - enqueue the page whether the update was successfully
+                // scheduled or not, because the page is sealed and can't be written to anymore.
+                TryEnqueuePage(P);
+            }
+        };
+
+        // Prevent ABA-style bug when pages are reset and reused
+        // * Thread T1 loads P0 as current, then gets descheduled before TryBeginWriting().
+        // * Render thread swaps current to P1, seals P0, sees it has 0 ops, calls Reset(nullptr)
+        //   which clears SEALED_BIT, and pushes P0 to FreePages
+        // * T1 resumes and calls P0->TryBeginWriting() and succeeds (because SEALED_BIT is now cleared).
+        // * T1 schedules an update into a page that is not current and is simultaneously sitting in the free list,
+        //   potentially being popped/installed elsewhere.
         {
-            // We were the last writer - enqueue the page whether the update was successfully scheduled or not,
-            // because the page is sealed and can't be written to anymore.
-            TryEnqueuePage(P);
+            Page* CurNow = m_pCurrentPage.load(std::memory_order_acquire);
+            // Validate that the page is either:
+            // * still current OR
+            // * already sealed (meaning it got rotated after we began writing, which is fine because writer count protects it)
+            if (P != CurNow && !P->IsSealed())
+            {
+                EndWriting();
+                // Reload current and retry
+                continue;
+            }
         }
 
+        const bool UpdateScheduled = Writer.ScheduleBufferUpdate(pDstBuffer, DstOffset, NumBytes, pSrcData, Callback, pCallbackData);
+        EndWriting();
+
         if (UpdateScheduled)
         {
             break;
@@ -492,7 +516,7 @@ bool GPUUploadManagerImpl::TryRotatePage(IDeviceContext* pContext, Page* Expecte
 
 bool GPUUploadManagerImpl::TryEnqueuePage(Page* P)
 {
-    VERIFY_EXPR(P->DbgIsSealed());
+    VERIFY_EXPR(P->IsSealed());
     if (P->TryEnqueue())
     {
         if (P->GetNumPendingOps() > 0)

Tests/DiligentCoreTest/src/GraphicsTools/GPUUploadManagerPageTest.cpp

@@ -166,7 +166,7 @@ TEST(GPUUploadManagerPageTest, ParallelTryBeginWriting)
     LOG_INFO_MESSAGE("Total writes: ", TotalWrites.load(), " out of ", kNumThreads * kNumIterations);
 
     EXPECT_EQ(NumPagesAlreadySealed.load(), kNumThreads - 1) << "Only one thread should be able to seal the page";
-    EXPECT_EQ(NumLastWriters.load(), 1) << "Only one thread should be the last writer";
+    EXPECT_EQ(NumLastWriters.load(), 1u) << "Only one thread should be the last writer";
 }