From c5db98cdcdae7d71623ccdfdc61b0b4675f8f42f Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 18 Jul 2024 15:55:41 -0600 Subject: [PATCH] Initial Rust connection pool implementation (#7478) This PR migrates the connection pooling algorithm from Python to Rust. It attempts to maintain as much of the algorithm as possible, though there will be some differences as a result of how the code is structured. We maintain the same terminology of "blocks" from the existing implementation and the spirit of the overall algorithm is the same as before. There are some unavoidable differences due to code structure and language changes, however the QoS results show the similar behaviour in all cases. --- Cargo.lock | 1296 ++++++++++++++++++++++- Cargo.toml | 1 + edb/server/conn_pool/Cargo.toml | 50 + edb/server/conn_pool/README.md | 112 ++ edb/server/conn_pool/src/algo.rs | 852 +++++++++++++++ edb/server/conn_pool/src/block.rs | 934 +++++++++++++++++ edb/server/conn_pool/src/conn.rs | 348 ++++++ edb/server/conn_pool/src/lib.rs | 27 + edb/server/conn_pool/src/metrics.rs | 373 +++++++ edb/server/conn_pool/src/pool.rs | 1398 +++++++++++++++++++++++++ edb/server/conn_pool/src/python.rs | 6 + edb/server/conn_pool/src/test.rs | 511 +++++++++ edb/server/conn_pool/src/waitqueue.rs | 104 ++ edb/server/connpool/__init__.py | 4 +- edb/server/connpool/pool2.py | 5 + setup.py | 6 + 16 files changed, 6008 insertions(+), 19 deletions(-) create mode 100644 edb/server/conn_pool/Cargo.toml create mode 100644 edb/server/conn_pool/README.md create mode 100644 edb/server/conn_pool/src/algo.rs create mode 100644 edb/server/conn_pool/src/block.rs create mode 100644 edb/server/conn_pool/src/conn.rs create mode 100644 edb/server/conn_pool/src/lib.rs create mode 100644 edb/server/conn_pool/src/metrics.rs create mode 100644 edb/server/conn_pool/src/pool.rs create mode 100644 edb/server/conn_pool/src/python.rs create mode 100644 edb/server/conn_pool/src/test.rs create mode 100644 edb/server/conn_pool/src/waitqueue.rs create mode 100644 edb/server/connpool/pool2.py diff --git a/Cargo.lock b/Cargo.lock index 4b6901feb36..28bf46d3b18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,24 +2,156 @@ # It is not intended for manual editing. version = 3 +[[package]] +name = "addr2line" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e4503c46a5c0c7844e948c9a4d6acd9f50cccb4de1c48eb9e291ea17470c678" +dependencies = [ + "gimli", +] + +[[package]] +name = "adler" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" + +[[package]] +name = "ahash" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e89da841a80418a9b391ebaea17f5c112ffaaa96f621d2c285b5174da76b9011" +dependencies = [ + "cfg-if", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c6cb57a04249c6480766f7f7cef5467412af1490f8d1e243141daddada3264f" + +[[package]] +name = "anstream" +version = "0.6.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418c75fa768af9c03be99d17643f93f79bbba589895012a80e3452a19ddda15b" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "038dfcf04a5feb68e9c60b21c9625a54c2c0616e79b72b0fd87075a056ae1d1b" + +[[package]] +name = "anstyle-parse" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c03a11a9034d92058ceb6ee011ce58af4a9bf61491aa7e1e59ecd24bd40d22d4" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad186efb764318d35165f1758e7dcef3b10628e26d41a44bc5550652e6804391" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61a38449feb7068f52bb06c12759005cf459ee52bb4adc1d5a7c4322d716fb19" +dependencies = [ + "anstyle", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.86" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3d1d046238990b9cf5bcde22a3fb3584ee5cf65fb2765f454ed428c7a0063da" + [[package]] name = "append-only-vec" version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f3cb8f874ecf419dd8165d0279746de966cb8966636d028845e3bd65d519812a" +[[package]] +name = "approx" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cab112f0a86d568ea0e627cc1d6be74a1e9cd55214684db5561995f6dad897c6" +dependencies = [ + "num-traits", +] + [[package]] name = "ascii" version = "0.9.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eab1c04a571841102f5345a8fc0f6bb3d31c315dec879b5c6e42e40ce7ffa34e" +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi 0.1.19", + "libc", + "winapi", +] + [[package]] name = "autocfg" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" +[[package]] +name = "backtrace" +version = "0.3.73" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc23269a4f8976d0a4d2e7109211a419fe30e8d88d677cd60b6bc79c5732e0a" +dependencies = [ + "addr2line", + "cc", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", +] + [[package]] name = "base32" version = "0.4.0" @@ -83,6 +215,12 @@ version = "3.15.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ea184aa71bb362a1157c896979544cc23974e08fd265f29ea96b59f0b4a555b" +[[package]] +name = "bytemuck" +version = "1.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b236fc92302c97ed75b38da1f4917b5cdda4984745740f153a5d3059e48d725e" + [[package]] name = "byteorder" version = "1.5.0" @@ -95,12 +233,24 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" +[[package]] +name = "cc" +version = "1.0.99" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c51067fd44124faa7f870b4b1c969379ad32b2ba805aa959430ceaa384f695" + [[package]] name = "cfg-if" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" +[[package]] +name = "colorchoice" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b6a852b24ab71dffc585bcb46eaf7959d175cb865a7152e35b348d1b2960422" + [[package]] name = "combine" version = "3.8.1" @@ -114,6 +264,47 @@ dependencies = [ "unreachable", ] +[[package]] +name = "conn-pool" +version = "0.1.0" +dependencies = [ + "anyhow", + "consume_on_drop", + "derive_more", + "futures", + "genetic_algorithm", + "itertools 0.13.0", + "lru", + "pretty_assertions", + "pyo3", + "rand", + "rstest", + "scopeguard", + "smart-default", + "statrs", + "strum", + "test-log", + "thiserror", + "tokio", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "consume_on_drop" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d91586aac6b9f9e32032b8ebf58a8eb03f4b58d3e18ea493b8b2083cb7b88e" + +[[package]] +name = "convert_case" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec182b0ca2f35d8fc196cf3404988fd8b8c739a4d270ff118a398feb0cbec1ca" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "cpufeatures" version = "0.2.12" @@ -123,6 +314,62 @@ dependencies = [ "libc", ] +[[package]] +name = "crossbeam" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1137cd7e7fc0fb5d3c5a8678be38ec56e819125d8d7907411fe24ccb943faca8" +dependencies = [ + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33480d6946193aa8033910124896ca395333cae7e2d1113d1fef6c3272217df2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613f8cc01fe9cf1a3eb3d7f488fd2fa8388403e97039e2f73692932e291a770d" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df0346b5d5e76ac2fe4e327c5fd1118d6be7c51dfb18f9b7922923f287471e35" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" + [[package]] name = "crypto-common" version = "0.1.6" @@ -133,6 +380,28 @@ dependencies = [ "typenum", ] +[[package]] +name = "derive_more" +version = "1.0.0-beta.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7abbfc297053be59290e3152f8cbcd52c8642e0728b69ee187d991d4c1af08d" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "1.0.0-beta.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bba3e9872d7c58ce7ef0fcf1844fcc3e23ef2a58377b50df35dd98e42a5726e" +dependencies = [ + "convert_case", + "proc-macro2", + "quote", + "syn 2.0.52", + "unicode-xid", +] + [[package]] name = "diff" version = "0.1.13" @@ -240,12 +509,150 @@ version = "1.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a" +[[package]] +name = "env_filter" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a009aa4810eb158359dda09d0c87378e4bbb89b5a801f016885a4707ba24f7ea" +dependencies = [ + "log", +] + +[[package]] +name = "env_logger" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a12e6657c4c97ebab115a42dcee77225f7f482cdd841cf7088c657a42e9e00e7" +dependencies = [ + "atty", + "humantime", + "log", + "regex", + "termcolor", +] + +[[package]] +name = "env_logger" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b35839ba51819680ba087cd351788c9a3c476841207e0b8cee0b04722343b9" +dependencies = [ + "anstream", + "anstyle", + "env_filter", + "log", +] + [[package]] name = "equivalent" version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +[[package]] +name = "factorial" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6398219c33c5768705156b8a69c99c463a51847e2d3647f5adce32bb6e990b1c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "futures" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "645c6916888f6cb6350d2550b80fb63e734897a8498abe35cfb732b6487804b0" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eac8f7d7865dcb88bd4373ab671c8cf4508703796caa2b1985a9ca867b3fcb78" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfc6580bb841c5a68e9ef15c77ccc837b40a7504914d52e47b8b0e9bbda25a1d" + +[[package]] +name = "futures-executor" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a576fc72ae164fca6b9db127eaa9a9dda0d61316034f33a0a0d4eda41f02b01d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a44623e20b9681a318efdd71c299b6b222ed6f231972bfe2f224ebad6311f0c1" + +[[package]] +name = "futures-macro" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87750cf4b7a4c0625b1529e4c543c2182106e4dedc60a2a6455e00d212c489ac" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "futures-sink" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb8e00e87438d937621c1c6269e53f536c14d3fbd6a042bb24879e57d474fb5" + +[[package]] +name = "futures-task" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d84fa142264698cdce1a9f9172cf383a0c82de1bddcf3092901442c4097004" + +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + +[[package]] +name = "futures-util" +version = "0.3.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d6401deb83407ab3da39eba7e33987a73c3df0c82b4bb5813ee871c19c41d48" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -256,6 +663,47 @@ dependencies = [ "version_check", ] +[[package]] +name = "genetic_algorithm" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6310c399e71285ca3a06cedd9df0c5c7a53414316579b75883ca3d8521d594e5" +dependencies = [ + "crossbeam", + "env_logger 0.9.3", + "factorial", + "itertools 0.10.5", + "log", + "num", + "rand", + "rayon", + "streaming-stats", + "thread_local", +] + +[[package]] +name = "getrandom" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94b22e06ecb0110981051723910cbf0b5f5e09a2062dd7663334ee79a9d1286c" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "gimli" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40ecd4077b5ae9fd2e9e169b102c6c330d0605168eb0e8bf79952b256dbefffd" + +[[package]] +name = "glob" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2fabcfbdc87f4758337ca535fb41a6d701b65693ce38287d856d1674551ec9b" + [[package]] name = "graphql-rewrite" version = "0.1.0" @@ -280,6 +728,10 @@ name = "hashbrown" version = "0.14.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "290f1a1d9242c78d09ce40a5e87e7554ee637af1351968159f4952f028f75604" +dependencies = [ + "ahash", + "allocator-api2", +] [[package]] name = "heck" @@ -287,6 +739,33 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" +dependencies = [ + "libc", +] + +[[package]] +name = "hermit-abi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" + +[[package]] +name = "humantime" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" + [[package]] name = "indexmap" version = "1.9.3" @@ -313,18 +792,54 @@ version = "2.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e186cfbae8084e513daff4240b4797e342f988cecda4fb6c939150f96315fd8" +[[package]] +name = "is_terminal_polyfill" +version = "1.70.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8478577c03552c21db0e2724ffb8986a5ce7af88107e6be5d2ee6e158c12800" + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c" +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" + [[package]] name = "libc" version = "0.2.153" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c198f91728a82281a64e1f4f9eeb25d82cb32a5de251c6bd1b5154d63a8e7bd" +[[package]] +name = "libm" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058" + [[package]] name = "lock_api" version = "0.4.11" @@ -341,6 +856,34 @@ version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" +[[package]] +name = "lru" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3262e75e648fce39813cb56ac41f3c3e3f65217ebf3844d818d1f9398cfb0dc" +dependencies = [ + "hashbrown 0.14.3", +] + +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + +[[package]] +name = "matrixmultiply" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7574c1cf36da4798ab73da5b215bbf444f50718207754cb522201d78d1cd0ff2" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.7.1" @@ -356,6 +899,68 @@ dependencies = [ "autocfg", ] +[[package]] +name = "miniz_oxide" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87dfd01fe195c66b572b37921ad8803d010623c0aca821bea2302239d155cdae" +dependencies = [ + "adler", +] + +[[package]] +name = "nalgebra" +version = "0.32.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5c17de023a86f59ed79891b2e5d5a94c705dbe904a5b5c9c952ea6221b03e4" +dependencies = [ + "approx", + "matrixmultiply", + "nalgebra-macros", + "num-complex", + "num-rational", + "num-traits", + "rand", + "rand_distr", + "simba", + "typenum", +] + +[[package]] +name = "nalgebra-macros" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "254a5372af8fc138e36684761d3c0cdb758a4410e938babcff1c860ce14ddbfc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + +[[package]] +name = "num" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3135b08af27d103b0a51f2ae0f8632117b7b185ccf931445affa8df530576a41" +dependencies = [ + "num-bigint 0.4.4", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.2.6" @@ -390,6 +995,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-integer" version = "0.1.46" @@ -399,6 +1013,28 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-iter" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf" +dependencies = [ + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint 0.4.4", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.18" @@ -406,6 +1042,26 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi 0.3.9", + "libc", +] + +[[package]] +name = "object" +version = "0.36.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "576dfe1fc8f9df304abb159d767a29d0476f7750fbf8aa7ad07816004a207434" +dependencies = [ + "memchr", ] [[package]] @@ -414,6 +1070,12 @@ version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + [[package]] name = "parking_lot" version = "0.12.1" @@ -434,9 +1096,15 @@ dependencies = [ "libc", "redox_syscall", "smallvec", - "windows-targets", + "windows-targets 0.48.5", ] +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "phf" version = "0.11.2" @@ -479,12 +1147,30 @@ dependencies = [ "siphasher", ] +[[package]] +name = "pin-project-lite" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bda66fc9667c18cb2758a2ac84d1167245054bcf85d5d1aaa6923f45801bdd02" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + [[package]] name = "portable-atomic" version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7170ef9988bc169ba16dd36a7fa041e5c4cbeb6a35b76d4c03daded371eae7c0" +[[package]] +name = "ppv-lite86" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" + [[package]] name = "pretty_assertions" version = "1.4.0" @@ -495,6 +1181,15 @@ dependencies = [ "yansi", ] +[[package]] +name = "proc-macro-crate" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d37c51ca738a55da99dc0c4a34860fd675453b8b36209178c2249bb13651284" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.78" @@ -560,7 +1255,7 @@ version = "0.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c7e9b68bb9c3149c5b0cade5d07f953d6d125eb4337723c4ccdb665f1f96185" dependencies = [ - "heck", + "heck 0.4.1", "proc-macro2", "pyo3-build-config", "quote", @@ -582,36 +1277,203 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ + "libc", + "rand_chacha", "rand_core", ] [[package]] -name = "rand_core" -version = "0.6.4" +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +dependencies = [ + "bitflags 1.3.2", +] + +[[package]] +name = "regex" +version = "1.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b91213439dad192326a0d7c6ee3955910425f441d7038e0d6933b0aec5c4517f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata 0.4.7", + "regex-syntax 0.8.4", +] + +[[package]] +name = "regex-automata" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", +] + +[[package]] +name = "regex-automata" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38caf58cc5ef2fed281f89292ef23f6365465ed9a41b7a7754eb4e26496c92df" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax 0.8.4", +] + +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + +[[package]] +name = "regex-syntax" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a66a03ae7c801facd77a29370b4faec201768915ac14a721ba36f20bc9c209b" + +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + +[[package]] +name = "rstest" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afd55a67069d6e434a95161415f5beeada95a01c7b815508a82dcb0e1593682" +dependencies = [ + "futures", + "futures-timer", + "rstest_macros", + "rustc_version", +] + +[[package]] +name = "rstest_macros" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4165dfae59a39dd41d8dec720d3cbfbc71f69744efb480a3920f5d4e0cc6798d" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn 2.0.52", + "unicode-ident", +] + +[[package]] +name = "rustc-demangle" +version = "0.1.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +checksum = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" [[package]] -name = "redox_syscall" -version = "0.4.1" +name = "rustc_version" +version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4722d768eff46b75989dd134e5c353f0d6296e5aaa3132e776cbdb56be7731aa" +checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366" dependencies = [ - "bitflags 1.3.2", + "semver", ] +[[package]] +name = "rustversion" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "955d28af4278de8121b7ebeb796b6a45735dc01436d898801014aced2773a3d6" + [[package]] name = "ryu" version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e86697c916019a8588c99b5fac3cead74ec0b4b819707a682fd4d23fa0ce1ba1" +[[package]] +name = "safe_arch" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3460605018fdc9612bce72735cba0d27efbcd9904780d44c7e3a9948f96148a" +dependencies = [ + "bytemuck", +] + [[package]] name = "scopeguard" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "semver" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61697e0a1c7e512e84a621326239844a24d8207b4669b41bc18b32ea5cbf988b" + [[package]] name = "serde" version = "1.0.197" @@ -655,18 +1517,60 @@ dependencies = [ "digest", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "simba" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "061507c94fc6ab4ba1c9a0305018408e312e17c041eb63bef8aa726fa33aceae" +dependencies = [ + "approx", + "num-complex", + "num-traits", + "paste", + "wide", +] + [[package]] name = "siphasher" version = "0.3.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" +[[package]] +name = "slab" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" +dependencies = [ + "autocfg", +] + [[package]] name = "smallvec" version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6ecd384b10a64542d77071bd64bd7b231f4ed5940fba55e98c3de13824cf3d7" +[[package]] +name = "smart-default" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eb01866308440fc64d6c44d9e86c5cc17adfe33c4d6eed55da9145044d0ffc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + [[package]] name = "snafu" version = "0.8.1" @@ -682,9 +1586,52 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4651148226ec36010993fcba6c3381552e8463e9f3e337b75af202b0688b5274" dependencies = [ - "heck", + "heck 0.4.1", + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "statrs" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f697a07e4606a0a25c044de247e583a330dbb1731d11bc7350b81f48ad567255" +dependencies = [ + "approx", + "nalgebra", + "num-traits", + "rand", +] + +[[package]] +name = "streaming-stats" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0d670ce4e348a2081843569e0f79b21c99c91bb9028b3b3ecb0f050306de547" +dependencies = [ + "num-traits", +] + +[[package]] +name = "strum" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d8cec3501a5194c432b2b7976db6b7d10ec95c253208b45f83f7136aa985e29" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck 0.5.0", "proc-macro2", "quote", + "rustversion", "syn 2.0.52", ] @@ -722,6 +1669,37 @@ version = "0.12.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1fc403891a21bcfb7c37834ba66a547a8f402146eba7265b5a6d88059c9ff2f" +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "test-log" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dffced63c2b5c7be278154d76b479f9f9920ed34e7574201407f0b14e2bbb93" +dependencies = [ + "env_logger 0.11.3", + "test-log-macros", + "tracing-subscriber", +] + +[[package]] +name = "test-log-macros" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5999e24eaa32083191ba4e425deb75cdf25efefabe5aaccb7446dd0d4122a3f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + [[package]] name = "thiserror" version = "1.0.57" @@ -742,6 +1720,117 @@ dependencies = [ "syn 2.0.52", ] +[[package]] +name = "thread_local" +version = "1.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b9ef9bad013ada3808854ceac7b46812a6465ba368859a37e2100283d2d719c" +dependencies = [ + "cfg-if", + "once_cell", +] + +[[package]] +name = "tokio" +version = "1.38.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba4f4a02a7a80d6f274636f0aa95c7e383b912d41fe721a31f29e29698585a4a" +dependencies = [ + "backtrace", + "num_cpus", + "pin-project-lite", + "tokio-macros", +] + +[[package]] +name = "tokio-macros" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f5ae998a069d4b5aba8ee9dad856af7d520c3699e6159b185c2acd48155d39a" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "toml_datetime" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4badfd56924ae69bcc9039335b2e017639ce3f9b001c393c1b2d1ef846ce2cbf" + +[[package]] +name = "toml_edit" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8534fd7f78b5405e860340ad6575217ce99f38d4d5c8f2442cb5ecb50090e1" +dependencies = [ + "indexmap 2.2.5", + "toml_datetime", + "winnow", +] + +[[package]] +name = "tracing" +version = "0.1.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34704c8d6ebcbc939824180af020566b01a7c01f80641264eba0999f6c2b6be7" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] + +[[package]] +name = "tracing-core" +version = "0.1.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c06d3da6113f116aaee68e4d601191614c9053067f9ab7f6edbcb161237daa54" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + [[package]] name = "typenum" version = "1.17.0" @@ -754,12 +1843,24 @@ version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" +[[package]] +name = "unicode-segmentation" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4c87d22b6e3f4a18d4d40ef354e97c90fcb14dd91d7dc0aa9d8a1172ebf7202" + [[package]] name = "unicode-width" version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +[[package]] +name = "unicode-xid" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f962df74c8c05a667b5ee8bcf162993134c104e96440b663c8daa176dc772d8c" + [[package]] name = "unindent" version = "0.2.3" @@ -775,12 +1876,24 @@ dependencies = [ "void", ] +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + [[package]] name = "uuid" version = "1.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f00cc9702ca12d3c81455259621e676d0f7251cec66a21e98fe2e9a37db93b2a" +[[package]] +name = "valuable" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b7e5d4d90034032940e4ace0d9a9a057e7a45cd94e6c007832e39edb82f6d" + [[package]] name = "version_check" version = "0.9.4" @@ -793,6 +1906,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + [[package]] name = "wasm-bindgen" version = "0.2.91" @@ -849,19 +1968,85 @@ version = "0.2.91" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4f186bd2dcf04330886ce82d6f33dd75a7bfcf69ecf5763b89fcde53b6ac9838" +[[package]] +name = "wide" +version = "0.7.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2caba658a80831539b30698ae9862a72db6697dfdd7151e46920f5f2755c3ce2" +dependencies = [ + "bytemuck", + "safe_arch", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d4cc384e1e73b93bafa6fb4f1df8c41695c8a91cf9c4c64358067d15a7b6c6b" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.5", +] + [[package]] name = "windows-targets" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" dependencies = [ - "windows_aarch64_gnullvm", - "windows_aarch64_msvc", - "windows_i686_gnu", - "windows_i686_msvc", - "windows_x86_64_gnu", - "windows_x86_64_gnullvm", - "windows_x86_64_msvc", + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", +] + +[[package]] +name = "windows-targets" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f0713a46559409d202e70e28227288446bf7841d3211583a4b53e3f6d96e7eb" +dependencies = [ + "windows_aarch64_gnullvm 0.52.5", + "windows_aarch64_msvc 0.52.5", + "windows_i686_gnu 0.52.5", + "windows_i686_gnullvm", + "windows_i686_msvc 0.52.5", + "windows_x86_64_gnu 0.52.5", + "windows_x86_64_gnullvm 0.52.5", + "windows_x86_64_msvc 0.52.5", ] [[package]] @@ -870,44 +2055,121 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7088eed71e8b8dda258ecc8bac5fb1153c5cffaf2578fc8ff5d61e23578d3263" + [[package]] name = "windows_aarch64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9985fd1504e250c615ca5f281c3f7a6da76213ebd5ccc9561496568a2752afb6" + [[package]] name = "windows_i686_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" +[[package]] +name = "windows_i686_gnu" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88ba073cf16d5372720ec942a8ccbf61626074c6d4dd2e745299726ce8b89670" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87f4261229030a858f36b459e748ae97545d6f1ec60e5e0d6a3d32e0dc232ee9" + [[package]] name = "windows_i686_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" +[[package]] +name = "windows_i686_msvc" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db3c2bf3d13d5b658be73463284eaf12830ac9a26a90c717b7f771dfe97487bf" + [[package]] name = "windows_x86_64_gnu" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e4246f76bdeff09eb48875a0fd3e2af6aada79d409d33011886d3e1581517d9" + [[package]] name = "windows_x86_64_gnullvm" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "852298e482cd67c356ddd9570386e2862b5673c85bd5f88df9ab6802b334c596" + [[package]] name = "windows_x86_64_msvc" version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bec47e5bfd1bff0eeaf6d8b485cc1074891a197ab4225d504cb7a1ab88b02bf0" + +[[package]] +name = "winnow" +version = "0.5.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f593a95398737aeed53e489c785df13f3618e41dbcd6718c6addbf1395aa6876" +dependencies = [ + "memchr", +] + [[package]] name = "yansi" version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09041cd90cf85f7f8b2df60c646f853b7f535ce68f85244eb6731cf89fa498ec" + +[[package]] +name = "zerocopy" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.7.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa4f8080344d4671fb4e831a13ad1e68092748387dfc4f55e356242fae12ce3e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.52", +] diff --git a/Cargo.toml b/Cargo.toml index 6c9a653dff0..18a287a159f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "edb/edgeql-parser/edgeql-parser-derive", "edb/edgeql-parser/edgeql-parser-python", "edb/graphql-rewrite", + "edb/server/conn_pool", ] resolver = "2" diff --git a/edb/server/conn_pool/Cargo.toml b/edb/server/conn_pool/Cargo.toml new file mode 100644 index 00000000000..ebf697449eb --- /dev/null +++ b/edb/server/conn_pool/Cargo.toml @@ -0,0 +1,50 @@ +[package] +name = "conn-pool" +version = "0.1.0" +license = "MIT/Apache-2.0" +authors = ["MagicStack Inc. "] +rust-version = "1.65.0" +edition = "2021" + +[features] +python_extension = ["pyo3/extension-module"] +optimizer = [] + +[dependencies] +futures = "0" +scopeguard = "1" +pyo3 = "0" +itertools = "0" +thiserror = "1" +tracing = "0" +tracing-subscriber = "0" +strum = { version = "0.26", features = ["derive"] } +consume_on_drop = "0" +smart-default = "0" + +[dependencies.tokio] +version = "1" +features = ["rt", "time", "sync"] + +[dependencies.derive_more] +version = "1.0.0-beta.6" +features = ["full"] + +[dev-dependencies] +pretty_assertions = "1.2.0" +test-log = { version = "0", features = ["trace"] } +anyhow = "1" +rstest = "0" +rand = "0" +statrs = "0" +genetic_algorithm = "0" +lru = "0" + +[dev-dependencies.tokio] +version = "1" +features = ["macros", "rt-multi-thread", "time", "test-util"] + +[lib] +crate-type = ["lib", "cdylib"] +name = "conn_pool" +path = "src/lib.rs" diff --git a/edb/server/conn_pool/README.md b/edb/server/conn_pool/README.md new file mode 100644 index 00000000000..37784975495 --- /dev/null +++ b/edb/server/conn_pool/README.md @@ -0,0 +1,112 @@ +# Connection Pool + +## Overview + +The load-balancing algorithm is designed to optimize the allocation and +management of database connections in a way that maximizes Quality of Service +(QoS). This involves minimizing the overall time spent on connecting and +reconnecting (connection efficiency) while ensuring that latencies remain +similar across different streams of connections (fairness). + +## Architecture + +This library is split into four major components: + + 1. The low-level blocks/block, connections, and metrics code. This code + creates, destroys and transfers connections without understanding of + policies, quotas or any sort of algorithm. We ensure that the blocks and + metrics are reliable, and use this as a building block for our pool. + 2. The algorithm. This performs planning operations for acquisition, release + and rebalancing of the pool. The algorithm does not perform operations, but + rather informs that caller what it should do. + 3. The pool itself. This drives the blocks and the connector interface, and + polls the algorithm to plan next steps during acquisition, release and + during the timer-based planning callback. + 4. The Python integration code. This is behind an optional feature, and exposes + PyO3-based interface that allows a connection factory to be implemented in + Python. + +## Details + +Demand for connections is measured in terms of “database time,” which is +calculated as the product of the number of connections and the average hold time +of these connections. This metric provides a basis for determining how resources +should be distributed among different database blocks to meet their needs +effectively. + +To maximize QoS, the algorithm aims to minimize the time spent on managing +connections and keep the latencies low and uniform across various connection +streams. This involves allocation strategies that balance the immediate needs of +different database blocks with the overall system capacity and future demand +predictions. + +When a connection is acquired, the system may be in a state where the pool is +not currently constrained by demand. In such cases, connections can be allocated +greedily without complex balancing, as there are sufficient resources to meet +all demands. This allows for quick connection handling without additional +overhead. + +When the pool is constrained, the “stealing” algorithm aims to transfer +connections from less utilized or idle database blocks (victims) to those +experiencing high demand (hunger) to ensure efficient resource use and maintain +QoS. A victim block is chosen based on its idle state, characterized by holding +connections but having low or no immediate demand for them. + +Upon releasing a connection, the algorithm evaluates which backend (database +block) needs the connection the most (the hungriest). This decision is based on +current demand, wait times, and historical usage patterns. By reallocating +connections to the blocks that need them most, the algorithm ensures that +resources are utilized efficiently and effectively. + +Unused connection capacity is eventually reclaimed to prevent wastage. The +algorithm includes mechanisms to identify and collect these idle connections, +redistributing them to blocks with higher demand or returning them to the pool +for future use. This helps maintain an optimal number of active connections, +reducing unnecessary resource consumption. + +To avoid excessive thrashing, the algorithm ensures that connections are held +for a minimum period, which is longer than the time it takes to reconnect to a +database or a configured minimum threshold. This reduces the frequency of +reallocation, preventing performance degradation due to constant connection +churn and ensuring that blocks can maintain stable and predictable access to +resource + +## Detailed Algorithm + +The algorithm is designed to 1) maximize time spent running queries in a +database and 2) minimize latency of queries waiting for their turn to run. These +goals may be in conflict at times. We do this by optimizing the time spent +switching between databases, which is considered "dead time" -- as the database +is not actively performing operations. + +The demand for a connection is based on estimated total sequential processing +time. We use the average time that a connection is held, times the number of +connections in demand as a rough idea of how much total sequential time a +certain block demands in the future. + +At a regular interval, we compute two items for each block: a quota, and a +"hunger" metric. The hunger metric may indicate that a block is "hungry" +(wanting more connections), satisfied (having the expected number of +connections) or overfull (holding more connections than it should). The "hungry" +score is determined by the estimated total sequential time needed for a block. +The "overfull" score is determined by the number of extra connections held by +this block, in combination with how old the longest-held connection is. Quota is +determined by the connection rate. + +We then use the hunger metric and quota in an attempt to rebalance the pool +proactively to ensure that the connection capacity of each block reflects its +most recent demand profile. Blocks are sorted into a list of hungry and overfull +blocks, and we attempt to transfer from the most hungry to the most overfull +until we run out of either list. We may not be able to perform the rebalance +fully because of block activity that cannot be interrupted. + +If a connection is requested for a block that is hungry, it is allowed to steal +a connection from the block that most overfull and has idle connections. As the +"overfull" score is calculated in part by the longest-held connection's age, we +minimize context switching. + +When a connection is released, we choose what happens based on its state. If +more connections are waiting on this block, we return the connection to the +block to be re-used immediately. If no connections are waiting but the block is +hungry, we return it. If the block is satisfied or overfull and we have hungry +blocks waiting, we transfer it to a hungry block that has waiters. diff --git a/edb/server/conn_pool/src/algo.rs b/edb/server/conn_pool/src/algo.rs new file mode 100644 index 00000000000..61be7e07331 --- /dev/null +++ b/edb/server/conn_pool/src/algo.rs @@ -0,0 +1,852 @@ +use scopeguard::defer; +use std::cell::{Cell, RefCell}; +use tracing::trace; + +use crate::{ + block::Name, + metrics::{MetricVariant, RollingAverageU32}, +}; + +/// The historical length of data we'll maintain for demand. +const DEMAND_HISTORY_LENGTH: usize = 16; + +#[cfg(not(feature = "optimizer"))] +#[derive(Clone, Copy, derive_more::From)] +pub struct Knob(&'static str, T); + +#[cfg(not(feature = "optimizer"))] +impl Knob { + pub const fn new(name: &'static str, value: T) -> Self { + Self(name, value) + } + + pub fn get(&self) -> T { + self.1 + } +} + +#[cfg(feature = "optimizer")] +pub struct Knob( + &'static str, + &'static std::thread::LocalKey>, + Option>, +); + +impl + std::fmt::Display + std::fmt::Debug> std::fmt::Debug for Knob { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!("{}={:?}", self.0, self.get())) + } +} + +#[cfg(feature = "optimizer")] +impl + std::fmt::Display + std::fmt::Debug> Knob { + pub const fn new( + name: &'static str, + value: &'static std::thread::LocalKey>, + bounds: &[std::ops::RangeInclusive], + ) -> Self { + let copy = if !bounds.is_empty() { + Some(*bounds[0].start()..=*bounds[0].end()) + } else { + None + }; + Self(name, value, copy) + } + + pub fn name(&self) -> &'static str { + self.0 + } + + pub fn get(&self) -> T { + self.1.with_borrow(|t| *t) + } + + pub fn set(&self, value: T) -> Result<(), String> { + if let Some(range) = &self.2 { + if range.contains(&value) { + self.1.with_borrow_mut(|t| *t = value); + Ok(()) + } else { + Err(format!("{value} is out of range of {range:?}")) + } + } else { + self.1.with_borrow_mut(|t| *t = value); + Ok(()) + } + } + + pub fn clamp(&self, value: &mut T) { + if let Some(range) = &self.2 { + if !range.contains(value) { + if *value < *range.start() { + *value = *range.start() + } else { + *value = *range.end() + } + } + } + } +} + +macro_rules! constants { + ($( + $( #[doc=$doc:literal] )* + $( #[range $range:tt] )? + const $name:ident: $type:ty = $value:literal; + )*) => { + #[cfg(feature="optimizer")] + pub mod knobs { + pub use super::Knob; + mod locals { + $( + thread_local! { + pub static $name: std::cell::RefCell<$type> = std::cell::RefCell::new($value); + } + )* + } + + $( + $( #[doc=$doc] )* + pub static $name: Knob<$type> = Knob::new(stringify!($name), &locals::$name, &[$($range)?]); + )* + + pub const ALL_KNOB_COUNT: usize = [$(stringify!($name)),*].len(); + pub static ALL_KNOBS: [&Knob; ALL_KNOB_COUNT] = [ + $(&$name),* + ]; + } + #[cfg(not(feature="optimizer"))] + pub mod knobs { + pub use super::Knob; + $( + $( #[doc=$doc] )* + pub const $name: Knob<$type> = Knob::new(stringify!($name), $value); + )* + } + pub use knobs::*; + }; +} + +// Note: these constants are tuned via the generic algorithm optimizer. +constants! { + /// The maximum number of connections to create or destroy during a rebalance. + #[range(0..=10)] + const MAX_REBALANCE_OPS: usize = 5; + /// The minimum headroom in a block between its current total and its target + /// for us to pre-create connections for it. + #[range(0..=10)] + const MIN_REBALANCE_HEADROOM_TO_CREATE: usize = 2; + /// The maximum number of excess connections (> target) we'll keep around during + /// a rebalance if there is still some demand. + #[range(0..=10)] + const MAX_REBALANCE_EXCESS_IDLE_CONNECTIONS: usize = 2; + + /// The minimum amount of time we'll consider for an active connection. + #[range(1..=100)] + const MIN_TIME: usize = 1; + + /// The weight we apply to waiting connections. + const DEMAND_WEIGHT_WAITING: usize = 3; + /// The weight we apply to active connections. + const DEMAND_WEIGHT_ACTIVE: usize = 277; + /// The minimum non-zero demand. This makes the demand calculations less noisy + /// when we are competing at lower levels of demand, allowing for more + /// reproducable results. + #[range(1..=256)] + const DEMAND_MINIMUM: usize = 168; + + /// The maximum-minimum connection count we'll allocate to connections if there + /// is more capacity than backends. + const MAXIMUM_SHARED_TARGET: usize = 1; + + /// The boost we apply to our own apparent hunger when releasing a connection. + /// This prevents excessive swapping when hunger is similar across various + /// backends. + const SELF_HUNGER_BOOST_FOR_RELEASE: usize = 160; + /// The weight we apply to the difference between the target and required + /// connections when determining overfullness. + const HUNGER_DIFF_WEIGHT: usize = 20; + /// The weight we apply to waiters when determining hunger. + const HUNGER_WAITER_WEIGHT: usize = 0; + const HUNGER_WAITER_ACTIVE_WEIGHT: usize = 0; + const HUNGER_ACTIVE_WEIGHT_DIVIDEND: usize = 9650; + /// The weight we apply to the oldest waiter's age in milliseconds (as a divisor). + #[range(1..=2000)] + const HUNGER_AGE_DIVISOR_WEIGHT: usize = 1360; + + /// The weight we apply to the difference between the target and required + /// connections when determining overfullness. + const OVERFULL_DIFF_WEIGHT: usize = 20; + /// The weight we apply to idle connections when determining overfullness. + const OVERFULL_IDLE_WEIGHT: usize = 100; + /// This is divided by the youngest connection metric to penalize switching from + /// a backend which has changed recently. + const OVERFULL_CHANGE_WEIGHT_DIVIDEND: usize = 4690; + /// The weight we apply to waiters when determining overfullness. + const OVERFULL_WAITER_WEIGHT: usize = 4460; + const OVERFULL_WAITER_ACTIVE_WEIGHT: usize = 1300; + const OVERFULL_ACTIVE_WEIGHT_DIVIDEND: usize = 6620; +} + +/// Determines the rebalance plan based on the current pool state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RebalanceOp { + /// Transfer from one block to another + Transfer { to: Name, from: Name }, + /// Create a block + Create(Name), + /// Garbage collect a block. + Close(Name), +} + +/// Determines the acquire plan based on the current pool state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum AcquireOp { + /// Create a new connection. + Create, + /// Steal a connection from another block. + Steal(Name), + /// Wait for a connection. + Wait, +} + +/// Determines the release plan based on the current pool state. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReleaseOp { + /// Release this connection back to the same database. + Release, + /// Reopen this connection. + Reopen, + /// Discard this connection. + Discard, + /// Release this connection to a different database. + ReleaseTo(Name), +} + +/// The type of release to perform. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub enum ReleaseType { + /// A normal release + #[default] + Normal, + /// A release of a poisoned connection. + Poison, + /// A release of a drained connection. + Drain, +} + +/// Generic trait to decouple the algorithm from the underlying pool blocks. +/// This minimizes the interface between the algorithm and the blocks to keep +/// coupling between the two at the right level. +pub trait VisitPoolAlgoData: PoolAlgorithmDataPool { + type Block: PoolAlgorithmDataBlock; + + /// Ensure that the given block is available, inserting it with the default + /// demand if necessary. + fn ensure_block(&self, db: &str, default_demand: usize) -> bool; + /// Iterates all the blocks, garbage collecting any idle blocks with no demand. + fn with_all(&self, f: impl FnMut(&Name, &Self::Block)); + /// Retreives a single block, returning `None` if the block doesn't exist. + fn with(&self, db: &str, f: impl Fn(&Self::Block) -> T) -> Option; + + #[inline] + fn target(&self, db: &str) -> usize { + self.with(db, |data| data.target()).unwrap_or_default() + } +} + +pub trait PoolAlgorithmDataMetrics { + fn total(&self) -> usize; + fn count(&self, variant: MetricVariant) -> usize; + fn total_max(&self) -> usize; + fn max(&self, variant: MetricVariant) -> usize; + fn avg_ms(&self, variant: MetricVariant) -> usize; +} + +pub trait PoolAlgorithmDataBlock: PoolAlgorithmDataMetrics { + fn target(&self) -> usize; + fn set_target(&self, target: usize); + fn insert_demand(&self, demand: u32); + fn demand(&self) -> u32; + + fn oldest_ms(&self, variant: MetricVariant) -> usize; + fn youngest_ms(&self) -> usize; + + /// Calculates the hunger score for the current state. + /// + /// The score is determined based on the difference between the target and current metrics, + /// and the number of waiting elements. It uses weights for each component to compute the final score. + /// If the current state exceeds the target, the function returns `None`. + /// + /// # Parameters + /// + /// - `will_release`: A boolean indicating whether an element will be released. + /// + /// # Returns + /// + /// Returns an `Option` containing the hunger score if the current state is below the target + /// and there are waiting elements; otherwise, returns `None`. + fn hunger_score(&self, will_release: bool) -> Option { + let waiting = self.count(MetricVariant::Waiting); + let connecting = self.count(MetricVariant::Connecting); + let waiters = waiting.saturating_sub(connecting); + let current = self.total() - if will_release { 1 } else { 0 }; + let target = self.target(); + let active_ms = self.avg_ms(MetricVariant::Active).max(MIN_TIME.get()); + + // Waiters become more hungry as they age + let age_score = + self.oldest_ms(MetricVariant::Waiting) / HUNGER_AGE_DIVISOR_WEIGHT.get().max(1); + let waiter_score = waiters * HUNGER_WAITER_WEIGHT.get() + + (waiters * HUNGER_WAITER_ACTIVE_WEIGHT.get() / active_ms) + + (HUNGER_ACTIVE_WEIGHT_DIVIDEND.get() / active_ms); + let base_score = age_score + waiter_score; + + // If we have more connections than our target, we are not hungry. We + // may still be hungry if current <= target if we have waiters, however. + if current > target || (target == current && waiters < 1) { + None + } else { + let diff = target - current; + Some((base_score + diff * HUNGER_DIFF_WEIGHT.get()) as _) + } + } + + /// Calculates the overfull score for the current state. + /// + /// The score is determined based on the difference between the current and target metrics, + /// the idle count, and the age of the youngest element. It uses weights for each component to + /// compute the final score. If the current state is not overfull or there are no idle elements, + /// the function returns `None`. + /// + /// # Parameters + /// + /// - `will_release`: A boolean indicating whether an element will be released. + /// + /// # Returns + /// + /// Returns an `Option` containing the overfull score if the current state is overfull + /// and there are idle elements; otherwise, returns `None`. + fn overfull_score(&self, will_release: bool) -> Option { + let idle = self.count(MetricVariant::Idle) + if will_release { 1 } else { 0 }; + let current = self.total(); + let target = self.target(); + let connecting = self.count(MetricVariant::Connecting); + let waiting = self.count(MetricVariant::Waiting); + let waiters = waiting.saturating_sub(connecting); + let active_ms = self.avg_ms(MetricVariant::Active).max(MIN_TIME.get()); + let connecting_ms = self.avg_ms(MetricVariant::Connecting).max(MIN_TIME.get()); + let youngest_ms = self.youngest_ms().max(MIN_TIME.get()); + + // If we have no idle connections, or we don't have enough connections we're not overfull. + if target >= current || idle == 0 { + None + } else { + // The more idle connections we have, the more overfull this block is. + let idle_score = (idle * OVERFULL_IDLE_WEIGHT.get()) as isize; + // We take the ratio of youngest/connecting and divide + // `OVERFULL_CHANGE_WEIGHT_DIVIDEND` by that to give an overfullness + // "negative" penalty to blocks that have newly acquired a connection. + let youngest_score = + ((OVERFULL_CHANGE_WEIGHT_DIVIDEND.get() * connecting_ms) / youngest_ms) as isize; + // The number of waiters and the amount of time we expect to spend + // active on these waiters also acts as a "negative" penalty. + let waiter_score = (waiters * OVERFULL_WAITER_WEIGHT.get() + + (waiters * OVERFULL_WAITER_ACTIVE_WEIGHT.get() / active_ms) + + (OVERFULL_ACTIVE_WEIGHT_DIVIDEND.get() / active_ms)) + as isize; + + let base_score = idle_score - youngest_score - waiter_score; + if current > target { + let diff = current - target; + let diff_score = (diff * OVERFULL_DIFF_WEIGHT.get()) as isize; + Some(diff_score + base_score) + } else { + Some(base_score) + } + } + } + + /// We calculate demand based on the estimated connection active time + /// multiplied by the active + waiting counts. This gives us an + /// estimated database time statistic we can use for relative + /// weighting. + fn demand_score(&self) -> usize { + let active = self.max(MetricVariant::Active); + let active_ms = self.avg_ms(MetricVariant::Active).max(MIN_TIME.get()); + let waiting = self.max(MetricVariant::Waiting); + let idle = active == 0 && waiting == 0; + + if idle { + 0 + } else { + let waiting_score = waiting * DEMAND_WEIGHT_WAITING.get(); + let active_score = active * DEMAND_WEIGHT_ACTIVE.get(); + // Note that we clamp to DEMAND_MINIMUM to ensure the average is non-zero + (active_ms * (waiting_score + active_score)) + .max(DEMAND_MINIMUM.get() * DEMAND_HISTORY_LENGTH) + } + } +} + +pub trait PoolAlgorithmDataPool: PoolAlgorithmDataMetrics { + fn reset_max(&self); +} + +#[derive(Default, Debug)] +pub struct PoolAlgoTargetData { + /// A numeric score representing hunger or overfullness. + target_size: Cell, + avg_demand: RefCell>, +} + +impl PoolAlgoTargetData { + pub fn set_target(&self, target: usize) { + self.target_size.set(target); + } + pub fn target(&self) -> usize { + self.target_size.get() + } + pub fn insert_demand(&self, demand: u32) { + self.avg_demand.borrow_mut().accum(demand) + } + pub fn demand(&self) -> u32 { + self.avg_demand.borrow().avg() + } +} + +/// The pool algorithm constraints. +#[derive(Debug)] +pub struct PoolConstraints { + /// Maximum pool size. + pub max: usize, +} + +impl PoolConstraints { + /// Recalculate the quota targets for each block within the pool/ + pub fn recalculate_shares(&self, it: &impl VisitPoolAlgoData) { + // First, compute the overall request load and number of backend targets + let mut total_demand = 0; + let mut total_target = 0; + let mut s = "".to_owned(); + + it.with_all(|name, data| { + let demand_avg = data.demand(); + + if tracing::enabled!(tracing::Level::TRACE) { + s += &format!("{name}={demand_avg} ",); + } + + total_demand += demand_avg as usize; + if demand_avg > 0 { + total_target += 1; + } else { + data.set_target(0); + } + }); + + if tracing::enabled!(tracing::Level::TRACE) { + trace!("Demand: {total_target} {}", s); + } + + self.allocate_demand(it, total_target, total_demand); + } + + /// Adjust the quota targets for each block within the pool. + pub fn adjust(&self, it: &impl VisitPoolAlgoData) { + // Once we've adjusted the constraints, reset the max settings + defer!(it.reset_max()); + + // First, compute the overall request load and number of backend targets + let mut total_demand = 0_usize; + let mut total_target = 0; + let mut s = "".to_owned(); + + it.with_all(|name, data| { + let demand = data.demand_score(); + data.insert_demand(demand as _); + let demand_avg = data.demand(); + + if tracing::enabled!(tracing::Level::TRACE) { + s += &format!("{name}={demand_avg}/{demand}",); + } + + total_demand += demand_avg as usize; + if demand_avg > 0 { + total_target += 1; + } else { + data.set_target(0); + } + }); + + if tracing::enabled!(tracing::Level::TRACE) { + trace!("Demand: {total_target} {}", s); + } + + self.allocate_demand(it, total_target, total_demand); + } + + /// Allocate the calculated demand to target quotas. + fn allocate_demand( + &self, + it: &impl VisitPoolAlgoData, + total_target: usize, + total_demand: usize, + ) { + // Empty pool, no math + if total_target == 0 || total_demand == 0 { + return; + } + + let mut allocated = 0; + // This is the minimum number of connections we'll allocate to any particular + // backend regardless of demand if there are less backends than the capacity. + let min = (self.max / total_target).min(MAXIMUM_SHARED_TARGET.get()); + // The remaining capacity after we allocated the `min` value above. + let capacity = self.max - min * total_target; + + if min == 0 { + it.with_all(|_name, data| { + data.set_target(0); + }); + } else { + it.with_all(|_name, data| { + let demand = data.demand(); + if demand == 0 { + return; + } + + // Give everyone what they requested, plus a share of the spare + // capacity. If there is leftover spare capacity, that capacity + // may float between whoever needs it the most. + let target = + (demand as f32 * capacity as f32 / total_demand as f32).floor() as usize + min; + + data.set_target(target); + allocated += target; + }); + } + + debug_assert!( + allocated <= self.max, + "Attempted to allocate more than we were allowed: {allocated} > {} \ + (req={total_demand}, target={total_target})", + self.max + ); + } + + /// Plan a rebalance to better match the target quotas of the blocks in the + /// pool. + pub fn plan_rebalance(&self, it: &impl VisitPoolAlgoData) -> Vec { + let mut current_pool_size = it.total(); + let max_pool_size = self.max; + + // If there's room in the pool, we can be more aggressive in + // how we allocate. + if current_pool_size < max_pool_size { + let mut changes = vec![]; + let mut made_changes = false; + + for i in 0..MAX_REBALANCE_OPS.get() { + it.with_all(|name, block| { + // If there's room in the block, and room in the pool, and + // the block is bumping up against its current headroom, we'll grab + // another one. + if block.target() > block.total() + && current_pool_size < max_pool_size + && (block.max(MetricVariant::Active) + block.max(MetricVariant::Waiting)) + > (block.total() + i) + .saturating_sub(MIN_REBALANCE_HEADROOM_TO_CREATE.get()) + { + changes.push(RebalanceOp::Create(name.clone())); + current_pool_size += 1; + made_changes = true; + } else if block.total() > block.target() + && block.count(MetricVariant::Idle) > i + && (i > MAX_REBALANCE_EXCESS_IDLE_CONNECTIONS.get() || block.demand() == 0) + { + // If we're holding on to too many connections, we'll + // release some of them. If there is still some demand + // around, we'll try to keep a few excess connections if + // nobody else wants them. Otherwise, we'll just try to close + // all the idle connections over time. + changes.push(RebalanceOp::Close(name.clone())); + made_changes = true; + } + }); + if !made_changes { + break; + } + } + + return changes; + } + + // For any block with less connections than its quota that has + // waiters, we want to transfer from the most overloaded block. + let mut overloaded = vec![]; + let mut hungriest = vec![]; + let mut idle = vec![]; + + let mut s1 = "".to_owned(); + let mut s2 = "".to_owned(); + + it.with_all(|name, block| { + if let Some(value) = block.hunger_score(false) { + if tracing::enabled!(tracing::Level::TRACE) { + s1 += &format!("{name}={value} "); + } + hungriest.push((value, name.clone())) + } else if let Some(value) = block.overfull_score(false) { + if tracing::enabled!(tracing::Level::TRACE) { + s2 += &format!("{name}={value} "); + } + if block.demand() == 0 { + idle.push(name.clone()); + } else { + overloaded.push((value, name.clone())) + } + } + }); + + if tracing::enabled!(tracing::Level::TRACE) { + trace!("Hunger: {s1}"); + trace!("Overfullness: {s2}"); + } + overloaded.sort(); + hungriest.sort(); + + let mut tasks = vec![]; + + for _ in 0..MAX_REBALANCE_OPS.get() { + let Some((_, to)) = hungriest.pop() else { + // TODO: close more than one? + if let Some(idle) = idle.pop() { + tasks.push(RebalanceOp::Close(idle.clone())); + } + break; + }; + + // Prefer rebalancing from idle connections, otherwise take from + // overloaded ones. + if let Some(from) = idle.pop() { + tasks.push(RebalanceOp::Transfer { to, from }); + } else if let Some((_, from)) = overloaded.pop() { + tasks.push(RebalanceOp::Transfer { to, from }); + } else { + break; + } + } + + tasks + } + + /// Plan a connection acquisition. + pub fn plan_acquire(&self, db: &str, it: &impl VisitPoolAlgoData) -> AcquireOp { + // If the block is new, we need to perform an initial adjustment to + // ensure this block gets some capacity. + if it.ensure_block(db, DEMAND_MINIMUM.get() * DEMAND_HISTORY_LENGTH) { + self.recalculate_shares(it); + } + + let target_block_size = it.target(db); + let current_block_size = it.with(db, |data| data.total()).unwrap_or_default(); + let current_pool_size = it.total(); + let max_pool_size = self.max; + + let pool_is_full = current_pool_size >= max_pool_size; + if !pool_is_full { + trace!("Pool has room, acquiring new connection for {db}"); + return AcquireOp::Create; + } + + let block_has_room = current_block_size < target_block_size || target_block_size == 0; + trace!("Acquiring {db}: {current_pool_size}/{max_pool_size} {current_block_size}/{target_block_size}"); + if pool_is_full && block_has_room { + let mut max = isize::MIN; + let mut which = None; + it.with_all(|name, block| { + if let Some(overfullness) = block.overfull_score(false) { + if overfullness > max { + which = Some(name.clone()); + max = overfullness; + } + } + }); + match which { + Some(name) => AcquireOp::Steal(name), + None => AcquireOp::Wait, + } + } else if block_has_room { + AcquireOp::Create + } else { + AcquireOp::Wait + } + } + + /// Plan a connection release. + pub fn plan_release( + &self, + db: &str, + release_type: ReleaseType, + it: &impl VisitPoolAlgoData, + ) -> ReleaseOp { + if release_type == ReleaseType::Poison { + return ReleaseOp::Reopen; + } + if release_type == ReleaseType::Drain { + return ReleaseOp::Discard; + } + + let current_pool_size = it.total(); + let max_pool_size = self.max; + if current_pool_size < max_pool_size { + trace!("Pool has room, keeping connection"); + return ReleaseOp::Release; + } + + // We only want to consider a release elsewhere if this block is overfull + if let Some(Some(overfull)) = it.with(db, |block| block.overfull_score(true)) { + trace!("Block {db} is overfull ({overfull}), trying to release"); + let mut max = isize::MIN; + let mut which = None; + let mut s = "".to_owned(); + it.with_all(|name, block| { + let is_self = &**name == db; + if let Some(mut hunger) = block.hunger_score(is_self) { + // Penalize switching by boosting the current database's relative hunger here + if is_self { + hunger += SELF_HUNGER_BOOST_FOR_RELEASE.get() as isize; + } + + if tracing::enabled!(tracing::Level::TRACE) { + s += &format!("{name}={hunger} "); + } + + if hunger > max { + which = if is_self { None } else { Some(name.clone()) }; + max = hunger; + } + } + }); + + if tracing::enabled!(tracing::Level::TRACE) { + trace!("Hunger: {s}"); + } + + match which { + Some(name) => { + trace!("Releasing to {name:?} with score {max}"); + ReleaseOp::ReleaseTo(name) + } + None => { + trace!("Keeping connection"); + ReleaseOp::Release + } + } + } else { + trace!("Block {db} is not overfull, keeping"); + ReleaseOp::Release + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{block::Blocks, test::BasicConnector, PoolConfig}; + use anyhow::{Ok, Result}; + use futures::{stream::FuturesUnordered, StreamExt}; + use test_log::test; + use tokio::task::LocalSet; + + #[test(tokio::test(flavor = "current_thread"))] + async fn test_pool_normal() -> Result<()> { + let future = async { + let connector = BasicConnector::no_delay(); + let config = PoolConfig::suggested_default_for(10); + let blocks = Blocks::default(); + let algo = &config.constraints; + + let futures = FuturesUnordered::new(); + for i in (0..algo.max).map(Name::from) { + assert_eq!(algo.plan_acquire(&i, &blocks), AcquireOp::Create); + futures.push(blocks.create_if_needed(&connector, &i)); + } + let conns: Vec<_> = futures.collect().await; + let futures = FuturesUnordered::new(); + for i in (0..algo.max).map(Name::from) { + assert_eq!(algo.plan_acquire(&i, &blocks), AcquireOp::Wait); + futures.push(blocks.queue(&i)); + } + for conn in conns { + assert_eq!( + algo.plan_release(&conn?.state.db_name, ReleaseType::Normal, &blocks), + ReleaseOp::Release + ); + } + let conns: Vec<_> = futures.collect().await; + for conn in conns { + assert_eq!( + algo.plan_release(&conn?.state.db_name, ReleaseType::Normal, &blocks), + ReleaseOp::Release + ); + } + Ok(()) + }; + LocalSet::new().run_until(future).await + } + + /// Ensures that when a pool is starved for connections because there are + /// more backends than connections, we release connections to other to + /// ensure fairness. + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_pool_starved() -> Result<()> { + let future = async { + let connector = BasicConnector::no_delay(); + let config = PoolConfig::suggested_default_for(10); + let algo = &config.constraints; + let blocks = Blocks::default(); + + // Room for these + let futures = FuturesUnordered::new(); + for db in (0..5).map(Name::from) { + assert_eq!(algo.plan_acquire(&db, &blocks), AcquireOp::Create); + futures.push(blocks.create_if_needed(&connector, &db)); + } + // ... and these + let futures2 = FuturesUnordered::new(); + for db in (5..10).map(Name::from) { + assert_eq!(algo.plan_acquire(&db, &blocks), AcquireOp::Create); + futures2.push(blocks.create_if_needed(&connector, &db)); + } + // But not these (yet) + let futures3 = FuturesUnordered::new(); + for db in (10..15).map(Name::from) { + assert_eq!(algo.plan_acquire(&db, &blocks), AcquireOp::Wait); + futures3.push(blocks.queue(&db)); + } + let conns: Vec<_> = futures.collect().await; + let conns2: Vec<_> = futures2.collect().await; + // These are released to 10..15 + for conn in conns { + let conn = conn?; + let res = algo.plan_release(&conn.state.db_name, ReleaseType::Normal, &blocks); + let ReleaseOp::ReleaseTo(to) = res else { + panic!("Wrong release: {res:?}"); + }; + blocks.task_move_to(&connector, conn, &to).await?; + } + // These don't have anywhere to go + for conn in conns2 { + let conn = conn?; + let res = algo.plan_release(&conn.state.db_name, ReleaseType::Normal, &blocks); + let ReleaseOp::Release = res else { + panic!("Wrong release: {res:?}"); + }; + } + Ok(()) + }; + LocalSet::new().run_until(future).await + } +} diff --git a/edb/server/conn_pool/src/block.rs b/edb/server/conn_pool/src/block.rs new file mode 100644 index 00000000000..374675e019f --- /dev/null +++ b/edb/server/conn_pool/src/block.rs @@ -0,0 +1,934 @@ +use crate::{ + algo::{ + PoolAlgoTargetData, PoolAlgorithmDataBlock, PoolAlgorithmDataMetrics, + PoolAlgorithmDataPool, VisitPoolAlgoData, + }, + conn::*, + metrics::{MetricVariant, MetricsAccum, PoolMetrics}, + time::Instant, + waitqueue::WaitQueue, +}; +use futures::future::Either; +use std::{ + cell::{Cell, RefCell}, + collections::HashMap, + future::{poll_fn, ready, Future}, + rc::Rc, +}; +use tracing::trace; + +/// Perform a consistency check on entry and exit for this function. +macro_rules! consistency_check { + ($self:ident) => { + // On entry + $self.check_consistency(); + // On exit + scopeguard::defer!($self.check_consistency()); + }; +} + +/// A cheaply cloneable name string. +#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct Name(Rc); + +impl std::fmt::Display for Name { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} +impl std::fmt::Debug for Name { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.0.fmt(f) + } +} + +impl PartialEq for Name { + fn eq(&self, other: &str) -> bool { + self.0.as_str() == other + } +} + +impl From<&str> for Name { + fn from(value: &str) -> Self { + Name(Rc::new(String::from(value))) + } +} + +impl From for Name { + fn from(value: String) -> Self { + Name(Rc::new(value)) + } +} + +#[cfg(test)] +impl From for Name { + fn from(value: usize) -> Self { + Name::from(format!("db-{value}")) + } +} + +impl AsRef for Name { + fn as_ref(&self) -> &str { + self.0.as_str() + } +} + +impl std::ops::Deref for Name { + type Target = str; + fn deref(&self) -> &Self::Target { + self.0.as_str() + } +} + +impl std::borrow::Borrow for Name { + fn borrow(&self) -> &str { + self.0.as_str() + } +} + +/// Manages the connection state for a single backend database. +/// +/// This is only a set of connections, and does not understand policy, capacity, +/// balancing, or anything outside of a request to make, transfer or discard a +/// connection. It also manages connection statistics for higher layers of code +/// to make decisions. +/// +/// The block provides a number of tasks related to connections. The task +/// methods provide futures, but run the accounting "up-front" to ensure that we +/// keep a handle on quotas, even if running the task async. +/// +/// The block has an associated data generic parameter that may be provided, +/// where additional metadata for this block can live. +pub struct Block { + pub db_name: Name, + conns: RefCell>>, + state: Rc, + youngest: Cell, + /// Associated data for this block useful for statistics, quotas or other + /// information. This is provided by the algorithm in this crate. + data: D, +} + +impl Block { + pub fn new(db_name: Name, parent_metrics: Option>) -> Self { + let metrics = Rc::new(MetricsAccum::new(parent_metrics)); + let state = ConnState { + db_name: db_name.clone(), + waiters: WaitQueue::new(), + metrics, + } + .into(); + Self { + db_name, + conns: Vec::new().into(), + state, + data: Default::default(), + youngest: Cell::new(Instant::now()), + } + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn len(&self) -> usize { + self.state.metrics.total() + } + + fn conn(&self, conn: Conn) -> ConnHandle { + ConnHandle::new(conn, self.state.clone()) + } + + #[track_caller] + pub fn check_consistency(&self) { + if cfg!(debug_assertions) { + assert_eq!( + self.len(), + self.conns.borrow().len(), + "Block {} failed consistency check. Total connection count was wrong.", + self.db_name + ); + let conn_metrics = MetricsAccum::default(); + for conn in &*self.conns.borrow() { + conn_metrics.insert(conn.variant()) + } + conn_metrics.set_value(MetricVariant::Waiting, self.state.waiters.lock.get()); + assert_eq!( + self.metrics().summary().value, + conn_metrics.summary().value, + "Connection metrics are incorrect. Left: actual, right: expected" + ); + } + } + + #[inline] + pub fn metrics(&self) -> Rc { + self.state.metrics.clone() + } + + fn try_acquire_used(&self) -> Option> { + for conn in &*self.conns.borrow() { + if conn.try_lock(&self.state.metrics) { + return Some(conn.clone()); + } + } + None + } + + fn try_get_used(&self) -> Option> { + for conn in &*self.conns.borrow() { + if conn.variant() == MetricVariant::Idle { + return Some(conn.clone()); + } + } + None + } + + fn try_take_used(&self) -> Option> { + let mut lock = self.conns.borrow_mut(); + let pos = lock + .iter() + .position(|conn| conn.variant() == MetricVariant::Idle); + if let Some(index) = pos { + let conn = lock.remove(index); + return Some(conn); + } + None + } + + /// Creates a connection from this block. + #[cfg(test)] + fn create(self: Rc, connector: &C) -> impl Future>> { + let conn = { + consistency_check!(self); + let conn = Conn::new(connector.connect(&self.db_name), &self.state.metrics); + self.youngest.set(Instant::now()); + self.conns.borrow_mut().push(conn.clone()); + conn + }; + async move { + consistency_check!(self); + poll_fn(|cx| conn.poll_ready(cx, &self.state.metrics, MetricVariant::Active)).await?; + Ok(self.conn(conn)) + } + } + + /// Awaits a connection from this block. + #[cfg(test)] + fn create_if_needed( + self: Rc, + connector: &C, + ) -> impl Future>> { + if let Some(conn) = self.try_acquire_used() { + return Either::Left(ready(Ok(self.conn(conn)))); + } + Either::Right(self.create(connector)) + } + + /// Awaits a connection from this block. + fn queue(self: Rc) -> impl Future>> { + // If someone else is waiting, we have to queue, even if there's a connection + if self.state.waiters.len() == 0 { + if let Some(conn) = self.try_acquire_used() { + trace!("Got a connection"); + return Either::Left(ready(Ok(self.conn(conn)))); + } + } + // Update the metrics now before we actually queue + self.state.waiters.lock(); + self.state.metrics.insert(MetricVariant::Waiting); + let state = self.state.clone(); + let now = Instant::now(); + let guard = scopeguard::guard((), move |_| { + state.waiters.unlock(); + state + .metrics + .remove_time(MetricVariant::Waiting, now.elapsed()); + }); + Either::Right(async move { + consistency_check!(self); + loop { + if let Some(conn) = self.try_acquire_used() { + trace!("Got a connection"); + drop(guard); + return Ok(self.conn(conn)); + } + trace!("Queueing for a connection"); + self.state.waiters.queue().await; + } + }) + } + + /// Creates a connection from this block. + fn task_create(self: Rc, connector: &C) -> impl Future> { + let conn = { + consistency_check!(self); + let conn = Conn::new(connector.connect(&self.db_name), &self.state.metrics); + self.youngest.set(Instant::now()); + self.conns.borrow_mut().push(conn.clone()); + conn + }; + async move { + consistency_check!(self); + poll_fn(|cx| conn.poll_ready(cx, &self.state.metrics, MetricVariant::Idle)).await?; + self.state.waiters.trigger(); + Ok(()) + } + } + + /// Close one of idle connections in this block + /// + /// ## Panics + /// + /// If there are no idle connections, this function will panic. + fn task_close_one(self: Rc, connector: &C) -> impl Future> { + let conn = { + consistency_check!(self); + let conn = self.try_get_used().expect("Could not acquire a connection"); + conn.close(connector, &self.state.metrics); + conn + }; + async move { + consistency_check!(self); + poll_fn(|cx| conn.poll_ready(cx, &self.state.metrics, MetricVariant::Closed)).await?; + // TODO: this can be replaced by moving the final item of the list into the + // empty spot to avoid reshuffling + self.conns.borrow_mut().retain(|other| other != &conn); + conn.untrack(&self.state.metrics); + Ok(()) + } + } + + /// Steals a connection from one block to another. + /// + /// ## Panics + /// + /// If there are no idle connections, this function will panic. + fn task_reconnect( + from: Rc, + to: Rc, + connector: &C, + ) -> impl Future> { + let conn = { + consistency_check!(from); + consistency_check!(to); + + let conn = from + .try_take_used() + .expect("Could not acquire a connection"); + to.youngest.set(Instant::now()); + to.conns.borrow_mut().push(conn.clone()); + conn.transfer( + connector, + &from.state.metrics, + &to.state.metrics, + &to.db_name, + ); + conn + }; + async move { + consistency_check!(from); + consistency_check!(to); + poll_fn(|cx| conn.poll_ready(cx, &to.state.metrics, MetricVariant::Idle)).await?; + to.state.waiters.trigger(); + Ok(()) + } + } + + /// Moves a connection to a different block than it was acquired from + /// without giving any wakers on the old block a chance to get it. + fn task_reconnect_conn( + from: Rc, + to: Rc, + conn: ConnHandle, + connector: &C, + ) -> impl Future> { + let conn = { + consistency_check!(from); + consistency_check!(to); + + // TODO: this can be replaced by moving the final item of the list into the + // empty spot to avoid reshuffling + let conn = conn.into_inner(); + from.conns.borrow_mut().retain(|other| other != &conn); + to.youngest.set(Instant::now()); + to.conns.borrow_mut().push(conn.clone()); + conn.transfer( + connector, + &from.state.metrics, + &to.state.metrics, + &to.db_name, + ); + conn + }; + async move { + consistency_check!(from); + consistency_check!(to); + poll_fn(|cx| conn.poll_ready(cx, &to.state.metrics, MetricVariant::Idle)).await?; + to.state.waiters.trigger(); + Ok(()) + } + } + + /// Marks a connection as requiring re-open. + fn task_reopen( + self: Rc, + conn: ConnHandle, + connector: &C, + ) -> impl Future> { + let conn = { + consistency_check!(self); + let conn = conn.into_inner(); + conn.reopen(connector, &self.state.metrics, &self.db_name); + conn + }; + async move { + consistency_check!(self); + poll_fn(|cx| conn.poll_ready(cx, &self.state.metrics, MetricVariant::Idle)).await?; + self.state.waiters.trigger(); + Ok(()) + } + } + + /// Marks a connection as requiring a discard. + fn task_discard( + self: Rc, + conn: ConnHandle, + connector: &C, + ) -> impl Future> { + let conn = { + consistency_check!(self); + let conn = conn.into_inner(); + conn.discard(connector, &self.state.metrics); + conn + }; + async move { + consistency_check!(self); + poll_fn(|cx| conn.poll_ready(cx, &self.state.metrics, MetricVariant::Closed)).await?; + // TODO: this can be replaced by moving the final item of the list into the + // empty spot to avoid reshuffling + self.conns.borrow_mut().retain(|other| other != &conn); + conn.untrack(&self.state.metrics); + Ok(()) + } + } +} + +/// Manages the connection state for a number of backend databases. See +/// the notes on [`Block`] for the scope of responsibility of this struct. +pub struct Blocks { + map: RefCell>>>, + metrics: Rc, +} + +impl Default for Blocks { + fn default() -> Self { + Self { + map: RefCell::new(HashMap::default()), + metrics: Rc::new(MetricsAccum::default()), + } + } +} + +impl PoolAlgorithmDataMetrics for Block { + #[inline(always)] + fn avg_ms(&self, variant: MetricVariant) -> usize { + self.state.metrics.avg_ms(variant) + } + #[inline(always)] + fn count(&self, variant: MetricVariant) -> usize { + self.state.metrics.count(variant) + } + #[inline(always)] + fn max(&self, variant: MetricVariant) -> usize { + self.state.metrics.max(variant) + } + #[inline(always)] + fn total(&self) -> usize { + self.state.metrics.total() + } + #[inline(always)] + fn total_max(&self) -> usize { + self.state.metrics.total_max() + } +} + +impl PoolAlgorithmDataBlock for Block { + #[inline(always)] + fn target(&self) -> usize { + self.data.target() + } + #[inline(always)] + fn set_target(&self, target: usize) { + self.data.set_target(target); + } + #[inline(always)] + fn insert_demand(&self, demand: u32) { + self.data.insert_demand(demand) + } + #[inline(always)] + fn demand(&self) -> u32 { + self.data.demand() + } + #[inline(always)] + fn oldest_ms(&self, variant: MetricVariant) -> usize { + assert_eq!(variant, MetricVariant::Waiting); + self.state.waiters.oldest().as_millis() as _ + } + #[inline(always)] + fn youngest_ms(&self) -> usize { + self.youngest.get().elapsed().as_millis() as _ + } +} + +impl PoolAlgorithmDataMetrics for Blocks { + #[inline(always)] + fn avg_ms(&self, variant: MetricVariant) -> usize { + self.metrics.avg_ms(variant) + } + #[inline(always)] + fn count(&self, variant: MetricVariant) -> usize { + self.metrics.count(variant) + } + #[inline(always)] + fn max(&self, variant: MetricVariant) -> usize { + self.metrics.max(variant) + } + #[inline(always)] + fn total(&self) -> usize { + self.metrics.total() + } + #[inline(always)] + fn total_max(&self) -> usize { + self.metrics.total_max() + } +} + +impl PoolAlgorithmDataPool for Blocks { + #[inline(always)] + fn reset_max(&self) { + self.metrics.reset_max(); + for block in self.map.borrow().values() { + block.metrics().reset_max() + } + } +} + +impl VisitPoolAlgoData for Blocks { + type Block = Block; + + fn ensure_block(&self, db: &str, default_demand: usize) -> bool { + if self.map.borrow().contains_key(db) { + false + } else { + let block: Rc> = + Rc::new(Block::new(db.into(), Some(self.metrics.clone()))); + block.data.insert_demand(default_demand as _); + self.map.borrow_mut().insert(db.into(), block); + true + } + } + + fn with(&self, db: &str, f: impl Fn(&Block) -> T) -> Option { + self.map.borrow().get(db).map(|block| f(block)) + } + + fn with_all(&self, mut f: impl FnMut(&Name, &Block)) { + self.map.borrow_mut().retain(|name, block| { + if block.is_empty() && block.data.demand() == 0 { + false + } else { + f(name, block); + true + } + }); + } +} + +impl Blocks { + /// To ensure that we can trust our summary statistics, we run a consistency check in + /// debug mode on most operations. This is cheap enough to run all the time, but we + /// assume confidence in this code and disable the checks in release mode. + /// + /// See [`consistency_check!`] for the macro that calls this on entry and exit. + #[track_caller] + pub fn check_consistency(&self) { + if cfg!(debug_assertions) { + let mut total = 0; + for block in self.map.borrow().values() { + block.check_consistency(); + total += block.len(); + } + if total != self.metrics.total() && tracing::enabled!(tracing::Level::TRACE) { + for block in self.map.borrow().values() { + trace!( + "{}: {} {:?}", + block.db_name, + block.len(), + block.metrics().summary() + ); + } + } + assert_eq!( + total, + self.metrics.total(), + "Blocks failed consistency check. Total connection count was wrong." + ); + } + } + + pub fn name(&self, db: &str) -> Option { + if let Some((name, _)) = self.map.borrow().get_key_value(db) { + Some(name.clone()) + } else { + None + } + } + + pub fn contains(&self, db: &str) -> bool { + self.map.borrow().contains_key(db) + } + + pub fn block_count(&self) -> usize { + self.map.borrow().len() + } + + pub fn conn_count(&self) -> usize { + self.metrics.total() + } + + pub fn metrics(&self, db: &str) -> Rc { + self.map + .borrow_mut() + .get(db) + .map(|b| b.metrics()) + .unwrap_or_default() + } + + pub fn summary(&self) -> PoolMetrics { + let mut metrics = PoolMetrics::default(); + metrics.pool = self.metrics.summary(); + metrics.all_time = self.metrics.all_time(); + for (name, block) in self.map.borrow().iter() { + metrics + .blocks + .insert(name.clone(), block.metrics().summary()); + } + metrics + } + + fn block(&self, db: &str) -> Rc> { + let mut lock = self.map.borrow_mut(); + if let Some(block) = lock.get(db) { + block.clone() + } else { + let db = Name(Rc::new(db.to_owned())); + let block = Rc::new(Block::new(db.clone(), Some(self.metrics.clone()))); + lock.insert(db, block.clone()); + block + } + } + + /// Create and acquire a connection. Only used for tests. + #[cfg(test)] + pub fn create( + &self, + connector: &C, + db: &str, + ) -> impl Future>> { + consistency_check!(self); + let block = self.block(db); + block.create(connector) + } + + /// Create and acquire a connection. If a connection is free, skips + /// creation. Only used for tests. + #[cfg(test)] + pub fn create_if_needed( + &self, + connector: &C, + db: &str, + ) -> impl Future>> { + consistency_check!(self); + let block = self.block(db); + block.create_if_needed(connector) + } + + /// Queue for a connection. + pub fn queue(&self, db: &str) -> impl Future>> { + consistency_check!(self); + let block = self.block(db); + block.queue() + } + + /// Creates one connection in a block. + pub fn task_create_one(&self, connector: &C, db: &str) -> impl Future> { + consistency_check!(self); + let block = self.block(db); + block.task_create(connector) + } + + /// Closes one connection in a block. + pub fn task_close_one(&self, connector: &C, db: &str) -> impl Future> { + consistency_check!(self); + let block = self.block(db); + block.task_close_one(connector) + } + + /// Steals a connection from one block to another. + pub fn task_steal( + &self, + connector: &C, + db: &str, + from: &str, + ) -> impl Future> { + let from_block = self.block(from); + let to_block = self.block(db); + Block::task_reconnect(from_block, to_block, connector) + } + + /// Moves a connection to a different block than it was acquired from + /// without giving any wakers on the old block a chance to get it. + pub fn task_move_to( + &self, + connector: &C, + conn: ConnHandle, + db: &str, + ) -> impl Future> { + let from_block = self.block(&conn.state.db_name); + let to_block = self.block(db); + Block::task_reconnect_conn(from_block, to_block, conn, connector) + } + + /// Marks a connection as requiring a discard. + pub fn task_discard( + &self, + connector: &C, + conn: ConnHandle, + ) -> impl Future> { + let block = self.block(&conn.state.db_name); + block.task_discard(conn, connector) + } + + /// Marks a connection as requiring re-open. + pub fn task_reopen( + &self, + connector: &C, + conn: ConnHandle, + ) -> impl Future> { + let block = self.block(&conn.state.db_name); + block.task_reopen(conn, connector) + } + + /// Do we have any live blocks? + pub fn is_empty(&self) -> bool { + self.conn_count() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metrics::{MetricVariant, VariantArray}; + use crate::test::*; + use anyhow::{Ok, Result}; + use pretty_assertions::assert_eq; + use test_log::test; + use tokio::task::LocalSet; + + /// Tiny DSL to make the tests more readable. + macro_rules! assert_block { + ($block:ident has $($count:literal $type:ident),+) => { + assert_eq!( + $block.metrics().summary().value, + [$(VariantArray::with(MetricVariant::$type, $count)),+].into_iter().sum(), + stringify!(Expected block has $($count $type),+) + ); + }; + ($block:ident $db:literal is empty) => { + assert_eq!($block.metrics($db).summary().value, VariantArray::default(), stringify!(Expected block is empty)); + }; + ($block:ident $db:literal has $count:literal $type:ident) => { + assert_eq!( + $block.metrics($db).summary().value, + VariantArray::with(MetricVariant::$type, $count), + stringify!(Expected block has $count $type) + ); + }; + } + + #[test(tokio::test)] + async fn test_counts_updated() -> Result<()> { + let connector = BasicConnector::no_delay(); + let block = Rc::new(Block::::new(Name::from("db"), None)); + let f = block.clone().create(&connector); + assert_block!(block has 1 Connecting); + let conn = f.await?; + assert_block!(block has 1 Active); + let f = block.clone().queue(); + assert_block!(block has 1 Waiting, 1 Active); + drop(conn); + assert_block!(block has 1 Waiting, 1 Idle); + let conn = f.await?; + assert_block!(block has 1 Active); + drop(conn); + assert_block!(block has 1 Idle); + + Ok(()) + } + + #[test(tokio::test)] + async fn test_block() -> Result<()> { + let connector = BasicConnector::no_delay(); + let block = Rc::new(Block::::new(Name::from("db"), None)); + let conn = block.clone().create(&connector).await?; + assert_block!(block has 1 Active); + let local = LocalSet::new(); + let block2 = block.clone(); + local.spawn_local(async move { + assert_block!(block2 has 1 Active); + let conn = block2.clone().queue().await?; + assert_block!(block2 has 1 Active); + drop(conn); + Ok(()) + }); + local.spawn_local(async move { + tokio::task::yield_now().await; + drop(conn); + }); + local.await; + Ok(()) + } + + #[test(tokio::test)] + async fn test_block_parallel_acquire() -> Result<()> { + let connector = BasicConnector::no_delay(); + let block = Rc::new(Block::::new(Name::from("db"), None)); + block.clone().create(&connector).await?; + block.clone().create(&connector).await?; + block.clone().create(&connector).await?; + assert_block!(block has 3 Idle); + + let local = LocalSet::new(); + for i in 0..100 { + let block2 = block.clone(); + local.spawn_local(async move { + for _ in 0..i % 10 { + tokio::task::yield_now().await; + } + block2.clone().queue().await + }); + } + local.await; + assert_block!(block has 3 Idle); + Ok(()) + } + + #[test(tokio::test)] + async fn test_steal() -> Result<()> { + let connector = BasicConnector::no_delay(); + let blocks = Blocks::<_, ()>::default(); + assert_eq!(0, blocks.block_count()); + blocks.create(&connector, "db").await?; + blocks.create(&connector, "db").await?; + blocks.create(&connector, "db").await?; + blocks.metrics("db").reset_max(); + blocks.metrics("db2").reset_max(); + assert_eq!(1, blocks.block_count()); + assert_block!(blocks "db" has 3 Idle); + assert_block!(blocks "db2" is empty); + blocks.task_steal(&connector, "db2", "db").await?; + blocks.task_steal(&connector, "db2", "db").await?; + blocks.task_steal(&connector, "db2", "db").await?; + // Block hasn't been GC'd yet + assert_eq!(2, blocks.block_count()); + assert_block!(blocks "db" is empty); + assert_block!(blocks "db2" has 3 Idle); + // Should not activate a connection to steal it + assert_eq!(0, blocks.metrics("db").max(MetricVariant::Active)); + assert_eq!(0, blocks.metrics("db2").max(MetricVariant::Active)); + Ok(()) + } + + #[test(tokio::test)] + async fn test_move() -> Result<()> { + let connector = BasicConnector::no_delay(); + let blocks = Blocks::<_, ()>::default(); + assert_eq!(0, blocks.block_count()); + blocks.create(&connector, "db").await?; + blocks.create(&connector, "db").await?; + blocks.create(&connector, "db").await?; + blocks.metrics("db").reset_max(); + blocks.metrics("db2").reset_max(); + assert_eq!(1, blocks.block_count()); + assert_block!(blocks "db" has 3 Idle); + assert_block!(blocks "db2" is empty); + let conn = blocks.queue("db").await?; + blocks.task_move_to(&connector, conn, "db2").await?; + assert_eq!(2, blocks.block_count()); + assert_block!(blocks "db" has 2 Idle); + assert_block!(blocks "db2" has 1 Idle); + // Should not activate a connection to move it + assert_eq!(1, blocks.metrics("db").max(MetricVariant::Active)); + assert_eq!(0, blocks.metrics("db2").max(MetricVariant::Active)); + Ok(()) + } + + #[test(tokio::test)] + async fn test_close() -> Result<()> { + let connector = BasicConnector::no_delay(); + let blocks = Blocks::<_, ()>::default(); + assert_eq!(0, blocks.block_count()); + blocks.create(&connector, "db").await?; + blocks.create(&connector, "db").await?; + blocks.metrics("db").reset_max(); + assert_eq!(1, blocks.block_count()); + assert_block!(blocks "db" has 2 Idle); + blocks.task_close_one(&connector, "db").await?; + blocks.task_close_one(&connector, "db").await?; + assert_block!(blocks "db" is empty); + // Hasn't GC'd yet + assert_eq!(1, blocks.block_count()); + // Should not activate a connection to close it + assert_eq!(0, blocks.metrics("db").max(MetricVariant::Active)); + Ok(()) + } + + #[test(tokio::test)] + async fn test_reopen() -> Result<()> { + let connector = BasicConnector::no_delay(); + let blocks = Blocks::<_, ()>::default(); + assert_eq!(0, blocks.block_count()); + let conn = blocks.create(&connector, "db").await?; + blocks.task_reopen(&connector, conn).await?; + assert_block!(blocks "db" has 1 Idle); + assert_eq!( + blocks.metrics("db").all_time()[MetricVariant::Connecting], + 2 + ); + assert_eq!( + blocks.metrics("db").all_time()[MetricVariant::Disconnecting], + 1 + ); + Ok(()) + } + + #[test(tokio::test)] + async fn test_discard() -> Result<()> { + let connector = BasicConnector::no_delay(); + let blocks = Blocks::<_, ()>::default(); + assert_eq!(0, blocks.block_count()); + let conn = blocks.create(&connector, "db").await?; + blocks.task_discard(&connector, conn).await?; + assert_block!(blocks "db" is empty); + assert_eq!( + blocks.metrics("db").all_time()[MetricVariant::Connecting], + 1 + ); + assert_eq!( + blocks.metrics("db").all_time()[MetricVariant::Disconnecting], + 1 + ); + Ok(()) + } +} diff --git a/edb/server/conn_pool/src/conn.rs b/edb/server/conn_pool/src/conn.rs new file mode 100644 index 00000000000..b2459581ec5 --- /dev/null +++ b/edb/server/conn_pool/src/conn.rs @@ -0,0 +1,348 @@ +use crate::{ + block::Name, + metrics::{MetricVariant, MetricsAccum}, + time::Instant, + waitqueue::WaitQueue, +}; +use futures::FutureExt; +use std::{ + borrow::Cow, + cell::{Cell, RefCell}, + future::Future, + pin::Pin, + rc::Rc, + task::{ready, Poll}, +}; +use tracing::error; + +pub struct ConnState { + pub db_name: Name, + pub waiters: WaitQueue, + pub metrics: Rc, +} + +#[derive(Debug, thiserror::Error)] +pub enum ConnError { + #[error("Shutdown")] + Shutdown, + #[error("{0}")] + Other(Cow<'static, str>), +} + +pub type ConnResult = Result; + +pub trait Connector: std::fmt::Debug + 'static { + /// The type of connection associated with this [`Connector`]. + type Conn; + + /// Perform a connect operation to the given database. + fn connect(&self, db: &str) -> impl Future> + 'static; + + /// Perform a graceful reconnect operation from the existing connection to a new database. + fn reconnect( + &self, + conn: Self::Conn, + db: &str, + ) -> impl Future> + 'static; + + /// Perform a graceful disconnect operation on the given connection. + fn disconnect(&self, conn: Self::Conn) -> impl Future> + 'static; +} + +#[derive(Debug)] +pub struct Conn { + inner: Rc>>, +} + +impl PartialEq for Conn { + fn eq(&self, other: &Self) -> bool { + Rc::ptr_eq(&self.inner, &other.inner) + } +} + +impl Eq for Conn {} + +impl Clone for Conn { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl Conn { + pub fn new( + f: impl Future> + 'static, + metrics: &MetricsAccum, + ) -> Self { + metrics.insert(MetricVariant::Connecting); + Self { + inner: Rc::new(RefCell::new(ConnInner::Connecting( + Instant::now(), + f.boxed_local(), + ))), + } + } + + #[inline(always)] + pub fn with_handle(&self, f: impl Fn(&C::Conn) -> T) -> Option { + match &*self.inner.borrow() { + ConnInner::Active(_, conn, ..) => Some(f(conn)), + _ => None, + } + } + + #[inline] + fn transition(&self, f: impl FnOnce(ConnInner) -> ConnInner) { + let mut lock = self.inner.borrow_mut(); + let inner = std::mem::replace(&mut *lock, ConnInner::Transition); + *lock = f(inner); + } + + pub fn close(&self, connector: &C, metrics: &MetricsAccum) { + self.transition(|inner| match inner { + ConnInner::Idle(t, conn, ..) => { + metrics.transition( + MetricVariant::Idle, + MetricVariant::Disconnecting, + t.elapsed(), + ); + let f = connector.disconnect(conn).boxed_local(); + ConnInner::Disconnecting(Instant::now(), f) + } + _ => unreachable!(), + }); + } + + pub fn discard(&self, connector: &C, metrics: &MetricsAccum) { + self.transition(|inner| match inner { + ConnInner::Active(t, conn, ..) => { + metrics.transition( + MetricVariant::Active, + MetricVariant::Disconnecting, + t.elapsed(), + ); + let f = connector.disconnect(conn).boxed_local(); + ConnInner::Disconnecting(Instant::now(), f) + } + _ => unreachable!(), + }); + } + + pub fn transfer(&self, connector: &C, from: &MetricsAccum, to: &MetricsAccum, db: &str) { + self.untrack(from); + self.transition(|inner| match inner { + ConnInner::Idle(_t, conn, ..) | ConnInner::Active(_t, conn, ..) => { + from.inc_all_time(MetricVariant::Disconnecting); + from.inc_all_time(MetricVariant::Closed); + to.insert(MetricVariant::Connecting); + let f = connector.reconnect(conn, db).boxed_local(); + ConnInner::Connecting(Instant::now(), f) + } + _ => unreachable!(), + }); + } + + pub fn reopen(&self, connector: &C, metrics: &MetricsAccum, db: &str) { + self.transition(|inner| match inner { + ConnInner::Active(t, conn) => { + metrics.inc_all_time(MetricVariant::Disconnecting); + metrics.inc_all_time(MetricVariant::Closed); + metrics.transition( + MetricVariant::Active, + MetricVariant::Connecting, + t.elapsed(), + ); + let f = connector.reconnect(conn, db).boxed_local(); + ConnInner::Connecting(Instant::now(), f) + } + _ => unreachable!(), + }); + } + + pub fn poll_ready( + &self, + cx: &mut std::task::Context, + metrics: &MetricsAccum, + to: MetricVariant, + ) -> Poll> { + let mut lock = self.inner.borrow_mut(); + + let res = match &mut *lock { + ConnInner::Idle(..) => Ok(()), + ConnInner::Connecting(t, f) => match ready!(f.poll_unpin(cx)) { + Ok(c) => { + debug_assert!(to == MetricVariant::Active || to == MetricVariant::Idle); + metrics.transition(MetricVariant::Connecting, to, t.elapsed()); + if to == MetricVariant::Active { + *lock = ConnInner::Active(Instant::now(), c); + } else { + *lock = ConnInner::Idle(Instant::now(), c); + } + Ok(()) + } + Err(err) => { + metrics.transition( + MetricVariant::Connecting, + MetricVariant::Failed, + t.elapsed(), + ); + *lock = ConnInner::Failed; + Err(err) + } + }, + ConnInner::Disconnecting(t, f) => match ready!(f.poll_unpin(cx)) { + Ok(_) => { + debug_assert_eq!(to, MetricVariant::Closed); + metrics.transition(MetricVariant::Disconnecting, to, t.elapsed()); + *lock = ConnInner::Closed; + Ok(()) + } + Err(err) => { + metrics.transition( + MetricVariant::Disconnecting, + MetricVariant::Failed, + t.elapsed(), + ); + *lock = ConnInner::Failed; + Err(err) + } + }, + ConnInner::Failed => Err(ConnError::Other("Failed".into())), + _ => unreachable!(), + }; + Poll::Ready(res) + } + + pub fn try_lock(&self, metrics: &MetricsAccum) -> bool { + let mut lock = self.inner.borrow_mut(); + + let res: bool; + let old = std::mem::replace(&mut *lock, ConnInner::Transition); + (*lock, res) = match old { + ConnInner::Idle(t, conn) => { + metrics.transition(MetricVariant::Idle, MetricVariant::Active, t.elapsed()); + (ConnInner::Active(Instant::now(), conn), true) + } + other => (other, false), + }; + res + } + + pub fn variant(&self) -> MetricVariant { + (&*self.inner.borrow()).into() + } + + pub fn untrack(&self, metrics: &MetricsAccum) { + match &*self.inner.borrow() { + ConnInner::Active(t, _) + | ConnInner::Idle(t, _) + | ConnInner::Connecting(t, _) + | ConnInner::Disconnecting(t, _) => metrics.remove_time(self.variant(), t.elapsed()), + other => metrics.remove(other.into()), + } + } +} + +/// Connection state diagram: +/// +/// ```text +/// v-------------+ +/// S -> Connecting -> Idle -> Active -+ +/// -> Failed +-> Disconnecting -> Closed +/// ``` +enum ConnInner { + /// Connecting connections hold a spot in the pool as they count towards quotas + Connecting(Instant, Pin>>>), + /// Disconnecting connections hold a spot in the pool as they count towards quotas + Disconnecting(Instant, Pin>>>), + /// The connection is alive, but it is not being held. + Idle(Instant, C::Conn), + /// The connection is alive, and is being held. + Active(Instant, C::Conn), + /// The connection is in a failed state. + Failed, + /// The connection is in a closed state. + Closed, + /// Transitioning between states. Used internally, never escapes an internal + /// function. + Transition, +} + +impl From<&ConnInner> for MetricVariant { + fn from(val: &ConnInner) -> Self { + match val { + ConnInner::Connecting(..) => MetricVariant::Connecting, + ConnInner::Disconnecting(..) => MetricVariant::Disconnecting, + ConnInner::Idle(..) => MetricVariant::Idle, + ConnInner::Active(..) => MetricVariant::Active, + ConnInner::Failed => MetricVariant::Failed, + ConnInner::Closed => MetricVariant::Closed, + ConnInner::Transition => unreachable!(), + } + } +} + +impl std::fmt::Debug for ConnInner { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!("ConnInner({:?})", MetricVariant::from(self))) + } +} + +pub struct ConnHandle { + pub(crate) conn: Conn, + pub(crate) state: Rc, + pub(crate) dropped: Cell, +} + +impl std::fmt::Debug for ConnHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_fmt(format_args!( + "Connection({:?}, {:?})", + self.state.db_name, + self.conn.variant() + )) + } +} + +impl ConnHandle { + pub fn new(conn: Conn, state: Rc) -> Self { + Self { + conn, + state, + dropped: Cell::default(), + } + } + + pub(crate) fn into_inner(mut self) -> Conn { + self.dropped.set(true); + std::mem::replace( + &mut self.conn, + Conn { + inner: Rc::new(RefCell::new(ConnInner::Closed)), + }, + ) + } +} + +impl Drop for ConnHandle { + fn drop(&mut self) { + if self.dropped.get() { + return; + } + self.conn.transition(|inner| match inner { + ConnInner::Active(t, c) => { + self.state.metrics.transition( + MetricVariant::Active, + MetricVariant::Idle, + t.elapsed(), + ); + self.state.waiters.trigger(); + ConnInner::Idle(Instant::now(), c) + } + _ => { + unreachable!("Impossible state: {:?}", MetricVariant::from(&inner)); + } + }); + } +} diff --git a/edb/server/conn_pool/src/lib.rs b/edb/server/conn_pool/src/lib.rs new file mode 100644 index 00000000000..36aa76f4fa8 --- /dev/null +++ b/edb/server/conn_pool/src/lib.rs @@ -0,0 +1,27 @@ +pub(crate) mod algo; +pub(crate) mod block; +pub(crate) mod conn; +pub(crate) mod metrics; +pub(crate) mod pool; +pub(crate) mod waitqueue; + +mod time { + #[cfg(not(test))] + pub use std::time::Instant; + #[cfg(test)] + pub use tokio::time::Instant; +} + +#[cfg(feature = "optimizer")] +pub use algo::knobs; + +// Public interface + +pub use conn::Connector; +pub use pool::{Pool, PoolConfig, PoolHandle}; + +#[cfg(test)] +pub mod test; + +#[cfg(feature = "python_extension")] +mod python; diff --git a/edb/server/conn_pool/src/metrics.rs b/edb/server/conn_pool/src/metrics.rs new file mode 100644 index 00000000000..bd82a33f8d2 --- /dev/null +++ b/edb/server/conn_pool/src/metrics.rs @@ -0,0 +1,373 @@ +use std::collections::BTreeMap; +use std::{cell::RefCell, rc::Rc, time::Duration}; +use strum::EnumCount; +use strum::IntoEnumIterator; + +use crate::algo::PoolAlgorithmDataMetrics; +use crate::block::Name; + +#[derive( + Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumCount, strum::EnumIter, strum::AsRefStr, +)] +pub enum MetricVariant { + Connecting, + Disconnecting, + Idle, + Active, + Failed, + Closed, + Waiting, +} + +/// Maintains a rolling average of `u32` values. Note that this struct attempts +/// to optimize `SIZE == 1`. +#[derive(Debug, PartialEq, Eq)] +pub struct RollingAverageU32 { + values: [u32; SIZE], + cumulative: u64, + ptr: u8, + /// If we've never rolled over, we cannot divide the entire array by `SIZE` + /// and have to use `ptr` instead. + rollover: bool, +} + +impl Default for RollingAverageU32 { + fn default() -> Self { + assert!(SIZE <= u8::MAX as _); + Self { + values: [0; SIZE], + ptr: 0, + cumulative: 0, + rollover: false, + } + } +} + +impl RollingAverageU32 { + pub fn accum(&mut self, new: u32) { + if SIZE == 1 { + self.values[0] = new; + } else { + let size = SIZE as u8; + let old = std::mem::replace(&mut self.values[self.ptr as usize], new); + self.cumulative -= old as u64; + self.cumulative += new as u64; + self.ptr = (self.ptr + 1) % size; + if self.ptr == 0 { + self.rollover = true; + } + } + } + + #[inline] + pub fn avg(&self) -> u32 { + if SIZE == 1 { + self.values[0] + } else if self.rollover { + (self.cumulative / SIZE as u64) as u32 + } else if self.ptr == 0 { + 0 + } else { + (self.cumulative / self.ptr as u64) as u32 + } + } +} + +#[derive(Debug, Default)] +pub struct PoolMetrics { + pub pool: ConnMetrics, + pub all_time: VariantArray, + pub blocks: BTreeMap, +} + +/// An array indexed by [`MetricVariant`]. +#[derive(Default, Clone, Copy, PartialEq, Eq)] +pub struct VariantArray([T; MetricVariant::COUNT]); + +impl std::ops::Index for VariantArray { + type Output = T; + fn index(&self, index: MetricVariant) -> &Self::Output { + &self.0[index as usize] + } +} + +impl std::ops::IndexMut for VariantArray { + fn index_mut(&mut self, index: MetricVariant) -> &mut Self::Output { + &mut self.0[index as usize] + } +} + +impl std::ops::Add for VariantArray { + type Output = VariantArray; + fn add(self, rhs: Self) -> Self::Output { + let mut out = self; + for i in MetricVariant::iter() { + out[i] += rhs[i]; + } + out + } +} + +impl std::ops::AddAssign for VariantArray { + fn add_assign(&mut self, rhs: Self) { + for i in MetricVariant::iter() { + self[i] += rhs[i]; + } + } +} + +impl std::iter::Sum for VariantArray { + fn sum>(iter: I) -> Self { + let mut sum = Default::default(); + for i in iter { + sum += i; + } + sum + } +} + +impl std::fmt::Debug for VariantArray { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut d = f.debug_struct(""); + for variant in MetricVariant::iter() { + if self[variant] != T::default() { + d.field(variant.as_ref(), &self[variant]); + } + } + d.finish() + } +} + +impl VariantArray { + #[cfg(test)] + pub fn with(variant: MetricVariant, count: T) -> Self { + let mut summary = Self::default(); + summary[variant] = count; + summary + } +} + +#[derive(Default)] +#[allow(unused)] +pub struct ConnMetrics { + pub(crate) value: VariantArray, + pub(crate) max: VariantArray, + pub(crate) avg_time: VariantArray, + pub(crate) total: usize, + pub(crate) total_max: usize, +} + +impl std::fmt::Debug for ConnMetrics { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("ConnMetrics {\n")?; + for variant in MetricVariant::iter() { + f.write_fmt(format_args!( + " {variant:?}: {} (max={}, avg={}ms)\n", + self.value[variant], self.max[variant], self.avg_time[variant] + ))?; + } + f.write_str("}")?; + Ok(()) + } +} + +#[derive(Debug, Default)] +struct RawMetrics { + /// The total number of non-waiting connections. + total: usize, + /// The max total number of non-waiting connections. + total_max: usize, + /// The number of connections per state. + counts: VariantArray, + /// The total number of transitions into this state for all time. + all_time: VariantArray, + /// The max number of connections per state. + max: VariantArray, + /// The time spent in each state. + times: VariantArray>, +} + +impl RawMetrics { + #[inline(always)] + fn reset_max(&mut self) { + self.max = self.counts; + self.total_max = self.total; + } + + #[inline(always)] + fn inc_all_time(&mut self, to: MetricVariant) { + self.all_time[to] += 1; + } + + #[inline(always)] + fn inc(&mut self, to: MetricVariant) { + self.counts[to] += 1; + self.max[to] = self.max[to].max(self.counts[to]); + self.inc_all_time(to) + } + + #[inline(always)] + fn inc_total(&mut self, to: MetricVariant) { + if to != MetricVariant::Waiting { + self.total += 1; + self.total_max = self.total_max.max(self.total); + } + } + + #[inline(always)] + fn dec(&mut self, from: MetricVariant) { + self.counts[from] -= 1; + } + + #[inline(always)] + fn time(&mut self, from: MetricVariant, time: Duration) { + self.times[from].accum(time.as_millis() as _); + } + + #[inline(always)] + fn dec_total(&mut self, from: MetricVariant) { + if from != MetricVariant::Waiting { + self.total -= 1; + } + } +} + +/// Metrics accumulator. Designed to be updated without a lock. +#[derive(Debug, Default)] +pub struct MetricsAccum { + raw: RefCell, + parent: Option>, +} + +impl PoolAlgorithmDataMetrics for MetricsAccum { + #[inline(always)] + fn avg_ms(&self, variant: MetricVariant) -> usize { + self.raw.borrow().times[variant].avg() as _ + } + #[inline(always)] + fn count(&self, variant: MetricVariant) -> usize { + self.raw.borrow().counts[variant] + } + #[inline(always)] + fn max(&self, variant: MetricVariant) -> usize { + self.raw.borrow().max[variant] + } + #[inline(always)] + fn total(&self) -> usize { + self.raw.borrow().total + } + #[inline(always)] + fn total_max(&self) -> usize { + self.raw.borrow().total_max + } +} + +impl MetricsAccum { + pub fn new(parent: Option>) -> Self { + Self { + parent, + ..Default::default() + } + } + + /// Get the current total + #[inline(always)] + pub fn total(&self) -> usize { + self.raw.borrow().total + } + + /// Get the current value of a variant + #[inline(always)] + pub fn get(&self, variant: MetricVariant) -> usize { + self.raw.borrow().counts[variant] + } + + #[inline(always)] + pub fn reset_max(&self) { + self.raw.borrow_mut().reset_max(); + } + + pub fn summary(&self) -> ConnMetrics { + let lock = self.raw.borrow_mut(); + let mut avg_time = VariantArray::default(); + for i in MetricVariant::iter() { + avg_time[i] = lock.times[i].avg(); + } + ConnMetrics { + value: lock.counts, + max: lock.max, + avg_time, + total: lock.total, + total_max: lock.total_max, + } + } + + pub fn all_time(&self) -> VariantArray { + let lock = self.raw.borrow(); + lock.all_time + } + + #[inline] + pub fn inc_all_time(&self, to: MetricVariant) { + let mut lock = self.raw.borrow_mut(); + lock.inc_all_time(to); + if let Some(parent) = &self.parent { + parent.inc_all_time(to); + } + } + + #[inline] + pub fn insert(&self, to: MetricVariant) { + let mut lock = self.raw.borrow_mut(); + lock.inc(to); + lock.inc_total(to); + // trace!("None->{to:?} ({})", lock[to ]); + if let Some(parent) = &self.parent { + parent.insert(to); + } + } + + #[inline] + pub fn set_value(&self, to: MetricVariant, len: usize) { + let mut lock = self.raw.borrow_mut(); + debug_assert_eq!(lock.counts[to], 0); + lock.counts[to] = len; + lock.total += len; + lock.max[to] = lock.max[to].max(lock.counts[to]); + } + + #[inline] + pub fn transition(&self, from: MetricVariant, to: MetricVariant, time: Duration) { + // trace!("{from:?}->{to:?}: {time:?}"); + let mut lock = self.raw.borrow_mut(); + lock.dec(from); + lock.time(from, time); + lock.inc(to); + if let Some(parent) = &self.parent { + parent.transition(from, to, time); + } + } + + #[inline] + pub fn remove_time(&self, from: MetricVariant, time: Duration) { + let mut lock = self.raw.borrow_mut(); + lock.dec(from); + lock.time(from, time); + lock.dec_total(from); + // trace!("{from:?}->None ({time:?})"); + if let Some(parent) = &self.parent { + parent.remove_time(from, time); + } + } + + #[inline] + pub fn remove(&self, from: MetricVariant) { + let mut lock = self.raw.borrow_mut(); + lock.dec(from); + lock.dec_total(from); + // trace!("{from:?}->None"); + if let Some(parent) = &self.parent { + parent.remove(from); + } + } +} diff --git a/edb/server/conn_pool/src/pool.rs b/edb/server/conn_pool/src/pool.rs new file mode 100644 index 00000000000..63729898246 --- /dev/null +++ b/edb/server/conn_pool/src/pool.rs @@ -0,0 +1,1398 @@ +use crate::{ + algo::{ + AcquireOp, PoolAlgoTargetData, PoolAlgorithmDataBlock, PoolAlgorithmDataMetrics, + PoolConstraints, RebalanceOp, ReleaseOp, ReleaseType, VisitPoolAlgoData, + }, + block::{Blocks, Name}, + conn::{ConnError, ConnHandle, ConnResult, Connector}, + metrics::{MetricVariant, PoolMetrics}, +}; +use consume_on_drop::{Consume, ConsumeOnDrop}; +use derive_more::Debug; +use std::{ + cell::{Cell, RefCell}, + collections::HashMap, + rc::Rc, + time::Duration, +}; +use tracing::trace; + +#[derive(Debug)] +pub struct PoolConfig { + pub constraints: PoolConstraints, + pub adjustment_interval: Duration, +} + +impl PoolConfig { + pub fn assert_valid(&self) { + assert!(self.constraints.max > 0); + } + + /// Generate suggested default configurations for the expected number of connections with an + /// unknown number of databases. + pub fn suggested_default_for(connections: usize) -> Self { + Self::suggested_default_for_databases(connections, usize::MAX) + } + + /// Generate suggested default configurations for the expected number of connections and databases. + pub fn suggested_default_for_databases(connections: usize, databases: usize) -> Self { + assert!(connections > 0); + assert!(databases > 0); + Self { + adjustment_interval: Duration::from_millis(10), + constraints: PoolConstraints { max: connections }, + } + } +} + +struct HandleAndPool(ConnHandle, Rc>, Cell); + +/// An opaque handle representing a RAII lock on a connection in the pool. The +/// underlying connection object may be +pub struct PoolHandle { + conn: ConsumeOnDrop>, +} + +impl Debug for PoolHandle { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.conn.0.fmt(f) + } +} + +impl Consume for HandleAndPool { + fn consume(self) { + self.1.release(self.0, self.2.get()) + } +} + +impl PoolHandle { + /// Marks this handle as poisoned, which will not allow it to be reused in the pool. The + /// most likely case for this is that the underlying connection's stream has closed, or + /// the remote end is no longer valid for some reason. + pub fn poison(&self) { + self.conn.2.set(true) + } + + /// Use this pool's handle temporarily. + #[inline(always)] + pub fn with_handle(&self, f: impl Fn(&C::Conn) -> T) -> T { + self.conn.0.conn.with_handle(f).unwrap() + } + + fn new(conn: ConnHandle, pool: Rc>) -> Self { + Self { + conn: ConsumeOnDrop::new(HandleAndPool(conn, pool, Cell::default())), + } + } +} + +impl PoolHandle +where + C::Conn: Copy, +{ + /// If the handle is `Copy`, copies this handle. + #[inline(always)] + pub fn handle(&self) -> C::Conn { + self.conn.0.conn.with_handle(|c| *c).unwrap() + } +} + +impl PoolHandle +where + C::Conn: Clone, +{ + /// If the handle is `Clone`, clones this handle. + #[inline(always)] + pub fn handle_clone(&self) -> C::Conn { + self.conn.0.conn.with_handle(|c| c.clone()).unwrap() + } +} + +#[derive(derive_more::Debug)] +/// A connection pool consists of a number of blocks, each with a target +/// connection count (aka a quota). Generally, a block may take up to its quota, +/// but no more, though the pool operating conditions may allow for this to vary +/// for optimal allocation of the limited connections. If a block is over quota, +/// one of its connections may be stolen to satisfy another block's needs. +pub struct Pool { + connector: C, + config: PoolConfig, + blocks: Blocks, + drain: Drain, + /// If the pool has been dirtied by acquiring or releasing a connection + dirty: Rc>, +} + +impl Pool { + pub fn new(config: PoolConfig, connector: C) -> Rc { + config.assert_valid(); + Rc::new(Self { + config, + blocks: Default::default(), + connector, + dirty: Default::default(), + drain: Default::default(), + }) + } +} + +impl Pool { + /// Runs the required async task that takes care of quota management, garbage collection, + /// and other important async tasks. This should happen only if something has changed in + /// the pool. + pub async fn run(&self) { + loop { + tokio::time::sleep(self.config.adjustment_interval).await; + self.run_once(); + } + } + + /// Runs the required async task that takes care of quota management, garbage collection, + /// and other important async tasks. This should happen only if we have live blocks. + pub fn run_once(&self) { + if self.blocks.is_empty() { + return; + } + + self.config.constraints.adjust(&self.blocks); + let mut s = String::new(); + self.blocks.with_all(|name, block| { + s += &format!("{name}={} ", block.target()); + }); + trace!("Targets: {s}"); + for op in self.config.constraints.plan_rebalance(&self.blocks) { + trace!("Rebalance: {op:?}"); + match op { + RebalanceOp::Transfer { from, to } => { + tokio::task::spawn_local(self.blocks.task_steal(&self.connector, &to, &from)); + } + RebalanceOp::Create(name) => { + tokio::task::spawn_local(self.blocks.task_create_one(&self.connector, &name)); + } + RebalanceOp::Close(name) => { + tokio::task::spawn_local(self.blocks.task_close_one(&self.connector, &name)); + } + } + } + } + + /// Acquire a handle from this connection pool. The returned [`PoolHandle`] + /// controls the lock for the connection and may be dropped to release it + /// back into the pool. + pub async fn acquire(self: &Rc, db: &str) -> ConnResult> { + if self.drain.shutdown.get() { + return Err(ConnError::Shutdown); + } + self.dirty.set(true); + let plan = self.config.constraints.plan_acquire(db, &self.blocks); + trace!("Acquire {db}: {plan:?}"); + match plan { + AcquireOp::Create => { + tokio::task::spawn_local(self.blocks.task_create_one(&self.connector, db)); + } + AcquireOp::Steal(from) => { + tokio::task::spawn_local(self.blocks.task_steal(&self.connector, db, &from)); + } + AcquireOp::Wait => {} + }; + let conn = self.blocks.queue(db).await?; + + Ok(PoolHandle::new(conn, self.clone())) + } + + /// Internal release method + fn release(self: Rc, conn: ConnHandle, poison: bool) { + let db = &conn.state.db_name; + self.dirty.set(true); + let release_type = if self.drain.is_draining(db) { + ReleaseType::Drain + } else if poison { + ReleaseType::Poison + } else { + ReleaseType::Normal + }; + let plan = self + .config + .constraints + .plan_release(db, release_type, &self.blocks); + trace!("Release: {conn:?} {plan:?}"); + match plan { + ReleaseOp::Release => {} + ReleaseOp::Discard => { + tokio::task::spawn_local(self.blocks.task_discard(&self.connector, conn)); + } + ReleaseOp::ReleaseTo(db) => { + tokio::task::spawn_local(self.blocks.task_move_to(&self.connector, conn, &db)); + } + ReleaseOp::Reopen => { + tokio::task::spawn_local(self.blocks.task_reopen(&self.connector, conn)); + } + } + } + + /// Retrieve the current pool metrics snapshot. + pub fn metrics(&self) -> PoolMetrics { + self.blocks.summary() + } + + /// Is this pool idle? + pub fn idle(&self) -> bool { + self.blocks.is_empty() + } + + /// Drain all connections to the given database. All connections will be + /// poisoned on return and this method will return when the given database + /// is idle. Multiple calls to this method with the same database are valid, + /// and the drain operation will be kept alive as long as one future has not + /// been dropped. + /// + /// It is valid, though unadvisable, to request a connection during this + /// period. The connection will be poisoned on return as well. + /// + /// Dropping this future cancels the drain operation. + pub async fn drain(self: Rc, db: &str) { + // If the block doesn't exist, we can return + let Some(name) = self.blocks.name(db) else { + return; + }; + + let lock = Drain::lock(self.clone(), name); + while self.blocks.metrics(db).total() > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + drop(lock); + } + + /// Drain all connections in the pool, returning when the pool is completely + /// empty. Multiple calls to this method with the same database are valid, + /// and the drain operation will be kept alive as long as one future has not + /// been dropped. + /// + /// It is valid, though unadvisable, to request a connection during this + /// period. The connection will be poisoned on return as well. + /// + /// Dropping this future cancels the drain operation. + pub async fn drain_all(self: Rc) { + let lock = Drain::lock_all(self.clone()); + while self.blocks.total() > 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + drop(lock); + } + + /// Shuts this pool down safely. Dropping this future does not cancel + /// the shutdown operation. + pub async fn shutdown(mut self: Rc) { + self.drain.shutdown(); + let pool = loop { + match Rc::try_unwrap(self) { + Ok(pool) => break pool, + Err(pool) => self = pool, + }; + tokio::time::sleep(Duration::from_millis(10)).await; + }; + while !pool.idle() { + pool.run_once(); + tokio::time::sleep(Duration::from_millis(10)).await; + } + if cfg!(debug_assertions) { + let all_time = &pool.metrics().all_time; + assert_eq!( + all_time[MetricVariant::Connecting], + all_time[MetricVariant::Disconnecting], + "Connecting != Disconnecting" + ); + assert_eq!( + all_time[MetricVariant::Disconnecting], + all_time[MetricVariant::Closed], + "Disconnecting != Closed" + ); + } + } +} + +impl AsRef for Rc> { + fn as_ref(&self) -> &Drain { + &self.drain + } +} + +/// Holds the current drainage and shutdown state for the `Pool`. +#[derive(Default, Debug)] +struct Drain { + drain_all: Cell, + drain: RefCell>, + shutdown: Cell, +} + +impl Drain { + pub fn shutdown(&self) { + self.shutdown.set(true) + } + + /// Lock all connections for draining. + pub fn lock_all>(this: T) -> DrainLock { + let drain = this.as_ref(); + drain.drain_all.set(drain.drain_all.get() + 1); + DrainLock { + db: None, + has_drain: this, + } + } + + // Lock a specific connection for draining. + pub fn lock>(this: T, db: Name) -> DrainLock { + { + let mut drain = this.as_ref().drain.borrow_mut(); + drain.entry(db.clone()).and_modify(|v| *v += 1).or_default(); + } + DrainLock { + db: Some(db), + has_drain: this, + } + } + + /// Is this connection draining? + fn is_draining(&self, db: &str) -> bool { + self.drain_all.get() > 0 || self.drain.borrow().contains_key(db) || self.shutdown.get() + } +} + +/// Provides a RAII lock for a db- or whole-pool drain operation. +struct DrainLock> { + db: Option, + has_drain: T, +} + +impl> Drop for DrainLock { + fn drop(&mut self) { + if let Some(name) = self.db.take() { + let mut drain = self.has_drain.as_ref().drain.borrow_mut(); + if let Some(count) = drain.get_mut(&name) { + if *count > 1 { + *count -= 1; + } else { + drain.remove(&name); + } + } else { + unreachable!() + } + } else { + let this = self.has_drain.as_ref(); + this.drain_all.set(this.drain_all.get() - 1); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::test::*; + use crate::time::Instant; + use anyhow::{Ok, Result}; + use itertools::Itertools; + use rstest::rstest; + + use test_log::test; + use tokio::task::LocalSet; + use tracing::{error, info, trace}; + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_pool_basic() -> Result<()> { + LocalSet::new() + .run_until(async { + let config = PoolConfig::suggested_default_for(10); + + let pool = Pool::new(config, BasicConnector::no_delay()); + let conn1 = pool.acquire("1").await?; + let conn2 = pool.acquire("1").await?; + + drop(conn1); + conn2.poison(); + drop(conn2); + + pool.shutdown().await; + + Ok(()) + }) + .await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_pool_eventually_idles() -> Result<()> { + let future = async { + let config = PoolConfig::suggested_default_for(10); + + let pool = Pool::new(config, BasicConnector::no_delay()); + let conn = pool.acquire("1").await?; + tokio::time::sleep(Duration::from_millis(10)).await; + drop(conn); + + while !pool.idle() { + tokio::time::sleep(Duration::from_millis(10)).await; + pool.run_once(); + } + trace!("Pool idle, shutting down"); + + pool.shutdown().await; + Ok(()) + }; + tokio::time::timeout(Duration::from_secs(120), LocalSet::new().run_until(future)).await? + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_pool_drains() -> Result<()> { + let future = async { + let config = PoolConfig::suggested_default_for(10); + + let pool = Pool::new(config, BasicConnector::no_delay()); + let conn = pool.acquire("1").await?; + tokio::task::spawn_local(pool.clone().drain_all()); + tokio::task::spawn_local(async { + tokio::time::sleep(Duration::from_millis(10)).await; + drop(conn); + }); + + while !pool.idle() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + trace!("Pool idle, shutting down"); + + pool.shutdown().await; + Ok(()) + }; + tokio::time::timeout(Duration::from_secs(120), LocalSet::new().run_until(future)).await? + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + #[rstest] + #[case::one(1)] + #[case::small(10)] + #[case::medium(12)] + #[case::large(20)] + async fn test_pool(#[case] databases: usize) -> Result<()> { + let spec = Spec { + name: format!("test_pool_{databases}").into(), + desc: "", + capacity: 10, + conn_cost: Triangle(0.05, 0.0), + score: vec![ + Score::new( + 0.8, + [2.0, 0.5, 0.25, 0.0], + LatencyDistribution { + group: 0..databases, + }, + ), + Score::new(0.2, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs: (0..databases) + .map(|db| DBSpec { + db, + start_at: 0.0, + end_at: 1.0, + qps: 1200, + query_cost: Triangle(0.001, 0.0), + }) + .collect_vec(), + ..Default::default() + }; + + run(spec).await.map(drop) + } + + async fn run(spec: Spec) -> Result { + let local = LocalSet::new(); + let res = local.run_until(run_local(spec)).await?; + local.await; + Ok(res) + } + + async fn run_local(spec: Spec) -> std::result::Result { + let start = Instant::now(); + let real_time = std::time::Instant::now(); + let config = PoolConfig::suggested_default_for(spec.capacity); + let disconnect_cost = spec.disconn_cost; + let connect_cost = spec.disconn_cost; + let pool = Pool::new( + config, + BasicConnector::delay(move |disconnect| { + if disconnect { + disconnect_cost.random_duration() + } else { + connect_cost.random_duration() + } + }), + ); + let mut tasks = vec![]; + let latencies = Latencies::default(); + + // Boot a task for each DBSpec in the Spec + for (i, db_spec) in spec.dbs.into_iter().enumerate() { + let interval = 1.0 / (db_spec.qps as f64); + info!("[{i:-2}] db {db_spec:?}"); + let db = format!("t{}", db_spec.db); + let pool = pool.clone(); + let latencies = latencies.clone(); + let local = async move { + let now = Instant::now(); + let count = ((db_spec.end_at - db_spec.start_at) * (db_spec.qps as f64)) as usize; + tokio::time::sleep(Duration::from_secs_f64(db_spec.start_at)).await; + info!( + "+[{i:-2}] Starting db {db} at {}qps (approx {}q·s/s from {}..{})...", + db_spec.qps, + db_spec.qps as f64 * db_spec.query_cost.0, + db_spec.start_at, + db_spec.end_at, + ); + let start_time = now.elapsed().as_secs_f64(); + // Boot one task for each expected query in a localset, with a + // sleep that schedules it for the appropriate time. + let local = LocalSet::new(); + for i in 0..count { + let pool = pool.clone(); + let latencies = latencies.clone(); + let duration = db_spec.query_cost.random_duration(); + let db = db.clone(); + local.spawn_local(async move { + tokio::time::sleep(Duration::from_secs_f64(i as f64 * interval)).await; + let now = Instant::now(); + let conn = pool.acquire(&db).await?; + let latency = now.elapsed(); + latencies.mark(&db, latency.as_secs_f64()); + tokio::time::sleep(duration).await; + drop(conn); + Ok(()) + }); + } + tokio::time::timeout(Duration::from_secs(120), local) + .await + .unwrap_or_else(move |_| error!("*[{i:-2}] DBSpec {i} for {db} timed out")); + let end_time = now.elapsed().as_secs_f64(); + info!("-[{i:-2}] Finished db t{} at {}qps. Load generated from {}..{}, processed from {}..{}", + db_spec.db, db_spec.qps, db_spec.start_at, db_spec.end_at, start_time, end_time); + }; + tasks.push(tokio::task::spawn_local(local)); + } + + // Boot the monitor the runs the pool algorithm and prints the current + // block connection stats. + let monitor = { + let pool = pool.clone(); + tokio::task::spawn_local(async move { + let mut orig = "".to_owned(); + loop { + pool.run_once(); + let mut s = "".to_owned(); + for (name, block) in pool.metrics().blocks { + s += &format!("{name}={} ", block.total); + } + if !s.is_empty() && s != orig { + trace!( + "Blocks: {}/{} {s}", + pool.metrics().pool.total, + pool.config.constraints.max + ); + orig = s; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + }; + + info!("Starting..."); + tokio::time::sleep(Duration::from_secs_f64(spec.duration)).await; + + for task in tasks { + _ = task.await; + } + + info!( + "Took {:?} of virtual time ({:?} real time)", + start.elapsed(), + real_time.elapsed() + ); + + monitor.abort(); + _ = monitor.await; + let metrics = pool.metrics(); + info!("{metrics:#?}"); + info!("{latencies:#?}"); + + let metrics = pool.metrics(); + let mut qos = 0.0; + let mut scores = vec![]; + for score in spec.score { + let scored = score.method.score(&latencies, &metrics, &pool.config); + + let score_component = score.calculate(scored.raw_value); + info!( + "[QoS: {}] {} = {:.2} -> {:.2} (weight {:.2})", + spec.name, scored.description, scored.raw_value, score_component, score.weight + ); + trace!( + "[QoS: {}] {} [detail]: {} = {:.3}", + spec.name, + scored.description, + (scored.detailed_calculation)(3), + scored.raw_value + ); + scores.push(WeightedScored { + scored, + weight: score.weight, + score: score_component, + }); + qos += score_component * score.weight; + } + info!("[QoS: {}] Score = {qos:0.02}", spec.name); + + info!("Shutting down..."); + pool.shutdown().await; + + Ok(QoS { scores, qos }) + } + + fn test_connpool_1() -> Spec { + let mut dbs = vec![]; + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.0, + end_at: 0.5, + qps: 50, + query_cost: Triangle(0.03, 0.005), + }) + } + for i in 6..12 { + dbs.push(DBSpec { + db: i, + start_at: 0.3, + end_at: 0.7, + qps: 50, + query_cost: Triangle(0.03, 0.005), + }) + } + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.6, + end_at: 0.8, + qps: 50, + query_cost: Triangle(0.03, 0.005), + }) + } + + Spec { + name: "test_connpool_1".into(), + desc: r#" + This is a test for Mode D, where 2 groups of blocks race for connections + in the pool with max capacity set to 6. The first group (0-5) has more + dedicated time with the pool, so it should have relatively lower latency + than the second group (6-11). But the QoS is focusing on the latency + distribution similarity, as we don't want to starve only a few blocks + because of the lack of capacity. Therefore, reconnection is a necessary + cost for QoS. + "#, + capacity: 6, + conn_cost: Triangle(0.05, 0.01), + score: vec![ + Score::new( + 0.18, + [2.0, 0.5, 0.25, 0.0], + LatencyDistribution { group: 0..6 }, + ), + Score::new( + 0.28, + [2.0, 0.3, 0.1, 0.0], + LatencyDistribution { group: 6..12 }, + ), + Score::new( + 0.48, + [2.0, 0.7, 0.45, 0.2], + LatencyDistribution { group: 0..12 }, + ), + Score::new(0.06, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs, + ..Default::default() + } + } + + fn test_connpool_2() -> Spec { + let mut dbs = vec![]; + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.0, + end_at: 0.5, + qps: 1500, + query_cost: Triangle(0.001, 0.005), + }) + } + for i in 6..12 { + dbs.push(DBSpec { + db: i, + start_at: 0.3, + end_at: 0.7, + qps: 700, + query_cost: Triangle(0.03, 0.001), + }) + } + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.6, + end_at: 0.8, + qps: 700, + query_cost: Triangle(0.06, 0.01), + }) + } + + Spec { + name: "test_connpool_2".into(), + desc: r#" + In this test, we have 6x1500qps connections that simulate fast + queries (0.001..0.006s), and 6x700qps connections that simulate + slow queries (~0.03s). The algorithm allocates connections + fairly to both groups, essentially using the + "demand = avg_query_time * avg_num_of_connection_waiters" + formula. The QoS is at the same level for all DBs. (Mode B / C) + "#, + capacity: 100, + conn_cost: Triangle(0.04, 0.011), + score: vec![ + Score::new( + 0.18, + [2.0, 0.5, 0.25, 0.0], + LatencyDistribution { group: 0..6 }, + ), + Score::new( + 0.28, + [2.0, 0.3, 0.1, 0.0], + LatencyDistribution { group: 6..12 }, + ), + Score::new( + 0.48, + [2.0, 0.7, 0.45, 0.2], + LatencyDistribution { group: 0..12 }, + ), + Score::new(0.06, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs, + ..Default::default() + } + } + + fn test_connpool_3() -> Spec { + let mut dbs = vec![]; + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.0, + end_at: 0.8, + qps: 5000, + query_cost: Triangle(0.01, 0.005), + }) + } + + Spec { + name: "test_connpool_3".into(), + desc: r#" + This test simply starts 6 same crazy requesters for 6 databases to + test the pool fairness in Mode C with max capacity of 100. + "#, + capacity: 100, + conn_cost: Triangle(0.04, 0.011), + score: vec![ + Score::new( + 0.85, + [1.0, 0.2, 0.1, 0.0], + LatencyDistribution { group: 0..6 }, + ), + Score::new(0.15, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs, + ..Default::default() + } + } + + fn test_connpool_4() -> Spec { + let mut dbs = vec![]; + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.0, + end_at: 0.8, + qps: 1000, + query_cost: Triangle(0.01 * (i as f64 + 1.0), 0.005 * (i as f64 + 1.0)), + }) + } + + Spec { + name: "test_connpool_4".into(), + desc: r#" + Similar to test 3, this test also has 6 requesters for 6 databases, + they have the same Q/s but with different query cost. In Mode C, + we should observe equal connection acquisition latency, fair and + stable connection distribution and reasonable reconnection cost. + "#, + capacity: 50, + conn_cost: Triangle(0.04, 0.011), + score: vec![ + Score::new( + 0.9, + [1.0, 0.2, 0.1, 0.0], + LatencyDistribution { group: 0..6 }, + ), + Score::new(0.1, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs, + ..Default::default() + } + } + + fn test_connpool_5() -> Spec { + let mut dbs = vec![]; + + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.0 + i as f64 / 10.0, + end_at: 0.5 + i as f64 / 10.0, + qps: 150, + query_cost: Triangle(0.020, 0.005), + }); + } + for i in 6..12 { + dbs.push(DBSpec { + db: i, + start_at: 0.3, + end_at: 0.7, + qps: 50, + query_cost: Triangle(0.008, 0.003), + }); + } + for i in 0..6 { + dbs.push(DBSpec { + db: i, + start_at: 0.6, + end_at: 0.8, + qps: 50, + query_cost: Triangle(0.003, 0.002), + }); + } + + Spec { + name: "test_connpool_5".into(), + desc: r#" + This is a mixed test with pool max capacity set to 6. Requests in + the first group (0-5) come and go alternatively as time goes on, + even with different query cost, so its latency similarity doesn't + matter much, as far as the latency distribution is not too crazy + and unstable. However the second group (6-11) has a stable + environment - pressure from the first group is quite even at the + time the second group works. So we should observe a high similarity + in the second group. Also due to a low query cost, the second group + should have a higher priority in connection acquisition, therefore + a much lower latency distribution comparing to the first group. + Pool Mode wise, we should observe a transition from Mode A to C, + then D and eventually back to C. One regression to be aware of is + that, the last D->C transition should keep the pool running at + a full capacity. + "#, + capacity: 6, + conn_cost: Triangle(0.15, 0.05), + score: vec![ + Score::new( + 0.05, + [2.0, 0.8, 0.4, 0.0], + LatencyDistribution { group: 0..6 }, + ), + Score::new( + 0.25, + [2.0, 0.8, 0.4, 0.0], + LatencyDistribution { group: 6..12 }, + ), + Score::new( + 0.45, + [1.0, 2.0, 5.0, 30.0], + LatencyRatio { + percentile: 75, + dividend: 0..6, + divisor: 6..12, + }, + ), + Score::new(0.15, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + Score::new(0.10, [3.0, 4.0, 5.0, 6.0], EndingCapacity {}), + ], + dbs, + ..Default::default() + } + } + + fn test_connpool_6() -> Spec { + let mut dbs = vec![]; + + for i in 0..6 { + dbs.push(DBSpec { + db: 0, + start_at: 0.0 + i as f64 / 10.0, + end_at: 0.5 + i as f64 / 10.0, + qps: 150, + query_cost: Triangle(0.020, 0.005), + }); + } + + Spec { + name: "test_connpool_6".into(), + desc: r#" + This is a simple test for Mode A. In this case, we don't want to + have lots of reconnection overhead. + "#, + capacity: 6, + conn_cost: Triangle(0.15, 0.05), + score: vec![Score::new(1.0, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {})], + dbs, + ..Default::default() + } + } + + fn test_connpool_7() -> Spec { + Spec { + name: "test_connpool_7".into(), + desc: r#" + The point of this test is to have one connection "t1" that + just has crazy demand for connections. Then the "t2" connections + are infrequent -- so they have a miniscule quota. + + Our goal is to make sure that "t2" has good QoS and gets + its queries processed as soon as they're submitted. Therefore, + "t2" should have way lower connection acquisition cost than "t1". + "#, + capacity: 6, + conn_cost: Triangle(0.15, 0.05), + score: vec![ + Score::new( + 0.2, + [1.0, 10.0, 50.0, 100.0], + LatencyRatio { + percentile: 99, + dividend: 1..2, + divisor: 2..3, + }, + ), + Score::new( + 0.4, + [1.0, 20.0, 100.0, 200.0], + LatencyRatio { + percentile: 75, + dividend: 1..2, + divisor: 2..3, + }, + ), + Score::new(0.4, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs: vec![ + DBSpec { + db: 1, + start_at: 0.0, + end_at: 1.0, + qps: 500, + query_cost: Triangle(0.040, 0.005), + }, + DBSpec { + db: 2, + start_at: 0.1, + end_at: 0.3, + qps: 30, + query_cost: Triangle(0.030, 0.005), + }, + DBSpec { + db: 2, + start_at: 0.6, + end_at: 0.9, + qps: 30, + query_cost: Triangle(0.010, 0.005), + }, + ], + ..Default::default() + } + } + + fn test_connpool_8() -> Spec { + let base_load = 200; + + Spec { + name: "test_connpool_8".into(), + desc: r#" + This test spec is to check the pool connection reusability with a + single block before the pool reaches its full capacity in Mode A. + We should observe just enough number of connects to serve the load, + while there can be very few disconnects because of GC. + "#, + capacity: 100, + conn_cost: Triangle(0.0, 0.0), + score: vec![Score::new(1.0, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {})], + dbs: vec![ + DBSpec { + db: 1, + start_at: 0.0, + end_at: 0.1, + qps: base_load / 4, + query_cost: Triangle(0.01, 0.0), + }, + DBSpec { + db: 1, + start_at: 0.1, + end_at: 0.2, + qps: base_load / 2, + query_cost: Triangle(0.01, 0.0), + }, + DBSpec { + db: 1, + start_at: 0.2, + end_at: 0.6, + qps: base_load, + query_cost: Triangle(0.01, 0.0), + }, + ], + ..Default::default() + } + } + + fn test_connpool_9() -> Spec { + let full_qps = 20000; + + Spec { + name: "test_connpool_9".into(), + desc: r#" + This test spec is to check the pool performance with low traffic + between 3 pre-heated blocks in Mode B. t1 is a reference block, + t2 has the same qps as t1, but t3 with doubled qps came in while t2 + is active. As the total throughput is low enough, we shouldn't have + a lot of connects and disconnects, nor a high acquire waiting time. + "#, + capacity: 100, + conn_cost: Triangle(0.01, 0.005), + score: vec![ + Score::new( + 0.1, + [2.0, 1.0, 0.5, 0.2], + LatencyDistribution { group: 1..4 }, + ), + Score::new( + 0.1, + [0.05, 0.004, 0.002, 0.001], + AbsoluteLatency { + group: 1..4, + percentile: 99, + }, + ), + Score::new( + 0.2, + [0.005, 0.0004, 0.0002, 0.0001], + AbsoluteLatency { + group: 1..4, + percentile: 75, + }, + ), + Score::new(0.6, [0.5, 0.2, 0.1, 0.0], ConnectionOverhead {}), + ], + dbs: vec![ + DBSpec { + db: 1, + start_at: 0.0, + end_at: 0.1, + qps: (full_qps / 32), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 1, + start_at: 0.1, + end_at: 0.4, + qps: (full_qps / 16), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 2, + start_at: 0.5, + end_at: 0.6, + qps: (full_qps / 32), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 2, + start_at: 0.6, + end_at: 1.0, + qps: (full_qps / 16), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 3, + start_at: 0.7, + end_at: 0.8, + qps: (full_qps / 16), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 3, + start_at: 0.8, + end_at: 0.9, + qps: (full_qps / 8), + query_cost: Triangle(0.01, 0.005), + }, + ], + ..Default::default() + } + } + + fn test_connpool_10() -> Spec { + let full_qps = 2000; + + Spec { + name: "test_connpool_10".into(), + desc: r#" + This test spec is to check the pool garbage collection feature. + t1 is a constantly-running reference block, t2 starts in the middle + with a full qps and ends early to leave enough time for the pool to + execute garbage collection. + "#, + timeout: 10, + duration: 2.0, + capacity: 100, + conn_cost: Triangle(0.01, 0.005), + score: vec![Score::new( + 1.0, + [100.0, 40.0, 20.0, 10.0], + EndingCapacity {}, + )], + dbs: vec![ + DBSpec { + db: 1, + start_at: 0.0, + end_at: 1.0, + qps: (full_qps / 32), + query_cost: Triangle(0.01, 0.005), + }, + DBSpec { + db: 2, + start_at: 0.4, + end_at: 0.6, + qps: ((full_qps / 32) * 31), + query_cost: Triangle(0.01, 0.005), + }, + ], + ..Default::default() + } + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn run_spec_tests() -> Result<()> { + spec_tests(None).await?; + Ok(()) + } + + async fn spec_tests(scale: Option) -> Result { + let mut results = SuiteQoS::default(); + for spec in SPEC_FUNCTIONS { + let mut spec = spec(); + if let Some(scale) = scale { + spec.scale(scale); + } + let name = spec.name.clone(); + let res = run(spec).await?; + results.insert(name, res); + } + for (name, QoS { qos, .. }) in &results { + info!("QoS[{name}] = [{qos:.02}]"); + } + info!( + "QoS = [{:.02}] (rms={:.02})", + results.qos(), + results.qos_rms_error() + ); + Ok(results) + } + + /// Runs the specs `count` times, returning the median run. + #[allow(unused)] + fn run_specs_tests_in_runtime(count: usize, scale: Option) -> Result { + let mut runs = vec![]; + for _ in 0..count { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let _guard = runtime.enter(); + tokio::time::pause(); + let qos = runtime.block_on(spec_tests(scale))?; + runs.push(qos); + } + runs.sort_by_cached_key(|run| (run.qos_rms_error() * 1_000_000.0) as usize); + let ret = runs.drain(count / 2..).next().unwrap(); + Ok(ret) + } + + #[test] + #[cfg(feature = "optimizer")] + fn optimizer() { + use crate::knobs::*; + use std::sync::atomic::AtomicIsize; + + use genetic_algorithm::strategy::evolve::prelude::*; + use lru::LruCache; + use rand::Rng; + + // the search goal to optimize towards (maximize or minimize) + #[derive(Clone, std::fmt::Debug, smart_default::SmartDefault)] + pub struct Optimizer { + #[default(std::sync::Arc::new(AtomicIsize::new(isize::MIN)))] + best: std::sync::Arc, + #[default(LruCache::new(100_000_000.try_into().unwrap()))] + lru: LruCache<[usize; ALL_KNOB_COUNT], isize>, + #[default(std::time::Instant::now())] + now: std::time::Instant, + } + + impl Fitness for Optimizer { + type Genotype = ContinuousGenotype; + fn calculate_for_chromosome( + &mut self, + chromosome: &Chromosome, + ) -> Option { + let mut knobs: [usize; ALL_KNOB_COUNT] = Default::default(); + for (knob, gene) in knobs.iter_mut().zip(&chromosome.genes) { + *knob = *gene as _; + } + if let Some(res) = self.lru.get(&knobs) { + return Some(*res); + } + + for (i, knob) in crate::knobs::ALL_KNOBS.iter().enumerate() { + if knob.set(knobs[i]).is_err() { + return None; + }; + } + + let real = rand::thread_rng().gen_range(0..1000) < 200; + let weights = if real { + [(1.0, 5, None), (0.5, 1, Some(10.0))] + } else { + [(1.0, 5, None), (0.5, 1, None)] + }; + let outputs = + weights.map(|(_, count, scale)| run_specs_tests_in_runtime(count, scale)); + let mut score = 0.0; + for ((weight, ..), output) in weights.iter().zip(&outputs) { + score += weight * output.as_ref().ok()?.qos_rms_error(); + } + let qos_i = (score * 1_000_000.0) as isize; + if real && qos_i > self.best.load(std::sync::atomic::Ordering::SeqCst) { + eprintln!("{:?} New best: {score:.02} {knobs:?}", self.now.elapsed()); + eprintln!("{:?}", crate::knobs::ALL_KNOBS); + for (weight, output) in weights.iter().zip(outputs) { + eprintln!("{weight:?}: {:?}", output.ok()?); + } + eprintln!("*****************************"); + self.best.store(qos_i, std::sync::atomic::Ordering::SeqCst); + } + self.lru.push(knobs, qos_i); + + Some(qos_i) + } + } + + let mut seeds: Vec> = vec![]; + + // The current state + seeds.push( + crate::knobs::ALL_KNOBS + .iter() + .map(|k| k.get() as _) + .collect(), + ); + + // A constant value for all knobs + for i in 0..100 { + seeds.push([i].repeat(crate::knobs::ALL_KNOBS.len())); + } + + // Some randomness + for _ in 0..100 { + seeds.push( + (0..crate::knobs::ALL_KNOBS.len()) + .map(|_| rand::thread_rng().gen_range(0..1000)) + .collect(), + ); + } + + let mut f32_seeds = vec![]; + for mut seed in seeds { + for (i, knob) in crate::knobs::ALL_KNOBS.iter().enumerate() { + let mut value = seed[i] as _; + if knob.set(value).is_err() { + knob.clamp(&mut value); + seed[i] = value as _; + }; + } + f32_seeds.push(seed.into_iter().map(|n| n as _).collect()); + } + + let genotype = ContinuousGenotype::builder() + .with_genes_size(crate::knobs::ALL_KNOBS.len()) + .with_allele_range(0.0..1000.0) + .with_allele_neighbour_ranges(vec![-50.0..50.0, -5.0..5.0]) + .with_seed_genes_list(f32_seeds) + .build() + .unwrap(); + + let mut rng = rand::thread_rng(); // a randomness provider implementing Trait rand::Rng + let evolve = Evolve::builder() + .with_multithreading(true) + .with_genotype(genotype) + .with_target_population_size(1000) + .with_target_fitness_score(100 * 1_000_000) + .with_max_stale_generations(1000) + .with_fitness(Optimizer::default()) + .with_crossover(CrossoverUniform::new(true)) + .with_mutate(MutateOnce::new(0.5)) + .with_compete(CompeteTournament::new(200)) + .with_extension(ExtensionMassInvasion::new(0.6, 0.6)) + .call(&mut rng) + .unwrap(); + println!("{}", evolve); + } + + macro_rules! run_spec { + ($($spec:ident),* $(,)?) => { + const SPEC_FUNCTIONS: [fn() -> Spec; [$( $spec ),*].len()] = [ + $( + $spec, + )* + ]; + + mod spec { + use super::*; + $( + #[super::test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn $spec() -> Result<()> { + run(super::$spec()).await.map(drop) + } + )* + } + }; + } + + run_spec!( + test_connpool_1, + test_connpool_2, + test_connpool_3, + test_connpool_4, + test_connpool_5, + test_connpool_6, + test_connpool_7, + test_connpool_8, + test_connpool_9, + test_connpool_10, + ); +} diff --git a/edb/server/conn_pool/src/python.rs b/edb/server/conn_pool/src/python.rs new file mode 100644 index 00000000000..8377b1779b9 --- /dev/null +++ b/edb/server/conn_pool/src/python.rs @@ -0,0 +1,6 @@ +use pyo3::{pymodule, types::PyModule, PyResult, Python}; + +#[pymodule] +fn _conn_pool(py: Python, m: &PyModule) -> PyResult<()> { + Ok(()) +} diff --git a/edb/server/conn_pool/src/test.rs b/edb/server/conn_pool/src/test.rs new file mode 100644 index 00000000000..1e8069e41e5 --- /dev/null +++ b/edb/server/conn_pool/src/test.rs @@ -0,0 +1,511 @@ +//! Test utilities. +use std::{ + borrow::Cow, + cell::{Cell, RefCell}, + collections::{BTreeMap, HashMap}, + future::Future, + ops::Range, + rc::Rc, + time::Duration, +}; + +use itertools::Itertools; +use rand::random; +use statrs::statistics::{Data, Distribution, OrderStatistics, Statistics}; + +use crate::{ + conn::{ConnResult, Connector}, + metrics::{MetricVariant, PoolMetrics}, + PoolConfig, +}; + +#[derive(derive_more::Debug)] +pub struct BasicConnector { + #[debug(skip)] + delay: Option Duration>>, +} + +impl BasicConnector { + pub fn no_delay() -> Self { + BasicConnector { delay: None } + } + pub fn delay(f: impl Fn(bool) -> Duration + 'static) -> Self { + BasicConnector { + delay: Some(Rc::new(f)), + } + } +} + +impl Connector for BasicConnector { + type Conn = (); + fn connect(&self, _db: &str) -> impl Future> + 'static { + let delay = self.delay.clone(); + async move { + if let Some(f) = delay { + tokio::time::sleep(f(false)).await; + } + Ok(()) + } + } + fn reconnect( + &self, + conn: Self::Conn, + _db: &str, + ) -> impl Future> + 'static { + let delay = self.delay.clone(); + async move { + if let Some(f) = delay { + tokio::time::sleep(f(true)).await; + tokio::time::sleep(f(false)).await; + } + Ok(conn) + } + } + fn disconnect(&self, _conn: Self::Conn) -> impl Future> + 'static { + let delay = self.delay.clone(); + async move { + if let Some(f) = delay { + tokio::time::sleep(f(true)).await; + } + Ok(()) + } + } +} + +#[derive(Clone, Default)] +pub struct Latencies { + data: Rc>>>, +} + +/// Helper function for [`Stats`] [`Debug`] impl. +#[allow(unused)] +fn m(v: &f64) -> Duration { + if *v <= 0.000_001 { + Duration::ZERO + } else { + Duration::from_secs_f64(*v) + } +} + +#[derive(derive_more::Debug)] +#[allow(unused)] +#[debug( + "#{count} %{{1,25,50,75,99}}: {:?}/{:?}/{:?}/{:?}/{:?}, x̄: {:?} Πx: {:?}", + m(p01), + m(p25), + m(p50), + m(p75), + m(p99), + m(mean), + m(geometric_mean) +)] +struct Stats { + p01: f64, + p25: f64, + p50: f64, + p75: f64, + p99: f64, + geometric_mean: f64, + mean: f64, + count: usize, +} + +impl Latencies { + pub fn mark(&self, db: &str, latency: f64) { + self.data + .borrow_mut() + .entry(db.to_owned()) + .or_default() + .push(latency.max(0.000_001)); + } + + fn len(&self) -> usize { + let mut len = 0; + for values in self.data.borrow().values() { + len += values.len() + } + len + } +} + +impl std::fmt::Debug for Latencies { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut s = f.debug_struct("Latencies"); + let mut data = self.data.borrow_mut(); + let mut all = vec![]; + for key in data.keys().cloned().sorted() { + let data = data.get_mut(&key).unwrap(); + all.extend_from_slice(data); + let stats = stats(data); + s.field(&key, &stats); + } + let stats = stats(&mut all); + s.field("all", &stats); + s.finish() + } +} + +fn stats(data: &mut [f64]) -> Stats { + let geometric_mean = data.geometric_mean(); + let mut data = Data::new(data); + let mean = data.mean().unwrap(); + + Stats { + p01: data.percentile(1), + p25: data.percentile(25), + p50: data.percentile(50), + p75: data.percentile(75), + p99: data.percentile(99), + geometric_mean, + mean, + count: data.len(), + } +} + +#[derive(smart_default::SmartDefault)] +pub struct Spec { + pub name: Cow<'static, str>, + pub desc: &'static str, + #[default = 30] + pub timeout: usize, + #[default = 1.1] + pub duration: f64, + pub capacity: usize, + pub conn_cost: Triangle, + pub dbs: Vec, + #[default(Triangle(0.006, 0.0015))] + pub disconn_cost: Triangle, + pub score: Vec, +} + +impl Spec { + pub fn scale(&mut self, time_scale: f64) { + self.duration *= time_scale; + for db in &mut self.dbs { + db.scale(time_scale); + } + } +} + +#[derive(derive_more::Debug)] +pub struct Scored { + pub description: String, + #[debug(skip)] + pub detailed_calculation: Box String>, + pub raw_value: f64, +} + +#[derive(Debug)] +pub struct WeightedScored { + pub weight: f64, + pub score: f64, + pub scored: Scored, +} + +#[derive(Debug)] +pub struct QoS { + pub scores: Vec, + pub qos: f64, +} + +#[derive(Default, derive_more::Deref, derive_more::DerefMut, derive_more::IntoIterator)] +pub struct SuiteQoS(#[into_iterator(owned, ref, ref_mut)] BTreeMap, QoS>); + +impl std::fmt::Debug for SuiteQoS { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let mut s = f.debug_struct("SuiteQos"); + for (name, qos) in self { + s.field(name, &format!("QoS = {:.02}", qos.qos)); + } + s.field("qos", &self.qos()); + s.field("qos_rms", &self.qos_rms_error()); + s.finish() + } +} + +impl SuiteQoS { + pub fn qos(&self) -> f64 { + let mut total = 0.0; + for qos in self.values() { + total += qos.qos; + } + total /= self.len() as f64; + if !total.is_normal() || total < 0.0 { + 0.0 + } else { + total + } + } + + /// Return the root-mean-square error QoS. The error between the QoS and 100 + /// is squared, averaged, and we subtract that from 100 for a final score. + pub fn qos_rms_error(&self) -> f64 { + let mut total = 0.0; + for qos in self.values() { + total += (100.0 - qos.qos).powf(2.0); + } + total /= self.len() as f64; + total = 100.0 - total.sqrt(); + if !total.is_normal() || total < 0.0 { + 0.0 + } else { + total + } + } +} + +pub trait ScoringMethod { + fn score(&self, latencies: &Latencies, metrics: &PoolMetrics, config: &PoolConfig) -> Scored; +} + +pub struct Score { + pub v100: f64, + pub v90: f64, + pub v60: f64, + pub v0: f64, + pub weight: f64, + pub method: Box, +} + +impl Score { + pub fn new(weight: f64, scores: [f64; 4], method: impl ScoringMethod + 'static) -> Self { + Self { + weight, + v0: scores[0], + v60: scores[1], + v90: scores[2], + v100: scores[3], + method: method.into(), + } + } + + pub fn calculate(&self, value: f64) -> f64 { + if value.is_nan() || value.is_infinite() { + return 0.0; + } + + let intervals = [ + (self.v100, self.v90, 90.0, 10.0), + (self.v90, self.v60, 60.0, 30.0), + (self.v60, self.v0, 0.0, 60.0), + ]; + + for &(v1, v2, base, diff) in &intervals { + let v_min = v1.min(v2); + let v_max = v1.max(v2); + if v_min <= value && value < v_max { + return base + (value - v2).abs() / (v_max - v_min) * diff; + } + } + + if self.v0 > self.v100 { + if value < self.v100 { + 100.0 + } else { + 0.0 + } + } else if value < self.v0 { + 0.0 + } else { + 100.0 + } + } +} + +pub struct LatencyDistribution { + pub group: Range, +} + +impl ScoringMethod for LatencyDistribution { + fn score(&self, latencies: &Latencies, _metrics: &PoolMetrics, _config: &PoolConfig) -> Scored { + let dbs = self.group.clone().map(|t| format!("t{t}")).collect_vec(); + let mut data = latencies.data.borrow_mut(); + let fail = Cell::new(false); + + // Calculates the average CV (coefficient of variation) of the given + // distributions. The result is a float ranging from zero indicating how + // different the given distributions are, where zero means no + // difference. Known defect: when the mean value is close to zero, the + // coefficient of variation will approach infinity and is therefore + // sensitive to small changes. + let values = (1..=9) + .map(move |n| { + let decile = Data::new( + dbs.iter() + .map(|db| { + let Some(data) = data.get_mut(db) else { + fail.set(true); + return 0.0; + }; + let mut data = Data::new(data.as_mut_slice()); + // This is equivalent to Python's statistics.quartile(n=10) + data.percentile(n * 10) + }) + .collect_vec(), + ); + let cv = decile.std_dev().unwrap_or_default() / decile.mean().unwrap_or_default(); + if cv.is_normal() { + cv + } else { + 0.0 + } + }) + .collect_vec(); + let mean = values.iter().geometric_mean(); + Scored { + description: format!("Average CV for range {:?}", self.group), + detailed_calculation: Box::new(move |precision| format!("{values:.precision$?}")), + raw_value: mean, + } + } +} + +impl From for Box { + fn from(value: T) -> Self { + Box::new(value) + } +} + +pub struct ConnectionOverhead {} + +impl ScoringMethod for ConnectionOverhead { + fn score(&self, latencies: &Latencies, metrics: &PoolMetrics, config: &PoolConfig) -> Scored { + let disconnects = metrics.all_time[MetricVariant::Disconnecting]; + // Calculate the GC + let max = config.constraints.max; + let total = metrics.pool.total; + let gc = max - total; + let disconnects_adj = disconnects.saturating_sub(gc); + let count = latencies.len(); + let raw_value = disconnects_adj as f64 / count as f64; + Scored { + description: "Num of disconnects/query".to_owned(), + detailed_calculation: Box::new(move |_precision| { + format!("({disconnects}-({max}-{total}))/{count}") + }), + raw_value, + } + } +} + +pub struct LatencyRatio { + pub percentile: u8, + pub dividend: Range, + pub divisor: Range, +} + +impl ScoringMethod for LatencyRatio { + fn score(&self, latencies: &Latencies, _metrics: &PoolMetrics, _config: &PoolConfig) -> Scored { + let mut data = latencies.data.borrow_mut(); + let dbs = self.divisor.clone().map(|t| format!("t{t}")).collect_vec(); + let divisor = dbs + .iter() + .map(|db| { + let Some(data) = data.get_mut(db) else { + return f64::NAN; + }; + let mut data = Data::new(data.as_mut_slice()); + data.percentile(self.percentile as _) + }) + .mean(); + let dbs = self.dividend.clone().map(|t| format!("t{t}")).collect_vec(); + let dividend = dbs + .iter() + .map(|db| { + let Some(data) = data.get_mut(db) else { + return f64::NAN; + }; + let mut data = Data::new(data.as_mut_slice()); + data.percentile(self.percentile as _) + }) + .mean(); + let raw_value = dividend / divisor; + Scored { + description: format!( + "P{} ratio {:?}/{:?}", + self.percentile, self.dividend, self.divisor + ), + detailed_calculation: Box::new(move |precision| { + format!("{dividend:.precision$}/{divisor:.precision$}") + }), + raw_value, + } + } +} + +pub struct EndingCapacity {} + +impl ScoringMethod for EndingCapacity { + fn score(&self, _latencies: &Latencies, metrics: &PoolMetrics, _config: &PoolConfig) -> Scored { + let total = metrics.pool.total; + let raw_value = total as _; + Scored { + description: "Ending capacity".to_string(), + detailed_calculation: Box::new(move |_| format!("{total}")), + raw_value, + } + } +} + +pub struct AbsoluteLatency { + pub percentile: u8, + pub group: Range, +} + +impl ScoringMethod for AbsoluteLatency { + fn score(&self, latencies: &Latencies, _metrics: &PoolMetrics, _config: &PoolConfig) -> Scored { + let mut data = latencies.data.borrow_mut(); + let dbs = self.group.clone().map(|t| format!("t{t}")).collect_vec(); + let raw_value = dbs + .iter() + .map(|db| { + let Some(data) = data.get_mut(db) else { + return f64::NAN; + }; + let mut data = Data::new(data.as_mut_slice()); + data.percentile(self.percentile as _) + }) + .mean(); + + Scored { + description: format!("Absolute P{} value {:?}", self.percentile, self.group), + detailed_calculation: Box::new(move |precision| format!("{raw_value:.precision$}")), + raw_value, + } + } +} + +#[derive(Debug)] +pub struct DBSpec { + pub db: usize, + pub start_at: f64, + pub end_at: f64, + pub qps: usize, + pub query_cost: Triangle, +} + +impl DBSpec { + pub fn scale(&mut self, time_scale: f64) { + self.start_at *= time_scale; + self.end_at *= time_scale; + } +} + +#[derive(Default, derive_more::Debug, Clone, Copy)] +#[debug("{0:?}±{1:?}", Duration::from_secs_f64(self.0), Duration::from_secs_f64(self.1))] +pub struct Triangle(pub f64, pub f64); + +impl Triangle { + pub fn random(&self) -> f64 { + self.0 + (random::() * 2.0 - 1.0) * self.1 + } + + pub fn random_duration(&self) -> Duration { + let r = self.random(); + if r <= 0.001 { + Duration::from_millis(1) + } else { + Duration::from_secs_f64(r) + } + } +} diff --git a/edb/server/conn_pool/src/waitqueue.rs b/edb/server/conn_pool/src/waitqueue.rs new file mode 100644 index 00000000000..5f00792c130 --- /dev/null +++ b/edb/server/conn_pool/src/waitqueue.rs @@ -0,0 +1,104 @@ +use crate::time::Instant; +use scopeguard::defer; +use std::{ + cell::{Cell, RefCell}, + collections::VecDeque, + future::poll_fn, + rc::Rc, + task::{Poll, Waker}, + time::Duration, +}; +use tracing::trace; + +struct WaitObject { + waker: Waker, + woke: Cell, + gc: Cell, + when: Instant, +} + +/// Maintains a list of waiters for a given object. Similar to tokio's `Notify` +/// but explicitly not thread-safe. +pub struct WaitQueue { + waiters: RefCell>>, + pub(crate) lock: Cell, +} + +impl Default for WaitQueue { + fn default() -> Self { + Self::new() + } +} + +impl WaitQueue { + pub fn new() -> Self { + Self { + waiters: RefCell::default(), + lock: Cell::default(), + } + } + + pub fn trigger(&self) { + loop { + if let Some(front) = self.waiters.borrow_mut().pop_front() { + if front.gc.get() { + trace!("Tossing away a GC'd entry"); + continue; + } + trace!("Triggered a waiter"); + front.woke.set(true); + front.waker.wake_by_ref(); + } else { + trace!("No waiters to trigger"); + } + break; + } + } + + pub async fn queue(&self) { + trace!("Queueing"); + let waker = poll_fn(|cx| Poll::Ready(cx.waker().clone())).await; + + let entry = Rc::new(WaitObject { + waker, + gc: Cell::default(), + woke: Cell::default(), + when: Instant::now(), + }); + + self.waiters.borrow_mut().push_back(entry.clone()); + + defer! { + entry.gc.set(true); + } + + poll_fn(|_cx| { + if entry.woke.get() { + Poll::Ready(()) + } else { + Poll::Pending + } + }) + .await; + } + + pub fn len(&self) -> usize { + self.waiters.borrow().len() + } + + pub(crate) fn lock(&self) { + self.lock.set(self.lock.get() + 1); + } + + pub(crate) fn unlock(&self) { + self.lock.set(self.lock.get() - 1); + } + + pub(crate) fn oldest(&self) -> Duration { + if let Some(entry) = self.waiters.borrow().front() { + entry.when.elapsed() + } else { + Duration::default() + } + } +} diff --git a/edb/server/connpool/__init__.py b/edb/server/connpool/__init__.py index 577dffa0fcf..63f400dfeec 100644 --- a/edb/server/connpool/__init__.py +++ b/edb/server/connpool/__init__.py @@ -17,6 +17,6 @@ # from .pool import Pool, _NaivePool # NoQA +from .pool2 import Pool as Pool2 - -__all__ = ('Pool',) +__all__ = ('Pool', 'Pool2') diff --git a/edb/server/connpool/pool2.py b/edb/server/connpool/pool2.py new file mode 100644 index 00000000000..d4f994ac448 --- /dev/null +++ b/edb/server/connpool/pool2.py @@ -0,0 +1,5 @@ +import edb.server._conn_pool # noqa: F401 + + +class Pool: + pass diff --git a/setup.py b/setup.py index ed5144e94a1..fa4ae077f32 100644 --- a/setup.py +++ b/setup.py @@ -1126,5 +1126,11 @@ def _version(): path="edb/graphql-rewrite/Cargo.toml", binding=setuptools_rust.Binding.PyO3, ), + setuptools_rust.RustExtension( + "edb.server._conn_pool", + path="edb/server/conn_pool/Cargo.toml", + features=["python_extension"], + binding=setuptools_rust.Binding.PyO3, + ), ], )