Update current metrics correctly

Signed-off-by: JaySon-Huang <tshent@qq.com>

Author: JaySon-Huang

dbms/src/Common/CurrentMetrics.cpp

@@ -65,6 +65,13 @@
     M(DT_SnapshotOfReplayVersionChain)          \
     M(DT_SnapshotOfBitmapFilter)                \
     M(DT_SnapshotOfDisaggReadNodeRead)          \
+    M(DT_NumKeyspace)                           \
+    M(DT_NumIStorages)                          \
+    M(DT_NumStorageDeltaMerges)                 \
+    M(DT_NumSegments)                           \
+    M(DT_NumMemTable)                           \
+    M(DT_BytesMemTable)                         \
+    M(DT_BytesMemTableAllocated)                \
     M(IOLimiterPendingBgWriteReq)               \
     M(IOLimiterPendingFgWriteReq)               \
     M(IOLimiterPendingBgReadReq)                \
@@ -88,10 +95,7 @@
     M(ConnectionPoolSize)                       \
     M(MemoryTrackingQueryStorageTask)           \
     M(MemoryTrackingFetchPages)                 \
-    M(MemoryTrackingSharedColumnData)           \
-    M(DT_NumSegments)                           \
-    M(DT_NumMemTable)                           \
-    M(DT_BytesMemTable)
+    M(MemoryTrackingSharedColumnData)
 
 namespace CurrentMetrics
 {

dbms/src/Storages/DeltaMerge/Delta/MemTableSet.cpp

@@ -26,26 +26,72 @@ namespace CurrentMetrics
 {
 extern const Metric DT_NumMemTable;
 extern const Metric DT_BytesMemTable;
+extern const Metric DT_BytesMemTableAllocated;
 } // namespace CurrentMetrics
 
 namespace DB::DM
 {
 
+/// Member functions of MemTableSet::Statistic ///
+
+MemTableSet::Statistic::Statistic()
+    : holder_bytes(CurrentMetrics::DT_BytesMemTable, 0)
+    , holder_allocated_bytes(CurrentMetrics::DT_BytesMemTableAllocated, 0)
+{}
+
+void MemTableSet::Statistic::append(
+    size_t rows_added,
+    size_t bytes_added,
+    size_t allocated_bytes_added,
+    size_t deletes_added,
+    size_t files_added)
+{
+    column_files_count += files_added;
+    rows += rows_added;
+    bytes += bytes_added;
+    allocated_bytes += allocated_bytes_added;
+    deletes += deletes_added;
+    // update the current metrics
+    holder_bytes.changeTo(bytes.load());
+    holder_allocated_bytes.changeTo(allocated_bytes.load());
+}
+
+void MemTableSet::Statistic::resetTo(
+    size_t new_column_files_count,
+    size_t new_rows,
+    size_t new_bytes,
+    size_t new_allocated_bytes,
+    size_t new_deletes)
+{
+    column_files_count = new_column_files_count;
+    rows = new_rows;
+    bytes = new_bytes;
+    allocated_bytes = new_allocated_bytes;
+    deletes = new_deletes;
+    // update the current metrics
+    holder_bytes.changeTo(bytes.load());
+    holder_allocated_bytes.changeTo(allocated_bytes.load());
+}
+
+/// Member functions of MemTableSet ///
+
 MemTableSet::MemTableSet(const ColumnFiles & in_memory_files)
     : holder_counter(CurrentMetrics::DT_NumMemTable, 1)
-    , holder_allocated_bytes(CurrentMetrics::DT_BytesMemTable, 0)
     , column_files(in_memory_files)
     , log(Logger::get())
 {
-    column_files_count = column_files.size();
+    size_t new_rows = 0;
+    size_t new_bytes = 0;
+    size_t new_alloc_bytes = 0;
+    size_t new_deletes = 0;
     for (const auto & file : column_files)
     {
-        rows += file->getRows();
-        bytes += file->getBytes();
-        allocated_bytes += file->getAllocateBytes();
-        deletes += file->getDeletes();
+        new_rows += file->getRows();
+        new_bytes += file->getBytes();
+        new_alloc_bytes += file->getAllocateBytes();
+        new_deletes += file->getDeletes();
     }
-    holder_allocated_bytes.changeTo(allocated_bytes.load());
+    stat.resetTo(column_files.size(), new_rows, new_bytes, new_alloc_bytes, new_deletes);
 }
 
 void MemTableSet::appendColumnFileInner(const ColumnFilePtr & column_file)
@@ -60,12 +106,12 @@ void MemTableSet::appendColumnFileInner(const ColumnFilePtr & column_file)
     }
 
     column_files.push_back(column_file);
-    column_files_count = column_files.size();
-
-    rows += column_file->getRows();
-    bytes += column_file->getBytes();
-    allocated_bytes += column_file->getAllocateBytes();
-    deletes += column_file->getDeletes();
+    stat.append(
+        column_file->getRows(),
+        column_file->getBytes(),
+        column_file->getAllocateBytes(),
+        column_file->getDeletes(),
+        /*files_added=*/1);
 }
 
 std::pair</* New */ ColumnFiles, /* Flushed */ ColumnFiles> MemTableSet::diffColumnFiles(
@@ -229,9 +275,13 @@ void MemTableSet::appendToCache(DMContext & context, const Block & block, size_t
         if (unlikely(!append_res.success))
             throw Exception("Write to MemTableSet failed", ErrorCodes::LOGICAL_ERROR);
     }
-    rows += limit;
-    bytes += append_bytes;
-    allocated_bytes += append_res.new_alloc_bytes;
+
+    stat.append( //
+        limit,
+        append_bytes,
+        append_res.new_alloc_bytes,
+        /*deletes_added*/ 0,
+        /*files_added*/ 0);
 }
 
 void MemTableSet::appendDeleteRange(const RowKeyRange & delete_range)
