From a69bcfcaf69f9869bd73afa5dceac90326dc3af2 Mon Sep 17 00:00:00 2001
From: Calvin Neo <CalvinNeo@users.noreply.github.com>
Date: Thu, 1 Aug 2024 14:11:26 +0800
Subject: [PATCH] *: Support generating adhoc heap profiling in svg format
 (#9273) (#252)

Signed-off-by: Calvin Neo <calvinneo1995@gmail.com>
---
 .../KVStore/tests/gtest_async_tasks.cpp       |  2 +-
 .../V3/Universal/UniversalWriteBatchImpl.h    | 10 +++++
 .../tests/gtest_universal_page_storage.cpp    | 44 +++++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/dbms/src/Storages/KVStore/tests/gtest_async_tasks.cpp b/dbms/src/Storages/KVStore/tests/gtest_async_tasks.cpp
index c370da390a2..464603b779c 100644
--- a/dbms/src/Storages/KVStore/tests/gtest_async_tasks.cpp
+++ b/dbms/src/Storages/KVStore/tests/gtest_async_tasks.cpp
@@ -228,7 +228,7 @@ TEST(AsyncTasksTest, AsyncTasksCommon)
     auto async_tasks = std::make_unique<TestAsyncTasks>(1, 1, 2);
 
     int total = 5;
-    int max_steps = 10;
+    int max_steps = 15;
     int current_step = 0;
     std::vector<bool> f(total, false);
     std::vector<bool> s(total, false);
diff --git a/dbms/src/Storages/Page/V3/Universal/UniversalWriteBatchImpl.h b/dbms/src/Storages/Page/V3/Universal/UniversalWriteBatchImpl.h
index 3ce03aa67e8..569a8993b57 100644
--- a/dbms/src/Storages/Page/V3/Universal/UniversalWriteBatchImpl.h
+++ b/dbms/src/Storages/Page/V3/Universal/UniversalWriteBatchImpl.h
@@ -28,6 +28,11 @@ namespace ErrorCodes
 extern const int LOGICAL_ERROR;
 } // namespace ErrorCodes
 
+namespace tests
+{
+class UniPageStorageIdTest;
+}
+
 class UniversalWriteBatch : private boost::noncopyable
 {
 private:
@@ -266,6 +271,7 @@ class UniversalWriteBatch : private boost::noncopyable
     void merge(UniversalWriteBatch & rhs)
     {
         writes.reserve(writes.size() + rhs.writes.size());
+        PS::PageStorageMemorySummary::universal_write_count.fetch_add(rhs.writes.size());
         for (const auto & r : rhs.writes)
         {
             writes.emplace_back(r);
@@ -276,6 +282,7 @@ class UniversalWriteBatch : private boost::noncopyable
 
     [[clang::reinitializes]] void clear()
     {
+        PS::PageStorageMemorySummary::universal_write_count.fetch_sub(writes.size());
         Writes tmp;
         writes.swap(tmp);
         total_data_size = 0;
@@ -302,6 +309,9 @@ class UniversalWriteBatch : private boost::noncopyable
         std::swap(remote_lock_disabled, o.remote_lock_disabled);
     }
 
+private:
+    friend class UniPageStorageIdTest;
+
 private:
     String prefix;
     Writes writes;
diff --git a/dbms/src/Storages/Page/V3/Universal/tests/gtest_universal_page_storage.cpp b/dbms/src/Storages/Page/V3/Universal/tests/gtest_universal_page_storage.cpp
index 3015c2d6132..1d9f7c197e2 100644
--- a/dbms/src/Storages/Page/V3/Universal/tests/gtest_universal_page_storage.cpp
+++ b/dbms/src/Storages/Page/V3/Universal/tests/gtest_universal_page_storage.cpp
@@ -619,6 +619,50 @@ TEST(UniPageStorageIdTest, UniversalWriteBatchMemory)
         ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 1);
     }
     ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 0);
+    {
+        UniversalWriteBatch wb;
+        wb.putPage(
+            UniversalPageIdFormat::toFullPageId(prefix, 0),
+            tag,
+            std::make_shared<ReadBufferFromMemory>(c_buff, buf_sz),
+            buf_sz);
+        UniversalWriteBatch wb2;
+        wb2.merge(wb);
+        ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 2);
+    }
+    {
+        UniversalWriteBatch wb;
+        wb.putPage(
+            UniversalPageIdFormat::toFullPageId(prefix, 0),
+            tag,
+            std::make_shared<ReadBufferFromMemory>(c_buff, buf_sz),
+            buf_sz);
+        wb.clear();
+        ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 0);
+    }
+    {
+        UniversalWriteBatch wb;
+        wb.putPage(
+            UniversalPageIdFormat::toFullPageId(prefix, 0),
+            tag,
+            std::make_shared<ReadBufferFromMemory>(c_buff, buf_sz),
+            buf_sz);
+        ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 1);
+        UniversalWriteBatch wb2;
+        wb.putPage(
+            UniversalPageIdFormat::toFullPageId(prefix, 1),
+            tag,
+            std::make_shared<ReadBufferFromMemory>(c_buff, buf_sz),
+            buf_sz);
+        wb.putPage(
+            UniversalPageIdFormat::toFullPageId(prefix, 2),
+            tag,
+            std::make_shared<ReadBufferFromMemory>(c_buff, buf_sz),
+            buf_sz);
+        ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 3);
+        wb.swap(wb2);
+        ASSERT_EQ(PageStorageMemorySummary::universal_write_count.load(), 3);
+    }
 }
 
 } // namespace PS::universal::tests