KVStore: Record table-wise memory usage and warn when it goes beyond the threshold

dbms/src/Debug/MockKVStore/MockRaftStoreProxy.cpp

@@ -183,7 +183,7 @@ void MockRaftStoreProxy::debugAddRegions(
             auto region = tests::makeRegion(region_ids[i], ranges[i].first, ranges[i].second, kvs.getProxyHelper());
             lock.regions.emplace(region_ids[i], region);
             lock.index.add(region);
-            tmt.getRegionTable().updateRegion(*region);
+            tmt.getRegionTable().addRegion(*region);
         }
     }
 }
@@ -601,7 +601,7 @@ std::tuple<RegionPtr, PrehandleResult> MockRaftStoreProxy::snapshot(
     uint64_t index,
     uint64_t term,
     std::optional<uint64_t> deadline_index,
-    bool cancel_after_prehandle)
+    std::optional<std::function<void()>> cancel_after_prehandle)
 {
     auto old_kv_region = kvs.getRegion(region_id);
     RUNTIME_CHECK(old_kv_region != nullptr);
@@ -628,7 +628,7 @@ std::tuple<RegionPtr, PrehandleResult> MockRaftStoreProxy::snapshot(
     uint64_t index,
     uint64_t term,
     std::optional<uint64_t> deadline_index,
-    bool cancel_after_prehandle)
+    std::optional<std::function<void()>> cancel_after_prehandle)
 {
     auto region = getRegion(region_id);
     RUNTIME_CHECK(region != nullptr);
@@ -640,10 +640,10 @@ std::tuple<RegionPtr, PrehandleResult> MockRaftStoreProxy::snapshot(
         term = region->getLatestCommitTerm();
     }
 
-    auto new_kv_region = kvs.genRegionPtr(std::move(region_meta), peer_id, index, term);
     // The new entry is committed on Proxy's side.
     region->updateCommitIndex(index);
-    new_kv_region->setApplied(index, term);
+    // Would set `appleid_index` in genRegionPtr.
+    auto new_kv_region = kvs.genRegionPtr(std::move(region_meta), peer_id, index, term, tmt, true);
 
     std::vector<SSTView> ssts;
     for (auto & cf : cfs)
@@ -663,6 +663,7 @@ std::tuple<RegionPtr, PrehandleResult> MockRaftStoreProxy::snapshot(
         auto rg = RegionPtrWithSnapshotFiles{new_kv_region, std::vector(prehandle_result.ingest_ids)};
         if (cancel_after_prehandle)
         {
+            cancel_after_prehandle.value()();
             kvs.releasePreHandledSnapshot(rg, tmt);
             return std::make_tuple(kvs.getRegion(region_id), prehandle_result);
         }

dbms/src/Debug/MockKVStore/MockRaftStoreProxy.h

@@ -242,7 +242,7 @@ struct MockRaftStoreProxy : MutexLockWrap
         uint64_t index,
         uint64_t term,
         std::optional<uint64_t> deadline_index,
-        bool cancel_after_prehandle);
+        std::optional<std::function<void()>> cancel_after_prehandle);
     std::tuple<RegionPtr, PrehandleResult> snapshot(
         KVStore & kvs,
         TMTContext & tmt,
@@ -251,7 +251,7 @@ struct MockRaftStoreProxy : MutexLockWrap
         uint64_t index,
         uint64_t term,
         std::optional<uint64_t> deadline_index,
-        bool cancel_after_prehandle = false);
+        std::optional<std::function<void()>> cancel_after_prehandle = std::nullopt);
 
     void doApply(
         KVStore & kvs,
@@ -266,6 +266,57 @@ struct MockRaftStoreProxy : MutexLockWrap
 
     void clear();
 
+    struct WriteCmdsViewHolder
+    {
+        WriteCmdsViewHolder(
+            std::vector<std::string> kk,
+            std::vector<std::string> vv,
+            std::vector<WriteCmdType> tt,
+            std::vector<ColumnFamilyType> ff)
+        {
+            kholder = std::move(kk);
+            vholder = std::move(vv);
+            for (size_t i = 0; i < kholder.size(); i++)
+            {
+                kbuff.push_back(strIntoView(&kholder[i]));
+            }
+            for (size_t i = 0; i < vholder.size(); i++)
+            {
+                vbuff.push_back(strIntoView(&vholder[i]));
+            }
+            cmd_type = std::move(tt);
+            cmd_cf = std::move(ff);
+        }
+
+        const BaseBuffView * getKeys() const { return kbuff.data(); }
+
+        const BaseBuffView * getVals() const { return vbuff.data(); }
+
+        const WriteCmdType * getTypes() const { return cmd_type.data(); }
+
+        const ColumnFamilyType * getCfs() const { return cmd_cf.data(); }
+
+        std::vector<std::string> kholder;
+        std::vector<std::string> vholder;
+        std::vector<BaseBuffView> kbuff;
+        std::vector<BaseBuffView> vbuff;
+        std::vector<WriteCmdType> cmd_type;
+        std::vector<ColumnFamilyType> cmd_cf;
+    };
+
+    static std::tuple<WriteCmdsView, std::shared_ptr<WriteCmdsViewHolder>> createWriteCmdsView(
+        std::vector<std::string> keys,
+        std::vector<std::string> vals,
+        std::vector<WriteCmdType> cmd_types,
+        std::vector<ColumnFamilyType> cmd_cf)
+    {
+        std::shared_ptr<WriteCmdsViewHolder> holder
+            = std::make_shared<WriteCmdsViewHolder>(keys, vals, cmd_types, cmd_cf);
+        return std::make_tuple(
+            WriteCmdsView{holder->getKeys(), holder->getVals(), holder->getTypes(), holder->getCfs(), keys.size()},
+            holder);
+    }
+
     std::pair<std::string, std::string> generateTiKVKeyValue(uint64_t tso, int64_t t) const;
 
     MockRaftStoreProxy()

dbms/src/Debug/dbgKVStore/dbgFuncMockRaftSnapshot.cpp

@@ -154,7 +154,11 @@ RegionPtr GenDbgRegionSnapshotWithData(Context & context, const ASTs & args)
                                          : RecordKVFormat::encodeWriteCfValue(Region::PutFlag, prewrite_ts, value);
             TiKVKey commit_key = RecordKVFormat::appendTs(key, commit_ts);
 
-            region->insert(ColumnFamilyType::Write, std::move(commit_key), std::move(commit_value));
+            region->insertFromSnap(
+                context.getTMTContext(),
+                ColumnFamilyType::Write,
+                std::move(commit_key),
+                std::move(commit_value));
         }
         MockTiKV::instance().getRaftIndex(region_id);
     }

