intel · bbernhar · Jul 7, 2022 · Jul 7, 2022
diff --git a/src/gpgmm/common/EventTraceWriter.cpp b/src/gpgmm/common/EventTraceWriter.cpp
@@ -31,20 +31,36 @@ namespace gpgmm {
     // Trace buffer that flushes and unlinks itself from the cache once destroyed.
     class ScopedTraceBufferInTLS {
       public:
-        ScopedTraceBufferInTLS(EventTraceWriter* writer) : mWriter(writer) {
-            ASSERT(writer != nullptr);
+        ScopedTraceBufferInTLS(std::shared_ptr<EventTraceWriter> writer)
+            : mWriter(std::move(writer)) {
+            ASSERT(mWriter != nullptr);
         }
 
         ~ScopedTraceBufferInTLS() {
-            mWriter->FlushAndRemoveBufferEntry(GetBuffer());
+            mWriter->FlushAndRemoveBufferEntry(&mBuffer);
         }
 
-        std::vector<TraceEvent>* GetBuffer() {
-            return &mBuffer;
+        void AddEvent(const TraceEvent& event) {
+            std::unique_lock<std::mutex> lock(mMutex);
+            mBuffer.push_back(event);
+        }
+
+        std::vector<TraceEvent> GetAndClearBuffer() {
+            std::unique_lock<std::mutex> lock(mMutex);
+            std::vector<TraceEvent> tmp = mBuffer;
+            mBuffer.clear();
+            return tmp;
+        }
+
+        size_t GetBufferSize() const {
+            std::unique_lock<std::mutex> lock(mMutex);
+            return mBuffer.size();
         }
 
       private:
-        EventTraceWriter* mWriter = nullptr;
+        std::shared_ptr<EventTraceWriter> mWriter;
+
+        mutable std::mutex mMutex;  // Protect access for members below.
         std::vector<TraceEvent> mBuffer;
     };
 
@@ -53,20 +69,17 @@ namespace gpgmm {
     }
 
     void EventTraceWriter::SetConfiguration(const std::string& traceFile,
-                                            const TraceEventPhase& ignoreMask,
-                                            bool flushOnDestruct) {
+                                            const TraceEventPhase& ignoreMask) {
         mTraceFile = (traceFile.empty()) ? mTraceFile : traceFile;
         mIgnoreMask = ignoreMask;
-        mFlushOnDestruct = flushOnDestruct;
     }
 
     EventTraceWriter::~EventTraceWriter() {
-        if (mFlushOnDestruct) {
-            FlushQueuedEventsToDisk();
-        }
+        FlushQueuedEventsToDisk();
     }
 
-    void EventTraceWriter::EnqueueTraceEvent(char phase,
+    void EventTraceWriter::EnqueueTraceEvent(std::shared_ptr<EventTraceWriter> writer,
+                                             char phase,
                                              TraceEventCategory category,
                                              const char* name,
                                              uint64_t id,
@@ -75,8 +88,8 @@ namespace gpgmm {
         const double timestampInSeconds = mPlatformTime->GetRelativeTime();
         const uint32_t threadID = std::stoi(ToString(std::this_thread::get_id()));
         if (timestampInSeconds != 0) {
-            GetOrCreateBufferFromTLS()->push_back(
-                {phase, category, name, id, threadID, timestampInSeconds, flags, args});
+            GetOrCreateBufferFromTLS(std::move(writer))
+                ->AddEvent({phase, category, name, id, threadID, timestampInSeconds, flags, args});
         }
     }
 
@@ -201,16 +214,17 @@ namespace gpgmm {
         DebugLog() << "Flushed " << mergedBuffer.size() << " events to disk.";
     }
 
-    std::vector<TraceEvent>* EventTraceWriter::GetOrCreateBufferFromTLS() {
+    ScopedTraceBufferInTLS* EventTraceWriter::GetOrCreateBufferFromTLS(
+        std::shared_ptr<EventTraceWriter> writer) {
         thread_local std::unique_ptr<ScopedTraceBufferInTLS> bufferInTLS;
         if (bufferInTLS == nullptr) {
-            bufferInTLS.reset(new ScopedTraceBufferInTLS(this));
+            bufferInTLS.reset(new ScopedTraceBufferInTLS(std::move(writer)));
 
             std::lock_guard<std::mutex> mutex(mMutex);
             mBufferPerThread[std::this_thread::get_id()] = bufferInTLS.get();
         }
         ASSERT(bufferInTLS != nullptr);
-        return bufferInTLS->GetBuffer();
+        return bufferInTLS.get();
     }
 
     void EventTraceWriter::FlushAndRemoveBufferEntry(std::vector<TraceEvent>* buffer) {
@@ -227,9 +241,9 @@ namespace gpgmm {
         mUnmergedBuffer.clear();
 
         for (auto& bufferOfThread : mBufferPerThread) {
-            std::vector<TraceEvent>* bufferToMerge = bufferOfThread.second->GetBuffer();
-            mergedBuffer.insert(mergedBuffer.end(), bufferToMerge->begin(), bufferToMerge->end());
-            bufferToMerge->clear();
+            std::vector<TraceEvent> bufferToMerge = bufferOfThread.second->GetAndClearBuffer();
+            mergedBuffer.insert(mergedBuffer.end(), bufferToMerge.begin(), bufferToMerge.end());
+            bufferToMerge.clear();
         }
         return mergedBuffer;
     }
@@ -239,7 +253,7 @@ namespace gpgmm {
         size_t numOfEvents = 0;
         numOfEvents += mUnmergedBuffer.size();
         for (auto& bufferOfThread : mBufferPerThread) {
-            numOfEvents += bufferOfThread.second->GetBuffer()->size();
+            numOfEvents += bufferOfThread.second->GetBufferSize();
         }
         return numOfEvents;
     }

diff --git a/src/gpgmm/common/EventTraceWriter.h b/src/gpgmm/common/EventTraceWriter.h
@@ -32,11 +32,10 @@ namespace gpgmm {
         EventTraceWriter();
         ~EventTraceWriter();
 
-        void SetConfiguration(const std::string& traceFile,
-                              const TraceEventPhase& ignoreMask,
-                              bool flushOnDestruct);
+        void SetConfiguration(const std::string& traceFile, const TraceEventPhase& ignoreMask);
 
-        void EnqueueTraceEvent(char phase,
+        void EnqueueTraceEvent(std::shared_ptr<EventTraceWriter> writer,
+                               char phase,
                                TraceEventCategory category,
                                const char* name,
                                uint64_t id,
@@ -50,14 +49,13 @@ namespace gpgmm {
         size_t GetQueuedEventsForTesting() const;
 
       private:
-        std::vector<TraceEvent>* GetOrCreateBufferFromTLS();
+        ScopedTraceBufferInTLS* GetOrCreateBufferFromTLS(std::shared_ptr<EventTraceWriter> writer);
         std::vector<TraceEvent> MergeAndClearBuffers();
 
         std::string mTraceFile;
-        std::unique_ptr<PlatformTime> mPlatformTime;
-
         TraceEventPhase mIgnoreMask;
-        bool mFlushOnDestruct = true;
+
+        std::unique_ptr<PlatformTime> mPlatformTime;
 
         mutable std::mutex mMutex;
 

diff --git a/src/gpgmm/common/TraceEvent.cpp b/src/gpgmm/common/TraceEvent.cpp
@@ -22,26 +22,24 @@
 
 namespace gpgmm {
 
-    static std::unique_ptr<EventTraceWriter> gEventTrace;
+    static std::shared_ptr<EventTraceWriter> gEventTrace;
     static std::mutex mMutex;
 
     static EventTraceWriter* GetInstance() {
         std::lock_guard<std::mutex> lock(mMutex);
         if (gEventTrace == nullptr) {
-            gEventTrace = std::make_unique<EventTraceWriter>();
+            gEventTrace = std::make_shared<EventTraceWriter>();
         }
         return gEventTrace.get();
     }
 
-    void StartupEventTrace(const std::string& traceFile,
-                           const TraceEventPhase& ignoreMask,
-                           bool flushOnDestruct) {
+    void StartupEventTrace(const std::string& traceFile, const TraceEventPhase& ignoreMask) {
 #if defined(GPGMM_DISABLE_TRACING)
         gpgmm::WarningLog()
             << "Event tracing enabled but unable to record due to GPGMM_DISABLE_TRACING.";
 #endif
 
-        GetInstance()->SetConfiguration(traceFile, ignoreMask, flushOnDestruct);
+        GetInstance()->SetConfiguration(traceFile, ignoreMask);
         TRACE_EVENT_METADATA1(TraceEventCategory::Metadata, "thread_name", "name",
                               "GPGMM_MainThread");
     }
@@ -90,7 +88,7 @@ namespace gpgmm {
                                     uint32_t flags,
                                     const JSONDict& args) {
         if (IsEventTraceEnabled()) {
-            GetInstance()->EnqueueTraceEvent(phase, category, name, id, flags, args);
+            GetInstance()->EnqueueTraceEvent(gEventTrace, phase, category, name, id, flags, args);
         }
     }
 }  // namespace gpgmm
diff --git a/src/gpgmm/common/TraceEvent.h b/src/gpgmm/common/TraceEvent.h
@@ -176,8 +176,7 @@ namespace gpgmm {
     class PlatformTime;
 
     void StartupEventTrace(const std::string& traceFile,
-                           const TraceEventPhase& ignoreMask,
-                           bool flushOnDestruct);
+                           const TraceEventPhase& ignoreMask);
 
     void FlushEventTraceToDisk();
 

diff --git a/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp b/src/gpgmm/d3d12/ResidencyManagerD3D12.cpp
@@ -164,8 +164,7 @@ namespace gpgmm::d3d12 {
 
         if (descriptor.RecordOptions.Flags != EVENT_RECORD_FLAG_NONE) {
             StartupEventTrace(descriptor.RecordOptions.TraceFile,
-                              static_cast<TraceEventPhase>(~descriptor.RecordOptions.Flags | 0),
-                              descriptor.RecordOptions.EventScope & EVENT_RECORD_SCOPE_PER_PROCESS);
+                              static_cast<TraceEventPhase>(~descriptor.RecordOptions.Flags | 0));
 
             SetEventMessageLevel(GetLogSeverity(descriptor.RecordOptions.MinMessageLevel));
         }

diff --git a/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp b/src/gpgmm/d3d12/ResourceAllocatorD3D12.cpp
@@ -342,10 +342,8 @@ namespace gpgmm::d3d12 {
 
         if (pResidencyManager == nullptr &&
             newDescriptor.RecordOptions.Flags != EVENT_RECORD_FLAG_NONE) {
-            StartupEventTrace(
-                allocatorDescriptor.RecordOptions.TraceFile,
-                static_cast<TraceEventPhase>(~newDescriptor.RecordOptions.Flags | 0),
-                allocatorDescriptor.RecordOptions.EventScope & EVENT_RECORD_SCOPE_PER_PROCESS);
+            StartupEventTrace(allocatorDescriptor.RecordOptions.TraceFile,
+                              static_cast<TraceEventPhase>(~newDescriptor.RecordOptions.Flags | 0));
 
             SetEventMessageLevel(GetLogSeverity(newDescriptor.RecordOptions.MinMessageLevel));
         } else {

diff --git a/src/tests/unittests/EventTraceWriterTests.cpp b/src/tests/unittests/EventTraceWriterTests.cpp
@@ -26,7 +26,7 @@ using namespace gpgmm;
 class EventTraceWriterTests : public testing::Test {
   public:
     void SetUp() override {
-        StartupEventTrace(kDummyTrace, TraceEventPhase::None, /*flushOnDestruct*/ true);
+        StartupEventTrace(kDummyTrace, TraceEventPhase::None);
     }
 
     void TearDown() override {