Pick getLockByKey to master (#8696)

close #8697
pingcap · Jan 17, 2024 · 684d89a · 684d89a
1 parent ec0c363
commit 684d89a
Show file tree

Hide file tree

Showing 7 changed files with 92 additions and 1 deletion.
diff --git a/contrib/tiflash-proxy b/contrib/tiflash-proxy
diff --git a/dbms/src/Storages/KVStore/FFI/ProxyFFI.cpp b/dbms/src/Storages/KVStore/FFI/ProxyFFI.cpp
@@ -989,4 +989,37 @@ BaseBuffView cppStringAsBuff(const std::string & s)
 {
     return BaseBuffView{.data = s.data(), .len = s.size()};
 }
+
+BaseBuffView GetLockByKey(const EngineStoreServerWrap * server, uint64_t region_id, BaseBuffView key)
+{
+    auto tikv_key = TiKVKey(key.data, key.len);
+    try
+    {
+        auto & kvstore = server->tmt->getKVStore();
+        auto region = kvstore->getRegion(region_id);
+        auto value = region->getLockByKey(tikv_key);
+        if (!value)
+        {
+            // key not exist
+            LOG_WARNING(
+                Logger::get(),
+                "Failed to get lock by key {}, region_id={}",
+                tikv_key.toDebugString(),
+                region_id);
+            return BaseBuffView{};
+        }
+
+        return BaseBuffView{value->data(), value->dataSize()};
+    }
+    catch (...)
+    {
+        LOG_WARNING( //
+            Logger::get(),
+            "Failed to get lock by key {}, region_id={}",
+            tikv_key.toDebugString(),
+            region_id);
+        return BaseBuffView{};
+    }
+}
+
 } // namespace DB
diff --git a/dbms/src/Storages/KVStore/FFI/ProxyFFI.h b/dbms/src/Storages/KVStore/FFI/ProxyFFI.h
@@ -172,6 +172,7 @@ RawCppPtr PreHandleSnapshot(
 void ApplyPreHandledSnapshot(EngineStoreServerWrap * server, void * res, RawCppPtrType type);
 void AbortPreHandledSnapshot(EngineStoreServerWrap * server, uint64_t region_id, uint64_t peer_id);
 void ReleasePreHandledSnapshot(EngineStoreServerWrap * server, void * res, RawCppPtrType type);
+BaseBuffView GetLockByKey(const EngineStoreServerWrap * server, uint64_t region_id, BaseBuffView key);
 HttpRequestRes HandleHttpRequest(EngineStoreServerWrap *, BaseBuffView path, BaseBuffView query, BaseBuffView body);
 uint8_t CheckHttpUriAvailable(BaseBuffView);
 void GcRawCppPtr(void * ptr, RawCppPtrType type);

diff --git a/dbms/src/Storages/KVStore/MultiRaft/RegionData.cpp b/dbms/src/Storages/KVStore/MultiRaft/RegionData.cpp
@@ -262,6 +262,28 @@ DecodedLockCFValuePtr RegionData::getLockInfo(const RegionLockReadQuery & query)
     return nullptr;
 }
 
+std::shared_ptr<const TiKVValue> RegionData::getLockByKey(const TiKVKey & key) const
+{
+    const auto & map = lock_cf.getData();
+    const auto & lock_key = RegionLockCFDataTrait::Key{nullptr, std::string_view(key.data(), key.dataSize())};
+    if (auto lock_it = map.find(lock_key); lock_it != map.end())
+    {
+        const auto & [tikv_key, tikv_val, lock_info_ptr] = lock_it->second;
+        std::ignore = tikv_key;
+        std::ignore = lock_info_ptr;
+        return tikv_val;
+    }
+
+    // It is safe to ignore the missing lock key after restart, print a warning log and return nullptr
+    LOG_WARNING(
+        Logger::get(),
+        "Failed to get lock by key in region data, key={} map_size={} count={}",
+        key.toDebugString(),
+        map.size(),
+        map.count(lock_key));
+    return nullptr;
+}
+
 void RegionData::splitInto(const RegionRange & range, RegionData & new_region_data)
 {
     size_t size_changed = 0;

diff --git a/dbms/src/Storages/KVStore/MultiRaft/RegionData.h b/dbms/src/Storages/KVStore/MultiRaft/RegionData.h
@@ -55,6 +55,8 @@ class RegionData
 
     DecodedLockCFValuePtr getLockInfo(const RegionLockReadQuery & query) const;
 
+    std::shared_ptr<const TiKVValue> getLockByKey(const TiKVKey & key) const;
+
     void splitInto(const RegionRange & range, RegionData & new_region_data);
     void mergeFrom(const RegionData & ori_region_data);
 

diff --git a/dbms/src/Storages/KVStore/Region.h b/dbms/src/Storages/KVStore/Region.h
@@ -243,6 +243,7 @@ class Region : public std::enable_shared_from_this<Region>
         UInt64 term,
         TMTContext & tmt);
 
+    std::shared_ptr<const TiKVValue> getLockByKey(const TiKVKey & key) { return data.getLockByKey(key); }
     UInt64 getSnapshotEventFlag() const { return snapshot_event_flag; }
 
     // IngestSST will first be applied to the `temp_region`, then we need to

diff --git a/dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp b/dbms/src/Storages/KVStore/tests/gtest_new_kvstore.cpp
@@ -27,6 +27,38 @@ namespace DB
 namespace tests
 {
 
+TEST_F(RegionKVStoreTest, RegionStruct)
+try
+{
+    auto & ctx = TiFlashTestEnv::getGlobalContext();
+    initStorages();
+    MockRaftStoreProxy::FailCond cond;
+    KVStore & kvs = getKVS();
+    auto table_id = proxy_instance->bootstrapTable(ctx, kvs, ctx.getTMTContext());
+    auto start = RecordKVFormat::genKey(table_id, 0);
+    auto end = RecordKVFormat::genKey(table_id, 100);
+    auto str_key = RecordKVFormat::genKey(table_id, 1, 111);
+    auto [str_val_write, str_val_default] = proxy_instance->generateTiKVKeyValue(111, 999);
+    auto str_lock_value
+        = RecordKVFormat::encodeLockCfValue(RecordKVFormat::CFModifyFlag::PutFlag, "PK", 111, 999).toString();
+    proxy_instance->bootstrapWithRegion(kvs, ctx.getTMTContext(), 1, std::nullopt);
+    {
+        auto kvr1 = kvs.getRegion(1);
+        auto [index, term] = proxy_instance->rawWrite(
+            1,
+            {str_key, str_key},
+            {str_lock_value, str_val_default},
+            {WriteCmdType::Put, WriteCmdType::Put},
+            {ColumnFamilyType::Lock, ColumnFamilyType::Default});
+        UNUSED(term);
+        proxy_instance->doApply(kvs, ctx.getTMTContext(), cond, 1, index);
+        ASSERT_EQ(kvr1->getLockByKey(str_key)->dataSize(), str_lock_value.size());
+        ASSERT_EQ(kvr1->getLockByKey(RecordKVFormat::genKey(table_id, 1, 112)), nullptr);
+    }
+}
+CATCH
+
+
 TEST_F(RegionKVStoreTest, MemoryTracker)
 try
 {