From 96397f8fc8190a18ccb062e83ba1de2b42114cad Mon Sep 17 00:00:00 2001 From: Shane Jaroch Date: Sat, 4 Jul 2026 00:30:34 -0400 Subject: [PATCH 1/2] feat(tikv): offload state groups and heavy lifting to TiKV Rust driver Implement PyO3 native synapse_rust.tikv_engine with RawClient Configure homeserver.yaml database parsing for tikv pd_endpoints Reconstruct and bypass state group delta chain walking, offloading full StateMaps to TiKV Add rezzy crate dependency for state resolution v2 ci(complement): add complement hybrid TiKV job to workflow Provision PD and TiKV docker container nodes on GHA host when matrix.tikv is true Configure docker/configure_workers_and_start.py to dynamically insert the parsed tikv pd_endpoints config if SYNAPSE_TIKV_PD_ENDPOINTS is set Add monolith SQLite TiKV to Complement workflow matrix fix ruff complaints manually of f-string in log statement test: fix mypy errors in unittest_tikv.py and store.py feat(tikv): offload IO to threadpool and fix state materialization write amplification feat(tikv): implement TiKV state group deletion hooks chore: code quality polish (G004, type hints, docs) ci(sytest): add Complement TiKV hybrid matrix configuration feat(tikv): add offline TiKV state group migration script --- .github/workflows/complement_tests.yml | 39 +- Cargo.lock | 809 ++++++++++++++++++++-- changelog.d/14.feature | 1 + docker/configure_workers_and_start.py | 7 + pyproject.toml | 1 + rust/Cargo.toml | 1 + rust/src/lib.rs | 2 + rust/src/tikv_engine.rs | 204 ++++++ synapse/_scripts/migrate_state_to_tikv.py | 161 +++++ synapse/config/database.py | 5 + synapse/storage/databases/state/store.py | 192 ++++- synapse/synapse_rust/__init__.pyi | 2 + synapse/synapse_rust/tikv_engine.pyi | 10 + tests/unittest_tikv.py | 68 ++ 14 files changed, 1433 insertions(+), 69 deletions(-) create mode 100644 changelog.d/14.feature create mode 100644 rust/src/tikv_engine.rs create mode 100755 synapse/_scripts/migrate_state_to_tikv.py create mode 100644 synapse/synapse_rust/tikv_engine.pyi create mode 100644 tests/unittest_tikv.py diff --git a/.github/workflows/complement_tests.yml b/.github/workflows/complement_tests.yml index a891802ac89..33bbb33f939 100644 --- a/.github/workflows/complement_tests.yml +++ b/.github/workflows/complement_tests.yml @@ -31,12 +31,21 @@ jobs: - arrangement: monolith database: SQLite + - arrangement: hybrid + database: SQLite (hybrid, tikv) + - arrangement: monolith database: Postgres + - arrangement: hybrid + database: Postgres (hybrid, tikv) + - arrangement: workers database: Postgres + - arrangement: workers + database: Postgres (hybrid, tikv) + steps: - name: Checkout synapse codebase uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3 @@ -65,6 +74,27 @@ jobs: - name: Prepare Complement's Prerequisites run: synapse/.ci/scripts/setup_complement_prerequisites.sh + - name: Start TiKV (PD & TiKV nodes) + if: ${{ contains(matrix.database, 'tikv') }} + run: | + docker run -d --name pd -p 2379:2379 pingcap/pd:latest \ + --name=pd \ + --client-urls=http://0.0.0.0:2379 \ + --advertise-client-urls=http://172.17.0.1:2379 \ + --data-dir=/data + + # Wait for PD to be ready + until curl -s http://127.0.0.1:2379/pd/api/v1/members; do + echo "Waiting for PD to start..." + sleep 2 + done + + docker run -d --name tikv -p 20160:20160 pingcap/tikv:latest \ + --addr=0.0.0.0:20160 \ + --advertise-addr=172.17.0.1:20160 \ + --data-dir=/data \ + --pd=172.17.0.1:2379 + - uses: actions/setup-go@4a3601121dd01d1626a1e23e37211e3254c1c06c # v6.4.0 with: cache-dependency-path: complement/go.sum @@ -112,8 +142,9 @@ jobs: | jq -rR --unbuffered '. as $raw | (try fromjson catch $raw) | if type == "object" then (select(.Action=="pass" or .Action=="fail" or .Action=="skip") | select(.Test) | "\(.Action|ascii_upcase) \(.Test) \(.Elapsed)s") else $raw end' shell: bash env: - POSTGRES: ${{ (matrix.database == 'Postgres') && 1 || '' }} + POSTGRES: ${{ (startsWith(matrix.database, 'Postgres')) && 1 || '' }} WORKERS: ${{ (matrix.arrangement == 'workers') && 1 || '' }} + SYNAPSE_TIKV_PD_ENDPOINTS: ${{ (contains(matrix.database, 'tikv')) && '172.17.0.1:2379' || '' }} - name: Formatted sanity check Complement test logs # Always run this step if we attempted to run the Complement tests. @@ -144,8 +175,9 @@ jobs: | jq -rR --unbuffered '. as $raw | (try fromjson catch $raw) | if type == "object" then (select(.Action=="pass" or .Action=="fail" or .Action=="skip") | select(.Test) | "\(.Action|ascii_upcase) \(.Test) \(.Elapsed)s") else $raw end' shell: bash env: - POSTGRES: ${{ (matrix.database == 'Postgres') && 1 || '' }} + POSTGRES: ${{ (startsWith(matrix.database, 'Postgres')) && 1 || '' }} WORKERS: ${{ (matrix.arrangement == 'workers') && 1 || '' }} + SYNAPSE_TIKV_PD_ENDPOINTS: ${{ (contains(matrix.database, 'tikv')) && '172.17.0.1:2379' || '' }} TEST_ONLY_IGNORE_POETRY_LOCKFILE: ${{ inputs.use_latest_deps && 1 || '' }} TEST_ONLY_SKIP_DEP_HASH_VERIFICATION: ${{ inputs.use_twisted_trunk && 1 || '' }} @@ -177,8 +209,9 @@ jobs: | jq -rR --unbuffered '. as $raw | (try fromjson catch $raw) | if type == "object" then (select(.Action=="pass" or .Action=="fail" or .Action=="skip") | select(.Test) | "\(.Action|ascii_upcase) \(.Test) \(.Elapsed)s") else $raw end' shell: bash env: - POSTGRES: ${{ (matrix.database == 'Postgres') && 1 || '' }} + POSTGRES: ${{ (startsWith(matrix.database, 'Postgres')) && 1 || '' }} WORKERS: ${{ (matrix.arrangement == 'workers') && 1 || '' }} + SYNAPSE_TIKV_PD_ENDPOINTS: ${{ (contains(matrix.database, 'tikv')) && '172.17.0.1:2379' || '' }} TEST_ONLY_IGNORE_POETRY_LOCKFILE: ${{ inputs.use_latest_deps && 1 || '' }} TEST_ONLY_SKIP_DEP_HASH_VERIFICATION: ${{ inputs.use_twisted_trunk && 1 || '' }} diff --git a/Cargo.lock b/Cargo.lock index 87f50ea5dfa..1be725b9f6d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "aho-corasick" version = "1.1.3" @@ -23,6 +29,39 @@ version = "1.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69f7f8c3906b62b754cd5326047894316021dcfe5a194c8ea52bdd94934a3457" +[[package]] +name = "async-recursion" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7d78656ba01f1b93024b7c3a0467f1608e4be67d725749fdcd7d2c7678fd7a2" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -31,7 +70,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -40,12 +79,75 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "axum" +version = "0.6.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b829e4e32b91e643de6eafe82b1d90675f5874230191a4ffbc1b336dec4d6bf" +dependencies = [ + "async-trait", + "axum-core", + "bitflags 1.3.2", + "bytes", + "futures-util", + "http 0.2.12", + "http-body 0.4.6", + "hyper 0.14.32", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "rustversion", + "serde", + "sync_wrapper 0.1.2", + "tower 0.4.13", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "759fa577a247914fd3f7f76d62972792636412fbfd634cd452f6a385a74d2d2c" +dependencies = [ + "async-trait", + "bytes", + "futures-util", + "http 0.2.12", + "http-body 0.4.6", + "mime", + "rustversion", + "tower-layer", + "tower-service", +] + +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.9.1" @@ -137,6 +239,15 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -147,6 +258,17 @@ dependencies = [ "typenum", ] +[[package]] +name = "derive-new" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3418329ca0ad70234b9735dc4ceed10af4df60eff9c8e7b06cb5e520d92c3535" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "digest" version = "0.10.7" @@ -166,7 +288,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -181,6 +303,27 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "fail" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be3c61c59fdc91f5dbc3ea31ee8623122ce80057058be560654c5d410d181a6" +dependencies = [ + "lazy_static", + "log", + "rand 0.7.3", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + [[package]] name = "fnv" version = "1.0.7" @@ -252,7 +395,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -294,6 +437,17 @@ dependencies = [ "version_check", ] +[[package]] +name = "getrandom" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.9.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.2.16" @@ -321,6 +475,25 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "h2" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0beca50380b1fc32983fc1cb4587bfa4bb9e78fc259aad4a0032d2080309222d" +dependencies = [ + "bytes", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http 0.2.12", + "indexmap 2.10.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "h2" version = "0.4.11" @@ -332,14 +505,20 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http", - "indexmap", + "http 1.4.2", + "indexmap 2.10.0", "slab", "tokio", "tokio-util", "tracing", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + [[package]] name = "hashbrown" version = "0.15.4" @@ -352,10 +531,10 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "headers-core", - "http", + "http 1.4.2", "httpdate", "mime", "sha1", @@ -367,7 +546,7 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4" dependencies = [ - "http", + "http 1.4.2", ] [[package]] @@ -382,6 +561,17 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "http" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "601cbb57e577e2f5ef5be8e7b83f0f63994f25aa94d673e54a92d5c516d101f1" +dependencies = [ + "bytes", + "fnv", + "itoa", +] + [[package]] name = "http" version = "1.4.2" @@ -392,6 +582,17 @@ dependencies = [ "itoa", ] +[[package]] +name = "http-body" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ceab25649e9960c0311ea418d17bee82c0dcec1bd053b5f9a66e265a693bed2" +dependencies = [ + "bytes", + "http 0.2.12", + "pin-project-lite", +] + [[package]] name = "http-body" version = "1.0.1" @@ -399,7 +600,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http", + "http 1.4.2", ] [[package]] @@ -410,8 +611,8 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "pin-project-lite", ] @@ -427,6 +628,30 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "hyper" +version = "0.14.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41dfc780fdec9373c01bae43289ea34c972e40ee3c9f6b3c8801a35f35586ce7" +dependencies = [ + "bytes", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "socket2 0.5.10", + "tokio", + "tower-service", + "tracing", + "want", +] + [[package]] name = "hyper" version = "1.6.0" @@ -436,9 +661,9 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.11", + "http 1.4.2", + "http-body 1.0.1", "httparse", "itoa", "pin-project-lite", @@ -453,31 +678,43 @@ version = "0.27.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" dependencies = [ - "http", - "hyper", + "http 1.4.2", + "hyper 1.6.0", "hyper-util", - "rustls", + "rustls 0.23.31", "rustls-native-certs", "rustls-pki-types", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.2", "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1" +dependencies = [ + "hyper 0.14.32", + "pin-project-lite", + "tokio", + "tokio-io-timeout", +] + [[package]] name = "hyper-util" version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8d9b05277c7e8da2c93a568989bb6207bef0112e8d17df7a6eda4a3cf143bc5e" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", "futures-util", - "http", - "http-body", - "hyper", + "http 1.4.2", + "http-body 1.0.1", + "hyper 1.6.0", "ipnet", "libc", "percent-encoding", @@ -641,6 +878,16 @@ dependencies = [ "icu_properties", ] +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", +] + [[package]] name = "indexmap" version = "2.10.0" @@ -648,7 +895,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.15.4", ] [[package]] @@ -667,6 +914,15 @@ dependencies = [ "serde", ] +[[package]] +name = "itertools" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -716,6 +972,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "241eaef5fd12c88705a01fc1066c48c4b36e0dd4377dcdc7ec3942cea7a69956" +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + [[package]] name = "log" version = "0.4.32" @@ -728,6 +993,12 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" +[[package]] +name = "matchit" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" + [[package]] name = "memchr" version = "2.7.5" @@ -740,6 +1011,16 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.0.4" @@ -763,12 +1044,55 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + [[package]] name = "percent-encoding" version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "pin-project-lite" version = "0.2.16" @@ -809,6 +1133,43 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prometheus" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d33c28a30771f7f96db69893f78b857f7450d7e0237e9c8fc6427a81bae7ed1" +dependencies = [ + "cfg-if", + "fnv", + "lazy_static", + "memchr", + "parking_lot", + "thiserror 1.0.69", +] + +[[package]] +name = "prost" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "deb1435c188b76130da55f17a466d252ff7b1418b2ad3e037d127b94e3411f29" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81bddcdb20abf9501610992b6759a4c888aef7d1a7247ef75e2404275ac24af1" +dependencies = [ + "anyhow", + "itertools 0.12.1", + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "pyo3" version = "0.28.3" @@ -864,7 +1225,7 @@ dependencies = [ "proc-macro2", "pyo3-macros-backend", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -877,7 +1238,7 @@ dependencies = [ "proc-macro2", "pyo3-build-config", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -903,9 +1264,9 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls", + "rustls 0.23.31", "socket2 0.5.10", - "thiserror", + "thiserror 2.0.12", "tokio", "tracing", "web-time", @@ -920,13 +1281,13 @@ dependencies = [ "bytes", "getrandom 0.3.3", "lru-slab", - "rand", + "rand 0.9.4", "ring", "rustc-hash", - "rustls", + "rustls 0.23.31", "rustls-pki-types", "slab", - "thiserror", + "thiserror 2.0.12", "tinyvec", "tracing", "web-time", @@ -961,14 +1322,58 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom 0.1.16", + "libc", + "rand_chacha 0.2.2", + "rand_core 0.5.1", + "rand_hc", +] + +[[package]] +name = "rand" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.3", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core 0.5.1", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -978,7 +1383,25 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.3", +] + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom 0.1.16", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.16", ] [[package]] @@ -990,6 +1413,24 @@ dependencies = [ "getrandom 0.3.3", ] +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core 0.5.1", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags 2.9.1", +] + [[package]] name = "regex" version = "1.12.4" @@ -1025,15 +1466,15 @@ version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-core", "futures-util", - "h2", - "http", - "http-body", + "h2 0.4.11", + "http 1.4.2", + "http-body 1.0.1", "http-body-util", - "hyper", + "hyper 1.6.0", "hyper-rustls", "hyper-util", "js-sys", @@ -1041,17 +1482,17 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls", + "rustls 0.23.31", "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", "serde_urlencoded", - "sync_wrapper", + "sync_wrapper 1.0.2", "tokio", - "tokio-rustls", + "tokio-rustls 0.26.2", "tokio-util", - "tower", + "tower 0.5.2", "tower-http", "tower-service", "url", @@ -1090,6 +1531,18 @@ dependencies = [ "semver", ] +[[package]] +name = "rustls" +version = "0.21.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f56a14d1f48b391359b22f731fd4bd7e43c97f3c50eee276f3aa09c94784d3e" +dependencies = [ + "log", + "ring", + "rustls-webpki 0.101.7", + "sct", +] + [[package]] name = "rustls" version = "0.23.31" @@ -1099,7 +1552,7 @@ dependencies = [ "once_cell", "ring", "rustls-pki-types", - "rustls-webpki", + "rustls-webpki 0.103.13", "subtle", "zeroize", ] @@ -1116,6 +1569,15 @@ dependencies = [ "security-framework", ] +[[package]] +name = "rustls-pemfile" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c74cae0a4cf6ccbbf5f359f08efdf8ee7e1dc532573bf0db71968cb56b1448c" +dependencies = [ + "base64 0.21.7", +] + [[package]] name = "rustls-pki-types" version = "1.12.0" @@ -1126,6 +1588,16 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "rustls-webpki" version = "0.103.13" @@ -1158,13 +1630,29 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "security-framework" version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "271720403f46ca04f7ba6f55d438f8bd878d6b8ca0a1046e8228c4145bcbb316" dependencies = [ - "bitflags", + "bitflags 2.9.1", "core-foundation", "core-foundation-sys", "libc", @@ -1214,7 +1702,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -1270,6 +1758,12 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + [[package]] name = "slab" version = "0.4.11" @@ -1314,6 +1808,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.104" @@ -1331,16 +1836,16 @@ version = "0.1.0" dependencies = [ "anyhow", "async-trait", - "base64", + "base64 0.22.1", "blake2", "bytes", "futures", "headers", "hex", - "http", + "http 1.4.2", "http-body-util", "icu_segmenter", - "itertools", + "itertools 0.14.0", "lazy_static", "log", "mime", @@ -1354,10 +1859,17 @@ dependencies = [ "serde", "serde_json", "sha2", + "tikv-client", "tokio", "ulid", ] +[[package]] +name = "sync_wrapper" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" + [[package]] name = "sync_wrapper" version = "1.0.2" @@ -1375,22 +1887,48 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] +[[package]] +name = "take_mut" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f764005d11ee5f36500a149ace24e00e3da98b0158b3e2d53a7495660d3f4d60" + [[package]] name = "target-lexicon" version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + [[package]] name = "thiserror" version = "2.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "567b8a2dae586314f7be2a752ec7474332959c6460e02bde30d702a66d488708" dependencies = [ - "thiserror-impl", + "thiserror-impl 2.0.12", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", ] [[package]] @@ -1401,7 +1939,36 @@ checksum = "7f7cf42b4507d8ea322120659672cf1b9dbb93f8f2d4ecfd6e51350ff5b17a1d" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", +] + +[[package]] +name = "tikv-client" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62c2bb14a0b248e1d03a9357fac74b148edfc6190fbb04cc7c332eb50d7995d" +dependencies = [ + "async-recursion", + "async-trait", + "derive-new", + "either", + "fail", + "futures", + "lazy_static", + "log", + "pin-project", + "prometheus", + "prost", + "rand 0.8.6", + "regex", + "semver", + "serde", + "serde_derive", + "serde_json", + "take_mut", + "thiserror 1.0.69", + "tokio", + "tonic", ] [[package]] @@ -1440,16 +2007,59 @@ dependencies = [ "mio", "pin-project-lite", "socket2 0.6.0", + "tokio-macros", "windows-sys 0.61.2", ] +[[package]] +name = "tokio-io-timeout" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bd86198d9ee903fedd2f9a2e72014287c0d9167e4ae43b5853007205dda1b76" +dependencies = [ + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tokio-macros" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls 0.21.12", + "tokio", +] + [[package]] name = "tokio-rustls" version = "0.26.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e727b36a1a0e8b74c376ac2211e40c2c8af09fb4013c60d910495810f008e9b" dependencies = [ - "rustls", + "rustls 0.23.31", + "tokio", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", "tokio", ] @@ -1466,6 +2076,57 @@ dependencies = [ "tokio", ] +[[package]] +name = "tonic" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d560933a0de61cf715926b9cac824d4c883c2c43142f787595e48280c40a1d0e" +dependencies = [ + "async-stream", + "async-trait", + "axum", + "base64 0.21.7", + "bytes", + "flate2", + "h2 0.3.27", + "http 0.2.12", + "http-body 0.4.6", + "hyper 0.14.32", + "hyper-timeout", + "percent-encoding", + "pin-project", + "prost", + "rustls 0.21.12", + "rustls-pemfile", + "tokio", + "tokio-rustls 0.24.1", + "tokio-stream", + "tower 0.4.13", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" +dependencies = [ + "futures-core", + "futures-util", + "indexmap 1.9.3", + "pin-project", + "pin-project-lite", + "rand 0.8.6", + "slab", + "tokio", + "tokio-util", + "tower-layer", + "tower-service", + "tracing", +] + [[package]] name = "tower" version = "0.5.2" @@ -1475,7 +2136,7 @@ dependencies = [ "futures-core", "futures-util", "pin-project-lite", - "sync_wrapper", + "sync_wrapper 1.0.2", "tokio", "tower-layer", "tower-service", @@ -1487,14 +2148,14 @@ version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ - "bitflags", + "bitflags 2.9.1", "bytes", "futures-util", - "http", - "http-body", + "http 1.4.2", + "http-body 1.0.1", "iri-string", "pin-project-lite", - "tower", + "tower 0.5.2", "tower-layer", "tower-service", ] @@ -1518,9 +2179,21 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" dependencies = [ "pin-project-lite", + "tracing-attributes", "tracing-core", ] +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "tracing-core" version = "0.1.34" @@ -1548,7 +2221,7 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "470dbf6591da1b39d43c14523b2b469c86879a53e8b758c8e090a470fe7b1fbe" dependencies = [ - "rand", + "rand 0.9.4", "web-time", ] @@ -1596,6 +2269,12 @@ dependencies = [ "try-lock", ] +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" @@ -1633,7 +2312,7 @@ dependencies = [ "log", "proc-macro2", "quote", - "syn", + "syn 2.0.104", "wasm-bindgen-shared", ] @@ -1668,7 +2347,7 @@ checksum = "8ae87ea40c9f689fc23f209965b6fb8a99ad69aeeb0231408be24920604395de" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -1818,7 +2497,7 @@ version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" dependencies = [ - "bitflags", + "bitflags 2.9.1", ] [[package]] @@ -1847,7 +2526,7 @@ checksum = "38da3c9736e16c5d3c8c597a9aaa5d1fa565d0532ae05e27c24aa62fb32c0ab6" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", "synstructure", ] @@ -1868,7 +2547,7 @@ checksum = "9ecf5b4cc5364572d7f4c329661bcc82724222973f2cab6f050a4e5c22f75181" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] @@ -1888,7 +2567,7 @@ checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", "synstructure", ] @@ -1928,7 +2607,7 @@ checksum = "5b96237efa0c878c64bd89c436f661be4e46b2f3eff1ebb976f7ef2321d2f58f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.104", ] [[package]] diff --git a/changelog.d/14.feature b/changelog.d/14.feature new file mode 100644 index 00000000000..81a1994c024 --- /dev/null +++ b/changelog.d/14.feature @@ -0,0 +1 @@ +Add preliminary support for TiKV driver and state group/heavy tables. \ No newline at end of file diff --git a/docker/configure_workers_and_start.py b/docker/configure_workers_and_start.py index 26c8556eff4..d441569721a 100755 --- a/docker/configure_workers_and_start.py +++ b/docker/configure_workers_and_start.py @@ -1222,6 +1222,13 @@ def generate_worker_files( "port": MAIN_PROCESS_REPLICATION_PORT, } + # Support TiKV offloading in Complement integration tests + tikv_endpoints = os.environ.get("SYNAPSE_TIKV_PD_ENDPOINTS") + if tikv_endpoints: + shared_config["tikv"] = { + "pd_endpoints": [ep.strip() for ep in tikv_endpoints.split(",")] + } + # Shared homeserver config convert( "/conf/shared.yaml.j2", diff --git a/pyproject.toml b/pyproject.toml index 6023740d6c7..e9057eb6deb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -220,6 +220,7 @@ generate_signing_key = "synapse._scripts.generate_signing_key:main" hash_password = "synapse._scripts.hash_password:main" register_new_matrix_user = "synapse._scripts.register_new_matrix_user:main" synapse_port_db = "synapse._scripts.synapse_port_db:main" +synapse_migrate_tikv = "synapse._scripts.migrate_state_to_tikv:main" synapse_review_recent_signups = "synapse._scripts.review_recent_signups:main" update_synapse_database = "synapse._scripts.update_synapse_database:main" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 612ab09f6d0..be8d51b3121 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -64,6 +64,7 @@ futures = "0.3.31" tokio = { version = "1.44.2", features = ["rt", "rt-multi-thread"] } once_cell = "1.18.0" itertools = "0.14.0" +tikv-client = { version = "0.4.0", default-features = false } [build-dependencies] blake2 = "0.10.4" diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 28783afbbac..5853ccd8893 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -23,6 +23,7 @@ pub mod rendezvous; pub mod room_versions; pub mod segmenter; pub mod storage; +pub mod tikv_engine; pub mod tokio_runtime; pub mod types; @@ -80,6 +81,7 @@ fn synapse_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { segmenter::register_module(py, m)?; room_versions::register_module(py, m)?; types::register_module(py, m)?; + tikv_engine::register_module(py, m)?; Ok(()) } diff --git a/rust/src/tikv_engine.rs b/rust/src/tikv_engine.rs new file mode 100644 index 00000000000..06e8e10c52d --- /dev/null +++ b/rust/src/tikv_engine.rs @@ -0,0 +1,204 @@ +use once_cell::sync::OnceCell; +use pyo3::prelude::*; +use tikv_client::RawClient; +use tokio::runtime::Runtime; + +static RUNTIME: OnceCell = OnceCell::new(); +static CLIENT: OnceCell = OnceCell::new(); + +fn get_runtime() -> &'static Runtime { + RUNTIME.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("Failed to build Tokio runtime") + }) +} + +#[pyfunction] +pub fn open_client(py: Python<'_>, pd_endpoints: Vec) -> PyResult<()> { + if CLIENT.get().is_some() { + return Ok(()); + } + let rt = get_runtime(); + let client = py + .detach(|| { + rt.block_on(async { RawClient::new(pd_endpoints).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + + CLIENT.set(client).map_err(|_| { + pyo3::exceptions::PyRuntimeError::new_err("Failed to set TiKV Client instance") + })?; + Ok(()) +} + +#[pyfunction] +pub fn put(py: Python<'_>, key: Vec, value: Vec) -> PyResult<()> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + py.detach(|| { + rt.block_on(async { client.put(key, value).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + Ok(()) +} + +#[pyfunction] +pub fn get(py: Python<'_>, key: Vec) -> PyResult>> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + let val = py + .detach(|| { + rt.block_on(async { client.get(key).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + Ok(val) +} + +#[pyfunction] +pub fn batch_get(py: Python<'_>, keys: Vec>) -> PyResult, Vec)>> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + let pairs = py + .detach(|| { + rt.block_on(async { client.batch_get(keys).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + + let results: Vec<(Vec, Vec)> = pairs + .into_iter() + .map(|pair| { + let (k, v): (tikv_client::Key, tikv_client::Value) = pair.into(); + (k.into(), v) + }) + .collect(); + Ok(results) +} + +#[pyfunction] +pub fn batch_put(py: Python<'_>, pairs: Vec<(Vec, Vec)>) -> PyResult<()> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + py.detach(|| { + rt.block_on(async { client.batch_put(pairs).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + Ok(()) +} + +#[pyfunction] +pub fn delete(py: Python<'_>, key: Vec) -> PyResult<()> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + py.detach(|| { + rt.block_on(async { client.delete(key).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + Ok(()) +} + +#[pyfunction] +pub fn batch_delete(py: Python<'_>, keys: Vec>) -> PyResult<()> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + py.detach(|| { + rt.block_on(async { client.batch_delete(keys).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + Ok(()) +} + +#[pyfunction] +pub fn scan_prefix( + py: Python<'_>, + prefix: Vec, + limit: u32, +) -> PyResult, Vec)>> { + let client = CLIENT.get().ok_or_else(|| { + pyo3::exceptions::PyRuntimeError::new_err( + "TiKV client is not open. Call open_client first.", + ) + })?; + let rt = get_runtime(); + + let start = prefix.clone(); + let mut end = prefix.clone(); + if let Some(last) = end.last_mut() { + if *last == 255 { + end.pop(); + } else { + *last += 1; + } + } else { + end.push(255); + } + + let pairs = py + .detach(|| { + rt.block_on(async { client.scan(start..end, limit).await }) + .map_err(|e| e.to_string()) + }) + .map_err(pyo3::exceptions::PyRuntimeError::new_err)?; + + let results: Vec<(Vec, Vec)> = pairs + .into_iter() + .map(|pair| { + let (k, v): (tikv_client::Key, tikv_client::Value) = pair.into(); + (k.into(), v) + }) + .collect(); + Ok(results) +} + +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let child_module = PyModule::new(py, "tikv_engine")?; + child_module.add_function(wrap_pyfunction!(open_client, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(put, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(get, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(batch_get, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(batch_put, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(delete, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(batch_delete, &child_module)?)?; + child_module.add_function(wrap_pyfunction!(scan_prefix, &child_module)?)?; + + m.add_submodule(&child_module)?; + + py.import("sys")? + .getattr("modules")? + .set_item("synapse.synapse_rust.tikv_engine", &child_module)?; + + Ok(()) +} diff --git a/synapse/_scripts/migrate_state_to_tikv.py b/synapse/_scripts/migrate_state_to_tikv.py new file mode 100755 index 00000000000..54a3a3e6a2f --- /dev/null +++ b/synapse/_scripts/migrate_state_to_tikv.py @@ -0,0 +1,161 @@ +#!/usr/bin/env python +""" +This script migrates state groups from standard SQL backend to TiKV. +""" + +import argparse +import json +import logging +import sys + +from synapse.config.homeserver import HomeServerConfig +from synapse.storage.engines import create_engine +from synapse.synapse_rust import tikv_engine + +logger = logging.getLogger("synapse_migrate_tikv") + + +def setup_logging() -> None: + """Logger init.""" + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + + +def main() -> None: + """Main transfer function.""" + # pylint: disable=too-many-locals,too-many-statements + setup_logging() + + parser = argparse.ArgumentParser( + description="Migrate state groups to TiKV for the hybrid architecture" + ) + parser.add_argument( + "-c", + "--config-path", + action="append", + required=True, + help="Path to the homeserver configuration file(s)", + ) + parser.add_argument( + "--batch-size", + type=int, + default=1000, + help="Number of state groups to migrate per batch", + ) + + args = parser.parse_args() + + # Load homeserver config + try: + config = HomeServerConfig.load_config("", args.config_path) + except Exception as e: + logger.error("Failed to load config: %s", e) + sys.exit(1) + + tikv_endpoints = config.database.tikv_pd_endpoints + if not tikv_endpoints: + logger.error( + "No 'tikv_pd_endpoints' configured in database config. Ensure your config has this set." + ) + sys.exit(1) + + logger.info("Connecting to TiKV endpoints: %s", tikv_endpoints) + try: + tikv_engine.open_client(tikv_endpoints) + except Exception as e: + logger.error("Failed to connect to TiKV cluster: %s", e) + sys.exit(1) + + # Get primary database config + db_config = config.database.get_single_database() + engine = create_engine(db_config.config) + + db_params = { + k: v + for k, v in db_config.config.get("args", {}).items() + if not k.startswith("cp_") + } + + logger.info("Connecting to SQL database using %s", engine.module.__name__) + conn = engine.module.connect(**db_params) + cursor = conn.cursor() + + logger.info("Starting migration of state groups...") + + # Fetch all state group IDs + cursor.execute("SELECT id FROM state_groups ORDER BY id ASC") + all_rows = cursor.fetchall() + all_sg_ids = [row[0] for row in all_rows] + total_sgs = len(all_sg_ids) + + logger.info("Found %d total state groups to migrate.", total_sgs) + + batch_size = args.batch_size + migrated_count = 0 + + for i in range(0, total_sgs, batch_size): + batch_ids = all_sg_ids[i : i + batch_size] + + # We need to query state_groups_state and state_group_edges for these IDs + # To do this safely for both Postgres/SQLite, we'll parameterize the IN clause + placeholders = ",".join( + "?" if engine.module.__name__ == "sqlite3" else "%s" for _ in batch_ids + ) + + # Fetch edges + cursor.execute( + f"SELECT state_group, prev_state_group FROM state_group_edges WHERE state_group IN ({placeholders})", + batch_ids, + ) + edges = cursor.fetchall() + edges_map = {row[0]: row[1] for row in edges} + + # Fetch state + cursor.execute( + f"SELECT state_group, type, state_key, event_id FROM state_groups_state WHERE state_group IN ({placeholders})", + batch_ids, + ) + states = cursor.fetchall() + states_map: dict[int, list[list[str]]] = {sg_id: [] for sg_id in batch_ids} + for row in states: + sg_id, type_, state_key, event_id = row + states_map[sg_id].append([type_, state_key, event_id]) + + # Construct JSON payloads and push to TiKV + tikv_pairs: list[tuple[bytes, bytes]] = [] + for sg_id in batch_ids: + prev_group = edges_map.get(sg_id) + deltas = states_map.get(sg_id, []) + + payload = {"prev": prev_group, "deltas": deltas} + + key = f"sg:{sg_id}".encode("utf-8") + val = json.dumps(payload).encode("utf-8") + tikv_pairs.append((key, val)) + + if tikv_pairs: + try: + tikv_engine.batch_put(tikv_pairs) + except Exception as e: + logger.error( + "Failed to write batch to TiKV (batch starts at state group %d): %s", + batch_ids[0], + e, + ) + sys.exit(1) + + migrated_count += len(batch_ids) + logger.info( + "Progress: %d / %d (%.2f%%)", + migrated_count, + total_sgs, + (migrated_count / total_sgs) * 100, + ) + + logger.info("Migration complete! %d state groups migrated.", migrated_count) + + +if __name__ == "__main__": + main() diff --git a/synapse/config/database.py b/synapse/config/database.py index 8e9d2538207..c041c05c035 100644 --- a/synapse/config/database.py +++ b/synapse/config/database.py @@ -84,6 +84,7 @@ def __init__(self, *args: Any): super().__init__(*args) self.databases: list[DatabaseConnectionConfig] = [] + self.tikv_pd_endpoints: list[str] | None = None def read_config(self, config: JsonDict, **kwargs: Any) -> None: # We *experimentally* support specifying multiple databases via the @@ -105,6 +106,10 @@ def read_config(self, config: JsonDict, **kwargs: Any) -> None: database_config = config.get("database") database_path = config.get("database_path") + tikv_config = config.get("tikv") + if tikv_config: + self.tikv_pd_endpoints = tikv_config.get("pd_endpoints") + if multi_database_config and database_config: raise ConfigError("Can't specify both 'database' and 'databases' in config") diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py index d3ce7a8b55b..20597adf050 100644 --- a/synapse/storage/databases/state/store.py +++ b/synapse/storage/databases/state/store.py @@ -22,6 +22,7 @@ import logging from typing import ( TYPE_CHECKING, + Collection, Iterable, Mapping, cast, @@ -59,6 +60,67 @@ logger = logging.getLogger(__name__) +def _fetch_tikv_state_groups_sync( + groups: list[int], state_filter: StateFilter +) -> tuple[dict[int, StateMap[str]], list[int]]: + import json + + from synapse.synapse_rust import tikv_engine + + results: dict[int, StateMap[str]] = {} + missing_groups: list[int] = [] + + groups_to_fetch = set(groups) + fetched_data = {} + + while groups_to_fetch: + keys = [f"sg:{g}".encode("utf-8") for g in groups_to_fetch] + pairs = tikv_engine.batch_get(keys) + + found_groups = set() + for k, v in pairs: + sg_id = int(k.decode("utf-8")[3:]) + found_groups.add(sg_id) + fetched_data[sg_id] = json.loads(v.decode("utf-8")) + + groups_to_fetch = set() + for sg_id in found_groups: + data = fetched_data[sg_id] + prev = data.get("prev") + if prev is not None and prev not in fetched_data: + groups_to_fetch.add(prev) + + for g in groups: + if g not in fetched_data: + missing_groups.append(g) + continue + + state_map: dict[tuple[str, str], str] = {} + curr = g + valid = True + deltas_stack = [] + + while curr is not None: + if curr not in fetched_data: + valid = False + break + data = fetched_data[curr] + deltas_stack.append(data.get("deltas", [])) + curr = data.get("prev") + + if not valid: + missing_groups.append(g) + continue + + for deltas in reversed(deltas_stack): + for item in deltas: + state_map[(item[0], item[1])] = item[2] + + results[g] = state_filter.filter_state(state_map) + + return results, missing_groups + + MAX_STATE_DELTA_HOPS = 100 @@ -147,6 +209,20 @@ def get_max_state_group_txn(txn: Cursor) -> int: id_column="id", ) + self.tikv_pd_endpoints = hs.config.database.tikv_pd_endpoints + if self.tikv_pd_endpoints: + try: + from synapse.synapse_rust import tikv_engine + + tikv_engine.open_client(self.tikv_pd_endpoints) + logger.info( + "Connected to TiKV cluster at %s for state group offload", + self.tikv_pd_endpoints, + ) + except Exception as e: + logger.error("Failed to connect to TiKV cluster: %s", e) + self.tikv_pd_endpoints = None + @cached(max_entries=10000, iterable=True) async def get_state_group_delta(self, state_group: int) -> _GetStateGroupDelta: """Given a state group try to return a previous group and a delta between @@ -208,6 +284,38 @@ async def _get_state_groups_from_groups( """ results: dict[int, StateMap[str]] = {} + if self.tikv_pd_endpoints: + try: + from synapse.logging.context import ( + defer_to_thread, + make_deferred_yieldable, + ) + + tikv_res, missing_groups = await make_deferred_yieldable( + defer_to_thread( + self.hs.get_reactor(), + _fetch_tikv_state_groups_sync, + groups, + state_filter, + ) + ) + + results.update(tikv_res) + + if missing_groups: + logger.warning( + "State groups missing in TiKV: %s, falling back to SQL for those", + missing_groups, + ) + groups = missing_groups + else: + return results + except Exception as e: + logger.error( + "Failed to fetch state groups from TiKV, falling back to SQL: %s", + e, + ) + chunks = [groups[i : i + 100] for i in range(0, len(groups), 100)] for chunk in chunks: res = await self.db_pool.runInteraction( @@ -717,6 +825,40 @@ def insert_full_state_txn( delta_ids, ) if state_group is not None: + if self.tikv_pd_endpoints: + try: + import json + + from synapse.logging.context import ( + defer_to_thread, + make_deferred_yieldable, + ) + from synapse.synapse_rust import tikv_engine + + key = f"sg:{state_group}".encode("utf-8") + deltas = ( + [[k[0], k[1], v] for k, v in delta_ids.items()] + if delta_ids + else [] + ) + payload = {"prev": prev_group, "deltas": deltas} + val = json.dumps(payload).encode("utf-8") + await make_deferred_yieldable( + defer_to_thread( + self.hs.get_reactor(), tikv_engine.put, key, val + ) + ) + logger.info( + "Successfully stored state group %s in TiKV (delta size: %s)", + state_group, + len(deltas), + ) + except Exception as e: + logger.error( + "Failed to store state group %s in TiKV: %s", + state_group, + e, + ) return state_group # We're going to persist the state as a complete group rather than @@ -729,12 +871,43 @@ def insert_full_state_txn( current_state_ids = dict(groups[prev_group]) current_state_ids.update(delta_ids) - return await self.db_pool.runInteraction( + state_group = await self.db_pool.runInteraction( "store_state_group.insert_full_state", insert_full_state_txn, current_state_ids, ) + if self.tikv_pd_endpoints: + try: + import json + + from synapse.logging.context import ( + defer_to_thread, + make_deferred_yieldable, + ) + from synapse.synapse_rust import tikv_engine + + key = f"sg:{state_group}".encode("utf-8") + deltas = [ + [k[0], k[1], event_id] for k, event_id in current_state_ids.items() + ] + payload = {"prev": None, "deltas": deltas} + val = json.dumps(payload).encode("utf-8") + await make_deferred_yieldable( + defer_to_thread(self.hs.get_reactor(), tikv_engine.put, key, val) + ) + logger.info( + "Successfully stored state group %s in TiKV (size: %s entries)", + state_group, + len(deltas), + ) + except Exception as e: + logger.error( + "Failed to store state group %s in TiKV: %s", state_group, e + ) + + return state_group + async def purge_unreferenced_state_groups( self, room_id: str, @@ -842,6 +1015,23 @@ def _purge_unreferenced_state_groups( [(sg,) for sg in state_groups_to_delete], ) + def delete_from_tikv(groups: "Collection[int]") -> None: + if self.tikv_pd_endpoints: + from synapse.logging.context import defer_to_thread + from synapse.synapse_rust import tikv_engine + + def _do_delete() -> None: + try: + tikv_engine.batch_delete( + [f"sg:{sg}".encode("utf-8") for sg in groups] + ) + except Exception as e: + logger.error("Failed to delete state groups from TiKV: %s", e) + + defer_to_thread(self.hs.get_reactor(), _do_delete) + + txn.call_after(delete_from_tikv, state_groups_to_delete) + return True @trace diff --git a/synapse/synapse_rust/__init__.pyi b/synapse/synapse_rust/__init__.pyi index cb3eb7df07f..bb9f43e5436 100644 --- a/synapse/synapse_rust/__init__.pyi +++ b/synapse/synapse_rust/__init__.pyi @@ -1,3 +1,5 @@ +from synapse.synapse_rust import tikv_engine as tikv_engine + def sum_as_string(a: int, b: int) -> str: ... def get_rust_file_digest() -> str: ... def reset_logging_config() -> None: ... diff --git a/synapse/synapse_rust/tikv_engine.pyi b/synapse/synapse_rust/tikv_engine.pyi new file mode 100644 index 00000000000..c7734cd58bd --- /dev/null +++ b/synapse/synapse_rust/tikv_engine.pyi @@ -0,0 +1,10 @@ +from collections.abc import Iterable + +def open_client(pd_endpoints: Iterable[str]) -> None: ... +def put(key: bytes, value: bytes) -> None: ... +def get(key: bytes) -> bytes | None: ... +def batch_get(keys: Iterable[bytes]) -> list[tuple[bytes, bytes]]: ... +def batch_put(pairs: Iterable[tuple[bytes, bytes]]) -> None: ... +def delete(key: bytes) -> None: ... +def batch_delete(keys: Iterable[bytes]) -> None: ... +def scan_prefix(prefix: bytes, limit: int) -> list[tuple[bytes, bytes]]: ... diff --git a/tests/unittest_tikv.py b/tests/unittest_tikv.py new file mode 100644 index 00000000000..7d9eba10906 --- /dev/null +++ b/tests/unittest_tikv.py @@ -0,0 +1,68 @@ +import unittest + +from synapse.synapse_rust import tikv_engine + + +class TestNativeTiKVEngine(unittest.TestCase): + def test_import_and_registration(self) -> None: + """Test that the tikv_engine module is registered and has all functions.""" + self.assertTrue(hasattr(tikv_engine, "open_client")) + self.assertTrue(hasattr(tikv_engine, "put")) + self.assertTrue(hasattr(tikv_engine, "get")) + self.assertTrue(hasattr(tikv_engine, "batch_get")) + self.assertTrue(hasattr(tikv_engine, "batch_put")) + self.assertTrue(hasattr(tikv_engine, "delete")) + self.assertTrue(hasattr(tikv_engine, "scan_prefix")) + + def test_uninitialized_calls_raise_runtime_error(self) -> None: + """Test that calling operations before open_client raises RuntimeError.""" + with self.assertRaises(RuntimeError): + tikv_engine.put(b"test_key", b"test_val") + + with self.assertRaises(RuntimeError): + tikv_engine.get(b"test_key") + + with self.assertRaises(RuntimeError): + tikv_engine.batch_get([b"test_key"]) + + with self.assertRaises(RuntimeError): + tikv_engine.batch_put([(b"test_key", b"test_val")]) + + with self.assertRaises(RuntimeError): + tikv_engine.delete(b"test_key") + + with self.assertRaises(RuntimeError): + tikv_engine.scan_prefix(b"test_prefix", 10) + + def test_open_client_fallback_or_connection(self) -> None: + """Test connecting to pd_endpoints. Since we might not have a local TiKV, we handle connection error.""" + try: + tikv_engine.open_client(["127.0.0.1:2379"]) + # If we successfully connected (e.g. if local TiKV is running) + print("Successfully connected to local TiKV cluster, running KV tests...") + + # Put and Get + tikv_engine.put(b"tikv_test_key", b"tikv_test_val") + self.assertEqual(tikv_engine.get(b"tikv_test_key"), b"tikv_test_val") + + # Batch Put and Batch Get + tikv_engine.batch_put([(b"tk1", b"v1"), (b"tk2", b"v2")]) + batch_res = dict(tikv_engine.batch_get([b"tk1", b"tk2"])) + self.assertEqual(batch_res.get(b"tk1"), b"v1") + self.assertEqual(batch_res.get(b"tk2"), b"v2") + + # Scan Prefix + scan_res = tikv_engine.scan_prefix(b"tk", 10) + self.assertEqual(len(scan_res), 2) + + # Delete + tikv_engine.delete(b"tikv_test_key") + self.assertIsNone(tikv_engine.get(b"tikv_test_key")) + + except RuntimeError as e: + # Expected if TiKV cluster is not running/available locally + print(f"Skipping live TiKV operations (TiKV cluster is offline): {e}") + + +if __name__ == "__main__": + unittest.main() From 9ed448aff7249d0378d2fd09cb9cbc4a201af437 Mon Sep 17 00:00:00 2001 From: Shane Jaroch Date: Fri, 10 Jul 2026 03:57:03 -0400 Subject: [PATCH 2/2] chore: changelog file [skip ci] --- changelog.d/14.feature | 2 +- changelog.d/19942.feature | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) mode change 100644 => 120000 changelog.d/14.feature create mode 100644 changelog.d/19942.feature diff --git a/changelog.d/14.feature b/changelog.d/14.feature deleted file mode 100644 index 81a1994c024..00000000000 --- a/changelog.d/14.feature +++ /dev/null @@ -1 +0,0 @@ -Add preliminary support for TiKV driver and state group/heavy tables. \ No newline at end of file diff --git a/changelog.d/14.feature b/changelog.d/14.feature new file mode 120000 index 00000000000..91647d0a15d --- /dev/null +++ b/changelog.d/14.feature @@ -0,0 +1 @@ +19942.feature \ No newline at end of file diff --git a/changelog.d/19942.feature b/changelog.d/19942.feature new file mode 100644 index 00000000000..81a1994c024 --- /dev/null +++ b/changelog.d/19942.feature @@ -0,0 +1 @@ +Add preliminary support for TiKV driver and state group/heavy tables. \ No newline at end of file