From eb8cfdb1f186ab9d78a06601e543a02218ed06cd Mon Sep 17 00:00:00 2001 From: zainkabani <77307340+zainkabani@users.noreply.github.com> Date: Mon, 27 Feb 2023 01:16:30 -0500 Subject: [PATCH] Adds SHUTDOWN command as alternate option to sending SIGINT (#331) * Adds SHUTDOWN command to PgCat as alternate option to sending SIGINT * Check if we're already in SHUTDOWN sequence * Send signal directly from shutdown instead of using channel * Add tests * trigger build * Lowercase response and boolean change * Update tests * Fix tests * typo --- Cargo.lock | 249 +++++++++++++++++++++++++----------------- Cargo.toml | 1 + src/admin.rs | 36 +++++- src/main.rs | 152 +++++++++++++------------- tests/python/tests.py | 61 +++++++++++ 5 files changed, 327 insertions(+), 172 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7b18b2c..41b32ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -79,9 +79,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.11.1" +version = "3.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "572f695136211188308f16ad2ca5c851a712c464060ae6974944458eb83880ba" +checksum = "0d261e256854913907f67ed06efbc3338dfe6179796deefc1ff763fc1aee5535" [[package]] name = "bytes" @@ -91,9 +91,9 @@ checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" [[package]] name = "cc" -version = "1.0.78" +version = "1.0.79" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a20104e2335ce8a659d6dd92a51a767a0c062599c73b343fd152cb401e828c3d" +checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f" [[package]] name = "cfg-if" @@ -153,9 +153,9 @@ dependencies = [ [[package]] name = "cxx" -version = "1.0.86" +version = "1.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51d1075c37807dcf850c379432f0df05ba52cc30f279c5cfc43cc221ce7f8579" +checksum = "86d3488e7665a7a483b57e25bdd90d0aeb2bc7608c8d0346acf2ad3f1caf1d62" dependencies = [ "cc", "cxxbridge-flags", @@ -165,9 +165,9 @@ dependencies = [ [[package]] name = "cxx-build" -version = "1.0.86" +version = "1.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5044281f61b27bc598f2f6647d480aed48d2bf52d6eb0b627d84c0361b17aa70" +checksum = "48fcaf066a053a41a81dfb14d57d99738b767febb8b735c3016e469fac5da690" dependencies = [ "cc", "codespan-reporting", @@ -180,15 +180,15 @@ dependencies = [ [[package]] name = "cxxbridge-flags" -version = "1.0.86" +version = "1.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61b50bc93ba22c27b0d31128d2d130a0a6b3d267ae27ef7e4fae2167dfe8781c" +checksum = "a2ef98b8b717a829ca5603af80e1f9e2e48013ab227b68ef37872ef84ee479bf" [[package]] name = "cxxbridge-macro" -version = "1.0.86" +version = "1.0.91" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e61fda7e62115119469c7b3591fd913ecca96fb766cfd3f2e2502ab7bc87a5" +checksum = "086c685979a698443656e5cf7856c95c642295a38599f12fb1ff76fb28d19892" dependencies = [ "proc-macro2", "quote", @@ -252,12 +252,6 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" -[[package]] -name = "fs_extra" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2022715d62ab30faffd124d40b76f4134a550a87792276512b18d63272333394" - [[package]] name = "futures" version = "0.3.26" @@ -402,6 +396,12 @@ dependencies = [ "libc", ] +[[package]] +name = "hermit-abi" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fed44880c466736ef9a5c5b5facefb5ed0785676d0c02d612db14e54f0d84286" + [[package]] name = "hmac" version = "0.12.1" @@ -413,9 +413,9 @@ dependencies = [ [[package]] name = "http" -version = "0.2.8" +version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75f43d41e26995c17e71ee126451dd3941010b0514a81a9d11f3b341debc2399" +checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482" dependencies = [ "bytes", "fnv", @@ -511,24 +511,24 @@ dependencies = [ [[package]] name = "io-lifetimes" -version = "1.0.3" +version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46112a93252b123d31a119a8d1a1ac19deac4fac6e0e8b0df58f0d4e5870e63c" +checksum = "1abeb7a0dd0f8181267ff8adc397075586500b81b28a73e8a0208b00fc170fb3" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] name = "is-terminal" -version = "0.4.2" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28dfb6c8100ccc63462345b67d1bbc3679177c75ee4bf59bf29c8b1d110b8189" +checksum = "21b6b32576413a8e69b90e952e4a026476040d81017b80445deda5f2d3921857" dependencies = [ - "hermit-abi", + "hermit-abi 0.3.1", "io-lifetimes", "rustix", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] @@ -539,12 +539,11 @@ checksum = "fad582f4b9e86b6caa621cabeb0963332d92eea04729ab12892c2533951e6440" [[package]] name = "jemalloc-sys" -version = "0.5.2+5.3.0-patched" +version = "0.5.3+5.3.0-patched" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "134163979b6eed9564c98637b710b40979939ba351f59952708234ea11b5f3f8" +checksum = "f9bd5d616ea7ed58b571b2e209a65759664d7fb021a0819d7a790afc67e47ca1" dependencies = [ "cc", - "fs_extra", "libc", ] @@ -560,9 +559,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.60" +version = "0.3.61" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49409df3e3bf0856b916e2ceaca09ee28e6871cf7d9ce97a692cacfdb2a25a47" +checksum = "445dde2150c55e483f3d8416706b97ec8e8237c307e5b7b4b8dd15e6af2a0730" dependencies = [ "wasm-bindgen", ] @@ -623,24 +622,38 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" [[package]] -name = "mio" -version = "0.8.5" +name = "memoffset" +version = "0.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5d732bc30207a6423068df043e3d02e0735b155ad7ce1a6f76fe2baa5b158de" +checksum = "5de893c32cde5f383baa4c04c5d6dbdd735cfd4a794b0debdb2bb1b421da5ff4" +dependencies = [ + "autocfg", +] + +[[package]] +name = "mio" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b9d9a46eff5b4ff64b45a9e316a6d1e0bc719ef429cbec4dc630684212bfdf9" dependencies = [ "libc", "log", "wasi 0.11.0+wasi-snapshot-preview1", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] -name = "nom8" -version = "0.2.0" +name = "nix" +version = "0.26.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ae01545c9c7fc4486ab7debaf2aad7003ac19431791868fb2e8066df97fad2f8" +checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a" dependencies = [ - "memchr", + "bitflags", + "cfg-if", + "libc", + "memoffset", + "pin-utils", + "static_assertions", ] [[package]] @@ -668,7 +681,7 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fac9e2da13b5eb447a6ce3d392f23a29d8694bff781bf03a16cd9ac8697593b" dependencies = [ - "hermit-abi", + "hermit-abi 0.2.6", "libc", ] @@ -690,15 +703,15 @@ dependencies = [ [[package]] name = "parking_lot_core" -version = "0.9.5" +version = "0.9.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ff9f3fef3968a3ec5945535ed654cb38ff72d7495a25619e2247fb15a2ed9ba" +checksum = "9069cbb9f99e3a5083476ccb29ceb1de18b9118cafa53e90c9551235de2b9521" dependencies = [ "cfg-if", "libc", "redox_syscall", "smallvec", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] @@ -719,6 +732,7 @@ dependencies = [ "jemallocator", "log", "md-5", + "nix", "num_cpus", "once_cell", "parking_lot", @@ -800,9 +814,9 @@ checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" [[package]] name = "proc-macro2" -version = "1.0.49" +version = "1.0.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57a8eca9f9c4ffde41714334dee777596264c7825420f521abc92b5b5deb63a5" +checksum = "5d727cae5b39d21da60fa540906919ad737832fe0b1c165da3a34d6548c849d6" dependencies = [ "unicode-ident", ] @@ -889,23 +903,23 @@ dependencies = [ [[package]] name = "rustix" -version = "0.36.6" +version = "0.36.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4feacf7db682c6c329c4ede12649cd36ecab0f3be5b7d74e6a20304725db4549" +checksum = "f43abb88211988493c1abb44a70efa56ff0ce98f233b7b276146f1f3f7ba9644" dependencies = [ "bitflags", "errno", "io-lifetimes", "libc", "linux-raw-sys", - "windows-sys", + "windows-sys 0.45.0", ] [[package]] name = "rustls" -version = "0.20.7" +version = "0.20.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "539a2bfe908f471bfa933876bd1eb6a19cf2176d375f82ef7f99530a40e48c2c" +checksum = "fff78fc74d175294f4e83b28343315ffcfb114b156f0185e9741cb5570f50e2f" dependencies = [ "log", "ring", @@ -994,9 +1008,9 @@ dependencies = [ [[package]] name = "signal-hook-registry" -version = "1.4.0" +version = "1.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +checksum = "d8229b473baa5980ac72ef434c4415e70c4b5e71b423043adb4ba059f89c99a1" dependencies = [ "libc", ] @@ -1009,9 +1023,9 @@ checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de" [[package]] name = "slab" -version = "0.4.7" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4614a76b2a8be0058caa9dbbaf66d988527d86d003c11a94fbd335d7661edcef" +checksum = "6528351c9bc8ab22353f9d776db39a20288e8d6c37ef8cfe3317cf875eecfc2d" dependencies = [ "autocfg", ] @@ -1047,6 +1061,12 @@ dependencies = [ "log", ] +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "stringprep" version = "0.1.2" @@ -1065,9 +1085,9 @@ checksum = "6bdef32e8150c2a081110b42772ffe7d7c9032b606bc226c8260fd97e0976601" [[package]] name = "syn" -version = "1.0.107" +version = "1.0.109" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f4064b5b16e03ae50984a5a8ed5d4f8803e6bc1fd170a3cda91a1be4b18e3f5" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" dependencies = [ "proc-macro2", "quote", @@ -1076,9 +1096,9 @@ dependencies = [ [[package]] name = "termcolor" -version = "1.1.3" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755" +checksum = "be55cf8942feac5c765c2c993422806843c9a9a45d4d5c407ad6dd2ea95eb9b6" dependencies = [ "winapi-util", ] @@ -1105,9 +1125,9 @@ dependencies = [ [[package]] name = "tinyvec_macros" -version = "0.1.0" +version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" @@ -1126,7 +1146,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.42.0", ] [[package]] @@ -1153,9 +1173,9 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.4" +version = "0.7.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bb2e075f03b3d66d8d8785356224ba688d2906a371015e225beeb65ca92c740" +checksum = "5427d89453009325de0d8f342c9490009f76e999cb7672d77e46267448f7e6b2" dependencies = [ "bytes", "futures-core", @@ -1167,9 +1187,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.7.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "772c1426ab886e7362aedf4abc9c0d1348a979517efedfc25862944d10137af0" +checksum = "f7afcae9e3f0fe2c370fd4657108972cbb2fa9db1b9f84849cefd80741b01cb6" dependencies = [ "serde", "serde_spanned", @@ -1188,15 +1208,15 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.19.1" +version = "0.19.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90a238ee2e6ede22fb95350acc78e21dc40da00bb66c0334bde83de4ed89424e" +checksum = "9a1eb0622d28f4b9c90adc4ea4b2b46b47663fde9ac5fafcb14a1369d5508825" dependencies = [ "indexmap", - "nom8", "serde", "serde_spanned", "toml_datetime", + "winnow", ] [[package]] @@ -1239,9 +1259,9 @@ checksum = "497961ef93d974e23eb6f433eb5fe1b7930b659f06d12dec6fc44a8f554c0bba" [[package]] name = "unicode-bidi" -version = "0.3.8" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "099b7128301d285f79ddd55b9a83d5e6b9e97c92e0ea0daebee7263e932de992" +checksum = "d54675592c1dbefd78cbd98db9bacd89886e1ca50692a0692baefffdeb92dd58" [[package]] name = "unicode-ident" @@ -1300,9 +1320,9 @@ checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" [[package]] name = "wasm-bindgen" -version = "0.2.83" +version = "0.2.84" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eaf9f5aceeec8be17c128b2e93e031fb8a4d469bb9c4ae2d7dc1888b26887268" +checksum = "31f8dcbc21f30d9b8f2ea926ecb58f6b91192c17e9d33594b3df58b2007ca53b" dependencies = [ "cfg-if", "wasm-bindgen-macro", @@ -1310,9 +1330,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-backend" -version = "0.2.83" +version = "0.2.84" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c8ffb332579b0557b52d268b91feab8df3615f265d5270fec2a8c95b17c1142" +checksum = "95ce90fd5bcc06af55a641a86428ee4229e44e07033963a2290a8e241607ccb9" dependencies = [ "bumpalo", "log", @@ -1325,9 +1345,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.83" +version = "0.2.84" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "052be0f94026e6cbc75cdefc9bae13fd6052cdcaf532fa6c45e7ae33a1e6c810" +checksum = "4c21f77c0bedc37fd5dc21f897894a5ca01e7bb159884559461862ae90c0b4c5" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -1335,9 +1355,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.83" +version = "0.2.84" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07bc0c051dc5f23e307b13285f9d75df86bfdf816c5721e573dec1f9b8aa193c" +checksum = "2aff81306fcac3c7515ad4e177f521b5c9a15f2b08f4e32d823066102f35a5f6" dependencies = [ "proc-macro2", "quote", @@ -1348,15 +1368,15 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.83" +version = "0.2.84" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c38c045535d93ec4f0b4defec448e4291638ee608530863b1e2ba115d4fff7f" +checksum = "0046fef7e28c3804e5e38bfa31ea2a0f73905319b677e57ebe37e49358989b5d" [[package]] name = "web-sys" -version = "0.3.60" +version = "0.3.61" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bcda906d8be16e728fd5adc5b729afad4e444e106ab28cd1c7256e54fa61510f" +checksum = "e33b99f4b23ba3eec1a53ac264e35a755f00e966e0065077d6027c0f575b0b97" dependencies = [ "js-sys", "wasm-bindgen", @@ -1419,43 +1439,76 @@ dependencies = [ ] [[package]] -name = "windows_aarch64_gnullvm" -version = "0.42.0" +name = "windows-sys" +version = "0.45.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d2aa71f6f0cbe00ae5167d90ef3cfe66527d6f613ca78ac8024c3ccab9a19e" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2522491fbfcd58cc84d47aeb2958948c4b8982e9a2d8a2a35bbaed431390e7" +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", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c9864e83243fdec7fc9c5444389dcbbfd258f745e7853198f365e3c4968a608" [[package]] name = "windows_aarch64_msvc" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd0f252f5a35cac83d6311b2e795981f5ee6e67eb1f9a7f64eb4500fbc4dcdb4" +checksum = "4c8b1b673ffc16c47a9ff48570a9d85e25d265735c503681332589af6253c6c7" [[package]] name = "windows_i686_gnu" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbeae19f6716841636c28d695375df17562ca208b2b7d0dc47635a50ae6c5de7" +checksum = "de3887528ad530ba7bdbb1faa8275ec7a1155a45ffa57c37993960277145d640" [[package]] name = "windows_i686_msvc" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84c12f65daa39dd2babe6e442988fc329d6243fdce47d7d2d155b8d874862246" +checksum = "bf4d1122317eddd6ff351aa852118a2418ad4214e6613a50e0191f7004372605" [[package]] name = "windows_x86_64_gnu" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf7b1b21b5362cbc318f686150e5bcea75ecedc74dd157d874d754a2ca44b0ed" +checksum = "c1040f221285e17ebccbc2591ffdc2d44ee1f9186324dd3e84e99ac68d699c45" [[package]] name = "windows_x86_64_gnullvm" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09d525d2ba30eeb3297665bd434a54297e4170c7f1a44cad4ef58095b4cd2028" +checksum = "628bfdf232daa22b0d64fdb62b09fcc36bb01f05a3939e20ab73aaf9470d0463" [[package]] name = "windows_x86_64_msvc" -version = "0.42.0" +version = "0.42.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f40009d85759725a34da6d89a94e63d7bdc50a862acf0dbc7c8e488f1edcb6f5" +checksum = "447660ad36a13288b1db4d4248e857b510e8c3a225c822ba4fb748c0aafecffd" + +[[package]] +name = "winnow" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf09497b8f8b5ac5d3bb4d05c0a99be20f26fd3d5f2db7b0716e946d5103658" +dependencies = [ + "memchr", +] diff --git a/Cargo.toml b/Cargo.toml index 3cf0e7a..0c02027 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,7 @@ phf = { version = "0.11.1", features = ["macros"] } exitcode = "1.1.2" futures = "0.3" socket2 = { version = "0.4.7", features = ["all"] } +nix = "0.26.2" [target.'cfg(not(target_env = "msvc"))'.dependencies] jemallocator = "0.5.0" diff --git a/src/admin.rs b/src/admin.rs index 406b0fe..7f7ada5 100644 --- a/src/admin.rs +++ b/src/admin.rs @@ -1,6 +1,8 @@ /// Admin database. use bytes::{Buf, BufMut, BytesMut}; -use log::{info, trace}; +use log::{error, info, trace}; +use nix::sys::signal::{self, Signal}; +use nix::unistd::Pid; use std::collections::HashMap; use tokio::time::Instant; @@ -67,6 +69,10 @@ where trace!("RESUME"); resume(stream, query_parts[1]).await } + "SHUTDOWN" => { + trace!("SHUTDOWN"); + shutdown(stream).await + } "SHOW" => match query_parts[1].to_ascii_uppercase().as_str() { "CONFIG" => { trace!("SHOW CONFIG"); @@ -671,6 +677,34 @@ where } } +/// Send response packets for shutdown. +async fn shutdown(stream: &mut T) -> Result<(), Error> +where + T: tokio::io::AsyncWrite + std::marker::Unpin, +{ + let mut res = BytesMut::new(); + + res.put(row_description(&vec![("success", DataType::Text)])); + + let mut shutdown_success = "t"; + + let pid = std::process::id(); + if signal::kill(Pid::from_raw(pid.try_into().unwrap()), Signal::SIGINT).is_err() { + error!("Unable to send SIGINT to PID: {}", pid); + shutdown_success = "f"; + } + + res.put(data_row(&vec![shutdown_success.to_string()])); + + res.put(command_complete("SHUTDOWN")); + + res.put_u8(b'Z'); + res.put_i32(5); + res.put_u8(b'I'); + + write_all_half(stream, &res).await +} + /// Show Users. async fn show_users(stream: &mut T) -> Result<(), Error> where diff --git a/src/main.rs b/src/main.rs index cd408c4..33536b0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -232,7 +232,13 @@ fn main() -> Result<(), Box> { // Initiate graceful shutdown sequence on sig int _ = interrupt_signal.recv() => { - info!("Got SIGINT, waiting for client connection drain now"); + info!("Got SIGINT"); + + // Don't want this to happen more than once + if admin_only { + continue; + } + admin_only = true; // Broadcast that client tasks need to finish @@ -241,98 +247,98 @@ fn main() -> Result<(), Box> { let _ = drain_tx.send(0).await; tokio::task::spawn(async move { - let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout)); + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout)); - // First tick fires immediately. - interval.tick().await; + // First tick fires immediately. + interval.tick().await; - // Second one in the interval time. - interval.tick().await; + // Second one in the interval time. + interval.tick().await; - // We're done waiting. - error!("Graceful shutdown timed out. {} active clients being closed", total_clients); + // We're done waiting. + error!("Graceful shutdown timed out. {} active clients being closed", total_clients); - let _ = exit_tx.send(()).await; + let _ = exit_tx.send(()).await; }); }, _ = term_signal.recv() => { info!("Got SIGTERM, closing with {} clients active", total_clients); break; - }, - - new_client = listener.accept() => { - let (socket, addr) = match new_client { - Ok((socket, addr)) => (socket, addr), - Err(err) => { - error!("{:?}", err); - continue; - } - }; - - let shutdown_rx = shutdown_tx.subscribe(); - let drain_tx = drain_tx.clone(); - let client_server_map = client_server_map.clone(); - - let tls_certificate = config.general.tls_certificate.clone(); - - tokio::task::spawn(async move { - let start = chrono::offset::Utc::now().naive_utc(); - - match client::client_entrypoint( - socket, - client_server_map, - shutdown_rx, - drain_tx, - admin_only, - tls_certificate.clone(), - config.general.log_client_connections, - ) - .await - { - Ok(()) => { - let duration = chrono::offset::Utc::now().naive_utc() - start; - - if config.general.log_client_disconnections { - info!( - "Client {:?} disconnected, session duration: {}", - addr, - format_duration(&duration) - ); - } else { - debug!( - "Client {:?} disconnected, session duration: {}", - addr, - format_duration(&duration) - ); - } - } + }, + new_client = listener.accept() => { + let (socket, addr) = match new_client { + Ok((socket, addr)) => (socket, addr), Err(err) => { - match err { - errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err), - _ => warn!("Client disconnected with error {:?}", err), - } - + error!("{:?}", err); + continue; } }; - }); - } - _ = exit_rx.recv() => { - break; - } + let shutdown_rx = shutdown_tx.subscribe(); + let drain_tx = drain_tx.clone(); + let client_server_map = client_server_map.clone(); - client_ping = drain_rx.recv() => { - let client_ping = client_ping.unwrap(); - total_clients += client_ping; + let tls_certificate = config.general.tls_certificate.clone(); - if total_clients == 0 && admin_only { - let _ = exit_tx.send(()).await; + tokio::task::spawn(async move { + let start = chrono::offset::Utc::now().naive_utc(); + + match client::client_entrypoint( + socket, + client_server_map, + shutdown_rx, + drain_tx, + admin_only, + tls_certificate.clone(), + config.general.log_client_connections, + ) + .await + { + Ok(()) => { + let duration = chrono::offset::Utc::now().naive_utc() - start; + + if config.general.log_client_disconnections { + info!( + "Client {:?} disconnected, session duration: {}", + addr, + format_duration(&duration) + ); + } else { + debug!( + "Client {:?} disconnected, session duration: {}", + addr, + format_duration(&duration) + ); + } + } + + Err(err) => { + match err { + errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err), + _ => warn!("Client disconnected with error {:?}", err), + } + + } + }; + }); + } + + _ = exit_rx.recv() => { + break; + } + + client_ping = drain_rx.recv() => { + let client_ping = client_ping.unwrap(); + total_clients += client_ping; + + if total_clients == 0 && admin_only { + let _ = exit_tx.send(()).await; + } } } } - } info!("Shutting down..."); }); diff --git a/tests/python/tests.py b/tests/python/tests.py index 7c99079..6108ff2 100644 --- a/tests/python/tests.py +++ b/tests/python/tests.py @@ -110,6 +110,37 @@ def test_shutdown_logic(): cleanup_conn(conn, cur) pg_cat_send_signal(signal.SIGTERM) + # - - - - - - - - - - - - - - - - - - + # NO ACTIVE QUERIES ADMIN SHUTDOWN COMMAND + + # Start pgcat + pgcat_start() + + # Create client connection and begin transaction + conn, cur = connect_db() + admin_conn, admin_cur = connect_db(admin=True) + + cur.execute("BEGIN;") + cur.execute("SELECT 1;") + cur.execute("COMMIT;") + + # Send SHUTDOWN command pgcat while not in transaction + admin_cur.execute("SHUTDOWN;") + time.sleep(1) + + # Check that any new queries fail after SHUTDOWN command since server should close with no active transactions + try: + cur.execute("SELECT 1;") + except psycopg2.OperationalError as e: + pass + else: + # Fail if query execution succeeded + raise Exception("Server not closed after sigint") + + cleanup_conn(conn, cur) + cleanup_conn(admin_conn, admin_cur) + pg_cat_send_signal(signal.SIGTERM) + # - - - - - - - - - - - - - - - - - - # HANDLE TRANSACTION WITH SIGINT @@ -136,6 +167,36 @@ def test_shutdown_logic(): cleanup_conn(conn, cur) pg_cat_send_signal(signal.SIGTERM) + # - - - - - - - - - - - - - - - - - - + # HANDLE TRANSACTION WITH ADMIN SHUTDOWN COMMAND + + # Start pgcat + pgcat_start() + + # Create client connection and begin transaction + conn, cur = connect_db() + admin_conn, admin_cur = connect_db(admin=True) + + cur.execute("BEGIN;") + cur.execute("SELECT 1;") + + # Send SHUTDOWN command pgcat while still in transaction + admin_cur.execute("SHUTDOWN;") + if admin_cur.fetchall()[0][0] != "t": + raise Exception("PgCat unable to send signal") + time.sleep(1) + + # Check that any new queries succeed after SHUTDOWN command since server should still allow transaction to complete + try: + cur.execute("SELECT 1;") + except psycopg2.OperationalError as e: + # Fail if query fails since server closed + raise Exception("Server closed while in transaction", e.pgerror) + + cleanup_conn(conn, cur) + cleanup_conn(admin_conn, admin_cur) + pg_cat_send_signal(signal.SIGTERM) + # - - - - - - - - - - - - - - - - - - # NO NEW NON-ADMIN CONNECTIONS DURING SHUTDOWN # Start pgcat