@@ -294,13 +344,18 @@ ColumnFileSetSnapshotPtr MemTableSet::createSnapshot(
     // This may indicate that you forget to acquire a lock -- there are modifications
     // while this function is still running...
     RUNTIME_CHECK(
-        total_rows == rows && total_deletes == deletes,
+        total_rows == stat.rows && total_deletes == stat.deletes,
         total_rows,
-        rows.load(),
+        stat.rows.load(),
         total_deletes,
-        deletes.load());
-
-    return std::make_shared<ColumnFileSetSnapshot>(data_provider, std::move(column_files_snap), rows, bytes, deletes);
+        stat.deletes.load());
+
+    return std::make_shared<ColumnFileSetSnapshot>(
+        data_provider,
+        std::move(column_files_snap),
+        stat.rows,
+        stat.bytes,
+        stat.deletes);
 }
 
 ColumnFileFlushTaskPtr MemTableSet::buildFlushTask(
@@ -333,16 +388,16 @@ ColumnFileFlushTaskPtr MemTableSet::buildFlushTask(
         cur_rows_offset += column_file->getRows();
         cur_deletes_offset += column_file->getDeletes();
     }
-    if (unlikely(flush_task->getFlushRows() != rows || flush_task->getFlushDeletes() != deletes))
+    if (unlikely(flush_task->getFlushRows() != stat.rows || flush_task->getFlushDeletes() != stat.deletes))
     {
         LOG_ERROR(
             log,
             "Rows and deletes check failed. Actual: rows[{}], deletes[{}]. Expected: rows[{}], deletes[{}]. Column "
             "Files: {}",
             flush_task->getFlushRows(),
             flush_task->getFlushDeletes(),
-            rows.load(),
-            deletes.load(),
+            stat.rows.load(),
+            stat.deletes.load(),
             ColumnFile::filesToString(column_files));
         throw Exception("Rows and deletes check failed.", ErrorCodes::LOGICAL_ERROR);
     }
@@ -378,11 +433,12 @@ void MemTableSet::removeColumnFilesInFlushTask(const ColumnFileFlushTask & flush
         new_deletes += column_file->getDeletes();
     }
     column_files.swap(new_column_files);
-    column_files_count = column_files.size();
-    rows = new_rows;
-    bytes = new_bytes;
-    allocated_bytes = new_alloc_bytes;
-    deletes = new_deletes;
+    stat.resetTo( //
+        new_column_files.size(),
+        new_rows,
+        new_bytes,
+        new_alloc_bytes,
+        new_deletes);
 }
 
 

