readme

* client md5 auth and clean up scram

* add pw

* add user

* add user

* log

.circleci/config.yml

@@ -12,13 +12,15 @@ jobs:
       - image: cimg/rust:1.58.1
         environment:
           RUST_LOG: info
-      - image: cimg/postgres:14.0
+      - image: postgres:14
-        auth:
+        # auth:
-          username: mydockerhub-user
+        #   username: mydockerhub-user
-          password: $DOCKERHUB_PASSWORD
+        #   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:

.circleci/pgcat.toml

@@ -29,9 +29,6 @@ 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]

.circleci/run_tests.sh

@@ -12,7 +12,11 @@ function start_pgcat() {
 }
 
 # Setup the database with shards and user
-psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
+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
 
 # 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
@@ -27,26 +31,28 @@ 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 -i -h 127.0.0.1 -p 6432 && \
+pgbench -U sharding_user -i -h 127.0.0.1 -p 6432
-    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple && \
+pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple -f tests/pgbench/simple.sql
-    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
+pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
 
 # COPY TO STDOUT test
-psql -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
+psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
 
 # Query cancellation test
-(psql -h 127.0.0.1 -p 6432 -c 'SELECT pg_sleep(5)' || true) &
+(psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'SELECT pg_sleep(5)' || true) &
 killall psql -s SIGINT
 
 # Sharding insert
-psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_insert.sql
+psql -U sharding_user -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_insert.sql
 
 # Sharding select
-psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_select.sql > /dev/null
+psql -U sharding_user -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 -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
+psql -U sharding_user -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
 
 #
 # ActiveRecord tests
@@ -57,15 +63,26 @@ cd tests/ruby && \
     ruby tests.rb && \
 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)
+
 # Start PgCat in debug to demonstrate failover better
-start_pgcat "debug"
+start_pgcat "trace"
 
 # 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 -e -h 127.0.0.1 -p 6432 <<-EOF
+timeout 5 psql -U sharding_user -e -h 127.0.0.1 -p 6432 <<-EOF
 SELECT 1;
 SELECT 1;
 SELECT 1;
@@ -83,7 +100,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 -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol prepared
+pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol prepared
 
 # Attempt clean shut down
 killall pgcat -s SIGINT

Cargo.lock

