stash to split

Signed-off-by: Calvin Neo <[email protected]>
pingcap · Jan 17, 2025 · c42c6b1 · c42c6b1
1 parent 81a3f64
commit c42c6b1
Show file tree

Hide file tree

Showing 3 changed files with 47 additions and 16 deletions.
diff --git a/dbms/src/Storages/KVStore/Decode/RegionTable.cpp b/dbms/src/Storages/KVStore/Decode/RegionTable.cpp
@@ -48,6 +48,7 @@ extern const char force_set_num_regions_for_table[];
 
 int64_t RegionTable::InternalRegion::updateRegionCacheBytes(size_t cache_bytes_)
 {
+    LOG_INFO(DB::Logger::get(), "!!!!!! updateRegionCacheBytes cache_bytes {} cache_bytes_ {}", cache_bytes, cache_bytes_);
     auto diff = static_cast<int64_t>(cache_bytes_) - static_cast<int64_t>(cache_bytes);
     cache_bytes = cache_bytes_;
     return diff;
@@ -66,8 +67,15 @@ RegionTable::Table & RegionTable::getOrCreateTable(const KeyspaceID keyspace_id,
     return it->second;
 }
 
-bool RegionTable::hasTable(KeyspaceID keyspace_id, TableID table_id)
+RegionTable::Table & RegionTable::debugGetOrCreateTable(const KeyspaceID keyspace_id, const TableID table_id)
 {
+    std::lock_guard lock(mutex);
+    return getOrCreateTable(keyspace_id, table_id);
+}
+
+bool RegionTable::debugHasTable(KeyspaceID keyspace_id, TableID table_id)
+{
+    std::lock_guard lock(mutex);
     auto ks_table_id = KeyspaceTableID{keyspace_id, table_id};
     auto it = tables.find(ks_table_id);
     if (it == tables.end())
@@ -104,7 +112,8 @@ RegionTable::InternalRegion & RegionTable::insertRegion(
     // Insert region mapping.
     regions[region.id()] = KeyspaceTableID{keyspace_id, table.table_id};
 
-    table.updateTableCacheBytes(it->second.updateRegionCacheBytes(region.totalSize()));
+    LOG_INFO(DB::Logger::get(), "!!!!!! updateTableCacheBytes insert diff {}", region.totalSize());
+    table.sizeDiff(it->second.updateRegionCacheBytes(region.totalSize()));
 
     return it->second;
 }
@@ -114,6 +123,7 @@ RegionTable::InternalRegion & RegionTable::doGetInternalRegion(KeyspaceTableID k
     return tables.find(ks_table_id)->second.regions.find(region_id)->second;
 }
 
+
 RegionTable::InternalRegion & RegionTable::getOrInsertRegion(const Region & region)
 {
     auto keyspace_id = region.getKeyspaceID();
@@ -131,7 +141,9 @@ void RegionTable::updateTableCacheBytes(const Region & region, int64_t diff)
     auto keyspace_id = region.getKeyspaceID();
     auto table_id = region.getMappedTableID();
     auto & table = getOrCreateTable(keyspace_id, table_id);
-    table.updateTableCacheBytes(diff);
+
+    LOG_INFO(DB::Logger::get(), "!!!!!! updateTableCacheBytes diff {}", diff);
+    table.sizeDiff(diff);
 }
 
 RegionTable::RegionTable(Context & context_)

diff --git a/dbms/src/Storages/KVStore/Decode/RegionTable.h b/dbms/src/Storages/KVStore/Decode/RegionTable.h
@@ -97,7 +97,8 @@ class RegionTable : private boost::noncopyable
             : table_id(table_id_)
             , size(0)
         {}
-        void updateTableCacheBytes(UInt64 diff) { size.fetch_add(diff); }
+        void sizeDiff(Int64 diff) { size.fetch_add(diff); }
+        Int64 getSize() const { return size; }
         TableID table_id;
         InternalRegions regions;
 
@@ -158,7 +159,8 @@ class RegionTable : private boost::noncopyable
         const LoggerPtr & log);
 
     void clear();
-    bool hasTable(KeyspaceID keyspace_id, TableID table_id);
+    bool debugHasTable(KeyspaceID keyspace_id, TableID table_id);
+    Table & debugGetOrCreateTable(KeyspaceID keyspace_id, TableID table_id);
 
 public:
     // safe ts is maintained by check_leader RPC (https://github.com/tikv/tikv/blob/1ea26a2ac8761af356cc5c0825eb89a0b8fc9749/components/resolved_ts/src/advance.rs#L262),
@@ -187,9 +189,11 @@ class RegionTable : private boost::noncopyable
 private:
     friend class MockTiDB;
     friend class StorageDeltaMerge;
+    friend class RegionKVStoreTest_MemoryTracker;
+
 
-    Table & getOrCreateTable(KeyspaceID keyspace_id, TableID table_id);
     void removeTable(KeyspaceID keyspace_id, TableID table_id);
+    Table & getOrCreateTable(KeyspaceID keyspace_id, TableID table_id);
     InternalRegion & getOrInsertRegion(const Region & region);
     InternalRegion & insertRegion(Table & table, const RegionRangeKeys & region_range_keys, const Region & region);
     InternalRegion & insertRegion(Table & table, const Region & region);
@@ -205,7 +209,9 @@ class RegionTable : private boost::noncopyable
 
     Context * const context;
 
+    // For regions and tables.
     mutable std::mutex mutex;
+    // For `safe_ts_map`
     mutable std::shared_mutex rw_lock;
 
     LoggerPtr log;

diff --git a/dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp b/dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp
@@ -87,18 +87,15 @@ try
         = RecordKVFormat::encodeLockCfValue(RecordKVFormat::CFModifyFlag::PutFlag, "PK", 111, 999).toString();
     MockRaftStoreProxy::FailCond cond;
     const auto decoded_lock_size = sizeof(DecodedLockCFValue) + sizeof(DecodedLockCFValue::Inner);
-    proxy_instance->debugAddRegions(
-        kvs,
-        ctx.getTMTContext(),
-        {1, 2, 3},
-        {{RecordKVFormat::genKey(table_id, 0), RecordKVFormat::genKey(table_id, 10)},
-         {RecordKVFormat::genKey(table_id, 11), RecordKVFormat::genKey(table_id, 20)},
-         {RecordKVFormat::genKey(table_id, 21), RecordKVFormat::genKey(table_id, 30)}});
 
     auto & region_table = ctx.getTMTContext().getRegionTable();
-    ASSERT_TRUE(region_table.hasTable(NullspaceID, table_id));
     {
         // default
+        proxy_instance->debugAddRegions(
+            kvs,
+            ctx.getTMTContext(),
+            {1},
+            {{RecordKVFormat::genKey(table_id, 0), RecordKVFormat::genKey(table_id, 10)}});
         auto region_id = 1;
         auto kvr1 = kvs.getRegion(region_id);
         auto [index, term]
@@ -109,10 +106,20 @@ try
         ASSERT_EQ(root_of_kvstore_mem_trackers->get(), str_key.dataSize() + str_val_default.size());
         ASSERT_EQ(kvr1->dataSize(), root_of_kvstore_mem_trackers->get());
         ASSERT_EQ(kvr1->dataSize(), kvr1->getData().totalSize());
+        ASSERT_TRUE(region_table.debugHasTable(NullspaceID, table_id));
+        auto & table1 = region_table.debugGetOrCreateTable(NullspaceID, table_id);
+        LOG_INFO(log, "!!!!!! dataSize {}", kvr1->dataSize());
+        ASSERT_EQ(kvr1->totalSize(), table1.getSize());
     }
     {
-        // lock
         root_of_kvstore_mem_trackers->reset();
+        // lock
+        proxy_instance->debugAddRegions(
+            kvs,
+            ctx.getTMTContext(),
+            {2},
+            {{RecordKVFormat::genKey(table_id, 11), RecordKVFormat::genKey(table_id, 20)}});
+
         auto region_id = 2;
         auto kvr1 = kvs.getRegion(region_id);
         auto [index, term]
@@ -125,8 +132,14 @@ try
         ASSERT_EQ(kvr1->dataSize() + decoded_lock_size, kvr1->getData().totalSize());
     }
     {
-        // lock with largetxn
         root_of_kvstore_mem_trackers->reset();
+        // lock with largetxn
+        proxy_instance->debugAddRegions(
+            kvs,
+            ctx.getTMTContext(),
+            {3},
+            {{RecordKVFormat::genKey(table_id, 21), RecordKVFormat::genKey(table_id, 30)}});
+
         auto region_id = 3;
         auto kvr1 = kvs.getRegion(region_id);
         ASSERT_NE(kvr1, nullptr);