├── cmake ├── find_kvproto.cmake ├── find_grpc.cmake ├── find_protobuf.cmake ├── find_poco.cmake ├── find_gtest.cmake └── Modules │ ├── FindPackageMessage.cmake │ └── FindPoco.cmake ├── include └── pingcap │ ├── Log.h │ ├── pd │ ├── Types.h │ ├── Oracle.h │ ├── IClient.h │ ├── CodecClient.h │ ├── MockPDClient.h │ └── Client.h │ ├── common │ ├── FixedThreadPool.h │ ├── MPMCQueue.h │ └── MPPProber.h │ ├── RedactHelpers.h │ ├── kv │ ├── Snapshot.h │ ├── internal │ │ ├── conn.h │ │ └── type_traits.h │ ├── Scanner.h │ ├── Txn.h │ ├── Cluster.h │ ├── Rpc.h │ ├── Backoff.h │ ├── 2pc.h │ ├── RegionCache.h │ ├── RegionClient.h │ └── LockResolver.h │ ├── SetThreadName.h │ ├── Config.h │ ├── Exception.h │ └── coprocessor │ └── Client.h ├── src ├── test │ ├── real_tikv_test │ │ ├── CMakeLists.txt │ │ └── 2pc_test.cc │ ├── CMakeLists.txt │ ├── bank_test │ │ ├── CMakeLists.txt │ │ ├── bank_test_ut.cc │ │ ├── schrodinger_client.h │ │ ├── bank_test_schrodinger.cc │ │ └── bank_test.h │ ├── io_or_region_error_get_test.cc │ ├── region_split_test.cc │ ├── mock_tikv.h │ ├── test_helper.h │ ├── lock_resolve_test.cc │ └── coprocessor_test.cc ├── kv │ ├── Rpc.cc │ ├── Cluster.cc │ ├── Backoff.cc │ ├── Snapshot.cc │ ├── Scanner.cc │ └── RegionClient.cc ├── CMakeLists.txt ├── common │ ├── FixedThreadPool.cc │ └── MPPProber.cc └── RedactHelpers.cc ├── ci ├── Makefile ├── build-test.sh └── Dockerfile ├── OWNERS ├── .gitmodules ├── format.sh ├── README.md ├── .gitignore ├── CMakeLists.txt ├── third_party ├── CMakeLists.txt └── abseil-cpp-cmake │ └── CMakeLists.txt ├── OWNERS_ALIASES ├── .clang-format └── .clang-tidy /cmake/find_kvproto.cmake: -------------------------------------------------------------------------------- 1 | message(STATUS "Using kvproto: ${kvClient_SOURCE_DIR}/third_party/kvproto/cpp") 2 | -------------------------------------------------------------------------------- /include/pingcap/Log.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | 4 | namespace pingcap 5 | { 6 | using Poco::Logger; 7 | } 8 | -------------------------------------------------------------------------------- /cmake/find_grpc.cmake: -------------------------------------------------------------------------------- 1 | 2 | find_package(gRPC REQUIRED) 3 | message(STATUS "Using gRPC: ${gRPC_VERSION} : ${gRPC_INCLUDE_DIRS}, ${gRPC_LIBRARIES}, ${gRPC_CPP_PLUGIN}") 4 | -------------------------------------------------------------------------------- /src/test/real_tikv_test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_executable(2pc_test 2pc_test.cc) 2 | target_include_directories(2pc_test PUBLIC ${test_includes}) 3 | target_link_libraries(2pc_test ${test_libs}) 4 | -------------------------------------------------------------------------------- /cmake/find_protobuf.cmake: -------------------------------------------------------------------------------- 1 | find_package(Protobuf REQUIRED) 2 | message(STATUS "Using protobuf: ${Protobuf_VERSION} : ${Protobuf_INCLUDE_DIR}, ${Protobuf_LIBRARY}, ${Protobuf_PROTOC_EXECUTABLE}") 3 | -------------------------------------------------------------------------------- /cmake/find_poco.cmake: -------------------------------------------------------------------------------- 1 | find_package (Poco REQUIRED Foundation Net JSON Util) 2 | 3 | if (Poco_FOUND) 4 | message(STATUS "Using Poco: ${Poco_VERSION} : ${Poco_INCLUDE_DIRS}, ${Poco_LIBRARIES}") 5 | else () 6 | message(STATUS "Poco Not Found") 7 | endif() 8 | -------------------------------------------------------------------------------- /ci/Makefile: -------------------------------------------------------------------------------- 1 | default: image build-test 2 | 3 | BUILDER_IMAGE := $(or $(BUILDER_IMAGE),hub.pingcap.net/tikv/client-c) 4 | 5 | image: 6 | docker build -f Dockerfile -t $(BUILDER_IMAGE) . 7 | build-test: 8 | docker run --rm -v $(realpath ..):/client-c $(BUILDER_IMAGE) /client-c/ci/build-test.sh 9 | -------------------------------------------------------------------------------- /OWNERS: -------------------------------------------------------------------------------- 1 | # See the OWNERS docs at https://www.kubernetes.dev/docs/guide/owners/#owners 2 | # The members of 'sig-community-*' are synced from memberships defined in repository: https://github.com/tikv/community. 3 | filters: 4 | .*: 5 | approvers: 6 | - sig-community-approvers 7 | reviewers: 8 | - sig-community-reviewers 9 | -------------------------------------------------------------------------------- /include/pingcap/pd/Types.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | namespace pingcap::pd 6 | { 7 | 8 | using KeyspaceID = uint32_t; 9 | 10 | enum : KeyspaceID 11 | { 12 | // The size of KeyspaceID allocated for PD is 3 bytes. 13 | // The NullspaceID is preserved for TiDB API V1 compatibility. 14 | NullspaceID = 0xffffffff, 15 | }; 16 | 17 | } // namespace pingcap::pd 18 | -------------------------------------------------------------------------------- /cmake/find_gtest.cmake: -------------------------------------------------------------------------------- 1 | if (NOT EXISTS "${kvClient_SOURCE_DIR}/third_party/googletest/CMakeLists.txt") 2 | set(USE_INTERNAL_GTEST_LIBRARY 0) 3 | find_package(PkgConfig) 4 | pkg_search_module(GTEST REQUIRED gtest_main) 5 | else () 6 | set(BUILD_GMOCK 0) 7 | set(INSTALL_GTEST 0) 8 | set(USE_INTERNAL_GTEST_LIBRARY 1) 9 | endif() 10 | 11 | message(STATUS "Using Gtest: ${gtest_INCLUDE_DIRS} ${USE_INTERNAL_GTEST_LIBRARY}") 12 | -------------------------------------------------------------------------------- /ci/build-test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -xe 4 | 5 | 6 | SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" 7 | SRCPATH=$(cd $SCRIPTPATH/..; pwd -P) 8 | NPROC=$(nproc || grep -c ^processor /proc/cpuinfo) 9 | 10 | build_dir="$SRCPATH/build" 11 | mkdir -p $build_dir && cd $build_dir 12 | cmake "$SRCPATH" \ 13 | -DENABLE_TESTS=on 14 | make -j $NPROC 15 | 16 | nohup /mock-tikv/bin/mock-tikv & 17 | mock_kv_pid=$! 18 | 19 | cd "$build_dir" && make test 20 | 21 | kill -9 $mock_kv_pid 22 | 23 | 24 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "third_party/kvproto"] 2 | path = third_party/kvproto 3 | url = https://github.com/pingcap/kvproto.git 4 | [submodule "third_party/googletest"] 5 | path = third_party/googletest 6 | url = https://github.com/google/googletest.git 7 | [submodule "third_party/libfiu"] 8 | path = third_party/libfiu 9 | url = https://github.com/albertito/libfiu.git 10 | [submodule "third_party/abseil-cpp"] 11 | path = third_party/abseil-cpp 12 | url = https://github.com/abseil/abseil-cpp.git 13 | -------------------------------------------------------------------------------- /format.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | clang-format -i ./include/pingcap/kv/*.h 3 | clang-format -i ./include/pingcap/coprocessor/*.h 4 | clang-format -i ./include/pingcap/kv/internal/*.h 5 | clang-format -i ./include/pingcap/pd/*.h 6 | clang-format -i ./include/pingcap/common/*.h 7 | clang-format -i ./include/pingcap/*.h 8 | clang-format -i ./src/kv/*.cc 9 | clang-format -i ./src/coprocessor/*.cc 10 | clang-format -i ./src/pd/*.cc 11 | clang-format -i ./src/common/*.cc 12 | clang-format -i ./src/test/*.cc 13 | clang-format -i ./src/test/*.h 14 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Overview 2 | 3 | TiKV Client C++ allow you to access [TiKV](https://github.com/tikv/tikv) from C and C++ applications. 4 | 5 | The code is in an early alpha status. Currently, it is only used by TiFlash. 6 | 7 | # License 8 | 9 | [Apache-2.0 License](/LICENSE) 10 | 11 | # Docs 12 | 13 | The docs can be found [here](https://tikv.org/docs/dev/develop/clients/cpp/). 14 | 15 | # Building 16 | 17 | Install dependencies (adjust for your platform): 18 | ``` 19 | sudo dnf install cmake grpc-devel poco-devel abseil-cpp-devel gcc-c++ 20 | ``` 21 | 22 | Update submodules 23 | ``` 24 | git submodule update --init --recursive 25 | ``` 26 | 27 | Build: 28 | ``` 29 | mkdir build 30 | cd build 31 | cmake .. 32 | make 33 | ``` 34 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Prerequisites 2 | *.d 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | 34 | # build directory 35 | build 36 | cmake-build* 37 | 38 | # IDE or editor files 39 | .idea 40 | *.iml 41 | *.swp 42 | *.swo 43 | tags 44 | .vscode 45 | 46 | # vscode clangd cache 47 | .cache 48 | 49 | # JSON Compilation Database Format Specification 50 | # https://clang.llvm.org/docs/JSONCompilationDatabase.html 51 | compile_commands.json 52 | -------------------------------------------------------------------------------- /src/test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | list (APPEND test_includes ${kvClient_INCLUDE_DIR} ${fiu_include_dirs}) 2 | 3 | list(APPEND test_libs 4 | kvproto 5 | kv_client 6 | ${Protobuf_LIBRARY} 7 | ${gRPC_LIBRARIES} 8 | ${Poco_Net_LIBRARY} 9 | ${Poco_JSON_LIBRARY} 10 | gtest_main 11 | fiu) 12 | 13 | add_executable(kv_client_ut 14 | io_or_region_error_get_test.cc 15 | region_split_test.cc 16 | lock_resolve_test.cc 17 | coprocessor_test.cc 18 | batch_coprocessor_test.cc 19 | ) 20 | target_include_directories(kv_client_ut PUBLIC ${test_includes}) 21 | target_link_libraries(kv_client_ut ${test_libs} ) 22 | 23 | include(CTest) 24 | add_test(kv_client_test kv_client_ut) 25 | 26 | add_subdirectory(bank_test) 27 | add_subdirectory(real_tikv_test) 28 | -------------------------------------------------------------------------------- /src/test/bank_test/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | add_executable(bank_test_ut bank_test_ut.cc) 2 | target_include_directories(bank_test_ut PUBLIC ${test_includes}) 3 | target_link_libraries(bank_test_ut ${test_libs}) 4 | 5 | add_test(bank_test bank_test_ut) 6 | 7 | add_executable(bank_test_schrodinger bank_test_schrodinger.cc) 8 | set(CMAKE_SHARED_LINKER_FLAGS "-static") 9 | target_include_directories(bank_test_schrodinger PUBLIC ${test_includes}) 10 | 11 | if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU") 12 | target_link_libraries(bank_test_schrodinger -static-libgcc -static-libstdc++ ${test_libs} ${Poco_Util_LIBRARY} ${Poco_XML_LIBRARY} ${Poco_JSON_LIBRARY}) 13 | else() 14 | target_link_libraries(bank_test_schrodinger ${test_libs} ${Poco_Util_LIBRARY} ${Poco_XML_LIBRARY} ${Poco_JSON_LIBRARY}) 15 | endif() 16 | 17 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.10) 2 | project(kvClient) 3 | set (CMAKE_CXX_STANDARD 17) 4 | set (CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-narrowing") 5 | set (CMAKE_EXPORT_COMPILE_COMMANDS ON) 6 | set (CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_CURRENT_SOURCE_DIR}/cmake/Modules/") 7 | enable_testing() 8 | 9 | if (NOT Protobuf_INCLUDE_DIR) 10 | include (cmake/find_protobuf.cmake) 11 | endif () 12 | 13 | if (NOT gRPC_FOUND) 14 | include (cmake/find_grpc.cmake) 15 | endif () 16 | 17 | if (NOT Poco_Foundation_LIBRARY) 18 | include (cmake/find_poco.cmake) 19 | endif() 20 | 21 | if (NOT KVPROTO_FOUND) 22 | include (cmake/find_kvproto.cmake) 23 | endif () 24 | 25 | if (NOT GTEST_FOUND AND ENABLE_TESTS) 26 | include (cmake/find_gtest.cmake) 27 | endif() 28 | 29 | add_subdirectory (third_party) 30 | add_subdirectory (src) 31 | -------------------------------------------------------------------------------- /ci/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:23.04 2 | 3 | RUN apt update -y \ 4 | && apt install -y cmake build-essential \ 5 | wget git \ 6 | protobuf-compiler libprotobuf-dev libgrpc-dev libgrpc++-dev libc-ares-dev protobuf-compiler-grpc libpoco-dev 7 | 8 | RUN rm -rf /var/lib/apt/lists/* 9 | 10 | #back to root dir and download golang 11 | RUN cd / 12 | 13 | ENV GOLANG_VERSION 1.13.3 14 | 15 | RUN wget -O go.tgz "https://dl.google.com/go/go$GOLANG_VERSION.linux-amd64.tar.gz"; \ 16 | tar -C /usr/local -xzf go.tgz; \ 17 | rm go.tgz; \ 18 | export PATH="/usr/local/go/bin:$PATH"; \ 19 | go version 20 | 21 | ENV GOPATH /go 22 | ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH 23 | 24 | RUN cd / 25 | 26 | RUN git clone https://github.com/tikv/mock-tikv.git && cd mock-tikv && git checkout 60d5921028afd72e1aeba880b9052c40e932eef3 && make failpoint-enable && make 27 | -------------------------------------------------------------------------------- /third_party/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | 2 | if (NOT ABSL_ROOT_DIR) 3 | add_subdirectory(abseil-cpp-cmake) 4 | endif() 5 | 6 | if (NOT KVPROTO_FOUND) 7 | add_subdirectory (kvproto/cpp) 8 | endif() 9 | 10 | if (USE_INTERNAL_GTEST_LIBRARY AND ENABLE_TESTS) 11 | set(GTEST_INSTALL OFF) 12 | add_subdirectory (googletest) 13 | endif() 14 | 15 | set(fiu_include_dirs libfiu/libfiu) 16 | list(APPEND fiu_sources libfiu/libfiu/fiu.c) 17 | list(APPEND fiu_sources libfiu/libfiu/fiu-rc.c) 18 | list(APPEND fiu_sources libfiu/libfiu/backtrace.c) 19 | list(APPEND fiu_sources libfiu/libfiu/wtable.c) 20 | list(APPEND fiu_sources libfiu/libfiu/hash.c) 21 | # We need to compile fiu manully because fiu does not provide their cmakefiles. 22 | add_library(fiu ${fiu_sources}) 23 | target_include_directories(fiu PUBLIC ${fiu_include_dirs}) 24 | target_compile_options(fiu PUBLIC -fPIC -DDUMMY_BACKTRACE) 25 | -------------------------------------------------------------------------------- /include/pingcap/common/FixedThreadPool.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | namespace pingcap 11 | { 12 | namespace common 13 | { 14 | 15 | class FixedThreadPool 16 | { 17 | public: 18 | using Task = std::function; 19 | explicit FixedThreadPool(size_t num_) 20 | : num(num_) 21 | , stopped(false) 22 | {} 23 | ~FixedThreadPool() = default; 24 | 25 | void start(); 26 | 27 | void enqueue(const Task & task); 28 | 29 | void stop(); 30 | 31 | private: 32 | void loop(); 33 | 34 | size_t num; 35 | bool stopped; 36 | std::vector threads; 37 | std::queue tasks; 38 | std::mutex mu; 39 | std::condition_variable cond; 40 | }; 41 | 42 | } // namespace common 43 | } // namespace pingcap 44 | -------------------------------------------------------------------------------- /include/pingcap/RedactHelpers.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | namespace pingcap 7 | { 8 | enum class RedactMode 9 | { 10 | Disable, 11 | Enable, 12 | Marker, 13 | }; 14 | class Redact 15 | { 16 | public: 17 | static void setRedactLog(RedactMode v); 18 | 19 | // Format as a hex string for debugging. The value will be converted to '?' if `REDACT_LOG` is true 20 | static std::string keyToDebugString(const char * key, size_t size); 21 | static std::string keyToDebugString(const std::string & key) { return Redact::keyToDebugString(key.data(), key.size()); } 22 | 23 | static std::string keyToHexString(const char * key, size_t size); 24 | 25 | protected: 26 | Redact() = default; 27 | 28 | private: 29 | // Log user data to log only when this flag is set to false. 30 | static std::atomic REDACT_LOG; 31 | }; 32 | 33 | } // namespace pingcap 34 | -------------------------------------------------------------------------------- /include/pingcap/kv/Snapshot.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | namespace pingcap 7 | { 8 | namespace kv 9 | { 10 | struct Scanner; 11 | 12 | struct Snapshot 13 | { 14 | Cluster * cluster; 15 | const int64_t version; 16 | MinCommitTSPushed min_commit_ts_pushed; 17 | 18 | Snapshot(Cluster * cluster_, uint64_t version_) 19 | : cluster(cluster_) 20 | , version(version_) 21 | {} 22 | explicit Snapshot(Cluster * cluster_) 23 | : cluster(cluster_) 24 | , version(cluster_->pd_client->getTS()) 25 | {} 26 | 27 | std::string Get(const std::string & key); 28 | std::string Get(Backoffer & bo, const std::string & key); 29 | 30 | kvrpcpb::MvccInfo mvccGet(const std::string & key); 31 | 32 | kvrpcpb::MvccInfo mvccGet(Backoffer & bo, const std::string & key); 33 | 34 | Scanner Scan(const std::string & begin, const std::string & end); 35 | }; 36 | 37 | } // namespace kv 38 | } // namespace pingcap 39 | -------------------------------------------------------------------------------- /include/pingcap/SetThreadName.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #if defined(__APPLE__) 4 | #include 5 | #elif defined(__FreeBSD__) 6 | #include 7 | #include 8 | #else 9 | #include 10 | #endif 11 | 12 | #include 13 | 14 | #include 15 | #include 16 | 17 | namespace pingcap 18 | { 19 | inline void SetThreadName(const char * tname) 20 | { 21 | constexpr static auto MAX_LEN = 15; // thread name will be tname[:MAX_LEN] 22 | if (std::strlen(tname) > MAX_LEN) 23 | std::cerr << "set thread name " << tname << " is too long and will be truncated by system\n"; 24 | 25 | #if defined(__FreeBSD__) 26 | pthread_set_name_np(pthread_self(), tname); 27 | return; 28 | 29 | #elif defined(__APPLE__) 30 | if (0 != pthread_setname_np(tname)) 31 | #else 32 | if (0 != prctl(PR_SET_NAME, tname, 0, 0, 0)) 33 | #endif 34 | throw pingcap::Exception("Cannot set thread name " + std::string(tname), ErrorCodes::LogicalError); 35 | } 36 | } // namespace pingcap 37 | -------------------------------------------------------------------------------- /src/kv/Rpc.cc: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | namespace pingcap 4 | { 5 | namespace kv 6 | { 7 | ConnArray::ConnArray(size_t max_size, const std::string & addr, const ClusterConfig & config_) 8 | : address(addr) 9 | , index(0) 10 | { 11 | vec.resize(max_size); 12 | for (size_t i = 0; i < max_size; i++) 13 | { 14 | vec[i] = std::make_shared(addr, config_); 15 | } 16 | } 17 | 18 | std::shared_ptr ConnArray::get() 19 | { 20 | std::lock_guard lock(mutex); 21 | index = (index + 1) % vec.size(); 22 | return vec[index]; 23 | } 24 | 25 | ConnArrayPtr RpcClient::getConnArray(const std::string & addr) 26 | { 27 | std::lock_guard lock(mutex); 28 | auto it = conns.find(addr); 29 | if (it == conns.end()) 30 | { 31 | return createConnArray(addr); 32 | } 33 | return it->second; 34 | } 35 | 36 | ConnArrayPtr RpcClient::createConnArray(const std::string & addr) 37 | { 38 | auto conn_array = std::make_shared(5, addr, config); 39 | conns[addr] = conn_array; 40 | return conn_array; 41 | } 42 | 43 | } // namespace kv 44 | } // namespace pingcap 45 | -------------------------------------------------------------------------------- /third_party/abseil-cpp-cmake/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(ABSL_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../abseil-cpp") 2 | set(ABSL_ROOT_DIR "${ABSL_ROOT_DIR}" PARENT_SCOPE) 3 | if(NOT EXISTS "${ABSL_ROOT_DIR}/CMakeLists.txt") 4 | message(FATAL_ERROR " submodule third_party/abseil-cpp is missing. To fix try run: \n git submodule update --init --recursive") 5 | endif() 6 | set(BUILD_TESTING OFF) 7 | set(ABSL_PROPAGATE_CXX_STD ON) 8 | 9 | add_subdirectory("${ABSL_ROOT_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}/../abseil-cpp") 10 | 11 | add_library(abseil_swiss_tables INTERFACE) 12 | set_target_properties(abseil_swiss_tables PROPERTIES INTERFACE_COMPILE_OPTIONS "-Wno-array-parameter") 13 | 14 | target_link_libraries(abseil_swiss_tables INTERFACE 15 | absl::flat_hash_map 16 | absl::flat_hash_set 17 | ) 18 | 19 | get_target_property(FLAT_HASH_MAP_INCLUDE_DIR absl::flat_hash_map INTERFACE_INCLUDE_DIRECTORIES) 20 | target_include_directories (abseil_swiss_tables SYSTEM BEFORE INTERFACE ${FLAT_HASH_MAP_INCLUDE_DIR}) 21 | 22 | get_target_property(FLAT_HASH_SET_INCLUDE_DIR absl::flat_hash_set INTERFACE_INCLUDE_DIRECTORIES) 23 | target_include_directories (abseil_swiss_tables SYSTEM BEFORE INTERFACE ${FLAT_HASH_SET_INCLUDE_DIR}) 24 | 25 | message(STATUS "Using absl: dir=${ABSL_ROOT_DIR}") 26 | -------------------------------------------------------------------------------- /src/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(kvClient_sources) 2 | 3 | list(APPEND kvClient_sources kv/2pc.cc) 4 | list(APPEND kvClient_sources kv/Backoff.cc) 5 | list(APPEND kvClient_sources kv/Cluster.cc) 6 | list(APPEND kvClient_sources kv/LockResolver.cc) 7 | list(APPEND kvClient_sources kv/RegionCache.cc) 8 | list(APPEND kvClient_sources kv/RegionClient.cc) 9 | list(APPEND kvClient_sources kv/Rpc.cc) 10 | list(APPEND kvClient_sources kv/Snapshot.cc) 11 | list(APPEND kvClient_sources kv/Scanner.cc) 12 | list(APPEND kvClient_sources pd/Client.cc) 13 | list(APPEND kvClient_sources coprocessor/Client.cc) 14 | list(APPEND kvClient_sources RedactHelpers.cc) 15 | list(APPEND kvClient_sources common/FixedThreadPool.cc) 16 | list(APPEND kvClient_sources common/MPPProber.cc) 17 | 18 | set(kvClient_INCLUDE_DIR ${kvClient_SOURCE_DIR}/include) 19 | 20 | # Only enable fiu under test mode. 21 | if (ENABLE_TESTS) 22 | add_definitions(-DFIU_ENABLE) 23 | endif() 24 | 25 | add_library(kv_client ${kvClient_sources}) 26 | target_include_directories(kv_client PUBLIC ${kvClient_INCLUDE_DIR} ${fiu_include_dirs} ${Protobuf_INCLUDE_DIR} ${gRPC_INCLUDE_DIRS}) 27 | target_link_libraries(kv_client kvproto ${Poco_Foundation_LIBRARY} fiu ${Protobuf_LIBRARY} ${gRPC_LIBRARIES}) 28 | 29 | if (ENABLE_TESTS) 30 | add_subdirectory (test) 31 | endif() 32 | -------------------------------------------------------------------------------- /include/pingcap/kv/internal/conn.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | #include 5 | 6 | namespace pingcap 7 | { 8 | namespace kv 9 | { 10 | // create and destroy stub but not destroy channel may case memory leak, so we 11 | // bound channel and stub in same struct. 12 | struct KvConnClient 13 | { 14 | std::shared_ptr channel; 15 | std::unique_ptr stub; 16 | 17 | KvConnClient(std::string addr, const ClusterConfig & config) 18 | { 19 | grpc::ChannelArguments ch_args; 20 | // set max size that grpc client can receive to max value. 21 | ch_args.SetMaxReceiveMessageSize(-1); 22 | ch_args.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, 60 * 1000); 23 | ch_args.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, 120 * 1000); 24 | ch_args.SetInt(GRPC_ARG_MIN_RECONNECT_BACKOFF_MS, 1 * 1000); 25 | ch_args.SetInt(GRPC_ARG_MAX_RECONNECT_BACKOFF_MS, 3 * 1000); 26 | ch_args.SetInt(GRPC_ARG_INITIAL_RECONNECT_BACKOFF_MS, 1 * 1000); 27 | std::shared_ptr cred; 28 | if (config.hasTlsConfig()) 29 | { 30 | cred = grpc::SslCredentials(config.getGrpcCredentials()); 31 | } 32 | else 33 | { 34 | cred = grpc::InsecureChannelCredentials(); 35 | } 36 | channel = grpc::CreateCustomChannel(addr, cred, ch_args); 37 | stub = tikvpb::Tikv::NewStub(channel); 38 | } 39 | }; 40 | 41 | } // namespace kv 42 | } // namespace pingcap 43 | -------------------------------------------------------------------------------- /src/kv/Cluster.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | #include 5 | 6 | namespace pingcap 7 | { 8 | namespace kv 9 | { 10 | void Cluster::splitRegion(const std::string & split_key) 11 | { 12 | Backoffer bo(splitRegionBackoff); 13 | auto loc = region_cache->locateKey(bo, split_key); 14 | RegionClient client(this, loc.region); 15 | kvrpcpb::SplitRegionRequest req; 16 | req.set_split_key(split_key); 17 | kvrpcpb::SplitRegionResponse resp; 18 | client.sendReqToRegion(bo, req, &resp); 19 | if (resp.has_region_error()) 20 | { 21 | throw Exception(resp.region_error().message(), RegionUnavailable); 22 | } 23 | auto lr = resp.left(); 24 | auto rr = resp.right(); 25 | region_cache->dropRegion(loc.region); 26 | region_cache->getRegionByID(bo, RegionVerID(lr.id(), lr.region_epoch().conf_ver(), lr.region_epoch().version())); 27 | region_cache->getRegionByID(bo, RegionVerID(rr.id(), rr.region_epoch().conf_ver(), rr.region_epoch().version())); 28 | } 29 | 30 | void Cluster::startBackgroundTasks() 31 | { 32 | thread_pool->start(); 33 | 34 | thread_pool->enqueue([this] { 35 | mpp_prober->run(); 36 | }); 37 | if (region_cache) 38 | { 39 | // region_cache may not be inited if pd addr is not setup. 40 | // So skip update region cache periodically. 41 | thread_pool->enqueue([this] { 42 | region_cache->updateCachePeriodically(); 43 | }); 44 | } 45 | } 46 | 47 | } // namespace kv 48 | } // namespace pingcap 49 | -------------------------------------------------------------------------------- /src/common/FixedThreadPool.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | namespace pingcap 6 | { 7 | namespace common 8 | { 9 | 10 | void FixedThreadPool::start() 11 | { 12 | for (size_t i = 0; i < num; ++i) 13 | { 14 | threads.push_back(std::thread(&FixedThreadPool::loop, this)); 15 | } 16 | } 17 | 18 | void FixedThreadPool::loop() 19 | { 20 | auto & log = Logger::get("pingcap/fixed_thread_pool"); 21 | while (true) 22 | { 23 | Task task; 24 | { 25 | std::unique_lock lock(mu); 26 | cond.wait(lock, [this] { 27 | return !tasks.empty() || stopped; 28 | }); 29 | 30 | if (stopped) 31 | return; 32 | 33 | task = tasks.front(); 34 | tasks.pop(); 35 | } 36 | 37 | try 38 | { 39 | task(); 40 | } 41 | catch (...) 42 | { 43 | log.warning(getCurrentExceptionMsg("FixedThreadPool task failed: ")); 44 | } 45 | } 46 | } 47 | 48 | void FixedThreadPool::enqueue(const Task & task) 49 | { 50 | { 51 | std::unique_lock lock(mu); 52 | tasks.push(task); 53 | } 54 | cond.notify_one(); 55 | } 56 | 57 | void FixedThreadPool::stop() 58 | { 59 | { 60 | std::unique_lock lock(mu); 61 | stopped = true; 62 | } 63 | cond.notify_all(); 64 | for (auto & thr : threads) 65 | { 66 | thr.join(); 67 | } 68 | threads.clear(); 69 | } 70 | 71 | } // namespace common 72 | } // namespace pingcap 73 | -------------------------------------------------------------------------------- /src/test/io_or_region_error_get_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | #include "mock_tikv.h" 9 | #include "test_helper.h" 10 | 11 | namespace pingcap::tests 12 | { 13 | using namespace pingcap; 14 | using namespace pingcap::kv; 15 | 16 | class TestWithMockKV : public testing::TestWithParam> 17 | { 18 | public: 19 | void SetUp() override 20 | { 21 | mock_kv_cluster = mockkv::initCluster(); 22 | std::vector pd_addrs = mock_kv_cluster->pd_addrs; 23 | 24 | test_cluster = createCluster(pd_addrs); 25 | 26 | std::tie(fail_point, fail_arg) = GetParam(); 27 | } 28 | 29 | void TearDown() override {} 30 | 31 | mockkv::ClusterPtr mock_kv_cluster; 32 | 33 | ClusterPtr test_cluster; 34 | 35 | std::string_view fail_point; 36 | std::string_view fail_arg; 37 | }; 38 | 39 | TEST_P(TestWithMockKV, testGetInjectError) 40 | { 41 | Txn txn(test_cluster.get()); 42 | txn.set("abc", "edf"); 43 | txn.commit(); 44 | 45 | mock_kv_cluster->updateFailPoint(mock_kv_cluster->stores[0].id, fail_point.data(), fail_arg.data()); 46 | Snapshot snap(test_cluster.get(), test_cluster->pd_client->getTS()); 47 | 48 | std::string result = snap.Get("abc"); 49 | 50 | ASSERT_EQ(result, "edf"); 51 | } 52 | 53 | INSTANTIATE_TEST_SUITE_P(RunGetWithInjectedErr, 54 | TestWithMockKV, 55 | testing::Values(std::make_tuple("server-is-busy", "2*return()"), 56 | std::make_tuple("io-timeout", "8*return()"))); 57 | 58 | } // namespace pingcap::tests 59 | -------------------------------------------------------------------------------- /OWNERS_ALIASES: -------------------------------------------------------------------------------- 1 | # See the OWNERS docs at https://www.kubernetes.dev/docs/guide/owners/#owners_aliases 2 | # The members of 'sig-community-*' are synced from memberships defined in repository: https://github.com/tikv/community. 3 | aliases: 4 | sig-community-reviewers: 5 | - 3AceShowHand 6 | - 3pointer 7 | - CalvinNeo 8 | - Fullstop000 9 | - HuSharp 10 | - Jibbow 11 | - JmPotato 12 | - Leavrth 13 | - Mossaka 14 | - MrCroxx 15 | - Rustin170506 16 | - Xuanwo 17 | - ethercflow 18 | - fredchenbj 19 | - gozssky 20 | - haojinming 21 | - hbisheng 22 | - hhwyt 23 | - jayzhan211 24 | - lcwangchao 25 | - lhy1024 26 | - longfangsong 27 | - lzmhhh123 28 | - mittalrishabh 29 | - nolouch 30 | - rleungx 31 | - tier-cap 32 | - wjhuang2016 33 | - wshwsh12 34 | sig-community-approvers: 35 | - 5kbpers 36 | - AndreMouche 37 | - BusyJay 38 | - Connor1996 39 | - Little-Wallace 40 | - LykxSassinator 41 | - MyonKeminta 42 | - NingLin-P 43 | - SpadeA-Tang 44 | - TennyZhuang 45 | - YuJuncen 46 | - andylokandy 47 | - breezewish 48 | - brson 49 | - bufferflies 50 | - cfzjywxk 51 | - coocood 52 | - crazycs520 53 | - disksing 54 | - ekexium 55 | - gengliqi 56 | - glorv 57 | - hicqu 58 | - hunterlxt 59 | - imtbkcat 60 | - innerr 61 | - iosmanthus 62 | - jackysp 63 | - kennytm 64 | - liuzix 65 | - lonng 66 | - lysu 67 | - marsishandsome 68 | - niedhui 69 | - nrc 70 | - overvenus 71 | - pingyu 72 | - skyzh 73 | - sticnarf 74 | - sunxiaoguang 75 | - tabokie 76 | - tonyxuqqi 77 | - v01dstar 78 | - yiwu-arbug 79 | - you06 80 | - youjiali1995 81 | - zhangjinpeng87 82 | - zhongzc 83 | - zhouqiang-cl 84 | - zyguan 85 | -------------------------------------------------------------------------------- /include/pingcap/kv/Scanner.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | namespace pingcap 8 | { 9 | namespace kv 10 | { 11 | inline std::string prefixNext(const std::string & str) 12 | { 13 | auto new_str = str; 14 | for (int i = int(str.size()); i > 0; i--) 15 | { 16 | char & c = new_str[i - 1]; 17 | c++; 18 | if (c != 0) 19 | { 20 | return new_str; 21 | } 22 | } 23 | return ""; 24 | } 25 | 26 | inline std::string alphabeticalNext(std::string str) 27 | { 28 | str.push_back('\0'); 29 | return str; 30 | } 31 | 32 | struct Scanner 33 | { 34 | Snapshot snap; 35 | std::string next_start_key; 36 | std::string end_key; 37 | int batch; 38 | 39 | std::vector<::kvrpcpb::KvPair> cache; 40 | size_t idx; 41 | bool valid; 42 | bool eof; 43 | 44 | 45 | Logger * log; 46 | 47 | Scanner(Snapshot & snapshot_, std::string start_key_, std::string end_key_, int batch_) 48 | : snap(snapshot_) 49 | , next_start_key(start_key_) 50 | , end_key(end_key_) 51 | , batch(batch_) 52 | , idx(0) 53 | , valid(true) 54 | , eof(false) 55 | , log(&Logger::get("pingcap.tikv")) 56 | { 57 | next(); 58 | } 59 | 60 | void next(); 61 | 62 | std::string key() 63 | { 64 | if (valid) 65 | return cache[idx].key(); 66 | return ""; 67 | } 68 | 69 | std::string value() 70 | { 71 | if (valid) 72 | return cache[idx].value(); 73 | return ""; 74 | } 75 | 76 | private: 77 | void resolveCurrentLock(Backoffer & bo, kvrpcpb::KvPair &); 78 | void getData(Backoffer & bo); 79 | }; 80 | 81 | // end of namespace. 82 | } // namespace kv 83 | } // namespace pingcap 84 | -------------------------------------------------------------------------------- /src/test/bank_test/bank_test_ut.cc: -------------------------------------------------------------------------------- 1 | #include "../mock_tikv.h" 2 | #include "../test_helper.h" 3 | #include "bank_test.h" 4 | #include "fiu-control.h" 5 | #include "fiu.h" 6 | 7 | namespace pingcap::tests 8 | { 9 | 10 | using namespace pingcap; 11 | using namespace pingcap::kv; 12 | 13 | class TestBankLoop : public testing::Test 14 | { 15 | protected: 16 | void SetUp() override 17 | { 18 | mock_kv_cluster = mockkv::initCluster(); 19 | std::vector pd_addrs = mock_kv_cluster->pd_addrs; 20 | 21 | test_cluster = createCluster(pd_addrs); 22 | control_cluster = createCluster(pd_addrs); 23 | 24 | fiu_init(0); 25 | fiu_enable("use_async_commit", 1, nullptr, 0); 26 | } 27 | 28 | mockkv::ClusterPtr mock_kv_cluster; 29 | 30 | ClusterPtr test_cluster; 31 | ClusterPtr control_cluster; 32 | }; 33 | 34 | TEST_F(TestBankLoop, testBankForever) 35 | try 36 | { 37 | BankCase bank(test_cluster.get(), 3000, 6); 38 | 39 | bank.initialize(); 40 | 41 | for (int i = 100; i < 3000; i += 100) 42 | { 43 | try 44 | { 45 | control_cluster->splitRegion(bank.bank_key(i)); 46 | } 47 | catch (Exception & e) 48 | { 49 | std::cerr << e.displayText() << std::endl; 50 | } 51 | } 52 | 53 | auto close_thread = std::thread([&]() { 54 | try 55 | { 56 | std::this_thread::sleep_for(std::chrono::seconds(100)); 57 | bank.close(); 58 | } 59 | catch (Exception & e) 60 | { 61 | std::cerr << e.displayText() << std::endl; 62 | } 63 | }); 64 | 65 | bank.execute(); 66 | close_thread.join(); 67 | } 68 | catch (Exception & e) 69 | { 70 | std::cerr << e.displayText() << std::endl; 71 | } 72 | catch (std::exception & e) 73 | { 74 | std::cerr << e.what() << std::endl; 75 | } 76 | 77 | } // namespace 78 | -------------------------------------------------------------------------------- /cmake/Modules/FindPackageMessage.cmake: -------------------------------------------------------------------------------- 1 | # Distributed under the OSI-approved BSD 3-Clause License. See accompanying 2 | # file Copyright.txt or https://cmake.org/licensing for details. 3 | 4 | #[=======================================================================[.rst: 5 | FindPackageMessage 6 | ------------------ 7 | 8 | .. code-block:: cmake 9 | 10 | find_package_message( "message for user" "find result details") 11 | 12 | This function is intended to be used in FindXXX.cmake modules files. 13 | It will print a message once for each unique find result. This is 14 | useful for telling the user where a package was found. The first 15 | argument specifies the name (XXX) of the package. The second argument 16 | specifies the message to display. The third argument lists details 17 | about the find result so that if they change the message will be 18 | displayed again. The macro also obeys the QUIET argument to the 19 | find_package command. 20 | 21 | Example: 22 | 23 | .. code-block:: cmake 24 | 25 | if(X11_FOUND) 26 | find_package_message(X11 "Found X11: ${X11_X11_LIB}" 27 | "[${X11_X11_LIB}][${X11_INCLUDE_DIR}]") 28 | else() 29 | ... 30 | endif() 31 | #]=======================================================================] 32 | 33 | function(find_package_message pkg msg details) 34 | # Avoid printing a message repeatedly for the same find result. 35 | if(NOT ${pkg}_FIND_QUIETLY) 36 | string(REPLACE "\n" "" details "${details}") 37 | set(DETAILS_VAR FIND_PACKAGE_MESSAGE_DETAILS_${pkg}) 38 | if(NOT "${details}" STREQUAL "${${DETAILS_VAR}}") 39 | # The message has not yet been printed. 40 | message(STATUS "${msg}") 41 | 42 | # Save the find details in the cache to avoid printing the same 43 | # message again. 44 | set("${DETAILS_VAR}" "${details}" 45 | CACHE INTERNAL "Details about finding ${pkg}") 46 | endif() 47 | endif() 48 | endfunction() 49 | -------------------------------------------------------------------------------- /include/pingcap/kv/Txn.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | #include 9 | #include 10 | 11 | namespace pingcap 12 | { 13 | namespace kv 14 | { 15 | using Buffer = std::map; 16 | 17 | // Txn supports transaction operation for TiKV. 18 | // Note that this implementation is only used for TEST right now. 19 | struct Txn 20 | { 21 | Cluster * cluster; 22 | 23 | Buffer buffer; 24 | 25 | uint64_t start_ts; 26 | 27 | std::chrono::milliseconds start_time; 28 | 29 | explicit Txn(Cluster * cluster_) 30 | : cluster(cluster_) 31 | , start_ts(cluster_->pd_client->getTS()) 32 | , start_time(std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch())) 33 | { 34 | } 35 | 36 | void commit() 37 | { 38 | // Only enable async commit for test 39 | bool use_async_commit = false; 40 | fiu_do_on("use_async_commit", { use_async_commit = true; }); 41 | auto committer = std::make_shared(this, use_async_commit); 42 | committer->execute(); 43 | } 44 | 45 | void set(const std::string & key, const std::string & value) { buffer.emplace(key, value); } 46 | 47 | std::pair get(const std::string & key) 48 | { 49 | auto it = buffer.find(key); 50 | if (it != buffer.end()) 51 | { 52 | return std::make_pair(it->second, true); 53 | } 54 | Snapshot snapshot(cluster, start_ts); 55 | std::string value = snapshot.Get(key); 56 | if (value.empty()) 57 | return std::make_pair("", false); 58 | return std::make_pair(value, true); 59 | } 60 | 61 | void walkBuffer(std::function foo) 62 | { 63 | for (auto & it : buffer) 64 | { 65 | foo(it.first, it.second); 66 | } 67 | } 68 | }; 69 | 70 | } // namespace kv 71 | } // namespace pingcap 72 | -------------------------------------------------------------------------------- /include/pingcap/Config.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #pragma GCC diagnostic push 4 | #ifdef __clang__ 5 | #pragma GCC diagnostic ignored "-Wdeprecated-declarations" 6 | #endif 7 | #include 8 | #pragma GCC diagnostic pop 9 | 10 | #include 11 | 12 | #include 13 | #include 14 | #include 15 | 16 | namespace pingcap 17 | { 18 | struct ClusterConfig 19 | { 20 | std::string tiflash_engine_key; 21 | std::string tiflash_engine_value; 22 | std::string ca_path; 23 | std::string cert_path; 24 | std::string key_path; 25 | ::kvrpcpb::APIVersion api_version = ::kvrpcpb::APIVersion::V1; 26 | 27 | ClusterConfig() = default; 28 | 29 | ClusterConfig(const std::string & engine_key_, 30 | const std::string & engine_value_, 31 | const std::string & ca_path_, 32 | const std::string & cert_path_, 33 | const std::string & key_path_, 34 | const ::kvrpcpb::APIVersion & api_version_) 35 | : tiflash_engine_key(engine_key_) 36 | , tiflash_engine_value(engine_value_) 37 | , ca_path(ca_path_) 38 | , cert_path(cert_path_) 39 | , key_path(key_path_) 40 | , api_version(api_version_) 41 | {} 42 | 43 | bool hasTlsConfig() const { return !ca_path.empty(); } 44 | 45 | 46 | grpc::SslCredentialsOptions getGrpcCredentials() const 47 | { 48 | if (hasTlsConfig()) 49 | { 50 | grpc::SslCredentialsOptions options; 51 | options.pem_root_certs = readFile(ca_path); 52 | options.pem_cert_chain = readFile(cert_path); 53 | options.pem_private_key = readFile(key_path); 54 | return options; 55 | } 56 | return {}; 57 | } 58 | 59 | private: 60 | static std::string readFile(const std::string & path) 61 | { 62 | std::ifstream t(path.data()); 63 | std::string str((std::istreambuf_iterator(t)), std::istreambuf_iterator()); 64 | return str; 65 | } 66 | }; 67 | 68 | } // namespace pingcap 69 | -------------------------------------------------------------------------------- /include/pingcap/Exception.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | namespace pingcap 9 | { 10 | enum ErrorCodes : int 11 | { 12 | MismatchClusterIDCode = 1, 13 | GRPCErrorCode = 2, 14 | InitClusterIDFailed = 3, 15 | UpdatePDLeaderFailed = 4, 16 | TimeoutError = 5, 17 | RegionUnavailable = 6, 18 | LogicalError = 7, 19 | LockError = 8, 20 | LeanerUnavailable = 9, 21 | StoreNotReady = 10, 22 | RaftEntryTooLarge = 11, 23 | ServerIsBusy = 12, 24 | NotLeader = 13, 25 | RegionEpochNotMatch = 14, 26 | CoprocessorError = 15, 27 | TxnNotFound = 16, 28 | NonAsyncCommit = 17, 29 | KeyspaceNotEnabled = 18, 30 | InternalError = 19, 31 | GRPCNotImplemented = 20, 32 | UnknownError = 21 33 | }; 34 | 35 | class Exception : public Poco::Exception 36 | { 37 | public: 38 | Exception() = default; /// For deferred initialization. 39 | explicit Exception(const std::string & msg, int code = 0) 40 | : Poco::Exception(msg, code) 41 | {} 42 | Exception(const std::string & msg, const std::string & arg, int code = 0) 43 | : Poco::Exception(msg, arg, code) 44 | {} 45 | Exception(const std::string & msg, const Exception & exc, int code = 0) 46 | : Poco::Exception(msg, exc, code) 47 | {} 48 | explicit Exception(const Poco::Exception & exc) 49 | : Poco::Exception(exc.displayText()) 50 | {} 51 | 52 | Exception * clone() const override { return new Exception(*this); } 53 | void rethrow() const override { throw *this; } 54 | 55 | bool empty() const { return code() == 0 && message().empty(); } 56 | }; 57 | 58 | inline std::string getCurrentExceptionMsg(const std::string & prefix_msg) 59 | { 60 | std::string msg = prefix_msg; 61 | try 62 | { 63 | throw; 64 | } 65 | catch (const Exception & e) 66 | { 67 | msg += e.message(); 68 | } 69 | catch (const std::exception & e) 70 | { 71 | msg += std::string(e.what()); 72 | } 73 | catch (...) 74 | { 75 | msg += "unknown exception"; 76 | } 77 | return msg; 78 | } 79 | 80 | } // namespace pingcap 81 | -------------------------------------------------------------------------------- /include/pingcap/pd/Oracle.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | #include 9 | 10 | namespace pingcap 11 | { 12 | namespace pd 13 | { 14 | constexpr int physicalShiftBits = 18; 15 | 16 | inline int64_t extractPhysical(uint64_t ts) 17 | { 18 | return ts >> physicalShiftBits; 19 | } 20 | 21 | // Oracle provides strictly ascending timestamps. 22 | class Oracle 23 | { 24 | ClientPtr pd_client; 25 | 26 | std::atomic_bool quit; 27 | 28 | std::atomic last_ts; 29 | 30 | std::thread work_thread; 31 | std::chrono::milliseconds update_interval; 32 | 33 | Logger * log; 34 | 35 | public: 36 | Oracle(ClientPtr pd_client_, std::chrono::milliseconds update_interval_) 37 | : pd_client(pd_client_) 38 | , update_interval(update_interval_) 39 | , log(&Logger::get("pd/oracle")) 40 | { 41 | quit = false; 42 | work_thread = std::thread([&]() { updateTS(update_interval); }); 43 | last_ts = 0; 44 | } 45 | 46 | ~Oracle() { close(); } 47 | 48 | void close() 49 | { 50 | quit = true; 51 | work_thread.join(); 52 | } 53 | 54 | int64_t untilExpired(uint64_t lock_ts, uint64_t ttl) { return extractPhysical(lock_ts) + ttl - extractPhysical(last_ts); } 55 | 56 | uint64_t getLowResolutionTimestamp() { return last_ts; } 57 | 58 | bool isExpired(uint64_t lock_ts, uint64_t ttl) { return untilExpired(lock_ts, ttl) <= 0; } 59 | 60 | private: 61 | void updateTS(std::chrono::milliseconds update_interval) 62 | { 63 | pingcap::SetThreadName("PDUpdateTS"); 64 | 65 | for (;;) 66 | { 67 | if (quit) 68 | { 69 | return; 70 | } 71 | try 72 | { 73 | last_ts = pd_client->getTS(); 74 | } 75 | catch (Exception & e) 76 | { 77 | log->warning("update ts error: " + e.displayText()); 78 | } 79 | std::this_thread::sleep_for(update_interval); 80 | } 81 | } 82 | }; 83 | 84 | using OraclePtr = std::unique_ptr; 85 | 86 | } // namespace pd 87 | } // namespace pingcap 88 | -------------------------------------------------------------------------------- /src/kv/Backoff.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | 4 | namespace pingcap 5 | { 6 | namespace kv 7 | { 8 | BackoffPtr newBackoff(BackoffType tp) 9 | { 10 | switch (tp) 11 | { 12 | case boTiKVRPC: 13 | return std::make_shared(100, 2000, EqualJitter); 14 | case boTxnLock: 15 | return std::make_shared(200, 3000, EqualJitter); 16 | case boTxnLockFast: 17 | return std::make_shared(100, 3000, EqualJitter); 18 | case boPDRPC: 19 | return std::make_shared(500, 3000, EqualJitter); 20 | case boRegionMiss: 21 | return std::make_shared(2, 500, NoJitter); 22 | case boRegionScheduling: 23 | return std::make_shared(2, 500, NoJitter); 24 | case boServerBusy: 25 | return std::make_shared(2000, 10000, EqualJitter); 26 | case boTiKVDiskFull: 27 | return std::make_shared(500, 5000, NoJitter); 28 | case boTxnNotFound: 29 | return std::make_shared(2, 500, NoJitter); 30 | case boMaxTsNotSynced: 31 | return std::make_shared(2, 500, NoJitter); 32 | case boMaxDataNotReady: 33 | return std::make_shared(100, 2000, NoJitter); 34 | case boMaxRegionNotInitialized: 35 | return std::make_shared(2, 1000, NoJitter); 36 | case boTiFlashRPC: 37 | return std::make_shared(100, 10000, EqualJitter); 38 | } 39 | return nullptr; 40 | } 41 | 42 | void Backoffer::backoff(pingcap::kv::BackoffType tp, const pingcap::Exception & exc) 43 | { 44 | backoffWithMaxSleep(tp, -1, exc); 45 | } 46 | 47 | void Backoffer::backoffWithMaxSleep(pingcap::kv::BackoffType tp, int max_sleep_time, const pingcap::Exception & exc) 48 | { 49 | if (exc.code() == MismatchClusterIDCode) 50 | { 51 | exc.rethrow(); 52 | } 53 | 54 | BackoffPtr bo; 55 | auto it = backoff_map.find(tp); 56 | if (it != backoff_map.end()) 57 | { 58 | bo = it->second; 59 | } 60 | else 61 | { 62 | bo = newBackoff(tp); 63 | backoff_map[tp] = bo; 64 | } 65 | total_sleep += bo->sleep(max_sleep_time); 66 | if (max_sleep > 0 && total_sleep > max_sleep) 67 | { 68 | // TODO:: Should Record all the errors!! 69 | throw exc; 70 | } 71 | } 72 | 73 | } // namespace kv 74 | } // namespace pingcap 75 | -------------------------------------------------------------------------------- /include/pingcap/pd/IClient.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | #include 7 | #include 8 | 9 | #pragma GCC diagnostic push 10 | #pragma GCC diagnostic ignored "-Wunused-parameter" 11 | #include 12 | #include 13 | #include 14 | #pragma GCC diagnostic pop 15 | 16 | namespace pingcap::pd 17 | { 18 | 19 | class IClient 20 | { 21 | public: 22 | virtual ~IClient() = default; 23 | 24 | virtual uint64_t getTS() = 0; 25 | 26 | virtual pdpb::GetRegionResponse getRegionByKey(const std::string & key) = 0; 27 | 28 | virtual pdpb::GetRegionResponse getRegionByID(uint64_t region_id) = 0; 29 | 30 | virtual metapb::Store getStore(uint64_t store_id) = 0; 31 | 32 | virtual bool isClusterBootstrapped() = 0; 33 | 34 | virtual std::vector getAllStores(bool exclude_tombstone) = 0; 35 | 36 | [[deprecated("Use getGCState instead")]] virtual uint64_t getGCSafePoint() = 0; 37 | 38 | // Return the gc safe point of given keyspace_id. 39 | [[deprecated("Use getGCState instead")]] virtual uint64_t getGCSafePointV2(KeyspaceID keyspace_id) = 0; 40 | 41 | virtual pdpb::GetGCStateResponse getGCState(KeyspaceID keyspace_id) = 0; 42 | 43 | virtual pdpb::GetAllKeyspacesGCStatesResponse getAllKeyspacesGCStates() = 0; 44 | 45 | virtual KeyspaceID getKeyspaceID(const std::string & keyspace_name) = 0; 46 | 47 | virtual void update(const std::vector & addrs, const ClusterConfig & config_) = 0; 48 | 49 | virtual bool isMock() = 0; 50 | 51 | virtual std::string getLeaderUrl() = 0; 52 | 53 | // ResourceControl related. 54 | virtual resource_manager::ListResourceGroupsResponse listResourceGroups(const resource_manager::ListResourceGroupsRequest &) = 0; 55 | 56 | virtual resource_manager::GetResourceGroupResponse getResourceGroup(const resource_manager::GetResourceGroupRequest &) = 0; 57 | 58 | virtual resource_manager::PutResourceGroupResponse addResourceGroup(const resource_manager::PutResourceGroupRequest &) = 0; 59 | 60 | virtual resource_manager::PutResourceGroupResponse modifyResourceGroup(const resource_manager::PutResourceGroupRequest &) = 0; 61 | 62 | virtual resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const resource_manager::DeleteResourceGroupRequest &) = 0; 63 | 64 | virtual resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest & req) = 0; 65 | }; 66 | 67 | using ClientPtr = std::shared_ptr; 68 | 69 | } // namespace pingcap::pd 70 | -------------------------------------------------------------------------------- /.clang-format: -------------------------------------------------------------------------------- 1 | --- 2 | BasedOnStyle: Google 3 | Language: Cpp 4 | AlignAfterOpenBracket: Align 5 | AlignEscapedNewlines: Left 6 | AlignTrailingComments: false 7 | AllowAllArgumentsOnNextLine: false 8 | AllowAllParametersOfDeclarationOnNextLine: false 9 | AllowShortBlocksOnASingleLine: Always 10 | AllowShortEnumsOnASingleLine: false 11 | AllowShortFunctionsOnASingleLine: Inline 12 | AllowShortLambdasOnASingleLine: Inline 13 | AlwaysBreakAfterReturnType: None 14 | AlwaysBreakTemplateDeclarations: true 15 | BinPackArguments: false 16 | BinPackParameters: false 17 | BreakBeforeBraces: Custom 18 | BraceWrapping: 19 | AfterCaseLabel: true 20 | AfterClass: true 21 | AfterControlStatement: true 22 | AfterEnum : true 23 | AfterFunction : true 24 | AfterNamespace : true 25 | AfterStruct : true 26 | AfterUnion : true 27 | BeforeCatch : true 28 | BeforeElse : true 29 | IndentBraces : false 30 | SplitEmptyFunction: false 31 | BreakConstructorInitializers: BeforeComma 32 | BreakInheritanceList : BeforeComma 33 | ColumnLimit: 0 34 | ConstructorInitializerAllOnOneLineOrOnePerLine: false 35 | Cpp11BracedListStyle: true 36 | EmptyLineBeforeAccessModifier: LogicalBlock 37 | ExperimentalAutoDetectBinPacking: true 38 | FixNamespaceComments: true 39 | IndentCaseLabels: false 40 | KeepEmptyLinesAtTheStartOfBlocks: false 41 | MaxEmptyLinesToKeep: 2 42 | #PPIndentWidth: 2 # clang-format version 12.0.0 doesn't support yet 43 | PointerAlignment: Middle 44 | ReflowComments: false 45 | SortIncludes: true 46 | SpaceAfterTemplateKeyword: true 47 | Standard: Cpp11 48 | TabWidth: 4 49 | UseTab: Never 50 | 51 | # Not changed: 52 | AccessModifierOffset: -4 53 | AlignConsecutiveAssignments: false 54 | AlignConsecutiveDeclarations: false 55 | AlignOperands: false 56 | AllowShortCaseLabelsOnASingleLine: false 57 | AllowShortIfStatementsOnASingleLine: false 58 | AllowShortLoopsOnASingleLine: false 59 | AlwaysBreakBeforeMultilineStrings: false 60 | BreakBeforeBinaryOperators: All 61 | BreakBeforeTernaryOperators: true 62 | CommentPragmas: '^ IWYU pragma:' 63 | ConstructorInitializerIndentWidth: 4 64 | ContinuationIndentWidth: 4 65 | DerivePointerAlignment: false 66 | DisableFormat: false 67 | ForEachMacros: [ foreach, Q_FOREACH, BOOST_FOREACH ] 68 | IndentWidth: 4 69 | IndentWrappedFunctionNames: false 70 | MacroBlockBegin: '' 71 | MacroBlockEnd: '' 72 | NamespaceIndentation: None 73 | ObjCBlockIndentWidth: 4 74 | ObjCSpaceAfterProperty: true 75 | ObjCSpaceBeforeProtocolList: true 76 | PenaltyBreakBeforeFirstCallParameter: 19 77 | PenaltyBreakComment: 300 78 | PenaltyBreakFirstLessLess: 120 79 | PenaltyBreakString: 1000 80 | PenaltyExcessCharacter: 1000000 81 | PenaltyReturnTypeOnItsOwnLine: 60 82 | SpaceAfterCStyleCast: false 83 | SpaceBeforeAssignmentOperators: true 84 | SpaceBeforeParens: ControlStatements 85 | SpaceInEmptyParentheses: false 86 | SpacesBeforeTrailingComments: 1 87 | SpacesInContainerLiterals: true 88 | SpacesInCStyleCastParentheses: false 89 | SpacesInParentheses: false 90 | SpacesInSquareBrackets: false 91 | ... 92 | -------------------------------------------------------------------------------- /src/RedactHelpers.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include 6 | #include 7 | 8 | namespace pingcap 9 | { 10 | 11 | std::atomic Redact::REDACT_LOG = RedactMode::Disable; 12 | 13 | void Redact::setRedactLog(RedactMode v) 14 | { 15 | Redact::REDACT_LOG.store(v, std::memory_order_relaxed); 16 | } 17 | 18 | constexpr auto hex_byte_to_char_uppercase_table = "000102030405060708090A0B0C0D0E0F" 19 | "101112131415161718191A1B1C1D1E1F" 20 | "202122232425262728292A2B2C2D2E2F" 21 | "303132333435363738393A3B3C3D3E3F" 22 | "404142434445464748494A4B4C4D4E4F" 23 | "505152535455565758595A5B5C5D5E5F" 24 | "606162636465666768696A6B6C6D6E6F" 25 | "707172737475767778797A7B7C7D7E7F" 26 | "808182838485868788898A8B8C8D8E8F" 27 | "909192939495969798999A9B9C9D9E9F" 28 | "A0A1A2A3A4A5A6A7A8A9AAABACADAEAF" 29 | "B0B1B2B3B4B5B6B7B8B9BABBBCBDBEBF" 30 | "C0C1C2C3C4C5C6C7C8C9CACBCCCDCECF" 31 | "D0D1D2D3D4D5D6D7D8D9DADBDCDDDEDF" 32 | "E0E1E2E3E4E5E6E7E8E9EAEBECEDEEEF" 33 | "F0F1F2F3F4F5F6F7F8F9FAFBFCFDFEFF"; 34 | 35 | inline void writeHexByteUppercase(uint8_t byte, void * out) 36 | { 37 | memcpy(out, &hex_byte_to_char_uppercase_table[static_cast(byte) * 2], 2); 38 | } 39 | 40 | std::string Redact::keyToHexString(const char * key, size_t size) 41 | { 42 | // Encode as upper hex string 43 | std::string buf(size * 2, '\0'); 44 | char * pos = buf.data(); 45 | for (size_t i = 0; i < size; ++i) 46 | { 47 | writeHexByteUppercase(static_cast(key[i]), pos); 48 | pos += 2; 49 | } 50 | return buf; 51 | } 52 | 53 | std::string Redact::keyToDebugString(const char * key, const size_t size) 54 | { 55 | const auto v = Redact::REDACT_LOG.load(std::memory_order_relaxed); 56 | switch (v) 57 | { 58 | case RedactMode::Enable: 59 | return "?"; 60 | case RedactMode::Disable: 61 | // Encode as string 62 | return Redact::keyToHexString(key, size); 63 | case RedactMode::Marker: 64 | { 65 | // Note: the `s` must be hexadecimal string so we don't need to care 66 | // about escaping here. 67 | auto s = Redact::keyToHexString(key, size); 68 | return std::string("‹") + s + "›"; 69 | } 70 | default: 71 | throw Exception(std::string("Should not reach here, v=") + std::to_string(static_cast(v))); 72 | } 73 | } 74 | 75 | } // namespace pingcap 76 | -------------------------------------------------------------------------------- /src/test/region_split_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include 7 | 8 | #include "mock_tikv.h" 9 | #include "test_helper.h" 10 | 11 | namespace pingcap::tests 12 | { 13 | using namespace pingcap; 14 | using namespace pingcap::kv; 15 | 16 | class TestWithMockKVRegionSplit : public testing::Test 17 | { 18 | protected: 19 | void SetUp() override 20 | { 21 | mock_kv_cluster = mockkv::initCluster(); 22 | std::vector pd_addrs = mock_kv_cluster->pd_addrs; 23 | 24 | test_cluster = createCluster(pd_addrs); 25 | control_cluster = createCluster(pd_addrs); 26 | } 27 | 28 | mockkv::ClusterPtr mock_kv_cluster; 29 | 30 | ClusterPtr test_cluster; 31 | ClusterPtr control_cluster; 32 | }; 33 | 34 | TEST_F(TestWithMockKVRegionSplit, testSplitRegionGet) 35 | { 36 | { 37 | Txn txn(test_cluster.get()); 38 | 39 | txn.set("abc", "1"); 40 | txn.set("abd", "2"); 41 | txn.set("abe", "3"); 42 | txn.set("abf", "4"); 43 | txn.set("abg", "5"); 44 | txn.set("abz", "6"); 45 | txn.commit(); 46 | Snapshot snap(test_cluster.get(), test_cluster->pd_client->getTS()); 47 | 48 | std::string result = snap.Get("abf"); 49 | 50 | ASSERT_EQ(result, "4"); 51 | 52 | control_cluster->splitRegion("abf"); 53 | 54 | result = snap.Get("abc"); 55 | 56 | ASSERT_EQ(result, "1"); 57 | 58 | result = snap.Get("abf"); 59 | 60 | ASSERT_EQ(result, "4"); 61 | } 62 | 63 | 64 | { 65 | Txn txn(test_cluster.get()); 66 | 67 | txn.set("abf", "6"); 68 | txn.set("abg", "5"); 69 | txn.set("abz", "4"); 70 | txn.commit(); 71 | 72 | Snapshot snap(test_cluster.get(), test_cluster->pd_client->getTS()); 73 | std::string result = snap.Get("abf"); 74 | 75 | ASSERT_EQ(result, "6"); 76 | } 77 | } 78 | 79 | TEST_F(TestWithMockKVRegionSplit, testSplitRegionScan) 80 | { 81 | Txn txn(test_cluster.get()); 82 | 83 | txn.set("abc", "1"); 84 | txn.set("abd", "2"); 85 | txn.set("abe", "3"); 86 | txn.set("abf", "4"); 87 | txn.set("abg", "5"); 88 | txn.set("abh", "6"); 89 | txn.set("zzz", "7"); 90 | txn.commit(); 91 | 92 | Snapshot snap(test_cluster.get(), test_cluster->pd_client->getTS()); 93 | 94 | auto scanner = snap.Scan("", ""); 95 | 96 | int answer = 0; 97 | while (scanner.valid) 98 | { 99 | ASSERT_EQ(scanner.value(), std::to_string(++answer)); 100 | scanner.next(); 101 | } 102 | 103 | ASSERT_EQ(answer, 7); 104 | 105 | answer = 0; 106 | 107 | control_cluster->splitRegion("abe"); 108 | 109 | auto scanner1 = snap.Scan("ab", "ac"); 110 | 111 | while (scanner1.valid) 112 | { 113 | ASSERT_EQ(scanner1.value(), std::to_string(++answer)); 114 | scanner1.next(); 115 | } 116 | ASSERT_EQ(answer, 6); 117 | } 118 | 119 | } // namespace pingcap::tests 120 | -------------------------------------------------------------------------------- /src/test/mock_tikv.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | #include 12 | 13 | namespace pingcap 14 | { 15 | namespace kv 16 | { 17 | namespace mockkv 18 | { 19 | using namespace Poco::Net; 20 | 21 | constexpr char mock_server[] = "http://127.0.0.1:2378"; 22 | 23 | struct Store 24 | { 25 | int id; 26 | void deser(Poco::JSON::Object::Ptr json_obj) { id = json_obj->getValue("id"); } 27 | }; 28 | 29 | struct Cluster 30 | { 31 | int id; 32 | std::vector pd_addrs; 33 | std::vector stores; 34 | 35 | // See https://github.com/pingcap/failpoint 36 | // Mock tikv use go failpoint to inject faults. 37 | void updateFailPoint(int store_id, std::string fail_point, std::string term) const 38 | { 39 | HTTPClientSession sess("127.0.0.1", 2378); 40 | HTTPRequest req(HTTPRequest::HTTP_POST, 41 | std::string(mock_server) + "/mock-tikv/api/v1/clusters/" + std::to_string(id) + "/stores/" + std::to_string(store_id) 42 | + "/failpoints/" + fail_point); 43 | req.setContentLength(term.size()); 44 | auto & ostream = sess.sendRequest(req); 45 | ostream << term; 46 | } 47 | }; 48 | 49 | using ClusterPtr = std::shared_ptr; 50 | 51 | inline ClusterPtr initCluster() 52 | { 53 | HTTPClientSession sess("127.0.0.1", 2378); 54 | HTTPRequest req(HTTPRequest::HTTP_POST, std::string(mock_server) + "/mock-tikv/api/v1/clusters"); 55 | req.setContentType("application/json"); 56 | req.setContentLength(2); 57 | auto & ostream = sess.sendRequest(req); 58 | ostream << "{}"; 59 | HTTPResponse res; 60 | auto & is = sess.receiveResponse(res); 61 | char buffer[1024]; 62 | std::string json_str; 63 | for (;;) 64 | { 65 | is.read(buffer, 1024); 66 | if (is) 67 | { 68 | json_str.append(buffer, 1024); 69 | } 70 | else 71 | { 72 | json_str.append(buffer, is.gcount()); 73 | break; 74 | } 75 | } 76 | Poco::JSON::Parser parser; 77 | Poco::Dynamic::Var result = parser.parse(json_str); 78 | auto json_obj = result.extract(); 79 | auto pd_arr = json_obj->getArray("members"); 80 | std::vector urls; 81 | for (size_t i = 0; i < pd_arr->size(); i++) 82 | { 83 | urls.push_back(pd_arr->getObject(i)->getArray("client_urls")->getElement(0)); 84 | } 85 | ClusterPtr cluster = std::make_shared(); 86 | cluster->pd_addrs = std::move(urls); 87 | auto json_stores = json_obj->getArray("stores"); 88 | for (int i = 0; i < json_stores->size(); i++) 89 | { 90 | Store store; 91 | store.deser(json_stores->getObject(i)); 92 | cluster->stores.push_back(store); 93 | } 94 | cluster->id = json_obj->getValue("id"); 95 | return cluster; 96 | } 97 | 98 | } // namespace mockkv 99 | } // namespace kv 100 | } // namespace pingcap 101 | -------------------------------------------------------------------------------- /include/pingcap/pd/CodecClient.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | namespace pingcap 11 | { 12 | namespace pd 13 | { 14 | struct CodecClient : public Client 15 | { 16 | CodecClient(const std::vector & addrs, const ClusterConfig & config) 17 | : Client(addrs, config) 18 | {} 19 | 20 | pdpb::GetRegionResponse getRegionByKey(const std::string & key) override 21 | { 22 | auto resp = Client::getRegionByKey(encodeBytes(key)); 23 | processRegionResult(*resp.mutable_region()); 24 | return resp; 25 | } 26 | 27 | pdpb::GetRegionResponse getRegionByID(uint64_t region_id) override 28 | { 29 | auto resp = Client::getRegionByID(region_id); 30 | processRegionResult(*resp.mutable_region()); 31 | return resp; 32 | } 33 | 34 | static metapb::Region processRegionResult(metapb::Region & region) 35 | { 36 | region.set_start_key(decodeBytes(region.start_key())); 37 | region.set_end_key(decodeBytes(region.end_key())); 38 | return region; 39 | } 40 | 41 | private: 42 | static constexpr uint8_t ENC_MARKER = 0xff; 43 | static constexpr uint8_t ENC_GROUP_SIZE = 8; 44 | static constexpr char ENC_ASC_PADDING[ENC_GROUP_SIZE] = {0}; 45 | 46 | static std::string encodeBytes(const std::string & raw) 47 | { 48 | if (raw.empty()) 49 | return ""; 50 | std::stringstream ss; 51 | size_t len = raw.size(); 52 | size_t index = 0; 53 | while (index <= len) 54 | { 55 | size_t remain = len - index; 56 | size_t pad = 0; 57 | if (remain >= ENC_GROUP_SIZE) 58 | { 59 | ss.write(raw.data() + index, ENC_GROUP_SIZE); 60 | } 61 | else 62 | { 63 | pad = ENC_GROUP_SIZE - remain; 64 | ss.write(raw.data() + index, remain); 65 | ss.write(ENC_ASC_PADDING, pad); 66 | } 67 | ss.put(static_cast(ENC_MARKER - static_cast(pad))); 68 | index += ENC_GROUP_SIZE; 69 | } 70 | return ss.str(); 71 | } 72 | 73 | static std::string decodeBytes(const std::string & raw) 74 | { 75 | if (raw.empty()) 76 | return ""; 77 | std::stringstream ss; 78 | int cursor = 0; 79 | while (true) 80 | { 81 | size_t next_cursor = cursor + 9; 82 | if (next_cursor > raw.size()) 83 | throw Exception("Wrong format, cursor over buffer size. (DecodeBytes)", ErrorCodes::LogicalError); 84 | auto marker = static_cast(raw[cursor + 8]); 85 | uint8_t pad_size = ENC_MARKER - marker; 86 | 87 | if (pad_size > 8) 88 | throw Exception("Wrong format, too many padding bytes. (DecodeBytes)", ErrorCodes::LogicalError); 89 | ss.write(&raw[cursor], 8 - pad_size); 90 | cursor = next_cursor; 91 | if (pad_size != 0) 92 | break; 93 | } 94 | return ss.str(); 95 | } 96 | }; 97 | 98 | } // namespace pd 99 | } // namespace pingcap 100 | -------------------------------------------------------------------------------- /src/test/bank_test/schrodinger_client.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | 15 | #include 16 | 17 | namespace 18 | { 19 | 20 | using namespace Poco::Net; 21 | 22 | class Client 23 | { 24 | public: 25 | Client() 26 | : manager_addr(std::getenv("MANAGER_ADDR")), box_id(std::stoi(std::getenv("BOX_ID"))), self_id(std::stoi(std::getenv("TEST_ID"))) 27 | { 28 | std::cerr << "manager addr: " << manager_addr << std::endl; 29 | std::cerr << "box id: " << box_id << std::endl; 30 | std::cerr << "self id: " << self_id << std::endl; 31 | getConfig(); 32 | } 33 | 34 | std::string debug_string(std::string str) 35 | { 36 | for (size_t i = 0; i < str.size(); i++) 37 | { 38 | if (str[i] == '\n') 39 | str[i] = ' '; 40 | } 41 | return str; 42 | } 43 | 44 | void getConfig() 45 | { 46 | std::string manager_addr_str(manager_addr); 47 | int begin = manager_addr_str.find("://") + 3; 48 | std::string rest_info = manager_addr_str.substr(begin); 49 | std::string host = rest_info.substr(0, rest_info.find(":")); 50 | int port = std::stoi(rest_info.substr(rest_info.find(":") + 1)); 51 | std::string url = manager_addr + std::string("/testConfig/") + std::to_string(self_id); 52 | SocketAddress addr(host, port); 53 | HTTPClientSession sess(addr); 54 | HTTPRequest req(HTTPRequest::HTTP_GET, url); 55 | sess.sendRequest(req); 56 | HTTPResponse res; 57 | auto & is = sess.receiveResponse(res); 58 | char buffer[1200]; 59 | std::string json_str; 60 | for (;;) 61 | { 62 | is.read(buffer, 256); 63 | if (is) 64 | { 65 | json_str.append(buffer, 256); 66 | } 67 | else 68 | { 69 | json_str.append(buffer, is.gcount()); 70 | break; 71 | } 72 | } 73 | config = json_str; 74 | } 75 | 76 | std::vector PDs() 77 | { 78 | Poco::JSON::Parser parser; 79 | Poco::Dynamic::Var result = parser.parse(config); 80 | auto json_obj = result.extract(); 81 | auto data = json_obj->getObject("data"); 82 | auto cat = data->getObject("cat"); 83 | auto pds = cat->getArray("pds"); 84 | std::vector rets; 85 | for (size_t i = 0; i < pds->size(); i++) 86 | { 87 | auto ip = pds->getObject(i)->getValue("ip"); 88 | auto port = pds->getObject(i)->getValue("service_port"); 89 | std::string addr = ip + ":" + std::to_string(port); 90 | std::cerr << "pd addr: " << addr << std::endl; 91 | rets.push_back(addr); 92 | } 93 | return rets; 94 | } 95 | 96 | const char * manager_addr; 97 | int box_id; 98 | int self_id; 99 | std::string config; 100 | }; 101 | } // namespace 102 | -------------------------------------------------------------------------------- /include/pingcap/kv/Cluster.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | namespace pingcap 14 | { 15 | namespace kv 16 | { 17 | constexpr int oracle_update_interval = 2000; 18 | // Cluster represents a tikv+pd cluster. 19 | 20 | struct Cluster 21 | { 22 | pd::ClientPtr pd_client; 23 | RegionCachePtr region_cache; 24 | RpcClientPtr rpc_client; 25 | 26 | pd::OraclePtr oracle; 27 | 28 | LockResolverPtr lock_resolver; 29 | 30 | ::kvrpcpb::APIVersion api_version = ::kvrpcpb::APIVersion::V1; 31 | 32 | std::unique_ptr thread_pool; 33 | std::unique_ptr mpp_prober; 34 | 35 | Cluster() 36 | : pd_client(std::make_shared()) 37 | , rpc_client(std::make_unique()) 38 | , thread_pool(std::make_unique(1)) 39 | , mpp_prober(std::make_unique(this)) 40 | { 41 | startBackgroundTasks(); 42 | } 43 | 44 | Cluster(const std::vector & pd_addrs, const ClusterConfig & config) 45 | : pd_client(std::make_shared(pd_addrs, config)) 46 | , region_cache(std::make_unique(pd_client, config)) 47 | , rpc_client(std::make_unique(config)) 48 | , oracle(std::make_unique(pd_client, std::chrono::milliseconds(oracle_update_interval))) 49 | , lock_resolver(std::make_unique(this)) 50 | , api_version(config.api_version) 51 | , thread_pool(std::make_unique(2)) 52 | , mpp_prober(std::make_unique(this)) 53 | { 54 | startBackgroundTasks(); 55 | } 56 | 57 | void update(const std::vector & pd_addrs, const ClusterConfig & config) const 58 | { 59 | pd_client->update(pd_addrs, config); 60 | rpc_client->update(config); 61 | } 62 | 63 | // TODO: When the cluster is closed, we should release all the resources 64 | // (e.g. background threads) that cluster object holds so as to exit elegantly. 65 | ~Cluster() 66 | { 67 | mpp_prober->stop(); 68 | if (region_cache) 69 | region_cache->stop(); 70 | thread_pool->stop(); 71 | } 72 | 73 | // Only used by Test and this is not safe ! 74 | void splitRegion(const std::string & split_key); 75 | 76 | void startBackgroundTasks(); 77 | }; 78 | 79 | struct MinCommitTSPushed 80 | { 81 | std::unordered_set container; 82 | 83 | mutable std::mutex mutex; 84 | 85 | MinCommitTSPushed() = default; 86 | 87 | MinCommitTSPushed(MinCommitTSPushed &) {} 88 | 89 | inline void addTimestamps(std::vector & tss) 90 | { 91 | std::lock_guard guard{mutex}; 92 | container.insert(tss.begin(), tss.end()); 93 | } 94 | 95 | inline std::vector getTimestamps() const 96 | { 97 | std::lock_guard guard{mutex}; 98 | std::vector result; 99 | std::copy(container.begin(), container.end(), std::back_inserter(result)); 100 | return result; 101 | } 102 | }; 103 | 104 | using ClusterPtr = std::unique_ptr; 105 | 106 | } // namespace kv 107 | } // namespace pingcap 108 | -------------------------------------------------------------------------------- /include/pingcap/kv/Rpc.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | 11 | namespace pingcap 12 | { 13 | namespace kv 14 | { 15 | struct ConnArray 16 | { 17 | std::mutex mutex; 18 | std::string address; 19 | 20 | size_t index = 0; 21 | std::vector> vec; 22 | 23 | ConnArray() = default; 24 | 25 | ConnArray(size_t max_size, const std::string & addr, const ClusterConfig & config_); 26 | 27 | std::shared_ptr get(); 28 | }; 29 | 30 | using ConnArrayPtr = std::shared_ptr; 31 | using GRPCMetaData = std::multimap; 32 | 33 | struct RpcClient 34 | { 35 | ClusterConfig config; 36 | 37 | std::mutex mutex; 38 | 39 | std::map conns; 40 | 41 | RpcClient() = default; 42 | 43 | explicit RpcClient(const ClusterConfig & config_) 44 | : config(config_) 45 | {} 46 | 47 | void update(const ClusterConfig & config_) 48 | { 49 | std::unique_lock lk(mutex); 50 | config = config_; 51 | conns.clear(); 52 | } 53 | 54 | ConnArrayPtr getConnArray(const std::string & addr); 55 | 56 | ConnArrayPtr createConnArray(const std::string & addr); 57 | }; 58 | 59 | using RpcClientPtr = std::unique_ptr; 60 | 61 | // RpcCall holds the request and response, and delegates RPC calls. 62 | template 63 | class RpcCall 64 | { 65 | public: 66 | RpcCall(const RpcClientPtr & client_, const std::string & addr_) 67 | : client(client_) 68 | , addr(addr_) 69 | , log(&Logger::get("pingcap.tikv")) 70 | {} 71 | 72 | template 73 | void setRequestCtx(REQ & req, RPCContextPtr rpc_ctx, kvrpcpb::APIVersion api_version) 74 | { 75 | ::kvrpcpb::Context * context = req.mutable_context(); 76 | // Set api_version to this context, it's caller's duty to ensure the api_version. 77 | // Besides, the tikv will check api_version and key mode in server-side. 78 | context->set_api_version(api_version); 79 | context->set_region_id(rpc_ctx->region.id); 80 | context->set_allocated_region_epoch(new metapb::RegionEpoch(rpc_ctx->meta.region_epoch())); 81 | context->set_allocated_peer(new metapb::Peer(rpc_ctx->peer)); 82 | } 83 | 84 | void setClientContext(::grpc::ClientContext & context, int timeout, const GRPCMetaData & meta_data = {}) 85 | { 86 | context.set_deadline(std::chrono::system_clock::now() + std::chrono::seconds(timeout)); 87 | for (const auto & it : meta_data) 88 | context.AddMetadata(it.first, it.second); 89 | } 90 | 91 | template 92 | auto call(Args &&... args) 93 | { 94 | ConnArrayPtr conn_array = client->getConnArray(addr); 95 | auto conn_client = conn_array->get(); 96 | return T::call(conn_client, std::forward(args)...); 97 | } 98 | 99 | std::string errMsg(const ::grpc::Status & status, const std::string & extra_msg) 100 | { 101 | auto msg = std::string(T::errMsg()) + " " + std::to_string(status.error_code()) + ": " + status.error_message(); 102 | if (!extra_msg.empty()) 103 | { 104 | msg += " " + extra_msg; 105 | } 106 | return msg; 107 | } 108 | 109 | private: 110 | const RpcClientPtr & client; 111 | const std::string & addr; 112 | Logger * log; 113 | }; 114 | 115 | } // namespace kv 116 | } // namespace pingcap 117 | -------------------------------------------------------------------------------- /include/pingcap/kv/Backoff.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | namespace pingcap 12 | { 13 | namespace kv 14 | { 15 | enum Jitter 16 | { 17 | NoJitter = 1, 18 | FullJitter, 19 | EqualJitter, 20 | DecorrJitter 21 | }; 22 | 23 | enum BackoffType 24 | { 25 | boTiKVRPC = 0, 26 | boTxnLock, 27 | boTxnLockFast, 28 | boPDRPC, 29 | boRegionMiss, 30 | boRegionScheduling, 31 | boServerBusy, 32 | boTiKVDiskFull, 33 | boTxnNotFound, 34 | boMaxTsNotSynced, 35 | boMaxDataNotReady, 36 | boMaxRegionNotInitialized, 37 | boTiFlashRPC, 38 | }; 39 | 40 | inline int expo(int base, int cap, int n) 41 | { 42 | return std::min(double(cap), double(base) * std::pow(2.0, double(n))); 43 | } 44 | 45 | struct Backoff 46 | { 47 | int base; 48 | int cap; 49 | int jitter; 50 | int last_sleep; 51 | int attempts; 52 | 53 | Backoff(int base_, int cap_, Jitter jitter_) 54 | : base(base_) 55 | , cap(cap_) 56 | , jitter(jitter_) 57 | , attempts(0) 58 | { 59 | if (base < 2) 60 | { 61 | base = 2; 62 | } 63 | last_sleep = base; 64 | } 65 | 66 | int sleep(int max_sleep_time) 67 | { 68 | int sleep_time = 0; 69 | int v = 0; 70 | switch (jitter) 71 | { 72 | case NoJitter: 73 | sleep_time = expo(base, cap, attempts); 74 | break; 75 | case FullJitter: 76 | v = expo(base, cap, attempts); 77 | sleep_time = rand() % v; 78 | break; 79 | case EqualJitter: 80 | v = expo(base, cap, attempts); 81 | sleep_time = v / 2 + rand() % (v / 2); 82 | break; 83 | case DecorrJitter: 84 | sleep_time = int(std::min(double(cap), double(base + rand() % (last_sleep * 3 - base)))); 85 | } 86 | if (max_sleep_time >= 0 && max_sleep_time < sleep_time) 87 | sleep_time = max_sleep_time; 88 | std::this_thread::sleep_for(std::chrono::milliseconds(sleep_time)); 89 | attempts++; 90 | last_sleep = sleep_time; 91 | return last_sleep; 92 | } 93 | }; 94 | 95 | constexpr int GetMaxBackoff = 20000; 96 | constexpr int scanMaxBackoff = 20000; 97 | constexpr int prewriteMaxBackoff = 20000; 98 | constexpr int commitMaxBackoff = 41000; 99 | constexpr int splitRegionBackoff = 20000; 100 | constexpr int cleanupMaxBackoff = 20000; 101 | constexpr int copBuildTaskMaxBackoff = 5000; 102 | constexpr int copNextMaxBackoff = 60000; 103 | constexpr int pessimisticLockMaxBackoff = 20000; 104 | 105 | using BackoffPtr = std::shared_ptr; 106 | 107 | struct Backoffer 108 | { 109 | size_t max_sleep; // ms 110 | size_t total_sleep; // ms 111 | std::map backoff_map; 112 | 113 | explicit Backoffer(size_t max_sleep_, size_t total_sleep_ = 0) 114 | : max_sleep(max_sleep_) 115 | , total_sleep(total_sleep_) 116 | {} 117 | 118 | Backoffer clone() const 119 | { 120 | Backoffer res(max_sleep, total_sleep); 121 | for (auto && [k, v] : backoff_map) 122 | { 123 | res.backoff_map.emplace(k, new Backoff(*v)); 124 | } 125 | return res; 126 | } 127 | 128 | Backoffer(const Backoffer &) = delete; 129 | Backoffer(Backoffer &&) = default; 130 | 131 | void backoff(BackoffType tp, const Exception & exc); 132 | void backoffWithMaxSleep(BackoffType tp, int max_sleep_time, const Exception & exc); 133 | }; 134 | 135 | } // namespace kv 136 | } // namespace pingcap 137 | -------------------------------------------------------------------------------- /src/test/test_helper.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | 11 | namespace pingcap 12 | { 13 | namespace tests 14 | { 15 | using namespace pingcap::kv; 16 | 17 | inline ClusterPtr createCluster(const std::vector & pd_addrs) 18 | { 19 | ClusterConfig config; 20 | config.tiflash_engine_key = "engine"; 21 | config.tiflash_engine_value = "tiflash"; 22 | return std::make_unique(pd_addrs, config); 23 | } 24 | 25 | inline std::string toString(const pingcap::coprocessor::KeyRange & range) 26 | { 27 | return "[" + range.start_key + "," + range.end_key + ")"; 28 | } 29 | 30 | inline std::string toString(const pingcap::coprocessor::KeyRanges & ranges) 31 | { 32 | if (ranges.empty()) 33 | return "[]"; 34 | 35 | std::string res("["); 36 | res += toString(ranges[0]); 37 | for (size_t i = 1; i < ranges.size(); ++i) 38 | res += "," + toString(ranges[i]); 39 | res += "]"; 40 | return res; 41 | } 42 | 43 | /// helper functions for comparing KeyRanges 44 | inline ::testing::AssertionResult keyRangesCompare( 45 | const char * lhs_expr, 46 | const char * rhs_expr, 47 | const coprocessor::KeyRanges & lhs, 48 | const coprocessor::KeyRanges & rhs) 49 | { 50 | if (lhs.size() != rhs.size()) 51 | { 52 | auto l_expr = "(" + std::string(lhs_expr) + ").size()"; 53 | auto r_expr = "(" + std::string(rhs_expr) + ").size()"; 54 | return ::testing::internal::EqFailure( 55 | l_expr.c_str(), 56 | r_expr.c_str(), 57 | "[size=" + std::to_string(lhs.size()) + "] [ranges=" + toString(lhs) + "]", 58 | "[size=" + std::to_string(rhs.size()) + "] [ranges=" + toString(rhs) + "]", 59 | false); 60 | } 61 | 62 | for (size_t i = 0; i < lhs.size(); ++i) 63 | { 64 | const auto & lkr = lhs[i]; 65 | const auto & rkr = rhs[i]; 66 | if (lkr.start_key != rkr.start_key || lkr.end_key != rkr.end_key) 67 | { 68 | auto l_expr = "(" + std::string(lhs_expr) + ")[" + std::to_string(i) + "]"; 69 | auto r_expr = "(" + std::string(rhs_expr) + ")[" + std::to_string(i) + "]"; 70 | return ::testing::internal::EqFailure( 71 | l_expr.c_str(), 72 | r_expr.c_str(), 73 | "[r=" + toString(lkr) + "] [ranges=" + toString(lhs) + "]", 74 | "[r=" + toString(rkr) + "] [ranges=" + toString(rhs) + "]", 75 | false); 76 | } 77 | } 78 | 79 | return ::testing::AssertionSuccess(); 80 | } 81 | #define ASSERT_KEY_RANGES_EQ(val1, val2) ASSERT_PRED_FORMAT2(::pingcap::tests::keyRangesCompare, val1, val2) 82 | #define EXPECT_KEY_RANGES_EQ(val1, val2) EXPECT_PRED_FORMAT2(::pingcap::tests::keyRangesCompare, val1, val2) 83 | 84 | // helper functions for comparing KeyRanges 85 | inline ::testing::AssertionResult locationRangeCompare( 86 | const char * lhs_expr, 87 | const char * rhs_expr, 88 | const kv::KeyLocation & lhs, 89 | const coprocessor::KeyRange & rhs) 90 | { 91 | if (lhs.start_key == rhs.start_key && lhs.end_key == rhs.end_key) 92 | return ::testing::AssertionSuccess(); 93 | 94 | return ::testing::internal::EqFailure( 95 | lhs_expr, 96 | rhs_expr, 97 | toString(coprocessor::KeyRange{lhs.start_key, lhs.end_key}), 98 | toString(rhs), 99 | false); 100 | } 101 | #define ASSERT_LOC_KEY_RANGES_EQ(val1, val2) ASSERT_PRED_FORMAT2(::pingcap::tests::locationRangeCompare, val1, val2) 102 | #define EXPECT_LOC_KEY_RANGES_EQ(val1, val2) EXPECT_PRED_FORMAT2(::pingcap::tests::locationRangeCompare, val1, val2) 103 | 104 | } // namespace tests 105 | } // namespace pingcap 106 | -------------------------------------------------------------------------------- /include/pingcap/common/MPMCQueue.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | namespace pingcap 9 | { 10 | namespace common 11 | { 12 | 13 | enum class MPMCQueueResult 14 | { 15 | OK, 16 | CANCELLED, 17 | FINISHED, 18 | EMPTY, 19 | FULL, 20 | }; 21 | 22 | enum class MPMCQueueStatus 23 | { 24 | NORMAL, 25 | CANCELLED, 26 | FINISHED, 27 | }; 28 | 29 | template 30 | class IMPMCQueue 31 | { 32 | public: 33 | virtual ~IMPMCQueue() = default; 34 | 35 | virtual MPMCQueueResult tryPush(T &&) = 0; 36 | virtual MPMCQueueResult push(T &&) = 0; 37 | 38 | virtual MPMCQueueResult tryPop(T &) = 0; 39 | virtual MPMCQueueResult pop(T &) = 0; 40 | 41 | virtual bool cancel() = 0; 42 | 43 | virtual bool finish() = 0; 44 | }; 45 | 46 | template 47 | class MPMCQueue : public IMPMCQueue 48 | { 49 | public: 50 | MPMCQueue() 51 | : status(MPMCQueueStatus::NORMAL) 52 | {} 53 | 54 | ~MPMCQueue() override = default; 55 | 56 | MPMCQueueResult tryPush(T && t) override 57 | { 58 | return push(std::move(t)); 59 | } 60 | 61 | MPMCQueueResult push(T && t) override 62 | { 63 | std::lock_guard lk(mu); 64 | switch (status) 65 | { 66 | case MPMCQueueStatus::NORMAL: 67 | data.push(std::move(t)); 68 | cond_var.notify_all(); 69 | return MPMCQueueResult::OK; 70 | case MPMCQueueStatus::CANCELLED: 71 | return MPMCQueueResult::CANCELLED; 72 | case MPMCQueueStatus::FINISHED: 73 | return MPMCQueueResult::FINISHED; 74 | } 75 | __builtin_unreachable(); 76 | } 77 | 78 | MPMCQueueResult tryPop(T & t) override 79 | { 80 | std::lock_guard lk(mu); 81 | if (status == MPMCQueueStatus::CANCELLED) 82 | return MPMCQueueResult::CANCELLED; 83 | if (data.empty()) 84 | { 85 | if (status == MPMCQueueStatus::FINISHED) 86 | return MPMCQueueResult::FINISHED; 87 | return MPMCQueueResult::EMPTY; 88 | } 89 | t = std::move(data.front()); 90 | data.pop(); 91 | return MPMCQueueResult::OK; 92 | } 93 | 94 | MPMCQueueResult pop(T & t) override 95 | { 96 | std::unique_lock lk(mu); 97 | 98 | cond_var.wait(lk, [this] { return status != MPMCQueueStatus::NORMAL || !data.empty(); }); 99 | 100 | if (status == MPMCQueueStatus::CANCELLED) 101 | return MPMCQueueResult::CANCELLED; 102 | 103 | if (data.empty()) 104 | { 105 | assert(status == MPMCQueueStatus::FINISHED); 106 | return MPMCQueueResult::FINISHED; 107 | } 108 | t = std::move(data.front()); 109 | data.pop(); 110 | return MPMCQueueResult::OK; 111 | } 112 | 113 | bool cancel() override 114 | { 115 | std::lock_guard lk(mu); 116 | if (status == MPMCQueueStatus::NORMAL) 117 | { 118 | status = MPMCQueueStatus::CANCELLED; 119 | cond_var.notify_all(); 120 | return true; 121 | } 122 | return false; 123 | } 124 | 125 | bool finish() override 126 | { 127 | std::lock_guard lk(mu); 128 | if (status == MPMCQueueStatus::NORMAL) 129 | { 130 | status = MPMCQueueStatus::FINISHED; 131 | cond_var.notify_all(); 132 | return true; 133 | } 134 | return false; 135 | } 136 | 137 | private: 138 | MPMCQueueStatus status; 139 | 140 | std::mutex mu; 141 | std::condition_variable cond_var; 142 | 143 | std::queue data; 144 | }; 145 | 146 | } // namespace common 147 | } // namespace pingcap 148 | -------------------------------------------------------------------------------- /src/kv/Snapshot.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | 7 | namespace pingcap 8 | { 9 | namespace kv 10 | { 11 | constexpr int scan_batch_size = 256; 12 | 13 | //bool extractLockFromKeyErr() 14 | 15 | kvrpcpb::MvccInfo Snapshot::mvccGet(const std::string & key) 16 | { 17 | Backoffer bo(GetMaxBackoff); 18 | return mvccGet(bo, key); 19 | } 20 | 21 | kvrpcpb::MvccInfo Snapshot::mvccGet(Backoffer & bo, const std::string & key) 22 | { 23 | for (;;) 24 | { 25 | kvrpcpb::MvccGetByKeyRequest request; 26 | request.set_key(key); 27 | ::kvrpcpb::Context * context = request.mutable_context(); 28 | context->set_priority(::kvrpcpb::Normal); 29 | context->set_not_fill_cache(false); 30 | for (auto ts : min_commit_ts_pushed.getTimestamps()) 31 | { 32 | context->add_resolved_locks(ts); 33 | } 34 | 35 | auto location = cluster->region_cache->locateKey(bo, key); 36 | auto region_client = RegionClient(cluster, location.region); 37 | 38 | ::kvrpcpb::MvccGetByKeyResponse response; 39 | try 40 | { 41 | region_client.sendReqToRegion(bo, request, &response); 42 | } 43 | catch (Exception & e) 44 | { 45 | bo.backoff(boRegionMiss, e); 46 | continue; 47 | } 48 | if (!response.error().empty()) 49 | { 50 | Logger * log(&Logger::get("Snapshot::mvccGet")); 51 | log->warning("response error is " + response.error()); 52 | continue; 53 | } 54 | return response.info(); 55 | } 56 | } 57 | 58 | std::string Snapshot::Get(const std::string & key) 59 | { 60 | Backoffer bo(GetMaxBackoff); 61 | return Get(bo, key); 62 | } 63 | 64 | std::string Snapshot::Get(Backoffer & bo, const std::string & key) 65 | { 66 | for (;;) 67 | { 68 | kvrpcpb::GetRequest request; 69 | request.set_key(key); 70 | request.set_version(version); 71 | ::kvrpcpb::Context * context = request.mutable_context(); 72 | context->set_priority(::kvrpcpb::Normal); 73 | context->set_not_fill_cache(false); 74 | for (auto ts : min_commit_ts_pushed.getTimestamps()) 75 | { 76 | context->add_resolved_locks(ts); 77 | } 78 | 79 | auto location = cluster->region_cache->locateKey(bo, key); 80 | auto region_client = RegionClient(cluster, location.region); 81 | 82 | ::kvrpcpb::GetResponse response; 83 | try 84 | { 85 | region_client.sendReqToRegion(bo, request, &response); 86 | } 87 | catch (Exception & e) 88 | { 89 | bo.backoff(boRegionMiss, e); 90 | continue; 91 | } 92 | if (response.has_error()) 93 | { 94 | auto lock = extractLockFromKeyErr(response.error()); 95 | std::vector locks{lock}; 96 | std::vector pushed; 97 | auto before_expired = cluster->lock_resolver->resolveLocks(bo, version, locks, pushed); 98 | 99 | if (!pushed.empty()) 100 | { 101 | min_commit_ts_pushed.addTimestamps(pushed); 102 | } 103 | if (before_expired > 0) 104 | { 105 | bo.backoffWithMaxSleep( 106 | boTxnLockFast, 107 | before_expired, 108 | Exception("key error : " + response.error().ShortDebugString(), LockError)); 109 | } 110 | continue; 111 | } 112 | return response.value(); 113 | } 114 | } 115 | 116 | Scanner Snapshot::Scan(const std::string & begin, const std::string & end) 117 | { 118 | return Scanner(*this, begin, end, scan_batch_size); 119 | } 120 | 121 | } // namespace kv 122 | } // namespace pingcap 123 | -------------------------------------------------------------------------------- /include/pingcap/common/MPPProber.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | 6 | #include 7 | #include 8 | #include 9 | #include 10 | 11 | 12 | namespace pingcap 13 | { 14 | namespace kv 15 | { 16 | struct Cluster; 17 | } // namespace kv 18 | namespace common 19 | { 20 | 21 | using TimePoint = std::chrono::time_point; 22 | static constexpr TimePoint INVALID_TIME_POINT = std::chrono::steady_clock::time_point::max(); 23 | static constexpr auto MAX_RECOVERY_TIME_LIMIT = std::chrono::minutes(15); 24 | static constexpr auto MAX_OBSOLETE_TIME_LIMIT = std::chrono::hours(1); 25 | static constexpr auto SCAN_INTERVAL = std::chrono::seconds(1); // scan per 1s. 26 | static constexpr auto DETECT_PERIOD = std::chrono::seconds(3); // do real alive rpc per 3s. 27 | static constexpr size_t DETECT_RPC_TIMEOUT = 2; 28 | 29 | inline std::chrono::seconds getElapsed(const TimePoint & ago) 30 | { 31 | auto now = std::chrono::steady_clock::now(); 32 | return std::chrono::duration_cast(now - ago); 33 | } 34 | 35 | 36 | bool detectStore(kv::RpcClientPtr & rpc_client, const std::string & store_addr, int rpc_timeout, Logger * log); 37 | 38 | struct ProbeState 39 | { 40 | ProbeState(const std::string & store_addr_, pingcap::kv::Cluster * cluster_) 41 | : store_addr(store_addr_) 42 | , cluster(cluster_) 43 | , log(&Logger::get("pingcap.ProbeState")) 44 | , recovery_timepoint(INVALID_TIME_POINT) 45 | , last_lookup_timepoint(INVALID_TIME_POINT) 46 | , last_detect_timepoint(INVALID_TIME_POINT) 47 | {} 48 | 49 | std::string store_addr; 50 | pingcap::kv::Cluster * cluster; 51 | Logger * log; 52 | TimePoint recovery_timepoint; 53 | TimePoint last_lookup_timepoint; 54 | TimePoint last_detect_timepoint; 55 | std::mutex state_lock; 56 | 57 | void detectAndUpdateState(const std::chrono::seconds & detect_period, size_t detect_rpc_timeout); 58 | }; 59 | 60 | // The main purpose of MPPProber is to prevent excessive delays caused by repeatedly probing failed stores. 61 | // 62 | // MPPProber continuously probes failed_stores in the background. If a probe succeeds, it sets the store's recovery_time. 63 | // Callers check this recovery_time and determine whether the store has recovered based on a TTL using MPPProber::isRecovery(). 64 | // 65 | // If a store is considered recovered, the caller will re-trigger a liveness probe. 66 | // If the probe fails again, the store will be re-added to failed_stores; otherwise, it can be used directly. 67 | // 68 | // If a store was previously attempted to be used (last_lookup_time exists) but couldn't actually be used, 69 | // and this state persists beyond MAX_OBSOLETE_TIME_LIMIT, it will be removed from failed_stores to avoid being continuously probed in the background. 70 | class MPPProber 71 | { 72 | public: 73 | explicit MPPProber(pingcap::kv::Cluster * cluster_) 74 | : cluster(cluster_) 75 | , scan_interval(SCAN_INTERVAL) 76 | , detect_period(DETECT_PERIOD) 77 | , detect_rpc_timeout(DETECT_RPC_TIMEOUT) 78 | , log(&Logger::get("pingcap.MPPProber")) 79 | , stopped(false) 80 | {} 81 | 82 | void run(); 83 | void stop(); 84 | 85 | // Return true is this store is alive, false if dead. 86 | bool isRecovery(const std::string & store_addr, const std::chrono::seconds & recovery_ttl); 87 | // Tag store as dead. 88 | void add(const std::string & store_addr); 89 | 90 | private: 91 | using FailedStoreMap = std::unordered_map>; 92 | 93 | void scan(); 94 | void detect(); 95 | 96 | pingcap::kv::Cluster * cluster; 97 | std::chrono::seconds scan_interval; 98 | std::chrono::seconds detect_period; 99 | size_t detect_rpc_timeout; 100 | Logger * log; 101 | std::atomic stopped; 102 | FailedStoreMap failed_stores; 103 | std::mutex store_lock; 104 | 105 | std::mutex scan_mu; 106 | std::condition_variable scan_cv; 107 | }; 108 | } // namespace common 109 | } // namespace pingcap 110 | -------------------------------------------------------------------------------- /src/test/lock_resolve_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include "mock_tikv.h" 9 | #include "test_helper.h" 10 | 11 | namespace pingcap::kv 12 | { 13 | extern BackoffPtr newBackoff(BackoffType); 14 | } 15 | 16 | namespace pingcap::tests 17 | { 18 | using namespace pingcap; 19 | using namespace pingcap::kv; 20 | 21 | class TestWithLockResolve : public testing::Test 22 | { 23 | protected: 24 | void SetUp() override 25 | { 26 | fiu_init(0); 27 | mock_kv_cluster = mockkv::initCluster(); 28 | std::vector pd_addrs = mock_kv_cluster->pd_addrs; 29 | 30 | test_cluster = createCluster(pd_addrs); 31 | control_cluster = createCluster(pd_addrs); 32 | } 33 | 34 | mockkv::ClusterPtr mock_kv_cluster; 35 | 36 | ClusterPtr test_cluster; 37 | ClusterPtr control_cluster; 38 | }; 39 | 40 | TEST_F(TestWithLockResolve, testResolveLockGet) 41 | { 42 | // Write First Time and Split int two regions. 43 | { 44 | Txn txn(test_cluster.get()); 45 | 46 | txn.set("abc", "1"); 47 | txn.set("abd", "2"); 48 | txn.set("abe", "3"); 49 | txn.set("abf", "4"); 50 | txn.set("abg", "5"); 51 | txn.set("abz", "6"); 52 | txn.commit(); 53 | control_cluster->splitRegion("abf"); 54 | } 55 | 56 | // and write again, but second region commits failed. 57 | { 58 | fiu_enable("rest commit fail", 1, nullptr, FIU_ONETIME); 59 | Txn txn(test_cluster.get()); 60 | 61 | txn.set("abc", "6"); 62 | txn.set("abd", "5"); 63 | txn.set("abe", "4"); 64 | txn.set("abf", "3"); 65 | txn.set("abg", "2"); 66 | txn.set("abz", "1"); 67 | txn.commit(); 68 | 69 | Snapshot snap(test_cluster.get()); 70 | 71 | std::string result = snap.Get("abe"); 72 | 73 | ASSERT_EQ(result, "4"); 74 | 75 | result = snap.Get("abz"); 76 | 77 | ASSERT_EQ(result, "1"); 78 | } 79 | 80 | // and write again, all commits succeed 81 | { 82 | Txn txn(test_cluster.get()); 83 | 84 | txn.set("abc", "1"); 85 | txn.set("abd", "2"); 86 | txn.set("abe", "3"); 87 | txn.set("abf", "4"); 88 | txn.set("abg", "5"); 89 | txn.set("abz", "6"); 90 | txn.commit(); 91 | 92 | Snapshot snap(test_cluster.get()); 93 | std::string result = snap.Get("abe"); 94 | 95 | ASSERT_EQ(result, "3"); 96 | 97 | result = snap.Get("abz"); 98 | 99 | ASSERT_EQ(result, "6"); 100 | } 101 | 102 | { 103 | fiu_enable("all commit fail", 1, nullptr, FIU_ONETIME); 104 | Txn txn(test_cluster.get()); 105 | 106 | txn.set("abc", "6"); 107 | txn.set("abd", "5"); 108 | txn.set("abe", "4"); 109 | txn.set("abf", "3"); 110 | txn.set("abg", "2"); 111 | txn.set("abz", "1"); 112 | txn.commit(); 113 | 114 | Snapshot snap(test_cluster.get()); 115 | 116 | std::string result = snap.Get("abe"); 117 | 118 | ASSERT_EQ(result, "3"); 119 | 120 | result = snap.Get("abz"); 121 | 122 | ASSERT_EQ(result, "6"); 123 | } 124 | } 125 | 126 | 127 | TEST_F(TestWithLockResolve, testResolveLockBase) 128 | { 129 | { 130 | Backoffer bo(kv::copNextMaxBackoff); 131 | for (int i = 0; i <= 12; ++i) 132 | { 133 | auto t = static_cast(i); 134 | bo.backoff_map.emplace(t, newBackoff(t)); 135 | } 136 | ASSERT_EQ(bo.backoff_map.size(), 13); 137 | for (int i = 0; i <= 12; ++i) 138 | { 139 | auto t = static_cast(i); 140 | bo.backoff(t, {}); 141 | } 142 | 143 | auto && new_bo = bo.clone(); 144 | ASSERT_EQ(new_bo.max_sleep, bo.max_sleep); 145 | ASSERT_EQ(new_bo.total_sleep, bo.total_sleep); 146 | ASSERT_EQ(new_bo.backoff_map.size(), bo.backoff_map.size()); 147 | for (auto && [k, v] : bo.backoff_map) 148 | { 149 | ASSERT_NE(v.get(), new_bo.backoff_map.at(k).get()); 150 | ASSERT_EQ(std::memcmp(v.get(), new_bo.backoff_map.at(k).get(), sizeof(Backoff)), 0); 151 | } 152 | } 153 | } 154 | 155 | } // namespace pingcap::tests 156 | -------------------------------------------------------------------------------- /src/kv/Scanner.cc: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | namespace pingcap 4 | { 5 | namespace kv 6 | { 7 | void Scanner::next() 8 | { 9 | Backoffer bo(scanMaxBackoff); 10 | if (!valid) 11 | { 12 | throw Exception("the scanner is invalid", LogicalError); 13 | } 14 | 15 | for (;;) 16 | { 17 | idx++; 18 | if (idx >= cache.size()) 19 | { 20 | if (eof) 21 | { 22 | valid = false; 23 | return; 24 | } 25 | getData(bo); 26 | if (idx >= cache.size()) 27 | { 28 | continue; 29 | } 30 | } 31 | 32 | auto & current = cache[idx]; 33 | if (!end_key.empty() && current.key() >= end_key) 34 | { 35 | eof = true; 36 | valid = false; 37 | } 38 | 39 | if (current.has_error()) 40 | { 41 | resolveCurrentLock(bo, current); 42 | } 43 | return; 44 | } 45 | } 46 | 47 | void Scanner::resolveCurrentLock(pingcap::kv::Backoffer & bo, kvrpcpb::KvPair & current) 48 | { 49 | auto value = snap.Get(bo, current.key()); 50 | current.set_allocated_error(nullptr); 51 | current.set_value(value); 52 | } 53 | 54 | void Scanner::getData(Backoffer & bo) 55 | { 56 | log->trace("get data for scanner"); 57 | for (;;) 58 | { 59 | auto loc = snap.cluster->region_cache->locateKey(bo, next_start_key); 60 | auto req_end_key = end_key; 61 | if (!req_end_key.empty() && !loc.end_key.empty() && loc.end_key < req_end_key) 62 | req_end_key = loc.end_key; 63 | 64 | 65 | auto region_client = RegionClient(snap.cluster, loc.region); 66 | kvrpcpb::ScanRequest request; 67 | request.set_start_key(next_start_key); 68 | request.set_end_key(req_end_key); 69 | request.set_limit(batch); 70 | request.set_version(snap.version); 71 | request.set_key_only(false); 72 | 73 | auto * context = request.mutable_context(); 74 | context->set_priority(::kvrpcpb::Normal); 75 | context->set_not_fill_cache(false); 76 | 77 | kvrpcpb::ScanResponse response; 78 | try 79 | { 80 | region_client.sendReqToRegion(bo, request, &response); 81 | } 82 | catch (Exception & e) 83 | { 84 | bo.backoff(boRegionMiss, e); 85 | continue; 86 | } 87 | 88 | // TODO Check safe point. 89 | 90 | // TiKV will only return locked keys if there is response level error 91 | if (response.has_error()) 92 | { 93 | auto lock = extractLockFromKeyErr(response.error()); 94 | std::vector locks{lock}; 95 | std::vector pushed{}; 96 | auto ms_before_expired = snap.cluster->lock_resolver->resolveLocks(bo, snap.version, locks, pushed); 97 | if (ms_before_expired > 0) 98 | { 99 | bo.backoffWithMaxSleep( 100 | BackoffType::boTxnLockFast, 101 | ms_before_expired, 102 | Exception("key is locked during scanning", ErrorCodes::LockError)); 103 | } 104 | continue; 105 | } 106 | 107 | int pairs_size = response.pairs_size(); 108 | idx = 0; 109 | cache.clear(); 110 | for (int i = 0; i < pairs_size; i++) 111 | { 112 | auto current = response.pairs(i); 113 | if (current.has_error()) 114 | { 115 | auto lock = extractLockFromKeyErr(current.error()); 116 | current.set_key(lock->key); 117 | } 118 | cache.push_back(current); 119 | } 120 | 121 | log->trace("get pair size: " + std::to_string(pairs_size)); 122 | 123 | if (pairs_size < batch) 124 | { 125 | next_start_key = loc.end_key; 126 | 127 | // If the end key is empty, it infers this region is last and should stop scan. 128 | if (loc.end_key.empty() || (next_start_key) >= end_key) 129 | { 130 | eof = true; 131 | } 132 | 133 | return; 134 | } 135 | 136 | auto last_key = cache.back(); 137 | next_start_key = alphabeticalNext(last_key.key()); 138 | log->trace("scan next key: " + next_start_key); 139 | return; 140 | } 141 | } 142 | } // namespace kv 143 | } // namespace pingcap 144 | -------------------------------------------------------------------------------- /include/pingcap/pd/MockPDClient.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | namespace pingcap::pd 9 | { 10 | using Clock = std::chrono::system_clock; 11 | 12 | class MockPDClient : public IClient 13 | { 14 | public: 15 | static constexpr uint64_t MOCKED_GC_SAFE_POINT = 10000000; 16 | 17 | public: 18 | MockPDClient() = default; 19 | 20 | ~MockPDClient() override = default; 21 | 22 | uint64_t getGCSafePoint() override { return MOCKED_GC_SAFE_POINT; } 23 | 24 | uint64_t getGCSafePointV2(KeyspaceID) override { return MOCKED_GC_SAFE_POINT; } 25 | 26 | pdpb::GetGCStateResponse getGCState(KeyspaceID keyspace_id) override 27 | { 28 | pdpb::GetGCStateResponse gc_state; 29 | auto * hdr = gc_state.mutable_header(); 30 | hdr->set_cluster_id(1); 31 | hdr->mutable_error()->set_type(pdpb::ErrorType::OK); 32 | auto * state = gc_state.mutable_gc_state(); 33 | state->mutable_keyspace_scope()->set_keyspace_id(keyspace_id); 34 | state->set_is_keyspace_level_gc(true); 35 | state->set_txn_safe_point(MOCKED_GC_SAFE_POINT); 36 | state->set_gc_safe_point(MOCKED_GC_SAFE_POINT); 37 | return gc_state; 38 | } 39 | 40 | pdpb::GetAllKeyspacesGCStatesResponse getAllKeyspacesGCStates() override 41 | { 42 | pdpb::GetAllKeyspacesGCStatesResponse all_states; 43 | auto * hdr = all_states.mutable_header(); 44 | hdr->set_cluster_id(1); 45 | hdr->mutable_error()->set_type(pdpb::ErrorType::OK); 46 | auto * state = all_states.add_gc_states(); 47 | state->mutable_keyspace_scope()->set_keyspace_id(1); 48 | state->set_is_keyspace_level_gc(true); 49 | state->set_txn_safe_point(MOCKED_GC_SAFE_POINT); 50 | state->set_gc_safe_point(MOCKED_GC_SAFE_POINT); 51 | return all_states; 52 | } 53 | 54 | uint64_t getTS() override { return Clock::now().time_since_epoch().count(); } 55 | 56 | pdpb::GetRegionResponse getRegionByKey(const std::string &) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 57 | 58 | pdpb::GetRegionResponse getRegionByID(uint64_t) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 59 | 60 | metapb::Store getStore(uint64_t) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 61 | std::vector getAllStores(bool) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 62 | 63 | bool isClusterBootstrapped() override { return true; } 64 | 65 | KeyspaceID getKeyspaceID(const std::string & /*keyspace_name*/) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 66 | 67 | void update(const std::vector & /*addrs*/, const ClusterConfig & /*config_*/) override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 68 | 69 | bool isMock() override { return true; } 70 | 71 | std::string getLeaderUrl() override { throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); } 72 | 73 | ::resource_manager::ListResourceGroupsResponse listResourceGroups(const ::resource_manager::ListResourceGroupsRequest &) override 74 | { 75 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 76 | } 77 | 78 | ::resource_manager::GetResourceGroupResponse getResourceGroup(const ::resource_manager::GetResourceGroupRequest &) override 79 | { 80 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 81 | } 82 | 83 | ::resource_manager::PutResourceGroupResponse addResourceGroup(const ::resource_manager::PutResourceGroupRequest &) override 84 | { 85 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 86 | } 87 | 88 | ::resource_manager::PutResourceGroupResponse modifyResourceGroup(const ::resource_manager::PutResourceGroupRequest &) override 89 | { 90 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 91 | } 92 | 93 | ::resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const ::resource_manager::DeleteResourceGroupRequest &) override 94 | { 95 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 96 | } 97 | 98 | resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest &) override 99 | { 100 | throw Exception("not implemented", pingcap::ErrorCodes::UnknownError); 101 | } 102 | }; 103 | 104 | } // namespace pingcap::pd 105 | -------------------------------------------------------------------------------- /src/test/bank_test/bank_test_schrodinger.cc: -------------------------------------------------------------------------------- 1 | #include "../test_helper.h" 2 | #include "bank_test.h" 3 | #include "schrodinger_client.h" 4 | 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | using namespace pingcap::tests; 11 | 12 | void RunBankCaseOnline(int account, int con, int run_time) 13 | { 14 | ::sleep(10); 15 | Client client; 16 | ::sleep(2); 17 | ClusterPtr cluster = createCluster(client.PDs()); 18 | std::cerr << "end create cluster\n"; 19 | BankCase bank(cluster.get(), account, con); 20 | bank.initialize(); 21 | auto close_thread = std::thread([&]() { 22 | std::this_thread::sleep_for(std::chrono::seconds(run_time)); 23 | bank.close(); 24 | }); 25 | bank.execute(); 26 | close_thread.join(); 27 | } 28 | 29 | void RunBankCaseCheck(const std::vector & pd_addr, int account) 30 | { 31 | ClusterPtr cluster = createCluster(pd_addr); 32 | BankCase bank(cluster.get(), account, 10); 33 | bank.verify(); 34 | } 35 | 36 | void RunBankCaseLocal(const std::vector & pd_addr, int account, int con, int run_time, bool init = true) 37 | { 38 | ClusterPtr cluster = createCluster(pd_addr); 39 | BankCase bank(cluster.get(), account, con); 40 | if (init) 41 | { 42 | bank.initialize(); 43 | } 44 | else 45 | { 46 | bank.enable_check(); 47 | } 48 | auto close_thread = std::thread([&]() { 49 | std::this_thread::sleep_for(std::chrono::seconds(run_time)); 50 | bank.close(); 51 | }); 52 | bank.execute(); 53 | close_thread.join(); 54 | } 55 | 56 | using namespace Poco::Util; 57 | 58 | class BankApp : public Application 59 | { 60 | private: 61 | std::vector pd_addrs; 62 | int concurrency; 63 | int account; 64 | int run_time; 65 | bool check_only; 66 | bool check_exec; 67 | 68 | void setConcurrency(const std::string &, const std::string & con) { concurrency = std::stoi(con); } 69 | 70 | void setPDAddr(const std::string &, const std::string & pd_addr) 71 | { 72 | pd_addrs.push_back(pd_addr); 73 | } 74 | 75 | void setAccount(const std::string &, const std::string & acc) { account = std::stoi(acc); } 76 | 77 | void setRunTime(const std::string &, const std::string & time) { run_time = std::stoi(time); } 78 | 79 | void setCheckOnly(const std::string &, const std::string &) { check_only = true; } 80 | 81 | void setCheckAndExec(const std::string &, const std::string &) { check_exec = true; } 82 | 83 | 84 | public: 85 | BankApp() : Application(), concurrency(10), account(100000), run_time(600), check_only(false), check_exec(false) 86 | { 87 | Logger::get("pingcap.tikv").setLevel("debug"); 88 | Logger::get("pingcap.pd").setLevel("debug"); 89 | } 90 | 91 | void defineOptions(OptionSet & options) override 92 | { 93 | options.addOption(Option("pd-address", "p") 94 | .repeatable(true) 95 | .argument("pd_address", false) 96 | .callback(OptionCallback(this, &BankApp::setPDAddr))); 97 | options.addOption(Option("concurrency", "c") 98 | .validator(new IntValidator(1, 1000)) 99 | .argument("concurrency", false) 100 | .callback(OptionCallback(this, &BankApp::setConcurrency))); 101 | options.addOption(Option("account", "a").argument("account", false).callback(OptionCallback(this, &BankApp::setAccount))); 102 | options.addOption(Option("time", "t").argument("time", false).callback(OptionCallback(this, &BankApp::setRunTime))); 103 | options.addOption(Option("check-only", "C").callback(OptionCallback(this, &BankApp::setCheckOnly))); 104 | options.addOption(Option("check-exec", "E").callback(OptionCallback(this, &BankApp::setCheckAndExec))); 105 | } 106 | int main(const std::vector &) override 107 | { 108 | 109 | if (pd_addrs.size() == 0) 110 | { 111 | RunBankCaseOnline(account, concurrency, run_time); 112 | } 113 | else 114 | { 115 | if (check_only) 116 | { 117 | RunBankCaseCheck(pd_addrs, account); 118 | } 119 | else if (check_exec) 120 | { 121 | RunBankCaseLocal(pd_addrs, account, concurrency, run_time, false); 122 | } 123 | else 124 | { 125 | RunBankCaseLocal(pd_addrs, account, concurrency, run_time); 126 | } 127 | } 128 | return 0; 129 | } 130 | }; 131 | 132 | POCO_APP_MAIN(BankApp) 133 | -------------------------------------------------------------------------------- /src/kv/RegionClient.cc: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | namespace pingcap 4 | { 5 | namespace kv 6 | { 7 | void RegionClient::onRegionError(Backoffer & bo, RPCContextPtr rpc_ctx, const errorpb::Error & err) const 8 | { 9 | if (err.has_not_leader()) 10 | { 11 | const auto & not_leader = err.not_leader(); 12 | if (not_leader.has_leader()) 13 | { 14 | // don't backoff if a new leader is returned. 15 | log->information("not leader but has leader, region_ver_id=" + rpc_ctx->region.toString() + ", new leader peer_id=" + std::to_string(not_leader.leader().id()) 16 | + ", store_id=" + std::to_string(not_leader.leader().store_id())); 17 | if (!cluster->region_cache->updateLeader(rpc_ctx->region, not_leader.leader())) 18 | { 19 | bo.backoff(boRegionScheduling, Exception("not leader, ctx: " + rpc_ctx->toString(), NotLeader)); 20 | } 21 | } 22 | else 23 | { 24 | // The peer doesn't know who is the current leader. Generally it's because 25 | // the Raft group is in an election, but it's possible that the peer is 26 | // isolated and removed from the Raft group. So it's necessary to reload 27 | // the region from PD. 28 | log->information("not leader but doesn't have new leader, region_ver_id=" + rpc_ctx->region.toString()); 29 | cluster->region_cache->dropRegion(rpc_ctx->region); 30 | bo.backoff(boRegionScheduling, Exception("not leader, ctx: " + rpc_ctx->toString(), NotLeader)); 31 | } 32 | return; 33 | } 34 | 35 | if (err.has_disk_full()) 36 | { 37 | bo.backoff(boTiKVDiskFull, Exception("tikv disk full: " + err.disk_full().DebugString() + ", ctx: " + rpc_ctx->toString())); 38 | } 39 | 40 | if (err.has_store_not_match()) 41 | { 42 | cluster->region_cache->dropStore(rpc_ctx->peer.store_id()); 43 | cluster->region_cache->dropRegion(rpc_ctx->region); 44 | return; 45 | } 46 | 47 | if (err.has_epoch_not_match()) 48 | { 49 | cluster->region_cache->onRegionStale(bo, rpc_ctx, err.epoch_not_match()); 50 | // Epoch not match should not retry, throw exception directly !! 51 | throw Exception("Region epoch not match for region " + rpc_ctx->region.toString() + ".", RegionEpochNotMatch); 52 | } 53 | 54 | if (err.has_server_is_busy()) 55 | { 56 | bo.backoff(boServerBusy, Exception("server is busy: " + err.server_is_busy().reason(), ServerIsBusy)); 57 | return; 58 | } 59 | 60 | if (err.has_stale_command()) 61 | { 62 | return; 63 | } 64 | 65 | if (err.has_raft_entry_too_large()) 66 | { 67 | throw Exception("entry too large", RaftEntryTooLarge); 68 | } 69 | 70 | if (err.has_max_timestamp_not_synced()) 71 | { 72 | bo.backoff(boMaxTsNotSynced, Exception("max timestamp not synced, ctx: " + rpc_ctx->toString())); 73 | return; 74 | } 75 | 76 | // A read request may be sent to a peer which has not been initialized yet, we should retry in this case. 77 | if (err.has_region_not_initialized()) 78 | { 79 | bo.backoff(boMaxRegionNotInitialized, Exception("region not initialized, ctx: " + rpc_ctx->toString())); 80 | return; 81 | } 82 | 83 | // The read-index can't be handled timely because the region is splitting or merging. 84 | if (err.has_read_index_not_ready()) 85 | { 86 | // The region can't provide service until split or merge finished, so backoff. 87 | bo.backoff(boRegionScheduling, Exception("read index not ready, ctx: " + rpc_ctx->toString())); 88 | return; 89 | } 90 | 91 | if (err.has_proposal_in_merging_mode()) 92 | { 93 | // The region is merging and it can't provide service until merge finished, so backoff. 94 | bo.backoff(boRegionScheduling, Exception("region is merging, ctx: " + rpc_ctx->toString())); 95 | return; 96 | } 97 | 98 | // A stale read request may be sent to a peer which the data is not ready yet, we should retry in this case. 99 | // This error is specific to stale read and the target replica is randomly selected. If the request is sent 100 | // to the leader, the data must be ready, so we don't backoff here. 101 | if (err.has_data_is_not_ready()) 102 | { 103 | bo.backoff(boMaxDataNotReady, Exception("data is not ready, ctx: " + rpc_ctx->toString())); 104 | } 105 | 106 | cluster->region_cache->dropRegion(rpc_ctx->region); 107 | } 108 | 109 | void RegionClient::onSendFail(Backoffer & bo, const Exception & e, RPCContextPtr rpc_ctx) const 110 | { 111 | cluster->region_cache->onSendReqFail(rpc_ctx, e); 112 | // Retry on send request failure when it's not canceled. 113 | // When a store is not available, the leader of related region should be elected quickly. 114 | bo.backoff(boTiKVRPC, e); 115 | } 116 | 117 | } // namespace kv 118 | } // namespace pingcap 119 | -------------------------------------------------------------------------------- /src/test/coprocessor_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | #include "mock_tikv.h" 6 | #include "test_helper.h" 7 | 8 | namespace pingcap::tests 9 | { 10 | using namespace pingcap; 11 | using namespace pingcap::kv; 12 | 13 | class TestCoprocessor : public testing::Test 14 | { 15 | protected: 16 | void SetUp() override 17 | { 18 | fiu_init(0); 19 | 20 | mock_kv_cluster = mockkv::initCluster(); 21 | std::vector pd_addrs = mock_kv_cluster->pd_addrs; 22 | 23 | test_cluster = createCluster(pd_addrs); 24 | control_cluster = createCluster(pd_addrs); 25 | } 26 | 27 | mockkv::ClusterPtr mock_kv_cluster; 28 | 29 | ClusterPtr test_cluster; 30 | ClusterPtr control_cluster; 31 | }; 32 | 33 | 34 | TEST_F(TestCoprocessor, BuildTask) 35 | { 36 | Backoffer bo(copBuildTaskMaxBackoff); 37 | 38 | control_cluster->splitRegion("a"); 39 | control_cluster->splitRegion("b"); 40 | control_cluster->splitRegion("z"); 41 | 42 | pingcap::coprocessor::KeyRanges ranges{{"a", "z"}}; 43 | 44 | std::shared_ptr req = std::make_shared(); 45 | req->tp = pingcap::coprocessor::DAG; 46 | req->start_ts = test_cluster->pd_client->getTS(); 47 | 48 | auto tasks = pingcap::coprocessor::buildCopTasks( 49 | bo, 50 | test_cluster.get(), 51 | ranges, 52 | req, 53 | kv::StoreType::TiKV, 54 | pd::NullspaceID, 55 | 0, 56 | "", 57 | &Logger::get("pingcap/coprocessor")); 58 | 59 | ASSERT_EQ(tasks.size(), 2); 60 | 61 | ASSERT_EQ(tasks[0].ranges[0].start_key, "a"); 62 | ASSERT_EQ(tasks[0].ranges[0].end_key, "b"); 63 | ASSERT_EQ(tasks[1].ranges[0].start_key, "b"); 64 | ASSERT_EQ(tasks[1].ranges[0].end_key, "z"); 65 | 66 | control_cluster->splitRegion("d"); 67 | control_cluster->splitRegion("e"); 68 | 69 | fiu_enable("sleep_before_push_result", 1, nullptr, 0); 70 | auto queue = std::make_unique>(); 71 | auto tasks2 = tasks; 72 | pingcap::coprocessor::ResponseIter iter(std::move(queue), std::move(tasks), test_cluster.get(), 8, &Logger::get("pingcap/coprocessor")); 73 | iter.open(); 74 | 75 | for (int i = 0; i < 4; i++) 76 | { 77 | ASSERT_EQ(iter.next().second, true); 78 | } 79 | ASSERT_EQ(iter.next().second, false); 80 | 81 | auto queue2 = std::make_unique>(); 82 | pingcap::coprocessor::ResponseIter iter2(std::move(queue2), std::move(tasks2), test_cluster.get(), 8, &Logger::get("pingcap/coprocessor")); 83 | iter2.open(); 84 | 85 | ASSERT_EQ(iter2.next().second, true); 86 | ASSERT_EQ(iter2.next().second, true); 87 | iter2.cancel(); 88 | ASSERT_EQ(iter2.next().second, false); 89 | } 90 | 91 | TEST_F(TestCoprocessor, BuildTaskStream) 92 | { 93 | Backoffer bo(copBuildTaskMaxBackoff); 94 | 95 | control_cluster->splitRegion("a"); 96 | control_cluster->splitRegion("b"); 97 | control_cluster->splitRegion("z"); 98 | 99 | pingcap::coprocessor::KeyRanges ranges{{"a", "z"}}; 100 | 101 | std::shared_ptr req = std::make_shared(); 102 | req->tp = pingcap::coprocessor::DAG; 103 | req->start_ts = test_cluster->pd_client->getTS(); 104 | 105 | auto tasks = pingcap::coprocessor::buildCopTasks( 106 | bo, 107 | test_cluster.get(), 108 | ranges, 109 | req, 110 | kv::StoreType::TiKV, 111 | pd::NullspaceID, 112 | 0, 113 | "", 114 | &Logger::get("pingcap/coprocessor")); 115 | 116 | ASSERT_EQ(tasks.size(), 2); 117 | 118 | ASSERT_EQ(tasks[0].ranges[0].start_key, "a"); 119 | ASSERT_EQ(tasks[0].ranges[0].end_key, "b"); 120 | ASSERT_EQ(tasks[1].ranges[0].start_key, "b"); 121 | ASSERT_EQ(tasks[1].ranges[0].end_key, "z"); 122 | 123 | control_cluster->splitRegion("d"); 124 | control_cluster->splitRegion("e"); 125 | 126 | fiu_enable("sleep_before_push_result", 1, nullptr, 0); 127 | auto queue = std::make_unique>(); 128 | auto tasks2 = tasks; 129 | pingcap::coprocessor::ResponseIter iter(std::move(queue), std::move(tasks), test_cluster.get(), 8, &Logger::get("pingcap/coprocessor")); 130 | iter.open(); 131 | 132 | for (int i = 0; i < 4; i++) 133 | { 134 | ASSERT_EQ(iter.next().second, true); 135 | } 136 | ASSERT_EQ(iter.next().second, false); 137 | 138 | auto queue2 = std::make_unique>(); 139 | pingcap::coprocessor::ResponseIter iter2(std::move(queue2), std::move(tasks2), test_cluster.get(), 8, &Logger::get("pingcap/coprocessor")); 140 | iter2.open(); 141 | 142 | ASSERT_EQ(iter2.next().second, true); 143 | ASSERT_EQ(iter2.next().second, true); 144 | iter2.cancel(); 145 | ASSERT_EQ(iter2.next().second, false); 146 | } 147 | 148 | } // namespace pingcap::tests 149 | -------------------------------------------------------------------------------- /include/pingcap/pd/Client.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | #include 13 | 14 | #include 15 | #include 16 | #include 17 | #include 18 | #include 19 | 20 | namespace pingcap::pd 21 | { 22 | 23 | class Client : public IClient 24 | { 25 | const int max_init_cluster_retries; 26 | 27 | const std::chrono::seconds pd_timeout; 28 | 29 | const std::chrono::microseconds loop_interval; 30 | 31 | const std::chrono::seconds update_leader_interval; 32 | 33 | void init(const std::vector & addrs, const ClusterConfig & config_); 34 | 35 | void uninit(); 36 | 37 | public: 38 | Client(const std::vector & addrs, const ClusterConfig & config); 39 | 40 | ~Client() override; 41 | 42 | void update(const std::vector & addrs, const ClusterConfig & config) override; 43 | 44 | // only implement a weak get ts. 45 | uint64_t getTS() override; 46 | 47 | pdpb::GetRegionResponse getRegionByKey(const std::string & key) override; 48 | 49 | pdpb::GetRegionResponse getRegionByID(uint64_t region_id) override; 50 | 51 | metapb::Store getStore(uint64_t store_id) override; 52 | 53 | std::vector getAllStores(bool exclude_tombstone) override; 54 | 55 | bool isClusterBootstrapped() override; 56 | 57 | uint64_t getGCSafePoint() override; 58 | 59 | uint64_t getGCSafePointV2(KeyspaceID keyspace_id) override; 60 | 61 | pdpb::GetGCStateResponse getGCState(KeyspaceID keyspace_id) override; 62 | 63 | pdpb::GetAllKeyspacesGCStatesResponse getAllKeyspacesGCStates() override; 64 | 65 | KeyspaceID getKeyspaceID(const std::string & keyspace_name) override; 66 | 67 | bool isMock() override; 68 | 69 | std::string getLeaderUrl() override; 70 | 71 | // ResourceControl related. 72 | resource_manager::ListResourceGroupsResponse listResourceGroups(const resource_manager::ListResourceGroupsRequest &) override; 73 | 74 | resource_manager::GetResourceGroupResponse getResourceGroup(const resource_manager::GetResourceGroupRequest &) override; 75 | 76 | resource_manager::PutResourceGroupResponse addResourceGroup(const resource_manager::PutResourceGroupRequest &) override; 77 | 78 | resource_manager::PutResourceGroupResponse modifyResourceGroup(const resource_manager::PutResourceGroupRequest &) override; 79 | 80 | resource_manager::DeleteResourceGroupResponse deleteResourceGroup(const resource_manager::DeleteResourceGroupRequest &) override; 81 | 82 | resource_manager::TokenBucketsResponse acquireTokenBuckets(const resource_manager::TokenBucketsRequest & req) override; 83 | 84 | private: 85 | void initClusterID(); 86 | 87 | void updateLeader(); 88 | 89 | void initLeader(); 90 | 91 | void updateURLs(const ::google::protobuf::RepeatedPtrField<::pdpb::Member> & members); 92 | 93 | void leaderLoop(); 94 | 95 | void switchLeader(const ::google::protobuf::RepeatedPtrField &); 96 | 97 | struct PDConnClient 98 | { 99 | std::shared_ptr channel; 100 | std::unique_ptr stub; 101 | std::unique_ptr keyspace_stub; 102 | std::unique_ptr resource_manager_stub; 103 | PDConnClient(std::string addr, const ClusterConfig & config) 104 | { 105 | if (config.hasTlsConfig()) 106 | { 107 | channel = grpc::CreateChannel(addr, grpc::SslCredentials(config.getGrpcCredentials())); 108 | } 109 | else 110 | { 111 | channel = grpc::CreateChannel(addr, grpc::InsecureChannelCredentials()); 112 | } 113 | stub = pdpb::PD::NewStub(channel); 114 | keyspace_stub = keyspacepb::Keyspace::NewStub(channel); 115 | resource_manager_stub = resource_manager::ResourceManager::NewStub(channel); 116 | } 117 | }; 118 | 119 | std::shared_ptr leaderClient(); 120 | 121 | pdpb::GetMembersResponse getMembers(const std::string &); 122 | 123 | pdpb::RequestHeader * requestHeader() const; 124 | 125 | std::shared_ptr getOrCreateGRPCConn(const std::string &); 126 | 127 | private: 128 | std::shared_mutex leader_mutex; 129 | 130 | std::mutex channel_map_mutex; 131 | 132 | std::mutex update_leader_mutex; 133 | 134 | std::unordered_map> channel_map; 135 | 136 | std::vector urls; 137 | std::unordered_set failed_urls; 138 | 139 | uint64_t cluster_id; 140 | 141 | std::string leader; 142 | 143 | std::atomic work_threads_stop; 144 | 145 | std::thread work_thread; 146 | 147 | std::condition_variable update_leader_cv; 148 | 149 | std::atomic check_leader; 150 | 151 | ClusterConfig config; 152 | 153 | Logger * log; 154 | }; 155 | 156 | 157 | } // namespace pingcap::pd 158 | -------------------------------------------------------------------------------- /src/common/MPPProber.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | namespace pingcap 7 | { 8 | namespace common 9 | { 10 | 11 | bool MPPProber::isRecovery(const std::string & store_addr, const std::chrono::seconds & recovery_ttl) 12 | { 13 | FailedStoreMap copy_failed_stores; 14 | { 15 | std::lock_guard lock(store_lock); 16 | copy_failed_stores = failed_stores; 17 | } 18 | auto iter = copy_failed_stores.find(store_addr); 19 | if (iter == copy_failed_stores.end()) 20 | return true; 21 | 22 | { 23 | std::lock_guard lock(iter->second->state_lock); 24 | iter->second->last_lookup_timepoint = std::chrono::steady_clock::now(); 25 | return iter->second->recovery_timepoint != INVALID_TIME_POINT && getElapsed(iter->second->recovery_timepoint) > recovery_ttl; 26 | } 27 | } 28 | 29 | void MPPProber::add(const std::string & store_addr) 30 | { 31 | std::lock_guard lock(store_lock); 32 | auto iter = failed_stores.find(store_addr); 33 | if (iter == failed_stores.end()) 34 | { 35 | auto state = std::make_shared(store_addr, cluster); 36 | state->last_lookup_timepoint = std::chrono::steady_clock::now(); 37 | failed_stores[store_addr] = state; 38 | } 39 | else 40 | { 41 | iter->second->last_lookup_timepoint = std::chrono::steady_clock::now(); 42 | } 43 | } 44 | 45 | void MPPProber::run() 46 | { 47 | while (!stopped.load()) 48 | { 49 | { 50 | std::unique_lock lock(scan_mu); 51 | scan_cv.wait_for(lock, std::chrono::seconds(scan_interval), [this] () { 52 | return stopped.load(); 53 | }); 54 | } 55 | 56 | try 57 | { 58 | scan(); 59 | } 60 | catch (...) 61 | { 62 | log->warning(getCurrentExceptionMsg("MPPProber scan failed: ")); 63 | } 64 | } 65 | } 66 | 67 | void MPPProber::stop() 68 | { 69 | stopped.store(true); 70 | std::lock_guard lock(scan_mu); 71 | scan_cv.notify_all(); 72 | } 73 | 74 | void MPPProber::scan() 75 | { 76 | FailedStoreMap copy_failed_stores; 77 | { 78 | std::lock_guard guard(store_lock); 79 | copy_failed_stores = failed_stores; 80 | } 81 | std::vector recovery_stores; 82 | recovery_stores.reserve(copy_failed_stores.size()); 83 | 84 | for (const auto & ele : copy_failed_stores) 85 | { 86 | if (!ele.second->state_lock.try_lock()) 87 | continue; 88 | 89 | ele.second->detectAndUpdateState(detect_period, detect_rpc_timeout); 90 | 91 | const auto & recovery_timepoint = ele.second->recovery_timepoint; 92 | if (recovery_timepoint != INVALID_TIME_POINT) 93 | { 94 | // Means this store is now alive. 95 | if (getElapsed(recovery_timepoint) > std::chrono::duration_cast(MAX_RECOVERY_TIME_LIMIT)) 96 | { 97 | recovery_stores.push_back(ele.first); 98 | } 99 | } 100 | else 101 | { 102 | // Store is dead, we want to check if this store has not used for MAX_OBSOLETE_TIME. 103 | if (ele.second->last_lookup_timepoint != INVALID_TIME_POINT && getElapsed(ele.second->last_lookup_timepoint) > std::chrono::duration_cast(MAX_OBSOLETE_TIME_LIMIT)) 104 | recovery_stores.push_back(ele.first); 105 | } 106 | ele.second->state_lock.unlock(); 107 | } 108 | 109 | { 110 | std::lock_guard guard(store_lock); 111 | for (const auto & store : recovery_stores) 112 | { 113 | failed_stores.erase(store); 114 | } 115 | } 116 | } 117 | 118 | void ProbeState::detectAndUpdateState(const std::chrono::seconds & detect_period, size_t detect_rpc_timeout) 119 | { 120 | if (last_detect_timepoint != INVALID_TIME_POINT && getElapsed(last_detect_timepoint) < detect_period) 121 | return; 122 | 123 | last_detect_timepoint = std::chrono::steady_clock::now(); 124 | bool is_alive = detectStore(cluster->rpc_client, store_addr, detect_rpc_timeout, log); 125 | if (!is_alive) 126 | { 127 | log->debug("got dead store: " + store_addr); 128 | recovery_timepoint = INVALID_TIME_POINT; 129 | } 130 | else if (recovery_timepoint == INVALID_TIME_POINT) 131 | { 132 | // Alive store, and first recovery, set its recovery time. 133 | recovery_timepoint = std::chrono::steady_clock::now(); 134 | } 135 | } 136 | 137 | bool detectStore(kv::RpcClientPtr & rpc_client, const std::string & store_addr, int rpc_timeout, Logger * log) 138 | { 139 | kv::RpcCall rpc(rpc_client, store_addr); 140 | grpc::ClientContext context; 141 | rpc.setClientContext(context, rpc_timeout); 142 | ::mpp::IsAliveRequest req; 143 | ::mpp::IsAliveResponse resp; 144 | auto status = rpc.call(&context, req, &resp); 145 | if (!status.ok()) 146 | { 147 | log->warning("detect failed: " + store_addr + " error: " + rpc.errMsg(status, "")); 148 | return false; 149 | } 150 | 151 | return resp.available(); 152 | } 153 | 154 | } // namespace common 155 | } // namespace pingcap 156 | -------------------------------------------------------------------------------- /include/pingcap/kv/internal/type_traits.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | namespace pingcap 9 | { 10 | namespace kv 11 | { 12 | 13 | #define RPC_NAME(METHOD) RpcTrait##METHOD 14 | 15 | #define GRPC_TRAIT(METHOD) \ 16 | struct RPC_NAME(METHOD) \ 17 | { \ 18 | static const char * errMsg() \ 19 | { \ 20 | return #METHOD " Failed"; \ 21 | } \ 22 | template \ 23 | static auto call( \ 24 | std::shared_ptr client, \ 25 | Args &&... args) \ 26 | { \ 27 | return client->stub->METHOD(std::forward(args)...); \ 28 | } \ 29 | }; \ 30 | struct RPC_NAME(Async##METHOD) \ 31 | { \ 32 | static const char * err_msg() \ 33 | { \ 34 | return "Async" #METHOD " Failed"; \ 35 | } \ 36 | template \ 37 | static auto call( \ 38 | std::shared_ptr client, \ 39 | Args &&... args) \ 40 | { \ 41 | return client->stub->Async##METHOD(std::forward(args)...); \ 42 | } \ 43 | }; 44 | 45 | #define M(method) GRPC_TRAIT(method) 46 | // Commands using a transactional interface. 47 | M(KvGet) 48 | M(KvScan) 49 | M(KvPrewrite) 50 | M(KvPessimisticLock) 51 | M(KVPessimisticRollback) 52 | M(KvTxnHeartBeat) 53 | M(KvCheckTxnStatus) 54 | M(KvCheckSecondaryLocks) 55 | M(KvCommit) 56 | M(KvCleanup) 57 | M(KvBatchGet) 58 | M(KvBatchRollback) 59 | M(KvScanLock) 60 | M(KvResolveLock) 61 | M(KvGC) 62 | M(KvDeleteRange) 63 | M(KvPrepareFlashbackToVersion) 64 | M(KvFlashbackToVersion) 65 | 66 | // Raw commands; no transaction support. 67 | M(RawGet) 68 | M(RawBatchGet) 69 | M(RawPut) 70 | M(RawBatchPut) 71 | M(RawDelete) 72 | M(RawBatchDelete) 73 | M(RawScan) 74 | M(RawDeleteRange) 75 | M(RawBatchScan) 76 | // Get TTL of the key. Returns 0 if TTL is not set for the key. 77 | M(RawGetKeyTTL) 78 | // Compare if the value in database equals to `RawCASRequest.previous_value` before putting the new value. If not, this request will have no effect and the value in the database will be returned. 79 | M(RawCompareAndSwap) 80 | M(RawChecksum) 81 | 82 | // Store commands (sent to a each TiKV node in a cluster, rather than a certain region). 83 | M(UnsafeDestroyRange) 84 | M(RegisterLockObserver) 85 | M(CheckLockObserver) 86 | M(RemoveLockObserver) 87 | M(PhysicalScanLock) 88 | 89 | // Commands for executing SQL in the TiKV coprocessor (i.e., 'pushed down' to TiKV rather than 90 | // executed in TiDB). 91 | M(Coprocessor) 92 | M(CoprocessorStream) 93 | M(BatchCoprocessor) 94 | 95 | // Command for executing custom user requests in TiKV coprocessor_v2. 96 | M(RawCoprocessor) 97 | 98 | // Raft commands (sent between TiKV nodes). 99 | M(Raft) 100 | M(BatchRaft) 101 | M(Snapshot) 102 | M(TabletSnapshot) 103 | 104 | // Sent from PD or TiDB to a TiKV node. 105 | M(SplitRegion) 106 | // Sent from TiFlash or TiKV to a TiKV node. 107 | M(ReadIndex) 108 | 109 | // Commands for debugging transactions. 110 | M(MvccGetByKey) 111 | M(MvccGetByStartTs) 112 | 113 | // Batched commands. 114 | M(BatchCommands) 115 | 116 | // These are for mpp execution. 117 | M(DispatchMPPTask) 118 | M(CancelMPPTask) 119 | M(EstablishMPPConnection) 120 | M(IsAlive) 121 | M(ReportMPPTaskStatus) 122 | 123 | /// CheckLeader sends all information (includes region term and epoch) to other stores. 124 | /// Once a store receives a request, it checks term and epoch for each region, and sends the regions whose 125 | /// term and epoch match with local information in the store. 126 | /// After the client collected all responses from all stores, it checks if got a quorum of responses from 127 | /// other stores for every region, and decides to advance resolved ts from these regions. 128 | M(CheckLeader) 129 | 130 | /// Get the minimal `safe_ts` from regions at the store 131 | M(GetStoreSafeTS) 132 | 133 | /// Get the information about lock waiting from TiKV. 134 | M(GetLockWaitInfo) 135 | 136 | /// Compact a specified key range. This request is not restricted to raft leaders and will not be replicated. 137 | /// It only compacts data on this node. 138 | /// TODO: Currently this RPC is designed to be only compatible with TiFlash. 139 | /// Shall be move out in https://github.com/pingcap/kvproto/issues/912 140 | M(Compact) 141 | /// Get the information about history lock waiting from TiKV. 142 | M(GetLockWaitHistory) 143 | 144 | /// Get system table from TiFlash 145 | M(GetTiFlashSystemTable) 146 | 147 | // These are for TiFlash disaggregated architecture 148 | /// Try to lock a S3 object, atomically 149 | M(tryAddLock) 150 | /// Try to delete a S3 object, atomically 151 | M(tryMarkDelete) 152 | /// Build the disaggregated task on TiFlash write node 153 | M(EstablishDisaggTask) 154 | /// Cancel the disaggregated task on TiFlash write node 155 | M(CancelDisaggTask) 156 | /// Exchange page data between TiFlash write node and compute node 157 | M(FetchDisaggPages) 158 | /// Compute node get configuration from Write node 159 | M(GetDisaggConfig) 160 | 161 | #undef M 162 | 163 | } // namespace kv 164 | } // namespace pingcap 165 | -------------------------------------------------------------------------------- /cmake/Modules/FindPoco.cmake: -------------------------------------------------------------------------------- 1 | # - finds the Poco C++ libraries 2 | # This module finds the Applied Informatics Poco libraries. 3 | # It supports the following components: 4 | # 5 | # Util (loaded by default) 6 | # Foundation (loaded by default) 7 | # JSON 8 | # XML 9 | # Zip 10 | # Crypto 11 | # Data 12 | # Net 13 | # NetSSL 14 | # OSP 15 | # 16 | # Usage: 17 | # set(ENV{Poco_DIR} path/to/poco/sdk) 18 | # find_package(Poco REQUIRED OSP Data Crypto) 19 | # 20 | # On completion, the script defines the following variables: 21 | # 22 | # - Compound variables: 23 | # Poco_FOUND 24 | # - true if all requested components were found. 25 | # Poco_LIBRARIES 26 | # - contains release (and debug if available) libraries for all requested components. 27 | # It has the form "optimized LIB1 debug LIBd1 optimized LIB2 ...", ready for use with the target_link_libraries command. 28 | # Poco_INCLUDE_DIRS 29 | # - Contains include directories for all requested components. 30 | # 31 | # - Component variables: 32 | # Poco_Xxx_FOUND 33 | # - Where Xxx is the properly cased component name (eg. 'Util', 'OSP'). 34 | # True if a component's library or debug library was found successfully. 35 | # Poco_Xxx_LIBRARY 36 | # - Library for component Xxx. 37 | # Poco_Xxx_LIBRARY_DEBUG 38 | # - debug library for component Xxx 39 | # Poco_Xxx_INCLUDE_DIR 40 | # - include directory for component Xxx 41 | # 42 | # - OSP BundleCreator variables: (i.e. bundle.exe on windows, bundle on unix-likes) 43 | # (is only discovered if OSP is a requested component) 44 | # Poco_OSP_Bundle_EXECUTABLE_FOUND 45 | # - true if the bundle-creator executable was found. 46 | # Poco_OSP_Bundle_EXECUTABLE 47 | # - the path to the bundle-creator executable. 48 | # 49 | # Author: Andreas Stahl andreas.stahl@tu-dresden.de 50 | 51 | set(Poco_HINTS 52 | /usr/local 53 | C:/AppliedInformatics 54 | ${Poco_DIR} 55 | $ENV{Poco_DIR} 56 | ) 57 | 58 | if(NOT Poco_ROOT_DIR) 59 | # look for the root directory, first for the source-tree variant 60 | find_path(Poco_ROOT_DIR 61 | NAMES Foundation/include/Poco/Poco.h 62 | HINTS ${Poco_HINTS} 63 | ) 64 | if(NOT Poco_ROOT_DIR) 65 | # this means poco may have a different directory structure, maybe it was installed, let's check for that 66 | message(STATUS "Looking for Poco install directory structure.") 67 | find_path(Poco_ROOT_DIR 68 | NAMES include/Poco/Poco.h 69 | HINTS ${Poco_HINTS} 70 | ) 71 | if(NOT Poco_ROOT_DIR) 72 | # poco was still not found -> Fail 73 | if(Poco_FIND_REQUIRED) 74 | message(FATAL_ERROR "Poco: Could not find Poco install directory") 75 | endif() 76 | if(NOT Poco_FIND_QUIETLY) 77 | message(STATUS "Poco: Could not find Poco install directory") 78 | endif() 79 | return() 80 | else() 81 | # poco was found with the make install directory structure 82 | message(STATUS "Assuming Poco install directory structure at ${Poco_ROOT_DIR}.") 83 | set(Poco_INSTALLED true) 84 | endif() 85 | endif() 86 | endif() 87 | 88 | # add dynamic library directory 89 | if(WIN32) 90 | find_path(Poco_RUNTIME_LIBRARY_DIRS 91 | NAMES PocoFoundation.dll 92 | HINTS ${Poco_ROOT_DIR} 93 | PATH_SUFFIXES 94 | bin 95 | lib 96 | ) 97 | endif() 98 | 99 | # if installed directory structure, set full include dir 100 | if(Poco_INSTALLED) 101 | set(Poco_INCLUDE_DIRS ${Poco_ROOT_DIR}/include/ CACHE PATH "The global include path for Poco") 102 | endif() 103 | 104 | # append the default minimum components to the list to find 105 | list(APPEND components 106 | ${Poco_FIND_COMPONENTS} 107 | # default components: 108 | "Util" 109 | "Foundation" 110 | ) 111 | list(REMOVE_DUPLICATES components) # remove duplicate defaults 112 | 113 | foreach( component ${components} ) 114 | #if(NOT Poco_${component}_FOUND) 115 | 116 | # include directory for the component 117 | if(NOT Poco_${component}_INCLUDE_DIR) 118 | if(component STREQUAL "NetSSL") 119 | set(component_INCLUDE_NAMES 120 | "Poco/Net/NetSSL.h" 121 | ) 122 | else() 123 | set(component_INCLUDE_NAMES 124 | "Poco/${component}.h" # e.g. Foundation.h 125 | "Poco/${component}/${component}.h" # e.g. OSP/OSP.h Util/Util.h 126 | ) 127 | endif() 128 | find_path(Poco_${component}_INCLUDE_DIR 129 | NAMES 130 | ${component_INCLUDE_NAMES} 131 | HINTS 132 | ${Poco_ROOT_DIR} 133 | PATH_SUFFIXES 134 | include 135 | ${component}/include 136 | ) 137 | endif() 138 | if(NOT Poco_${component}_INCLUDE_DIR) 139 | message(FATAL_ERROR "Poco_${component}_INCLUDE_DIR NOT FOUND") 140 | else() 141 | list(APPEND Poco_INCLUDE_DIRS ${Poco_${component}_INCLUDE_DIR}) 142 | endif() 143 | 144 | # release library 145 | if(NOT Poco_${component}_LIBRARY) 146 | find_library( 147 | Poco_${component}_LIBRARY 148 | NAMES Poco${component} 149 | HINTS ${Poco_ROOT_DIR} 150 | PATH_SUFFIXES 151 | lib 152 | bin 153 | ) 154 | if(Poco_${component}_LIBRARY) 155 | message(STATUS "Found Poco ${component}: ${Poco_${component}_LIBRARY}") 156 | endif() 157 | endif() 158 | if(Poco_${component}_LIBRARY) 159 | list(APPEND Poco_LIBRARIES "optimized" ${Poco_${component}_LIBRARY} ) 160 | mark_as_advanced(Poco_${component}_LIBRARY) 161 | endif() 162 | 163 | # debug library 164 | if(NOT Poco_${component}_LIBRARY_DEBUG) 165 | find_library( 166 | Poco_${component}_LIBRARY_DEBUG 167 | Names Poco${component}d 168 | HINTS ${Poco_ROOT_DIR} 169 | PATH_SUFFIXES 170 | lib 171 | bin 172 | ) 173 | if(Poco_${component}_LIBRARY_DEBUG) 174 | message(STATUS "Found Poco ${component} (debug): ${Poco_${component}_LIBRARY_DEBUG}") 175 | endif() 176 | endif(NOT Poco_${component}_LIBRARY_DEBUG) 177 | if(Poco_${component}_LIBRARY_DEBUG) 178 | list(APPEND Poco_LIBRARIES "debug" ${Poco_${component}_LIBRARY_DEBUG}) 179 | mark_as_advanced(Poco_${component}_LIBRARY_DEBUG) 180 | endif() 181 | 182 | # mark component as found or handle not finding it 183 | if(Poco_${component}_LIBRARY_DEBUG OR Poco_${component}_LIBRARY) 184 | set(Poco_${component}_FOUND TRUE) 185 | elseif(NOT Poco_FIND_QUIETLY) 186 | message(FATAL_ERROR "Could not find Poco component ${component}!") 187 | endif() 188 | endforeach() 189 | 190 | if(DEFINED Poco_LIBRARIES) 191 | set(Poco_FOUND true) 192 | endif() 193 | 194 | if(${Poco_OSP_FOUND}) 195 | # find the osp bundle program 196 | find_program( 197 | Poco_OSP_Bundle_EXECUTABLE 198 | NAMES bundle 199 | HINTS 200 | ${Poco_RUNTIME_LIBRARY_DIRS} 201 | ${Poco_ROOT_DIR} 202 | PATH_SUFFIXES 203 | bin 204 | OSP/BundleCreator/bin/Darwin/x86_64 205 | OSP/BundleCreator/bin/Darwin/i386 206 | DOC "The executable that bundles OSP packages according to a .bndlspec specification." 207 | ) 208 | if(Poco_OSP_Bundle_EXECUTABLE) 209 | set(Poco_OSP_Bundle_EXECUTABLE_FOUND true) 210 | endif() 211 | # include bundle script file 212 | find_file(Poco_OSP_Bundles_file NAMES PocoBundles.cmake HINTS ${CMAKE_MODULE_PATH}) 213 | if(${Poco_OSP_Bundles_file}) 214 | include(${Poco_OSP_Bundles_file}) 215 | endif() 216 | endif() 217 | 218 | message(STATUS "Found Poco: ${Poco_LIBRARIES}") 219 | 220 | 221 | -------------------------------------------------------------------------------- /include/pingcap/kv/2pc.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | 9 | #include 10 | #include 11 | #include 12 | #include 13 | #include 14 | #include 15 | 16 | namespace pingcap 17 | { 18 | namespace kv 19 | { 20 | constexpr uint32_t txnCommitBatchSize = 16 * 1024; 21 | 22 | struct Txn; 23 | 24 | struct TwoPhaseCommitter; 25 | 26 | using TwoPhaseCommitterPtr = std::shared_ptr; 27 | 28 | uint64_t sendTxnHeartBeat(Backoffer & bo, Cluster * cluster, std::string & primary_key, uint64_t start_ts, uint64_t ttl); 29 | 30 | uint64_t txnLockTTL(std::chrono::milliseconds start, uint64_t txn_size); 31 | 32 | class TTLManager 33 | { 34 | private: 35 | enum TTLManagerState 36 | { 37 | StateUninitialized = 0, 38 | StateRunning, 39 | StateClosed 40 | }; 41 | 42 | std::atomic state; 43 | 44 | bool worker_running; 45 | std::thread * worker; 46 | 47 | public: 48 | TTLManager() 49 | : state{StateUninitialized} 50 | , worker_running{false} 51 | , worker{nullptr} 52 | {} 53 | 54 | void run(TwoPhaseCommitterPtr committer) 55 | { 56 | // Run only once and start a background thread to refresh lock ttl 57 | uint32_t expected = StateUninitialized; 58 | if (!state.compare_exchange_strong(expected, StateRunning, std::memory_order_acquire, std::memory_order_relaxed)) 59 | { 60 | return; 61 | } 62 | 63 | worker_running = true; 64 | worker = new std::thread{&TTLManager::keepAlive, this, committer}; 65 | } 66 | 67 | void close() 68 | { 69 | uint32_t expected = StateRunning; 70 | state.compare_exchange_strong(expected, StateClosed, std::memory_order_acq_rel); 71 | if (worker_running && worker->joinable()) 72 | { 73 | worker->join(); 74 | worker_running = false; 75 | delete worker; 76 | } 77 | } 78 | 79 | void keepAlive(TwoPhaseCommitterPtr committer); 80 | }; 81 | 82 | struct TwoPhaseCommitter : public std::enable_shared_from_this 83 | { 84 | private: 85 | std::unordered_map mutations; 86 | 87 | std::vector keys; 88 | uint64_t start_ts = 0; 89 | 90 | std::shared_mutex commit_ts_mu; 91 | uint64_t commit_ts = 0; 92 | uint64_t min_commit_ts = 0; 93 | uint64_t max_commit_ts = 0; 94 | 95 | // Used to calculate max_commit_ts 96 | std::chrono::milliseconds start_time; 97 | 98 | Cluster * cluster; 99 | 100 | std::unordered_map region_txn_size; 101 | uint64_t txn_size = 0; 102 | 103 | int lock_ttl = 0; 104 | 105 | std::string primary_lock; 106 | // commited means primary key has been written to kv stores. 107 | bool commited; 108 | 109 | // Only for test now 110 | bool use_async_commit; 111 | 112 | TTLManager ttl_manager; 113 | 114 | Logger * log; 115 | 116 | friend class TTLManager; 117 | 118 | friend class TestTwoPhaseCommitter; 119 | 120 | public: 121 | explicit TwoPhaseCommitter(Txn * txn, bool _use_async_commit = false); 122 | 123 | void execute(); 124 | 125 | private: 126 | enum Action 127 | { 128 | ActionPrewrite = 0, 129 | ActionCommit, 130 | ActionCleanUp 131 | }; 132 | 133 | struct BatchKeys 134 | { 135 | RegionVerID region; 136 | std::vector keys; 137 | bool is_primary; 138 | BatchKeys(const RegionVerID & region_, std::vector keys_, bool is_primary_ = false) 139 | : region(region_) 140 | , keys(std::move(keys_)) 141 | , is_primary(is_primary_) 142 | {} 143 | }; 144 | 145 | void calculateMaxCommitTS(); 146 | 147 | void prewriteKeys(Backoffer & bo, const std::vector & keys) { doActionOnKeys(bo, keys); } 148 | 149 | void commitKeys(Backoffer & bo, const std::vector & keys) { doActionOnKeys(bo, keys); } 150 | 151 | template 152 | void doActionOnKeys(Backoffer & bo, const std::vector & cur_keys) 153 | { 154 | auto [groups, first_region] = cluster->region_cache->groupKeysByRegion(bo, cur_keys); 155 | std::ignore = first_region; 156 | 157 | // TODO: presplit region when needed 158 | std::vector batches; 159 | uint64_t primary_idx = std::numeric_limits::max(); 160 | for (auto & group : groups) 161 | { 162 | uint32_t end = 0; 163 | for (uint32_t start = 0; start < group.second.size(); start = end) 164 | { 165 | uint64_t size = 0; 166 | std::vector sub_keys; 167 | for (end = start; end < group.second.size() && size < txnCommitBatchSize; end++) 168 | { 169 | auto & key = group.second[end]; 170 | size += key.size(); 171 | if constexpr (action == ActionPrewrite) 172 | size += mutations[key].size(); 173 | 174 | if (key == primary_lock) 175 | primary_idx = batches.size(); 176 | sub_keys.push_back(key); 177 | } 178 | batches.emplace_back(BatchKeys(group.first, sub_keys)); 179 | } 180 | } 181 | if (primary_idx != std::numeric_limits::max() && primary_idx != 0) 182 | { 183 | std::swap(batches[0], batches[primary_idx]); 184 | batches[0].is_primary = true; 185 | } 186 | 187 | if constexpr (action == ActionCommit || action == ActionCleanUp) 188 | { 189 | if constexpr (action == ActionCommit) 190 | { 191 | fiu_do_on("all commit fail", return); 192 | } 193 | doActionOnBatches(bo, std::vector(batches.begin(), batches.begin() + 1)); 194 | batches = std::vector(batches.begin() + 1, batches.end()); 195 | } 196 | if (action != ActionCommit || !fiu_fail("rest commit fail")) 197 | { 198 | doActionOnBatches(bo, batches); 199 | } 200 | } 201 | 202 | template 203 | void doActionOnBatches(Backoffer & bo, const std::vector & batches) 204 | { 205 | for (const auto & batch : batches) 206 | { 207 | if constexpr (action == ActionPrewrite) 208 | { 209 | region_txn_size[batch.region.id] = batch.keys.size(); 210 | prewriteSingleBatch(bo, batch); 211 | } 212 | else if constexpr (action == ActionCommit) 213 | { 214 | commitSingleBatch(bo, batch); 215 | } 216 | } 217 | } 218 | 219 | void prewriteSingleBatch(Backoffer & bo, const BatchKeys & batch); 220 | 221 | void commitSingleBatch(Backoffer & bo, const BatchKeys & batch); 222 | }; 223 | 224 | } // namespace kv 225 | } // namespace pingcap 226 | -------------------------------------------------------------------------------- /src/test/bank_test/bank_test.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | 10 | namespace 11 | { 12 | using namespace pingcap; 13 | using namespace pingcap::kv; 14 | 15 | struct BankCase 16 | { 17 | Cluster * cluster; 18 | int batch_size = 100; 19 | std::atomic_bool stop; 20 | int account_cnt; 21 | std::atomic_int start; 22 | int concurrency; 23 | 24 | std::thread check_thread; 25 | 26 | BankCase(Cluster * cluster_, int account_cnt_, int con_) 27 | : cluster(cluster_) 28 | , batch_size(100) 29 | , stop(false) 30 | , account_cnt(account_cnt_) 31 | , start(account_cnt_ / batch_size) 32 | , concurrency(con_) 33 | {} 34 | 35 | void close() 36 | { 37 | stop = true; 38 | check_thread.join(); 39 | } 40 | 41 | void enable_check() 42 | { 43 | check_thread = std::thread([&]() { 44 | try 45 | { 46 | verify(); 47 | } 48 | catch (Exception & e) 49 | { 50 | std::cerr << e.displayText() << std::endl; 51 | } 52 | }); 53 | } 54 | 55 | void initialize() 56 | { 57 | std::cerr << "bank case start to init\n"; 58 | std::vector threads; 59 | threads.reserve(concurrency); 60 | for (int i = 0; i < concurrency; i++) 61 | { 62 | threads.push_back(std::thread([&]() { 63 | try 64 | { 65 | initAccount(); 66 | } 67 | catch (Exception & e) 68 | { 69 | std::cerr << e.displayText() << std::endl; 70 | } 71 | })); 72 | } 73 | for (int i = 0; i < concurrency; i++) 74 | { 75 | threads[i].join(); 76 | } 77 | enable_check(); 78 | std::cerr << "bank case end init\n"; 79 | } 80 | 81 | void verify() 82 | { 83 | for (;;) 84 | { 85 | if (stop) 86 | { 87 | std::cerr << "end check\n"; 88 | return; 89 | } 90 | int total = 0; 91 | Snapshot snapshot(cluster); 92 | std::string prefix = "bankkey_"; 93 | auto scanner = snapshot.Scan(prefix, prefixNext(prefix)); 94 | int cnt = 0; 95 | std::map key_count; 96 | while (scanner.valid) 97 | { 98 | auto key = scanner.key(); 99 | auto key_index = get_bank_key_index(key); 100 | 101 | key_count[key_index]++; 102 | 103 | auto value = scanner.value(); 104 | try 105 | { 106 | total += std::stoi(value); 107 | } 108 | catch (std::exception & e) 109 | { 110 | std::cerr << "Invalid value: " << value << " , key: " << key << ", version: " << snapshot.version << std::endl; 111 | } 112 | scanner.next(); 113 | cnt++; 114 | } 115 | 116 | if (account_cnt != cnt) 117 | { 118 | std::cerr << "read ts: " << snapshot.version << std::endl; 119 | for (int i = 0; i < account_cnt; i++) 120 | if (key_count[i] != 1) 121 | { 122 | std::cerr << "key idx: " << i << " " << key_count[i] << std::endl; 123 | } 124 | } 125 | 126 | std::cerr << "total: " << total << " account " << account_cnt << " cnt " << cnt << " version " << snapshot.version << std::endl; 127 | assert(total == account_cnt * 1000); 128 | std::this_thread::sleep_for(std::chrono::seconds(2)); 129 | } 130 | } 131 | 132 | void execute() 133 | { 134 | std::cerr << "bank case start to execute\n"; 135 | std::vector threads; 136 | threads.reserve(concurrency); 137 | for (int i = 0; i < concurrency; i++) 138 | { 139 | threads.push_back(std::thread([&]() { 140 | try 141 | { 142 | moveMoney(); 143 | } 144 | catch (Exception & e) 145 | { 146 | std::cerr << e.displayText() << std::endl; 147 | } 148 | })); 149 | } 150 | for (int i = 0; i < concurrency; i++) 151 | { 152 | threads[i].join(); 153 | } 154 | std::cerr << "bank case end execute\n"; 155 | } 156 | 157 | void moveMoneyOnce(std::mt19937 & generator) 158 | { 159 | int from, to; 160 | for (;;) 161 | { 162 | from = generator() % account_cnt; 163 | to = generator() % account_cnt; 164 | if (to != from) 165 | break; 166 | } 167 | Txn txn(cluster); 168 | 169 | std::string rest; 170 | bool exists; 171 | std::tie(rest, exists) = txn.get(bank_key(from)); 172 | assert(exists); 173 | if (rest.empty()) 174 | return; 175 | int rest_money = std::stoi(rest); 176 | int money = generator() % rest_money; 177 | 178 | txn.set(bank_key(from), bank_value(rest_money - money)); 179 | 180 | std::tie(rest, exists) = txn.get(bank_key(to)); 181 | assert(exists); 182 | if (rest.empty()) 183 | return; 184 | 185 | rest_money = std::stoi(rest); 186 | txn.set(bank_key(to), bank_value(rest_money + money)); 187 | txn.commit(); 188 | } 189 | 190 | void moveMoney() 191 | { 192 | static thread_local std::mt19937 generator; 193 | for (;;) 194 | { 195 | if (stop) 196 | return; 197 | try 198 | { 199 | moveMoneyOnce(generator); 200 | } 201 | catch (Exception & e) 202 | { 203 | std::cerr << "move money failed: " << e.displayText() << std::endl; 204 | } 205 | } 206 | } 207 | 208 | int get_bank_key_index(std::string key) { return std::stoi(key.substr(key.find("_") + 1)); } 209 | std::string bank_key(int idx) 210 | { 211 | char buff[12] = ""; 212 | assert(idx < 10000 && idx >= 0); 213 | std::snprintf(buff, sizeof(buff), "%04d", idx); 214 | return "bankkey_" + std::string(buff); 215 | } 216 | std::string bank_value(int money) { return std::to_string(money); } 217 | 218 | void initAccount() 219 | { 220 | for (;;) 221 | { 222 | if (stop) 223 | return; 224 | int start_idx = start.fetch_sub(1) - 1; 225 | if (start_idx < 0) 226 | return; 227 | Txn txn(cluster); 228 | for (int i = start_idx * batch_size; i < start_idx * batch_size + batch_size; i++) 229 | { 230 | txn.set(bank_key(i), bank_value(1000)); 231 | } 232 | txn.commit(); 233 | } 234 | } 235 | }; 236 | 237 | } // namespace 238 | -------------------------------------------------------------------------------- /.clang-tidy: -------------------------------------------------------------------------------- 1 | Checks: '-*, 2 | misc-throw-by-value-catch-by-reference, 3 | misc-misplaced-const, 4 | misc-unconventional-assign-operator, 5 | misc-redundant-expression, 6 | misc-static-assert, 7 | misc-unconventional-assign-operator, 8 | misc-uniqueptr-reset-release, 9 | misc-unused-alias-decls, 10 | misc-unused-parameters, 11 | misc-unused-using-decls, 12 | 13 | modernize-avoid-bind, 14 | modernize-loop-convert, 15 | modernize-make-shared, 16 | modernize-make-unique, 17 | modernize-raw-string-literal, 18 | modernize-redundant-void-arg, 19 | modernize-replace-auto-ptr, 20 | modernize-replace-random-shuffle, 21 | modernize-use-auto, 22 | modernize-use-bool-literals, 23 | modernize-use-nullptr, 24 | modernize-use-using, 25 | modernize-use-override, 26 | modernize-use-equals-default, 27 | modernize-use-equals-delete, 28 | 29 | performance-faster-string-find, 30 | performance-for-range-copy, 31 | performance-implicit-conversion-in-loop, 32 | performance-inefficient-algorithm, 33 | performance-inefficient-vector-operation, 34 | performance-move-constructor-init, 35 | performance-no-automatic-move, 36 | performance-trivially-destructible, 37 | performance-unnecessary-copy-initialization, 38 | 39 | readability-avoid-const-params-in-decls, 40 | readability-const-return-type, 41 | readability-container-size-empty, 42 | readability-convert-member-functions-to-static, 43 | readability-delete-null-pointer, 44 | readability-deleted-default, 45 | readability-make-member-function-const, 46 | readability-misplaced-array-index, 47 | readability-non-const-parameter, 48 | readability-qualified-auto, 49 | readability-redundant-control-flow, 50 | readability-redundant-function-ptr-dereference, 51 | readability-redundant-smartptr-get, 52 | readability-redundant-string-cstr, 53 | readability-redundant-string-init, 54 | readability-static-definition-in-anonymous-namespace, 55 | readability-string-compare, 56 | readability-uniqueptr-delete-release, 57 | readability-redundant-member-init, 58 | readability-simplify-subscript-expr, 59 | readability-simplify-boolean-expr, 60 | readability-inconsistent-declaration-parameter-name, 61 | readability-identifier-naming, 62 | 63 | bugprone-undelegated-constructor, 64 | bugprone-argument-comment, 65 | bugprone-bad-signal-to-kill-thread, 66 | bugprone-bool-pointer-implicit-conversion, 67 | bugprone-copy-constructor-init, 68 | bugprone-dangling-handle, 69 | bugprone-forward-declaration-namespace, 70 | bugprone-fold-init-type, 71 | bugprone-inaccurate-erase, 72 | bugprone-incorrect-roundings, 73 | bugprone-infinite-loop, 74 | bugprone-integer-division, 75 | bugprone-macro-parentheses, 76 | bugprone-macro-repeated-side-effects, 77 | bugprone-misplaced-operator-in-strlen-in-alloc, 78 | bugprone-misplaced-pointer-artithmetic-in-alloc, 79 | bugprone-misplaced-widening-cast, 80 | bugprone-move-forwarding-reference, 81 | bugprone-multiple-statement-macro, 82 | bugprone-parent-virtual-call, 83 | bugprone-posix-return, 84 | bugprone-reserved-identifier, 85 | bugprone-signed-char-misuse, 86 | bugprone-sizeof-container, 87 | bugprone-sizeof-expression, 88 | bugprone-string-constructor, 89 | bugprone-string-integer-assignment, 90 | bugprone-string-literal-with-embedded-nul, 91 | bugprone-suspicious-enum-usage, 92 | bugprone-suspicious-include, 93 | bugprone-suspicious-memset-usage, 94 | bugprone-suspicious-missing-comma, 95 | bugprone-suspicious-string-compare, 96 | bugprone-swapped-arguments, 97 | bugprone-terminating-continue, 98 | bugprone-throw-keyword-missing, 99 | bugprone-too-small-loop-variable, 100 | bugprone-undefined-memory-manipulation, 101 | bugprone-unhandled-self-assignment, 102 | bugprone-unused-raii, 103 | bugprone-unused-return-value, 104 | bugprone-use-after-move, 105 | bugprone-virtual-near-miss, 106 | 107 | cert-dcl21-cpp, 108 | cert-dcl50-cpp, 109 | cert-env33-c, 110 | cert-err34-c, 111 | cert-err52-cpp, 112 | cert-flp30-c, 113 | cert-mem57-cpp, 114 | cert-msc50-cpp, 115 | cert-oop58-cpp, 116 | 117 | google-build-explicit-make-pair, 118 | google-build-namespaces, 119 | google-default-arguments, 120 | google-explicit-constructor, 121 | google-readability-casting, 122 | google-readability-avoid-underscore-in-googletest-name, 123 | google-runtime-int, 124 | google-runtime-operator, 125 | 126 | hicpp-exception-baseclass, 127 | 128 | clang-analyzer-core.CallAndMessage, 129 | clang-analyzer-core.DivideZero, 130 | clang-analyzer-core.NonNullParamChecker, 131 | clang-analyzer-core.NullDereference, 132 | clang-analyzer-core.StackAddressEscape, 133 | clang-analyzer-core.UndefinedBinaryOperatorResult, 134 | clang-analyzer-core.VLASize, 135 | clang-analyzer-core.uninitialized.ArraySubscript, 136 | clang-analyzer-core.uninitialized.Assign, 137 | clang-analyzer-core.uninitialized.Branch, 138 | clang-analyzer-core.uninitialized.CapturedBlockVariable, 139 | clang-analyzer-core.uninitialized.UndefReturn, 140 | clang-analyzer-cplusplus.InnerPointer, 141 | clang-analyzer-cplusplus.NewDelete, 142 | clang-analyzer-cplusplus.NewDeleteLeaks, 143 | clang-analyzer-cplusplus.PlacementNewChecker, 144 | clang-analyzer-cplusplus.SelfAssignment, 145 | clang-analyzer-deadcode.DeadStores, 146 | clang-analyzer-optin.cplusplus.VirtualCall, 147 | clang-analyzer-security.insecureAPI.UncheckedReturn, 148 | clang-analyzer-security.insecureAPI.bcmp, 149 | clang-analyzer-security.insecureAPI.bcopy, 150 | clang-analyzer-security.insecureAPI.bzero, 151 | clang-analyzer-security.insecureAPI.getpw, 152 | clang-analyzer-security.insecureAPI.gets, 153 | clang-analyzer-security.insecureAPI.mkstemp, 154 | clang-analyzer-security.insecureAPI.mktemp, 155 | clang-analyzer-security.insecureAPI.rand, 156 | clang-analyzer-security.insecureAPI.strcpy, 157 | clang-analyzer-unix.Malloc, 158 | clang-analyzer-unix.MallocSizeof, 159 | clang-analyzer-unix.MismatchedDeallocator, 160 | clang-analyzer-unix.Vfork, 161 | clang-analyzer-unix.cstring.BadSizeArg, 162 | clang-analyzer-unix.cstring.NullArg, 163 | 164 | boost-use-to-string, 165 | 166 | cppcoreguidelines-pro-type-cstyle-cast, 167 | cppcoreguidelines-pro-type-member-init, 168 | cppcoreguidelines-no-malloc, 169 | cppcoreguidelines-virtual-class-destructor, 170 | ' 171 | WarningsAsErrors: '*' 172 | 173 | CheckOptions: 174 | - key: readability-identifier-naming.ClassCase 175 | value: CamelCase 176 | - key: readability-identifier-naming.EnumCase 177 | value: CamelCase 178 | - key: readability-identifier-naming.LocalVariableCase 179 | value: lower_case 180 | - key: readability-identifier-naming.StaticConstantCase 181 | value: aNy_CasE 182 | - key: readability-identifier-naming.MemberCase 183 | value: lower_case 184 | - key: readability-identifier-naming.PrivateMemberPrefix 185 | value: '' 186 | - key: readability-identifier-naming.ProtectedMemberPrefix 187 | value: '' 188 | - key: readability-identifier-naming.PublicMemberCase 189 | value: lower_case 190 | - key: readability-identifier-naming.MethodCase 191 | value: camelBack 192 | - key: readability-identifier-naming.PrivateMethodPrefix 193 | value: '' 194 | - key: readability-identifier-naming.ProtectedMethodPrefix 195 | value: '' 196 | - key: readability-identifier-naming.ParameterPackCase 197 | value: lower_case 198 | - key: readability-identifier-naming.StructCase 199 | value: CamelCase 200 | - key: readability-identifier-naming.TemplateTemplateParameterCase 201 | value: CamelCase 202 | - key: readability-identifier-naming.TemplateUsingCase 203 | value: lower_case 204 | - key: readability-identifier-naming.TypeTemplateParameterCase 205 | value: CamelCase 206 | - key: readability-identifier-naming.TypedefCase 207 | value: CamelCase 208 | - key: readability-identifier-naming.UnionCase 209 | value: CamelCase 210 | - key: readability-identifier-naming.UsingCase 211 | value: CamelCase 212 | -------------------------------------------------------------------------------- /include/pingcap/kv/RegionCache.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | #include 11 | #include 12 | 13 | namespace pingcap 14 | { 15 | namespace kv 16 | { 17 | enum class StoreType 18 | { 19 | TiKV, 20 | TiFlash, 21 | }; 22 | 23 | struct Store 24 | { 25 | uint64_t id; 26 | std::string addr; 27 | std::string peer_addr; 28 | std::map labels; 29 | StoreType store_type; 30 | ::metapb::StoreState state; 31 | 32 | Store(uint64_t id_, const std::string & addr_, const std::string & peer_addr_, const std::map & labels_, StoreType store_type_, const ::metapb::StoreState state_) 33 | : id(id_) 34 | , addr(addr_) 35 | , peer_addr(peer_addr_) 36 | , labels(labels_) 37 | , store_type(store_type_) 38 | , state(state_) 39 | {} 40 | }; 41 | 42 | struct RegionVerID 43 | { 44 | uint64_t id; 45 | uint64_t conf_ver; 46 | uint64_t ver; 47 | 48 | RegionVerID() 49 | : RegionVerID(0, 0, 0) 50 | {} 51 | RegionVerID(uint64_t id_, uint64_t conf_ver_, uint64_t ver_) 52 | : id(id_) 53 | , conf_ver(conf_ver_) 54 | , ver(ver_) 55 | {} 56 | 57 | bool operator==(const RegionVerID & rhs) const { return id == rhs.id && conf_ver == rhs.conf_ver && ver == rhs.ver; } 58 | 59 | // for debug output 60 | std::string toString() const 61 | { 62 | return "{" + std::to_string(id) + "," + std::to_string(conf_ver) + "," + std::to_string(ver) + "}"; 63 | } 64 | }; 65 | 66 | } // namespace kv 67 | } // namespace pingcap 68 | 69 | namespace std 70 | { 71 | template <> 72 | struct hash 73 | { 74 | using argument_type = pingcap::kv::RegionVerID; 75 | using result_type = size_t; 76 | size_t operator()(const pingcap::kv::RegionVerID & key) const { return key.id; } 77 | }; 78 | } // namespace std 79 | 80 | namespace pingcap 81 | { 82 | namespace kv 83 | { 84 | struct Region 85 | { 86 | metapb::Region meta; 87 | metapb::Peer leader_peer; 88 | std::vector pending_peers; 89 | std::atomic_uint work_tiflash_peer_idx; 90 | 91 | Region(const metapb::Region & meta_, const metapb::Peer & peer_) 92 | : meta(meta_) 93 | , leader_peer(peer_) 94 | , work_tiflash_peer_idx(0) 95 | {} 96 | 97 | Region(const metapb::Region & meta_, const metapb::Peer & peer_, const std::vector & pending_peers_) 98 | : meta(meta_) 99 | , leader_peer(peer_) 100 | , pending_peers(pending_peers_) 101 | , work_tiflash_peer_idx(0) 102 | {} 103 | 104 | const std::string & startKey() const { return meta.start_key(); } 105 | 106 | const std::string & endKey() const { return meta.end_key(); } 107 | 108 | bool contains(const std::string & key) const { return key >= startKey() && (key < endKey() || meta.end_key().empty()); } 109 | 110 | RegionVerID verID() const 111 | { 112 | return RegionVerID{ 113 | meta.id(), 114 | meta.region_epoch().conf_ver(), 115 | meta.region_epoch().version(), 116 | }; 117 | } 118 | 119 | bool switchPeer(uint64_t peer_id) 120 | { 121 | for (const auto & peer : meta.peers()) 122 | { 123 | if (peer.id() == peer_id) 124 | { 125 | leader_peer = peer; 126 | return true; 127 | } 128 | } 129 | return false; 130 | } 131 | }; 132 | 133 | using RegionPtr = std::shared_ptr; 134 | using LabelFilter = bool (*)(const std::map &); 135 | using StoreFilter = std::function; 136 | 137 | struct KeyLocation 138 | { 139 | RegionVerID region; 140 | std::string start_key; 141 | std::string end_key; 142 | 143 | KeyLocation() = default; 144 | KeyLocation(const RegionVerID & region_, const std::string & start_key_, const std::string & end_key_) 145 | : region(region_) 146 | , start_key(start_key_) 147 | , end_key(end_key_) 148 | {} 149 | 150 | bool contains(const std::string & key) const { return key >= start_key && (key < end_key || end_key.empty()); } 151 | }; 152 | 153 | struct RPCContext 154 | { 155 | RegionVerID region; 156 | metapb::Region meta; 157 | metapb::Peer peer; 158 | Store store; 159 | std::string addr; 160 | 161 | RPCContext(const RegionVerID & region_, const metapb::Region & meta_, const metapb::Peer & peer_, const Store & store_, const std::string & addr_) 162 | : region(region_) 163 | , meta(meta_) 164 | , peer(peer_) 165 | , store(store_) 166 | , addr(addr_) 167 | {} 168 | 169 | std::string toString() const 170 | { 171 | return "region id: " + std::to_string(region.id) + ", meta: " + meta.DebugString() + ", peer: " 172 | + peer.DebugString() + ", addr: " + addr; 173 | } 174 | }; 175 | 176 | using RPCContextPtr = std::shared_ptr; 177 | 178 | class RegionCache 179 | { 180 | public: 181 | RegionCache(pd::ClientPtr pdClient_, const ClusterConfig & config) 182 | : pd_client(pdClient_) 183 | , tiflash_engine_key(config.tiflash_engine_key) 184 | , tiflash_engine_value(config.tiflash_engine_value) 185 | , log(&Logger::get("pingcap.tikv")) 186 | {} 187 | 188 | RPCContextPtr getRPCContext(Backoffer & bo, 189 | const RegionVerID & id, 190 | StoreType store_type, 191 | bool load_balance, 192 | const LabelFilter & tiflash_label_filter, 193 | const std::unordered_set * store_id_blocklist = nullptr, 194 | uint64_t prefer_store_id = 0); 195 | 196 | bool updateLeader(const RegionVerID & region_id, const metapb::Peer & leader); 197 | 198 | KeyLocation locateKey(Backoffer & bo, const std::string & key); 199 | 200 | void dropRegion(const RegionVerID &); 201 | 202 | void dropStore(uint64_t failed_store_id); 203 | 204 | void onSendReqFail(RPCContextPtr & ctx, const Exception & exc); 205 | 206 | void onSendReqFailForBatchRegions(const std::vector & region_ids, uint64_t store_id); 207 | 208 | void onRegionStale(Backoffer & bo, RPCContextPtr ctx, const errorpb::EpochNotMatch & stale_epoch); 209 | 210 | RegionPtr getRegionByID(Backoffer & bo, const RegionVerID & id); 211 | 212 | Store getStore(Backoffer & bo, uint64_t id); 213 | void forceReloadAllStores(); 214 | 215 | // Return values: 216 | // 1. all stores of this region. 217 | // 2. stores of non pending peers of this region. 218 | std::pair, std::vector> getAllValidTiFlashStores( 219 | Backoffer & bo, 220 | const RegionVerID & region_id, 221 | const Store & current_store, 222 | const LabelFilter & label_filter, 223 | const std::unordered_set * store_id_blocklist = nullptr); 224 | 225 | std::pair>, RegionVerID> 226 | groupKeysByRegion(Backoffer & bo, 227 | const std::vector & keys); 228 | 229 | std::map getAllTiFlashStores(const LabelFilter & label_filter, bool exclude_tombstone); 230 | 231 | void updateCachePeriodically(); 232 | 233 | void stop() 234 | { 235 | stopped.store(true); 236 | std::lock_guard lock(update_cache_mu); 237 | update_cache_cv.notify_all(); 238 | } 239 | 240 | private: 241 | RegionPtr loadRegionByKey(Backoffer & bo, const std::string & key); 242 | 243 | RegionPtr getRegionByIDFromCache(const RegionVerID & region_id); 244 | 245 | RegionPtr loadRegionByID(Backoffer & bo, uint64_t region_id); 246 | 247 | metapb::Store loadStore(Backoffer & bo, uint64_t id); 248 | 249 | Store reloadStoreWithoutLock(const metapb::Store & store); 250 | 251 | RegionPtr searchCachedRegion(const std::string & key); 252 | 253 | std::vector selectTiFlashPeers(Backoffer & bo, const metapb::Region & meta, const LabelFilter & label_filter); 254 | 255 | void insertRegionToCache(RegionPtr region); 256 | 257 | std::map regions_map; 258 | 259 | std::unordered_map regions; 260 | 261 | /// stores the last work_flash_index when a region is dropped, this value is 262 | /// used to initialize the work_flash_index when a region with the same id 263 | /// is added next time. 264 | std::unordered_map region_last_work_flash_index; 265 | 266 | std::map stores; 267 | 268 | pd::ClientPtr pd_client; 269 | 270 | std::shared_mutex region_mutex; 271 | 272 | std::mutex store_mutex; 273 | 274 | const std::string tiflash_engine_key; 275 | 276 | const std::string tiflash_engine_value; 277 | 278 | Logger * log; 279 | 280 | std::atomic stopped = false; 281 | std::mutex update_cache_mu; 282 | std::condition_variable update_cache_cv; 283 | }; 284 | 285 | using RegionCachePtr = std::unique_ptr; 286 | static const std::string DCLabelKey = "zone"; 287 | static const std::string EngineLabelKey = "engine"; 288 | static const std::string EngineLabelTiFlash = "tiflash"; 289 | static const std::string EngineRoleLabelKey = "engine_role"; 290 | static const std::string EngineRoleWrite = "write"; 291 | 292 | bool hasLabel(const std::map & labels, const std::string & key, const std::string & val); 293 | // Returns false means label doesn't match, and will ignore this store. 294 | bool labelFilterOnlyTiFlashWriteNode(const std::map & labels); 295 | bool labelFilterNoTiFlashWriteNode(const std::map & labels); 296 | bool labelFilterAllTiFlashNode(const std::map & labels); 297 | bool labelFilterAllNode(const std::map &); 298 | bool labelFilterInvalid(const std::map &); 299 | 300 | } // namespace kv 301 | } // namespace pingcap 302 | -------------------------------------------------------------------------------- /include/pingcap/kv/RegionClient.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | namespace pingcap 9 | { 10 | namespace kv 11 | { 12 | constexpr int dailTimeout = 5; 13 | constexpr int copTimeout = 20; 14 | 15 | // RegionClient sends KV/Cop requests to tikv/tiflash server (corresponding to `RegionRequestSender` in go-client). It handles network errors and some region errors internally. 16 | // 17 | // Typically, a KV/Cop requests is bind to a region, all keys that are involved in the request should be located in the region. 18 | // The sending process begins with looking for the address of leader (maybe learners for ReadIndex request) store's address of the target region from cache, 19 | // and the request is then sent to the destination TiKV server over TCP connection. 20 | // If region is updated, can be caused by leader transfer, region split, region merge, or region balance, tikv server may not able to process request and send back a RegionError. 21 | // RegionClient takes care of errors that does not relevant to region range, such as 'I/O timeout', 'NotLeader', and 'ServerIsBusy'. 22 | // For other errors, since region range have changed, the request may need to split, so we simply return the error to caller. 23 | 24 | struct RegionClient 25 | { 26 | Cluster * cluster; 27 | 28 | const RegionVerID region_id; 29 | 30 | Logger * log; 31 | 32 | RegionClient(Cluster * cluster_, const RegionVerID & id) 33 | : cluster(cluster_) 34 | , region_id(id) 35 | , log(&Logger::get("pingcap.tikv")) 36 | {} 37 | 38 | // This method send a request to region, but is NOT Thread-Safe !! 39 | template 40 | void sendReqToRegion(Backoffer & bo, 41 | REQ & req, 42 | RESP * resp, 43 | const LabelFilter & tiflash_label_filter = kv::labelFilterInvalid, 44 | int timeout = dailTimeout, 45 | StoreType store_type = StoreType::TiKV, 46 | const kv::GRPCMetaData & meta_data = {}, 47 | const std::unordered_set * store_id_blocklist = nullptr, 48 | const std::string & source_zone_label = "", 49 | bool * same_zone_flag = nullptr, 50 | uint64_t prefer_store_id = 0) 51 | { 52 | if (store_type == kv::StoreType::TiFlash && tiflash_label_filter == kv::labelFilterInvalid) 53 | { 54 | throw Exception("should setup proper label_filter for tiflash"); 55 | } 56 | for (;;) 57 | { 58 | RPCContextPtr ctx = cluster->region_cache->getRPCContext(bo, region_id, store_type, /*load_balance=*/true, tiflash_label_filter, store_id_blocklist, prefer_store_id); 59 | if (ctx == nullptr) 60 | { 61 | // If the region is not found in cache, it must be out 62 | // of date and already be cleaned up. We can skip the 63 | // RPC by returning RegionError directly. 64 | auto s = store_id_blocklist != nullptr ? ", store_filter_size=" + std::to_string(store_id_blocklist->size()) + "." : std::string("."); 65 | throw Exception("Region epoch not match after retries: Region " + region_id.toString() + " not in region cache" + s, RegionEpochNotMatch); 66 | } 67 | RpcCall rpc(cluster->rpc_client, ctx->addr); 68 | rpc.setRequestCtx(req, ctx, cluster->api_version); 69 | 70 | grpc::ClientContext context; 71 | rpc.setClientContext(context, timeout, meta_data); 72 | 73 | auto status = rpc.call(&context, req, resp); 74 | if (!status.ok()) 75 | { 76 | auto extra_msg = "region_id: " + region_id.toString() + ", addr: " + ctx->addr; 77 | if (status.error_code() == ::grpc::StatusCode::UNIMPLEMENTED) 78 | { 79 | // The rpc is not implemented on this service. 80 | throw Exception("rpc is not implemented: " + rpc.errMsg(status, extra_msg), GRPCNotImplemented); 81 | } 82 | std::string err_msg = rpc.errMsg(status, extra_msg); 83 | log->warning(err_msg); 84 | onSendFail(bo, Exception(err_msg, GRPCErrorCode), ctx); 85 | continue; 86 | } 87 | if (resp->has_region_error()) 88 | { 89 | log->warning("region_id " + region_id.toString() + " find error: " + resp->region_error().DebugString()); 90 | onRegionError(bo, ctx, resp->region_error()); 91 | continue; 92 | } 93 | if (same_zone_flag && !source_zone_label.empty()) { 94 | auto iter = ctx->store.labels.find(DCLabelKey); 95 | if (iter != ctx->store.labels.end()) { 96 | *same_zone_flag = iter->second == source_zone_label; 97 | } 98 | } 99 | return; 100 | } 101 | } 102 | 103 | template 104 | class StreamReader 105 | { 106 | public: 107 | StreamReader() = default; 108 | StreamReader(StreamReader && other) = default; 109 | StreamReader & operator=(StreamReader && other) = default; 110 | 111 | bool read(RESP * msg) 112 | { 113 | if (no_resp) 114 | return false; 115 | if (is_first_read) 116 | { 117 | is_first_read = false; 118 | msg->Swap(&first_resp); 119 | return true; 120 | } 121 | return reader->Read(msg); 122 | } 123 | 124 | ::grpc::Status finish() 125 | { 126 | if (no_resp) 127 | return ::grpc::Status::OK; 128 | return reader->Finish(); 129 | } 130 | 131 | private: 132 | friend struct RegionClient; 133 | ::grpc::ClientContext context; 134 | std::unique_ptr<::grpc::ClientReader> reader; 135 | bool no_resp = false; 136 | bool is_first_read = true; 137 | RESP first_resp; 138 | }; 139 | 140 | template 141 | std::unique_ptr> sendStreamReqToRegion(Backoffer & bo, 142 | REQ & req, 143 | const LabelFilter & tiflash_label_filter = kv::labelFilterInvalid, 144 | int timeout = dailTimeout, 145 | StoreType store_type = StoreType::TiKV, 146 | const kv::GRPCMetaData & meta_data = {}, 147 | const std::unordered_set * store_id_blocklist = nullptr, 148 | const std::string & source_zone_label = "", 149 | bool * same_zone_flag = nullptr, 150 | uint64_t prefer_store_id = 0) 151 | { 152 | if (store_type == kv::StoreType::TiFlash && tiflash_label_filter == kv::labelFilterInvalid) 153 | { 154 | throw Exception("should setup proper label_filter for tiflash"); 155 | } 156 | for (;;) 157 | { 158 | RPCContextPtr ctx = cluster->region_cache->getRPCContext(bo, region_id, store_type, /*load_balance=*/true, tiflash_label_filter, store_id_blocklist, prefer_store_id); 159 | if (ctx == nullptr) 160 | { 161 | // If the region is not found in cache, it must be out 162 | // of date and already be cleaned up. We can skip the 163 | // RPC by returning RegionError directly. 164 | throw Exception("Region epoch not match after retries: Region " + region_id.toString() + " not in region cache.", RegionEpochNotMatch); 165 | } 166 | 167 | auto stream_reader = std::make_unique>(); 168 | RpcCall rpc(cluster->rpc_client, ctx->addr); 169 | rpc.setRequestCtx(req, ctx, cluster->api_version); 170 | rpc.setClientContext(stream_reader->context, timeout, meta_data); 171 | 172 | stream_reader->reader = rpc.call(&stream_reader->context, req); 173 | if (stream_reader->reader->Read(&stream_reader->first_resp)) 174 | { 175 | if (stream_reader->first_resp.has_region_error()) 176 | { 177 | log->warning("region_id " + region_id.toString() + " find error: " + stream_reader->first_resp.region_error().message()); 178 | onRegionError(bo, ctx, stream_reader->first_resp.region_error()); 179 | continue; 180 | } 181 | return stream_reader; 182 | } 183 | auto status = stream_reader->reader->Finish(); 184 | if (status.ok()) 185 | { 186 | if (same_zone_flag && !source_zone_label.empty()) { 187 | auto iter = ctx->store.labels.find(DCLabelKey); 188 | if (iter != ctx->store.labels.end()) { 189 | *same_zone_flag = iter->second == source_zone_label; 190 | } 191 | } 192 | // No response msg. 193 | stream_reader->no_resp = true; 194 | return stream_reader; 195 | } 196 | auto extra_msg = "region_id: " + region_id.toString() + ", addr: " + ctx->addr; 197 | if (status.error_code() == ::grpc::StatusCode::UNIMPLEMENTED) 198 | { 199 | 200 | // The rpc is not implemented on this service. 201 | throw Exception("rpc is not implemented: " + rpc.errMsg(status, extra_msg), GRPCNotImplemented); 202 | } 203 | std::string err_msg = rpc.errMsg(status, extra_msg); 204 | log->warning(err_msg); 205 | onSendFail(bo, Exception(err_msg, GRPCErrorCode), ctx); 206 | } 207 | } 208 | 209 | protected: 210 | void onRegionError(Backoffer & bo, RPCContextPtr rpc_ctx, const errorpb::Error & err) const; 211 | 212 | // Normally, it happens when machine down or network partition between tidb and kv or process crash. 213 | void onSendFail(Backoffer & bo, const Exception & e, RPCContextPtr rpc_ctx) const; 214 | }; 215 | 216 | } // namespace kv 217 | } // namespace pingcap 218 | -------------------------------------------------------------------------------- /src/test/real_tikv_test/2pc_test.cc: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | #include 5 | 6 | #include 7 | 8 | #include "../test_helper.h" 9 | 10 | namespace pingcap 11 | { 12 | namespace kv 13 | { 14 | 15 | // Just for test purpose 16 | struct TestTwoPhaseCommitter 17 | { 18 | private: 19 | TwoPhaseCommitterPtr committer; 20 | 21 | public: 22 | TestTwoPhaseCommitter(Txn * txn) : committer(std::make_shared(txn)) {} 23 | 24 | void prewriteKeys(Backoffer & bo, const std::vector & keys) { committer->prewriteKeys(bo, keys); } 25 | 26 | void commitKeys(Backoffer & bo, const std::vector & keys) { committer->commitKeys(bo, keys); } 27 | 28 | std::vector keys() { return committer->keys; } 29 | 30 | void setCommitTS(int64_t commit_ts) { committer->commit_ts = commit_ts; } 31 | }; 32 | 33 | } // namespace kv 34 | } // namespace pingcap 35 | 36 | namespace pingcap::tests 37 | { 38 | 39 | using namespace pingcap; 40 | using namespace pingcap::kv; 41 | 42 | struct TestUtil 43 | { 44 | static std::string get_random_string(size_t length) 45 | { 46 | auto randchar = []() -> char { 47 | const char charset[] = "0123456789" 48 | "ABCDEFGHIJKLMNOPQRSTUVWXYZ" 49 | "abcdefghijklmnopqrstuvwxyz"; 50 | const size_t max_index = (sizeof(charset) - 1); 51 | return charset[rand() % max_index]; 52 | }; 53 | std::string str(length, 0); 54 | std::generate_n(str.begin(), length, randchar); 55 | return str; 56 | } 57 | }; 58 | 59 | class TestWith2PCRealTiKV : public testing::Test 60 | { 61 | protected: 62 | void SetUp() override 63 | { 64 | std::vector pd_addrs{"127.0.0.1:2379"}; 65 | 66 | test_cluster = createCluster(pd_addrs); 67 | } 68 | 69 | ClusterPtr test_cluster; 70 | }; 71 | 72 | TEST_F(TestWith2PCRealTiKV, testCommitRollback) 73 | { 74 | 75 | // Commit. 76 | { 77 | Txn txn(test_cluster.get()); 78 | txn.set("a", "a"); 79 | txn.set("b", "b"); 80 | txn.set("c", "c"); 81 | txn.commit(); 82 | 83 | Snapshot snap(test_cluster.get()); 84 | ASSERT_EQ(snap.Get("a"), "a"); 85 | ASSERT_EQ(snap.Get("b"), "b"); 86 | ASSERT_EQ(snap.Get("c"), "c"); 87 | } 88 | 89 | // Write conflict. 90 | { 91 | Txn txn1(test_cluster.get()); 92 | txn1.set("a", "a1"); 93 | txn1.set("b", "b1"); 94 | txn1.set("c", "c1"); 95 | 96 | Txn txn2(test_cluster.get()); 97 | txn2.set("c", "c2"); 98 | txn2.commit(); 99 | 100 | txn1.commit(); 101 | 102 | Snapshot snap(test_cluster.get()); 103 | ASSERT_EQ(snap.Get("a"), "a"); 104 | ASSERT_EQ(snap.Get("b"), "b"); 105 | ASSERT_EQ(snap.Get("c"), "c2"); 106 | } 107 | } 108 | 109 | TEST_F(TestWith2PCRealTiKV, commitAfterReadByOtherTxn) 110 | { 111 | 112 | // Commit. 113 | { 114 | Txn txn(test_cluster.get()); 115 | txn.set("a", "a"); 116 | txn.set("b", "b"); 117 | txn.set("c", "c"); 118 | txn.commit(); 119 | 120 | Snapshot snap(test_cluster.get()); 121 | ASSERT_EQ(snap.Get("a"), "a"); 122 | ASSERT_EQ(snap.Get("b"), "b"); 123 | ASSERT_EQ(snap.Get("c"), "c"); 124 | } 125 | 126 | // Prewrite and commit after read by other txn. 127 | { 128 | Txn txn1(test_cluster.get()); 129 | txn1.set("a", "a1"); 130 | txn1.set("b", "b1"); 131 | txn1.set("c", "c1"); 132 | TestTwoPhaseCommitter committer{&txn1}; 133 | Backoffer prewrite_bo(prewriteMaxBackoff); 134 | committer.prewriteKeys(prewrite_bo, committer.keys()); 135 | 136 | // read by other txn after prewrite 137 | Txn txn2(test_cluster.get()); 138 | auto result = txn2.get("a"); 139 | ASSERT_EQ(result.second, true); 140 | ASSERT_EQ(result.first, "a"); 141 | auto result2 = txn2.get("b"); 142 | ASSERT_EQ(result2.second, true); 143 | ASSERT_EQ(result2.first, "b"); 144 | auto result3 = txn2.get("c"); 145 | ASSERT_EQ(result3.second, true); 146 | ASSERT_EQ(result3.first, "c"); 147 | 148 | // commit after read by other txn 149 | committer.setCommitTS(test_cluster->pd_client->getTS()); 150 | Backoffer commit_bo(commitMaxBackoff); 151 | committer.commitKeys(commit_bo, committer.keys()); 152 | 153 | Snapshot snap2(test_cluster.get()); 154 | ASSERT_EQ(snap2.Get("a"), "a1"); 155 | ASSERT_EQ(snap2.Get("b"), "b1"); 156 | ASSERT_EQ(snap2.Get("c"), "c1"); 157 | } 158 | } 159 | 160 | TEST_F(TestWith2PCRealTiKV, testLargeTxn) 161 | { 162 | // Commit. 163 | { 164 | Txn txn(test_cluster.get()); 165 | txn.set("a", "a0"); 166 | txn.set("b", "b0"); 167 | txn.set("c", "c0"); 168 | txn.commit(); 169 | 170 | Snapshot snap(test_cluster.get()); 171 | ASSERT_EQ(snap.Get("a"), "a0"); 172 | ASSERT_EQ(snap.Get("b"), "b0"); 173 | ASSERT_EQ(snap.Get("c"), "c0"); 174 | } 175 | 176 | // Prewrite. 177 | { 178 | Txn txn1(test_cluster.get()); 179 | txn1.set("a", "a1"); 180 | txn1.set("b", "b1"); 181 | txn1.set("c", "c1"); 182 | std::unordered_set inserted_keys; 183 | for (size_t i = 0; i < 33 * 1024 * 1024; i++) 184 | { 185 | if (i % 1000000 == 0) 186 | { 187 | std::cout << "process to " << std::to_string(i) << std::endl; 188 | } 189 | for (;;) 190 | { 191 | std::string rand_str = TestUtil::get_random_string(rand() % 30 + 10); 192 | if (inserted_keys.find(rand_str) == inserted_keys.end()) 193 | { 194 | txn1.set(rand_str, rand_str); 195 | break; 196 | } 197 | } 198 | } 199 | 200 | TestTwoPhaseCommitter committer{&txn1}; 201 | Backoffer prewrite_bo(prewriteMaxBackoff); 202 | try 203 | { 204 | committer.prewriteKeys(prewrite_bo, committer.keys()); 205 | } 206 | catch (Exception & e) 207 | { 208 | std::cout << "Prewrite meet exception: " << e.message() << std::endl; 209 | } 210 | 211 | 212 | Snapshot snap1(test_cluster.get()); 213 | ASSERT_EQ(snap1.Get("a"), "a0"); 214 | ASSERT_EQ(snap1.Get("b"), "b0"); 215 | ASSERT_EQ(snap1.Get("c"), "c0"); 216 | 217 | std::this_thread::sleep_for(std::chrono::milliseconds(10000)); 218 | 219 | try 220 | { 221 | committer.setCommitTS(test_cluster->pd_client->getTS()); 222 | Backoffer commit_bo(commitMaxBackoff); 223 | committer.commitKeys(commit_bo, committer.keys()); 224 | } 225 | catch (Exception & e) 226 | { 227 | std::cout << "Commit meet exception: " << e.message() << std::endl; 228 | } 229 | 230 | Snapshot snap2(test_cluster.get()); 231 | ASSERT_EQ(snap2.Get("a"), "a1"); 232 | ASSERT_EQ(snap2.Get("b"), "b1"); 233 | ASSERT_EQ(snap2.Get("c"), "c1"); 234 | } 235 | } 236 | 237 | TEST_F(TestWith2PCRealTiKV, testScanWithLargeTxn) 238 | { 239 | std::vector keys; 240 | std::string prefix = "test_scan_"; 241 | std::string candidates = "abcdefghijklmnopqrstuvwxyz"; 242 | for (size_t i = 0; i <= candidates.size(); i++) 243 | { 244 | keys.push_back(prefix + candidates[i]); 245 | } 246 | // Commit. 247 | { 248 | Txn txn(test_cluster.get()); 249 | for (auto & key : keys) 250 | { 251 | txn.set(key, key); 252 | } 253 | txn.commit(); 254 | 255 | Snapshot snap(test_cluster.get()); 256 | for (auto & key : keys) 257 | { 258 | ASSERT_EQ(snap.Get(key), key); 259 | } 260 | } 261 | 262 | // Prewrite. 263 | { 264 | Txn txn1(test_cluster.get()); 265 | for (auto & key : keys) 266 | { 267 | txn1.set(key, key + "1"); 268 | } 269 | std::unordered_set inserted_keys; 270 | for (size_t i = 0; i < 33 * 1024 * 1024; i++) 271 | { 272 | if (i % 1000000 == 0) 273 | { 274 | std::cout << "process to " << std::to_string(i) << std::endl; 275 | } 276 | for (;;) 277 | { 278 | std::string rand_str = TestUtil::get_random_string(rand() % 30 + 10); 279 | if (inserted_keys.find(rand_str) == inserted_keys.end()) 280 | { 281 | txn1.set(rand_str, rand_str); 282 | break; 283 | } 284 | } 285 | } 286 | 287 | std::cout << "Begin to prewrite\n"; 288 | TestTwoPhaseCommitter committer{&txn1}; 289 | Backoffer prewrite_bo(prewriteMaxBackoff); 290 | try 291 | { 292 | committer.prewriteKeys(prewrite_bo, committer.keys()); 293 | } 294 | catch (Exception & e) 295 | { 296 | std::cout << "Prewrite meet exception: " << e.message() << std::endl; 297 | } 298 | 299 | for (size_t i = 0; i < 10; i++) 300 | { 301 | Snapshot snap1(test_cluster.get()); 302 | auto scanner = snap1.Scan(prefix, prefixNext(prefix)); 303 | size_t count = 0; 304 | while (scanner.valid) 305 | { 306 | auto key = scanner.key(); 307 | if (std::find(keys.begin(), keys.end(), key) == keys.end()) 308 | { 309 | scanner.next(); 310 | continue; 311 | } 312 | ASSERT_EQ(key, scanner.value()); 313 | scanner.next(); 314 | count += 1; 315 | } 316 | ASSERT_EQ(count, keys.size()); 317 | std::this_thread::sleep_for(std::chrono::milliseconds(1000)); 318 | } 319 | 320 | std::cout << "Begin to commit\n"; 321 | try 322 | { 323 | committer.setCommitTS(test_cluster->pd_client->getTS()); 324 | Backoffer commit_bo(commitMaxBackoff); 325 | committer.commitKeys(commit_bo, committer.keys()); 326 | } 327 | catch (Exception & e) 328 | { 329 | std::cout << "Commit meet exception: " << e.message() << std::endl; 330 | } 331 | 332 | Snapshot snap(test_cluster.get()); 333 | for (auto & key : keys) 334 | { 335 | ASSERT_EQ(snap.Get(key), key + "1"); 336 | } 337 | } 338 | } 339 | 340 | } // namespace 341 | -------------------------------------------------------------------------------- /include/pingcap/coprocessor/Client.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include 3 | #include 4 | #include 5 | #include 6 | 7 | #include 8 | #include 9 | #include 10 | #include 11 | #include 12 | 13 | namespace pingcap 14 | { 15 | namespace coprocessor 16 | { 17 | enum ReqType : int16_t 18 | { 19 | Select = 101, 20 | Index = 102, 21 | DAG = 103, 22 | Analyze = 104, 23 | Checksum = 105, 24 | }; 25 | 26 | struct KeyRange 27 | { 28 | std::string start_key; 29 | std::string end_key; 30 | KeyRange(const std::string & start_key_, const std::string & end_key_) 31 | : start_key(start_key_) 32 | , end_key(end_key_) 33 | {} 34 | 35 | void setKeyRange(::coprocessor::KeyRange * range) const 36 | { 37 | range->set_start(start_key); 38 | range->set_end(end_key); 39 | } 40 | 41 | KeyRange(KeyRange &&) = default; 42 | KeyRange(const KeyRange &) = default; 43 | KeyRange & operator=(const KeyRange &) = default; 44 | KeyRange & operator=(KeyRange &&) = default; 45 | 46 | bool operator<(const KeyRange & rhs) const { return start_key < rhs.start_key; } 47 | }; 48 | using KeyRanges = std::vector; 49 | 50 | struct Request 51 | { 52 | int64_t tp; 53 | uint64_t start_ts; 54 | std::string data; 55 | int64_t schema_version; 56 | std::string resource_group_name; 57 | }; 58 | 59 | using RequestPtr = std::shared_ptr; 60 | 61 | struct CopTask 62 | { 63 | kv::RegionVerID region_id; 64 | KeyRanges ranges; 65 | RequestPtr req; 66 | kv::StoreType store_type; 67 | int64_t partition_index; 68 | kv::GRPCMetaData meta_data; 69 | // call before send request, can be used to collect TiFlash metrics. 70 | std::function before_send; 71 | pd::KeyspaceID keyspace_id; 72 | uint64_t connection_id; 73 | std::string connection_alias; 74 | uint64_t prefer_store_id = 0; 75 | }; 76 | 77 | struct RegionInfo 78 | { 79 | kv::RegionVerID region_id; 80 | // meta; 81 | KeyRanges ranges; 82 | std::vector all_stores; 83 | // Used by PartitionTableScan, indicates the n-th partition of the partition table. 84 | int64_t partition_index; 85 | }; 86 | 87 | struct TableRegions 88 | { 89 | int64_t physical_table_id; 90 | std::vector region_infos; 91 | }; 92 | 93 | struct BatchCopTask 94 | { 95 | std::string store_addr; 96 | std::vector region_infos; 97 | // Used by PartitionTableScan, indicates region infos for each partition. 98 | std::vector table_regions; 99 | RequestPtr req; 100 | kv::StoreType store_type; 101 | 102 | uint64_t store_id; 103 | std::map store_labels; 104 | }; 105 | 106 | /// A iterator dedicated to send coprocessor(stream) request and receive responses. 107 | /// All functions are thread-safe. 108 | class ResponseIter 109 | { 110 | public: 111 | struct Result 112 | { 113 | std::shared_ptr<::coprocessor::Response> resp; 114 | bool same_zone{true}; 115 | Exception error; 116 | bool finished{false}; 117 | 118 | Result() = default; 119 | explicit Result(std::shared_ptr<::coprocessor::Response> resp_) 120 | : resp(resp_) 121 | {} 122 | explicit Result(const Exception & err) 123 | : error(err) 124 | {} 125 | explicit Result(bool finished_) 126 | : finished(finished_) 127 | {} 128 | Result(std::shared_ptr<::coprocessor::Response> resp_, bool same_zone_) 129 | : resp(resp_) 130 | , same_zone(same_zone_) 131 | {} 132 | 133 | const std::string & data() const { return resp->data(); } 134 | }; 135 | 136 | ResponseIter(std::unique_ptr> && queue_, 137 | std::vector && tasks_, 138 | kv::Cluster * cluster_, 139 | int concurrency_, 140 | Logger * log_, 141 | int timeout_ = kv::copTimeout, 142 | const kv::LabelFilter & tiflash_label_filter_ = kv::labelFilterInvalid, 143 | const std::string source_zone_label_ = "", 144 | uint64_t prefer_store_id_ = 0) 145 | : queue(std::move(queue_)) 146 | , tasks(std::move(tasks_)) 147 | , cluster(cluster_) 148 | , concurrency(concurrency_) 149 | , timeout(timeout_) 150 | , unfinished_thread(0) 151 | , tiflash_label_filter(tiflash_label_filter_) 152 | , source_zone_label(source_zone_label_) 153 | , log(log_) 154 | , prefer_store_id(prefer_store_id_) 155 | {} 156 | 157 | ~ResponseIter() 158 | { 159 | is_cancelled = true; 160 | for (auto & worker_thread : worker_threads) 161 | { 162 | worker_thread.join(); 163 | } 164 | } 165 | 166 | template 167 | void open() 168 | { 169 | bool old_val = false; 170 | if (!is_opened.compare_exchange_strong(old_val, true)) 171 | return; 172 | 173 | unfinished_thread = concurrency; 174 | for (int i = 0; i < concurrency; i++) 175 | { 176 | std::thread worker(&ResponseIter::thread, this); 177 | worker_threads.push_back(std::move(worker)); 178 | } 179 | 180 | log->debug("coprocessor has " + std::to_string(tasks.size()) + " tasks."); 181 | } 182 | 183 | void cancel() 184 | { 185 | bool old_val = false; 186 | if (!is_cancelled.compare_exchange_strong(old_val, true)) 187 | return; 188 | queue->cancel(); 189 | } 190 | 191 | std::pair nonBlockingNext() 192 | { 193 | assert(is_opened); 194 | Result res; 195 | switch (queue->tryPop(res)) 196 | { 197 | case common::MPMCQueueResult::OK: 198 | return {std::move(res), true}; 199 | case common::MPMCQueueResult::CANCELLED: 200 | case common::MPMCQueueResult::FINISHED: 201 | return {Result(true), false}; 202 | case common::MPMCQueueResult::EMPTY: 203 | return {Result(), false}; 204 | default: 205 | __builtin_unreachable(); 206 | } 207 | } 208 | 209 | std::pair next() 210 | { 211 | assert(is_opened); 212 | Result res; 213 | switch (queue->pop(res)) 214 | { 215 | case common::MPMCQueueResult::OK: 216 | return {std::move(res), true}; 217 | case common::MPMCQueueResult::CANCELLED: 218 | case common::MPMCQueueResult::FINISHED: 219 | return {Result(true), false}; 220 | default: 221 | __builtin_unreachable(); 222 | } 223 | } 224 | 225 | private: 226 | template 227 | void thread() 228 | { 229 | log->information("thread start."); 230 | while (true) 231 | { 232 | if (is_cancelled || meet_error) 233 | { 234 | const char * msg = is_cancelled ? "has been cancelled" : "already meet error"; 235 | log->information("cop task exit because " + std::string(msg)); 236 | return; 237 | } 238 | std::unique_lock lk(task_mutex); 239 | if (tasks.size() == task_index) 240 | { 241 | lk.unlock(); 242 | if (unfinished_thread.fetch_sub(1) == 1) 243 | queue->finish(); 244 | return; 245 | } 246 | const CopTask & task = tasks[task_index]; 247 | task_index++; 248 | lk.unlock(); 249 | handleTask(task); 250 | } 251 | } 252 | 253 | template 254 | std::vector handleTaskImpl(kv::Backoffer & bo, const CopTask & task); 255 | template 256 | void handleTask(const CopTask & task); 257 | 258 | std::unique_ptr> queue; 259 | 260 | std::mutex task_mutex; 261 | size_t task_index = 0; 262 | const std::vector tasks; 263 | 264 | std::vector worker_threads; 265 | 266 | kv::Cluster * cluster; 267 | const int concurrency; 268 | const int timeout; 269 | kv::MinCommitTSPushed min_commit_ts_pushed; 270 | 271 | std::atomic_int unfinished_thread; 272 | 273 | std::atomic_bool is_cancelled = false; 274 | std::atomic_bool meet_error = false; 275 | std::atomic_bool is_opened = false; 276 | 277 | const kv::LabelFilter tiflash_label_filter; 278 | const std::string source_zone_label; 279 | 280 | Logger * log; 281 | 282 | const uint64_t prefer_store_id; 283 | }; 284 | 285 | std::vector buildCopTasks( 286 | kv::Backoffer & bo, 287 | kv::Cluster * cluster, 288 | KeyRanges ranges, 289 | RequestPtr cop_req, 290 | kv::StoreType store_type, 291 | pd::KeyspaceID keyspace_id, 292 | uint64_t connection_id, 293 | const std::string & connection_alias, 294 | Logger * log, 295 | kv::GRPCMetaData meta_data = {}, 296 | std::function before_send = {}); 297 | 298 | 299 | /* 300 | centralized_schedule is a parameter that controls the distribution of regions across BatchCopTasks. 301 | 302 | When set to false (default): 303 | - Regions are distributed as evenly as possible among BatchCopTasks 304 | - This balances the workload across tasks 305 | 306 | When set to true: 307 | - Regions are concentrated into as few BatchCopTasks as possible 308 | - This is currently only used for ANN (Approximate Nearest Neighbor) queries in vector search 309 | 310 | The centralized scheduling approach can potentially improve performance for certain types of queries 311 | by reducing the number of tasks and minimizing coordination overhead. 312 | */ 313 | std::vector buildBatchCopTasks( 314 | kv::Backoffer & bo, 315 | kv::Cluster * cluster, 316 | bool is_mpp, 317 | bool is_partition_table_scan, 318 | const std::vector & physical_table_ids, 319 | const std::vector & ranges_for_each_physical_table, 320 | const std::unordered_set * store_id_blocklist, 321 | kv::StoreType store_type, 322 | const kv::LabelFilter & label_filter, 323 | Logger * log, 324 | bool centralized_schedule = false); 325 | 326 | namespace details 327 | { 328 | // LocationKeyRanges wraps a real Location in PD and its logical ranges info. 329 | struct LocationKeyRanges 330 | { 331 | // The read location in PD. 332 | kv::KeyLocation location; 333 | // The logic ranges the current Location contains. 334 | KeyRanges ranges; 335 | }; 336 | 337 | std::vector splitKeyRangesByLocations( 338 | const kv::RegionCachePtr & cache, 339 | kv::Backoffer & bo, 340 | std::vector<::pingcap::coprocessor::KeyRange> ranges); 341 | } // namespace details 342 | 343 | } // namespace coprocessor 344 | } // namespace pingcap 345 | -------------------------------------------------------------------------------- /include/pingcap/kv/LockResolver.h: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | 8 | #include 9 | #include 10 | #include 11 | 12 | namespace pingcap 13 | { 14 | namespace kv 15 | { 16 | struct Cluster; 17 | 18 | // TxnStatus represents a txn's final status. It should be Lock or Commit or Rollback. 19 | struct TxnStatus 20 | { 21 | uint64_t ttl = 0; 22 | uint64_t commit_ts = 0; 23 | ::kvrpcpb::Action action; 24 | std::optional<::kvrpcpb::LockInfo> primary_lock; 25 | bool isCommitted() const { return ttl == 0 && commit_ts > 0; } 26 | 27 | bool isCacheable() const 28 | { 29 | if (isCommitted()) 30 | { 31 | return true; 32 | } 33 | if (ttl == 0) 34 | { 35 | if (action == kvrpcpb::Action::NoAction || action == kvrpcpb::Action::LockNotExistRollback 36 | || action == kvrpcpb::Action::TTLExpireRollback) 37 | { 38 | return true; 39 | } 40 | } 41 | return false; 42 | } 43 | }; 44 | 45 | constexpr size_t resolvedCacheSize = 2048; 46 | constexpr int bigTxnThreshold = 16; 47 | 48 | const uint64_t defaultLockTTL = 3000; 49 | 50 | const uint64_t maxLockTTL = 120000; 51 | 52 | const uint64_t ttlFactor = 6000; 53 | 54 | // Lock represents a lock from tikv server. 55 | struct Lock 56 | { 57 | std::string key; 58 | std::string primary; 59 | uint64_t txn_id; 60 | uint64_t ttl; 61 | uint64_t txn_size; 62 | ::kvrpcpb::Op lock_type; 63 | bool use_async_commit; 64 | uint64_t lock_for_update_ts; 65 | uint64_t min_commit_ts; 66 | bool is_txn_file; 67 | 68 | explicit Lock(const ::kvrpcpb::LockInfo & l) 69 | : key(l.key()) 70 | , primary(l.primary_lock()) 71 | , txn_id(l.lock_version()) 72 | , ttl(l.lock_ttl()) 73 | , txn_size(l.txn_size()) 74 | , lock_type(l.lock_type()) 75 | , use_async_commit(l.use_async_commit()) 76 | , lock_for_update_ts(l.lock_for_update_ts()) 77 | , min_commit_ts(l.min_commit_ts()) 78 | , is_txn_file(l.is_txn_file()) 79 | {} 80 | 81 | std::string toDebugString() const; 82 | }; 83 | 84 | using LockPtr = std::shared_ptr; 85 | 86 | inline LockPtr extractLockFromKeyErr(const ::kvrpcpb::KeyError & key_err) 87 | { 88 | if (key_err.has_locked()) 89 | { 90 | return std::make_shared(key_err.locked()); 91 | } 92 | throw Exception("unknown error : " + key_err.ShortDebugString(), ErrorCodes::UnknownError); 93 | } 94 | 95 | struct TxnExpireTime 96 | { 97 | bool initialized; 98 | int64_t txn_expire; 99 | 100 | TxnExpireTime() 101 | : initialized{false} 102 | , txn_expire{0} 103 | {} 104 | 105 | void update(int64_t lock_expire) 106 | { 107 | if (lock_expire <= 0) 108 | lock_expire = 0; 109 | if (!initialized) 110 | { 111 | txn_expire = lock_expire; 112 | initialized = true; 113 | } 114 | else if (lock_expire < txn_expire) 115 | txn_expire = lock_expire; 116 | } 117 | 118 | int64_t value() const { return initialized ? txn_expire : 0; } 119 | }; 120 | 121 | // AsyncResolveData is data contributed by multiple threads when resolving locks using the async commit protocol. All 122 | // data should be protected by the mu field. 123 | struct AsyncResolveData 124 | { 125 | std::mutex mu; 126 | uint64_t commit_ts = 0; 127 | std::vector keys; 128 | bool missing_lock = false; 129 | 130 | explicit AsyncResolveData(uint64_t _commit_ts, bool _missing_lock = false) 131 | : commit_ts(_commit_ts) 132 | , missing_lock(_missing_lock) 133 | {} 134 | 135 | // addKeys adds the keys from locks to data, keeping other fields up to date. start_ts and _commit_ts are for the 136 | // transaction being resolved. 137 | // 138 | // In the async commit protocol when checking locks, we send a list of keys to check and get back a list of locks. There 139 | // will be a lock for every key which is locked. If there are fewer locks than keys, then a lock is missing because it 140 | // has been committed, rolled back, or was never locked. 141 | // 142 | // In this function, resp->locks is the list of locks, and expected is the number of keys. AsyncResolveData.missing_lock will be 143 | // set to true if the lengths don't match. If the lengths do match, then the locks are added to asyncResolveData.locks 144 | // and will need to be resolved by the caller. 145 | void addKeys(::kvrpcpb::CheckSecondaryLocksResponse * resp, int expected, uint64_t start_ts) 146 | { 147 | std::lock_guard l(mu); 148 | 149 | // Check locks to see if any has been commited or rolled back. 150 | if (resp->locks_size() < expected) 151 | { 152 | // A lock is missing - the transaction must either have been rolled back or committed. 153 | if (!missing_lock) 154 | { 155 | if (resp->commit_ts() != 0 && resp->commit_ts() < commit_ts) 156 | { 157 | // commitTS == 0 => lock has been rolled back. 158 | throw Exception("commit TS must be greater than or equal to min commit TS: commit ts: " 159 | + std::to_string(resp->commit_ts()) + ", min commit ts: " + std::to_string(commit_ts), 160 | ErrorCodes::UnknownError); 161 | } 162 | commit_ts = resp->commit_ts(); 163 | } 164 | missing_lock = true; 165 | if (commit_ts != resp->commit_ts()) 166 | { 167 | throw Exception("commit TS mismatch in async commit recovery: " + std::to_string(commit_ts) + " and " 168 | + std::to_string(resp->commit_ts()), 169 | ErrorCodes::UnknownError); 170 | } 171 | 172 | // We do not need to resolve the remaining locks because TiKV will have resolved them as appropriate. 173 | return; 174 | } 175 | 176 | for (int i = 0; i < resp->locks_size(); i++) 177 | { 178 | const auto & lock = resp->locks(i); 179 | if (lock.lock_version() != start_ts) 180 | { 181 | throw Exception( 182 | "unexpected timestamp, expected: " + std::to_string(start_ts) + ", found: " + std::to_string(lock.lock_version()), 183 | ErrorCodes::UnknownError); 184 | } 185 | 186 | if (!lock.use_async_commit()) 187 | { 188 | throw Exception("CheckSecondaryLocks receives a non-async-commit lock", ErrorCodes::NonAsyncCommit); 189 | } 190 | 191 | if (!missing_lock && lock.min_commit_ts() > commit_ts) 192 | { 193 | commit_ts = lock.min_commit_ts(); 194 | } 195 | keys.push_back(lock.key()); 196 | } 197 | } 198 | }; 199 | 200 | using AsyncResolveDataPtr = std::shared_ptr; 201 | 202 | // LockResolver resolves locks and also caches resolved txn status. 203 | class LockResolver 204 | { 205 | public: 206 | explicit LockResolver(Cluster * cluster_) 207 | : cluster(cluster_) 208 | , log(&Logger::get("pingcap/resolve_lock")) 209 | {} 210 | 211 | void update(Cluster * cluster_) 212 | { 213 | cluster = cluster_; 214 | } 215 | 216 | // resolveLocks tries to resolve Locks. The resolving process is in 3 steps: 217 | // 1) Use the `lockTTL` to pick up all expired locks. Only locks that are too 218 | // old are considered orphan locks and will be handled later. If all locks 219 | // are expired then all locks will be resolved so the returned `ok` will be 220 | // true, otherwise caller should sleep a while before retry. 221 | // 2) For each lock, query the primary key to get txn(which left the lock)'s 222 | // commit status. 223 | // 3) Send `ResolveLock` cmd to the lock's region to resolve all locks belong to 224 | // the same transaction. 225 | 226 | int64_t resolveLocks(Backoffer & bo, uint64_t caller_start_ts, std::vector & locks, std::vector & pushed); 227 | 228 | int64_t resolveLocks( 229 | Backoffer & bo, 230 | uint64_t caller_start_ts, 231 | std::vector & locks, 232 | std::vector & pushed, 233 | bool for_write); 234 | 235 | int64_t resolveLocksForWrite(Backoffer & bo, uint64_t caller_start_ts, std::vector & locks); 236 | 237 | private: 238 | void saveResolved(uint64_t txn_id, const TxnStatus & status) 239 | { 240 | std::unique_lock lk(mu); 241 | 242 | if (resolved.find(txn_id) != resolved.end()) 243 | return; 244 | cached.push(txn_id); 245 | resolved.emplace(txn_id, status); 246 | if (cached.size() > resolvedCacheSize) 247 | { 248 | auto to_remove = cached.front(); 249 | cached.pop(); 250 | resolved.erase(to_remove); 251 | } 252 | } 253 | 254 | TxnStatus * getResolved(uint64_t txn_id) 255 | { 256 | std::shared_lock lk(mu); 257 | 258 | auto it = resolved.find(txn_id); 259 | if (it == resolved.end()) 260 | return nullptr; 261 | return &(it->second); 262 | } 263 | 264 | 265 | void resolveLock(Backoffer & bo, LockPtr lock, TxnStatus & status, std::unordered_set & set); 266 | 267 | 268 | void resolvePessimisticLock(Backoffer & bo, LockPtr lock, std::unordered_set & set); 269 | 270 | 271 | void resolveLockAsync(Backoffer & bo, LockPtr lock, TxnStatus & status); 272 | 273 | 274 | // Resolve locks in a region with pre-grouped keys. Only used by resolveLockAsync. 275 | void resolveRegionLocks(Backoffer & bo, LockPtr lock, RegionVerID region_id, std::vector & keys, TxnStatus & status); 276 | 277 | 278 | AsyncResolveDataPtr checkAllSecondaries(Backoffer & bo, LockPtr lock, TxnStatus & status); 279 | 280 | 281 | void checkSecondaries( 282 | Backoffer & bo, 283 | uint64_t txn_id, 284 | const std::vector & cur_keys, 285 | RegionVerID cur_region_id, 286 | AsyncResolveDataPtr shared_data); 287 | 288 | 289 | TxnStatus getTxnStatusFromLock(Backoffer & bo, LockPtr lock, uint64_t caller_start_ts, bool force_sync_commit); 290 | 291 | 292 | TxnStatus getTxnStatus(Backoffer & bo, uint64_t txn_id, const std::string & primary, uint64_t caller_start_ts, uint64_t current_ts, bool rollback_if_not_exists, bool force_sync_commit, bool is_txn_file); 293 | 294 | Cluster * cluster; 295 | std::shared_mutex mu; 296 | std::unordered_map resolved; 297 | std::queue cached; 298 | 299 | Logger * log; 300 | }; 301 | 302 | using LockResolverPtr = std::unique_ptr; 303 | 304 | } // namespace kv 305 | } // namespace pingcap 306 | --------------------------------------------------------------------------------