Fix AV upon merging events from exited thread.

.github/workflows/win_msvc_rel_x64_cmake.yaml

@@ -71,7 +71,7 @@ jobs:
       shell: cmd
       run: |
         cd test
-        cmake . -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF -DCMAKE_TOOLCHAIN_FILE=..\vcpkg\scripts\buildsystems\vcpkg.cmake
+        cmake . -DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF -DGPGMM_FORCE_TRACING=ON -DCMAKE_TOOLCHAIN_FILE=..\vcpkg\scripts\buildsystems\vcpkg.cmake
 
     - name: Build for main branch (with patch)
       shell: cmd

src/gpgmm/common/EventTraceWriter.cpp

@@ -28,6 +28,26 @@
 
 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() {
+            mWriter->FlushAndRemoveBufferEntry(GetBuffer());
+        }
+
+        std::vector<TraceEvent>* GetBuffer() {
+            return &mBuffer;
+        }
+
+      private:
+        EventTraceWriter* mWriter = nullptr;
+        std::vector<TraceEvent> mBuffer;
+    };
+
     EventTraceWriter::EventTraceWriter()
         : mTraceFile(kDefaultTraceFile), mPlatformTime(CreatePlatformTime()) {
     }
@@ -182,25 +202,46 @@ namespace gpgmm {
     }
 
     std::vector<TraceEvent>* EventTraceWriter::GetOrCreateBufferFromTLS() {
-        thread_local std::unique_ptr<std::vector<TraceEvent>> bufferInTLS;
+        thread_local std::unique_ptr<ScopedTraceBufferInTLS> bufferInTLS;
         if (bufferInTLS == nullptr) {
-            bufferInTLS.reset(new std::vector<TraceEvent>());
+            bufferInTLS.reset(new ScopedTraceBufferInTLS(this));
 
             std::lock_guard<std::mutex> mutex(mMutex);
             mBufferPerThread[std::this_thread::get_id()] = bufferInTLS.get();
         }
         ASSERT(bufferInTLS != nullptr);
-        return bufferInTLS.get();
+        return bufferInTLS->GetBuffer();
     }
 
-    std::vector<TraceEvent> EventTraceWriter::MergeAndClearBuffers() const {
+    void EventTraceWriter::FlushAndRemoveBufferEntry(std::vector<TraceEvent>* buffer) {
+        std::lock_guard<std::mutex> mutex(mMutex);
+        const size_t removed = mBufferPerThread.erase(std::this_thread::get_id());
+        ASSERT(removed == 1);
+        mUnmergedBuffer.insert(mUnmergedBuffer.end(), buffer->begin(), buffer->end());
+    }
+
+    std::vector<TraceEvent> EventTraceWriter::MergeAndClearBuffers() {
         std::vector<TraceEvent> mergedBuffer;
+
+        mergedBuffer.insert(mergedBuffer.end(), mUnmergedBuffer.begin(), mUnmergedBuffer.end());
+        mUnmergedBuffer.clear();
+
         for (auto& bufferOfThread : mBufferPerThread) {
-            mergedBuffer.insert(mergedBuffer.end(), bufferOfThread.second->begin(),
-                                bufferOfThread.second->end());
-            bufferOfThread.second->clear();
+            std::vector<TraceEvent>* bufferToMerge = bufferOfThread.second->GetBuffer();
+            mergedBuffer.insert(mergedBuffer.end(), bufferToMerge->begin(), bufferToMerge->end());
+            bufferToMerge->clear();
         }
         return mergedBuffer;
     }
 
+    size_t EventTraceWriter::GetQueuedEventsForTesting() const {
+        std::lock_guard<std::mutex> mutex(mMutex);
+        size_t numOfEvents = 0;
+        numOfEvents += mUnmergedBuffer.size();
+        for (auto& bufferOfThread : mBufferPerThread) {
+            numOfEvents += bufferOfThread.second->GetBuffer()->size();
+        }
+        return numOfEvents;
+    }
+
 }  // namespace gpgmm

src/gpgmm/common/EventTraceWriter.h

