From 513e08da1c7479c3cc76c28fa061cd65fe0e5657 Mon Sep 17 00:00:00 2001 From: Matt Mastracci Date: Thu, 13 Jun 2024 13:52:19 -0600 Subject: [PATCH] wip --- Cargo.lock | 1043 +++++++++++++++++++++- Cargo.toml | 1 + edb/server/conn_pool/Cargo.toml | 47 + edb/server/conn_pool/README.md | 112 +++ edb/server/conn_pool/src/algo.rs | 670 ++++++++++++++ edb/server/conn_pool/src/block.rs | 914 +++++++++++++++++++ edb/server/conn_pool/src/conn.rs | 336 +++++++ edb/server/conn_pool/src/lib.rs | 24 + edb/server/conn_pool/src/metrics.rs | 373 ++++++++ edb/server/conn_pool/src/pool.rs | 1184 +++++++++++++++++++++++++ edb/server/conn_pool/src/python.rs | 380 ++++++++ edb/server/conn_pool/src/test.rs | 417 +++++++++ edb/server/conn_pool/src/waitqueue.rs | 104 +++ edb/server/connpool/__init__.py | 4 +- edb/server/connpool/pool.py | 10 + edb/server/connpool/pool2.py | 275 ++++++ setup.py | 6 + tests/test_server_pool.py | 68 +- 18 files changed, 5913 insertions(+), 55 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 4b6901feb369..852ceea87163 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,12 +2,100 @@ # 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 = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[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" @@ -20,6 +108,21 @@ 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 +186,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 +204,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 +235,45 @@ dependencies = [ "unreachable", ] +[[package]] +name = "conn-pool" +version = "0.1.0" +dependencies = [ + "anyhow", + "consume_on_drop", + "derive_more", + "futures", + "itertools", + "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" @@ -133,6 +293,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 +422,128 @@ 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.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 = "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 +554,29 @@ dependencies = [ "version_check", ] +[[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" @@ -287,6 +608,18 @@ 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.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" + [[package]] name = "indexmap" version = "1.9.3" @@ -313,18 +646,45 @@ 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.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 +701,25 @@ version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "90ed8c1e510134f979dbc4f070f87d4313098b704861a105fe34231c70a3901c" +[[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 +735,54 @@ 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-bigint" version = "0.2.6" @@ -390,6 +817,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 +835,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.18" @@ -406,26 +852,52 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", + "libm", ] [[package]] -name = "once_cell" -version = "1.19.0" +name = "num_cpus" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb12b2476b595f9358c5161aa467c2438859caa136dec86c26fdd2efe17b92" +checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" +dependencies = [ + "hermit-abi", + "libc", +] [[package]] -name = "parking_lot" -version = "0.12.1" +name = "object" +version = "0.36.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +checksum = "576dfe1fc8f9df304abb159d767a29d0476f7750fbf8aa7ad07816004a207434" dependencies = [ - "lock_api", - "parking_lot_core", + "memchr", ] [[package]] -name = "parking_lot_core" +name = "once_cell" +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" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3742b2c103b9f06bc9fff0a37ff4912935851bee6d36f3c02bcc755bcfec228f" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" version = "0.9.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4c42a9226546d68acdd9c0a280d17ce19bfe27a46bf68784e4066115788d008e" @@ -434,9 +906,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 +957,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 +991,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 +1065,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,6 +1087,18 @@ 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_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", "rand_core", ] @@ -590,6 +1107,25 @@ 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 = "redox_syscall" @@ -600,18 +1136,134 @@ 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 = "719b953e2095829ee67db738b3bfa9fa368c94900df327b3f07fe6e794d2fe1f" + +[[package]] +name = "rustc_version" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366" +dependencies = [ + "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 +1307,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 +1376,43 @@ 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 = "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 +1450,28 @@ version = "0.12.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1fc403891a21bcfb7c37834ba66a547a8f402146eba7265b5a6d88059c9ff2f" +[[package]] +name = "test-log" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dffced63c2b5c7be278154d76b479f9f9920ed34e7574201407f0b14e2bbb93" +dependencies = [ + "env_logger", + "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 +1492,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 +1615,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 +1648,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 +1678,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 +1740,76 @@ 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-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,42 +1818,99 @@ 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" diff --git a/Cargo.toml b/Cargo.toml index 6c9a653dff0f..18a287a159fa 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 000000000000..ea6ba28ecb49 --- /dev/null +++ b/edb/server/conn_pool/Cargo.toml @@ -0,0 +1,47 @@ +[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"] + +[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" + +[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 000000000000..37784975495a --- /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 000000000000..73872d6e1ed3 --- /dev/null +++ b/edb/server/conn_pool/src/algo.rs @@ -0,0 +1,670 @@ +use scopeguard::defer; +use std::{ + cell::{Cell, RefCell}, + num::NonZeroUsize, +}; +use tracing::trace; + +use crate::{ + block::Name, + metrics::{MetricVariant, RollingAverageU32}, +}; + +/// The maximum number of connections to create during a rebalance. +const MAX_REBALANCE_CREATE: usize = 5; +/// The maximum number of excess connections (> target) we'll keep around during +/// a rebalance if there is still some demand. +const MAX_EXCESS_IDLE_CONNECTIONS: usize = 2; +/// The minimum amount of time we'll consider for an active connection. +const MIN_ACTIVE_TIME: usize = 1; + +/// The weight we apply to waiting connections. +const DEMAND_WEIGHT_WAITING: usize = 1; +/// The weight we apply to active connections. +const DEMAND_WEIGHT_ACTIVE: usize = 1; +/// The historical length of data we'll maintain for demand. +const DEMAND_HISTORY_LENGTH: usize = 4; +/// 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. +const DEMAND_MINIMUM: usize = 16; + +/// 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 = 10; + +/// The weight we apply to the difference between the target and required +/// connections when determining overfullness. +const HUNGER_DIFF_WEIGHT: usize = 100; +/// The weight we apply to waiters when determining hunger. +const HUNGER_WAITER_WEIGHT: usize = 1; +/// The weight we apply to the oldest waiter's age in milliseconds (as a divisor). +const HUNGER_AGE_DIVISOR_WEIGHT: usize = 10; + +/// The weight we apply to the difference between the target and required +/// connections when determining overfullness. +const OVERFULL_DIFF_WEIGHT: usize = 100; +/// The weight we apply to idle connections when determining overfullness. +const OVERFULL_IDLE_WEIGHT: usize = 1; +/// This is divided by the youngest connection metric to penalize switching from +/// a backend which has changed recently. +const OVERFULL_CHANGE_WEIGHT_DIVIDEND: usize = 1000; + +/// 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(); + + // Waiters become more hungry as they age + let age_score = self.oldest_ms(MetricVariant::Waiting) / HUNGER_AGE_DIVISOR_WEIGHT; + let base_score = age_score + waiters * HUNGER_WAITER_WEIGHT; + + // 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 { + None + } else if target > current { + let diff = target - current; + (base_score + diff * HUNGER_DIFF_WEIGHT).try_into().ok() + } else if waiters > 0 { + base_score.try_into().ok() + } else { + None + } + } + + /// 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 youngest_ms = self.youngest_ms().max(1); + // If we have no idle connections, or we don't have enough connections we're not overfull. + if target >= current || idle == 0 { + None + } else { + let base_score = + idle * OVERFULL_IDLE_WEIGHT + OVERFULL_CHANGE_WEIGHT_DIVIDEND / youngest_ms; + if current > target { + let diff = current - target; + (diff * OVERFULL_DIFF_WEIGHT + base_score).try_into().ok() + } else { + base_score.try_into().ok() + } + } + } +} + +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; + 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; + let mut total_target = 0; + let mut s = "".to_owned(); + + it.with_all(|name, data| { + // 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. + let active = data.max(MetricVariant::Active); + let active_ms = data.avg_ms(MetricVariant::Active).max(MIN_ACTIVE_TIME); + let waiting = data.max(MetricVariant::Waiting); + let idle = active == 0 && waiting == 0; + let demand = if idle { + 0 + } else { + // Note that we clamp to DEMAND_MINIMUM to ensure the average is non-zero + (active_ms * (waiting * DEMAND_WEIGHT_WAITING + active * DEMAND_WEIGHT_ACTIVE)) + .max(DEMAND_MINIMUM) + }; + data.insert_demand(demand as _); + let demand_avg = data.demand(); + + if tracing::enabled!(tracing::Level::TRACE) { + s += &format!( + "{name}={demand_avg}/{demand} (a={},w={},t={}ms) ", + active, waiting, active_ms + ); + } + + total_demand += demand_avg; + 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: u32) { + // 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); + // The remaining capacity after we allocated the `min` value above. + let capacity = self.max - min * total_target; + + 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_CREATE { + it.with_all(|name, block| { + if block.target() > block.total() && current_pool_size < max_pool_size { + // If we are allocated more connections than we currently have, + // we'll try to grab some more. + 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_EXCESS_IDLE_CONNECTIONS || 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 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} "); + } + 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![]; + + // TODO: rebalance more than one? + loop { + let Some((_, to)) = hungriest.pop() else { + // TODO: close more than one? + if let Some((_, from)) = overloaded.pop() { + tasks.push(RebalanceOp::Close(from.clone())); + } + break; + }; + + let Some((_, from)) = overloaded.pop() else { + break; + }; + + tasks.push(RebalanceOp::Transfer { to, from }); + } + + 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) { + 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 = 0; + let mut which = None; + it.with_all(|name, block| { + if let Some(overfullness) = block.overfull_score(false) { + let overfullness: usize = overfullness.into(); + 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 = 0; + let mut which = None; + let mut s = "".to_owned(); + it.with_all(|name, block| { + let is_self = &**name == db; + if let Some(hunger) = block.hunger_score(is_self) { + let mut hunger: usize = hunger.into(); + // Penalize switching by boosting the current database's relative hunger here + if is_self { + hunger += SELF_HUNGER_BOOST_FOR_RELEASE; + } + + if tracing::enabled!(tracing::Level::TRACE) { + s += &format!("{name}={hunger} "); + } + // If this current block has equal hunger to the hungriest, it takes priority + 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 000000000000..dd0be3563ada --- /dev/null +++ b/edb/server/conn_pool/src/block.rs @@ -0,0 +1,914 @@ +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); + } + } + + #[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 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 + 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(()) + } + } + + 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(()) + } + } + + 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(()) + } + } + + 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(()) + } + } + + 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 000000000000..1ca399a56d24 --- /dev/null +++ b/edb/server/conn_pool/src/conn.rs @@ -0,0 +1,336 @@ +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); + 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.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(); + match &mut *lock { + ConnInner::Idle(..) => Poll::Ready(Ok(())), + ConnInner::Connecting(t, f) => Poll::Ready(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) => Poll::Ready(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 => Poll::Ready(Err(ConnError::Other("Failed".into()))), + _ => unreachable!(), + } + } + + 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()), + } + } +} + +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. + 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 000000000000..feeb984379dc --- /dev/null +++ b/edb/server/conn_pool/src/lib.rs @@ -0,0 +1,24 @@ +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; +} + +// 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 000000000000..bd82a33f8d2c --- /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 000000000000..5ed51af7ed22 --- /dev/null +++ b/edb/server/conn_pool/src/pool.rs @@ -0,0 +1,1184 @@ +use crate::{ + algo::{ + AcquireOp, PoolAlgoTargetData, PoolAlgorithmDataBlock, PoolAlgorithmDataMetrics, + PoolConstraints, RebalanceOp, ReleaseOp, ReleaseType, VisitPoolAlgoData, + }, + block::{Blocks, Name}, + conn::{ConnError, ConnHandle, ConnResult, Connector}, + metrics::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(), + }) + } + + /// 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(); + while let Err(pool_) = Rc::try_unwrap(self) { + self = pool_; + tokio::time::sleep(Duration::from_millis(10)).await; + } + } +} + +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::{info, trace}; + + #[test(tokio::test)] + 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 + } + + async fn run(spec: Spec) -> Result<()> { + let local = LocalSet::new(); + local.run_until(run_local(spec)).await?; + local.await; + Ok(()) + } + + async fn run_local(spec: Spec) -> std::result::Result<(), anyhow::Error> { + 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; + virtual_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 { + virtual_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()); + virtual_sleep(duration).await; + drop(conn); + Ok(()) + }); + } + tokio::time::timeout(Duration::from_secs(120), local) + .await + .unwrap_or_else(move |_| panic!("*[{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; + } + virtual_sleep(Duration::from_millis(10)).await; + } + }) + }; + + info!("Starting..."); + virtual_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; + 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 + ); + qos += score_component * score.weight; + } + info!("[QoS: {}] Score = {qos:0.02}", spec.name); + + info!("Shutting down..."); + pool.shutdown().await; + + Ok(()) + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_1() -> Result<()> { + 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), + }) + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_2() -> Result<()> { + 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), + }) + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_3() -> Result<()> { + 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), + }) + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_4() -> Result<()> { + 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)), + }) + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_5() -> Result<()> { + 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), + }); + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_6() -> Result<()> { + 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), + }); + } + + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_7() -> Result<()> { + let 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_8() -> Result<()> { + let base_load = 200; + let spec = 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() + }; + + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_9() -> Result<()> { + let full_qps = 20000; + let spec = 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() + }; + run(spec).await + } + + #[test(tokio::test(flavor = "current_thread", start_paused = true))] + async fn test_connpool_10() -> Result<()> { + let full_qps = 2000; + let spec = 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() + }; + run(spec).await + } +} diff --git a/edb/server/conn_pool/src/python.rs b/edb/server/conn_pool/src/python.rs new file mode 100644 index 000000000000..c84ff1ada154 --- /dev/null +++ b/edb/server/conn_pool/src/python.rs @@ -0,0 +1,380 @@ +use crate::{ + conn::{ConnError, Connector}, + pool::{Pool, PoolConfig}, + PoolHandle, +}; +use futures::TryFutureExt; +use pyo3::{ + exceptions::PyException, + prelude::*, + types::{PyDict, PyTuple}, +}; +use std::{ + cell::RefCell, + collections::HashMap, + rc::Rc, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, RwLock, + }, + time::Duration, +}; +use tokio::task::LocalSet; +use tracing::{error, trace}; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; + +pyo3::create_exception!(_conn_pool, InternalError, PyException); + +#[derive(Debug)] +#[repr(u8)] +enum ConnectOp { + Connect, + Disconnect, + Reconnect, +} + +#[derive(Debug, Default)] +struct PythonConnectionMap { + /// Connection : [`PoolHandle`] (to keep the handle alive) + handle: HashMap>, + py_dict: Option>, + next_id: usize, +} + +impl PythonConnectionMap { + pub fn insert(&mut self, py: Python, handle: PoolHandle) { + let py_dict = self + .py_dict + .get_or_insert_with(|| PyDict::new(py).into()) + .as_ref(py); + _ = handle.with_handle(|conn| py_dict.set_item(conn, self.next_id)); + self.handle.insert(self.next_id, handle); + self.next_id += 1; + } + + pub fn remove( + &mut self, + py: Python, + conn: PyObject, + ) -> Option> { + let Some(py_dict) = &mut self.py_dict else { + return None; + }; + let py_dict = py_dict.as_ref(py); + let item = py_dict.get_item(conn.clone_ref(py)).ok()??; + _ = py_dict.del_item(conn); + let key = item.extract::().ok()?; + self.handle.remove(&key) + } +} + +/// Implementation of the [`Connector`] interface. We don't pass the pool or Python objects +/// between threads, but rather use a usize ID that allows us to keep two maps in sync on +/// both sides of this interface. +#[derive(Debug)] +struct PythonConnectionFactory { + /// The _callback method that triggers the correctly-threaded task for the + /// connection operation. + callback: PyObject, + /// RPC callbacks. + responses: Arc>>>, + /// Next RPC ID. + next_response_id: Arc, +} + +impl PythonConnectionFactory { + fn send( + &self, + op: ConnectOp, + args: impl IntoPy>, + ) -> impl futures::Future> + 'static { + let (sender, receiver) = tokio::sync::oneshot::channel::(); + let response_id = self.next_response_id.fetch_add(1, Ordering::SeqCst); + self.responses.write().unwrap().insert(response_id, sender); + let success = Python::with_gil(|py| { + let args0: Py = (op as u8, response_id).into_py(py); + let args = args.into_py(py); + + let Ok(result) = self.callback.call(py, (args0, args), None) else { + error!("Unexpected failure in _callback"); + return false; + }; + let Ok(result) = result.is_true(py) else { + error!("Unexpected return value from _callback"); + return false; + }; + if !result { + return false; + } + true + }); + async move { + if success { + let conn = receiver.await.unwrap(); + let conn = Python::with_gil(|py| conn.to_object(py)); + trace!("Thread received {response_id} {}", conn); + Ok(conn) + } else { + Err(ConnError::Shutdown) + } + } + } +} + +impl Connector for PythonConnectionFactory { + type Conn = PyObject; + + fn connect( + &self, + db: &str, + ) -> impl futures::Future> + 'static { + self.send(ConnectOp::Connect, (db,)) + } + + fn disconnect( + &self, + conn: Self::Conn, + ) -> impl futures::Future> + 'static { + self.send(ConnectOp::Disconnect, (conn,)).map_ok(|_| ()) + } + + fn reconnect( + &self, + conn: Self::Conn, + db: &str, + ) -> impl futures::Future> + 'static { + self.send(ConnectOp::Reconnect, (conn, db)) + } +} + +impl PythonConnectionFactory { + fn new(callback: PyObject) -> Self { + Self { + callback, + responses: Default::default(), + next_response_id: Default::default(), + } + } +} + +#[derive(Debug)] +enum PoolRPC { + Acquire(String, PyObject), + Release(PyObject, bool), +} + +#[pyclass] +struct ConnPool { + connector: RwLock>, + responses: Arc>>>, + rpc_tx: RwLock>>, +} + +fn internal_error(py: Python, message: &str) { + error!("{message}"); + InternalError::new_err(()).restore(py); +} + +async fn run_and_block( + connector: PythonConnectionFactory, + mut rpc_rx: tokio::sync::mpsc::UnboundedReceiver, +) { + let pool = Rc::new(Pool::::new( + PoolConfig::suggested_default_for(100), + connector, + )); + let conns = Rc::new(RefCell::new(PythonConnectionMap::default())); + + let pool_task = { + let pool = pool.clone(); + tokio::task::spawn_local(async move { + loop { + pool.run_once(); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + }; + + loop { + let Some(rpc) = rpc_rx.recv().await else { + pool_task.abort(); + break; + }; + let pool = pool.clone(); + let conns = conns.clone(); + trace!("Received RPC: {rpc:?}"); + tokio::task::spawn_local(async move { + match rpc { + PoolRPC::Acquire(db, callback) => { + let conn = pool.acquire(&db).await.unwrap(); + trace!("Acquired a handle to return to Python!"); + Python::with_gil(|py| { + let handle = conn.handle_clone(); + conns.borrow_mut().insert(py, conn); + callback.call1(py, (handle,)).unwrap(); + }); + } + PoolRPC::Release(conn, dispose) => { + Python::with_gil(|py| { + let Some(conn) = conns.borrow_mut().remove(py, conn) else { + error!("Attempted to dispose a connection that does not exist"); + return; + }; + + if dispose { + conn.poison(); + } + + drop(conn); + }); + } + } + }); + } +} + +#[pymethods] +impl ConnPool { + #[new] + fn new(callback: PyObject) -> Self { + let connector = PythonConnectionFactory::new(callback); + let responses = connector.responses.clone(); + ConnPool { + connector: RwLock::new(Some(connector)), + responses, + rpc_tx: Default::default(), + } + } + + fn _respond(&self, py: Python, response_id: usize, object: PyObject) { + trace!("_respond({response_id}, {object})"); + let response = self.responses.write().unwrap().remove(&response_id); + if let Some(response) = response { + response.send(object).unwrap(); + } else { + internal_error(py, "Missing response sender"); + } + } + + fn halt(&self, _py: Python) { + self.rpc_tx.write().unwrap().take(); + } + + /// Asynchronously acquires a connection, returning it to the callback + fn acquire(&self, db: &str, callback: PyObject) { + self.rpc_tx + .read() + .unwrap() + .as_ref() + .unwrap() + .send(PoolRPC::Acquire(db.to_owned(), callback)) + .unwrap(); + } + + /// Releases a connection when possible, potentially discarding it + fn release(&self, conn: PyObject, discard: bool) { + self.rpc_tx + .read() + .unwrap() + .as_ref() + .unwrap() + .send(PoolRPC::Release(conn, discard)) + .unwrap(); + } + + /// Boot the connection pool on this thread. + fn run_and_block(&self, py: Python) { + let connector = self.connector.write().unwrap().take().unwrap(); + let (rpc_tx, rpc_rx) = tokio::sync::mpsc::unbounded_channel(); + *self.rpc_tx.write().unwrap() = Some(rpc_tx); + py.allow_threads(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_time() + .build() + .unwrap(); + let local = LocalSet::new(); + local.block_on(&rt, run_and_block(connector, rpc_rx)); + }) + } +} + +#[pymodule] +fn _conn_pool(py: Python, m: &PyModule) -> PyResult<()> { + m.add_class::()?; + m.add("InternalError", py.get_type::())?; + + let logging = py.import("logging")?; + let logger = logging + .getattr("getLogger")? + .call(("edb.server.connpool",), None)?; + let level = logger + .getattr("getEffectiveLevel")? + .call((), None)? + .extract::()?; + let logger = logger.to_object(py); + + struct PythonSubscriber { + logger: Py, + } + + impl tracing_subscriber::Layer for PythonSubscriber { + fn on_event(&self, event: &tracing::Event, _ctx: tracing_subscriber::layer::Context) { + let mut message = format!("[{}] ", event.metadata().target()); + + #[derive(Default)] + struct Visitor(String); + impl tracing::field::Visit for Visitor { + fn record_debug( + &mut self, + field: &tracing::field::Field, + value: &dyn std::fmt::Debug, + ) { + if field.name() == "message" { + self.0 += &format!("{value:?} "); + } else { + self.0 += &format!("{}={:?} ", field.name(), value) + } + } + } + + let mut visitor = Visitor::default(); + event.record(&mut visitor); + message += &visitor.0; + + Python::with_gil(|py| { + let log = match *event.metadata().level() { + tracing::Level::TRACE => self.logger.getattr(py, "debug").unwrap(), + tracing::Level::DEBUG => self.logger.getattr(py, "warning").unwrap(), + tracing::Level::INFO => self.logger.getattr(py, "info").unwrap(), + tracing::Level::WARN => self.logger.getattr(py, "warning").unwrap(), + tracing::Level::ERROR => self.logger.getattr(py, "error").unwrap(), + }; + log.call1(py, (message,)).unwrap(); + }) + } + } + + let level = if level < 10 { + tracing_subscriber::filter::LevelFilter::TRACE + } else if level <= 10 { + tracing_subscriber::filter::LevelFilter::DEBUG + } else if level <= 20 { + tracing_subscriber::filter::LevelFilter::INFO + } else if level <= 30 { + tracing_subscriber::filter::LevelFilter::WARN + } else if level <= 40 { + tracing_subscriber::filter::LevelFilter::ERROR + } else { + tracing_subscriber::filter::LevelFilter::OFF + }; + + let subscriber = PythonSubscriber { logger }; + tracing_subscriber::registry() + .with(level) + .with(subscriber) + .init(); + + tracing::info!("ConnPool initialized (level = {level})"); + + Ok(()) +} diff --git a/edb/server/conn_pool/src/test.rs b/edb/server/conn_pool/src/test.rs new file mode 100644 index 000000000000..d4f3353b76f2 --- /dev/null +++ b/edb/server/conn_pool/src/test.rs @@ -0,0 +1,417 @@ +//! Test utilities. +use std::{ + borrow::Cow, cell::RefCell, collections::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)), + } + } +} + +/// Perform a virtual async sleep that advances all of the virtual clocks (`mock_instant` and `tokio`). +pub async fn virtual_sleep(duration: Duration) { + // Perform the mock sleep, assumes that the tokio time is paused which will + // auto-advance the paused clock. + tokio::time::sleep(duration).await; +} + +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 { + virtual_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 { + virtual_sleep(f(true)).await; + virtual_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 { + virtual_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, +} + +pub struct Scored { + pub description: String, + pub detailed_calculation: Box String>, + pub raw_value: f64, +} + +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 { + 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(); + // 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 mut data = Data::new(data.get_mut(db).expect(db).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 mut data = Data::new(data.get_mut(db).expect(db).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 mut data = Data::new(data.get_mut(db).expect(db).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 mut data = Data::new(data.get_mut(db).expect(db).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, +} + +#[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 000000000000..5f00792c1306 --- /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 577dffa0fcf4..63f400dfeec3 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/pool.py b/edb/server/connpool/pool.py index 0f129948f5f2..821f8d11677a 100644 --- a/edb/server/connpool/pool.py +++ b/edb/server/connpool/pool.py @@ -419,6 +419,16 @@ def failed_connects(self) -> int: def failed_disconnects(self) -> int: return self._failed_disconnects + + async def __aenter__(self) -> typing.Self: + return self + + async def __aexit__(self, + exc_type: typing.Optional[type], + exc_val: typing.Optional[BaseException], + exc_tb: typing.Optional[typing.Any]) -> None: + pass + def get_pending_conns(self) -> int: return sum( block.count_pending_conns() for block in self._blocks.values() diff --git a/edb/server/connpool/pool2.py b/edb/server/connpool/pool2.py new file mode 100644 index 000000000000..8b4dee5b0675 --- /dev/null +++ b/edb/server/connpool/pool2.py @@ -0,0 +1,275 @@ +import edb.server._conn_pool +import asyncio +import threading +import typing +import dataclasses + +# Connections must be hashable because we use them to reverse-lookup +# an internal ID. +C = typing.TypeVar("C", bound=typing.Hashable) + +CP1 = typing.TypeVar('CP1', covariant=True) +CP2 = typing.TypeVar('CP2', contravariant=True) + + +class Connector(typing.Protocol[CP1]): + + def __call__(self, dbname: str) -> typing.Awaitable[CP1]: + pass + + +class Disconnector(typing.Protocol[CP2]): + + def __call__(self, conn: CP2) -> typing.Awaitable[None]: + pass + + +@dataclasses.dataclass +class Snapshot: + timestamp: float + capacity: int + + failed_connects: int + failed_disconnects: int + successful_connects: int + successful_disconnects: int + + +class StatsCollector(typing.Protocol): + + def __call__(self, stats: Snapshot) -> None: + pass + + +class ConnectionFactory(typing.Protocol[C]): + """The async interface to create and destroy database connections. + + All connections returned from successful calls to `connect` or reconnect + are guaranteed to be `disconnect`ed or `reconnect`ed.""" + + async def connect(self, db: str) -> C: + """Create a new connection asynchronously. + + This method must retry exceptions internally. If an exception is thrown + from this method, the database is considered to be failed.""" + ... + + async def disconnect(self, conn: C) -> None: + """Gracefully disconnect a connection asynchronously. + + If an exception is thrown from this method, the connection is simply + forgotten.""" + ... + + async def reconnect(self, conn: C, db: str) -> C: + """Reconnects a connection to the given database. If this is not + possible, it is permissable to return a new connection and gracefully + disconnect the other connection in parallel or in the background. + + This method must retry exceptions internally. If an exception is thrown + from this method, the database is considered to be failed.""" + ... + + +class ConnPool(typing.Generic[C]): + _connection_factory: ConnectionFactory[C] + _pool: edb.server._conn_pool.ConnPool + _loop: asyncio.AbstractEventLoop + _completion: asyncio.Future[bool] + _ready: asyncio.Future[bool] + _active_conns: set[C] + _cur_capacity: int + _cur_waiters: int + + def __init__(self, connection_factory: ConnectionFactory[C]): + self._connection_factory = connection_factory + self._loop = asyncio.get_event_loop() + self._pool = None + self._completion = self._loop.create_future() + self._ready = self._loop.create_future() + self._active_conns = set() + self._cur_capacity = 0 + self._cur_waiters = 0 + + def _callback(self, args0: typing.Any, args: typing.Any) -> bool: + """Receives the callback from the Rust connection pool. + + Required to call pool._respond on the main thread with the result of + this callback. + """ + (kind, response_id) = args0 + if self._loop.is_closed(): + return False + else: + self._loop.call_soon_threadsafe( + self._loop.create_task, + self._perform_async(kind, response_id, *args), + ) + return True + + async def _perform_async(self, kind: int, + response_id: int, + *args: typing.Any) -> None: + """Delegates the callback from Rust to the appropriate connection + factory method.""" + if kind == 0: + self._cur_capacity += 1 + response = await self._connection_factory.connect(*args) + elif kind == 1: + await self._connection_factory.disconnect(*args) + self._cur_capacity -= 1 + response = None + elif kind == 2: + response = await self._connection_factory.reconnect(*args) + if self._pool is not None: + self._pool._respond(response_id, response) + + def _thread_main(self) -> None: + self._loop.call_soon_threadsafe(self._ready.set_result, True) + self._pool.run_and_block() + if not self._loop.is_closed(): + self._loop.call_soon_threadsafe(self._completion.set_result, True) + + async def run(self) -> None: + """Creates a long-lived task that manages the connection pool. Required + before any connections may be acquired.""" + if self._pool is not None: + raise RuntimeError(f"pool already started") from None + + self._pool = edb.server._conn_pool.ConnPool(self._callback) + threading.Thread(target=self._thread_main, daemon=True).start() + try: + await self._completion + except asyncio.exceptions.CancelledError: + self._pool.halt() + self._pool = None + + async def acquire(self, db: str) -> C: + """Acquire a connection from the database. This connection must be + released.""" + await self._ready + future: asyncio.Future[C] = self._loop.create_future() + # Note that this callback is called on the internal pool's thread + self._cur_waiters += 1 + try: + self._pool.acquire( + db, + lambda res: self._loop.call_soon_threadsafe(future.set_result, + res), + ) + conn = await future + finally: + self._cur_waiters -= 1 + self._active_conns.add(conn) + return conn + + def release(self, _db: str, conn: C, discard: bool = False) -> None: + """Releases a connection back into the pool, discarding or returning it + in the background.""" + self._active_conns.remove(conn) + self._pool.release(conn, discard) + pass + + def count_waiters(self) -> int: + return self._cur_waiters + + +class FactoryAdapter(typing.Generic[C]): + _connect: Connector[C] + _disconnect: Disconnector[C] + + def __init__(self, + connect: Connector[C], + disconnect: Disconnector[C]) -> None: + self._connect = connect + self._disconnect = disconnect + + async def connect(self, db: str) -> C: + return await self._connect(db) + + async def disconnect(self, conn: C) -> None: + await self._disconnect(conn) + + async def reconnect(self, conn: C, db: str) -> C: + await self._disconnect(conn) + return await self._connect(db) + + +class Pool(typing.Generic[C]): + _pool: ConnPool[C] + _failed_connects: int + _failed_disconnects: int + _successful_connects: int + _successful_disconnects: int + _cur_capacity: int + _max_capacity: int + _task: typing.Optional[asyncio.Task[None]] + + def __init__(self, *, connect: Connector[C], + disconnect: Disconnector[C], + stats_collector: typing.Optional[StatsCollector], + max_capacity: int) -> None: + self._pool = ConnPool(connection_factory=FactoryAdapter(connect, + disconnect,)) + self._failed_connects = 0 + self._failed_disconnects = 0 + self._successful_connects = 0 + self._successful_disconnects = 0 + self._task = None + + if stats_collector: + stats_collector(Snapshot( + timestamp=0, + capacity=10, + failed_connects=0, + failed_disconnects=0, + successful_connects=0, + successful_disconnects=0)) + pass + + async def __aenter__(self) -> typing.Self: + self._task = asyncio.create_task(self._pool.run()) + return self + + async def __aexit__(self, + exc_type: typing.Optional[type], + exc_val: typing.Optional[BaseException], + exc_tb: typing.Optional[typing.Any]) -> None: + if self._task: + self._task.cancel() + try: + await self._task + except asyncio.CancelledError: + print("Task has been successfully cancelled") + self._task = None + print("Exiting context") + + async def acquire(self, db: str) -> C: + """Acquire a connection from the database. This connection must be + released.""" + if self._task is None: + raise RuntimeError("Not entered") + return await self._pool.acquire(db) + + def release(self, _db: str, conn: C, discard: bool = False) -> None: + """Releases a connection back into the pool, discarding or returning it + in the background.""" + if self._task is None: + raise RuntimeError("Not entered") + self._pool.release(_db, conn, discard) + + @property + def max_capacity(self) -> int: + return self._max_capacity + + @property + def current_capacity(self) -> int: + return self._cur_capacity + + @property + def failed_connects(self) -> int: + return self._failed_connects + + @property + def failed_disconnects(self) -> int: + return self._failed_disconnects diff --git a/setup.py b/setup.py index a5fd3d394450..b4fae21bccd2 100644 --- a/setup.py +++ b/setup.py @@ -1121,5 +1121,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, + ), ], ) diff --git a/tests/test_server_pool.py b/tests/test_server_pool.py index f5ae94bea058..77ba66030ccd 100644 --- a/tests/test_server_pool.py +++ b/tests/test_server_pool.py @@ -511,54 +511,54 @@ def on_stats(stat): stat = dataclasses.asdict(stat) sim.stats.append(stat) - pool = pool_cls( + async with pool_cls( connect=self.make_fake_connect( sim, spec.conn_cost_base, spec.conn_cost_var), disconnect=self.make_fake_disconnect( sim, spec.disconn_cost_base, spec.disconn_cost_var), stats_collector=on_stats if collect_stats else None, max_capacity=spec.capacity, - ) - if hasattr(pool, '_gc_interval'): - pool._gc_interval = 0.1 * TIME_SCALE + ) as pool: + if hasattr(pool, '_gc_interval'): + pool._gc_interval = 0.1 * TIME_SCALE started_at = time.monotonic() async with asyncio.TaskGroup() as g: for db in spec.dbs: g.create_task(self.simulate_db(sim, pool, g, db)) - self.assertEqual(sim.failed_disconnects, 0) - self.assertEqual(sim.failed_queries, 0) + self.assertEqual(sim.failed_disconnects, 0) + self.assertEqual(sim.failed_queries, 0) - self.assertEqual(pool.failed_disconnects, 0) - self.assertEqual(pool.failed_connects, 0) + self.assertEqual(pool.failed_disconnects, 0) + self.assertEqual(pool.failed_connects, 0) - try: - for db in sim.latencies: - int(db[1:]) - except ValueError: - key_func = lambda x: x - else: - key_func = lambda x: int(x[0][1:]) - - if collect_stats: - pn = f'{type(pool).__module__}.{type(pool).__qualname__}' - score = int(round(sum(sm.calculate(sim) for sm in spec.score))) - print('weighted score:'.rjust(68), score) - js_data = { - 'test_started_at': started_at, - 'total_lats': calc_total_percentiles(sim.latencies), - "score": score, - 'scores': sim.scores, - 'lats': { - db: calc_percentiles(lats) - for db, lats in sorted(sim.latencies.items(), key=key_func) - }, - 'pool_name': pn, - 'stats': sim.stats, - } + try: + for db in sim.latencies: + int(db[1:]) + except ValueError: + key_func = lambda x: x + else: + key_func = lambda x: int(x[0][1:]) + + if collect_stats: + pn = f'{type(pool).__module__}.{type(pool).__qualname__}' + score = int(round(sum(sm.calculate(sim) for sm in spec.score))) + print('weighted score:'.rjust(68), score) + js_data = { + 'test_started_at': started_at, + 'total_lats': calc_total_percentiles(sim.latencies), + "score": score, + 'scores': sim.scores, + 'lats': { + db: calc_percentiles(lats) + for db, lats in sorted(sim.latencies.items(), key=key_func) + }, + 'pool_name': pn, + 'stats': sim.stats, + } - return js_data + return js_data async def simulate(self, testname, spec): if os.environ.get('EDGEDB_TEST_DEBUG_POOL'): @@ -578,7 +578,7 @@ async def simulate(self, testname, spec): ) async def simulate_and_collect_stats(self, testname, spec): - pools = [connpool.Pool, connpool._NaivePool] + pools = [connpool.Pool2, connpool.Pool, connpool._NaivePool] js_data = [] for pool_cls in pools: