mirror of
https://github.com/postgresml/pgcat.git
synced 2026-03-23 17:36:28 +00:00
Compare commits
1 Commits
v0.2.1-bet
...
sven_md5_a
| Author | SHA1 | Date | |
|---|---|---|---|
|
5f745859c0 |
@@ -12,15 +12,13 @@ jobs:
|
||||
- image: cimg/rust:1.58.1
|
||||
environment:
|
||||
RUST_LOG: info
|
||||
- image: postgres:14
|
||||
# auth:
|
||||
# username: mydockerhub-user
|
||||
# password: $DOCKERHUB_PASSWORD
|
||||
- image: cimg/postgres:14.0
|
||||
auth:
|
||||
username: mydockerhub-user
|
||||
password: $DOCKERHUB_PASSWORD
|
||||
environment:
|
||||
POSTGRES_USER: postgres
|
||||
POSTGRES_DB: postgres
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_HOST_AUTH_METHOD: scram-sha-256
|
||||
# Add steps to the job
|
||||
# See: https://circleci.com/docs/2.0/configuration-reference/#steps
|
||||
steps:
|
||||
|
||||
@@ -29,6 +29,9 @@ healthcheck_timeout = 100
|
||||
# For how long to ban a server if it fails a health check (seconds).
|
||||
ban_time = 60 # Seconds
|
||||
|
||||
# Stats will be sent here
|
||||
statsd_address = "127.0.0.1:8125"
|
||||
|
||||
#
|
||||
# User to use for authentication against the server.
|
||||
[user]
|
||||
|
||||
@@ -12,11 +12,7 @@ function start_pgcat() {
|
||||
}
|
||||
|
||||
# Setup the database with shards and user
|
||||
PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
|
||||
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard0 -i
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i
|
||||
PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i
|
||||
psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
|
||||
|
||||
# Install Toxiproxy to simulate a downed/slow database
|
||||
wget -O toxiproxy-2.1.4.deb https://github.com/Shopify/toxiproxy/releases/download/v2.1.4/toxiproxy_2.1.4_amd64.deb
|
||||
@@ -31,28 +27,26 @@ toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica
|
||||
|
||||
start_pgcat "info"
|
||||
|
||||
export PGPASSWORD=sharding_user
|
||||
|
||||
# pgbench test
|
||||
pgbench -U sharding_user -i -h 127.0.0.1 -p 6432
|
||||
pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple -f tests/pgbench/simple.sql
|
||||
pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
|
||||
pgbench -i -h 127.0.0.1 -p 6432 && \
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple && \
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
|
||||
|
||||
# COPY TO STDOUT test
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
|
||||
|
||||
# Query cancellation test
|
||||
(psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'SELECT pg_sleep(5)' || true) &
|
||||
(psql -h 127.0.0.1 -p 6432 -c 'SELECT pg_sleep(5)' || true) &
|
||||
killall psql -s SIGINT
|
||||
|
||||
# Sharding insert
|
||||
psql -U sharding_user -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_insert.sql
|
||||
psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_insert.sql
|
||||
|
||||
# Sharding select
|
||||
psql -U sharding_user -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_select.sql > /dev/null
|
||||
psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_select.sql > /dev/null
|
||||
|
||||
# Replica/primary selection & more sharding tests
|
||||
psql -U sharding_user -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
|
||||
psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
|
||||
|
||||
#
|
||||
# ActiveRecord tests
|
||||
@@ -64,25 +58,25 @@ cd tests/ruby && \
|
||||
cd ../..
|
||||
|
||||
# Admin tests
|
||||
psql -U sharding_user -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
|
||||
psql -U sharding_user -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
|
||||
(! psql -U sharding_user -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
|
||||
psql -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
|
||||
(! psql -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
|
||||
|
||||
# Start PgCat in debug to demonstrate failover better
|
||||
start_pgcat "trace"
|
||||
start_pgcat "debug"
|
||||
|
||||
# Add latency to the replica at port 5433 slightly above the healthcheck timeout
|
||||
toxiproxy-cli toxic add -t latency -a latency=300 postgres_replica
|
||||
sleep 1
|
||||
|
||||
# Note the failover in the logs
|
||||
timeout 5 psql -U sharding_user -e -h 127.0.0.1 -p 6432 <<-EOF
|
||||
timeout 5 psql -e -h 127.0.0.1 -p 6432 <<-EOF
|
||||
SELECT 1;
|
||||
SELECT 1;
|
||||
SELECT 1;
|
||||
@@ -100,7 +94,7 @@ sed -i 's/pool_mode = "transaction"/pool_mode = "session"/' pgcat.toml
|
||||
kill -SIGHUP $(pgrep pgcat)
|
||||
|
||||
# Prepared statements that will only work in session mode
|
||||
pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol prepared
|
||||
pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol prepared
|
||||
|
||||
# Attempt clean shut down
|
||||
killall pgcat -s SIGINT
|
||||
|
||||
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1,2 +1,5 @@
|
||||
/target
|
||||
*.deb
|
||||
.idea/*
|
||||
tests/ruby/.bundle/*
|
||||
tests/ruby/vendor/*
|
||||
117
Cargo.lock
generated
117
Cargo.lock
generated
@@ -45,12 +45,6 @@ version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cdb031dd78e28731d87d56cc8ffef4a8f36ca26c38fe2de700543e627f8a464a"
|
||||
|
||||
[[package]]
|
||||
name = "base64"
|
||||
version = "0.13.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd"
|
||||
|
||||
[[package]]
|
||||
name = "bb8"
|
||||
version = "0.7.1"
|
||||
@@ -115,23 +109,22 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "crypto-common"
|
||||
version = "0.1.3"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "57952ca27b5e3606ff4dd79b0020231aaf9d6aa76dc05fd30137538c50bd3ce8"
|
||||
checksum = "683d6b536309245c849479fba3da410962a43ed8e51c26b729208ec0ac2798d0"
|
||||
dependencies = [
|
||||
"generic-array",
|
||||
"typenum",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.3"
|
||||
version = "0.10.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f2fb860ca6fafa5552fb6d0e816a69c8e49f0908bf524e30a90d97c85892d506"
|
||||
checksum = "b697d66081d42af4fba142d56918a3cb21dc8eb63372c6b85d14f44fb9c5979b"
|
||||
dependencies = [
|
||||
"block-buffer",
|
||||
"crypto-common",
|
||||
"subtle",
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -212,15 +205,6 @@ dependencies = [
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e"
|
||||
dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "humantime"
|
||||
version = "2.1.0"
|
||||
@@ -236,6 +220,12 @@ dependencies = [
|
||||
"cfg-if",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1aab8fc367588b89dcee83ab0fd66b72b50b72fa1904d7095045ace2b0c81c35"
|
||||
|
||||
[[package]]
|
||||
name = "libc"
|
||||
version = "0.2.117"
|
||||
@@ -368,16 +358,14 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "pgcat"
|
||||
version = "0.2.0-beta1"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"async-trait",
|
||||
"base64",
|
||||
"bb8",
|
||||
"bytes",
|
||||
"chrono",
|
||||
"env_logger",
|
||||
"hmac",
|
||||
"log",
|
||||
"md-5",
|
||||
"num_cpus",
|
||||
@@ -387,10 +375,10 @@ dependencies = [
|
||||
"regex",
|
||||
"serde",
|
||||
"serde_derive",
|
||||
"serde_json",
|
||||
"sha-1",
|
||||
"sha2",
|
||||
"sqlparser",
|
||||
"stringprep",
|
||||
"statsd",
|
||||
"tokio",
|
||||
"toml",
|
||||
]
|
||||
@@ -482,9 +470,9 @@ dependencies = [
|
||||
|
||||
[[package]]
|
||||
name = "regex"
|
||||
version = "1.5.5"
|
||||
version = "1.5.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1a11647b6b25ff05a515cb92c365cec08801e83423a235b51e231e1808747286"
|
||||
checksum = "d07a8629359eb56f1e2fb1652bb04212c072a87ba68546a04065d525673ac461"
|
||||
dependencies = [
|
||||
"aho-corasick",
|
||||
"memchr",
|
||||
@@ -497,6 +485,12 @@ version = "0.6.25"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b"
|
||||
|
||||
[[package]]
|
||||
name = "ryu"
|
||||
version = "1.0.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "73b4b750c782965c211b42f022f59af1fbceabdd026623714f104152f1ec149f"
|
||||
|
||||
[[package]]
|
||||
name = "scopeguard"
|
||||
version = "1.1.0"
|
||||
@@ -520,6 +514,17 @@ dependencies = [
|
||||
"syn",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "serde_json"
|
||||
version = "1.0.79"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e8d9fa5c3b304765ce1fd9c4c8a3de2c8db365a5b91be52f186efc675681d95"
|
||||
dependencies = [
|
||||
"itoa",
|
||||
"ryu",
|
||||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha-1"
|
||||
version = "0.10.0"
|
||||
@@ -531,17 +536,6 @@ dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sha2"
|
||||
version = "0.10.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "55deaec60f81eefe3cce0dc50bda92d6d8e88f2a27df7c5033b42afeb1ed2676"
|
||||
dependencies = [
|
||||
"cfg-if",
|
||||
"cpufeatures",
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "signal-hook-registry"
|
||||
version = "1.4.0"
|
||||
@@ -573,21 +567,14 @@ dependencies = [
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "stringprep"
|
||||
version = "0.1.2"
|
||||
name = "statsd"
|
||||
version = "0.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ee348cb74b87454fff4b551cbf727025810a004f88aeacae7f85b87f4e9a1c1"
|
||||
checksum = "df1efceb4bf2c0b5ebec94354285a43bbbed1375605bdf2ebe4132299434a330"
|
||||
dependencies = [
|
||||
"unicode-bidi",
|
||||
"unicode-normalization",
|
||||
"rand",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "subtle"
|
||||
version = "2.4.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6bdef32e8150c2a081110b42772ffe7d7c9032b606bc226c8260fd97e0976601"
|
||||
|
||||
[[package]]
|
||||
name = "syn"
|
||||
version = "1.0.86"
|
||||
@@ -619,21 +606,6 @@ dependencies = [
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "87cc5ceb3875bb20c2890005a4e226a4651264a5c75edb2421b52861a0a0cb50"
|
||||
dependencies = [
|
||||
"tinyvec_macros",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinyvec_macros"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c"
|
||||
|
||||
[[package]]
|
||||
name = "tokio"
|
||||
version = "1.16.1"
|
||||
@@ -679,21 +651,6 @@ version = "1.15.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dcf81ac59edc17cc8697ff311e8f5ef2d99fcbd9817b34cec66f90b6c3dfd987"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-bidi"
|
||||
version = "0.3.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "099b7128301d285f79ddd55b9a83d5e6b9e97c92e0ea0daebee7263e932de992"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-normalization"
|
||||
version = "0.1.19"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d54590932941a9e9266f0832deed84ebe1bf2e4c9e4a3554d393d18f5e854bf9"
|
||||
dependencies = [
|
||||
"tinyvec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode-xid"
|
||||
version = "0.2.2"
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "pgcat"
|
||||
version = "0.2.1-beta1"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
|
||||
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
|
||||
@@ -17,15 +17,13 @@ sha-1 = "0.10"
|
||||
toml = "0.5"
|
||||
serde = "1"
|
||||
serde_derive = "1"
|
||||
serde_json = "1"
|
||||
regex = "1"
|
||||
num_cpus = "1"
|
||||
once_cell = "1"
|
||||
statsd = "0.15"
|
||||
sqlparser = "0.14"
|
||||
log = "0.4"
|
||||
arc-swap = "1"
|
||||
env_logger = "0.9"
|
||||
parking_lot = "0.11"
|
||||
hmac = "0.12"
|
||||
sha2 = "0.10"
|
||||
base64 = "0.13"
|
||||
stringprep = "0.1"
|
||||
@@ -8,4 +8,4 @@ COPY --from=builder /app/target/release/pgcat /usr/bin/pgcat
|
||||
COPY --from=builder /app/pgcat.toml /etc/pgcat/pgcat.toml
|
||||
WORKDIR /etc/pgcat
|
||||
ENV RUST_LOG=info
|
||||
CMD ["pgcat"]
|
||||
ENTRYPOINT ["/usr/bin/pgcat"]
|
||||
|
||||
53
README.md
53
README.md
@@ -4,37 +4,28 @@
|
||||
|
||||
|
||||
|
||||
PostgreSQL pooler (like PgBouncer) with sharding, load balancing and failover support.
|
||||
Meow. PgBouncer rewritten in Rust, with sharding, load balancing and failover support.
|
||||
|
||||
**Beta**: looking for beta testers, see [#35](https://github.com/levkk/pgcat/issues/35).
|
||||
**Alpha**: looking for alpha testers, see [#35](https://github.com/levkk/pgcat/issues/35).
|
||||
|
||||
## Features
|
||||
| **Feature** | **Status** | **Comments** |
|
||||
|--------------------------------|-----------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Transaction pooling | :white_check_mark: | Identical to PgBouncer. |
|
||||
| Session pooling | :white_check_mark: | Identical to PgBouncer. |
|
||||
| `COPY` support | :white_check_mark: | Both `COPY TO` and `COPY FROM` are supported. |
|
||||
| Query cancellation | :white_check_mark: | Supported both in transaction and session pooling modes. |
|
||||
| Load balancing of read queries | :white_check_mark: | Using round-robin between replicas. Primary is included when `primary_reads_enabled` is enabled (default). |
|
||||
| Sharding | :white_check_mark: | Transactions are sharded using `SET SHARD TO` and `SET SHARDING KEY TO` syntax extensions; see examples below. |
|
||||
| Failover | :white_check_mark: | Replicas are tested with a health check. If a health check fails, remaining replicas are attempted; see below for algorithm description and examples. |
|
||||
| Statistics | :white_check_mark: | Statistics available in the admin database (`pgcat` and `pgbouncer`) with `SHOW STATS`, `SHOW POOLS` and others. |
|
||||
| Live configuration reloading | :white_check_mark: | Reload supported settings with a `SIGHUP` to the process, e.g. `kill -s SIGHUP $(pgrep pgcat)` or `RELOAD` query issued to the admin database. |
|
||||
| Client authentication | :white_check_mark: :wrench: | MD5 password authentication is supported, SCRAM is on the roadmap; one user is used to connect to Postgres with both SCRAM and MD5 supported. |
|
||||
| Admin database | :white_check_mark: | The admin database, similar to PgBouncer's, allows to query for statistics and reload the configuration. |
|
||||
| **Feature** | **Status** | **Comments** |
|
||||
|--------------------------------|-----------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|
|
||||
| Transaction pooling | :white_check_mark: | Identical to PgBouncer. |
|
||||
| Session pooling | :white_check_mark: | Identical to PgBouncer. |
|
||||
| `COPY` support | :white_check_mark: | Both `COPY TO` and `COPY FROM` are supported. |
|
||||
| Query cancellation | :white_check_mark: | Supported both in transaction and session pooling modes. |
|
||||
| Load balancing of read queries | :white_check_mark: | Using round-robin between replicas. Primary is included when `primary_reads_enabled` is enabled (default). |
|
||||
| Sharding | :white_check_mark: | Transactions are sharded using `SET SHARD TO` and `SET SHARDING KEY TO` syntax extensions; see examples below. |
|
||||
| Failover | :white_check_mark: | Replicas are tested with a health check. If a health check fails, remaining replicas are attempted; see below for algorithm description and examples. |
|
||||
| Statistics reporting | :white_check_mark: | Statistics similar to PgBouncers are reported via StatsD. |
|
||||
| Live configuration reloading | :construction_worker: | Reload config with a `SIGHUP` to the process, e.g. `kill -s SIGHUP $(pgrep pgcat)`. Not all settings can be reloaded without a restart. |
|
||||
| Client authentication | :x: :wrench: | On the roadmap; currently all clients are allowed to connect and one user is used to connect to Postgres. |
|
||||
|
||||
## Deployment
|
||||
|
||||
See `Dockerfile` for example deployment using Docker. The pooler is configured to spawn 4 workers so 4 CPUs are recommended for optimal performance. That setting can be adjusted to spawn as many (or as little) workers as needed.
|
||||
|
||||
For quick local example, use the Docker Compose environment provided:
|
||||
|
||||
```bash
|
||||
docker-compose up
|
||||
|
||||
# In a new terminal:
|
||||
psql -h 127.0.0.1 -p 6432 -c 'SELECT 1'
|
||||
```
|
||||
See `Dockerfile` for example deployment using Docker. The pooler is configured to spawn 4 workers so 4 CPUs are recommended for optimal performance.
|
||||
That setting can be adjusted to spawn as many (or as little) workers as needed.
|
||||
|
||||
### Config
|
||||
|
||||
@@ -48,6 +39,7 @@ psql -h 127.0.0.1 -p 6432 -c 'SELECT 1'
|
||||
| `connect_timeout` | Maximum time to establish a connection to a server (milliseconds). If reached, the server is banned and the next target is attempted. | `5000` |
|
||||
| `healthcheck_timeout` | Maximum time to pass a health check (`SELECT 1`, milliseconds). If reached, the server is banned and the next target is attempted. | `1000` |
|
||||
| `ban_time` | Ban time for a server (seconds). It won't be allowed to serve transactions until the ban expires; failover targets will be used instead. | `60` |
|
||||
| `statsd_address` | StatsD host and port. Statistics will be sent there every 15 seconds. | `127.0.0.1:8125` |
|
||||
| | | |
|
||||
| **`user`** | | |
|
||||
| `name` | The user name. | `sharding_user` |
|
||||
@@ -90,7 +82,7 @@ See [sharding README](./tests/sharding/README.md) for sharding logic testing.
|
||||
| Load balancing | :white_check_mark: | :white_check_mark: | We could test this by emitting statistics for each replica and compare them. |
|
||||
| Failover | :white_check_mark: | :white_check_mark: | Misconfigure a replica in `pgcat.toml` and watch it forward queries to spares. CI testing is using Toxiproxy. |
|
||||
| Sharding | :white_check_mark: | :white_check_mark: | See `tests/sharding` and `tests/ruby` for an Rails/ActiveRecord example. |
|
||||
| Statistics | :white_check_mark: | :white_check_mark: | Query the admin database with `psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS'`. |
|
||||
| Statistics reporting | :x: | :white_check_mark: | Run `nc -l -u 8125` and watch the stats come in every 15 seconds. |
|
||||
| Live config reloading | :white_check_mark: | :white_check_mark: | Run `kill -s SIGHUP $(pgrep pgcat)` and watch the config reload. |
|
||||
|
||||
## Usage
|
||||
@@ -233,15 +225,11 @@ SELECT * FROM users WHERE email = 'test@example.com'; -- shard setting lasts unt
|
||||
|
||||
### Statistics reporting
|
||||
|
||||
The stats are very similar to what Pgbouncer reports and the names are kept to be comparable. They are accessible by querying the admin database `pgcat`, and `pgbouncer` for compatibility.
|
||||
|
||||
```
|
||||
psql -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES'
|
||||
```
|
||||
Stats are reported using StatsD every 15 seconds. The address is configurable with `statsd_address`, the default is `127.0.0.1:8125`. The stats are very similar to what Pgbouncer reports and the names are kept to be comparable.
|
||||
|
||||
### Live configuration reloading
|
||||
|
||||
The config can be reloaded by sending a `kill -s SIGHUP` to the process or by querying `RELOAD` to the admin database. Not all settings are currently supported by live reload:
|
||||
The config can be reloaded by sending a `kill -s SIGHUP` to the process. Not all settings are currently supported by live reload:
|
||||
|
||||
| **Config** | **Requires restart** |
|
||||
|-------------------------|----------------------|
|
||||
@@ -251,6 +239,7 @@ The config can be reloaded by sending a `kill -s SIGHUP` to the process or by qu
|
||||
| `connect_timeout` | yes |
|
||||
| `healthcheck_timeout` | no |
|
||||
| `ban_time` | no |
|
||||
| `statsd_address` | yes |
|
||||
| `user` | yes |
|
||||
| `shards` | yes |
|
||||
| `default_role` | no |
|
||||
|
||||
@@ -1,16 +0,0 @@
|
||||
version: "3"
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:13
|
||||
environment:
|
||||
POSTGRES_PASSWORD: postgres
|
||||
POSTGRES_HOST_AUTH_METHOD: md5
|
||||
pgcat:
|
||||
build: .
|
||||
command:
|
||||
- "pgcat"
|
||||
- "/etc/pgcat/pgcat.toml"
|
||||
volumes:
|
||||
- "${PWD}/examples/docker/pgcat.toml:/etc/pgcat/pgcat.toml"
|
||||
ports:
|
||||
- "6432:6432"
|
||||
@@ -1,105 +0,0 @@
|
||||
#
|
||||
# PgCat config example.
|
||||
#
|
||||
|
||||
#
|
||||
# General pooler settings
|
||||
[general]
|
||||
|
||||
# What IP to run on, 0.0.0.0 means accessible from everywhere.
|
||||
host = "0.0.0.0"
|
||||
|
||||
# Port to run on, same as PgBouncer used in this example.
|
||||
port = 6432
|
||||
|
||||
# How many connections to allocate per server.
|
||||
pool_size = 15
|
||||
|
||||
# Pool mode (see PgBouncer docs for more).
|
||||
# session: one server connection per connected client
|
||||
# transaction: one server connection per client transaction
|
||||
pool_mode = "transaction"
|
||||
|
||||
# How long to wait before aborting a server connection (ms).
|
||||
connect_timeout = 5000
|
||||
|
||||
# How much time to give `SELECT 1` health check query to return with a result (ms).
|
||||
healthcheck_timeout = 1000
|
||||
|
||||
# For how long to ban a server if it fails a health check (seconds).
|
||||
ban_time = 60 # Seconds
|
||||
|
||||
#
|
||||
# User to use for authentication against the server.
|
||||
[user]
|
||||
name = "postgres"
|
||||
password = "postgres"
|
||||
|
||||
|
||||
#
|
||||
# Shards in the cluster
|
||||
[shards]
|
||||
|
||||
# Shard 0
|
||||
[shards.0]
|
||||
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "postgres", 5432, "primary" ],
|
||||
[ "postgres", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
# Database name (e.g. "postgres")
|
||||
database = "postgres"
|
||||
|
||||
[shards.1]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "postgres", 5432, "primary" ],
|
||||
[ "postgres", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "postgres"
|
||||
|
||||
[shards.2]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "postgres", 5432, "primary" ],
|
||||
[ "postgres", 5432, "replica" ],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "postgres"
|
||||
|
||||
|
||||
# Settings for our query routing layer.
|
||||
[query_router]
|
||||
|
||||
# If the client doesn't specify, route traffic to
|
||||
# this role by default.
|
||||
#
|
||||
# any: round-robin between primary and replicas,
|
||||
# replica: round-robin between replicas only without touching the primary,
|
||||
# primary: all queries go to the primary unless otherwise specified.
|
||||
default_role = "any"
|
||||
|
||||
|
||||
# Query parser. If enabled, we'll attempt to parse
|
||||
# every incoming query to determine if it's a read or a write.
|
||||
# If it's a read query, we'll direct it to a replica. Otherwise, if it's a write,
|
||||
# we'll direct it to the primary.
|
||||
query_parser_enabled = false
|
||||
|
||||
# If the query parser is enabled and this setting is enabled, the primary will be part of the pool of databases used for
|
||||
# load balancing of read queries. Otherwise, the primary will only be used for write
|
||||
# queries. The primary can always be explicitely selected with our custom protocol.
|
||||
primary_reads_enabled = true
|
||||
|
||||
# So what if you wanted to implement a different hashing function,
|
||||
# or you've already built one and you want this pooler to use it?
|
||||
#
|
||||
# Current options:
|
||||
#
|
||||
# pg_bigint_hash: PARTITION BY HASH (Postgres hashing function)
|
||||
# sha1: A hashing function based on SHA1
|
||||
#
|
||||
sharding_function = "pg_bigint_hash"
|
||||
15
pgcat.toml
15
pgcat.toml
@@ -29,6 +29,9 @@ healthcheck_timeout = 1000
|
||||
# For how long to ban a server if it fails a health check (seconds).
|
||||
ban_time = 60 # Seconds
|
||||
|
||||
# Stats will be sent here
|
||||
statsd_address = "127.0.0.1:8125"
|
||||
|
||||
#
|
||||
# User to use for authentication against the server.
|
||||
[user]
|
||||
@@ -45,8 +48,8 @@ password = "sharding_user"
|
||||
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
# Database name (e.g. "postgres")
|
||||
@@ -55,8 +58,8 @@ database = "shard0"
|
||||
[shards.1]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "shard1"
|
||||
@@ -64,8 +67,8 @@ database = "shard1"
|
||||
[shards.2]
|
||||
# [ host, port, role ]
|
||||
servers = [
|
||||
[ "127.0.0.1", 5432, "primary" ],
|
||||
[ "localhost", 5432, "replica" ],
|
||||
["127.0.0.1", 5432, "primary"],
|
||||
["localhost", 5432, "replica"],
|
||||
# [ "127.0.1.1", 5432, "replica" ],
|
||||
]
|
||||
database = "shard2"
|
||||
|
||||
101
src/admin.rs
101
src/admin.rs
@@ -1,16 +1,16 @@
|
||||
/// Admin database.
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use log::{info, trace};
|
||||
use std::collections::HashMap;
|
||||
use tokio::net::tcp::OwnedWriteHalf;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::config::{get_config, parse};
|
||||
use crate::errors::Error;
|
||||
use crate::messages::*;
|
||||
use crate::pool::ConnectionPool;
|
||||
use crate::stats::get_stats;
|
||||
|
||||
/// Handle admin client.
|
||||
/// Handle admin client
|
||||
pub async fn handle_admin(
|
||||
stream: &mut OwnedWriteHalf,
|
||||
mut query: BytesMut,
|
||||
@@ -31,7 +31,7 @@ pub async fn handle_admin(
|
||||
|
||||
if query.starts_with("SHOW STATS") {
|
||||
trace!("SHOW STATS");
|
||||
show_stats(stream, &pool).await
|
||||
show_stats(stream).await
|
||||
} else if query.starts_with("RELOAD") {
|
||||
trace!("RELOAD");
|
||||
reload(stream).await
|
||||
@@ -58,7 +58,7 @@ pub async fn handle_admin(
|
||||
}
|
||||
}
|
||||
|
||||
/// Column-oriented statistics.
|
||||
/// SHOW LISTS
|
||||
async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
let stats = get_stats();
|
||||
|
||||
@@ -77,19 +77,11 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
])); // but admin tools that work with pgbouncer want this
|
||||
res.put(data_row(&vec![
|
||||
"free_clients".to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["cl_idle"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
stats["cl_idle"].to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"used_clients".to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["cl_active"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
stats["cl_active"].to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"login_clients".to_string(),
|
||||
@@ -97,19 +89,11 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"free_servers".to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["sv_idle"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
stats["sv_idle"].to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec![
|
||||
"used_servers".to_string(),
|
||||
stats
|
||||
.keys()
|
||||
.map(|address_id| stats[&address_id]["sv_active"])
|
||||
.sum::<i64>()
|
||||
.to_string(),
|
||||
stats["sv_active"].to_string(),
|
||||
]));
|
||||
res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
|
||||
res.put(data_row(&vec!["dns_zones".to_string(), "0".to_string()]));
|
||||
@@ -125,7 +109,7 @@ async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
write_all_half(stream, res).await
|
||||
}
|
||||
|
||||
/// Show PgCat version.
|
||||
/// SHOW VERSION
|
||||
async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
let mut res = BytesMut::new();
|
||||
|
||||
@@ -140,8 +124,8 @@ async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
write_all_half(stream, res).await
|
||||
}
|
||||
|
||||
/// Show utilization of connection pools for each shard and replicas.
|
||||
async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
/// SHOW POOLS
|
||||
async fn show_pools(stream: &mut OwnedWriteHalf, _pool: &ConnectionPool) -> Result<(), Error> {
|
||||
let stats = get_stats();
|
||||
let config = {
|
||||
let guard = get_config();
|
||||
@@ -167,29 +151,18 @@ async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
let mut res = BytesMut::new();
|
||||
res.put(row_description(&columns));
|
||||
|
||||
for shard in 0..pool.shards() {
|
||||
for server in 0..pool.servers(shard) {
|
||||
let address = pool.address(shard, server);
|
||||
let stats = match stats.get(&address.id) {
|
||||
Some(stats) => stats.clone(),
|
||||
None => HashMap::new(),
|
||||
};
|
||||
let mut row = vec![String::from("all"), config.user.name.clone()];
|
||||
|
||||
let mut row = vec![address.name(), config.user.name.clone()];
|
||||
|
||||
for column in &columns[2..columns.len() - 1] {
|
||||
let value = stats.get(column.0).unwrap_or(&0).to_string();
|
||||
row.push(value);
|
||||
}
|
||||
|
||||
row.push(config.general.pool_mode.to_string());
|
||||
res.put(data_row(&row));
|
||||
}
|
||||
for column in &columns[2..columns.len() - 1] {
|
||||
let value = stats.get(column.0).unwrap_or(&0).to_string();
|
||||
row.push(value);
|
||||
}
|
||||
|
||||
row.push(config.general.pool_mode.to_string());
|
||||
|
||||
res.put(data_row(&row));
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
// ReadyForQuery
|
||||
res.put_u8(b'Z');
|
||||
res.put_i32(5);
|
||||
res.put_u8(b'I');
|
||||
@@ -197,7 +170,7 @@ async fn show_pools(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
write_all_half(stream, res).await
|
||||
}
|
||||
|
||||
/// Show shards and replicas.
|
||||
/// SHOW DATABASES
|
||||
async fn show_databases(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
let guard = get_config();
|
||||
let config = &*guard.clone();
|
||||
@@ -222,6 +195,7 @@ async fn show_databases(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> R
|
||||
|
||||
let mut res = BytesMut::new();
|
||||
|
||||
// RowDescription
|
||||
res.put(row_description(&columns));
|
||||
|
||||
for shard in 0..pool.shards() {
|
||||
@@ -265,7 +239,7 @@ async fn ignore_set(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
custom_protocol_response_ok(stream, "SET").await
|
||||
}
|
||||
|
||||
/// Reload the configuration file without restarting the process.
|
||||
/// RELOAD
|
||||
async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
info!("Reloading config");
|
||||
|
||||
@@ -280,6 +254,7 @@ async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
|
||||
let mut res = BytesMut::new();
|
||||
|
||||
// CommandComplete
|
||||
res.put(command_complete("RELOAD"));
|
||||
|
||||
// ReadyForQuery
|
||||
@@ -290,14 +265,13 @@ async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
write_all_half(stream, res).await
|
||||
}
|
||||
|
||||
/// Shows current configuration.
|
||||
async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
let guard = get_config();
|
||||
let config = &*guard.clone();
|
||||
let config: HashMap<String, String> = config.into();
|
||||
drop(guard);
|
||||
|
||||
// Configs that cannot be changed without restarting.
|
||||
// Configs that cannot be changed dynamically.
|
||||
let immutables = ["host", "port", "connect_timeout"];
|
||||
|
||||
// Columns
|
||||
@@ -327,7 +301,6 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
// ReadyForQuery
|
||||
res.put_u8(b'Z');
|
||||
res.put_i32(5);
|
||||
res.put_u8(b'I');
|
||||
@@ -335,8 +308,8 @@ async fn show_config(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
write_all_half(stream, res).await
|
||||
}
|
||||
|
||||
/// Show shard and replicas statistics.
|
||||
async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
|
||||
/// SHOW STATS
|
||||
async fn show_stats(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
|
||||
let columns = vec![
|
||||
("database", DataType::Text),
|
||||
("total_xact_count", DataType::Numeric),
|
||||
@@ -359,27 +332,17 @@ async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Resul
|
||||
let mut res = BytesMut::new();
|
||||
res.put(row_description(&columns));
|
||||
|
||||
for shard in 0..pool.shards() {
|
||||
for server in 0..pool.servers(shard) {
|
||||
let address = pool.address(shard, server);
|
||||
let stats = match stats.get(&address.id) {
|
||||
Some(stats) => stats.clone(),
|
||||
None => HashMap::new(),
|
||||
};
|
||||
let mut row = vec![
|
||||
String::from("all"), // TODO: per-database stats,
|
||||
];
|
||||
|
||||
let mut row = vec![address.name()];
|
||||
|
||||
for column in &columns[1..] {
|
||||
row.push(stats.get(column.0).unwrap_or(&0).to_string());
|
||||
}
|
||||
|
||||
res.put(data_row(&row));
|
||||
}
|
||||
for column in &columns[1..] {
|
||||
row.push(stats.get(column.0).unwrap_or(&0).to_string());
|
||||
}
|
||||
|
||||
res.put(data_row(&row));
|
||||
res.put(command_complete("SHOW"));
|
||||
|
||||
// ReadyForQuery
|
||||
res.put_u8(b'Z');
|
||||
res.put_i32(5);
|
||||
res.put_u8(b'I');
|
||||
|
||||
229
src/client.rs
229
src/client.rs
@@ -1,13 +1,16 @@
|
||||
/// Handle clients by pretending to be a PostgreSQL server.
|
||||
/// Implementation of the PostgreSQL client.
|
||||
/// We are pretending to the server in this scenario,
|
||||
/// and this module implements that.
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use log::{debug, error, trace};
|
||||
use std::collections::HashMap;
|
||||
use tokio::io::{AsyncReadExt, BufReader};
|
||||
use tokio::net::{
|
||||
tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
TcpStream,
|
||||
};
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use crate::admin::handle_admin;
|
||||
use crate::config::get_config;
|
||||
use crate::constants::*;
|
||||
@@ -20,70 +23,67 @@ use crate::stats::Reporter;
|
||||
|
||||
/// The client state. One of these is created per client.
|
||||
pub struct Client {
|
||||
/// The reads are buffered (8K by default).
|
||||
// The reads are buffered (8K by default).
|
||||
read: BufReader<OwnedReadHalf>,
|
||||
|
||||
/// We buffer the writes ourselves because we know the protocol
|
||||
/// better than a stock buffer.
|
||||
// We buffer the writes ourselves because we know the protocol
|
||||
// better than a stock buffer.
|
||||
write: OwnedWriteHalf,
|
||||
|
||||
/// Internal buffer, where we place messages until we have to flush
|
||||
/// them to the backend.
|
||||
// Internal buffer, where we place messages until we have to flush
|
||||
// them to the backend.
|
||||
buffer: BytesMut,
|
||||
|
||||
/// The client was started with the sole reason to cancel another running query.
|
||||
// The client was started with the sole reason to cancel another running query.
|
||||
cancel_mode: bool,
|
||||
|
||||
/// In transaction mode, the connection is released after each transaction.
|
||||
/// Session mode has slightly higher throughput per client, but lower capacity.
|
||||
// In transaction mode, the connection is released after each transaction.
|
||||
// Session mode has slightly higher throughput per client, but lower capacity.
|
||||
transaction_mode: bool,
|
||||
|
||||
/// For query cancellation, the client is given a random process ID and secret on startup.
|
||||
// For query cancellation, the client is given a random process ID and secret on startup.
|
||||
process_id: i32,
|
||||
secret_key: i32,
|
||||
|
||||
/// Clients are mapped to servers while they use them. This allows a client
|
||||
/// to connect and cancel a query.
|
||||
// Clients are mapped to servers while they use them. This allows a client
|
||||
// to connect and cancel a query.
|
||||
client_server_map: ClientServerMap,
|
||||
|
||||
/// Client parameters, e.g. user, client_encoding, etc.
|
||||
// Client parameters, e.g. user, client_encoding, etc.
|
||||
#[allow(dead_code)]
|
||||
parameters: HashMap<String, String>,
|
||||
|
||||
/// Statistics
|
||||
// Statistics
|
||||
stats: Reporter,
|
||||
|
||||
/// Clients want to talk to admin database.
|
||||
// Clients want to talk to admin
|
||||
admin: bool,
|
||||
|
||||
/// Last address the client talked to.
|
||||
last_address_id: Option<usize>,
|
||||
|
||||
/// Last server process id we talked to.
|
||||
last_server_id: Option<i32>,
|
||||
}
|
||||
|
||||
impl Client {
|
||||
/// Perform client startup sequence.
|
||||
/// See docs: <https://www.postgresql.org/docs/12/protocol-flow.html#id-1.10.5.7.3>
|
||||
/// Given a TCP socket, trick the client into thinking we are
|
||||
/// the Postgres server. Perform the authentication and place
|
||||
/// the client in query-ready mode.
|
||||
pub async fn startup(
|
||||
mut stream: TcpStream,
|
||||
client_server_map: ClientServerMap,
|
||||
server_info: BytesMut,
|
||||
stats: Reporter,
|
||||
) -> Result<Client, Error> {
|
||||
let config = get_config().clone();
|
||||
let config = get_config();
|
||||
let transaction_mode = config.general.pool_mode.starts_with("t");
|
||||
// drop(config);
|
||||
drop(config);
|
||||
loop {
|
||||
trace!("Waiting for StartupMessage");
|
||||
|
||||
// Could be StartupMessage, SSLRequest or CancelRequest.
|
||||
// Could be StartupMessage or SSLRequest
|
||||
// which makes this variable length.
|
||||
let len = match stream.read_i32().await {
|
||||
Ok(len) => len,
|
||||
Err(_) => return Err(Error::ClientBadStartup),
|
||||
};
|
||||
|
||||
// Read whatever is left.
|
||||
let mut startup = vec![0u8; len as usize - 4];
|
||||
|
||||
match stream.read_exact(&mut startup).await {
|
||||
@@ -109,51 +109,17 @@ impl Client {
|
||||
PROTOCOL_VERSION_NUMBER => {
|
||||
trace!("Got StartupMessage");
|
||||
let parameters = parse_startup(bytes.clone())?;
|
||||
let mut user_name: String = String::new();
|
||||
match parameters.get(&"user") {
|
||||
Some(&user) => user_name = user,
|
||||
None => return Err(Error::ClientBadStartup),
|
||||
}
|
||||
start_auth(&mut stream, &user_name).await?;
|
||||
|
||||
// Generate random backend ID and secret key
|
||||
let process_id: i32 = rand::random();
|
||||
let secret_key: i32 = rand::random();
|
||||
|
||||
// Perform MD5 authentication.
|
||||
// TODO: Add SASL support.
|
||||
let salt = md5_challenge(&mut stream).await?;
|
||||
|
||||
let code = match stream.read_u8().await {
|
||||
Ok(p) => p,
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
// PasswordMessage
|
||||
if code as char != 'p' {
|
||||
debug!("Expected p, got {}", code as char);
|
||||
return Err(Error::ProtocolSyncError);
|
||||
}
|
||||
|
||||
let len = match stream.read_i32().await {
|
||||
Ok(len) => len,
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
let mut password_response = vec![0u8; (len - 4) as usize];
|
||||
|
||||
match stream.read_exact(&mut password_response).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
// Compare server and client hashes.
|
||||
let password_hash =
|
||||
md5_hash_password(&config.user.name, &config.user.password, &salt);
|
||||
|
||||
if password_hash != password_response {
|
||||
debug!("Password authentication failed");
|
||||
wrong_password(&mut stream, &config.user.name).await?;
|
||||
return Err(Error::ClientError);
|
||||
}
|
||||
|
||||
debug!("Password authentication successful");
|
||||
|
||||
auth_ok(&mut stream).await?;
|
||||
write_all(&mut stream, server_info).await?;
|
||||
backend_key_data(&mut stream, process_id, secret_key).await?;
|
||||
ready_for_query(&mut stream).await?;
|
||||
@@ -184,8 +150,6 @@ impl Client {
|
||||
parameters: parameters,
|
||||
stats: stats,
|
||||
admin: admin,
|
||||
last_address_id: None,
|
||||
last_server_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -208,8 +172,6 @@ impl Client {
|
||||
parameters: HashMap::new(),
|
||||
stats: stats,
|
||||
admin: false,
|
||||
last_address_id: None,
|
||||
last_server_id: None,
|
||||
});
|
||||
}
|
||||
|
||||
@@ -220,7 +182,7 @@ impl Client {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle a connected and authenticated client.
|
||||
/// Client loop. We handle all messages between the client and the database here.
|
||||
pub async fn handle(&mut self, mut pool: ConnectionPool) -> Result<(), Error> {
|
||||
// The client wants to cancel a query it has issued previously.
|
||||
if self.cancel_mode {
|
||||
@@ -256,14 +218,17 @@ impl Client {
|
||||
|
||||
// Our custom protocol loop.
|
||||
// We expect the client to either start a transaction with regular queries
|
||||
// or issue commands for our sharding and server selection protocol.
|
||||
// or issue commands for our sharding and server selection protocols.
|
||||
loop {
|
||||
trace!("Client idle, waiting for message");
|
||||
|
||||
// Client idle, waiting for messages.
|
||||
self.stats.client_idle(self.process_id);
|
||||
|
||||
// Read a complete message from the client, which normally would be
|
||||
// either a `Q` (query) or `P` (prepare, extended protocol).
|
||||
// We can parse it here before grabbing a server from the pool,
|
||||
// in case the client is sending some custom protocol messages, e.g.
|
||||
// in case the client is sending some control messages, e.g.
|
||||
// SET SHARDING KEY TO 'bigint';
|
||||
let mut message = read_message(&mut self.read).await?;
|
||||
|
||||
@@ -273,48 +238,43 @@ impl Client {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Handle admin database queries.
|
||||
// Handle admin database real quick
|
||||
if self.admin {
|
||||
trace!("Handling admin command");
|
||||
handle_admin(&mut self.write, message, pool.clone()).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// Handle all custom protocol commands, if any.
|
||||
// Handle all custom protocol commands here.
|
||||
match query_router.try_execute_command(message.clone()) {
|
||||
// Normal query, not a custom command.
|
||||
// Normal query
|
||||
None => {
|
||||
// Attempt to infer which server we want to query, i.e. primary or replica.
|
||||
if query_router.query_parser_enabled() && query_router.role() == None {
|
||||
query_router.infer_role(message.clone());
|
||||
}
|
||||
}
|
||||
|
||||
// SET SHARD TO
|
||||
Some((Command::SetShard, _)) => {
|
||||
custom_protocol_response_ok(&mut self.write, "SET SHARD").await?;
|
||||
custom_protocol_response_ok(&mut self.write, &format!("SET SHARD")).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// SET SHARDING KEY TO
|
||||
Some((Command::SetShardingKey, _)) => {
|
||||
custom_protocol_response_ok(&mut self.write, "SET SHARDING KEY").await?;
|
||||
custom_protocol_response_ok(&mut self.write, &format!("SET SHARDING KEY"))
|
||||
.await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// SET SERVER ROLE TO
|
||||
Some((Command::SetServerRole, _)) => {
|
||||
custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// SHOW SERVER ROLE
|
||||
Some((Command::ShowServerRole, value)) => {
|
||||
show_response(&mut self.write, "server role", &value).await?;
|
||||
continue;
|
||||
}
|
||||
|
||||
// SHOW SHARD
|
||||
Some((Command::ShowShard, value)) => {
|
||||
show_response(&mut self.write, "shard", &value).await?;
|
||||
continue;
|
||||
@@ -326,7 +286,7 @@ impl Client {
|
||||
error_response(
|
||||
&mut self.write,
|
||||
&format!(
|
||||
"shard {} is more than configured {}",
|
||||
"shard '{}' is more than configured '{}'",
|
||||
query_router.shard(),
|
||||
pool.shards()
|
||||
),
|
||||
@@ -335,13 +295,13 @@ impl Client {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Waiting for server connection.
|
||||
self.stats.client_waiting(self.process_id);
|
||||
|
||||
debug!("Waiting for connection from pool");
|
||||
|
||||
// Grab a server from the pool.
|
||||
let connection = match pool
|
||||
.get(query_router.shard(), query_router.role(), self.process_id)
|
||||
.await
|
||||
{
|
||||
// Grab a server from the pool: the client issued a regular query.
|
||||
let connection = match pool.get(query_router.shard(), query_router.role()).await {
|
||||
Ok(conn) => {
|
||||
debug!("Got connection from pool");
|
||||
conn
|
||||
@@ -355,23 +315,15 @@ impl Client {
|
||||
};
|
||||
|
||||
let mut reference = connection.0;
|
||||
let address = connection.1;
|
||||
let _address = connection.1;
|
||||
let server = &mut *reference;
|
||||
|
||||
// Server is assigned to the client in case the client wants to
|
||||
// cancel a query later.
|
||||
// Claim this server as mine for query cancellation.
|
||||
server.claim(self.process_id, self.secret_key);
|
||||
|
||||
// Update statistics.
|
||||
if let Some(last_address_id) = self.last_address_id {
|
||||
self.stats
|
||||
.client_disconnecting(self.process_id, last_address_id);
|
||||
}
|
||||
self.stats.client_active(self.process_id, address.id);
|
||||
self.stats.server_active(server.process_id(), address.id);
|
||||
|
||||
self.last_address_id = Some(address.id);
|
||||
self.last_server_id = Some(server.process_id());
|
||||
// Client active & server active
|
||||
self.stats.client_active(self.process_id);
|
||||
self.stats.server_active(server.process_id());
|
||||
|
||||
debug!(
|
||||
"Client {:?} talking to server {:?}",
|
||||
@@ -379,20 +331,9 @@ impl Client {
|
||||
server.address()
|
||||
);
|
||||
|
||||
// Set application_name if any.
|
||||
// TODO: investigate other parameters and set them too.
|
||||
if self.parameters.contains_key("application_name") {
|
||||
server
|
||||
.set_name(&self.parameters["application_name"])
|
||||
.await?;
|
||||
}
|
||||
|
||||
// Transaction loop. Multiple queries can be issued by the client here.
|
||||
// The connection belongs to the client until the transaction is over,
|
||||
// or until the client disconnects if we are in session mode.
|
||||
//
|
||||
// If the client is in session mode, no more custom protocol
|
||||
// commands will be accepted.
|
||||
loop {
|
||||
let mut message = if message.len() == 0 {
|
||||
trace!("Waiting for message inside transaction or in session mode");
|
||||
@@ -400,13 +341,11 @@ impl Client {
|
||||
match read_message(&mut self.read).await {
|
||||
Ok(message) => message,
|
||||
Err(err) => {
|
||||
// Client disconnected inside a transaction.
|
||||
// Clean up the server and re-use it.
|
||||
// This prevents connection thrashing by bad clients.
|
||||
// Client disconnected without warning.
|
||||
if server.in_transaction() {
|
||||
server.query("ROLLBACK").await?;
|
||||
server.query("DISCARD ALL").await?;
|
||||
server.set_name("pgcat").await?;
|
||||
// Client left dirty server. Clean up and proceed
|
||||
// without thrashing this connection.
|
||||
server.query("ROLLBACK; DISCARD ALL;").await?;
|
||||
}
|
||||
|
||||
return Err(err);
|
||||
@@ -432,11 +371,13 @@ impl Client {
|
||||
'Q' => {
|
||||
debug!("Sending query to server");
|
||||
|
||||
// TODO: implement retries here for read-only transactions.
|
||||
server.send(original).await?;
|
||||
|
||||
// Read all data the server has to offer, which can be multiple messages
|
||||
// buffered in 8196 bytes chunks.
|
||||
loop {
|
||||
// TODO: implement retries here for read-only transactions.
|
||||
let response = server.recv().await?;
|
||||
|
||||
// Send server reply to the client.
|
||||
@@ -454,16 +395,17 @@ impl Client {
|
||||
}
|
||||
|
||||
// Report query executed statistics.
|
||||
self.stats.query(self.process_id, address.id);
|
||||
self.stats.query();
|
||||
|
||||
// The transaction is over, we can release the connection back to the pool.
|
||||
if !server.in_transaction() {
|
||||
// Report transaction executed statistics.
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
self.stats.transaction();
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
self.stats.server_idle(server.process_id());
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -475,14 +417,11 @@ impl Client {
|
||||
// connection before releasing into the pool.
|
||||
// Pgbouncer closes the connection which leads to
|
||||
// connection thrashing when clients misbehave.
|
||||
// This pool will protect the database. :salute:
|
||||
if server.in_transaction() {
|
||||
server.query("ROLLBACK").await?;
|
||||
server.query("DISCARD ALL").await?;
|
||||
server.set_name("pgcat").await?;
|
||||
server.query("ROLLBACK; DISCARD ALL;").await?;
|
||||
}
|
||||
|
||||
self.release();
|
||||
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
@@ -517,6 +456,7 @@ impl Client {
|
||||
|
||||
self.buffer.put(&original[..]);
|
||||
|
||||
// TODO: retries for read-only transactions.
|
||||
server.send(self.buffer.clone()).await?;
|
||||
|
||||
self.buffer.clear();
|
||||
@@ -524,6 +464,7 @@ impl Client {
|
||||
// Read all data the server has to offer, which can be multiple messages
|
||||
// buffered in 8196 bytes chunks.
|
||||
loop {
|
||||
// TODO: retries for read-only transactions
|
||||
let response = server.recv().await?;
|
||||
|
||||
match write_all_half(&mut self.write, response).await {
|
||||
@@ -540,15 +481,15 @@ impl Client {
|
||||
}
|
||||
|
||||
// Report query executed statistics.
|
||||
self.stats.query(self.process_id, address.id);
|
||||
self.stats.query();
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if !server.in_transaction() {
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
self.stats.transaction();
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
self.stats.server_idle(server.process_id());
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -576,13 +517,13 @@ impl Client {
|
||||
}
|
||||
};
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if !server.in_transaction() {
|
||||
self.stats.transaction(self.process_id, address.id);
|
||||
self.stats.transaction();
|
||||
|
||||
// Release server back to the pool if we are in transaction mode.
|
||||
// If we are in session mode, we keep the server until the client disconnects.
|
||||
if self.transaction_mode {
|
||||
self.stats.server_idle(server.process_id(), address.id);
|
||||
self.stats.server_idle(server.process_id());
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -599,11 +540,10 @@ impl Client {
|
||||
// The server is no longer bound to us, we can't cancel it's queries anymore.
|
||||
debug!("Releasing server back into the pool");
|
||||
self.release();
|
||||
self.stats.client_idle(self.process_id, address.id);
|
||||
}
|
||||
}
|
||||
|
||||
/// Release the server from the client: it can't cancel its queries anymore.
|
||||
/// Release the server from being mine. I can't cancel its queries anymore.
|
||||
pub fn release(&self) {
|
||||
let mut guard = self.client_server_map.lock();
|
||||
guard.remove(&(self.process_id, self.secret_key));
|
||||
@@ -612,13 +552,6 @@ impl Client {
|
||||
|
||||
impl Drop for Client {
|
||||
fn drop(&mut self) {
|
||||
// Update statistics.
|
||||
if let Some(address_id) = self.last_address_id {
|
||||
self.stats.client_disconnecting(self.process_id, address_id);
|
||||
|
||||
if let Some(process_id) = self.last_server_id {
|
||||
self.stats.server_idle(process_id, address_id);
|
||||
}
|
||||
}
|
||||
self.stats.client_disconnecting(self.process_id);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,20 +1,18 @@
|
||||
/// Parse the configuration file.
|
||||
use arc_swap::{ArcSwap, Guard};
|
||||
use log::{error, info};
|
||||
use once_cell::sync::Lazy;
|
||||
use serde_derive::Deserialize;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
use tokio::fs::File;
|
||||
use tokio::io::AsyncReadExt;
|
||||
use toml;
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::errors::Error;
|
||||
|
||||
/// Globally available configuration.
|
||||
static CONFIG: Lazy<ArcSwap<Config>> = Lazy::new(|| ArcSwap::from_pointee(Config::default()));
|
||||
|
||||
/// Server role: primary or replica.
|
||||
#[derive(Clone, PartialEq, Deserialize, Hash, std::cmp::Eq, Debug, Copy)]
|
||||
pub enum Role {
|
||||
Primary,
|
||||
@@ -48,10 +46,8 @@ impl PartialEq<Role> for Option<Role> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Address identifying a PostgreSQL server uniquely.
|
||||
#[derive(Clone, PartialEq, Hash, std::cmp::Eq, Debug)]
|
||||
pub struct Address {
|
||||
pub id: usize,
|
||||
pub host: String,
|
||||
pub port: String,
|
||||
pub shard: usize,
|
||||
@@ -62,7 +58,6 @@ pub struct Address {
|
||||
impl Default for Address {
|
||||
fn default() -> Address {
|
||||
Address {
|
||||
id: 0,
|
||||
host: String::from("127.0.0.1"),
|
||||
port: String::from("5432"),
|
||||
shard: 0,
|
||||
@@ -73,7 +68,6 @@ impl Default for Address {
|
||||
}
|
||||
|
||||
impl Address {
|
||||
/// Address name (aka database) used in `SHOW STATS`, `SHOW DATABASES`, and `SHOW POOLS`.
|
||||
pub fn name(&self) -> String {
|
||||
match self.role {
|
||||
Role::Primary => format!("shard_{}_primary", self.shard),
|
||||
@@ -83,7 +77,6 @@ impl Address {
|
||||
}
|
||||
}
|
||||
|
||||
/// PostgreSQL user.
|
||||
#[derive(Clone, PartialEq, Hash, std::cmp::Eq, Deserialize, Debug)]
|
||||
pub struct User {
|
||||
pub name: String,
|
||||
@@ -99,7 +92,6 @@ impl Default for User {
|
||||
}
|
||||
}
|
||||
|
||||
/// General configuration.
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct General {
|
||||
pub host: String,
|
||||
@@ -109,6 +101,7 @@ pub struct General {
|
||||
pub connect_timeout: u64,
|
||||
pub healthcheck_timeout: u64,
|
||||
pub ban_time: i64,
|
||||
pub statsd_address: String,
|
||||
}
|
||||
|
||||
impl Default for General {
|
||||
@@ -121,11 +114,11 @@ impl Default for General {
|
||||
connect_timeout: 5000,
|
||||
healthcheck_timeout: 1000,
|
||||
ban_time: 60,
|
||||
statsd_address: String::from("127.0.0.1:8125"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Shard configuration.
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct Shard {
|
||||
pub servers: Vec<(String, u16, String)>,
|
||||
@@ -141,7 +134,6 @@ impl Default for Shard {
|
||||
}
|
||||
}
|
||||
|
||||
/// Query Router configuration.
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct QueryRouter {
|
||||
pub default_role: String,
|
||||
@@ -161,7 +153,6 @@ impl Default for QueryRouter {
|
||||
}
|
||||
}
|
||||
|
||||
/// Configuration wrapper.
|
||||
#[derive(Deserialize, Debug, Clone)]
|
||||
pub struct Config {
|
||||
pub path: Option<String>,
|
||||
@@ -205,6 +196,10 @@ impl From<&Config> for std::collections::HashMap<String, String> {
|
||||
config.general.healthcheck_timeout.to_string(),
|
||||
),
|
||||
("ban_time".to_string(), config.general.ban_time.to_string()),
|
||||
(
|
||||
"statsd_address".to_string(),
|
||||
config.general.statsd_address.to_string(),
|
||||
),
|
||||
(
|
||||
"default_role".to_string(),
|
||||
config.query_router.default_role.to_string(),
|
||||
@@ -226,7 +221,6 @@ impl From<&Config> for std::collections::HashMap<String, String> {
|
||||
}
|
||||
|
||||
impl Config {
|
||||
/// Print current configuration.
|
||||
pub fn show(&self) {
|
||||
info!("Pool size: {}", self.general.pool_size);
|
||||
info!("Pool mode: {}", self.general.pool_mode);
|
||||
@@ -241,14 +235,11 @@ impl Config {
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a read-only instance of the configuration
|
||||
/// from anywhere in the app.
|
||||
/// ArcSwap makes this cheap and quick.
|
||||
pub fn get_config() -> Guard<Arc<Config>> {
|
||||
CONFIG.load()
|
||||
}
|
||||
|
||||
/// Parse the configuration file located at the path.
|
||||
/// Parse the config.
|
||||
pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
let mut contents = String::new();
|
||||
let mut file = match File::open(path).await {
|
||||
@@ -359,7 +350,6 @@ pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
|
||||
config.path = Some(path.to_string());
|
||||
|
||||
// Update the configuration globally.
|
||||
CONFIG.store(Arc::new(config.clone()));
|
||||
|
||||
Ok(())
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
/// Various protocol constants, as defined in
|
||||
/// <https://www.postgresql.org/docs/12/protocol-message-formats.html>
|
||||
/// https://www.postgresql.org/docs/12/protocol-message-formats.html
|
||||
/// and elsewhere in the source code.
|
||||
/// Also other constants we use elsewhere.
|
||||
|
||||
// Used in the StartupMessage to indicate regular handshake.
|
||||
pub const PROTOCOL_VERSION_NUMBER: i32 = 196608;
|
||||
@@ -14,13 +15,6 @@ pub const CANCEL_REQUEST_CODE: i32 = 80877102;
|
||||
// AuthenticationMD5Password
|
||||
pub const MD5_ENCRYPTED_PASSWORD: i32 = 5;
|
||||
|
||||
// SASL
|
||||
pub const SASL: i32 = 10;
|
||||
pub const SASL_CONTINUE: i32 = 11;
|
||||
pub const SASL_FINAL: i32 = 12;
|
||||
pub const SCRAM_SHA_256: &str = "SCRAM-SHA-256";
|
||||
pub const NONCE_LENGTH: usize = 24;
|
||||
|
||||
// AuthenticationOk
|
||||
pub const AUTHENTICATION_SUCCESSFUL: i32 = 0;
|
||||
|
||||
|
||||
@@ -1,13 +1,14 @@
|
||||
/// Errors.
|
||||
|
||||
/// Various errors.
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub enum Error {
|
||||
SocketError,
|
||||
// ClientDisconnected,
|
||||
ClientBadStartup,
|
||||
ProtocolSyncError,
|
||||
ServerError,
|
||||
// ServerTimeout,
|
||||
// DirtyServer,
|
||||
BadConfig,
|
||||
BadUserList,
|
||||
AllServersDown,
|
||||
ClientError,
|
||||
AuthenticationError
|
||||
}
|
||||
|
||||
60
src/main.rs
60
src/main.rs
@@ -31,6 +31,7 @@ extern crate once_cell;
|
||||
extern crate serde;
|
||||
extern crate serde_derive;
|
||||
extern crate sqlparser;
|
||||
extern crate statsd;
|
||||
extern crate tokio;
|
||||
extern crate toml;
|
||||
|
||||
@@ -49,25 +50,29 @@ use std::sync::Arc;
|
||||
mod admin;
|
||||
mod client;
|
||||
mod config;
|
||||
mod userlist;
|
||||
mod constants;
|
||||
mod errors;
|
||||
mod messages;
|
||||
mod pool;
|
||||
mod query_router;
|
||||
mod scram;
|
||||
mod server;
|
||||
mod sharding;
|
||||
mod stats;
|
||||
|
||||
// Support for query cancellation: this maps our process_ids and
|
||||
// secret keys to the backend's.
|
||||
use config::get_config;
|
||||
use pool::{ClientServerMap, ConnectionPool};
|
||||
use stats::{Collector, Reporter};
|
||||
|
||||
/// Main!
|
||||
#[tokio::main(worker_threads = 4)]
|
||||
async fn main() {
|
||||
env_logger::init();
|
||||
info!("Welcome to PgCat! Meow.");
|
||||
|
||||
// Prepare regexes
|
||||
if !query_router::QueryRouter::setup() {
|
||||
error!("Could not setup query router");
|
||||
return;
|
||||
@@ -81,6 +86,7 @@ async fn main() {
|
||||
String::from("pgcat.toml")
|
||||
};
|
||||
|
||||
// Prepare the config
|
||||
match config::parse(&config_file).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
@@ -89,9 +95,18 @@ async fn main() {
|
||||
}
|
||||
};
|
||||
|
||||
let config = get_config();
|
||||
let addr = format!("{}:{}", config.general.host, config.general.port);
|
||||
// Prepare user list
|
||||
match userlist::parse("userlist.json").await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
error!("Userlist parse error: {:?}", err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let config = get_config();
|
||||
|
||||
let addr = format!("{}:{}", config.general.host, config.general.port);
|
||||
let listener = match TcpListener::bind(&addr).await {
|
||||
Ok(sock) => sock,
|
||||
Err(err) => {
|
||||
@@ -101,28 +116,22 @@ async fn main() {
|
||||
};
|
||||
|
||||
info!("Running on {}", addr);
|
||||
|
||||
config.show();
|
||||
|
||||
// Tracks which client is connected to which server for query cancellation.
|
||||
let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
|
||||
|
||||
// Statistics reporting.
|
||||
// Collect statistics and send them to StatsD
|
||||
let (tx, rx) = mpsc::channel(100);
|
||||
let collector_tx = tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut stats_collector = Collector::new(rx, collector_tx);
|
||||
stats_collector.collect().await;
|
||||
});
|
||||
|
||||
// Connection pool that allows to query all shards and replicas.
|
||||
let mut pool =
|
||||
ConnectionPool::from_config(client_server_map.clone(), Reporter::new(tx.clone())).await;
|
||||
|
||||
// Statistics collector task.
|
||||
let collector_tx = tx.clone();
|
||||
let addresses = pool.databases();
|
||||
tokio::task::spawn(async move {
|
||||
let mut stats_collector = Collector::new(rx, collector_tx);
|
||||
stats_collector.collect(addresses).await;
|
||||
});
|
||||
|
||||
// Connect to all servers and validate their versions.
|
||||
let server_info = match pool.validate().await {
|
||||
Ok(info) => info,
|
||||
Err(err) => {
|
||||
@@ -133,7 +142,7 @@ async fn main() {
|
||||
|
||||
info!("Waiting for clients");
|
||||
|
||||
// Client connection loop.
|
||||
// Main app runs here.
|
||||
tokio::task::spawn(async move {
|
||||
loop {
|
||||
let pool = pool.clone();
|
||||
@@ -149,7 +158,7 @@ async fn main() {
|
||||
}
|
||||
};
|
||||
|
||||
// Handle client.
|
||||
// Client goes to another thread, bye.
|
||||
tokio::task::spawn(async move {
|
||||
let start = chrono::offset::Utc::now().naive_utc();
|
||||
match client::Client::startup(socket, client_server_map, server_info, reporter)
|
||||
@@ -183,7 +192,7 @@ async fn main() {
|
||||
}
|
||||
});
|
||||
|
||||
// Reload config:
|
||||
// Reload config
|
||||
// kill -SIGHUP $(pgrep pgcat)
|
||||
tokio::task::spawn(async move {
|
||||
let mut stream = unix_signal(SignalKind::hangup()).unwrap();
|
||||
@@ -203,15 +212,16 @@ async fn main() {
|
||||
}
|
||||
});
|
||||
|
||||
// Exit on Ctrl-C (SIGINT) and SIGTERM.
|
||||
let mut term_signal = unix_signal(SignalKind::terminate()).unwrap();
|
||||
// Setup shut down sequence
|
||||
match signal::ctrl_c().await {
|
||||
Ok(()) => {
|
||||
info!("Shutting down...");
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
_ = signal::ctrl_c() => (),
|
||||
_ = term_signal.recv() => (),
|
||||
Err(err) => {
|
||||
error!("Unable to listen for shutdown signal: {}", err);
|
||||
}
|
||||
};
|
||||
|
||||
info!("Shutting down...");
|
||||
}
|
||||
|
||||
/// Format chrono::Duration to be more human-friendly.
|
||||
|
||||
180
src/messages.rs
180
src/messages.rs
@@ -1,5 +1,7 @@
|
||||
/// Helper functions to send one-off protocol messages
|
||||
/// and handle TcpStream (TCP socket).
|
||||
|
||||
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
use md5::{Digest, Md5};
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt, BufReader};
|
||||
@@ -7,10 +9,16 @@ use tokio::net::{
|
||||
tcp::{OwnedReadHalf, OwnedWriteHalf},
|
||||
TcpStream,
|
||||
};
|
||||
use log::{error};
|
||||
|
||||
use crate::errors::Error;
|
||||
use std::collections::HashMap;
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
use crate::userlist::get_user_list;
|
||||
|
||||
|
||||
/// Postgres data type mappings
|
||||
/// used in RowDescription ('T') message.
|
||||
pub enum DataType {
|
||||
@@ -29,6 +37,98 @@ impl From<&DataType> for i32 {
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
1. Generate salt (4 bytes of random data)
|
||||
md5(concat(md5(concat(password, username)), random-salt)))
|
||||
2. Send md5 auth request
|
||||
3. recieve PasswordMessage with salt.
|
||||
4. refactor md5_password function to be reusable
|
||||
5. check username hash combo against file
|
||||
6. AuthenticationOk or ErrorResponse
|
||||
**/
|
||||
pub async fn start_auth(stream: &mut TcpStream, user_name: &String) -> Result<(), Error> {
|
||||
let mut rng = rand::thread_rng();
|
||||
|
||||
//Generate random 4 byte salt
|
||||
let salt = rng.gen::<u32>();
|
||||
|
||||
// Send AuthenticationMD5Password request
|
||||
send_md5_request(stream, salt).await?;
|
||||
|
||||
let code = match stream.read_u8().await {
|
||||
Ok(code) => code as char,
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
match code {
|
||||
// Password response
|
||||
'p' => {
|
||||
fetch_password_and_authenticate(stream, &user_name, &salt).await?;
|
||||
Ok(auth_ok(stream).await?)
|
||||
}
|
||||
_ => {
|
||||
error!("Unknown code: {}", code);
|
||||
return Err(Error::AuthenticationError);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn send_md5_request(stream: &mut TcpStream, salt: u32) -> Result<(), Error> {
|
||||
let mut authentication_md5password = BytesMut::with_capacity(12);
|
||||
authentication_md5password.put_u8(b'R');
|
||||
authentication_md5password.put_i32(12);
|
||||
authentication_md5password.put_i32(5);
|
||||
authentication_md5password.put_u32(salt);
|
||||
|
||||
// Send AuthenticationMD5Password request
|
||||
Ok(write_all(stream, authentication_md5password).await?)
|
||||
}
|
||||
|
||||
pub async fn fetch_password_and_authenticate(stream: &mut TcpStream, user_name: &String, salt: &u32) -> Result<(), Error> {
|
||||
/**
|
||||
1. How do I store the lists of users and paswords? clear text or hash?? wtf
|
||||
2. Add auth to tests
|
||||
**/
|
||||
|
||||
let len = match stream.read_i32().await {
|
||||
Ok(len) => len,
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
// Read whatever is left.
|
||||
let mut password_hash = vec![0u8; len as usize - 4];
|
||||
|
||||
match stream.read_exact(&mut password_hash).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::AuthenticationError),
|
||||
};
|
||||
|
||||
let user_list = get_user_list();
|
||||
let mut password: String = String::new();
|
||||
match user_list.get(&user_name) {
|
||||
Some(&p) => password = p,
|
||||
None => return Err(Error::AuthenticationError),
|
||||
}
|
||||
|
||||
let mut md5 = Md5::new();
|
||||
|
||||
// concat('md5', md5(concat(md5(concat(password, username)), random-salt)))
|
||||
// First pass
|
||||
md5.update(&password.as_bytes());
|
||||
md5.update(&user_name.as_bytes());
|
||||
let output = md5.finalize_reset();
|
||||
// Second pass
|
||||
md5.update(format!("{:x}", output));
|
||||
md5.update(salt.to_be_bytes().to_vec());
|
||||
|
||||
|
||||
let password_string: String = String::from_utf8(password_hash).expect("Could not get password hash");
|
||||
match format!("md5{:x}", md5.finalize()) == password_string {
|
||||
true => Ok(()),
|
||||
_ => Err(Error::AuthenticationError)
|
||||
}
|
||||
}
|
||||
|
||||
/// Tell the client that authentication handshake completed successfully.
|
||||
pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
|
||||
let mut auth_ok = BytesMut::with_capacity(9);
|
||||
@@ -40,26 +140,6 @@ pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
|
||||
Ok(write_all(stream, auth_ok).await?)
|
||||
}
|
||||
|
||||
/// Generate md5 password challenge.
|
||||
pub async fn md5_challenge(stream: &mut TcpStream) -> Result<[u8; 4], Error> {
|
||||
// let mut rng = rand::thread_rng();
|
||||
let salt: [u8; 4] = [
|
||||
rand::random(),
|
||||
rand::random(),
|
||||
rand::random(),
|
||||
rand::random(),
|
||||
];
|
||||
|
||||
let mut res = BytesMut::new();
|
||||
res.put_u8(b'R');
|
||||
res.put_i32(12);
|
||||
res.put_i32(5); // MD5
|
||||
res.put_slice(&salt[..]);
|
||||
|
||||
write_all(stream, res).await?;
|
||||
Ok(salt)
|
||||
}
|
||||
|
||||
/// Give the client the process_id and secret we generated
|
||||
/// used in query cancellation.
|
||||
pub async fn backend_key_data(
|
||||
@@ -180,8 +260,14 @@ pub fn parse_startup(bytes: BytesMut) -> Result<HashMap<String, String>, Error>
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Create md5 password hash given a salt.
|
||||
pub fn md5_hash_password(user: &str, password: &str, salt: &[u8]) -> Vec<u8> {
|
||||
/// Send password challenge response to the server.
|
||||
/// This is the MD5 challenge.
|
||||
pub async fn md5_password(
|
||||
stream: &mut TcpStream,
|
||||
user: &str,
|
||||
password: &str,
|
||||
salt: &[u8],
|
||||
) -> Result<(), Error> {
|
||||
let mut md5 = Md5::new();
|
||||
|
||||
// First pass
|
||||
@@ -200,19 +286,6 @@ pub fn md5_hash_password(user: &str, password: &str, salt: &[u8]) -> Vec<u8> {
|
||||
.collect::<Vec<u8>>();
|
||||
password.push(0);
|
||||
|
||||
password
|
||||
}
|
||||
|
||||
/// Send password challenge response to the server.
|
||||
/// This is the MD5 challenge.
|
||||
pub async fn md5_password(
|
||||
stream: &mut TcpStream,
|
||||
user: &str,
|
||||
password: &str,
|
||||
salt: &[u8],
|
||||
) -> Result<(), Error> {
|
||||
let password = md5_hash_password(user, password, salt);
|
||||
|
||||
let mut message = BytesMut::with_capacity(password.len() as usize + 5);
|
||||
|
||||
message.put_u8(b'p');
|
||||
@@ -249,7 +322,7 @@ pub async fn custom_protocol_response_ok(
|
||||
|
||||
/// Send a custom error message to the client.
|
||||
/// Tell the client we are ready for the next query and no rollback is necessary.
|
||||
/// Docs on error codes: <https://www.postgresql.org/docs/12/errcodes-appendix.html>.
|
||||
/// Docs on error codes: https://www.postgresql.org/docs/12/errcodes-appendix.html
|
||||
pub async fn error_response(stream: &mut OwnedWriteHalf, message: &str) -> Result<(), Error> {
|
||||
let mut error = BytesMut::new();
|
||||
|
||||
@@ -291,39 +364,6 @@ pub async fn error_response(stream: &mut OwnedWriteHalf, message: &str) -> Resul
|
||||
Ok(write_all_half(stream, res).await?)
|
||||
}
|
||||
|
||||
pub async fn wrong_password(stream: &mut TcpStream, user: &str) -> Result<(), Error> {
|
||||
let mut error = BytesMut::new();
|
||||
|
||||
// Error level
|
||||
error.put_u8(b'S');
|
||||
error.put_slice(&b"FATAL\0"[..]);
|
||||
|
||||
// Error level (non-translatable)
|
||||
error.put_u8(b'V');
|
||||
error.put_slice(&b"FATAL\0"[..]);
|
||||
|
||||
// Error code: not sure how much this matters.
|
||||
error.put_u8(b'C');
|
||||
error.put_slice(&b"28P01\0"[..]); // system_error, see Appendix A.
|
||||
|
||||
// The short error message.
|
||||
error.put_u8(b'M');
|
||||
error.put_slice(&format!("password authentication failed for user \"{}\"\0", user).as_bytes());
|
||||
|
||||
// No more fields follow.
|
||||
error.put_u8(0);
|
||||
|
||||
// Compose the two message reply.
|
||||
let mut res = BytesMut::new();
|
||||
|
||||
res.put_u8(b'E');
|
||||
res.put_i32(error.len() as i32 + 4);
|
||||
|
||||
res.put(error);
|
||||
|
||||
write_all(stream, res).await
|
||||
}
|
||||
|
||||
/// Respond to a SHOW SHARD command.
|
||||
pub async fn show_response(
|
||||
stream: &mut OwnedWriteHalf,
|
||||
@@ -399,7 +439,6 @@ pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
|
||||
res
|
||||
}
|
||||
|
||||
/// Create a DataRow message.
|
||||
pub fn data_row(row: &Vec<String>) -> BytesMut {
|
||||
let mut res = BytesMut::new();
|
||||
let mut data_row = BytesMut::new();
|
||||
@@ -419,7 +458,6 @@ pub fn data_row(row: &Vec<String>) -> BytesMut {
|
||||
res
|
||||
}
|
||||
|
||||
/// Create a CommandComplete message.
|
||||
pub fn command_complete(command: &str) -> BytesMut {
|
||||
let cmd = BytesMut::from(format!("{}\0", command).as_bytes());
|
||||
let mut res = BytesMut::new();
|
||||
|
||||
77
src/pool.rs
77
src/pool.rs
@@ -1,23 +1,24 @@
|
||||
/// Pooling, failover and banlist.
|
||||
/// Pooling and failover and banlist.
|
||||
use async_trait::async_trait;
|
||||
use bb8::{ManageConnection, Pool, PooledConnection};
|
||||
use bytes::BytesMut;
|
||||
use chrono::naive::NaiveDateTime;
|
||||
use log::{debug, error, info, warn};
|
||||
use parking_lot::{Mutex, RwLock};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
use crate::config::{get_config, Address, Role, User};
|
||||
use crate::errors::Error;
|
||||
use crate::server::Server;
|
||||
use crate::stats::Reporter;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
// Banlist: bad servers go in here.
|
||||
pub type BanList = Arc<RwLock<Vec<HashMap<Address, NaiveDateTime>>>>;
|
||||
pub type ClientServerMap = Arc<Mutex<HashMap<(i32, i32), (i32, i32, String, String)>>>;
|
||||
|
||||
/// The globally accessible connection pool.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ConnectionPool {
|
||||
databases: Vec<Vec<Pool<ServerPool>>>,
|
||||
@@ -28,7 +29,7 @@ pub struct ConnectionPool {
|
||||
}
|
||||
|
||||
impl ConnectionPool {
|
||||
/// Construct the connection pool from the configuration.
|
||||
/// Construct the connection pool from a config file.
|
||||
pub async fn from_config(
|
||||
client_server_map: ClientServerMap,
|
||||
stats: Reporter,
|
||||
@@ -37,7 +38,6 @@ impl ConnectionPool {
|
||||
let mut shards = Vec::new();
|
||||
let mut addresses = Vec::new();
|
||||
let mut banlist = Vec::new();
|
||||
let mut address_id = 0;
|
||||
let mut shard_ids = config
|
||||
.shards
|
||||
.clone()
|
||||
@@ -63,7 +63,6 @@ impl ConnectionPool {
|
||||
};
|
||||
|
||||
let address = Address {
|
||||
id: address_id,
|
||||
host: server.0.clone(),
|
||||
port: server.1.to_string(),
|
||||
role: role,
|
||||
@@ -71,8 +70,6 @@ impl ConnectionPool {
|
||||
shard: shard_idx.parse::<usize>().unwrap(),
|
||||
};
|
||||
|
||||
address_id += 1;
|
||||
|
||||
if role == Role::Replica {
|
||||
replica_number += 1;
|
||||
}
|
||||
@@ -124,13 +121,9 @@ impl ConnectionPool {
|
||||
pub async fn validate(&mut self) -> Result<BytesMut, Error> {
|
||||
let mut server_infos = Vec::new();
|
||||
|
||||
let stats = self.stats.clone();
|
||||
for shard in 0..self.shards() {
|
||||
for _ in 0..self.servers(shard) {
|
||||
// To keep stats consistent.
|
||||
let fake_process_id = 0;
|
||||
|
||||
let connection = match self.get(shard, None, fake_process_id).await {
|
||||
let connection = match self.get(shard, None).await {
|
||||
Ok(conn) => conn,
|
||||
Err(err) => {
|
||||
error!("Shard {} down or misconfigured: {:?}", shard, err);
|
||||
@@ -144,8 +137,6 @@ impl ConnectionPool {
|
||||
|
||||
let server_info = server.server_info();
|
||||
|
||||
stats.client_disconnecting(fake_process_id, address.id);
|
||||
|
||||
if server_infos.len() > 0 {
|
||||
// Compare against the last server checked.
|
||||
if server_info != server_infos[server_infos.len() - 1] {
|
||||
@@ -174,7 +165,6 @@ impl ConnectionPool {
|
||||
&mut self,
|
||||
shard: usize,
|
||||
role: Option<Role>,
|
||||
process_id: i32,
|
||||
) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
|
||||
let now = Instant::now();
|
||||
let addresses = &self.addresses[shard];
|
||||
@@ -203,9 +193,10 @@ impl ConnectionPool {
|
||||
}
|
||||
|
||||
while allowed_attempts > 0 {
|
||||
// Round-robin replicas.
|
||||
// Round-robin each client's queries.
|
||||
// If a client only sends one query and then disconnects, it doesn't matter
|
||||
// which replica it'll go to.
|
||||
self.round_robin += 1;
|
||||
|
||||
let index = self.round_robin % addresses.len();
|
||||
let address = &addresses[index];
|
||||
|
||||
@@ -222,27 +213,21 @@ impl ConnectionPool {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Indicate we're waiting on a server connection from a pool.
|
||||
self.stats.client_waiting(process_id, address.id);
|
||||
|
||||
// Check if we can connect
|
||||
let mut conn = match self.databases[shard][index].get().await {
|
||||
Ok(conn) => conn,
|
||||
Err(err) => {
|
||||
error!("Banning replica {}, error: {:?}", index, err);
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// // Check if this server is alive with a health check.
|
||||
// // Check if this server is alive with a health check
|
||||
let server = &mut *conn;
|
||||
let healthcheck_timeout = get_config().general.healthcheck_timeout;
|
||||
|
||||
self.stats.server_tested(server.process_id(), address.id);
|
||||
self.stats.server_tested(server.process_id());
|
||||
|
||||
match tokio::time::timeout(
|
||||
tokio::time::Duration::from_millis(healthcheck_timeout),
|
||||
@@ -250,40 +235,29 @@ impl ConnectionPool {
|
||||
)
|
||||
.await
|
||||
{
|
||||
// Check if health check succeeded.
|
||||
// Check if health check succeeded
|
||||
Ok(res) => match res {
|
||||
Ok(_) => {
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
self.stats.server_idle(conn.process_id(), address.id);
|
||||
self.stats.checkout_time(now.elapsed().as_micros());
|
||||
self.stats.server_idle(conn.process_id());
|
||||
return Ok((conn, address.clone()));
|
||||
}
|
||||
|
||||
// Health check failed.
|
||||
Err(_) => {
|
||||
error!("Banning replica {} because of failed health check", index);
|
||||
|
||||
// Don't leave a bad connection in the pool.
|
||||
server.mark_bad();
|
||||
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
},
|
||||
|
||||
// Health check timed out.
|
||||
// Health check never came back, database is really really down
|
||||
Err(_) => {
|
||||
error!("Banning replica {} because of health check timeout", index);
|
||||
// Don't leave a bad connection in the pool.
|
||||
server.mark_bad();
|
||||
|
||||
self.ban(address, shard);
|
||||
self.stats.client_disconnecting(process_id, address.id);
|
||||
self.stats
|
||||
.checkout_time(now.elapsed().as_micros(), process_id, address.id);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
@@ -361,18 +335,14 @@ impl ConnectionPool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the number of configured shards.
|
||||
pub fn shards(&self) -> usize {
|
||||
self.databases.len()
|
||||
}
|
||||
|
||||
/// Get the number of servers (primary and replicas)
|
||||
/// configured for a shard.
|
||||
pub fn servers(&self, shard: usize) -> usize {
|
||||
self.addresses[shard].len()
|
||||
}
|
||||
|
||||
/// Get the total number of servers (databases) we are connected to.
|
||||
pub fn databases(&self) -> usize {
|
||||
let mut databases = 0;
|
||||
for shard in 0..self.shards() {
|
||||
@@ -381,18 +351,15 @@ impl ConnectionPool {
|
||||
databases
|
||||
}
|
||||
|
||||
/// Get pool state for a particular shard server as reported by bb8.
|
||||
pub fn pool_state(&self, shard: usize, server: usize) -> bb8::State {
|
||||
self.databases[shard][server].state()
|
||||
}
|
||||
|
||||
/// Get the address information for a shard server.
|
||||
pub fn address(&self, shard: usize, server: usize) -> &Address {
|
||||
&self.addresses[shard][server]
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper for the bb8 connection pool.
|
||||
pub struct ServerPool {
|
||||
address: Address,
|
||||
user: User,
|
||||
@@ -428,16 +395,14 @@ impl ManageConnection for ServerPool {
|
||||
async fn connect(&self) -> Result<Self::Connection, Self::Error> {
|
||||
info!(
|
||||
"Creating a new connection to {:?} using user {:?}",
|
||||
self.address.name(),
|
||||
self.user.name
|
||||
self.address, self.user.name
|
||||
);
|
||||
|
||||
// Put a temporary process_id into the stats
|
||||
// for server login.
|
||||
let process_id = rand::random::<i32>();
|
||||
self.stats.server_login(process_id, self.address.id);
|
||||
self.stats.server_login(process_id);
|
||||
|
||||
// Connect to the PostgreSQL server.
|
||||
match Server::startup(
|
||||
&self.address,
|
||||
&self.user,
|
||||
@@ -449,12 +414,12 @@ impl ManageConnection for ServerPool {
|
||||
{
|
||||
Ok(conn) => {
|
||||
// Remove the temporary process_id from the stats.
|
||||
self.stats.server_disconnecting(process_id, self.address.id);
|
||||
self.stats.server_disconnecting(process_id);
|
||||
Ok(conn)
|
||||
}
|
||||
Err(err) => {
|
||||
// Remove the temporary process_id from the stats.
|
||||
self.stats.server_disconnecting(process_id, self.address.id);
|
||||
self.stats.server_disconnecting(process_id);
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use crate::config::{get_config, Role};
|
||||
use crate::sharding::{Sharder, ShardingFunction};
|
||||
/// Route queries automatically based on explicitely requested
|
||||
/// or implied query characteristics.
|
||||
use bytes::{Buf, BytesMut};
|
||||
@@ -8,19 +10,14 @@ use sqlparser::ast::Statement::{Query, StartTransaction};
|
||||
use sqlparser::dialect::PostgreSqlDialect;
|
||||
use sqlparser::parser::Parser;
|
||||
|
||||
use crate::config::{get_config, Role};
|
||||
use crate::sharding::{Sharder, ShardingFunction};
|
||||
|
||||
/// Regexes used to parse custom commands.
|
||||
const CUSTOM_SQL_REGEXES: [&str; 5] = [
|
||||
r"(?i)^ *SET SHARDING KEY TO '?([0-9]+)'? *;? *$",
|
||||
r"(?i)^ *SET SHARD TO '?([0-9]+|ANY)'? *;? *$",
|
||||
r"(?i)^ *SET SHARD TO '?([0-9]+)'? *;? *$",
|
||||
r"(?i)^ *SHOW SHARD *;? *$",
|
||||
r"(?i)^ *SET SERVER ROLE TO '(PRIMARY|REPLICA|ANY|AUTO|DEFAULT)' *;? *$",
|
||||
r"(?i)^ *SHOW SERVER ROLE *;? *$",
|
||||
];
|
||||
|
||||
/// Custom commands.
|
||||
#[derive(PartialEq, Debug)]
|
||||
pub enum Command {
|
||||
SetShardingKey,
|
||||
@@ -30,39 +27,37 @@ pub enum Command {
|
||||
ShowServerRole,
|
||||
}
|
||||
|
||||
/// Quickly test for match when a query is received.
|
||||
// Quick test
|
||||
static CUSTOM_SQL_REGEX_SET: OnceCell<RegexSet> = OnceCell::new();
|
||||
|
||||
// Get the value inside the custom command.
|
||||
// Capture value
|
||||
static CUSTOM_SQL_REGEX_LIST: OnceCell<Vec<Regex>> = OnceCell::new();
|
||||
|
||||
/// The query router.
|
||||
pub struct QueryRouter {
|
||||
/// By default, queries go here, unless we have better information
|
||||
/// about what the client wants.
|
||||
// By default, queries go here, unless we have better information
|
||||
// about what the client wants.
|
||||
default_server_role: Option<Role>,
|
||||
|
||||
/// Number of shards in the cluster.
|
||||
// Number of shards in the cluster.
|
||||
shards: usize,
|
||||
|
||||
/// Which shard we should be talking to right now.
|
||||
// Which shard we should be talking to right now.
|
||||
active_shard: Option<usize>,
|
||||
|
||||
/// Which server should we be talking to.
|
||||
// Should we be talking to a primary or a replica?
|
||||
active_role: Option<Role>,
|
||||
|
||||
/// Include the primary into the replica pool for reads.
|
||||
// Include the primary into the replica pool?
|
||||
primary_reads_enabled: bool,
|
||||
|
||||
/// Should we try to parse queries to route them to replicas or primary automatically.
|
||||
// Should we try to parse queries?
|
||||
query_parser_enabled: bool,
|
||||
|
||||
/// Which sharding function we're using.
|
||||
// Which sharding function are we using?
|
||||
sharding_function: ShardingFunction,
|
||||
}
|
||||
|
||||
impl QueryRouter {
|
||||
/// One-time initialization of regexes.
|
||||
pub fn setup() -> bool {
|
||||
let set = match RegexSet::new(&CUSTOM_SQL_REGEXES) {
|
||||
Ok(rgx) => rgx,
|
||||
@@ -93,7 +88,6 @@ impl QueryRouter {
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a new instance of the query router. Each client gets its own.
|
||||
pub fn new() -> QueryRouter {
|
||||
let config = get_config();
|
||||
|
||||
@@ -126,7 +120,6 @@ impl QueryRouter {
|
||||
pub fn try_execute_command(&mut self, mut buf: BytesMut) -> Option<(Command, String)> {
|
||||
let code = buf.get_u8() as char;
|
||||
|
||||
// Only simple protocol supported for commands.
|
||||
if code != 'Q' {
|
||||
return None;
|
||||
}
|
||||
@@ -165,7 +158,8 @@ impl QueryRouter {
|
||||
// figured out a better way just yet. I think I can write a single Regex
|
||||
// that matches all 5 custom SQL patterns, but maybe that's not very legible?
|
||||
//
|
||||
// I think this is faster than running the Regex engine 5 times.
|
||||
// I think this is faster than running the Regex engine 5 times, so
|
||||
// this is a strong maybe for me so far.
|
||||
match regex_list[matches[0]].captures(&query) {
|
||||
Some(captures) => match captures.get(1) {
|
||||
Some(value) => value.as_str().to_string(),
|
||||
@@ -198,10 +192,7 @@ impl QueryRouter {
|
||||
}
|
||||
|
||||
Command::SetShard => {
|
||||
self.active_shard = match value.to_ascii_uppercase().as_ref() {
|
||||
"ANY" => Some(rand::random::<usize>() % self.shards),
|
||||
_ => Some(value.parse::<usize>().unwrap()),
|
||||
};
|
||||
self.active_shard = Some(value.parse::<usize>().unwrap());
|
||||
}
|
||||
|
||||
Command::SetServerRole => {
|
||||
@@ -227,6 +218,7 @@ impl QueryRouter {
|
||||
}
|
||||
|
||||
"default" => {
|
||||
// TODO: reset query parser to default here.
|
||||
self.active_role = self.default_server_role;
|
||||
self.query_parser_enabled = get_config().query_router.query_parser_enabled;
|
||||
self.active_role
|
||||
@@ -248,14 +240,12 @@ impl QueryRouter {
|
||||
let len = buf.get_i32() as usize;
|
||||
|
||||
let query = match code {
|
||||
// Query
|
||||
'Q' => {
|
||||
let query = String::from_utf8_lossy(&buf[..len - 5]).to_string();
|
||||
debug!("Query: '{}'", query);
|
||||
query
|
||||
}
|
||||
|
||||
// Parse (prepared statement)
|
||||
'P' => {
|
||||
let mut start = 0;
|
||||
let mut end;
|
||||
@@ -278,7 +268,6 @@ impl QueryRouter {
|
||||
|
||||
query.replace("$", "") // Remove placeholders turning them into "values"
|
||||
}
|
||||
|
||||
_ => return false,
|
||||
};
|
||||
|
||||
@@ -342,7 +331,6 @@ impl QueryRouter {
|
||||
self.query_parser_enabled
|
||||
}
|
||||
|
||||
/// Allows to toggle primary reads in tests.
|
||||
#[allow(dead_code)]
|
||||
pub fn toggle_primary_reads(&mut self, value: bool) {
|
||||
self.primary_reads_enabled = value;
|
||||
|
||||
320
src/scram.rs
320
src/scram.rs
@@ -1,320 +0,0 @@
|
||||
// SCRAM-SHA-256 authentication. Heavily inspired by
|
||||
// https://github.com/sfackler/rust-postgres/
|
||||
// SASL implementation.
|
||||
|
||||
use bytes::BytesMut;
|
||||
use hmac::{Hmac, Mac};
|
||||
use rand::{self, Rng};
|
||||
use sha2::digest::FixedOutput;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use std::fmt::Write;
|
||||
|
||||
use crate::constants::*;
|
||||
use crate::errors::Error;
|
||||
|
||||
/// Normalize a password string. Postgres
|
||||
/// passwords don't have to be UTF-8.
|
||||
fn normalize(pass: &[u8]) -> Vec<u8> {
|
||||
let pass = match std::str::from_utf8(pass) {
|
||||
Ok(pass) => pass,
|
||||
Err(_) => return pass.to_vec(),
|
||||
};
|
||||
|
||||
match stringprep::saslprep(pass) {
|
||||
Ok(pass) => pass.into_owned().into_bytes(),
|
||||
Err(_) => pass.as_bytes().to_vec(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Keep the SASL state through the exchange.
|
||||
/// It takes 3 messages to complete the authentication.
|
||||
pub struct ScramSha256 {
|
||||
password: String,
|
||||
salted_password: [u8; 32],
|
||||
auth_message: String,
|
||||
message: BytesMut,
|
||||
nonce: String,
|
||||
}
|
||||
|
||||
impl ScramSha256 {
|
||||
/// Create the Scram state from a password. It'll automatically
|
||||
/// generate a nonce.
|
||||
pub fn new(password: &str) -> ScramSha256 {
|
||||
let mut rng = rand::thread_rng();
|
||||
let nonce = (0..NONCE_LENGTH)
|
||||
.map(|_| {
|
||||
let mut v = rng.gen_range(0x21u8..0x7e);
|
||||
if v == 0x2c {
|
||||
v = 0x7e
|
||||
}
|
||||
v as char
|
||||
})
|
||||
.collect::<String>();
|
||||
|
||||
Self::from_nonce(password, &nonce)
|
||||
}
|
||||
|
||||
/// Used for testing.
|
||||
pub fn from_nonce(password: &str, nonce: &str) -> ScramSha256 {
|
||||
let message = BytesMut::from(&format!("{}n=,r={}", "n,,", nonce).as_bytes()[..]);
|
||||
|
||||
ScramSha256 {
|
||||
password: password.to_string(),
|
||||
nonce: String::from(nonce),
|
||||
message,
|
||||
salted_password: [0u8; 32],
|
||||
auth_message: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the current state of the SASL authentication.
|
||||
pub fn message(&mut self) -> BytesMut {
|
||||
self.message.clone()
|
||||
}
|
||||
|
||||
/// Update the state with message received from server.
|
||||
pub fn update(&mut self, message: &BytesMut) -> Result<BytesMut, Error> {
|
||||
let server_message = Message::parse(message)?;
|
||||
|
||||
if !server_message.nonce.starts_with(&self.nonce) {
|
||||
return Err(Error::ProtocolSyncError);
|
||||
}
|
||||
|
||||
let salt = match base64::decode(&server_message.salt) {
|
||||
Ok(salt) => salt,
|
||||
Err(_) => return Err(Error::ProtocolSyncError),
|
||||
};
|
||||
|
||||
let salted_password = Self::hi(
|
||||
&normalize(&self.password.as_bytes()[..]),
|
||||
&salt,
|
||||
server_message.iterations,
|
||||
);
|
||||
|
||||
// Save for verification of final server message.
|
||||
self.salted_password = salted_password;
|
||||
|
||||
let mut hmac = match Hmac::<Sha256>::new_from_slice(&salted_password) {
|
||||
Ok(hmac) => hmac,
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
|
||||
hmac.update(b"Client Key");
|
||||
|
||||
let client_key = hmac.finalize().into_bytes();
|
||||
|
||||
let mut hash = Sha256::default();
|
||||
hash.update(client_key.as_slice());
|
||||
|
||||
let stored_key = hash.finalize_fixed();
|
||||
let mut cbind_input = vec![];
|
||||
cbind_input.extend("n,,".as_bytes());
|
||||
|
||||
let cbind_input = base64::encode(&cbind_input);
|
||||
|
||||
self.message.clear();
|
||||
|
||||
// Start writing the client reply.
|
||||
match write!(
|
||||
&mut self.message,
|
||||
"c={},r={}",
|
||||
cbind_input, server_message.nonce
|
||||
) {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
|
||||
let auth_message = format!(
|
||||
"n=,r={},{},{}",
|
||||
self.nonce,
|
||||
String::from_utf8_lossy(&message[..]),
|
||||
String::from_utf8_lossy(&self.message[..])
|
||||
);
|
||||
|
||||
let mut hmac = match Hmac::<Sha256>::new_from_slice(&stored_key) {
|
||||
Ok(hmac) => hmac,
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
hmac.update(auth_message.as_bytes());
|
||||
|
||||
// Save the auth message for server final message verification.
|
||||
self.auth_message = auth_message;
|
||||
|
||||
let client_signature = hmac.finalize().into_bytes();
|
||||
|
||||
// Sign the client proof.
|
||||
let mut client_proof = client_key;
|
||||
for (proof, signature) in client_proof.iter_mut().zip(client_signature) {
|
||||
*proof ^= signature;
|
||||
}
|
||||
|
||||
match write!(&mut self.message, ",p={}", base64::encode(&*client_proof)) {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
|
||||
Ok(self.message.clone())
|
||||
}
|
||||
|
||||
/// Verify final server message.
|
||||
pub fn finish(&mut self, message: &BytesMut) -> Result<(), Error> {
|
||||
let final_message = FinalMessage::parse(message)?;
|
||||
|
||||
let verifier = match base64::decode(&final_message.value) {
|
||||
Ok(verifier) => verifier,
|
||||
Err(_) => return Err(Error::ProtocolSyncError),
|
||||
};
|
||||
|
||||
let mut hmac = match Hmac::<Sha256>::new_from_slice(&self.salted_password) {
|
||||
Ok(hmac) => hmac,
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
hmac.update(b"Server Key");
|
||||
let server_key = hmac.finalize().into_bytes();
|
||||
|
||||
let mut hmac = match Hmac::<Sha256>::new_from_slice(&server_key) {
|
||||
Ok(hmac) => hmac,
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
};
|
||||
hmac.update(self.auth_message.as_bytes());
|
||||
|
||||
match hmac.verify_slice(&verifier) {
|
||||
Ok(_) => Ok(()),
|
||||
Err(_) => return Err(Error::ServerError),
|
||||
}
|
||||
}
|
||||
|
||||
/// Hash the password with the salt i-times.
|
||||
fn hi(str: &[u8], salt: &[u8], i: u32) -> [u8; 32] {
|
||||
let mut hmac =
|
||||
Hmac::<Sha256>::new_from_slice(str).expect("HMAC is able to accept all key sizes");
|
||||
hmac.update(salt);
|
||||
hmac.update(&[0, 0, 0, 1]);
|
||||
let mut prev = hmac.finalize().into_bytes();
|
||||
|
||||
let mut hi = prev;
|
||||
|
||||
for _ in 1..i {
|
||||
let mut hmac = Hmac::<Sha256>::new_from_slice(str).expect("already checked above");
|
||||
hmac.update(&prev);
|
||||
prev = hmac.finalize().into_bytes();
|
||||
|
||||
for (hi, prev) in hi.iter_mut().zip(prev) {
|
||||
*hi ^= prev;
|
||||
}
|
||||
}
|
||||
|
||||
hi.into()
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse the server challenge.
|
||||
struct Message {
|
||||
nonce: String,
|
||||
salt: String,
|
||||
iterations: u32,
|
||||
}
|
||||
|
||||
impl Message {
|
||||
/// Parse the server SASL challenge.
|
||||
fn parse(message: &BytesMut) -> Result<Message, Error> {
|
||||
let parts = String::from_utf8_lossy(&message[..])
|
||||
.split(",")
|
||||
.map(|s| s.to_string())
|
||||
.collect::<Vec<String>>();
|
||||
|
||||
if parts.len() != 3 {
|
||||
return Err(Error::ProtocolSyncError);
|
||||
}
|
||||
|
||||
let nonce = str::replace(&parts[0], "r=", "");
|
||||
let salt = str::replace(&parts[1], "s=", "");
|
||||
let iterations = match str::replace(&parts[2], "i=", "").parse::<u32>() {
|
||||
Ok(iterations) => iterations,
|
||||
Err(_) => return Err(Error::ProtocolSyncError),
|
||||
};
|
||||
|
||||
Ok(Message {
|
||||
nonce,
|
||||
salt,
|
||||
iterations,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse server final validation message.
|
||||
struct FinalMessage {
|
||||
value: String,
|
||||
}
|
||||
|
||||
impl FinalMessage {
|
||||
/// Parse the server final validation message.
|
||||
pub fn parse(message: &BytesMut) -> Result<FinalMessage, Error> {
|
||||
if !message.starts_with(b"v=") || message.len() < 4 {
|
||||
return Err(Error::ProtocolSyncError);
|
||||
}
|
||||
|
||||
Ok(FinalMessage {
|
||||
value: String::from_utf8_lossy(&message[2..]).to_string(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_server_first_message() {
|
||||
let message = BytesMut::from(
|
||||
&"r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,s=QSXCR+Q6sek8bf92,i=4096".as_bytes()[..],
|
||||
);
|
||||
let message = Message::parse(&message).unwrap();
|
||||
assert_eq!(message.nonce, "fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j");
|
||||
assert_eq!(message.salt, "QSXCR+Q6sek8bf92");
|
||||
assert_eq!(message.iterations, 4096);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_server_last_message() {
|
||||
let f = FinalMessage::parse(&BytesMut::from(
|
||||
&"v=U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw".as_bytes()[..],
|
||||
))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
f.value,
|
||||
"U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw".to_string()
|
||||
);
|
||||
}
|
||||
|
||||
// recorded auth exchange from psql
|
||||
#[test]
|
||||
fn exchange() {
|
||||
let password = "foobar";
|
||||
let nonce = "9IZ2O01zb9IgiIZ1WJ/zgpJB";
|
||||
|
||||
let client_first = "n,,n=,r=9IZ2O01zb9IgiIZ1WJ/zgpJB";
|
||||
let server_first =
|
||||
"r=9IZ2O01zb9IgiIZ1WJ/zgpJBjx/oIRLs02gGSHcw1KEty3eY,s=fs3IXBy7U7+IvVjZ,i\
|
||||
=4096";
|
||||
let client_final =
|
||||
"c=biws,r=9IZ2O01zb9IgiIZ1WJ/zgpJBjx/oIRLs02gGSHcw1KEty3eY,p=AmNKosjJzS3\
|
||||
1NTlQYNs5BTeQjdHdk7lOflDo5re2an8=";
|
||||
let server_final = "v=U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw=";
|
||||
|
||||
let mut scram = ScramSha256::from_nonce(password, nonce);
|
||||
|
||||
let message = scram.message();
|
||||
assert_eq!(std::str::from_utf8(&message).unwrap(), client_first);
|
||||
|
||||
let result = scram
|
||||
.update(&BytesMut::from(&server_first.as_bytes()[..]))
|
||||
.unwrap();
|
||||
assert_eq!(std::str::from_utf8(&result).unwrap(), client_final);
|
||||
|
||||
scram
|
||||
.finish(&BytesMut::from(&server_final.as_bytes()[..]))
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
165
src/server.rs
165
src/server.rs
@@ -1,6 +1,6 @@
|
||||
/// Implementation of the PostgreSQL server (database) protocol.
|
||||
/// Here we are pretending to the a Postgres client.
|
||||
use bytes::{Buf, BufMut, BytesMut};
|
||||
///! Implementation of the PostgreSQL server (database) protocol.
|
||||
///! Here we are pretending to the a Postgres client.
|
||||
use log::{debug, error, info, trace};
|
||||
use tokio::io::{AsyncReadExt, BufReader};
|
||||
use tokio::net::{
|
||||
@@ -12,52 +12,48 @@ use crate::config::{Address, User};
|
||||
use crate::constants::*;
|
||||
use crate::errors::Error;
|
||||
use crate::messages::*;
|
||||
use crate::scram::ScramSha256;
|
||||
use crate::stats::Reporter;
|
||||
use crate::ClientServerMap;
|
||||
|
||||
/// Server state.
|
||||
pub struct Server {
|
||||
/// Server host, e.g. localhost,
|
||||
/// port, e.g. 5432, and role, e.g. primary or replica.
|
||||
// Server host, e.g. localhost,
|
||||
// port, e.g. 5432, and role, e.g. primary or replica.
|
||||
address: Address,
|
||||
|
||||
/// Buffered read socket.
|
||||
// Buffered read socket.
|
||||
read: BufReader<OwnedReadHalf>,
|
||||
|
||||
/// Unbuffered write socket (our client code buffers).
|
||||
// Unbuffered write socket (our client code buffers).
|
||||
write: OwnedWriteHalf,
|
||||
|
||||
/// Our server response buffer. We buffer data before we give it to the client.
|
||||
// Our server response buffer. We buffer data before we give it to the client.
|
||||
buffer: BytesMut,
|
||||
|
||||
/// Server information the server sent us over on startup.
|
||||
// Server information the server sent us over on startup.
|
||||
server_info: BytesMut,
|
||||
|
||||
/// Backend id and secret key used for query cancellation.
|
||||
// Backend id and secret key used for query cancellation.
|
||||
process_id: i32,
|
||||
secret_key: i32,
|
||||
|
||||
/// Is the server inside a transaction or idle.
|
||||
// Is the server inside a transaction or idle.
|
||||
in_transaction: bool,
|
||||
|
||||
/// Is there more data for the client to read.
|
||||
// Is there more data for the client to read.
|
||||
data_available: bool,
|
||||
|
||||
/// Is the server broken? We'll remote it from the pool if so.
|
||||
// Is the server broken? We'll remote it from the pool if so.
|
||||
bad: bool,
|
||||
|
||||
/// Mapping of clients and servers used for query cancellation.
|
||||
// Mapping of clients and servers used for query cancellation.
|
||||
client_server_map: ClientServerMap,
|
||||
|
||||
/// Server connected at.
|
||||
// Server connected at.
|
||||
connected_at: chrono::naive::NaiveDateTime,
|
||||
|
||||
/// Reports various metrics, e.g. data sent & received.
|
||||
// Reports various metrics, e.g. data sent & received.
|
||||
stats: Reporter,
|
||||
|
||||
/// Application name using the server at the moment.
|
||||
application_name: String,
|
||||
}
|
||||
|
||||
impl Server {
|
||||
@@ -81,7 +77,7 @@ impl Server {
|
||||
|
||||
trace!("Sending StartupMessage");
|
||||
|
||||
// StartupMessage
|
||||
// Send the startup packet telling the server we're a normal Postgres client.
|
||||
startup(&mut stream, &user.name, database).await?;
|
||||
|
||||
let mut server_info = BytesMut::new();
|
||||
@@ -90,8 +86,6 @@ impl Server {
|
||||
|
||||
// We'll be handling multiple packets, but they will all be structured the same.
|
||||
// We'll loop here until this exchange is complete.
|
||||
let mut scram = ScramSha256::new(&user.password);
|
||||
|
||||
loop {
|
||||
let code = match stream.read_u8().await {
|
||||
Ok(code) => code as char,
|
||||
@@ -133,91 +127,6 @@ impl Server {
|
||||
|
||||
AUTHENTICATION_SUCCESSFUL => (),
|
||||
|
||||
SASL => {
|
||||
debug!("Starting SASL authentication");
|
||||
let sasl_len = (len - 8) as usize;
|
||||
let mut sasl_auth = vec![0u8; sasl_len];
|
||||
|
||||
match stream.read_exact(&mut sasl_auth).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
let sasl_type = String::from_utf8_lossy(&sasl_auth[..sasl_len - 2]);
|
||||
|
||||
if sasl_type == SCRAM_SHA_256 {
|
||||
debug!("Using {}", SCRAM_SHA_256);
|
||||
|
||||
// Generate client message.
|
||||
let sasl_response = scram.message();
|
||||
|
||||
// SASLInitialResponse (F)
|
||||
let mut res = BytesMut::new();
|
||||
res.put_u8(b'p');
|
||||
|
||||
// length + String length + length + length of sasl response
|
||||
res.put_i32(
|
||||
4 // i32 size
|
||||
+ SCRAM_SHA_256.len() as i32 // length of SASL version string,
|
||||
+ 1 // Null terminator for the SASL version string,
|
||||
+ 4 // i32 size
|
||||
+ sasl_response.len() as i32, // length of SASL response
|
||||
);
|
||||
|
||||
res.put_slice(&format!("{}\0", SCRAM_SHA_256).as_bytes()[..]);
|
||||
res.put_i32(sasl_response.len() as i32);
|
||||
res.put(sasl_response);
|
||||
|
||||
write_all(&mut stream, res).await?;
|
||||
} else {
|
||||
error!("Unsupported SCRAM version: {}", sasl_type);
|
||||
return Err(Error::ServerError);
|
||||
}
|
||||
}
|
||||
|
||||
SASL_CONTINUE => {
|
||||
trace!("Continuing SASL");
|
||||
|
||||
let mut sasl_data = vec![0u8; (len - 8) as usize];
|
||||
|
||||
match stream.read_exact(&mut sasl_data).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
let msg = BytesMut::from(&sasl_data[..]);
|
||||
let sasl_response = scram.update(&msg)?;
|
||||
|
||||
// SASLResponse
|
||||
let mut res = BytesMut::new();
|
||||
res.put_u8(b'p');
|
||||
res.put_i32(4 + sasl_response.len() as i32);
|
||||
res.put(sasl_response);
|
||||
|
||||
write_all(&mut stream, res).await?;
|
||||
}
|
||||
|
||||
SASL_FINAL => {
|
||||
trace!("Final SASL");
|
||||
|
||||
let mut sasl_final = vec![0u8; len as usize - 8];
|
||||
match stream.read_exact(&mut sasl_final).await {
|
||||
Ok(_) => (),
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
match scram.finish(&BytesMut::from(&sasl_final[..])) {
|
||||
Ok(_) => {
|
||||
debug!("SASL authentication successful");
|
||||
}
|
||||
|
||||
Err(err) => {
|
||||
debug!("SASL authentication failed");
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
_ => {
|
||||
error!("Unsupported authentication mechanism: {}", auth_code);
|
||||
return Err(Error::ServerError);
|
||||
@@ -278,7 +187,7 @@ impl Server {
|
||||
// BackendKeyData
|
||||
'K' => {
|
||||
// The frontend must save these values if it wishes to be able to issue CancelRequest messages later.
|
||||
// See: <https://www.postgresql.org/docs/12/protocol-message-formats.html>.
|
||||
// See: https://www.postgresql.org/docs/12/protocol-message-formats.html
|
||||
process_id = match stream.read_i32().await {
|
||||
Ok(id) => id,
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
@@ -299,9 +208,11 @@ impl Server {
|
||||
Err(_) => return Err(Error::SocketError),
|
||||
};
|
||||
|
||||
// This is the last step in the client-server connection setup,
|
||||
// and indicates the server is ready for to query it.
|
||||
let (read, write) = stream.into_split();
|
||||
|
||||
let mut server = Server {
|
||||
return Ok(Server {
|
||||
address: address.clone(),
|
||||
read: BufReader::new(read),
|
||||
write: write,
|
||||
@@ -315,12 +226,7 @@ impl Server {
|
||||
client_server_map: client_server_map,
|
||||
connected_at: chrono::offset::Utc::now().naive_utc(),
|
||||
stats: stats,
|
||||
application_name: String::new(),
|
||||
};
|
||||
|
||||
server.set_name("pgcat").await?;
|
||||
|
||||
return Ok(server);
|
||||
});
|
||||
}
|
||||
|
||||
// We have an unexpected message from the server during this exchange.
|
||||
@@ -362,8 +268,7 @@ impl Server {
|
||||
|
||||
/// Send messages to the server from the client.
|
||||
pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
|
||||
self.stats
|
||||
.data_sent(messages.len(), self.process_id, self.address.id);
|
||||
self.stats.data_sent(messages.len());
|
||||
|
||||
match write_all_half(&mut self.write, messages).await {
|
||||
Ok(_) => Ok(()),
|
||||
@@ -436,7 +341,8 @@ impl Server {
|
||||
// More data is available after this message, this is not the end of the reply.
|
||||
self.data_available = true;
|
||||
|
||||
// Don't flush yet, the more we buffer, the faster this goes...up to a limit.
|
||||
// Don't flush yet, the more we buffer, the faster this goes...
|
||||
// up to a limit of course.
|
||||
if self.buffer.len() >= 8196 {
|
||||
break;
|
||||
}
|
||||
@@ -468,8 +374,7 @@ impl Server {
|
||||
let bytes = self.buffer.clone();
|
||||
|
||||
// Keep track of how much data we got from the server for stats.
|
||||
self.stats
|
||||
.data_received(bytes.len(), self.process_id, self.address.id);
|
||||
self.stats.data_received(bytes.len());
|
||||
|
||||
// Clear the buffer for next query.
|
||||
self.buffer.clear();
|
||||
@@ -504,7 +409,7 @@ impl Server {
|
||||
|
||||
/// Indicate that this server connection cannot be re-used and must be discarded.
|
||||
pub fn mark_bad(&mut self) {
|
||||
error!("Server {:?} marked bad", self.address);
|
||||
error!("Server marked bad");
|
||||
self.bad = true;
|
||||
}
|
||||
|
||||
@@ -544,14 +449,9 @@ impl Server {
|
||||
/// A shorthand for `SET application_name = $1`.
|
||||
#[allow(dead_code)]
|
||||
pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
|
||||
if self.application_name != name {
|
||||
self.application_name = name.to_string();
|
||||
Ok(self
|
||||
.query(&format!("SET application_name = '{}'", name))
|
||||
.await?)
|
||||
} else {
|
||||
Ok(())
|
||||
}
|
||||
Ok(self
|
||||
.query(&format!("SET application_name = '{}'", name))
|
||||
.await?)
|
||||
}
|
||||
|
||||
/// Get the servers address.
|
||||
@@ -560,7 +460,6 @@ impl Server {
|
||||
self.address.clone()
|
||||
}
|
||||
|
||||
/// Get the server's unique identifier.
|
||||
pub fn process_id(&self) -> i32 {
|
||||
self.process_id
|
||||
}
|
||||
@@ -571,8 +470,7 @@ impl Drop for Server {
|
||||
/// the socket is in non-blocking mode, so it may not be ready
|
||||
/// for a write.
|
||||
fn drop(&mut self) {
|
||||
self.stats
|
||||
.server_disconnecting(self.process_id(), self.address.id);
|
||||
self.stats.server_disconnecting(self.process_id());
|
||||
|
||||
let mut bytes = BytesMut::with_capacity(4);
|
||||
bytes.put_u8(b'X');
|
||||
@@ -580,10 +478,9 @@ impl Drop for Server {
|
||||
|
||||
match self.write.try_write(&bytes) {
|
||||
Ok(_) => (),
|
||||
Err(_) => debug!("Dirty shutdown"),
|
||||
Err(_) => (),
|
||||
};
|
||||
|
||||
// Should not matter.
|
||||
self.bad = true;
|
||||
|
||||
let now = chrono::offset::Utc::now().naive_utc();
|
||||
|
||||
@@ -1,27 +1,20 @@
|
||||
/// Implements various sharding functions.
|
||||
use sha1::{Digest, Sha1};
|
||||
|
||||
/// See: <https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/include/catalog/partition.h#L20>.
|
||||
// https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/include/catalog/partition.h#L20
|
||||
const PARTITION_HASH_SEED: u64 = 0x7A5B22367996DCFD;
|
||||
|
||||
/// The sharding functions we support.
|
||||
#[derive(Debug, PartialEq, Copy, Clone)]
|
||||
pub enum ShardingFunction {
|
||||
PgBigintHash,
|
||||
Sha1,
|
||||
}
|
||||
|
||||
/// The sharder.
|
||||
pub struct Sharder {
|
||||
/// Number of shards in the cluster.
|
||||
shards: usize,
|
||||
|
||||
/// The sharding function in use.
|
||||
sharding_function: ShardingFunction,
|
||||
}
|
||||
|
||||
impl Sharder {
|
||||
/// Create new instance of the sharder.
|
||||
pub fn new(shards: usize, sharding_function: ShardingFunction) -> Sharder {
|
||||
Sharder {
|
||||
shards,
|
||||
@@ -29,7 +22,6 @@ impl Sharder {
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute the shard given sharding key.
|
||||
pub fn shard(&self, key: i64) -> usize {
|
||||
match self.sharding_function {
|
||||
ShardingFunction::PgBigintHash => self.pg_bigint_hash(key),
|
||||
@@ -39,7 +31,7 @@ impl Sharder {
|
||||
|
||||
/// Hash function used by Postgres to determine which partition
|
||||
/// to put the row in when using HASH(column) partitioning.
|
||||
/// Source: <https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/common/hashfn.c#L631>.
|
||||
/// Source: https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/common/hashfn.c#L631
|
||||
/// Supports only 1 bigint at the moment, but we can add more later.
|
||||
fn pg_bigint_hash(&self, key: i64) -> usize {
|
||||
let mut lohalf = key as u32;
|
||||
@@ -127,7 +119,6 @@ impl Sharder {
|
||||
a
|
||||
}
|
||||
|
||||
#[inline]
|
||||
fn pg_u32_hash(k: u32) -> u64 {
|
||||
let mut a: u32 = 0x9e3779b9 as u32 + std::mem::size_of::<u32>() as u32 + 3923095 as u32;
|
||||
let mut b = a;
|
||||
|
||||
275
src/stats.rs
275
src/stats.rs
@@ -1,20 +1,17 @@
|
||||
/// Statistics and reporting.
|
||||
use log::info;
|
||||
use log::{debug, info};
|
||||
use once_cell::sync::Lazy;
|
||||
use parking_lot::Mutex;
|
||||
use std::collections::HashMap;
|
||||
use statsd::Client;
|
||||
use tokio::sync::mpsc::{Receiver, Sender};
|
||||
|
||||
/// Latest stats updated every second; used in SHOW STATS and other admin commands.
|
||||
static LATEST_STATS: Lazy<Mutex<HashMap<usize, HashMap<String, i64>>>> =
|
||||
Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Statistics period used for average calculations.
|
||||
/// 15 seconds.
|
||||
static STAT_PERIOD: u64 = 15000;
|
||||
use crate::config::get_config;
|
||||
|
||||
// Stats used in SHOW STATS
|
||||
static LATEST_STATS: Lazy<Mutex<HashMap<String, i64>>> = Lazy::new(|| Mutex::new(HashMap::new()));
|
||||
static STAT_PERIOD: u64 = 15000; //15 seconds
|
||||
|
||||
/// The names for the events reported
|
||||
/// to the statistics collector.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum EventName {
|
||||
CheckoutTime,
|
||||
@@ -31,253 +28,186 @@ enum EventName {
|
||||
ServerTested,
|
||||
ServerLogin,
|
||||
ServerDisconnecting,
|
||||
UpdateStats,
|
||||
UpdateAverages,
|
||||
FlushStatsToStatsD,
|
||||
}
|
||||
|
||||
/// Event data sent to the collector
|
||||
/// from clients and servers.
|
||||
#[derive(Debug)]
|
||||
pub struct Event {
|
||||
/// The name of the event being reported.
|
||||
name: EventName,
|
||||
|
||||
/// The value being reported. Meaning differs based on event name.
|
||||
value: i64,
|
||||
|
||||
/// The client or server connection reporting the event.
|
||||
process_id: i32,
|
||||
|
||||
/// The server the client is connected to.
|
||||
address_id: usize,
|
||||
process_id: Option<i32>,
|
||||
}
|
||||
|
||||
/// The statistics reporter. An instance is given
|
||||
/// to each possible source of statistics,
|
||||
/// e.g. clients, servers, connection pool.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Reporter {
|
||||
tx: Sender<Event>,
|
||||
}
|
||||
|
||||
impl Reporter {
|
||||
/// Create a new Reporter instance.
|
||||
pub fn new(tx: Sender<Event>) -> Reporter {
|
||||
Reporter { tx: tx }
|
||||
}
|
||||
|
||||
/// Report a query executed by a client against
|
||||
/// a server identified by the `address_id`.
|
||||
pub fn query(&self, process_id: i32, address_id: usize) {
|
||||
pub fn query(&self) {
|
||||
let event = Event {
|
||||
name: EventName::Query,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: None,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Report a transaction executed by a client against
|
||||
/// a server identified by the `address_id`.
|
||||
pub fn transaction(&self, process_id: i32, address_id: usize) {
|
||||
pub fn transaction(&self) {
|
||||
let event = Event {
|
||||
name: EventName::Transaction,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: None,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Report data sent to a server identified by `address_id`.
|
||||
/// The `amount` is measured in bytes.
|
||||
pub fn data_sent(&self, amount: usize, process_id: i32, address_id: usize) {
|
||||
pub fn data_sent(&self, amount: usize) {
|
||||
let event = Event {
|
||||
name: EventName::DataSent,
|
||||
value: amount as i64,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: None,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Report data received from a server identified by `address_id`.
|
||||
/// The `amount` is measured in bytes.
|
||||
pub fn data_received(&self, amount: usize, process_id: i32, address_id: usize) {
|
||||
pub fn data_received(&self, amount: usize) {
|
||||
let event = Event {
|
||||
name: EventName::DataReceived,
|
||||
value: amount as i64,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: None,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Time spent waiting to get a healthy connection from the pool
|
||||
/// for a server identified by `address_id`.
|
||||
/// Measured in milliseconds.
|
||||
pub fn checkout_time(&self, ms: u128, process_id: i32, address_id: usize) {
|
||||
pub fn checkout_time(&self, ms: u128) {
|
||||
let event = Event {
|
||||
name: EventName::CheckoutTime,
|
||||
value: ms as i64,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: None,
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a client identified by `process_id` waiting for a connection
|
||||
/// to a server identified by `address_id`.
|
||||
pub fn client_waiting(&self, process_id: i32, address_id: usize) {
|
||||
pub fn client_waiting(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ClientWaiting,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a client identified by `process_id` is done waiting for a connection
|
||||
/// to a server identified by `address_id` and is about to query the server.
|
||||
pub fn client_active(&self, process_id: i32, address_id: usize) {
|
||||
pub fn client_active(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ClientActive,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a client identified by `process_id` is done querying the server
|
||||
/// identified by `address_id` and is no longer active.
|
||||
pub fn client_idle(&self, process_id: i32, address_id: usize) {
|
||||
pub fn client_idle(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ClientIdle,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a client identified by `process_id` is disconecting from the pooler.
|
||||
/// The last server it was connected to is identified by `address_id`.
|
||||
pub fn client_disconnecting(&self, process_id: i32, address_id: usize) {
|
||||
pub fn client_disconnecting(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ClientDisconnecting,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is actively used
|
||||
/// by a client.
|
||||
pub fn server_active(&self, process_id: i32, address_id: usize) {
|
||||
pub fn server_active(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ServerActive,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is no longer
|
||||
/// actively used by a client and is now idle.
|
||||
pub fn server_idle(&self, process_id: i32, address_id: usize) {
|
||||
pub fn server_idle(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ServerIdle,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is attempting
|
||||
/// to login.
|
||||
pub fn server_login(&self, process_id: i32, address_id: usize) {
|
||||
pub fn server_login(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ServerLogin,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a server connection identified by `process_id` for
|
||||
/// a configured server identified by `address_id` is being
|
||||
/// tested before being given to a client.
|
||||
pub fn server_tested(&self, process_id: i32, address_id: usize) {
|
||||
pub fn server_tested(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ServerTested,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
|
||||
/// Reports a server connection identified by `process_id` is disconecting from the pooler.
|
||||
/// The configured server it was connected to is identified by `address_id`.
|
||||
pub fn server_disconnecting(&self, process_id: i32, address_id: usize) {
|
||||
pub fn server_disconnecting(&self, process_id: i32) {
|
||||
let event = Event {
|
||||
name: EventName::ServerDisconnecting,
|
||||
value: 1,
|
||||
process_id: process_id,
|
||||
address_id: address_id,
|
||||
process_id: Some(process_id),
|
||||
};
|
||||
|
||||
let _ = self.tx.try_send(event);
|
||||
}
|
||||
}
|
||||
|
||||
/// The statistics collector which is receiving statistics
|
||||
/// from clients, servers, and the connection pool. There is
|
||||
/// only one collector (kind of like a singleton).
|
||||
/// The collector can trigger events on its own, e.g.
|
||||
/// it updates aggregates every second and averages every
|
||||
/// 15 seconds.
|
||||
pub struct Collector {
|
||||
rx: Receiver<Event>,
|
||||
tx: Sender<Event>,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
impl Collector {
|
||||
/// Create a new collector instance. There should only be one instance
|
||||
/// at a time. This is ensured by mpsc which allows only one receiver.
|
||||
pub fn new(rx: Receiver<Event>, tx: Sender<Event>) -> Collector {
|
||||
Collector { rx, tx }
|
||||
Collector {
|
||||
rx,
|
||||
tx,
|
||||
client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The statistics collection handler. It will collect statistics
|
||||
/// for `address_id`s starting at 0 up to `addresses`.
|
||||
pub async fn collect(&mut self, addresses: usize) {
|
||||
pub async fn collect(&mut self) {
|
||||
info!("Events reporter started");
|
||||
|
||||
let stats_template = HashMap::from([
|
||||
let mut stats = HashMap::from([
|
||||
("total_query_count", 0),
|
||||
("total_xact_count", 0),
|
||||
("total_sent", 0),
|
||||
@@ -302,47 +232,25 @@ impl Collector {
|
||||
("sv_tested", 0),
|
||||
]);
|
||||
|
||||
let mut stats = HashMap::new();
|
||||
|
||||
// Stats saved after each iteration of the flush event. Used in calculation
|
||||
// of averages in the last flush period.
|
||||
let mut old_stats: HashMap<usize, HashMap<String, i64>> = HashMap::new();
|
||||
let mut old_stats: HashMap<String, i64> = HashMap::new();
|
||||
|
||||
// Track which state the client and server are at any given time.
|
||||
let mut client_server_states: HashMap<usize, HashMap<i32, EventName>> = HashMap::new();
|
||||
let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
|
||||
|
||||
// Flush stats to StatsD and calculate averages every 15 seconds.
|
||||
let tx = self.tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut interval =
|
||||
tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD / 15));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
for address_id in 0..addresses {
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::UpdateStats,
|
||||
value: 0,
|
||||
process_id: -1,
|
||||
address_id: address_id,
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let tx = self.tx.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let mut interval =
|
||||
tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
for address_id in 0..addresses {
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::UpdateAverages,
|
||||
value: 0,
|
||||
process_id: -1,
|
||||
address_id: address_id,
|
||||
});
|
||||
}
|
||||
let _ = tx.try_send(Event {
|
||||
name: EventName::FlushStatsToStatsD,
|
||||
value: 0,
|
||||
process_id: None,
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
@@ -356,14 +264,6 @@ impl Collector {
|
||||
}
|
||||
};
|
||||
|
||||
let stats = stats
|
||||
.entry(stat.address_id)
|
||||
.or_insert(stats_template.clone());
|
||||
let client_server_states = client_server_states
|
||||
.entry(stat.address_id)
|
||||
.or_insert(HashMap::new());
|
||||
let old_stats = old_stats.entry(stat.address_id).or_insert(HashMap::new());
|
||||
|
||||
// Some are counters, some are gauges...
|
||||
match stat.name {
|
||||
EventName::Query => {
|
||||
@@ -413,16 +313,16 @@ impl Collector {
|
||||
| EventName::ServerIdle
|
||||
| EventName::ServerTested
|
||||
| EventName::ServerLogin => {
|
||||
client_server_states.insert(stat.process_id, stat.name);
|
||||
client_server_states.insert(stat.process_id.unwrap(), stat.name);
|
||||
}
|
||||
|
||||
EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
|
||||
client_server_states.remove(&stat.process_id);
|
||||
client_server_states.remove(&stat.process_id.unwrap());
|
||||
}
|
||||
|
||||
EventName::UpdateStats => {
|
||||
EventName::FlushStatsToStatsD => {
|
||||
// Calculate connection states
|
||||
for (_, state) in client_server_states.iter() {
|
||||
for (_, state) in &client_server_states {
|
||||
match state {
|
||||
EventName::ClientActive => {
|
||||
let counter = stats.entry("cl_active").or_insert(0);
|
||||
@@ -434,6 +334,11 @@ impl Collector {
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ClientIdle => {
|
||||
let counter = stats.entry("cl_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ServerIdle => {
|
||||
let counter = stats.entry("sv_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
@@ -454,20 +359,39 @@ impl Collector {
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
EventName::ClientIdle => {
|
||||
let counter = stats.entry("cl_idle").or_insert(0);
|
||||
*counter += 1;
|
||||
}
|
||||
|
||||
_ => unreachable!(),
|
||||
};
|
||||
}
|
||||
|
||||
// Calculate averages
|
||||
for stat in &[
|
||||
"avg_query_count",
|
||||
"avgxact_count",
|
||||
"avg_sent",
|
||||
"avg_received",
|
||||
"avg_wait_time",
|
||||
] {
|
||||
let total_name = stat.replace("avg_", "total_");
|
||||
let old_value = old_stats.entry(total_name.clone()).or_insert(0);
|
||||
let new_value = stats.get(total_name.as_str()).unwrap_or(&0).to_owned();
|
||||
let avg = (new_value - *old_value) / (STAT_PERIOD as i64 / 1_000); // Avg / second
|
||||
|
||||
stats.insert(stat, avg);
|
||||
*old_value = new_value;
|
||||
}
|
||||
|
||||
debug!("{:?}", stats);
|
||||
|
||||
// Update latest stats used in SHOW STATS
|
||||
let mut guard = LATEST_STATS.lock();
|
||||
for (key, value) in &stats {
|
||||
guard.insert(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
let mut pipeline = self.client.pipeline();
|
||||
|
||||
for (key, value) in stats.iter() {
|
||||
let entry = guard.entry(stat.address_id).or_insert(HashMap::new());
|
||||
entry.insert(key.to_string(), value.clone());
|
||||
pipeline.gauge(key, *value as f64);
|
||||
}
|
||||
|
||||
// These are re-calculated every iteration of the loop, so we don't want to add values
|
||||
@@ -485,33 +409,14 @@ impl Collector {
|
||||
] {
|
||||
stats.insert(stat, 0);
|
||||
}
|
||||
}
|
||||
|
||||
EventName::UpdateAverages => {
|
||||
// Calculate averages
|
||||
for stat in &[
|
||||
"avg_query_count",
|
||||
"avgxact_count",
|
||||
"avg_sent",
|
||||
"avg_received",
|
||||
"avg_wait_time",
|
||||
] {
|
||||
let total_name = stat.replace("avg_", "total_");
|
||||
let old_value = old_stats.entry(total_name.clone()).or_insert(0);
|
||||
let new_value = stats.get(total_name.as_str()).unwrap_or(&0).to_owned();
|
||||
let avg = (new_value - *old_value) / (STAT_PERIOD as i64 / 1_000); // Avg / second
|
||||
|
||||
stats.insert(stat, avg);
|
||||
*old_value = new_value;
|
||||
}
|
||||
pipeline.send(&self.client);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a snapshot of statistics. Updated once a second
|
||||
/// by the `Collector`.
|
||||
pub fn get_stats() -> HashMap<usize, HashMap<String, i64>> {
|
||||
pub fn get_stats() -> HashMap<String, i64> {
|
||||
LATEST_STATS.lock().clone()
|
||||
}
|
||||
|
||||
4
src/userlist.json
Normal file
4
src/userlist.json
Normal file
@@ -0,0 +1,4 @@
|
||||
{
|
||||
"sven": "clear_text_password",
|
||||
"sharding_user": "sharding_user"
|
||||
}
|
||||
57
src/userlist.rs
Normal file
57
src/userlist.rs
Normal file
@@ -0,0 +1,57 @@
|
||||
use arc_swap::{ArcSwap, Guard};
|
||||
use log::{error};
|
||||
use once_cell::sync::Lazy;
|
||||
use tokio::fs::File;
|
||||
use tokio::io::AsyncReadExt;
|
||||
|
||||
use std::collections::{HashMap};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::errors::Error;
|
||||
|
||||
pub type UserList = HashMap<String, String>;
|
||||
static USER_LIST: Lazy<ArcSwap<UserList>> = Lazy::new(|| ArcSwap::from_pointee(HashMap::new()));
|
||||
|
||||
pub fn get_user_list() -> Guard<Arc<UserList>> {
|
||||
USER_LIST.load()
|
||||
}
|
||||
|
||||
/// Parse the user list.
|
||||
pub async fn parse(path: &str) -> Result<(), Error> {
|
||||
let mut contents = String::new();
|
||||
let mut file = match File::open(path).await {
|
||||
Ok(file) => file,
|
||||
Err(err) => {
|
||||
error!("Could not open '{}': {}", path, err.to_string());
|
||||
return Err(Error::BadConfig);
|
||||
}
|
||||
};
|
||||
|
||||
match file.read_to_string(&mut contents).await {
|
||||
Ok(_) => (),
|
||||
Err(err) => {
|
||||
error!("Could not read config file: {}", err.to_string());
|
||||
return Err(Error::BadConfig);
|
||||
}
|
||||
};
|
||||
|
||||
let map: HashMap<String, String> = serde_json::from_str(&contents).expect("JSON was not well-formatted");
|
||||
|
||||
|
||||
|
||||
USER_LIST.store(Arc::new(map.clone()));
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_config() {
|
||||
parse("userlist.json").await.unwrap();
|
||||
assert_eq!(get_user_list()["sven"], "clear_text_password");
|
||||
assert_eq!(get_user_list()["sharding_user"], "sharding_user");
|
||||
}
|
||||
}
|
||||
@@ -1,28 +0,0 @@
|
||||
|
||||
-- \setrandom aid 1 :naccounts
|
||||
\set aid random(1, 100000)
|
||||
-- \setrandom bid 1 :nbranches
|
||||
\set bid random(1, 100000)
|
||||
-- \setrandom tid 1 :ntellers
|
||||
\set tid random(1, 100000)
|
||||
-- \setrandom delta -5000 5000
|
||||
\set delta random(-5000,5000)
|
||||
|
||||
\set shard random(0, 2)
|
||||
|
||||
SET SHARD TO :shard;
|
||||
|
||||
BEGIN;
|
||||
|
||||
UPDATE pgbench_accounts SET abalance = abalance + :delta WHERE aid = :aid;
|
||||
|
||||
SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
|
||||
|
||||
UPDATE pgbench_tellers SET tbalance = tbalance + :delta WHERE tid = :tid;
|
||||
|
||||
UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid;
|
||||
|
||||
INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP);
|
||||
|
||||
END;
|
||||
|
||||
@@ -2,4 +2,3 @@ source "https://rubygems.org"
|
||||
|
||||
gem "pg"
|
||||
gem "activerecord"
|
||||
gem "rubocop"
|
||||
|
||||
@@ -11,33 +11,13 @@ GEM
|
||||
i18n (>= 1.6, < 2)
|
||||
minitest (>= 5.1)
|
||||
tzinfo (~> 2.0)
|
||||
ast (2.4.2)
|
||||
concurrent-ruby (1.1.9)
|
||||
i18n (1.10.0)
|
||||
concurrent-ruby (~> 1.0)
|
||||
minitest (5.15.0)
|
||||
parallel (1.22.1)
|
||||
parser (3.1.2.0)
|
||||
ast (~> 2.4.1)
|
||||
pg (1.3.2)
|
||||
rainbow (3.1.1)
|
||||
regexp_parser (2.3.1)
|
||||
rexml (3.2.5)
|
||||
rubocop (1.29.0)
|
||||
parallel (~> 1.10)
|
||||
parser (>= 3.1.0.0)
|
||||
rainbow (>= 2.2.2, < 4.0)
|
||||
regexp_parser (>= 1.8, < 3.0)
|
||||
rexml (>= 3.2.5, < 4.0)
|
||||
rubocop-ast (>= 1.17.0, < 2.0)
|
||||
ruby-progressbar (~> 1.7)
|
||||
unicode-display_width (>= 1.4.0, < 3.0)
|
||||
rubocop-ast (1.17.0)
|
||||
parser (>= 3.1.1.0)
|
||||
ruby-progressbar (1.11.0)
|
||||
tzinfo (2.0.4)
|
||||
concurrent-ruby (~> 1.0)
|
||||
unicode-display_width (2.1.0)
|
||||
|
||||
PLATFORMS
|
||||
x86_64-linux
|
||||
@@ -45,7 +25,6 @@ PLATFORMS
|
||||
DEPENDENCIES
|
||||
activerecord
|
||||
pg
|
||||
rubocop
|
||||
|
||||
BUNDLED WITH
|
||||
2.3.7
|
||||
|
||||
@@ -1,9 +1,6 @@
|
||||
# frozen_string_literal: true
|
||||
|
||||
require 'active_record'
|
||||
require 'pg'
|
||||
|
||||
$stdout.sync = true
|
||||
|
||||
# Uncomment these two to see all queries.
|
||||
# ActiveRecord.verbose_query_logs = true
|
||||
@@ -16,7 +13,6 @@ ActiveRecord::Base.establish_connection(
|
||||
username: 'sharding_user',
|
||||
password: 'sharding_user',
|
||||
database: 'rails_dev',
|
||||
application_name: 'testing_pgcat',
|
||||
prepared_statements: false, # Transaction mode
|
||||
advisory_locks: false # Same
|
||||
)
|
||||
@@ -114,17 +110,3 @@ begin
|
||||
rescue ActiveRecord::StatementInvalid
|
||||
puts 'OK'
|
||||
end
|
||||
|
||||
# Test evil clients
|
||||
def poorly_behaved_client
|
||||
conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/rails_dev?application_name=testing_pgcat")
|
||||
conn.async_exec 'BEGIN'
|
||||
conn.async_exec 'SELECT 1'
|
||||
|
||||
conn.close
|
||||
puts 'Bad client ok'
|
||||
end
|
||||
|
||||
25.times do
|
||||
poorly_behaved_client
|
||||
end
|
||||
|
||||
Reference in New Issue
Block a user