@@ -25,37 +25,44 @@
 namespace gpgmm {
 
     class PlatformTime;
+    class ScopedTraceBufferInTLS;
 
     class EventTraceWriter {
       public:
         EventTraceWriter();
+        ~EventTraceWriter();
 
         void SetConfiguration(const std::string& traceFile,
                               const TraceEventPhase& ignoreMask,
                               bool flushOnDestruct);
 
-        ~EventTraceWriter();
-
         void EnqueueTraceEvent(char phase,
                                TraceEventCategory category,
                                const char* name,
                                uint64_t id,
                                uint32_t flags,
                                const JSONDict& args);
+
         void FlushQueuedEventsToDisk();
 
+        void FlushAndRemoveBufferEntry(std::vector<TraceEvent>* buffer);
+
+        size_t GetQueuedEventsForTesting() const;
+
       private:
         std::vector<TraceEvent>* GetOrCreateBufferFromTLS();
-        std::vector<TraceEvent> MergeAndClearBuffers() const;
+        std::vector<TraceEvent> MergeAndClearBuffers();
 
         std::string mTraceFile;
         std::unique_ptr<PlatformTime> mPlatformTime;
-        mutable std::mutex mMutex;
-
-        std::unordered_map<std::thread::id, std::vector<TraceEvent>*> mBufferPerThread;
 
         TraceEventPhase mIgnoreMask;
         bool mFlushOnDestruct = true;
+
+        mutable std::mutex mMutex;
+
+        std::unordered_map<std::thread::id, ScopedTraceBufferInTLS*> mBufferPerThread;
+        std::vector<TraceEvent> mUnmergedBuffer;
     };
 
 }  // namespace gpgmm

src/gpgmm/common/TraceEvent.cpp

@@ -54,6 +54,13 @@ namespace gpgmm {
         return gEventTrace != nullptr;
     }
 
