Optimize kvstore memory

dbms/src/Common/TiFlashMetrics.h

@@ -80,6 +80,11 @@ namespace DB
         F(type_cancel_mpp_task, {{"type", "cancel_mpp_task"}}, ExpBuckets{0.0005, 2, 30}),                                                \
         F(type_run_mpp_task, {{"type", "run_mpp_task"}}, ExpBuckets{0.0005, 2, 30}))                                                      \
     M(tiflash_coprocessor_response_bytes, "Total bytes of response body", Counter)                                                        \
+    M(tiflash_kvstore_region_data_memory_size, "Total region data memory size of tiflash kvstore", Gauge)                                 \
+    M(tiflash_kvstore_region_data_consume_memory_size, "Total region data memory size of tiflash kvstore", Gauge)                                 \
+    M(tiflash_kvstore_region_cf_size, "Total region data memory size of tiflash kvstore", Gauge)                                 \
+    M(tiflash_kvstore_raft_command_size, "Total region data memory size of tiflash kvstore", Gauge)                                 \
+    M(tiflash_kvstore_data_approx_cache_memory_size, "Total region data memory size of tiflash memory data", Gauge)                                 \
     M(tiflash_schema_version, "Current version of tiflash cached schema", Gauge)                                                          \
     M(tiflash_schema_applying, "Whether the schema is applying or not (holding lock)", Gauge)                                             \
     M(tiflash_schema_apply_count, "Total number of each kinds of apply", Counter, F(type_diff, {"type", "diff"}),                         \

dbms/src/Storages/DeltaMerge/DeltaMergeStore.cpp

@@ -1240,7 +1240,7 @@ void DeltaMergeStore::checkSegmentUpdate(const DMContextPtr & dm_context, const
     bool should_background_flush = (unsaved_rows >= delta_cache_limit_rows || unsaved_bytes >= delta_cache_limit_bytes) //
         && (delta_rows - delta_last_try_flush_rows >= delta_cache_limit_rows
             || delta_bytes - delta_last_try_flush_bytes >= delta_cache_limit_bytes);
-    bool should_foreground_flush = unsaved_rows >= delta_cache_limit_rows * 3 || unsaved_bytes >= delta_cache_limit_bytes * 3;
+    bool should_foreground_flush = unsaved_rows >= delta_cache_limit_rows * 100 || unsaved_bytes >= delta_cache_limit_bytes * 100;
 
     bool should_background_merge_delta = ((delta_check_rows >= delta_limit_rows || delta_check_bytes >= delta_limit_bytes) //
                                           && (delta_rows - delta_last_try_merge_delta_rows >= delta_cache_limit_rows

dbms/src/Storages/Transaction/KVStore.cpp

@@ -275,6 +275,27 @@ EngineStoreApplyRes KVStore::handleWriteRaftCmd(
 
 EngineStoreApplyRes KVStore::handleWriteRaftCmd(const WriteCmdsView & cmds, UInt64 region_id, UInt64 index, UInt64 term, TMTContext & tmt)
 {
+    size_t cmd_size = 0;
+    cmd_size += sizeof(WriteCmdsView);
+    for (size_t i = 0; i < cmds.len; i++)
+    {
+        auto type = cmds.cmd_types[i];
+        switch (type)
+        {
+        case WriteCmdType::Put:
+        {
+            cmd_size += cmds.keys[i].len + sizeof(BaseBuffView);
+            cmd_size += cmds.vals[i].len + sizeof(BaseBuffView);
+            break;
+        }
+        case WriteCmdType::Del:
+        {
+            cmd_size += cmds.keys[i].len + sizeof(BaseBuffView);
+            break;
+        }
+        }
+    }
+    GET_METRIC(tiflash_kvstore_raft_command_size).Increment(cmd_size);
     auto region_persist_lock = region_manager.genRegionTaskLock(region_id);
 
     const RegionPtr region = getRegion(region_id);
@@ -284,6 +305,20 @@ EngineStoreApplyRes KVStore::handleWriteRaftCmd(const WriteCmdsView & cmds, UInt
     }
 
     auto res = region->handleWriteRaftCmd(cmds, index, term, tmt);
+    size_t region_data_size = 0;
+    size_t region_data_memory_size = 0;
+    size_t region_default_cf_num = 0;
+    size_t region_lock_cf_num = 0;
+    traverseRegions([&region_data_size, &region_data_memory_size, &region_default_cf_num, &region_lock_cf_num](RegionID, const RegionPtr & region) {
+        region_data_size += region->dataSize();
+        region_data_memory_size += region->memorySize();
+        region_default_cf_num += region->defaultCFCount();
+        region_lock_cf_num += region->lockCFCount();
+    });
+    GET_METRIC(tiflash_kvstore_region_data_memory_size).Set(region_data_size);
+    GET_METRIC(tiflash_kvstore_region_data_consume_memory_size).Set(region_data_memory_size);
+    GET_METRIC(tiflash_kvstore_region_cf_size).Set(region_default_cf_num);
+    LOG_FMT_DEBUG(&Poco::Logger::get(__PRETTY_FUNCTION__), "region total default cf count {} lock cf count {}", region_default_cf_num, region_lock_cf_num);
     return res;
 }
 

dbms/src/Storages/Transaction/PartitionStreams.cpp

@@ -344,6 +344,7 @@ void RegionTable::writeBlockByRegion(
     if (!data_list_read)
         return;
 
+    LOG_FMT_DEBUG(&Poco::Logger::get(__PRETTY_FUNCTION__), "write committed data rows {}", data_list_read->size());
     reportUpstreamLatency(*data_list_read);
     writeRegionDataToStorage(context, region, *data_list_read, log);
 

dbms/src/Storages/Transaction/Region.cpp

@@ -434,12 +434,29 @@ size_t Region::dataSize() const
     return data.dataSize();
 }
 
+size_t Region::memorySize() const
+{
+    return data.memorySize();
+}
+
 size_t Region::writeCFCount() const
 {
     std::shared_lock<std::shared_mutex> lock(mutex);
     return data.writeCF().getSize();
 }
 
+size_t Region::defaultCFCount() const
+{
+    std::shared_lock<std::shared_mutex> lock(mutex);
+    return data.defaultCF().getSize();
+}
+
+size_t Region::lockCFCount() const
+{
+    std::shared_lock<std::shared_mutex> lock(mutex);
+    return data.lockCF().getSize();
+}
+
 std::string Region::dataInfo() const
 {
     std::shared_lock<std::shared_mutex> lock(mutex);

dbms/src/Storages/Transaction/Region.h

@@ -147,7 +147,10 @@ class Region : public std::enable_shared_from_this<Region>
     void setStateApplying();
 
     size_t dataSize() const;
+    size_t memorySize() const;
     size_t writeCFCount() const;
+    size_t defaultCFCount() const;
+    size_t lockCFCount() const;
     std::string dataInfo() const;
 
     void markCompactLog() const;

dbms/src/Storages/Transaction/RegionCFDataBase.cpp

@@ -16,6 +16,7 @@
 #include <Storages/Transaction/RegionCFDataTrait.h>
 #include <Storages/Transaction/RegionData.h>
 #include <Storages/Transaction/RegionRangeKeys.h>
+#include <common/logger_useful.h>
 
 namespace DB
 {
@@ -43,7 +44,8 @@ template <typename Trait>
 RegionDataRes RegionCFDataBase<Trait>::insert(TiKVKey && key, TiKVValue && value)
 {
     const auto & raw_key = RecordKVFormat::decodeTiKVKey(key);
-    auto kv_pair = Trait::genKVPair(std::move(key), raw_key, std::move(value));
+    auto [kv_pair, memory_size] = Trait::genKVPair(std::move(key), raw_key, std::move(value));
+    std::ignore = memory_size;
     if (!kv_pair)
         return 0;
 
@@ -53,10 +55,23 @@ RegionDataRes RegionCFDataBase<Trait>::insert(TiKVKey && key, TiKVValue && value
 template <>
 RegionDataRes RegionCFDataBase<RegionLockCFDataTrait>::insert(TiKVKey && key, TiKVValue && value)
 {
+    auto allocator = data.get_allocator();
+    auto old_size = allocator.getMemoryAllocatedSize();
     Pair kv_pair = RegionLockCFDataTrait::genKVPair(std::move(key), std::move(value));
     // according to the process of pessimistic lock, just overwrite.
-    data.insert_or_assign(std::move(kv_pair.first), std::move(kv_pair.second));
-    return 0;
+    auto [it, ok] = data.insert_or_assign(std::move(kv_pair.first), std::move(kv_pair.second));
+    LOG_FMT_DEBUG(&Poco::Logger::get(__PRETTY_FUNCTION__), "Lock cf size {} bucket count {} allocator size {} entry size {}",
+                  data.size(), data.bucket_count(), allocator.getMemoryAllocatedSize(), sizeof(RegionLockCFDataTrait::Map::value_type));
+    if (ok)
+    {
+        auto new_size = allocator.getMemoryAllocatedSize();
+        return new_size - old_size + calcTiKVKeyValueSize(it->second);
+    }
+    else
+    {
+        LOG_FMT_DEBUG(&Poco::Logger::get(__PRETTY_FUNCTION__), "RegionCFDataBase<RegionLockCFDataTrait> override lock");
+        return 0;
+    }
 }
 
 template <typename Trait>
@@ -79,14 +94,13 @@ size_t RegionCFDataBase<Trait>::calcTiKVKeyValueSize(const Value & value)
 template <typename Trait>
 size_t RegionCFDataBase<Trait>::calcTiKVKeyValueSize(const TiKVKey & key, const TiKVValue & value)
 {
-    if constexpr (std::is_same<Trait, RegionLockCFDataTrait>::value)
-    {
-        std::ignore = key;
-        std::ignore = value;
-        return 0;
-    }
-    else
-        return key.dataSize() + value.dataSize();
+//    if constexpr (std::is_same<Trait, RegionLockCFDataTrait>::value)
+//    {
+//        return key.dataSize() + value.dataSize();
+//    }
+//    else
+//        return key.dataSize() + value.dataSize();
+    return key.dataSize() + value.dataSize();
 }
 
 
@@ -115,7 +129,20 @@ size_t RegionCFDataBase<Trait>::remove(const Key & key, bool quiet)
             return 0;
 
         size_t size = calcTiKVKeyValueSize(value);
+        size_t old_map_size = 0;
+        size_t new_map_size = 0;
+        if constexpr (std::is_same<Trait, RegionLockCFDataTrait>::value)
+        {
+            old_map_size = map.get_allocator().getMemoryAllocatedSize();
+        }
         map.erase(it);
+        if constexpr (std::is_same<Trait, RegionLockCFDataTrait>::value)
+        {
+            new_map_size = map.get_allocator().getMemoryAllocatedSize();
+            assert(new_map_size <= old_map_size);
+            size += old_map_size - new_map_size;
+        }
+
         return size;
     }
     else if (!quiet)
@@ -241,7 +268,8 @@ size_t RegionCFDataBase<Trait>::deserialize(ReadBuffer & buf, RegionCFDataBase &
     {
         auto key = TiKVKey::deserialize(buf);
         auto value = TiKVValue::deserialize(buf);
-        cf_data_size += new_region_data.insert(std::move(key), std::move(value));
+        auto data_size = new_region_data.insert(std::move(key), std::move(value));
+        cf_data_size += data_size;
     }
     return cf_data_size;
 }

dbms/src/Storages/Transaction/RegionCFDataTrait.h

@@ -38,17 +38,35 @@ struct RegionWriteCFDataTrait
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>, DecodedWriteCFValue>;
     using Map = std::map<Key, Value>;
 
-    static std::optional<Map::value_type> genKVPair(TiKVKey && key, const DecodedTiKVKey & raw_key, TiKVValue && value)
+    static std::pair<std::optional<Map::value_type>, size_t> genKVPair(TiKVKey && key, const DecodedTiKVKey & raw_key, TiKVValue && value)
     {
         auto decoded_val = RecordKVFormat::decodeWriteCfValue(value);
         if (!decoded_val)
-            return std::nullopt;
+            return std::pair(std::nullopt, 0);
 
         RawTiDBPK tidb_pk = RecordKVFormat::getRawTiDBPK(raw_key);
         Timestamp ts = RecordKVFormat::getTs(key);
-        return Map::value_type(Key(std::move(tidb_pk), ts),
+
+        size_t memory_size = 0;
+        memory_size += key.capacity();
+        memory_size += raw_key.capacity();
+        memory_size += value.capacity();
+        memory_size += sizeof(decoded_val);
+        memory_size += sizeof(RecordKVFormat::InnerDecodedWriteCFValue);
+        memory_size += sizeof(decoded_val->short_value);
+        memory_size += decoded_val->short_value->capacity();
+        memory_size += sizeof(tidb_pk);
+        memory_size += tidb_pk->capacity();
+        memory_size += sizeof(ts);
+        memory_size += sizeof(Key);
+        memory_size += sizeof(Value);
+        memory_size += sizeof(std::shared_ptr<const TiKVKey>);
+        memory_size += sizeof(std::shared_ptr<const TiKVValue>);
+        memory_size += sizeof(Map::value_type);
+
+        return std::make_pair(Map::value_type(Key(std::move(tidb_pk), ts),
             Value(std::make_shared<const TiKVKey>(std::move(key)), std::make_shared<const TiKVValue>(std::move(value)),
-                std::move(*decoded_val)));
+                std::move(*decoded_val))), memory_size);
     }
 
     static const std::shared_ptr<const TiKVValue> & getRecordRawValuePtr(const Value & value) { return std::get<2>(value).short_value; }
@@ -63,17 +81,63 @@ struct RegionDefaultCFDataTrait
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>>;
     using Map = std::map<Key, Value>;
 
-    static std::optional<Map::value_type> genKVPair(TiKVKey && key, const DecodedTiKVKey & raw_key, TiKVValue && value)
+    static std::pair<std::optional<Map::value_type>, size_t> genKVPair(TiKVKey && key, const DecodedTiKVKey & raw_key, TiKVValue && value)
     {
         RawTiDBPK tidb_pk = RecordKVFormat::getRawTiDBPK(raw_key);
         Timestamp ts = RecordKVFormat::getTs(key);
-        return Map::value_type(Key(std::move(tidb_pk), ts),
-            Value(std::make_shared<const TiKVKey>(std::move(key)), std::make_shared<const TiKVValue>(std::move(value))));
+        size_t memory_size = 0;
+        memory_size += key.capacity();
+        memory_size += raw_key.capacity();
+        memory_size += value.capacity();
+        memory_size += sizeof(tidb_pk);
+        memory_size += tidb_pk->capacity();
+        memory_size += sizeof(ts);
+        memory_size += sizeof(Key);
+        memory_size += sizeof(Value);
+        memory_size += sizeof(std::shared_ptr<const TiKVKey>);
+        memory_size += sizeof(std::shared_ptr<const TiKVValue>);
+        memory_size += sizeof(Map::value_type);
+        return std::make_pair(Map::value_type(Key(std::move(tidb_pk), ts),
+            Value(std::make_shared<const TiKVKey>(std::move(key)), std::make_shared<const TiKVValue>(std::move(value)))), memory_size);
     }
 
     static std::shared_ptr<const TiKVValue> getTiKVValue(const Map::const_iterator & it) { return std::get<1>(it->second); }
 };
 
+template <typename T>
+class MyAlloc {
+public:
+    typedef T value_type;
+
+    MyAlloc() =default;
+    template <typename U>
+    MyAlloc (const MyAlloc<U>& other) noexcept {
+        memory_allocated = other.getMemoryAllocatedSize();
+    }
+
+    T* allocate (std::size_t num)
+    {
+        memory_allocated += num * sizeof(T);
+        return static_cast<T*>(::operator new(num*sizeof(T)));
+    }
+    void deallocate (T* p, std::size_t num)
+    {
+        memory_allocated -= num * sizeof(T);
+        ::operator delete(p);
+    }
+
+    size_t getMemoryAllocatedSize() const noexcept { return memory_allocated; }
+
+private:
+    size_t memory_allocated = 0;
+};
+
+template <typename T1, typename T2>
+bool operator== (const MyAlloc<T1>&, const MyAlloc<T2>&) noexcept { return false; }
+
+template <typename T1, typename T2>
+bool operator!= (const MyAlloc<T1>&, const MyAlloc<T2>&) noexcept { return true; }
+
 struct RegionLockCFDataTrait
 {
     struct Key
@@ -88,14 +152,14 @@ struct RegionLockCFDataTrait
     };
     using DecodedLockCFValue = RecordKVFormat::DecodedLockCFValue;
     using Value = std::tuple<std::shared_ptr<const TiKVKey>, std::shared_ptr<const TiKVValue>, std::shared_ptr<const DecodedLockCFValue>>;
-    using Map = std::unordered_map<Key, Value, Key::Hash>;
+    using Map = std::unordered_map<Key, Value, Key::Hash, std::equal_to<Key>, MyAlloc<std::pair<const Key, Value>>>;
 
     static Map::value_type genKVPair(TiKVKey && key_, TiKVValue && value_)
     {
         auto key = std::make_shared<const TiKVKey>(std::move(key_));
         auto value = std::make_shared<const TiKVValue>(std::move(value_));
         return {{key, std::string_view(key->data(), key->dataSize())},
-            Value{key, value, std::make_shared<const DecodedLockCFValue>(key, value)}};
+            Value{key, value, nullptr}};
     }
 };