dbms/src/Storages/DeltaMerge/Delta/MemTableSet.h

@@ -32,25 +32,60 @@ class MemTableSet
     : public std::enable_shared_from_this<MemTableSet>
     , private boost::noncopyable
 {
+private:
+    struct Statistic
+    {
+        // TODO: check the proper memory_order when use this atomic variable
+        std::atomic<size_t> column_files_count = 0;
+        std::atomic<size_t> rows = 0;
+        std::atomic<size_t> bytes = 0;
+        std::atomic<size_t> allocated_bytes = 0;
+        std::atomic<size_t> deletes = 0;
+
+        CurrentMetrics::Increment holder_bytes;
+        CurrentMetrics::Increment holder_allocated_bytes;
+
+        Statistic();
+
+        String info() const
+        {
+            return fmt::format(
+                "MemTableSet: {} column files, {} rows, {} bytes, {} deletes",
+                column_files_count.load(),
+                rows.load(),
+                bytes.load(),
+                deletes.load());
+        }
+
+        void append(
+            size_t rows_added,
+            size_t bytes_added,
+            size_t allocated_bytes_added,
+            size_t deletes_added,
+            size_t files_added);
+
+        void resetTo(
+            size_t new_column_files_count,
+            size_t new_rows,
+            size_t new_bytes,
+            size_t new_allocated_bytes,
+            size_t new_deletes);
+    };
+
 #ifndef DBMS_PUBLIC_GTEST
 private:
 #else
 public:
 #endif
     // Keep track of the number of mem-table in memory.
     CurrentMetrics::Increment holder_counter;
-    CurrentMetrics::Increment holder_allocated_bytes;
 
     // Note that we must update `column_files_count` for outer thread-safe after `column_files` changed
     ColumnFiles column_files;
 
-    // In order to avoid data-race, we use atomic variables to track the state of this MemTableSet.
-    // TODO: check the proper memory_order when use this atomic variable
-    std::atomic<size_t> column_files_count;
-    std::atomic<size_t> rows = 0;
-    std::atomic<size_t> bytes = 0;
-    std::atomic<size_t> allocated_bytes = 0;
-    std::atomic<size_t> deletes = 0;
+    // In order to avoid data-race and make it lightweight for accessing the statistic
+    // of mem-table, we use atomic variables to track the state of this MemTableSet.
+    Statistic stat;
 
     LoggerPtr log;
 
@@ -68,21 +103,13 @@ class MemTableSet
     void resetLogger(const LoggerPtr & segment_log) { log = segment_log; }
 
     /// Thread safe part start
-    String info() const
-    {
-        return fmt::format(
-            "MemTableSet: {} column files, {} rows, {} bytes, {} deletes",
-            column_files_count.load(),
-            rows.load(),
-            bytes.load(),
-            deletes.load());
-    }
-
-    size_t getColumnFileCount() const { return column_files_count.load(); }
-    size_t getRows() const { return rows.load(); }
-    size_t getBytes() const { return bytes.load(); }
-    size_t getAllocatedBytes() const { return allocated_bytes.load(); }
-    size_t getDeletes() const { return deletes.load(); }
+    String info() const { return stat.info(); }
+
+    size_t getColumnFileCount() const { return stat.column_files_count.load(); }
+    size_t getRows() const { return stat.rows.load(); }
+    size_t getBytes() const { return stat.bytes.load(); }
+    size_t getAllocatedBytes() const { return stat.allocated_bytes.load(); }
+    size_t getDeletes() const { return stat.deletes.load(); }
     /// Thread safe part end
 
     /**

dbms/src/Storages/IStorage.cpp

@@ -14,6 +14,10 @@
 
 #include <Storages/IStorage.h>
 
+namespace CurrentMetrics
+{
+extern const Metric DT_NumIStorages;
+} // namespace CurrentMetrics
 
 namespace DB
 {
@@ -24,6 +28,14 @@ extern const int DEADLOCK_AVOIDED;
 extern const int TABLE_IS_DROPPED;
 } // namespace ErrorCodes
 
+IStorage::IStorage()
+    : holder_counter(CurrentMetrics::DT_NumIStorages, 1)
+{}
+
+IStorage::IStorage(ColumnsDescription columns_)
+    : ITableDeclaration(std::move(columns_))
+    , holder_counter(CurrentMetrics::DT_NumIStorages, 1)
+{}
 
 RWLock::LockHolder IStorage::tryLockTimed(
     const RWLockPtr & rwlock,

dbms/src/Storages/IStorage.h

@@ -63,6 +63,10 @@ class IStorage
     , public ITableDeclaration
 {
 public:
+    IStorage();
+
+    explicit IStorage(ColumnsDescription columns_);
+
     /// The main name of the table type (for example, StorageDeltaMerge).
     virtual std::string getName() const = 0;
 
@@ -329,6 +333,8 @@ class IStorage
     /// DROP-like queries take this lock for write (lockExclusively), to be sure
     /// that all table threads finished.
     mutable RWLockPtr drop_lock = RWLock::create();
+
+    CurrentMetrics::Increment holder_counter;
 };
 
 /// table name -> table

dbms/src/Storages/ITableDeclaration.h

@@ -37,7 +37,7 @@ class ITableDeclaration
     Block getSampleBlockNonMaterialized() const;
     Block getSampleBlockForColumns(const Names & column_names) const;
 
-    /** The hidden coloumns will not be returned. Mainly for INSERT query.
+    /** The hidden columns will not be returned. Mainly for INSERT query.
       */
     Block getSampleBlockNoHidden() const;
     Block getSampleBlockNonMaterializedNoHidden() const;

dbms/src/Storages/StorageDeltaMerge.cpp

@@ -62,6 +62,10 @@
 #include <TiDB/Schema/TiDB.h>
 #include <common/logger_useful.h>
 
+namespace CurrentMetrics
+{
+extern const Metric DT_NumStorageDeltaMerges;
+} // namespace CurrentMetrics
 
 namespace DB
 {
@@ -88,6 +92,7 @@ StorageDeltaMerge::StorageDeltaMerge(
     Timestamp tombstone,
     Context & global_context_)
     : IManageableStorage{columns_, tombstone}
+    , holder_counter(CurrentMetrics::DT_NumStorageDeltaMerges, 1)
     , store_inited(false)
     , max_column_id_used(0)
     , data_path_contains_database_name(db_engine != "TiFlash")

dbms/src/Storages/StorageDeltaMerge.h

@@ -275,6 +275,9 @@ class StorageDeltaMerge
         DM::ColumnDefine handle_column_define;
     };
 
+    // Keep track of the number of StorageDeltaMerge in memory.
+    CurrentMetrics::Increment holder_counter;
+
     mutable std::mutex store_mutex;
 
     std::unique_ptr<TableColumnInfo> table_column_info; // After create DeltaMergeStore object, it is deprecated.
@@ -303,12 +306,11 @@ class StorageDeltaMerge
     // Used to allocate new column-id when this table is NOT synced from TiDB
     ColumnID max_column_id_used;
 
-    std::atomic<bool> shutdown_called{false};
-
     // TODO: remove the following two members, which are only used for debug from ch-client.
     Strings pk_column_names;
     std::atomic<UInt64> next_version = 1;
 
+    std::atomic<bool> shutdown_called{false};
     bool is_common_handle = false;
     bool pk_is_handle = false;