From c0819af0a298efb258d2bcbd66645bc9c3ef6c26 Mon Sep 17 00:00:00 2001
From: "Bernhart, Bryan" <bryan.bernhart@intel.com>
Date: Thu, 6 Apr 2023 23:13:14 +0000
Subject: [PATCH] Fixes to enable UE5.1 CitySample

---
 src/gpgmm/common/EventTraceWriter.cpp | 14 ++++++++++----
 src/gpgmm/common/TraceEvent.cpp       | 23 +++++++++++++++++------
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/src/gpgmm/common/EventTraceWriter.cpp b/src/gpgmm/common/EventTraceWriter.cpp
index 2b0ce9b0b..58bd6ed7d 100644
--- a/src/gpgmm/common/EventTraceWriter.cpp
+++ b/src/gpgmm/common/EventTraceWriter.cpp
@@ -218,13 +218,19 @@ namespace gpgmm {
 
     ScopedTraceBufferInTLS* EventTraceWriter::GetOrCreateBufferFromTLS(
         std::shared_ptr<EventTraceWriter> writer) {
-        thread_local std::unique_ptr<ScopedTraceBufferInTLS> bufferInTLS;
-        if (bufferInTLS == nullptr) {
-            bufferInTLS.reset(new ScopedTraceBufferInTLS(std::move(writer)));
+        std::lock_guard<std::mutex> mutex(mMutex);
 
-            std::lock_guard<std::mutex> mutex(mMutex);
+        // Prevent |bufferInTLS| from being accessed after destruction by checking if it was removed
+        // from the thread id cache then always relinquishing ownership before creation to avoid a
+        // double delete. This was because the thread storage for |bufferInTLS| was found to be
+        // reused and never re-initialized to nullptr.
+        thread_local std::unique_ptr<ScopedTraceBufferInTLS> bufferInTLS;
+        if (mBufferPerThread.find(std::this_thread::get_id()) == mBufferPerThread.end()) {
+            bufferInTLS.release();
+            bufferInTLS = std::make_unique<ScopedTraceBufferInTLS>(std::move(writer));
             mBufferPerThread[std::this_thread::get_id()] = bufferInTLS.get();
         }
+
         ASSERT(bufferInTLS != nullptr);
         return bufferInTLS.get();
     }
diff --git a/src/gpgmm/common/TraceEvent.cpp b/src/gpgmm/common/TraceEvent.cpp
index 106be1038..fc8b394d6 100644
--- a/src/gpgmm/common/TraceEvent.cpp
+++ b/src/gpgmm/common/TraceEvent.cpp
@@ -22,15 +22,25 @@
 
 namespace gpgmm {
 
-    static std::shared_ptr<EventTraceWriter> gEventTrace;
+    static std::shared_ptr<EventTraceWriter> gEventTraceWriter;
     static std::mutex mMutex;
 
+    void OnFlushAtExit() {
+        std::lock_guard<std::mutex> lock(mMutex);
+        if (gEventTraceWriter != nullptr) {
+            gEventTraceWriter->FlushQueuedEventsToDisk();
+        }
+    }
+
     static EventTraceWriter* GetInstance() {
         std::lock_guard<std::mutex> lock(mMutex);
-        if (gEventTrace == nullptr) {
-            gEventTrace = std::make_shared<EventTraceWriter>();
+        if (gEventTraceWriter == nullptr) {
+            gEventTraceWriter = std::make_shared<EventTraceWriter>();
+
+            // Unmerged events never merge if all threads using the writer never terminate.
+            std::atexit(OnFlushAtExit);
         }
-        return gEventTrace.get();
+        return gEventTraceWriter.get();
     }
 
     void StartupEventTrace(const char* traceFile, const TraceEventPhase& ignoreMask) {
@@ -53,7 +63,7 @@ namespace gpgmm {
 
     bool IsEventTraceEnabled() {
         std::lock_guard<std::mutex> lock(mMutex);
-        return gEventTrace != nullptr;
+        return gEventTraceWriter != nullptr;
     }
 
     size_t GetQueuedEventsForTesting() {
@@ -88,7 +98,8 @@ namespace gpgmm {
                                     uint32_t flags,
                                     const JSONDict& args) {
         if (IsEventTraceEnabled()) {
-            GetInstance()->EnqueueTraceEvent(gEventTrace, phase, category, name, id, flags, args);
+            GetInstance()->EnqueueTraceEvent(gEventTraceWriter, phase, category, name, id, flags,
+                                             args);
         }
     }
 }  // namespace gpgmm