dbms/src/Debug/dbgTools.cpp

@@ -784,7 +784,7 @@ void handleApplySnapshot(
     std::optional<uint64_t> deadline_index,
     TMTContext & tmt)
 {
-    auto new_region = kvstore.genRegionPtr(std::move(region), peer_id, index, term);
+    auto new_region = kvstore.genRegionPtr(std::move(region), peer_id, index, term, tmt, true);
     auto prehandle_result = kvstore.preHandleSnapshotToFiles(new_region, snaps, index, term, deadline_index, tmt);
     kvstore.applyPreHandledSnapshot(
         RegionPtrWithSnapshotFiles{new_region, std::move(prehandle_result.ingest_ids)},

dbms/src/Server/Server.cpp

@@ -991,7 +991,10 @@ int Server::main(const std::vector<std::string> & /*args*/)
         /// Create TMTContext
         auto cluster_config = getClusterConfig(global_context->getSecurityConfig(), storage_config.api_version, log);
         global_context->createTMTContext(raft_config, std::move(cluster_config));
-        proxy_machine.initKVStore(global_context->getTMTContext(), store_ident);
+        proxy_machine.initKVStore(
+            global_context->getTMTContext(),
+            store_ident,
+            settings.max_memory_usage_for_all_queries.getActualBytes(server_info.memory_info.capacity));
 
         global_context->getTMTContext().reloadConfig(config());
         // setup the kv cluster for disagg compute node fetching config

dbms/src/Storages/DeltaMerge/Decode/SSTFilesToBlockInputStream.cpp

@@ -194,7 +194,7 @@ Block SSTFilesToBlockInputStream::read()
             BaseBuffView key = write_cf_reader->keyView();
             BaseBuffView value = write_cf_reader->valueView();
             auto tikv_key = TiKVKey(key.data, key.len);
-            region->insert(ColumnFamilyType::Write, std::move(tikv_key), TiKVValue(value.data, value.len));
+            region->insertFromSnap(tmt, ColumnFamilyType::Write, std::move(tikv_key), TiKVValue(value.data, value.len));
             ++process_keys.write_cf;
             process_keys.write_cf_bytes += (key.len + value.len);
             if (process_keys.write_cf % opts.expected_size == 0)
@@ -277,7 +277,12 @@ void SSTFilesToBlockInputStream::loadCFDataFromSST(
             BaseBuffView key = reader->keyView();
             BaseBuffView value = reader->valueView();
             // TODO: use doInsert to avoid locking
-            region->insert(cf, TiKVKey(key.data, key.len), TiKVValue(value.data, value.len), DupCheck::AllowSame);
+            region->insertFromSnap(
+                tmt,
+                cf,
+                TiKVKey(key.data, key.len),
+                TiKVValue(value.data, value.len),
+                DupCheck::AllowSame);
             (*p_process_keys) += 1;
             (*p_process_keys_bytes) += (key.len + value.len);
             reader->next();
@@ -337,7 +342,7 @@ void SSTFilesToBlockInputStream::loadCFDataFromSST(
                 BaseBuffView key = reader->keyView();
                 BaseBuffView value = reader->valueView();
                 // TODO: use doInsert to avoid locking
-                region->insert(cf, TiKVKey(key.data, key.len), TiKVValue(value.data, value.len));
+                region->insertFromSnap(tmt, cf, TiKVKey(key.data, key.len), TiKVValue(value.data, value.len));
                 (*p_process_keys) += 1;
                 (*p_process_keys_bytes) += (key.len + value.len);
                 if (*p_process_keys == process_keys_offset_end)