+    size_t GetQueuedEventsForTesting() {
+        if (!IsEventTraceEnabled()) {
+            return 0;
+        }
+        return GetInstance()->GetQueuedEventsForTesting();
+    }
+
     TraceEvent::TraceEvent(char phase,
                            TraceEventCategory category,
                            const std::string& name,

src/gpgmm/common/TraceEvent.h

@@ -39,7 +39,8 @@
 #ifdef GPGMM_DISABLE_TRACING
 
 #define TRACE_EVENT0(category_group, name) TRACE_EMPTY
-#define TRACE_EVENT_INSTANT0(category_group, name, scope, args) TRACE_EMPTY
+#define TRACE_EVENT_INSTANT0(category_group, name) TRACE_EMPTY
+#define TRACE_EVENT_INSTANT1(category_group, name, args) TRACE_EMPTY
 #define TRACE_COUNTER1(category_group, name, value) TRACE_EMPTY
 #define TRACE_EVENT_METADATA1(category_group, name, arg1_name, arg1_value) TRACE_EMPTY
 
@@ -67,6 +68,9 @@ const uint64_t kNoId = 0;
     INTERNAL_TRACE_EVENT_ADD(TRACE_EVENT_PHASE_COUNTER, category_group, name, TRACE_EVENT_FLAG_NONE, "value", \
                              static_cast<int>(value))
 
+#define TRACE_EVENT_INSTANT0(category_group, name) \
+    INTERNAL_TRACE_EVENT_ADD(TRACE_EVENT_PHASE_INSTANT, category_group, name, TRACE_EVENT_FLAG_NONE)
+
 #define TRACE_EVENT_INSTANT1(category_group, name, args) \
     INTERNAL_TRACE_EVENT_ADD(TRACE_EVENT_PHASE_INSTANT, category_group, name, TRACE_EVENT_FLAG_NONE, args)
 
@@ -179,6 +183,8 @@ namespace gpgmm {
 
     bool IsEventTraceEnabled();
 
+    size_t GetQueuedEventsForTesting();
+
     class TraceEventID {
       public:
         explicit TraceEventID(const void* id)

src/tests/BUILD.gn

@@ -107,6 +107,7 @@ test("gpgmm_unittests") {
     "unittests/BuddyBlockAllocatorTests.cpp",
     "unittests/BuddyMemoryAllocatorTests.cpp",
     "unittests/ConditionalMemoryAllocatorTests.cpp",
+    "unittests/EventTraceWriterTests.cpp",
     "unittests/FlagsTests.cpp",
     "unittests/LinkedListTests.cpp",
     "unittests/MathTests.cpp",

src/tests/CMakeLists.txt

@@ -23,22 +23,23 @@ target_link_libraries(gpgmm_unittests PRIVATE
 )
 
 target_sources(gpgmm_unittests PRIVATE
-    "DummyMemoryAllocator.h"
-    "unittests/BuddyBlockAllocatorTests.cpp"
-    "unittests/ConditionalMemoryAllocatorTests.cpp"
-    "unittests/FlagsTests.cpp"
-    "unittests/LinkedListTests.cpp"
-    "unittests/MathTests.cpp"
-    "unittests/MemoryAllocatorTests.cpp"
-    "unittests/MemoryCacheTests.cpp"
-    "unittests/MemoryPoolTests.cpp"
-    "unittests/PooledMemoryAllocatorTests.cpp"
-    "unittests/RefCountTests.cpp"
-    "unittests/SegmentedMemoryAllocatorTests.cpp"
-    "unittests/SlabBlockAllocatorTests.cpp"
-    "unittests/SlabMemoryAllocatorTests.cpp"
-    "unittests/UtilsTest.cpp"
-    "UnittestsMain.cpp"
+  "DummyMemoryAllocator.h"
+  "unittests/BuddyBlockAllocatorTests.cpp"
+  "unittests/ConditionalMemoryAllocatorTests.cpp"
+  "unittests/EventTraceWriterTests.cpp"
+  "unittests/FlagsTests.cpp"
+  "unittests/LinkedListTests.cpp"
+  "unittests/MathTests.cpp"
+  "unittests/MemoryAllocatorTests.cpp"
+  "unittests/MemoryCacheTests.cpp"
+  "unittests/MemoryPoolTests.cpp"
+  "unittests/PooledMemoryAllocatorTests.cpp"
+  "unittests/RefCountTests.cpp"
+  "unittests/SegmentedMemoryAllocatorTests.cpp"
+  "unittests/SlabBlockAllocatorTests.cpp"
+  "unittests/SlabMemoryAllocatorTests.cpp"
+  "unittests/UtilsTest.cpp"
+  "UnittestsMain.cpp"
 )
 
 target_link_libraries(gpgmm_unittests PRIVATE

src/tests/unittests/EventTraceWriterTests.cpp

@@ -0,0 +1,61 @@
+// Copyright 2021 The GPGMM Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <gtest/gtest.h>
+
+#include "gpgmm/common/TraceEvent.h"
+
+#include <thread>
+#include <vector>
+
+static constexpr const char* kDummyTrace = "DummyTrace.json";
+
+using namespace gpgmm;
+
+class EventTraceWriterTests : public testing::Test {
+  public:
+    void SetUp() override {
+        StartupEventTrace(kDummyTrace, TraceEventPhase::None, /*flushOnDestruct*/ true);
+    }
+
+    void TearDown() override {
+        ShutdownEventTrace();
+    }
+};
+
+TEST_F(EventTraceWriterTests, SingleThreadWrites) {
+    constexpr uint32_t kEventCount = 64u;
+    for (size_t i = 0; i < kEventCount; i++) {
+        TRACE_EVENT_INSTANT0(TraceEventCategory::Default, "InstantEvent");
+    }
+
+    // 1 event per thread + 1 metadata event for main thread name.
+    EXPECT_EQ(GetQueuedEventsForTesting(), 64 + 1u);
+}
+
+TEST_F(EventTraceWriterTests, MultiThreadWrites) {
+    constexpr uint32_t kThreadCount = 64u;
+    std::vector<std::thread> threads(kThreadCount);
+    for (size_t threadIdx = 0; threadIdx < threads.size(); threadIdx++) {
+        threads[threadIdx] = std::thread(
+            [&]() { TRACE_EVENT_INSTANT0(TraceEventCategory::Default, "InstantEvent"); });
+    }
+
+    for (std::thread& thread : threads) {
+        thread.join();
+    }
+
+    // 1 event per thread + 1 metadata event for main thread name.
+    EXPECT_EQ(GetQueuedEventsForTesting(), 64 + 1u);
+}