KVStore: LargeTxn spill: Store kv pairs of large txns in dedicated CFs

dbms/CMakeLists.txt

@@ -72,6 +72,7 @@ add_headers_and_sources(dbms src/Databases)
 add_headers_and_sources(dbms src/Debug)
 add_headers_and_sources(dbms src/Debug/MockExecutor)
 add_headers_and_sources(dbms src/Debug/MockKVStore)
+add_headers_and_sources(dbms src/Debug/dbgKVStore)
 add_headers_and_sources(dbms src/Dictionaries)
 # Encryption needs to link tiflash-proxy, but we do not want to link tiflash-proxy to tiflash_common_io.
 # So we put the encryption source files in dbms.

dbms/src/Common/FailPoint.cpp

@@ -111,6 +111,7 @@ namespace DB
     M(force_set_fap_candidate_store_id)                      \
     M(force_not_clean_fap_on_destroy)                        \
     M(delta_tree_create_node_fail)                           \
+    M(force_write_to_large_txn_default)                      \
     M(disable_flush_cache)
 
 #define APPLY_FOR_PAUSEABLE_FAILPOINTS_ONCE(M) \

dbms/src/Debug/MockKVStore/MockUtils.h

@@ -69,11 +69,13 @@ inline metapb::Region createRegionInfo(
 inline RegionMeta createRegionMeta(
     UInt64 id,
     DB::TableID table_id,
+    const HandleID from = 0,
+    const HandleID to = 300,
     std::optional<raft_serverpb::RaftApplyState> apply_state = std::nullopt)
 {
     return RegionMeta(
         /*peer=*/createPeer(31, true),
-        /*region=*/createRegionInfo(id, RecordKVFormat::genKey(table_id, 0), RecordKVFormat::genKey(table_id, 300)),
+        /*region=*/createRegionInfo(id, RecordKVFormat::genKey(table_id, from), RecordKVFormat::genKey(table_id, to)),
         /*apply_state_=*/apply_state.value_or(initialApplyState()));
 }
 

dbms/src/Debug/dbgKVStore/dbgKVStore.cpp

@@ -0,0 +1,26 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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 <Debug/dbgKVStore/dbgKVStore.h>
+#include <Storages/KVStore/MultiRaft/RegionPersister.h>
+
+namespace DB::RegionBench
+{
+
+RegionSerdeOpts & DebugKVStore::mutRegionSerdeOpts()
+{
+    return kvstore.region_persister->region_serde_opts;
+}
+
+} // namespace DB::RegionBench

dbms/src/Debug/dbgKVStore/dbgKVStore.h

@@ -0,0 +1,40 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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.
+
+#pragma once
+
+#include <Storages/KVStore/KVStore.h>
+#include <Storages/KVStore/MultiRaft/RegionSerde.h>
+
+namespace DB::RegionBench
+{
+struct DebugKVStore
+{
+    DebugKVStore(KVStore & kvs)
+        : kvstore(kvs)
+    {}
+    KVStore * operator->() { return &kvstore; }
+    KVStore * operator->() const { return &kvstore; }
+
+    RegionSerdeOpts & mutRegionSerdeOpts();
+    template <typename RegionPtrWrap>
+    void onSnapshot(const RegionPtrWrap & r, RegionPtr old_region, UInt64 old_region_index, TMTContext & tmt)
+    {
+        kvstore.onSnapshot<RegionPtrWrap>(r, old_region, old_region_index, tmt);
+    }
+
+private:
+    KVStore & kvstore;
+};
+} // namespace DB::RegionBench

dbms/src/Debug/dbgKVStore/dbgRegion.cpp

@@ -0,0 +1,29 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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 <Debug/dbgKVStore/dbgRegion.h>
+
+namespace DB::RegionBench
+{
+
+RegionPtr DebugRegion::debugSplitInto(RegionMeta && meta)
+{
+    return region.splitInto(std::move(meta));
+}
+RegionData & DebugRegion::debugData()
+{
+    return region.data;
+}
+
+} // namespace DB::RegionBench

dbms/src/Debug/dbgKVStore/dbgRegion.h

@@ -0,0 +1,34 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// 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.
+
+#pragma once
+
+#include <Storages/KVStore/Region.h>
+
+namespace DB::RegionBench
+{
+struct DebugRegion
+{
+    DebugRegion(RegionPtr region_ptr)
+        : region(*region_ptr)
+    {}
+    RegionPtr debugSplitInto(RegionMeta && meta);
+    RegionData & debugData();
+    Region * operator->() { return &region; }
+    Region * operator->() const { return &region; }
+
+private:
+    Region & region;
+};
+} // namespace DB::RegionBench

dbms/src/Debug/dbgTools.cpp

@@ -15,6 +15,7 @@
 #include <Common/typeid_cast.h>
 #include <Debug/MockTiDB.h>
 #include <Debug/MockTiKV.h>
+#include <Debug/dbgKVStore/dbgKVStore.h>
 #include <Debug/dbgTools.h>
 #include <Interpreters/Context.h>
 #include <Parsers/ASTLiteral.h>
@@ -559,10 +560,12 @@ void concurrentBatchInsert(
     Context & context)
 {
     TMTContext & tmt = context.getTMTContext();
+    auto & orig_kvs = tmt.getKVStore();
+    auto debug_kvs = DebugKVStore(*orig_kvs);
 
     RegionID curr_max_region_id(InvalidRegionID);
     HandleID curr_max_handle_id = 0;
-    tmt.getKVStore()->traverseRegions([&](const RegionID region_id, const RegionPtr & region) {
+    debug_kvs->traverseRegions([&](const RegionID region_id, const RegionPtr & region) {
         curr_max_region_id
             = (curr_max_region_id == InvalidRegionID) ? region_id : std::max<RegionID>(curr_max_region_id, region_id);
         const auto range = region->getRange();
@@ -575,7 +578,9 @@ void concurrentBatchInsert(
     Regions regions
         = createRegions(table_info.id, concurrent_num, key_num_each_region, handle_begin, curr_max_region_id + 1);
     for (const RegionPtr & region : regions)
-        tmt.getKVStore()->onSnapshot<RegionPtrWithBlock>(region, nullptr, 0, tmt);
+    {
+        debug_kvs.onSnapshot<RegionPtrWithBlock>(region, nullptr, 0, tmt);
+    }
 
     std::list<std::thread> threads;
     for (Int64 i = 0; i < concurrent_num; i++, handle_begin += key_num_each_region)

dbms/src/Storages/KVStore/KVStore.h

@@ -35,7 +35,7 @@ namespace DB
 {
 namespace RegionBench
 {
-extern void concurrentBatchInsert(const TiDB::TableInfo &, Int64, Int64, Int64, UInt64, UInt64, Context &);
+struct DebugKVStore;
 } // namespace RegionBench
 namespace DM
 {
@@ -46,7 +46,7 @@ struct ExternalDTFileInfo;
 namespace tests
 {
 class KVStoreTestBase;
-}
+} // namespace tests
 
 class IAST;
 using ASTPtr = std::shared_ptr<IAST>;
@@ -149,6 +149,7 @@ class KVStore final : private boost::noncopyable
     FileUsageStatistics getFileUsageStatistics() const;
     // Proxy will validate and refit the config items from the toml file.
     const ProxyConfigSummary & getProxyConfigSummay() const { return proxy_config_summary; }
+    const RegionPersister & getRegionPersister() const { return *region_persister; }
 
 public: // Region Management
     void restore(PathPool & path_pool, const TiFlashRaftProxyHelper *);
@@ -258,19 +259,12 @@ class KVStore final : private boost::noncopyable
     friend struct MockRaftCommand;
     friend class RegionMockTest;
     friend class NaturalDag;
-    friend void RegionBench::concurrentBatchInsert(
-        const TiDB::TableInfo &,
-        Int64,
-        Int64,
-        Int64,
-        UInt64,
-        UInt64,
-        Context &);
     using DBGInvokerPrinter = std::function<void(const std::string &)>;
     friend void dbgFuncRemoveRegion(Context &, const ASTs &, DBGInvokerPrinter);
     friend void dbgFuncPutRegion(Context &, const ASTs &, DBGInvokerPrinter);
     friend class tests::KVStoreTestBase;
     friend class ReadIndexStressTest;
+    friend struct RegionBench::DebugKVStore;
     struct StoreMeta
     {
         mutable std::mutex mu;

dbms/src/Storages/KVStore/MultiRaft/Disagg/CheckpointIngestInfo.cpp

@@ -140,7 +140,10 @@ FastAddPeerProto::CheckpointIngestInfoPersisted CheckpointIngestInfo::serializeM
         // Although the region is the first peer of this region in this store, we can't write it to formal KVStore for now.
         // Otherwise it could be uploaded and then overwritten.
         WriteBufferFromOwnString wb;
-        RegionPersister::computeRegionWriteBuffer(*region, wb);
+        RegionPersister::computeRegionWriteBuffer(
+            *region,
+            wb,
+            tmt.getKVStore()->getRegionPersister().getRegionSerdeOpts());
         ingest_info_persisted.set_region_info(wb.releaseStr());
     }
     ingest_info_persisted.set_remote_store_id(remote_store_id);

dbms/src/Storages/KVStore/MultiRaft/PrehandleSnapshot.cpp

@@ -421,7 +421,7 @@ static void runInParallel(
     std::shared_ptr<StorageDeltaMerge> dm_storage)
 {
     std::string limit_tag = part_limit.toDebugString();
-    auto part_new_region = std::make_shared<Region>(new_region->mutMeta().clone(), proxy_helper);
+    auto part_new_region = std::make_shared<Region>(new_region->getMeta().clone(), proxy_helper);
     auto part_sst_stream = std::make_shared<DM::SSTFilesToBlockInputStream>(
         part_new_region,
         index,

dbms/src/Storages/KVStore/MultiRaft/RegionCFDataBase.cpp

@@ -16,6 +16,7 @@
 #include <Storages/KVStore/MultiRaft/RegionCFDataTrait.h>
 #include <Storages/KVStore/MultiRaft/RegionData.h>
 #include <Storages/KVStore/MultiRaft/RegionRangeKeys.h>
+#include <Storages/KVStore/MultiRaft/Spill/LargeTxnDefaultCf.h>
 
 namespace DB
 {
@@ -39,6 +40,21 @@ const TiKVValue & RegionCFDataBase<Trait>::getTiKVValue(const Value & val)
     return *getTiKVValuePtr<Value>(val);
 }
 
+RegionDataRes insertWithTs(
+    RegionCFDataBase<RegionDefaultCFDataTrait> & default_cf,
+    TiKVKey && key,
+    TiKVValue && value,
+    Timestamp ts,
+    DupCheck mode)
+{
+    const auto & raw_key = RecordKVFormat::decodeTiKVKey(key);
+    auto kv_pair = RegionDefaultCFDataTrait::genKVPairWithTs(std::move(key), ts, raw_key, std::move(value));
+    if (!kv_pair)
+        return 0;
+
+    return default_cf.doInsert(std::move(*kv_pair), mode);
+}
+
 template <typename Trait>
 RegionDataRes RegionCFDataBase<Trait>::insert(TiKVKey && key, TiKVValue && value, DupCheck mode)
 {
@@ -47,7 +63,7 @@ RegionDataRes RegionCFDataBase<Trait>::insert(TiKVKey && key, TiKVValue && value
     if (!kv_pair)
         return 0;
 
-    return insert(std::move(*kv_pair), mode);
+    return doInsert(std::move(*kv_pair), mode);
 }
 
 template <>
@@ -62,7 +78,7 @@ RegionDataRes RegionCFDataBase<RegionLockCFDataTrait>::insert(TiKVKey && key, Ti
 }
 
 template <typename Trait>
-RegionDataRes RegionCFDataBase<Trait>::insert(std::pair<Key, Value> && kv_pair, DupCheck mode)
+RegionDataRes RegionCFDataBase<Trait>::doInsert(std::pair<Key, Value> && kv_pair, DupCheck mode)
 {
     auto & map = data;
     TiKVValue prev_value;
@@ -258,15 +274,12 @@ size_t RegionCFDataBase<Trait>::splitInto(const RegionRange & range, RegionCFDat
 {
     const auto & [start_key, end_key] = range;
     size_t size_changed = 0;
-
     {
         auto & ori_map = data;
         auto & tar_map = new_region_data.data;
-
         for (auto it = ori_map.begin(); it != ori_map.end();)
         {
             const auto & key = getTiKVKey(it->second);
-
             if (start_key.compare(key) <= 0 && end_key.compare(key) > 0)
             {
                 size_changed += calcTiKVKeyValueSize(it->second);
@@ -329,4 +342,5 @@ typename RegionCFDataBase<Trait>::Data & RegionCFDataBase<Trait>::getDataMut()
 template struct RegionCFDataBase<RegionWriteCFDataTrait>;
 template struct RegionCFDataBase<RegionDefaultCFDataTrait>;
 template struct RegionCFDataBase<RegionLockCFDataTrait>;
+template struct RegionCFDataBase<LargeDefaultCFDataTrait>;
 } // namespace DB

dbms/src/Storages/KVStore/MultiRaft/RegionCFDataBase.h

@@ -21,8 +21,9 @@
 
 namespace DB
 {
-
+struct RegionDefaultCFDataTrait;
 struct TiKVRangeKey;
+
 using RegionRange = RegionRangeKeys::RegionRange;
 using RegionDataRes = size_t;
 
@@ -43,7 +44,6 @@ struct RegionCFDataBase
     using Status = bool;
 
     static const TiKVKey & getTiKVKey(const Value & val);
-
     static const TiKVValue & getTiKVValue(const Value & val);
 
     RegionDataRes insert(TiKVKey && key, TiKVValue && value, DupCheck mode = DupCheck::Deny);
@@ -52,7 +52,7 @@ struct RegionCFDataBase
 
     static size_t calcTiKVKeyValueSize(const TiKVKey & key, const TiKVValue & value);
 
-    size_t remove(const Key & key, bool quiet = false);
+    size_t remove(const Key & key, bool quiet);
 
     static bool cmp(const Map & a, const Map & b);
 
@@ -68,19 +68,26 @@ struct RegionCFDataBase
     size_t mergeFrom(const RegionCFDataBase & ori_region_data);
 
     size_t serialize(WriteBuffer & buf) const;
-
     static size_t deserialize(ReadBuffer & buf, RegionCFDataBase & new_region_data);
 
     const Data & getData() const;
 
     Data & getDataMut();
 
+    RegionDataRes doInsert(std::pair<Key, Value> && kv_pair, DupCheck mode = DupCheck::Deny);
+
 private:
     static bool shouldIgnoreRemove(const Value & value);
-    RegionDataRes insert(std::pair<Key, Value> && kv_pair, DupCheck mode = DupCheck::Deny);
 
 private:
     Data data;
 };
 
+RegionDataRes insertWithTs(
+    RegionCFDataBase<RegionDefaultCFDataTrait> & default_cf,
+    TiKVKey && key,
+    TiKVValue && value,
+    Timestamp ts,
+    DupCheck mode);
+
 } // namespace DB