@@ -45,6 +45,12 @@ 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"
@@ -109,22 +115,23 @@ dependencies = [
 
 [[package]]
 name = "crypto-common"
-version = "0.1.1"
+version = "0.1.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "683d6b536309245c849479fba3da410962a43ed8e51c26b729208ec0ac2798d0"
+checksum = "57952ca27b5e3606ff4dd79b0020231aaf9d6aa76dc05fd30137538c50bd3ce8"
 dependencies = [
  "generic-array",
+ "typenum",
 ]
 
 [[package]]
 name = "digest"
-version = "0.10.1"
+version = "0.10.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b697d66081d42af4fba142d56918a3cb21dc8eb63372c6b85d14f44fb9c5979b"
+checksum = "f2fb860ca6fafa5552fb6d0e816a69c8e49f0908bf524e30a90d97c85892d506"
 dependencies = [
  "block-buffer",
  "crypto-common",
- "generic-array",
+ "subtle",
 ]
 
 [[package]]
@@ -205,6 +212,15 @@ 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"
@@ -352,14 +368,16 @@ dependencies = [
 
 [[package]]
 name = "pgcat"
-version = "0.1.0"
+version = "0.2.0-beta1"
 dependencies = [
  "arc-swap",
  "async-trait",
+ "base64",
  "bb8",
  "bytes",
  "chrono",
  "env_logger",
+ "hmac",
  "log",
  "md-5",
  "num_cpus",
@@ -370,8 +388,9 @@ dependencies = [
  "serde",
  "serde_derive",
  "sha-1",
+ "sha2",
  "sqlparser",
- "statsd",
+ "stringprep",
  "tokio",
  "toml",
 ]
@@ -463,9 +482,9 @@ dependencies = [
 
 [[package]]
 name = "regex"
-version = "1.5.4"
+version = "1.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d07a8629359eb56f1e2fb1652bb04212c072a87ba68546a04065d525673ac461"
+checksum = "1a11647b6b25ff05a515cb92c365cec08801e83423a235b51e231e1808747286"
 dependencies = [
  "aho-corasick",
  "memchr",
@@ -512,6 +531,17 @@ 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"
@@ -543,14 +573,21 @@ dependencies = [
 ]
 
 [[package]]
-name = "statsd"
+name = "stringprep"
-version = "0.15.0"
+version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "df1efceb4bf2c0b5ebec94354285a43bbbed1375605bdf2ebe4132299434a330"
+checksum = "8ee348cb74b87454fff4b551cbf727025810a004f88aeacae7f85b87f4e9a1c1"
 dependencies = [
- "rand",
+ "unicode-bidi",
+ "unicode-normalization",
 ]
 
+[[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"
@@ -582,6 +619,21 @@ 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"
@@ -627,6 +679,21 @@ 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"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pgcat"
-version = "0.1.0"
+version = "0.2.1-beta1"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -20,9 +20,12 @@ serde_derive = "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"

Dockerfile

@@ -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
-ENTRYPOINT ["/usr/bin/pgcat"]
+CMD ["pgcat"]

README.md

@@ -4,28 +4,37 @@
 
 ![PgCat](./pgcat3.png)
 
-Meow. PgBouncer rewritten in Rust, with sharding, load balancing and failover support.
+PostgreSQL pooler (like PgBouncer) with sharding, load balancing and failover support.
 
-**Alpha**: looking for alpha testers, see [#35](https://github.com/levkk/pgcat/issues/35).
+**Beta**: looking for beta testers, see [#35](https://github.com/levkk/pgcat/issues/35).
 
 ## Features
-| **Feature**                    | **Status**            | **Comments**                                                                                                                                          |
+| **Feature**                    | **Status**                  | **Comments**                                                                                                                                          |
-|--------------------------------|-----------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|
+|--------------------------------|-----------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Transaction pooling            | :white_check_mark:    | Identical to PgBouncer.                                                                                                                               |
+| Transaction pooling            | :white_check_mark:          | Identical to PgBouncer.                                                                                                                               |
-| Session 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.                                                                                                         |
+| `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.                                                                                              |
+| 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).                                            |
+| 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.                                        |
+| 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. |
+| 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.                                                                                             |
+| Statistics                     | :white_check_mark:          | Statistics available in the admin database (`pgcat` and `pgbouncer`) with `SHOW STATS`, `SHOW POOLS` and others.                                      |
-| 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.               |
+| 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          | :x: :wrench:          | On the roadmap; currently all clients are allowed to connect and one user is used to connect to Postgres.                                             |
+| 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.                                              |
 
 ## Deployment
 
-See `Dockerfile` for example deployment using Docker. The pooler is configured to spawn 4 workers so 4 CPUs are recommended for optimal performance.
+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.
-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'
+```
 
 ### Config
 
@@ -39,7 +48,6 @@ That setting can be adjusted to spawn as many (or as little) workers as needed.
 | `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`                  |
@@ -82,7 +90,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 reporting  | :x:                | :white_check_mark:  | Run `nc -l -u 8125` and watch the stats come in every 15 seconds.                                                        |
+| 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'`.                                  |
 | Live config reloading | :white_check_mark: | :white_check_mark:  | Run `kill -s SIGHUP $(pgrep pgcat)` and watch the config reload.                                                         |
 
 ## Usage
@@ -225,11 +233,15 @@ SELECT * FROM users WHERE email = 'test@example.com'; -- shard setting lasts unt
 
 ### Statistics reporting
 
-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.
+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'
+```
 
 ### Live configuration reloading
 
-The config can be reloaded by sending a `kill -s SIGHUP` to the process. Not all settings are currently supported by live reload:
+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:
 
 | **Config**              | **Requires restart** |
 |-------------------------|----------------------|
@@ -239,7 +251,6 @@ The config can be reloaded by sending a `kill -s SIGHUP` to the process. Not all
 | `connect_timeout`       | yes                  |
 | `healthcheck_timeout`   | no                   |
 | `ban_time`              | no                   |
-| `statsd_address`        | yes                  |
 | `user`                  | yes                  |
 | `shards`                | yes                  |
 | `default_role`          | no                   |

docker-compose.yml

@@ -0,0 +1,16 @@
+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"

examples/docker/pgcat.toml

@@ -0,0 +1,105 @@
+#
+# 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"

pgcat.toml

@@ -29,9 +29,6 @@ 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]

src/admin.rs

@@ -0,0 +1,388 @@
+/// Admin database.
+use bytes::{Buf, BufMut, BytesMut};
+use log::{info, trace};
+use std::collections::HashMap;
+use tokio::net::tcp::OwnedWriteHalf;
+
+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.
+pub async fn handle_admin(
+    stream: &mut OwnedWriteHalf,
+    mut query: BytesMut,
+    pool: ConnectionPool,
+) -> Result<(), Error> {
+    let code = query.get_u8() as char;
+
+    if code != 'Q' {
+        return Err(Error::ProtocolSyncError);
+    }
+
+    let len = query.get_i32() as usize;
+    let query = String::from_utf8_lossy(&query[..len - 5])
+        .to_string()
+        .to_ascii_uppercase();
+
+    trace!("Admin query: {}", query);
+
+    if query.starts_with("SHOW STATS") {
+        trace!("SHOW STATS");
+        show_stats(stream, &pool).await
+    } else if query.starts_with("RELOAD") {
+        trace!("RELOAD");
+        reload(stream).await
+    } else if query.starts_with("SHOW CONFIG") {
+        trace!("SHOW CONFIG");
+        show_config(stream).await
+    } else if query.starts_with("SHOW DATABASES") {
+        trace!("SHOW DATABASES");
+        show_databases(stream, &pool).await
+    } else if query.starts_with("SHOW POOLS") {
+        trace!("SHOW POOLS");
+        show_pools(stream, &pool).await
+    } else if query.starts_with("SHOW LISTS") {
+        trace!("SHOW LISTS");
+        show_lists(stream, &pool).await
+    } else if query.starts_with("SHOW VERSION") {
+        trace!("SHOW VERSION");
+        show_version(stream).await
+    } else if query.starts_with("SET ") {
+        trace!("SET");
+        ignore_set(stream).await
+    } else {
+        error_response(stream, "Unsupported query against the admin database").await
+    }
+}
+
+/// Column-oriented statistics.
+async fn show_lists(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
+    let stats = get_stats();
+
+    let columns = vec![("list", DataType::Text), ("items", DataType::Int4)];
+
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+    res.put(data_row(&vec![
+        "databases".to_string(),
+        (pool.databases() + 1).to_string(), // see comment below
+    ]));
+    res.put(data_row(&vec!["users".to_string(), "1".to_string()]));
+    res.put(data_row(&vec![
+        "pools".to_string(),
+        (pool.databases() + 1).to_string(), // +1 for the pgbouncer admin db pool which isn't real
+    ])); // 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(),
+    ]));
+    res.put(data_row(&vec![
+        "used_clients".to_string(),
+        stats
+            .keys()
+            .map(|address_id| stats[&address_id]["cl_active"])
+            .sum::<i64>()
+            .to_string(),
+    ]));
+    res.put(data_row(&vec![
+        "login_clients".to_string(),
+        "0".to_string(),
+    ]));
+    res.put(data_row(&vec![
+        "free_servers".to_string(),
+        stats
+            .keys()
+            .map(|address_id| stats[&address_id]["sv_idle"])
+            .sum::<i64>()
+            .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(),
+    ]));
+    res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
+    res.put(data_row(&vec!["dns_zones".to_string(), "0".to_string()]));
+    res.put(data_row(&vec!["dns_queries".to_string(), "0".to_string()]));
+    res.put(data_row(&vec!["dns_pending".to_string(), "0".to_string()]));
+
+    res.put(command_complete("SHOW"));
+
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, res).await
+}
+
+/// Show PgCat version.
+async fn show_version(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
+    let mut res = BytesMut::new();
+
+    res.put(row_description(&vec![("version", DataType::Text)]));
+    res.put(data_row(&vec!["PgCat 0.1.0".to_string()]));
+    res.put(command_complete("SHOW"));
+
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    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> {
+    let stats = get_stats();
+    let config = {
+        let guard = get_config();
+        &*guard.clone()
+    };
+
+    let columns = vec![
+        ("database", DataType::Text),
+        ("user", DataType::Text),
+        ("cl_active", DataType::Numeric),
+        ("cl_waiting", DataType::Numeric),
+        ("cl_cancel_req", DataType::Numeric),
+        ("sv_active", DataType::Numeric),
+        ("sv_idle", DataType::Numeric),
+        ("sv_used", DataType::Numeric),
+        ("sv_tested", DataType::Numeric),
+        ("sv_login", DataType::Numeric),
+        ("maxwait", DataType::Numeric),
+        ("maxwait_us", DataType::Numeric),
+        ("pool_mode", DataType::Text),
+    ];
+
+    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![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));
+        }
+    }
+
+    res.put(command_complete("SHOW"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, res).await
+}
+
+/// Show shards and replicas.
+async fn show_databases(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
+    let guard = get_config();
+    let config = &*guard.clone();
+    drop(guard);
+
+    // Columns
+    let columns = vec![
+        ("name", DataType::Text),
+        ("host", DataType::Text),
+        ("port", DataType::Text),
+        ("database", DataType::Text),
+        ("force_user", DataType::Text),
+        ("pool_size", DataType::Int4),
+        ("min_pool_size", DataType::Int4),
+        ("reserve_pool", DataType::Int4),
+        ("pool_mode", DataType::Text),
+        ("max_connections", DataType::Int4),
+        ("current_connections", DataType::Int4),
+        ("paused", DataType::Int4),
+        ("disabled", DataType::Int4),
+    ];
+
+    let mut res = BytesMut::new();
+
+    res.put(row_description(&columns));
+
+    for shard in 0..pool.shards() {
+        let database_name = &config.shards[&shard.to_string()].database;
+
+        for server in 0..pool.servers(shard) {
+            let address = pool.address(shard, server);
+            let pool_state = pool.pool_state(shard, server);
+
+            res.put(data_row(&vec![
+                address.name(),                       // name
+                address.host.to_string(),             // host
+                address.port.to_string(),             // port
+                database_name.to_string(),            // database
+                config.user.name.to_string(),         // force_user
+                config.general.pool_size.to_string(), // pool_size
+                "0".to_string(),                      // min_pool_size
+                "0".to_string(),                      // reserve_pool
+                config.general.pool_mode.to_string(), // pool_mode
+                config.general.pool_size.to_string(), // max_connections
+                pool_state.connections.to_string(),   // current_connections
+                "0".to_string(),                      // paused
+                "0".to_string(),                      // disabled
+            ]));
+        }
+    }
+
+    res.put(command_complete("SHOW"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, res).await
+}
+
+/// Ignore any SET commands the client sends.
+/// This is common initialization done by ORMs.
+async fn ignore_set(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
+    custom_protocol_response_ok(stream, "SET").await
+}
+
+/// Reload the configuration file without restarting the process.
+async fn reload(stream: &mut OwnedWriteHalf) -> Result<(), Error> {
+    info!("Reloading config");
+
+    let config = get_config();
+    let path = config.path.clone().unwrap();
+
+    parse(&path).await?;
+
+    let config = get_config();
+
+    config.show();
+
+    let mut res = BytesMut::new();
+
+    res.put(command_complete("RELOAD"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    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.
+    let immutables = ["host", "port", "connect_timeout"];
+
+    // Columns
+    let columns = vec![
+        ("key", DataType::Text),
+        ("value", DataType::Text),
+        ("default", DataType::Text),
+        ("changeable", DataType::Text),
+    ];
+
+    // Response data
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    // DataRow rows
+    for (key, value) in config {
+        let changeable = if immutables.iter().filter(|col| *col == &key).count() == 1 {
+            "no".to_string()
+        } else {
+            "yes".to_string()
+        };
+
+        let row = vec![key, value, "-".to_string(), changeable];
+
+        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');
+
+    write_all_half(stream, res).await
+}
+
+/// Show shard and replicas statistics.
+async fn show_stats(stream: &mut OwnedWriteHalf, pool: &ConnectionPool) -> Result<(), Error> {
+    let columns = vec![
+        ("database", DataType::Text),
+        ("total_xact_count", DataType::Numeric),
+        ("total_query_count", DataType::Numeric),
+        ("total_received", DataType::Numeric),
+        ("total_sent", DataType::Numeric),
+        ("total_xact_time", DataType::Numeric),
+        ("total_query_time", DataType::Numeric),
+        ("total_wait_time", DataType::Numeric),
+        ("avg_xact_count", DataType::Numeric),
+        ("avg_query_count", DataType::Numeric),
+        ("avg_recv", DataType::Numeric),
+        ("avg_sent", DataType::Numeric),
+        ("avg_xact_time", DataType::Numeric),
+        ("avg_query_time", DataType::Numeric),
+        ("avg_wait_time", DataType::Numeric),
+    ];
+
+    let stats = get_stats();
+    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![address.name()];
+
+            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');
+
+    write_all_half(stream, res).await
+}

src/client.rs

@@ -1,16 +1,14 @@
-/// Implementation of the PostgreSQL client.
+/// Handle clients by pretending to be a PostgreSQL server.
-/// 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::*;
 use crate::errors::Error;
@@ -22,64 +20,70 @@ 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
+    /// We buffer the writes ourselves because we know the protocol
-    // better than a stock buffer.
+    /// better than a stock buffer.
     write: OwnedWriteHalf,
 
-    // Internal buffer, where we place messages until we have to flush
+    /// Internal buffer, where we place messages until we have to flush
-    // them to the backend.
+    /// 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.
+    /// In transaction mode, the connection is released after each transaction.
-    // Session mode has slightly higher throughput per client, but lower capacity.
+    /// 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
+    /// Clients are mapped to servers while they use them. This allows a client
-    // to connect and cancel a query.
+    /// 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.
+    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 {
-    /// Given a TCP socket, trick the client into thinking we are
+    /// Perform client startup sequence.
-    /// the Postgres server. Perform the authentication and place
+    /// See docs: <https://www.postgresql.org/docs/12/protocol-flow.html#id-1.10.5.7.3>
-    /// 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();
+        let config = get_config().clone();
         let transaction_mode = config.general.pool_mode.starts_with("t");
-        drop(config);
+        // drop(config);
         loop {
             trace!("Waiting for StartupMessage");
 
-            // Could be StartupMessage or SSLRequest
+            // Could be StartupMessage, SSLRequest or CancelRequest.
-            // 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 {
@@ -104,20 +108,66 @@ impl Client {
                 // Regular startup message.
                 PROTOCOL_VERSION_NUMBER => {
                     trace!("Got StartupMessage");
-
-                    // TODO: perform actual auth.
                     let parameters = parse_startup(bytes.clone())?;
 
                     // 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?;
                     trace!("Startup OK");
 
+                    let database = parameters
+                        .get("database")
+                        .unwrap_or(parameters.get("user").unwrap());
+                    let admin = ["pgcat", "pgbouncer"]
+                        .iter()
+                        .filter(|db| *db == &database)
+                        .count()
+                        == 1;
+
                     // Split the read and write streams
                     // so we can control buffering.
                     let (read, write) = stream.into_split();
@@ -133,6 +183,9 @@ impl Client {
                         client_server_map: client_server_map,
                         parameters: parameters,
                         stats: stats,
+                        admin: admin,
+                        last_address_id: None,
+                        last_server_id: None,
                     });
                 }
 
@@ -154,6 +207,9 @@ impl Client {
                         client_server_map: client_server_map,
                         parameters: HashMap::new(),
                         stats: stats,
+                        admin: false,
+                        last_address_id: None,
+                        last_server_id: None,
                     });
                 }
 
@@ -164,7 +220,7 @@ impl Client {
         }
     }
 
-    /// Client loop. We handle all messages between the client and the database here.
+    /// Handle a connected and authenticated client.
     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 {
@@ -200,17 +256,14 @@ 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 protocols.
+        // or issue commands for our sharding and server selection protocol.
         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 control messages, e.g.
+            // in case the client is sending some custom protocol messages, e.g.
             // SET SHARDING KEY TO 'bigint';
             let mut message = read_message(&mut self.read).await?;
 
@@ -220,36 +273,48 @@ impl Client {
                 return Ok(());
             }
 
-            // Handle all custom protocol commands here.
+            // Handle admin database queries.
+            if self.admin {
+                trace!("Handling admin command");
+                handle_admin(&mut self.write, message, pool.clone()).await?;
+                continue;
+            }
+
+            // Handle all custom protocol commands, if any.
             match query_router.try_execute_command(message.clone()) {
-                // Normal query
+                // Normal query, not a custom command.
                 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, &format!("SET SHARD")).await?;
+                    custom_protocol_response_ok(&mut self.write, "SET SHARD").await?;
                     continue;
                 }
 
+                // SET SHARDING KEY TO
                 Some((Command::SetShardingKey, _)) => {
-                    custom_protocol_response_ok(&mut self.write, &format!("SET SHARDING KEY"))
+                    custom_protocol_response_ok(&mut self.write, "SET SHARDING KEY").await?;
-                        .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;
@@ -261,7 +326,7 @@ impl Client {
                 error_response(
                     &mut self.write,
                     &format!(
-                        "shard '{}' is more than configured '{}'",
+                        "shard {} is more than configured {}",
                         query_router.shard(),
                         pool.shards()
                     ),
@@ -270,13 +335,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: the client issued a regular query.
+            // Grab a server from the pool.
-            let connection = match pool.get(query_router.shard(), query_router.role()).await {
+            let connection = match pool
+                .get(query_router.shard(), query_router.role(), self.process_id)
+                .await
+            {
                 Ok(conn) => {
                     debug!("Got connection from pool");
                     conn
@@ -290,15 +355,23 @@ impl Client {
             };
 
             let mut reference = connection.0;
-            let _address = connection.1;
+            let address = connection.1;
             let server = &mut *reference;
 
-            // Claim this server as mine for query cancellation.
+            // Server is assigned to the client in case the client wants to
+            // cancel a query later.
             server.claim(self.process_id, self.secret_key);
 
-            // Client active & server active
+            // Update statistics.
-            self.stats.client_active(self.process_id);
+            if let Some(last_address_id) = self.last_address_id {
-            self.stats.server_active(server.process_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());
 
             debug!(
                 "Client {:?} talking to server {:?}",
@@ -306,9 +379,20 @@ 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");
@@ -316,11 +400,13 @@ impl Client {
                     match read_message(&mut self.read).await {
                         Ok(message) => message,
                         Err(err) => {
-                            // Client disconnected without warning.
+                            // Client disconnected inside a transaction.
+                            // Clean up the server and re-use it.
+                            // This prevents connection thrashing by bad clients.
                             if server.in_transaction() {
-                                // Client left dirty server. Clean up and proceed
+                                server.query("ROLLBACK").await?;
-                                // without thrashing this connection.
+                                server.query("DISCARD ALL").await?;
-                                server.query("ROLLBACK; DISCARD ALL;").await?;
+                                server.set_name("pgcat").await?;
                             }
 
                             return Err(err);
@@ -346,13 +432,11 @@ 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.
@@ -370,17 +454,16 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query();
+                        self.stats.query(self.process_id, address.id);
 
-                        // 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.stats.transaction(self.process_id, address.id);
 
                             // 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());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -392,11 +475,14 @@ 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; DISCARD ALL;").await?;
+                            server.query("ROLLBACK").await?;
+                            server.query("DISCARD ALL").await?;
+                            server.set_name("pgcat").await?;
                         }
 
+                        self.release();
+
                         return Ok(());
                     }
 
@@ -431,7 +517,6 @@ impl Client {
 
                         self.buffer.put(&original[..]);
 
-                        // TODO: retries for read-only transactions.
                         server.send(self.buffer.clone()).await?;
 
                         self.buffer.clear();
@@ -439,7 +524,6 @@ 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 {
@@ -456,15 +540,15 @@ impl Client {
                         }
 
                         // Report query executed statistics.
-                        self.stats.query();
+                        self.stats.query(self.process_id, address.id);
 
-                        // 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.stats.transaction(self.process_id, address.id);
 
+                            // 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());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -492,13 +576,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.stats.transaction(self.process_id, address.id);
 
+                            // 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());
+                                self.stats.server_idle(server.process_id(), address.id);
                                 break;
                             }
                         }
@@ -515,10 +599,11 @@ 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 being mine. I can't cancel its queries anymore.
+    /// Release the server from the client: it 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));
@@ -527,6 +612,13 @@ impl Client {
 
 impl Drop for Client {
     fn drop(&mut self) {
-        self.stats.client_disconnecting(self.process_id);
+        // 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);
+            }
+        }
     }
 }

src/config.rs

@@ -1,24 +1,35 @@
+/// 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,
     Replica,
 }
 
+impl ToString for Role {
+    fn to_string(&self) -> String {
+        match *self {
+            Role::Primary => "primary".to_string(),
+            Role::Replica => "replica".to_string(),
+        }
+    }
+}
+
 impl PartialEq<Option<Role>> for Role {
     fn eq(&self, other: &Option<Role>) -> bool {
         match other {
@@ -37,25 +48,42 @@ 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,
     pub role: Role,
+    pub replica_number: usize,
 }
 
 impl Default for Address {
     fn default() -> Address {
         Address {
+            id: 0,
             host: String::from("127.0.0.1"),
             port: String::from("5432"),
             shard: 0,
+            replica_number: 0,
             role: Role::Replica,
         }
     }
 }
 
+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),
+
+            Role::Replica => format!("shard_{}_replica_{}", self.shard, self.replica_number),
+        }
+    }
+}
+
+/// PostgreSQL user.
 #[derive(Clone, PartialEq, Hash, std::cmp::Eq, Deserialize, Debug)]
 pub struct User {
     pub name: String,
@@ -71,6 +99,7 @@ impl Default for User {
     }
 }
 
+/// General configuration.
 #[derive(Deserialize, Debug, Clone)]
 pub struct General {
     pub host: String,
@@ -80,7 +109,6 @@ pub struct General {
     pub connect_timeout: u64,
     pub healthcheck_timeout: u64,
     pub ban_time: i64,
-    pub statsd_address: String,
 }
 
 impl Default for General {
@@ -93,11 +121,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)>,
@@ -113,6 +141,7 @@ impl Default for Shard {
     }
 }
 
+/// Query Router configuration.
 #[derive(Deserialize, Debug, Clone)]
 pub struct QueryRouter {
     pub default_role: String,
@@ -132,8 +161,10 @@ impl Default for QueryRouter {
     }
 }
 
+/// Configuration wrapper.
 #[derive(Deserialize, Debug, Clone)]
 pub struct Config {
+    pub path: Option<String>,
     pub general: General,
     pub user: User,
     pub shards: HashMap<String, Shard>,
@@ -143,6 +174,7 @@ pub struct Config {
 impl Default for Config {
     fn default() -> Config {
         Config {
+            path: Some(String::from("pgcat.toml")),
             general: General::default(),
             user: User::default(),
             shards: HashMap::from([(String::from("1"), Shard::default())]),
@@ -151,7 +183,50 @@ impl Default for Config {
     }
 }
 
+impl From<&Config> for std::collections::HashMap<String, String> {
+    fn from(config: &Config) -> HashMap<String, String> {
+        HashMap::from([
+            ("host".to_string(), config.general.host.to_string()),
+            ("port".to_string(), config.general.port.to_string()),
+            (
+                "pool_size".to_string(),
+                config.general.pool_size.to_string(),
+            ),
+            (
+                "pool_mode".to_string(),
+                config.general.pool_mode.to_string(),
+            ),
+            (
+                "connect_timeout".to_string(),
+                config.general.connect_timeout.to_string(),
+            ),
+            (
+                "healthcheck_timeout".to_string(),
+                config.general.healthcheck_timeout.to_string(),
+            ),
+            ("ban_time".to_string(), config.general.ban_time.to_string()),
+            (
+                "default_role".to_string(),
+                config.query_router.default_role.to_string(),
+            ),
+            (
+                "query_parser_enabled".to_string(),
+                config.query_router.query_parser_enabled.to_string(),
+            ),
+            (
+                "primary_reads_enabled".to_string(),
+                config.query_router.primary_reads_enabled.to_string(),
+            ),
+            (
+                "sharding_function".to_string(),
+                config.query_router.sharding_function.to_string(),
+            ),
+        ])
+    }
+}
+
 impl Config {
+    /// Print current configuration.
     pub fn show(&self) {
         info!("Pool size: {}", self.general.pool_size);
         info!("Pool mode: {}", self.general.pool_mode);
@@ -166,11 +241,14 @@ 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 config.
+/// Parse the configuration file located at the path.
 pub async fn parse(path: &str) -> Result<(), Error> {
     let mut contents = String::new();
     let mut file = match File::open(path).await {
@@ -189,7 +267,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
-    let config: Config = match toml::from_str(&contents) {
+    let mut config: Config = match toml::from_str(&contents) {
         Ok(config) => config,
         Err(err) => {
             error!("Could not parse config file: {}", err.to_string());
@@ -279,6 +357,9 @@ 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(())
@@ -296,5 +377,6 @@ mod test {
         assert_eq!(get_config().shards["1"].servers[0].0, "127.0.0.1");
         assert_eq!(get_config().shards["0"].servers[0].2, "primary");
         assert_eq!(get_config().query_router.default_role, "any");
+        assert_eq!(get_config().path, Some("pgcat.toml".to_string()));
     }
 }

src/constants.rs

@@ -1,7 +1,6 @@
 /// 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;
@@ -15,8 +14,20 @@ 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;
 
 // ErrorResponse: A code identifying the field type; if zero, this is the message terminator and no string follows.
 pub const MESSAGE_TERMINATOR: u8 = 0;
+
+//
+// Data types
+//
+pub const _OID_INT8: i32 = 20; // bigint

src/errors.rs

@@ -1,12 +1,13 @@
+/// Errors.
+
+/// Various errors.
 #[derive(Debug, PartialEq)]
 pub enum Error {
     SocketError,
-    // ClientDisconnected,
     ClientBadStartup,
     ProtocolSyncError,
     ServerError,
-    // ServerTimeout,
-    // DirtyServer,
     BadConfig,
     AllServersDown,
+    ClientError,
 }

src/main.rs

@@ -31,7 +31,6 @@ extern crate once_cell;
 extern crate serde;
 extern crate serde_derive;
 extern crate sqlparser;
-extern crate statsd;
 extern crate tokio;
 extern crate toml;
 
@@ -47,6 +46,7 @@ use tokio::{
 use std::collections::HashMap;
 use std::sync::Arc;
 
+mod admin;
 mod client;
 mod config;
 mod constants;
@@ -54,23 +54,20 @@ 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;
@@ -84,7 +81,6 @@ async fn main() {
         String::from("pgcat.toml")
     };
 
-    // Prepare the config
     match config::parse(&config_file).await {
         Ok(_) => (),
         Err(err) => {
@@ -94,8 +90,8 @@ async fn main() {
     };
 
     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) => {
@@ -105,22 +101,28 @@ 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()));
 
-    // Collect statistics and send them to StatsD
+    // Statistics reporting.
     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) => {
@@ -131,7 +133,7 @@ async fn main() {
 
     info!("Waiting for clients");
 
-    // Main app runs here.
+    // Client connection loop.
     tokio::task::spawn(async move {
         loop {
             let pool = pool.clone();
@@ -147,7 +149,7 @@ async fn main() {
                 }
             };
 
-            // Client goes to another thread, bye.
+            // Handle client.
             tokio::task::spawn(async move {
                 let start = chrono::offset::Utc::now().naive_utc();
                 match client::Client::startup(socket, client_server_map, server_info, reporter)
@@ -181,7 +183,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();
@@ -201,16 +203,15 @@ async fn main() {
         }
     });
 
-    // Setup shut down sequence
+    // Exit on Ctrl-C (SIGINT) and SIGTERM.
-    match signal::ctrl_c().await {
+    let mut term_signal = unix_signal(SignalKind::terminate()).unwrap();
-        Ok(()) => {
-            info!("Shutting down...");
-        }
 
-        Err(err) => {
+    tokio::select! {
-            error!("Unable to listen for shutdown signal: {}", err);
+        _ = signal::ctrl_c() => (),
-        }
+        _ = term_signal.recv() => (),
     };
+
+    info!("Shutting down...");
 }
 
 /// Format chrono::Duration to be more human-friendly.

src/messages.rs

@@ -8,9 +8,26 @@ use tokio::net::{
     TcpStream,
 };
 
+use crate::errors::Error;
 use std::collections::HashMap;
 
-use crate::errors::Error;
+/// Postgres data type mappings
+/// used in RowDescription ('T') message.
+pub enum DataType {
+    Text,
+    Int4,
+    Numeric,
+}
+
+impl From<&DataType> for i32 {
+    fn from(data_type: &DataType) -> i32 {
+        match data_type {
+            DataType::Text => 25,
+            DataType::Int4 => 23,
+            DataType::Numeric => 1700,
+        }
+    }
+}
 
 /// Tell the client that authentication handshake completed successfully.
 pub async fn auth_ok(stream: &mut TcpStream) -> Result<(), Error> {
@@ -23,6 +40,26 @@ 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(
@@ -91,9 +128,8 @@ pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Resu
     }
 }
 
-/// Parse StartupMessage parameters.
+/// Parse the params the server sends as a key/value format.
-/// e.g. user, database, application_name, etc.
+pub fn parse_params(mut bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
-pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
     let mut result = HashMap::new();
     let mut buf = Vec::new();
     let mut tmp = String::new();
@@ -115,7 +151,7 @@ pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Err
 
     // Expect pairs of name and value
     // and at least one pair to be present.
-    if buf.len() % 2 != 0 && buf.len() >= 2 {
+    if buf.len() % 2 != 0 || buf.len() < 2 {
         return Err(Error::ClientBadStartup);
     }
 
@@ -127,6 +163,14 @@ pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Err
         i += 2;
     }
 
+    Ok(result)
+}
+
+/// Parse StartupMessage parameters.
+/// e.g. user, database, application_name, etc.
+pub fn parse_startup(bytes: BytesMut) -> Result<HashMap<String, String>, Error> {
+    let result = parse_params(bytes)?;
+
     // Minimum required parameters
     // I want to have the user at the very minimum, according to the protocol spec.
     if !result.contains_key("user") {
@@ -136,14 +180,8 @@ pub fn parse_startup(mut bytes: BytesMut) -> Result<HashMap<String, String>, Err
     Ok(result)
 }
 
-/// Send password challenge response to the server.
+/// Create md5 password hash given a salt.
-/// This is the MD5 challenge.
+pub fn md5_hash_password(user: &str, password: &str, salt: &[u8]) -> Vec<u8> {
-pub async fn md5_password(
-    stream: &mut TcpStream,
-    user: &str,
-    password: &str,
-    salt: &[u8],
-) -> Result<(), Error> {
     let mut md5 = Md5::new();
 
     // First pass
@@ -162,6 +200,19 @@ pub async fn md5_password(
         .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');
@@ -198,7 +249,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();
 
@@ -240,6 +291,39 @@ 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,
@@ -252,68 +336,17 @@ pub async fn show_response(
     // 3. CommandComplete
     // 4. ReadyForQuery
 
-    // RowDescription
-    let mut row_desc = BytesMut::new();
-
-    // Number of columns: 1
-    row_desc.put_i16(1);
-
-    // Column name
-    row_desc.put_slice(&format!("{}\0", name).as_bytes());
-
-    // Doesn't belong to any table
-    row_desc.put_i32(0);
-
-    // Doesn't belong to any table
-    row_desc.put_i16(0);
-
-    // Text
-    row_desc.put_i32(25);
-
-    // Text size = variable (-1)
-    row_desc.put_i16(-1);
-
-    // Type modifier: none that I know
-    row_desc.put_i32(0);
-
-    // Format being used: text (0), binary (1)
-    row_desc.put_i16(0);
-
-    // DataRow
-    let mut data_row = BytesMut::new();
-
-    // Number of columns
-    data_row.put_i16(1);
-
-    // Size of the column content (length of the string really)
-    data_row.put_i32(value.len() as i32);
-
-    // The content
-    data_row.put_slice(value.as_bytes());
-
-    // CommandComplete
-    let mut command_complete = BytesMut::new();
-
-    // Number of rows returned (just one)
-    command_complete.put_slice(&b"SELECT 1\0"[..]);
-
     // The final messages sent to the client
     let mut res = BytesMut::new();
 
     // RowDescription
-    res.put_u8(b'T');
+    res.put(row_description(&vec![(name, DataType::Text)]));
-    res.put_i32(row_desc.len() as i32 + 4);
-    res.put(row_desc);
 
     // DataRow
-    res.put_u8(b'D');
+    res.put(data_row(&vec![value.to_string()]));
-    res.put_i32(data_row.len() as i32 + 4);
-    res.put(data_row);
 
     // CommandComplete
-    res.put_u8(b'C');
+    res.put(command_complete("SELECT 1"));
-    res.put_i32(command_complete.len() as i32 + 4);
-    res.put(command_complete);
 
     // ReadyForQuery
     res.put_u8(b'Z');
@@ -323,6 +356,79 @@ pub async fn show_response(
     write_all_half(stream, res).await
 }
 
+pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
+    let mut res = BytesMut::new();
+    let mut row_desc = BytesMut::new();
+
+    // how many colums we are storing
+    row_desc.put_i16(columns.len() as i16);
+
+    for (name, data_type) in columns {
+        // Column name
+        row_desc.put_slice(&format!("{}\0", name).as_bytes());
+
+        // Doesn't belong to any table
+        row_desc.put_i32(0);
+
+        // Doesn't belong to any table
+        row_desc.put_i16(0);
+
+        // Text
+        row_desc.put_i32(data_type.into());
+
+        // Text size = variable (-1)
+        let type_size = match data_type {
+            DataType::Text => -1,
+            DataType::Int4 => 4,
+            DataType::Numeric => -1,
+        };
+
+        row_desc.put_i16(type_size);
+
+        // Type modifier: none that I know
+        row_desc.put_i32(-1);
+
+        // Format being used: text (0), binary (1)
+        row_desc.put_i16(0);
+    }
+
+    res.put_u8(b'T');
+    res.put_i32(row_desc.len() as i32 + 4);
+    res.put(row_desc);
+
+    res
+}
+
+/// Create a DataRow message.
+pub fn data_row(row: &Vec<String>) -> BytesMut {
+    let mut res = BytesMut::new();
+    let mut data_row = BytesMut::new();
+
+    data_row.put_i16(row.len() as i16);
+
+    for column in row {
+        let column = column.as_bytes();
+        data_row.put_i32(column.len() as i32);
+        data_row.put_slice(&column);
+    }
+
+    res.put_u8(b'D');
+    res.put_i32(data_row.len() as i32 + 4);
+    res.put(data_row);
+
+    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();
+    res.put_u8(b'C');
+    res.put_i32(cmd.len() as i32 + 4);
+    res.put(cmd);
+    res
+}
+
 /// Write all data in the buffer to the TcpStream.
 pub async fn write_all(stream: &mut TcpStream, buf: BytesMut) -> Result<(), Error> {
     match stream.write_all(&buf).await {

src/pool.rs

@@ -1,24 +1,23 @@
-/// Pooling and failover and banlist.
+/// Pooling, 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>>>,
@@ -29,7 +28,7 @@ pub struct ConnectionPool {
 }
 
 impl ConnectionPool {
-    /// Construct the connection pool from a config file.
+    /// Construct the connection pool from the configuration.
     pub async fn from_config(
         client_server_map: ClientServerMap,
         stats: Reporter,
@@ -38,6 +37,7 @@ 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()
@@ -49,7 +49,8 @@ impl ConnectionPool {
         for shard_idx in shard_ids {
             let shard = &config.shards[&shard_idx];
             let mut pools = Vec::new();
-            let mut replica_addresses = Vec::new();
+            let mut servers = Vec::new();
+            let mut replica_number = 0;
 
             for server in shard.servers.iter() {
                 let role = match server.2.as_ref() {
@@ -62,12 +63,20 @@ impl ConnectionPool {
                 };
 
                 let address = Address {
+                    id: address_id,
                     host: server.0.clone(),
                     port: server.1.to_string(),
                     role: role,
+                    replica_number,
                     shard: shard_idx.parse::<usize>().unwrap(),
                 };
 
+                address_id += 1;
+
+                if role == Role::Replica {
+                    replica_number += 1;
+                }
+
                 let manager = ServerPool::new(
                     address.clone(),
                     config.user.clone(),
@@ -87,11 +96,11 @@ impl ConnectionPool {
                     .unwrap();
 
                 pools.push(pool);
-                replica_addresses.push(address);
+                servers.push(address);
             }
 
             shards.push(pools);
-            addresses.push(replica_addresses);
+            addresses.push(servers);
             banlist.push(HashMap::new());
         }
 
@@ -115,9 +124,13 @@ 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) {
-                let connection = match self.get(shard, None).await {
+                // To keep stats consistent.
+                let fake_process_id = 0;
+
+                let connection = match self.get(shard, None, fake_process_id).await {
                     Ok(conn) => conn,
                     Err(err) => {
                         error!("Shard {} down or misconfigured: {:?}", shard, err);
@@ -126,10 +139,24 @@ impl ConnectionPool {
                 };
 
                 let mut proxy = connection.0;
-                let _address = connection.1;
+                let address = connection.1;
                 let server = &mut *proxy;
 
-                server_infos.push(server.server_info());
+                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] {
+                        warn!(
+                            "{:?} has different server configuration than the last server",
+                            address
+                        );
+                    }
+                }
+
+                server_infos.push(server_info);
             }
         }
 
@@ -147,6 +174,7 @@ 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];
@@ -175,10 +203,9 @@ impl ConnectionPool {
         }
 
         while allowed_attempts > 0 {
-            // Round-robin each client's queries.
+            // Round-robin replicas.
-            // 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];
 
@@ -195,21 +222,27 @@ 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());
+            self.stats.server_tested(server.process_id(), address.id);
 
             match tokio::time::timeout(
                 tokio::time::Duration::from_millis(healthcheck_timeout),
@@ -217,29 +250,40 @@ 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());
+                        self.stats
-                        self.stats.server_idle(conn.process_id());
+                            .checkout_time(now.elapsed().as_micros(), process_id, address.id);
+                        self.stats.server_idle(conn.process_id(), address.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 never came back, database is really really down
+
+                // Health check timed out.
                 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;
                 }
             }
@@ -317,15 +361,38 @@ 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() {
+            databases += self.servers(shard);
+        }
+        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,
@@ -361,14 +428,16 @@ impl ManageConnection for ServerPool {
     async fn connect(&self) -> Result<Self::Connection, Self::Error> {
         info!(
             "Creating a new connection to {:?} using user {:?}",
-            self.address, self.user.name
+            self.address.name(),
+            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.stats.server_login(process_id, self.address.id);
 
+        // Connect to the PostgreSQL server.
         match Server::startup(
             &self.address,
             &self.user,
@@ -380,12 +449,12 @@ impl ManageConnection for ServerPool {
         {
             Ok(conn) => {
                 // Remove the temporary process_id from the stats.
-                self.stats.server_disconnecting(process_id);
+                self.stats.server_disconnecting(process_id, self.address.id);
                 Ok(conn)
             }
             Err(err) => {
                 // Remove the temporary process_id from the stats.
-                self.stats.server_disconnecting(process_id);
+                self.stats.server_disconnecting(process_id, self.address.id);
                 Err(err)
             }
         }

src/query_router.rs

@@ -1,5 +1,3 @@
-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};
@@ -10,14 +8,19 @@ 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]+)'? *;? *$",
+    r"(?i)^ *SET SHARD TO '?([0-9]+|ANY)'? *;? *$",
     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,
@@ -27,37 +30,39 @@ pub enum Command {
     ShowServerRole,
 }
 
-// Quick test
+/// Quickly test for match when a query is received.
 static CUSTOM_SQL_REGEX_SET: OnceCell<RegexSet> = OnceCell::new();
 
-// Capture value
+// Get the value inside the custom command.
 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
+    /// By default, queries go here, unless we have better information
-    // about what the client wants.
+    /// 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>,
 
-    // Should we be talking to a primary or a replica?
+    /// Which server should we be talking to.
     active_role: Option<Role>,
 
-    // Include the primary into the replica pool?
+    /// Include the primary into the replica pool for reads.
     primary_reads_enabled: bool,
 
-    // Should we try to parse queries?
+    /// Should we try to parse queries to route them to replicas or primary automatically.
     query_parser_enabled: bool,
 
-    // Which sharding function are we using?
+    /// Which sharding function we're 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,
@@ -88,6 +93,7 @@ impl QueryRouter {
         }
     }
 
+    /// Create a new instance of the query router. Each client gets its own.
     pub fn new() -> QueryRouter {
         let config = get_config();
 
@@ -120,6 +126,7 @@ 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;
         }
@@ -158,8 +165,7 @@ 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, so
+                // I think this is faster than running the Regex engine 5 times.
-                // 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(),
@@ -192,7 +198,10 @@ impl QueryRouter {
             }
 
             Command::SetShard => {
-                self.active_shard = Some(value.parse::<usize>().unwrap());
+                self.active_shard = match value.to_ascii_uppercase().as_ref() {
+                    "ANY" => Some(rand::random::<usize>() % self.shards),
+                    _ => Some(value.parse::<usize>().unwrap()),
+                };
             }
 
             Command::SetServerRole => {
@@ -218,7 +227,6 @@ 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
@@ -240,12 +248,14 @@ 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;
@@ -268,6 +278,7 @@ impl QueryRouter {
 
                 query.replace("$", "") // Remove placeholders turning them into "values"
             }
+
             _ => return false,
         };
 
@@ -331,6 +342,7 @@ 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;

src/scram.rs

@@ -0,0 +1,320 @@
+// 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();
+    }
+}

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,48 +12,52 @@ 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,
+    /// Server host, e.g. localhost,
-    // port, e.g. 5432, and role, e.g. primary or replica.
+    /// 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 {
@@ -77,7 +81,7 @@ impl Server {
 
         trace!("Sending StartupMessage");
 
-        // Send the startup packet telling the server we're a normal Postgres client.
+        // StartupMessage
         startup(&mut stream, &user.name, database).await?;
 
         let mut server_info = BytesMut::new();
@@ -86,6 +90,8 @@ 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,
@@ -127,6 +133,91 @@ 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);
@@ -187,7 +278,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),
@@ -208,11 +299,9 @@ 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();
 
-                    return Ok(Server {
+                    let mut server = Server {
                         address: address.clone(),
                         read: BufReader::new(read),
                         write: write,
@@ -226,7 +315,12 @@ 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.
@@ -268,7 +362,8 @@ 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.stats
+            .data_sent(messages.len(), self.process_id, self.address.id);
 
         match write_all_half(&mut self.write, messages).await {
             Ok(_) => Ok(()),
@@ -341,8 +436,7 @@ 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...
+                    // Don't flush yet, the more we buffer, the faster this goes...up to a limit.
-                    // up to a limit of course.
                     if self.buffer.len() >= 8196 {
                         break;
                     }
@@ -374,7 +468,8 @@ 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.stats
+            .data_received(bytes.len(), self.process_id, self.address.id);
 
         // Clear the buffer for next query.
         self.buffer.clear();
@@ -409,7 +504,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");
+        error!("Server {:?} marked bad", self.address);
         self.bad = true;
     }
 
@@ -449,9 +544,14 @@ impl Server {
     /// A shorthand for `SET application_name = $1`.
     #[allow(dead_code)]
     pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
-        Ok(self
+        if self.application_name != name {
-            .query(&format!("SET application_name = '{}'", name))
+            self.application_name = name.to_string();
-            .await?)
+            Ok(self
+                .query(&format!("SET application_name = '{}'", name))
+                .await?)
+        } else {
+            Ok(())
+        }
     }
 
     /// Get the servers address.
@@ -460,6 +560,7 @@ impl Server {
         self.address.clone()
     }
 
+    /// Get the server's unique identifier.
     pub fn process_id(&self) -> i32 {
         self.process_id
     }
@@ -470,7 +571,8 @@ 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.stats
+            .server_disconnecting(self.process_id(), self.address.id);
 
         let mut bytes = BytesMut::with_capacity(4);
         bytes.put_u8(b'X');
@@ -478,9 +580,10 @@ impl Drop for Server {
 
         match self.write.try_write(&bytes) {
             Ok(_) => (),
-            Err(_) => (),
+            Err(_) => debug!("Dirty shutdown"),
         };
 
+        // Should not matter.
         self.bad = true;
 
         let now = chrono::offset::Utc::now().naive_utc();

src/sharding.rs

@@ -1,20 +1,27 @@
+/// Implements various sharding functions.
 use sha1::{Digest, Sha1};
 
-// https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/include/catalog/partition.h#L20
+/// See: <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,
@@ -22,6 +29,7 @@ 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),
@@ -31,7 +39,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;
@@ -119,6 +127,7 @@ 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;

src/stats.rs

@@ -1,12 +1,20 @@
+/// Statistics and reporting.
 use log::info;
-use statsd::Client;
+use once_cell::sync::Lazy;
-/// Events collector and publisher.
+use parking_lot::Mutex;
+use std::collections::HashMap;
 use tokio::sync::mpsc::{Receiver, Sender};
 
-use std::collections::HashMap;
+/// 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 crate::config::get_config;
+/// Statistics period used for average calculations.
+/// 15 seconds.
+static STAT_PERIOD: u64 = 15000;
 
+/// The names for the events reported
+/// to the statistics collector.
 #[derive(Debug, Clone, Copy)]
 enum EventName {
     CheckoutTime,
@@ -23,201 +31,266 @@ enum EventName {
     ServerTested,
     ServerLogin,
     ServerDisconnecting,
-    FlushStatsToStatsD,
+    UpdateStats,
+    UpdateAverages,
 }
 
+/// 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,
-    process_id: Option<i32>,
+
+    /// The client or server connection reporting the event.
+    process_id: i32,
+
+    /// The server the client is connected to.
+    address_id: usize,
 }
 
+/// 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 }
     }
 
-    pub fn query(&self) {
+    /// 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) {
         let event = Event {
             name: EventName::Query,
             value: 1,
-            process_id: None,
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn transaction(&self) {
+    /// 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) {
         let event = Event {
             name: EventName::Transaction,
             value: 1,
-            process_id: None,
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn data_sent(&self, amount: usize) {
+    /// 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) {
         let event = Event {
             name: EventName::DataSent,
             value: amount as i64,
-            process_id: None,
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn data_received(&self, amount: usize) {
+    /// 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) {
         let event = Event {
             name: EventName::DataReceived,
             value: amount as i64,
-            process_id: None,
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn checkout_time(&self, ms: u128) {
+    /// 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) {
         let event = Event {
             name: EventName::CheckoutTime,
             value: ms as i64,
-            process_id: None,
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn client_waiting(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ClientWaiting,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn client_active(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ClientActive,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn client_idle(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ClientIdle,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn client_disconnecting(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ClientDisconnecting,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn server_active(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ServerActive,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn server_idle(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ServerIdle,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn server_login(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ServerLogin,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn server_tested(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ServerTested,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
 
-    pub fn server_disconnecting(&self, process_id: i32) {
+    /// 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) {
         let event = Event {
             name: EventName::ServerDisconnecting,
             value: 1,
-            process_id: Some(process_id),
+            process_id: process_id,
+            address_id: address_id,
         };
 
         let _ = self.tx.try_send(event);
     }
-
-    // pub fn flush_to_statsd(&self) {
-    //     let event = Event {
-    //         name: EventName::FlushStatsToStatsD,
-    //         value: 0,
-    //         process_id: None,
-    //     };
-
-    //     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 {
+        Collector { rx, tx }
-            rx,
-            tx,
-            client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
-        }
     }
 
-    pub async fn collect(&mut self) {
+    /// 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) {
         info!("Events reporter started");
 
-        let mut stats = HashMap::from([
+        let stats_template = HashMap::from([
             ("total_query_count", 0),
             ("total_xact_count", 0),
             ("total_sent", 0),
             ("total_received", 0),
+            ("total_xact_time", 0),
+            ("total_query_time", 0),
             ("total_wait_time", 0),
+            ("avg_xact_time", 0),
+            ("avg_query_time", 0),
+            ("avg_xact_count", 0),
+            ("avg_sent", 0),
+            ("avg_received", 0),
+            ("avg_wait_time", 0),
             ("maxwait_us", 0),
             ("maxwait", 0),
             ("cl_waiting", 0),
@@ -229,21 +302,51 @@ impl Collector {
             ("sv_tested", 0),
         ]);
 
-        let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
+        let mut stats = HashMap::new();
-        let tx = self.tx.clone();
 
+        // 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();
+
+        // Track which state the client and server are at any given time.
+        let mut client_server_states: HashMap<usize, 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(15000));
+            let mut interval =
+                tokio::time::interval(tokio::time::Duration::from_millis(STAT_PERIOD / 15));
             loop {
                 interval.tick().await;
-                let _ = tx.try_send(Event {
+                for address_id in 0..addresses {
-                    name: EventName::FlushStatsToStatsD,
+                    let _ = tx.try_send(Event {
-                    value: 0,
+                        name: EventName::UpdateStats,
-                    process_id: None,
+                        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,
+                    });
+                }
+            }
+        });
+
+        // The collector loop
         loop {
             let stat = match self.rx.recv().await {
                 Some(stat) => stat,
@@ -253,6 +356,14 @@ 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 => {
@@ -280,10 +391,11 @@ impl Collector {
                     *counter += stat.value;
 
                     let counter = stats.entry("maxwait_us").or_insert(0);
+                    let mic_part = stat.value % 1_000_000;
 
                     // Report max time here
-                    if stat.value > *counter {
+                    if mic_part > *counter {
-                        *counter = stat.value;
+                        *counter = mic_part;
                     }
 
                     let counter = stats.entry("maxwait").or_insert(0);
@@ -301,15 +413,16 @@ impl Collector {
                 | EventName::ServerIdle
                 | EventName::ServerTested
                 | EventName::ServerLogin => {
-                    client_server_states.insert(stat.process_id.unwrap(), stat.name);
+                    client_server_states.insert(stat.process_id, stat.name);
                 }
 
                 EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
-                    client_server_states.remove(&stat.process_id.unwrap());
+                    client_server_states.remove(&stat.process_id);
                 }
 
-                EventName::FlushStatsToStatsD => {
+                EventName::UpdateStats => {
-                    for (_, state) in &client_server_states {
+                    // Calculate connection states
+                    for (_, state) in client_server_states.iter() {
                         match state {
                             EventName::ClientActive => {
                                 let counter = stats.entry("cl_active").or_insert(0);
@@ -321,11 +434,6 @@ 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;
@@ -346,22 +454,64 @@ impl Collector {
                                 *counter += 1;
                             }
 
+                            EventName::ClientIdle => {
+                                let counter = stats.entry("cl_idle").or_insert(0);
+                                *counter += 1;
+                            }
+
                             _ => unreachable!(),
                         };
                     }
 
-                    info!("{:?}", stats);
+                    // Update latest stats used in SHOW STATS
-
+                    let mut guard = LATEST_STATS.lock();
-                    let mut pipeline = self.client.pipeline();
+                    for (key, value) in stats.iter() {
-
+                        let entry = guard.entry(stat.address_id).or_insert(HashMap::new());
-                    for (key, value) in stats.iter_mut() {
+                        entry.insert(key.to_string(), value.clone());
-                        pipeline.gauge(key, *value as f64);
-                        *value = 0;
                     }
 
-                    pipeline.send(&self.client);
+                    // These are re-calculated every iteration of the loop, so we don't want to add values
+                    // from the last iteration.
+                    for stat in &[
+                        "cl_active",
+                        "cl_waiting",
+                        "cl_idle",
+                        "sv_idle",
+                        "sv_active",
+                        "sv_tested",
+                        "sv_login",
+                        "maxwait",
+                        "maxwait_us",
+                    ] {
+                        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;
+                    }
                 }
             };
         }
     }
 }
+
+/// Get a snapshot of statistics. Updated once a second
+/// by the `Collector`.
+pub fn get_stats() -> HashMap<usize, HashMap<String, i64>> {
+    LATEST_STATS.lock().clone()
+}

tests/pgbench/simple.sql

@@ -0,0 +1,28 @@
+
+-- \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;
+

tests/ruby/Gemfile

@@ -2,3 +2,4 @@ source "https://rubygems.org"
 
 gem "pg"
 gem "activerecord"
+gem "rubocop"

tests/ruby/Gemfile.lock

@@ -11,13 +11,33 @@ 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
@@ -25,6 +45,7 @@ PLATFORMS
 DEPENDENCIES
   activerecord
   pg
+  rubocop
 
 BUNDLED WITH
    2.3.7

tests/ruby/tests.rb

@@ -1,6 +1,9 @@
 # frozen_string_literal: true
 
 require 'active_record'
+require 'pg'
+
+$stdout.sync = true
 
 # Uncomment these two to see all queries.
 # ActiveRecord.verbose_query_logs = true
@@ -13,6 +16,7 @@ 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
 )
@@ -110,3 +114,17 @@ 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