wip

When recv is called in the mirroring client, we noticed an occasional panic when reading the message.

thread 'tokio-runtime-worker' panicked at 'slice index starts at 5 but ends at 0', src/messages.rs:522:18
We are still debugging the reason why this happens but adding a check for slice bounds seems like a good idea. Instead of panicking, this will return an Err to the caller which will close the connection.

.circleci/config.yml

@@ -9,39 +9,42 @@ jobs:
     # Specify the execution environment. You can specify an image from Dockerhub or use one of our Convenience Images from CircleCI's Developer Hub.
     # See: https://circleci.com/docs/2.0/configuration-reference/#docker-machine-macos-windows-executor
     docker:
-      - image: levkk/pgcat-ci:latest
+      - image: ghcr.io/levkk/pgcat-ci:1.67
         environment:
           RUST_LOG: info
-          RUSTFLAGS: "-C instrument-coverage"
-          LLVM_PROFILE_FILE: "pgcat-%m.profraw"
+          LLVM_PROFILE_FILE: /tmp/pgcat-%m-%p.profraw
+          RUSTC_BOOTSTRAP: 1
+          CARGO_INCREMENTAL: 0
+          RUSTFLAGS: "-Zprofile -Ccodegen-units=1 -Copt-level=0 -Clink-dead-code -Coverflow-checks=off -Zpanic_abort_tests -Cpanic=abort -Cinstrument-coverage"
+          RUSTDOCFLAGS: "-Cpanic=abort"
       - image: postgres:14
-        command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements"]
+        command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
         environment:
           POSTGRES_USER: postgres
           POSTGRES_DB: postgres
           POSTGRES_PASSWORD: postgres
-          POSTGRES_HOST_AUTH_METHOD: scram-sha-256
+          POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
       - image: postgres:14
-        command: ["postgres", "-p", "7432", "-c", "shared_preload_libraries=pg_stat_statements"]
+        command: ["postgres", "-p", "7432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
         environment:
           POSTGRES_USER: postgres
           POSTGRES_DB: postgres
           POSTGRES_PASSWORD: postgres
-          POSTGRES_HOST_AUTH_METHOD: scram-sha-256
+          POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
       - image: postgres:14
-        command: ["postgres", "-p", "8432", "-c", "shared_preload_libraries=pg_stat_statements"]
+        command: ["postgres", "-p", "8432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
         environment:
           POSTGRES_USER: postgres
           POSTGRES_DB: postgres
           POSTGRES_PASSWORD: postgres
-          POSTGRES_HOST_AUTH_METHOD: scram-sha-256
+          POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
       - image: postgres:14
-        command: ["postgres", "-p", "9432", "-c", "shared_preload_libraries=pg_stat_statements"]
+        command: ["postgres", "-p", "9432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
         environment:
           POSTGRES_USER: postgres
           POSTGRES_DB: postgres
           POSTGRES_PASSWORD: postgres
-          POSTGRES_HOST_AUTH_METHOD: scram-sha-256
+          POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
 
     # Add steps to the job
     # See: https://circleci.com/docs/2.0/configuration-reference/#steps
@@ -52,18 +55,9 @@ jobs:
       - run:
           name: "Lint"
           command: "cargo fmt --check"
-      - run:
-          name: "Install dependencies"
-          command: "sudo apt-get update && sudo apt-get install -y psmisc postgresql-contrib-12 postgresql-client-12 ruby ruby-dev libpq-dev python3 python3-pip lcov llvm-11 && sudo apt-get upgrade curl"
-      - run:
-          name: "Install rust tools"
-          command: "cargo install cargo-binutils rustfilt && rustup component add llvm-tools-preview"
-      - run:
-          name: "Build"
-          command: "cargo build"
       - run:
           name: "Tests"
-          command: "cargo test && bash .circleci/run_tests.sh && .circleci/generate_coverage.sh"
+          command: "cargo clean && cargo build && cargo test && bash .circleci/run_tests.sh && .circleci/generate_coverage.sh"
       - store_artifacts:
           path: /tmp/cov
           destination: coverage-data

.circleci/generate_coverage.sh

@@ -1,7 +1,15 @@
 #!/bin/bash
 
-rust-profdata merge -sparse pgcat-*.profraw -o pgcat.profdata
+# inspired by https://doc.rust-lang.org/rustc/instrument-coverage.html#tips-for-listing-the-binaries-automatically
+TEST_OBJECTS=$( \
+    for file in $(cargo test --no-run 2>&1 | grep "target/debug/deps/pgcat-[[:alnum:]]\+" -o); \
+    do \
+        printf "%s %s " --object $file; \
+    done \
+)
 
-rust-cov export -ignore-filename-regex="rustc|registry" -Xdemangler=rustfilt -instr-profile=pgcat.profdata --object ./target/debug/pgcat --format lcov > ./lcov.info
+rust-profdata merge -sparse /tmp/pgcat-*.profraw -o /tmp/pgcat.profdata
 
-genhtml lcov.info --output-directory /tmp/cov --prefix $(pwd)
+bash -c "rust-cov export -ignore-filename-regex='rustc|registry' -Xdemangler=rustfilt -instr-profile=/tmp/pgcat.profdata $TEST_OBJECTS --object ./target/debug/pgcat --format lcov > ./lcov.info"
+
+genhtml lcov.info  --title "PgCat Code Coverage" --css-file ./cov-style.css --no-function-coverage --highlight --ignore-errors source --legend  --output-directory /tmp/cov --prefix $(pwd)

.circleci/pgcat.toml

@@ -18,10 +18,10 @@ enable_prometheus_exporter = true
 prometheus_exporter_port = 9930
 
 # How long to wait before aborting a server connection (ms).
-connect_timeout = 100
+connect_timeout = 1000
 
 # How much time to give the health check query to return with a result (ms).
-healthcheck_timeout = 100
+healthcheck_timeout = 1000
 
 # How long to keep connection available for immediate re-use, without running a healthcheck query on it
 healthcheck_delay = 30000
@@ -32,6 +32,12 @@ shutdown_timeout = 5000
 # For how long to ban a server if it fails a health check (seconds).
 ban_time = 60 # Seconds
 
+# If we should log client connections
+log_client_connections = false
+
+# If we should log client disconnections
+log_client_disconnections = false
+
 # Reload config automatically if it changes.
 autoreload = true
 

.circleci/run_tests.sh

@@ -24,10 +24,6 @@ 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.4.0.deb https://github.com/Shopify/toxiproxy/releases/download/v2.4.0/toxiproxy_2.4.0_linux_$(dpkg --print-architecture).deb
-sudo dpkg -i toxiproxy-2.4.0.deb
-
 # Start Toxiproxy
 LOG_LEVEL=error toxiproxy-server &
 sleep 1
@@ -56,6 +52,17 @@ psql -U sharding_user -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accou
 sleep 1
 killall psql -s SIGINT
 
+# Pause/resume test.
+# Running benches before, during, and after pause/resume.
+pgbench -U sharding_user -t 500 -c 2 -h 127.0.0.1 -p 6432 --protocol extended &
+BENCH_ONE=$!
+PGPASSWORD=admin_pass psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'PAUSE sharded_db,sharding_user'
+pgbench -U sharding_user -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended &
+BENCH_TWO=$!
+PGPASSWORD=admin_pass psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RESUME sharded_db,sharding_user'
+wait ${BENCH_ONE}
+wait ${BENCH_TWO}
+
 # Reload pool (closing unused server connections)
 PGPASSWORD=admin_pass psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD'
 
@@ -85,13 +92,12 @@ sed -i 's/statement_timeout = 100/statement_timeout = 0/' .circleci/pgcat.toml
 kill -SIGHUP $(pgrep pgcat) # Reload config again
 
 #
-# ActiveRecord tests
+# Integration tests and ActiveRecord tests
 #
 cd tests/ruby
-sudo gem install bundler
-bundle install
-bundle exec ruby tests.rb || exit 1
-bundle exec rspec *_spec.rb || exit 1
+sudo bundle install
+bundle exec ruby tests.rb --format documentation || exit 1
+bundle exec rspec *_spec.rb --format documentation || exit 1
 cd ../..
 
 #

.dockerignore

@@ -2,3 +2,5 @@ target/
 tests/
 tracing/
 .circleci/
+.git/
+dev/

.github/dependabot.yml

@@ -0,0 +1,12 @@
+version: 2
+updates:
+  - package-ecosystem: "cargo"
+    directory: "/"
+    schedule:
+      interval: "daily"
+      time: "04:00" # UTC
+    labels:
+      - "dependencies"
+    commit-message:
+      prefix: "chore(deps)"
+    open-pull-requests-limit: 10

.github/workflows/build-and-push.yaml

@@ -0,0 +1,54 @@
+name: Build and Push
+
+on: push
+
+env:
+  registry: ghcr.io
+  image-name: ${{ github.repository }}
+
+jobs:
+  build-and-push:
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - name: Checkout Repository
+        uses: actions/checkout@v3
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+
+      - name: Determine tags
+        id: metadata
+        uses: docker/metadata-action@v4
+        with:
+          images: ${{ env.registry }}/${{ env.image-name }}
+          tags: |
+            type=sha,prefix=,format=long
+            type=schedule
+            type=ref,event=branch
+            type=ref,event=pr
+            type=raw,value=latest,enable={{ is_default_branch }}
+
+      - name: Log in to the Container registry
+        uses: docker/login-action@v2.1.0
+        with:
+          registry: ${{ env.registry }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Build and push ${{ env.image-name }}
+        uses: docker/build-push-action@v3
+        with:
+          push: true
+          tags: ${{ steps.metadata.outputs.tags }}
+          labels: ${{ steps.metadata.outputs.labels }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+concurrency:
+  group: ${{ github.ref }}
+  cancel-in-progress: true

.github/workflows/publish-ci-docker-image.yml

@@ -0,0 +1,20 @@
+name: publish-ci-docker-image
+on:
+  push:
+    branches: [ main ]
+jobs:
+  publish-ci-docker-image:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Login to GitHub Container Registry
+        uses: docker/login-action@v1
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+      - name: Build CI Docker image
+        run: |
+          docker build . -f Dockerfile.ci --tag ghcr.io/levkk/pgcat-ci:latest
+          docker run ghcr.io/levkk/pgcat-ci:latest
+          docker push ghcr.io/levkk/pgcat-ci:latest

.gitignore

@@ -1,4 +1,12 @@
 .idea
 /target
 *.deb
-.vscode
+.vscode
+*.profraw
+cov/
+lcov.info
+
+# Dev
+dev/.bash_history
+dev/cache
+!dev/cache/.keepme

CONFIG.md

@@ -0,0 +1,340 @@
+# PgCat Configurations 
+## `general` Section
+
+### host
+```
+path: general.host
+default: "0.0.0.0"
+```
+
+What IP to run on, 0.0.0.0 means accessible from everywhere.
+
+### port
+```
+path: general.port
+default: 6432
+```
+
+Port to run on, same as PgBouncer used in this example.
+
+### enable_prometheus_exporter
+```
+path: general.enable_prometheus_exporter
+default: true
+```
+
+Whether to enable prometheus exporter or not.
+
+### prometheus_exporter_port
+```
+path: general.prometheus_exporter_port
+default: 9930
+```
+
+Port at which prometheus exporter listens on.
+
+### connect_timeout
+```
+path: general.connect_timeout
+default: 5000 # milliseconds
+```
+
+How long to wait before aborting a server connection (ms).
+
+### idle_timeout
+```
+path: general.idle_timeout
+default: 30000 # milliseconds
+```
+
+How long an idle connection with a server is left open (ms).
+
+### healthcheck_timeout
+```
+path: general.healthcheck_timeout
+default: 1000 # milliseconds
+```
+
+How much time to give the health check query to return with a result (ms).
+
+### healthcheck_delay
+```
+path: general.healthcheck_delay
+default: 30000 # milliseconds
+```
+
+How long to keep connection available for immediate re-use, without running a healthcheck query on it
+
+### shutdown_timeout
+```
+path: general.shutdown_timeout
+default: 60000 # milliseconds
+```
+
+How much time to give clients during shutdown before forcibly killing client connections (ms).
+
+### ban_time
+```
+path: general.ban_time
+default: 60 # seconds
+```
+
+How long to ban a server if it fails a health check (seconds).
+
+### log_client_connections
+```
+path: general.log_client_connections
+default: false
+```
+
+If we should log client connections
+
+### log_client_disconnections
+```
+path: general.log_client_disconnections
+default: false
+```
+
+If we should log client disconnections
+
+### autoreload
+```
+path: general.autoreload
+default: false
+```
+
+When set to true, PgCat reloads configs if it detects a change in the config file.
+
+### worker_threads
+```
+path: general.worker_threads
+default: 5
+```
+
+Number of worker threads the Runtime will use (4 by default).
+
+### tcp_keepalives_idle
+```
+path: general.tcp_keepalives_idle
+default: 5
+```
+
+Number of seconds of connection idleness to wait before sending a keepalive packet to the server.
+
+### tcp_keepalives_count
+```
+path: general.tcp_keepalives_count
+default: 5
+```
+
+Number of unacknowledged keepalive packets allowed before giving up and closing the connection.
+
+### tcp_keepalives_interval
+```
+path: general.tcp_keepalives_interval
+default: 5
+```
+
+Number of seconds between keepalive packets.
+
+### tls_certificate
+```
+path: general.tls_certificate
+default: <UNSET>
+example: "server.cert"
+```
+
+Path to TLS Certficate file to use for TLS connections
+
+### tls_private_key
+```
+path: general.tls_private_key
+default: <UNSET>
+example: "server.key"
+```
+
+Path to TLS private key file to use for TLS connections
+
+### admin_username
+```
+path: general.admin_username
+default: "admin_user"
+```
+
+User name to access the virtual administrative database (pgbouncer or pgcat)
+Connecting to that database allows running commands like `SHOW POOLS`, `SHOW DATABASES`, etc..
+
+### admin_password
+```
+path: general.admin_password
+default: "admin_pass"
+```
+
+Password to access the virtual administrative database
+
+## `pools.<pool_name>` Section
+
+### pool_mode
+```
+path: pools.<pool_name>.pool_mode
+default: "transaction"
+```
+
+Pool mode (see PgBouncer docs for more).
+`session` one server connection per connected client
+`transaction` one server connection per client transaction
+
+### load_balancing_mode
+```
+path: pools.<pool_name>.load_balancing_mode
+default: "random"
+```
+
+Load balancing mode
+`random` selects the server at random
+`loc` selects the server with the least outstanding busy conncetions
+
+### default_role
+```
+path: pools.<pool_name>.default_role
+default: "any"
+```
+
+If the client doesn't specify, PgCat routes 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.
+
+### query_parser_enabled (experimental)
+```
+path: pools.<pool_name>.query_parser_enabled
+default: true
+```
+
+If Query Parser is 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.
+
+### primary_reads_enabled
+```
+path: pools.<pool_name>.primary_reads_enabled
+default: true
+```
+
+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 explicitly selected with our custom protocol.
+
+### sharding_key_regex (experimental)
+```
+path: pools.<pool_name>.sharding_key_regex
+default: <UNSET>
+example: '/\* sharding_key: (\d+) \*/'
+```
+
+Allow sharding commands to be passed as statement comments instead of
+separate commands. If these are unset this functionality is disabled.
+
+### sharding_function
+```
+path: pools.<pool_name>.sharding_function
+default: "pg_bigint_hash"
+```
+
+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
+
+### automatic_sharding_key (experimental)
+```
+path: pools.<pool_name>.automatic_sharding_key
+default: <UNSET>
+example: "data.id"
+```
+
+Automatically parse this from queries and route queries to the right shard!
+
+### idle_timeout
+```
+path: pools.<pool_name>.idle_timeout
+default: 40000
+```
+
+Idle timeout can be overwritten in the pool
+
+### connect_timeout
+```
+path: pools.<pool_name>.connect_timeout
+default: 3000
+```
+
+Connect timeout can be overwritten in the pool
+
+## `pools.<pool_name>.users.<user_index>` Section
+
+### username
+```
+path: pools.<pool_name>.users.<user_index>.username
+default: "sharding_user"
+```
+
+Postgresql username
+
+### password
+```
+path: pools.<pool_name>.users.<user_index>.password
+default: "sharding_user"
+```
+
+Postgresql password
+
+### pool_size
+```
+path: pools.<pool_name>.users.<user_index>.pool_size
+default: 9
+```
+
+Maximum number of server connections that can be established for this user
+The maximum number of connection from a single Pgcat process to any database in the cluster
+is the sum of pool_size across all users.
+
+### statement_timeout
+```
+path: pools.<pool_name>.users.<user_index>.statement_timeout
+default: 0
+```
+
+Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
+0 means it is disabled.
+
+## `pools.<pool_name>.shards.<shard_index>` Section
+
+### servers
+```
+path: pools.<pool_name>.shards.<shard_index>.servers
+default: [["127.0.0.1", 5432, "primary"], ["localhost", 5432, "replica"]]
+```
+
+Array of servers in the shard, each server entry is an array of `[host, port, role]`
+
+### mirrors (experimental)
+```
+path: pools.<pool_name>.shards.<shard_index>.mirrors
+default: <UNSET>
+example: [["1.2.3.4", 5432, 0], ["1.2.3.4", 5432, 1]]
+```
+
+Array of mirrors for the shard, each mirror entry is an array of `[host, port, index of server in servers array]`
+Traffic hitting the server identified by the index will be sent to the mirror.
+
+### database
+```
+path: pools.<pool_name>.shards.<shard_index>.database
+default: "shard0"
+```
+
+Database name (e.g. "postgres")
+

Cargo.toml

@@ -14,23 +14,29 @@ async-trait = "0.1"
 rand = "0.8"
 chrono = "0.4"
 sha-1 = "0.10"
-toml = "0.5"
+toml = "0.7"
 serde = "1"
 serde_derive = "1"
 regex = "1"
 num_cpus = "1"
 once_cell = "1"
-sqlparser = "0.23.0"
+sqlparser = "0.32.0"
 log = "0.4"
 arc-swap = "1"
-env_logger = "0.9"
+env_logger = "0.10"
 parking_lot = "0.12.1"
 hmac = "0.12"
 sha2 = "0.10"
-base64 = "0.13"
+base64 = "0.21"
 stringprep = "0.1"
 tokio-rustls = "0.23"
 rustls-pemfile = "1"
 hyper = { version = "0.14", features = ["full"] }
 phf = { version = "0.11.1", features = ["macros"] }
 exitcode = "1.1.2"
+futures = "0.3"
+socket2 = { version = "0.4.7", features = ["all"] }
+nix = "0.26.2"
+
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+jemallocator = "0.5.0"

Dockerfile.ci

@@ -1,8 +1,12 @@
-FROM cimg/rust:1.62.0
+FROM cimg/rust:1.67.1
 RUN sudo apt-get update && \
-	sudo apt-get install -y psmisc postgresql-contrib-12 postgresql-client-12 ruby ruby-dev libpq-dev python3 python3-pip lcov llvm-11 && \
-	sudo apt-get upgrade curl
-RUN cargo install cargo-binutils rustfilt && \
-	rustup component add llvm-tools-preview
-RUN pip3 install psycopg2 && \
-	sudo gem install bundler
+	sudo apt-get install -y \
+		psmisc postgresql-contrib-14 postgresql-client-14 libpq-dev \
+		ruby ruby-dev python3 python3-pip \
+		lcov llvm-11 iproute2 && \
+	sudo apt-get upgrade curl && \
+	cargo install cargo-binutils rustfilt && \
+	rustup component add llvm-tools-preview && \
+ 	pip3 install psycopg2 && sudo gem install bundler && \
+	wget -O /tmp/toxiproxy-2.4.0.deb https://github.com/Shopify/toxiproxy/releases/download/v2.4.0/toxiproxy_2.4.0_linux_$(dpkg --print-architecture).deb && \
+	sudo dpkg -i /tmp/toxiproxy-2.4.0.deb

README.md

@@ -1,5 +1,3 @@
-![PgCat](./pgcat3.png)
-
 ##### PgCat: PostgreSQL at petabyte scale
 
 [![CircleCI](https://circleci.com/gh/levkk/pgcat/tree/main.svg?style=svg)](https://circleci.com/gh/levkk/pgcat/tree/main)
@@ -36,40 +34,12 @@ For quick local example, use the Docker Compose environment provided:
 docker-compose up
 
 # In a new terminal:
-psql -h 127.0.0.1 -p 6432 -c 'SELECT 1'
+PGPASSWORD=postgres psql -h 127.0.0.1 -p 6432 -U postgres -c 'SELECT 1'
 ```
 
 ### Config
 
-| **Name**                     | **Description**                                                                                                                            | **Examples**                     |
-|------------------------------|--------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------|
-| **`general`**                |                                                                                                                                            |                                  |
-| `host`                       | The pooler will run on this host, 0.0.0.0 means accessible from everywhere.                                                                | `0.0.0.0`                        |
-| `port`                       | The pooler will run on this port.                                                                                                          | `6432`                           |
-| `enable_prometheus_exporter` | Enable prometheus exporter which will export metrics in prometheus exposition format.                                                      | `true`                           |
-| `prometheus_exporter_port`   | Port at which prometheus exporter listens on.                                                                                              | `9930`                           |
-| `pool_size`                  | Maximum allowed server connections per pool. Pools are separated for each user/shard/server role. The connections are allocated as needed. | `15`                             |
-| `pool_mode`                  | The pool mode to use, i.e. `session` or `transaction`.                                                                                     | `transaction`                    |
-| `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`                           |
-| `shutdown_timeout`           | Maximum time to give clients during shutdown before forcibly killing client connections (ms).                                              | `60000`                          |
-| `healthcheck_delay`          | How long to keep connection available for immediate re-use, without running a healthcheck query on it                                      | `30000`                          |
-| `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`                             |
-| `autoreload`                 | Enable auto-reload of config after fixed time-interval.                                                                                    | `false`                          |
-|                              |                                                                                                                                            |                                  |
-| **`user`**                   |                                                                                                                                            |                                  |
-| `name`                       | The user name.                                                                                                                             | `sharding_user`                  |
-| `password`                   | The user password in plaintext.                                                                                                            | `hunter2`                        |
-| `statement_timeout` | Timeout in milliseconds for how long a query takes to execute | `0 (disabled)`
-|                              |                                                                                                                                            |                                  |
-| **`shards`**                 | Shards are numerically numbered starting from 0; the order in the config is preserved by the pooler to route queries accordingly.          | `[shards.0]`                     |
-| `servers`                    | List of servers to connect to and their roles. A server is: `[host, port, role]`, where `role` is either `primary` or `replica`.           | `["127.0.0.1", 5432, "primary"]` |
-| `database`                   | The name of the database to connect to. This is the same on all servers that are part of one shard.                                        |                                  |
-|                              |                                                                                                                                            |                                  |
-| **`query_router`**           |                                                                                                                                            |                                  |
-| `default_role`               | Traffic is routed to this role by default (random), unless the client specifies otherwise. Default is `any`, for any role available.  | `any`, `primary`, `replica`      |
-| `query_parser_enabled`       | Enable the query parser which will inspect incoming queries and route them to a primary or replicas.                                       | `false`                          |
-| `primary_reads_enabled`      | Enable this to allow read queries on the primary; otherwise read queries are routed to the replicas.                                       | `true`                           |
+See [Configurations page](https://github.com/levkk/pgcat/blob/main/CONFIG.md)
 
 ## Local development
 
@@ -111,6 +81,17 @@ docker compose up --exit-code-from main # This will also produce coverage report
 | 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.                                                         |
 
+### Dev
+
+Also, you can open a 'dev' environment where you can debug tests easier by running the following command:
+
+```
+./dev/script/console
+```
+
+This will open a terminal in an environment similar to that used in tests. In there you can compile, run tests, do some debugging with the test environment, etc. Objects
+compiled inside the contaner (and bundled gems) will be placed in `dev/cache` so they don't interfere with what you have in your host.
+
 ## Usage
 
 ### Session mode
@@ -190,7 +171,7 @@ We use the `PARTITION BY HASH` hashing function, the same as used by Postgres fo
 To route queries to a particular shard, we use this custom SQL syntax:
 
 ```sql
--- To talk to a shard explicitely
+-- To talk to a shard explicitly
 SET SHARD TO '1';
 
 -- To let the pooler choose based on a value

cov-style.css

@@ -0,0 +1,158 @@
+/*
+ * Copyright 2021 Collabora, Ltd.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files (the
+ * "Software"), to deal in the Software without restriction, including
+ * without limitation the rights to use, copy, modify, merge, publish,
+ * distribute, sublicense, and/or sell copies of the Software, and to
+ * permit persons to whom the Software is furnished to do so, subject to
+ * the following conditions:
+ *
+ * The above copyright notice and this permission notice (including the
+ * next paragraph) shall be included in all copies or substantial
+ * portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+ * NONINFRINGEMENT.  IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
+ * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
+ * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+ * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+body {
+	background-color: #f2f2f2;
+	font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto,
+		"Noto Sans", Ubuntu, Cantarell, "Helvetica Neue", sans-serif,
+		"Apple Color Emoji", "Segoe UI Emoji", "Segoe UI Symbol",
+		"Noto Color Emoji";
+}
+
+.sourceHeading, .source, .coverFn,
+.testName, .testPer, .testNum,
+.coverLegendCovLo, .headerCovTableEntryLo, .coverPerLo, .coverNumLo,
+.coverLegendCovMed, .headerCovTableEntryMed, .coverPerMed, .coverNumMed,
+.coverLegendCovHi, .headerCovTableEntryHi, .coverPerHi, .coverNumHi,
+.coverFile {
+	font-family: "Menlo", "DejaVu Sans Mono", "Liberation Mono",
+		"Consolas", "Ubuntu Mono", "Courier New", "andale mono",
+		"lucida console", monospace;
+}
+
+pre {
+	font-size: 0.7875rem;
+}
+
+.headerCovTableEntry, .testPer, .testNum, .testName,
+.coverLegendCovLo, .headerCovTableEntryLo, .coverPerLo, .coverNumLo,
+.coverLegendCovMed, .headerCovTableEntryMed, .coverPerMed, .coverNumMed,
+.coverLegendCovHi, .headerCovTableEntryHi, .coverPerHi, .coverNumHi {
+	text-align: right;
+	white-space: nowrap;
+}
+
+.coverPerLo, .coverPerMed, .coverPerHi, .testPer {
+/*	font-weight: bold;*/
+}
+
+.coverNumLo, .coverNumMed, .coverNumHi, .testNum {
+	font-style: italic;
+	font-size: 90%;
+	padding-left: 1em;
+}
+
+.title {
+	font-size: 200%;
+}
+
+.tableHead {
+	text-align: center;
+	font-weight: bold;
+	background-color: #bfbfbf;
+}
+
+.coverFile, .coverBar, .coverFn {
+	background-color: #d9d9d9;
+}
+
+.headerCovTableHead {
+	font-weight: bold;
+	text-align: right;
+}
+
+.headerCovTableEntry {
+	background-color: #d9d9d9;
+}
+
+.coverFnLo,
+.coverLegendCovLo, .headerCovTableEntryLo, .coverPerLo, .coverNumLo {
+	background-color: #f2dada;
+}
+
+.coverFnHi,
+.coverLegendCovMed, .headerCovTableEntryMed, .coverPerMed, .coverNumMed {
+	background-color: #add9ad;
+}
+
+.coverLegendCovHi, .headerCovTableEntryHi, .coverPerHi, .coverNumHi {
+	background-color: #59b359;
+}
+
+.coverBarOutline {
+	border-style: solid;
+	border-width: 1px;
+	border-color: black;
+	padding: 0px;
+}
+
+.coverFnLo, .coverFnHi {
+	text-align: right;
+}
+
+.lineNum {
+	background-color: #d9d9d9;
+}
+
+.coverLegendCov, .lineCov, .branchCov {
+	background-image: url('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAABCAIAAABsYngUAAADAXpUWHRSYXcgcHJvZmlsZSB0eXBlIGV4aWYAAHjazZVbktwgDEX/WUWWgCSExHIwj6rsIMvPxcY9PY9MzVTyEVMNtCwkoYNwGL9+zvADDxHHkNQ8l5wjnlRS4YqJx+upZ08xnf313O/otTw8FBgzwShbP2/5gJyhz1vetp0KuT4ZKmO/OF6/qNsQ+3ZwO9yOhC4HcRsOdRsS3p7T9f+4thVzcXveQtv6sz5t1dfW0CUxzprJEvrE0SwXzJ1jMuStr0CPvhfqdvTmf7hVGTHxEJKI3leEsn4kFWNCT/CGfUnBXDEuyd4yaHGIhnm58/r581nk4Q59Y32N+p69Qc3xPelwJvRWkTeE8mP8UE76Ig/PSE9uT55z3jN+LZ/pJaibXLjxzdl9znHtrqaMLee9qXuL5wx6x8rWuSqjGX4afSV7tYLmKImGc9RxyA60RoUYGCcl6lRp0jjHRg0hJh4MjszcALcFCB0wCjcgJYBGo8kGzF0cB6DhOAik/IiFTrfldNfI4biTB5wegjHCkr9q4StKc66CIlq55CtXiItXwhHFIkeE6ocaiNDcSdUzwXd7+yyuAoJ6ptmxwRqPZQH4D6WXwyUnaIGiYrwKmKxvA0gRIlAEQwICMZMoZYrGHIwIiXQAqgidJfEBLKTKHUFyEsmAgyqAb6wxOlVZ+RLjIgQIlRzEwAaFCFgpKc6PJccZqiqaVDWrqWvRmiWvCsvZ8rpRq4klU8tm5lasBhdPrp7d3L14LVwEN64W1GPxUkqtcFphuWJ1hUKtBx9ypEOPfNjhRzlq49CkpaYtN2veSqudu3TUcc/duvfS66CBozTS0JGHDR9l1ImjNmWmqTNPmz5LmPVBbWN9175BjTY1PkktRXtQg9TsNkHrOtHFDMQ4EYDbIkASmBez6JQSL3KLWSyMqlBGkLrgdFrEQDANYp30YPdCToPkf8MtAAT/C3JhofsCuffcPqLW6/mhk5PQKsOV1CiovpHgnx3LcCvhwlnz9dF8P4Y/vfju+J8aQpZK+A373P3XzDqcKwAAAAZiS0dEAAAAAAAA+UO7fwAAAAlwSFlzAAAOxAAADsQBlSsOGwAAAAd0SU1FB+UEEQYyDQA04tUAAAAZdEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIEdJTVBXgQ4XAAAADklEQVQI12PULVBlwAYAEagAxGHRDdwAAAAASUVORK5CYII=');
+	background-repeat: repeat-y;
+	background-position: left top;
+	background-color: #c6ffb8;
+}
+
+.coverLegendNoCov, .lineNoCov, .branchNoCov, .branchNoExec {
+	background-image: url('data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAgAAAABCAIAAABsYngUAAAACXBIWXMAAA7EAAAOxAGVKw4bAAAAB3RJTUUH5QMUCiMidNgp2gAAABl0RVh0Q29tbWVudABDcmVhdGVkIHdpdGggR0lNUFeBDhcAAAAPSURBVAjXY/wZIcWADQAAIa4BbZaExr0AAAAASUVORK5CYII=');
+	background-repeat: repeat-y;
+	background-position: left top;
+	background-color: #ffcfbb;
+}
+
+.coverLegendCov, .coverLegendNoCov {
+	padding: 0em 1em 0em 1em;
+}
+
+.headerItem, .headerValue, .headerValueLeg {
+	white-space: nowrap;
+}
+
+.headerItem {
+	text-align: right;
+	font-weight: bold;
+}
+
+.ruler {
+	background-color: #d9d9d9;
+}
+
+.detail {
+	font-size: 80%;
+}
+
+.versionInfo {
+	font-size: 80%;
+	text-align: right;
+}
+

dev/Dockerfile

@@ -0,0 +1,33 @@
+FROM rust:bullseye
+
+# Dependencies
+RUN apt-get update -y \
+    && apt-get install -y \
+    llvm-11 psmisc postgresql-contrib postgresql-client \
+    ruby ruby-dev libpq-dev python3 python3-pip lcov curl sudo iproute2 \
+    strace ngrep iproute2 dnsutils lsof net-tools telnet
+
+# Rust
+RUN cargo install cargo-binutils rustfilt
+RUN rustup component add llvm-tools-preview
+
+# Ruby
+RUN sudo gem install bundler
+
+# Toxyproxy
+RUN wget -O toxiproxy-2.4.0.deb https://github.com/Shopify/toxiproxy/releases/download/v2.4.0/toxiproxy_2.4.0_linux_$(dpkg --print-architecture).deb && \
+    sudo dpkg -i toxiproxy-2.4.0.deb
+
+# Config
+ENV APP_ROOT=/app
+ARG APP_USER=pgcat
+COPY dev_bashrc /etc/bash.bashrc
+
+RUN useradd -m -o -u 999 ${APP_USER} || exit 0 && mkdir ${APP_ROOT} && chown ${APP_USER} ${APP_ROOT}
+RUN adduser ${APP_USER} sudo \
+    && echo "${APP_USER} ALL=NOPASSWD: ALL" > /etc/sudoers.d/${APP_USER} \
+    && chmod ugo+s /usr/sbin/usermod /usr/sbin/groupmod
+ENV HOME=${APP_ROOT}
+WORKDIR ${APP_ROOT}
+
+ENTRYPOINT ["/bin/bash"]

dev/dev_bashrc

@@ -0,0 +1,120 @@
+# ~/.bashrc: executed by bash(1) for non-login shells.
+# see /usr/share/doc/bash/examples/startup-files (in the package bash-doc)
+# for examples
+
+# FIX USER NEEDED SO WE CAN SHARE UID BETWEEN HOST AND DEV ENV
+usermod -o -u $(id -u) pgcat 
+groupmod -o -g $(id -g) pgcat 
+
+# We fix the setuid in those commands as we now have sudo
+sudo chmod ugo-s /usr/sbin/usermod /usr/sbin/groupmod 
+
+# Environment customization
+export DEV_ROOT="${APP_ROOT}/dev"
+export HISTFILE="${DEV_ROOT}/.bash_history"
+export CARGO_TARGET_DIR="${DEV_ROOT}/cache/target"
+export CARGO_HOME="${DEV_ROOT}/cache/target/.cargo"
+export BUNDLE_PATH="${DEV_ROOT}/cache/bundle"
+
+# Regular bashrc
+# If not running interactively, don't do anything
+case $- in
+  *i*) ;;
+  *) return;;
+esac
+
+# don't put duplicate lines or lines starting with space in the history.
+# See bash(1) for more options
+HISTCONTROL=ignoreboth
+
+# append to the history file, don't overwrite it
+shopt -s histappend
+
+# for setting history length see HISTSIZE and HISTFILESIZE in bash(1)
+HISTSIZE=1000
+HISTFILESIZE=2000
+
+# check the window size after each command and, if necessary,
+# update the values of LINES and COLUMNS.
+shopt -s checkwinsize
+
+# If set, the pattern "**" used in a pathname expansion context will
+# match all files and zero or more directories and subdirectories.
+#shopt -s globstar
+
+# make less more friendly for non-text input files, see lesspipe(1)
+[ -x /usr/bin/lesspipe ] && eval "$(SHELL=/bin/sh lesspipe)"
+
+# set variable identifying the chroot you work in (used in the prompt below)
+if [ -z "${debian_chroot:-}" ] && [ -r /etc/debian_chroot ]; then
+  debian_chroot=$(cat /etc/debian_chroot)
+fi
+
+# set a fancy prompt (non-color, unless we know we "want" color)
+case "$TERM" in
+  xterm-color|*-256color) color_prompt=yes;;
+esac
+
+# uncomment for a colored prompt, if the terminal has the capability; turned
+# off by default to not distract the user: the focus in a terminal window
+# should be on the output of commands, not on the prompt
+#force_color_prompt=yes
+
+if [ -n "$force_color_prompt" ]; then
+  if [ -x /usr/bin/tput ] && tput setaf 1 >&/dev/null; then
+    # We have color support; assume it's compliant with Ecma-48
+    # (ISO/IEC-6429). (Lack of such support is extremely rare, and such
+    # a case would tend to support setf rather than setaf.)
+    color_prompt=yes
+  else
+    color_prompt=
+  fi
+fi
+
+PS1='\[\e]0;pgcat@dev-container\h: \w\a\]${debian_chroot:+($debian_chroot)}\[\033[01;32m\]pgcat\[\033[00m\]@\[\033[01;32m\]dev-container\[\033[00m\]:\[\033[01;34m\]\w\[\033[00m\]\[\033[01;31m\]$(git branch &>/dev/null; if [ $? -eq 0 ]; then echo " ($(git branch | grep ^* |sed s/\*\ //))"; fi)\[\033[00m\]\$ '
+
+unset color_prompt force_color_prompt
+
+# enable color support of ls and also add handy aliases
+if [ -x /usr/bin/dircolors ]; then
+  test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)"
+  alias ls='ls --color=auto'
+  #alias dir='dir --color=auto'
+  #alias vdir='vdir --color=auto'
+
+  alias grep='grep --color=auto'
+  alias fgrep='fgrep --color=auto'
+  alias egrep='egrep --color=auto'
+fi
+
+# colored GCC warnings and errors
+#export GCC_COLORS='error=01;31:warning=01;35:note=01;36:caret=01;32:locus=01:quote=01'
+
+# some more ls aliases
+alias ll='ls -alF'
+alias la='ls -A'
+alias l='ls -CF'
+
+# Add an "alert" alias for long running commands.  Use like so:
+#   sleep 10; alert
+alias alert='notify-send --urgency=low -i "$([ $? = 0 ] && echo terminal || echo error)" "$(history|tail -n1|sed -e '\''s/^\s*[0-9]\+\s*//;s/[;&|]\s*alert$//'\'')"'
+
+# Alias definitions.
+# You may want to put all your additions into a separate file like
+# ~/.bash_aliases, instead of adding them here directly.
+# See /usr/share/doc/bash-doc/examples in the bash-doc package.
+
+if [ -f ~/.bash_aliases ]; then
+  . ~/.bash_aliases
+fi
+
+# enable programmable completion features (you don't need to enable
+# this, if it's already enabled in /etc/bash.bashrc and /etc/profile
+# sources /etc/bash.bashrc).
+if ! shopt -oq posix; then
+  if [ -f /usr/share/bash-completion/bash_completion ]; then
+    . /usr/share/bash-completion/bash_completion
+  elif [ -f /etc/bash_completion ]; then
+    . /etc/bash_completion
+  fi
+fi

dev/docker-compose.yaml

@@ -0,0 +1,84 @@
+version: "3"
+
+x-common-definition-pg:
+  &common-definition-pg
+  image: postgres:14
+  network_mode: "service:main"
+  healthcheck:
+    test: [ "CMD-SHELL", "pg_isready -U postgres -d postgres" ]
+    interval: 5s
+    timeout: 5s
+    retries: 5
+  volumes:
+    - type: bind
+      source: ../tests/sharding/query_routing_setup.sql
+      target: /docker-entrypoint-initdb.d/query_routing_setup.sql
+    - type: bind
+      source: ../tests/sharding/partition_hash_test_setup.sql
+      target: /docker-entrypoint-initdb.d/partition_hash_test_setup.sql
+
+x-common-env-pg:
+  &common-env-pg
+  POSTGRES_USER: postgres
+  POSTGRES_DB: postgres
+  POSTGRES_PASSWORD: postgres
+
+services:
+  main:
+    image: kubernetes/pause
+
+  pg1:
+    <<: *common-definition-pg
+    environment:
+      <<: *common-env-pg
+      POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
+      PGPORT: 5432
+    command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+
+  pg2:
+    <<: *common-definition-pg
+    environment:
+      <<: *common-env-pg
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+      PGPORT: 7432
+    command: ["postgres", "-p", "7432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+  pg3:
+    <<: *common-definition-pg
+    environment:
+      <<: *common-env-pg
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+      PGPORT: 8432
+    command: ["postgres", "-p", "8432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+  pg4:
+    <<: *common-definition-pg
+    environment:
+      <<: *common-env-pg
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+      PGPORT: 9432
+    command: ["postgres", "-p", "9432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+
+  toxiproxy:
+    build: .
+    network_mode: "service:main"
+    container_name: toxiproxy
+    environment:
+      LOG_LEVEL: info
+    entrypoint: toxiproxy-server
+    depends_on:
+      - pg1
+      - pg2
+      - pg3
+      - pg4
+
+  pgcat-shell:
+    stdin_open: true
+    user: "${HOST_UID}:${HOST_GID}"
+    build: .
+    network_mode: "service:main"
+    depends_on:
+      - toxiproxy
+    volumes:
+      - ../:/app/
+    entrypoint:
+      - /bin/bash
+      - -i

dev/script/console

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"
+export HOST_UID="$(id -u)"
+export HOST_GID="$(id -g)"
+
+if [[ "${1}" == "down" ]]; then
+	docker-compose -f "${DIR}/../docker-compose.yaml" down
+	exit 0
+else
+	docker-compose -f "${DIR}/../docker-compose.yaml" run --rm pgcat-shell
+fi

examples/docker/pgcat.toml

@@ -32,6 +32,12 @@ shutdown_timeout = 60000
 # For how long to ban a server if it fails a health check (seconds).
 ban_time = 60 # seconds
 
+# If we should log client connections
+log_client_connections = false
+
+# If we should log client disconnections
+log_client_disconnections = false
+
 # Reload config automatically if it changes.
 autoreload = false
 
@@ -48,7 +54,7 @@ admin_password = "postgres"
 # configs are structured as pool.<pool_name>
 # the pool_name is what clients use as database name when connecting
 # For the example below a client can connect using "postgres://sharding_user:sharding_user@pgcat_host:pgcat_port/sharded"
-[pools.sharded]
+[pools.postgres]
 # Pool mode (see PgBouncer docs for more).
 # session: one server connection per connected client
 # transaction: one server connection per client transaction
@@ -84,7 +90,7 @@ primary_reads_enabled = true
 sharding_function = "pg_bigint_hash"
 
 # Credentials for users that may connect to this cluster
-[pools.sharded.users.0]
+[pools.postgres.users.0]
 username = "postgres"
 password = "postgres"
 # Maximum number of server connections that can be established for this user
@@ -95,14 +101,8 @@ pool_size = 9
 # Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
 statement_timeout = 0
 
-[pools.sharded.users.1]
-username = "postgres"
-password = "postgres"
-pool_size = 21
-statement_timeout = 15000
-
 # Shard 0
-[pools.sharded.shards.0]
+[pools.postgres.shards.0]
 # [ host, port, role ]
 servers = [
     [ "postgres", 5432, "primary" ],
@@ -111,37 +111,16 @@ servers = [
 # Database name (e.g. "postgres")
 database = "postgres"
 
-[pools.sharded.shards.1]
+[pools.postgres.shards.1]
 servers = [
     [ "postgres", 5432, "primary" ],
     [ "postgres", 5432, "replica" ],
 ]
 database = "postgres"
 
-[pools.sharded.shards.2]
+[pools.postgres.shards.2]
 servers = [
     [ "postgres", 5432, "primary" ],
     [ "postgres", 5432, "replica" ],
 ]
 database = "postgres"
-
-
-[pools.simple_db]
-pool_mode = "session"
-default_role = "primary"
-query_parser_enabled = true
-primary_reads_enabled = true
-sharding_function = "pg_bigint_hash"
-
-[pools.simple_db.users.0]
-username = "postgres"
-password = "postgres"
-pool_size = 5
-statement_timeout = 0
-
-[pools.simple_db.shards.0]
-servers = [
-    [ "postgres", 5432, "primary" ],
-    [ "postgres", 5432, "replica" ]
-]
-database = "postgres"

pgcat.toml

@@ -18,51 +18,75 @@ enable_prometheus_exporter = true
 prometheus_exporter_port = 9930
 
 # How long to wait before aborting a server connection (ms).
-connect_timeout = 5000
+connect_timeout = 5000 # milliseconds
+
+# How long an idle connection with a server is left open (ms).
+idle_timeout = 30000 # milliseconds
 
 # How much time to give the health check query to return with a result (ms).
-healthcheck_timeout = 1000
+healthcheck_timeout = 1000 # milliseconds
 
 # How long to keep connection available for immediate re-use, without running a healthcheck query on it
-healthcheck_delay = 30000
+healthcheck_delay = 30000 # milliseconds
 
 # How much time to give clients during shutdown before forcibly killing client connections (ms).
-shutdown_timeout = 60000
+shutdown_timeout = 60000 # milliseconds
 
-# For how long to ban a server if it fails a health check (seconds).
+# How long to ban a server if it fails a health check (seconds).
 ban_time = 60 # seconds
 
-# Reload config automatically if it changes.
+# If we should log client connections
+log_client_connections = false
+
+# If we should log client disconnections
+log_client_disconnections = false
+
+# When set to true, PgCat reloads configs if it detects a change in the config file.
 autoreload = false
 
-# TLS
+# Number of worker threads the Runtime will use (4 by default).
+worker_threads = 5
+
+# Number of seconds of connection idleness to wait before sending a keepalive packet to the server.
+tcp_keepalives_idle = 5
+# Number of unacknowledged keepalive packets allowed before giving up and closing the connection.
+tcp_keepalives_count = 5
+# Number of seconds between keepalive packets.
+tcp_keepalives_interval = 5
+
+# Path to TLS Certficate file to use for TLS connections
 # tls_certificate = "server.cert"
+# Path to TLS private key file to use for TLS connections
 # tls_private_key = "server.key"
 
-# Credentials to access the virtual administrative database (pgbouncer or pgcat)
+# User name to access the virtual administrative database (pgbouncer or pgcat)
 # Connecting to that database allows running commands like `SHOW POOLS`, `SHOW DATABASES`, etc..
 admin_username = "admin_user"
+# Password to access the virtual administrative database
 admin_password = "admin_pass"
 
-# pool
-# configs are structured as pool.<pool_name>
-# the pool_name is what clients use as database name when connecting
-# For the example below a client can connect using "postgres://sharding_user:sharding_user@pgcat_host:pgcat_port/sharded_db"
+# pool configs are structured as pool.<pool_name>
+# the pool_name is what clients use as database name when connecting.
+# For a pool named `sharded_db`, clients access that pool using connection string like
+# `postgres://sharding_user:sharding_user@pgcat_host:pgcat_port/sharded_db`
 [pools.sharded_db]
 # Pool mode (see PgBouncer docs for more).
-# session: one server connection per connected client
-# transaction: one server connection per client transaction
+# `session` one server connection per connected client
+# `transaction` one server connection per client transaction
 pool_mode = "transaction"
 
-# 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.
+# Load balancing mode
+# `random` selects the server at random
+# `loc` selects the server with the least outstanding busy conncetions
+load_balancing_mode = "random"
+
+# If the client doesn't specify, PgCat routes 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
+# If Query Parser is 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.
@@ -73,19 +97,34 @@ query_parser_enabled = true
 # queries. The primary can always be explicitly selected with our custom protocol.
 primary_reads_enabled = true
 
+# Allow sharding commands to be passed as statement comments instead of
+# separate commands. If these are unset this functionality is disabled.
+# sharding_key_regex = '/\* sharding_key: (\d+) \*/'
+# shard_id_regex = '/\* shard_id: (\d+) \*/'
+# regex_search_limit = 1000 # only look at the first 1000 characters of SQL statements
+
 # 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
-#
+# `pg_bigint_hash`: PARTITION BY HASH (Postgres hashing function)
+# `sha1`: A hashing function based on SHA1
 sharding_function = "pg_bigint_hash"
 
-# Credentials for users that may connect to this cluster
+# Automatically parse this from queries and route queries to the right shard!
+# automatic_sharding_key = "data.id"
+
+# Idle timeout can be overwritten in the pool
+idle_timeout = 40000
+
+# Connect timeout can be overwritten in the pool
+connect_timeout = 3000
+
+# User configs are structured as pool.<pool_name>.users.<user_index>
+# This secion holds the credentials for users that may connect to this cluster
 [pools.sharded_db.users.0]
+# Postgresql username
 username = "sharding_user"
+# Postgresql password
 password = "sharding_user"
 # Maximum number of server connections that can be established for this user
 # The maximum number of connection from a single Pgcat process to any database in the cluster
@@ -93,6 +132,7 @@ password = "sharding_user"
 pool_size = 9
 
 # Maximum query duration. Dangerous, but protects against DBs that died in a non-obvious way.
+# 0 means it is disabled.
 statement_timeout = 0
 
 [pools.sharded_db.users.1]
@@ -101,28 +141,26 @@ password = "other_user"
 pool_size = 21
 statement_timeout = 15000
 
-# Shard 0
+# Shard configs are structured as pool.<pool_name>.shards.<shard_id>
+# Each shard config contains a list of servers that make up the shard
+# and the database name to use.
 [pools.sharded_db.shards.0]
-# [ host, port, role ]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ]
-]
+# Array of servers in the shard, each server entry is an array of `[host, port, role]`
+servers = [["127.0.0.1", 5432, "primary"], ["localhost", 5432, "replica"]]
+
+# Array of mirrors for the shard, each mirror entry is an array of `[host, port, index of server in servers array]`
+# Traffic hitting the server identified by the index will be sent to the mirror.
+# mirrors = [["1.2.3.4", 5432, 0], ["1.2.3.4", 5432, 1]]
+
 # Database name (e.g. "postgres")
 database = "shard0"
 
 [pools.sharded_db.shards.1]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ],
-]
+servers = [["127.0.0.1", 5432, "primary"], ["localhost", 5432, "replica"]]
 database = "shard1"
 
 [pools.sharded_db.shards.2]
-servers = [
-    [ "127.0.0.1", 5432, "primary" ],
-    [ "localhost", 5432, "replica" ],
-]
+servers = [["127.0.0.1", 5432, "primary" ], ["localhost", 5432, "replica" ]]
 database = "shard2"
 
 

src/admin.rs

@@ -1,13 +1,17 @@
+use crate::pool::BanReason;
 /// Admin database.
 use bytes::{Buf, BufMut, BytesMut};
-use log::{info, trace};
+use log::{error, info, trace};
+use nix::sys::signal::{self, Signal};
+use nix::unistd::Pid;
 use std::collections::HashMap;
+use std::time::{SystemTime, UNIX_EPOCH};
 use tokio::time::Instant;
 
 use crate::config::{get_config, reload_config, VERSION};
 use crate::errors::Error;
 use crate::messages::*;
-use crate::pool::get_all_pools;
+use crate::pool::{get_all_pools, get_pool};
 use crate::stats::{
     get_address_stats, get_client_stats, get_pool_stats, get_server_stats, ClientState, ServerState,
 };
@@ -22,7 +26,7 @@ pub fn generate_server_info_for_admin() -> BytesMut {
     server_info.put(server_parameter_message("server_version", VERSION));
     server_info.put(server_parameter_message("DateStyle", "ISO, MDY"));
 
-    return server_info;
+    server_info
 }
 
 /// Handle admin client.
@@ -37,19 +41,28 @@ where
     let code = query.get_u8() as char;
 
     if code != 'Q' {
-        return Err(Error::ProtocolSyncError);
+        return Err(Error::ProtocolSyncError(format!(
+            "Invalid code, expected 'Q' but got '{}'",
+            code
+        )));
     }
 
     let len = query.get_i32() as usize;
-    let query = String::from_utf8_lossy(&query[..len - 5])
-        .to_string()
-        .to_ascii_uppercase();
+    let query = String::from_utf8_lossy(&query[..len - 5]).to_string();
 
     trace!("Admin query: {}", query);
 
     let query_parts: Vec<&str> = query.trim_end_matches(';').split_whitespace().collect();
 
-    match query_parts[0] {
+    match query_parts[0].to_ascii_uppercase().as_str() {
+        "BAN" => {
+            trace!("BAN");
+            ban(stream, query_parts).await
+        }
+        "UNBAN" => {
+            trace!("UNBAN");
+            unban(stream, query_parts).await
+        }
         "RELOAD" => {
             trace!("RELOAD");
             reload(stream, client_server_map).await
@@ -58,7 +71,23 @@ where
             trace!("SET");
             ignore_set(stream).await
         }
-        "SHOW" => match query_parts[1] {
+        "PAUSE" => {
+            trace!("PAUSE");
+            pause(stream, query_parts[1]).await
+        }
+        "RESUME" => {
+            trace!("RESUME");
+            resume(stream, query_parts[1]).await
+        }
+        "SHUTDOWN" => {
+            trace!("SHUTDOWN");
+            shutdown(stream).await
+        }
+        "SHOW" => match query_parts[1].to_ascii_uppercase().as_str() {
+            "BANS" => {
+                trace!("SHOW BANS");
+                show_bans(stream).await
+            }
             "CONFIG" => {
                 trace!("SHOW CONFIG");
                 show_config(stream).await
@@ -91,6 +120,10 @@ where
                 trace!("SHOW VERSION");
                 show_version(stream).await
             }
+            "USERS" => {
+                trace!("SHOW USERS");
+                show_users(stream).await
+            }
             _ => error_response(stream, "Unsupported SHOW query against the admin database").await,
         },
         _ => error_response(stream, "Unsupported query against the admin database").await,
@@ -168,7 +201,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Show PgCat version.
@@ -179,14 +212,14 @@ where
     let mut res = BytesMut::new();
 
     res.put(row_description(&vec![("version", DataType::Text)]));
-    res.put(data_row(&vec![format!("PgCat {}", VERSION).to_string()]));
+    res.put(data_row(&vec![format!("PgCat {}", VERSION)]));
     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
+    write_all_half(stream, &res).await
 }
 
 /// Show utilization of connection pools for each shard and replicas.
@@ -247,7 +280,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Show shards and replicas.
@@ -283,7 +316,8 @@ where
             for server in 0..pool.servers(shard) {
                 let address = pool.address(shard, server);
                 let pool_state = pool.pool_state(shard, server);
-                let banned = pool.is_banned(address, Some(address.role));
+                let banned = pool.is_banned(address);
+                let paused = pool.paused();
 
                 res.put(data_row(&vec![
                     address.name(),                         // name
@@ -297,7 +331,11 @@ where
                     pool_config.pool_mode.to_string(),      // pool_mode
                     pool_config.user.pool_size.to_string(), // max_connections
                     pool_state.connections.to_string(),     // current_connections
-                    "0".to_string(),                        // paused
+                    match paused {
+                        // paused
+                        true => "1".to_string(),
+                        false => "0".to_string(),
+                    },
                     match banned {
                         // disabled
                         true => "1".to_string(),
@@ -314,7 +352,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Ignore any SET commands the client sends.
@@ -326,6 +364,163 @@ where
     custom_protocol_response_ok(stream, "SET").await
 }
 
+/// Bans a host from being used
+async fn ban<T>(stream: &mut T, tokens: Vec<&str>) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let host = match tokens.get(1) {
+        Some(host) => host,
+        None => return error_response(stream, "usage: BAN hostname duration_seconds").await,
+    };
+
+    let duration_seconds = match tokens.get(2) {
+        Some(duration_seconds) => match duration_seconds.parse::<i64>() {
+            Ok(duration_seconds) => duration_seconds,
+            Err(_) => {
+                return error_response(stream, "duration_seconds must be an integer").await;
+            }
+        },
+        None => return error_response(stream, "usage: BAN hostname duration_seconds").await,
+    };
+
+    if duration_seconds <= 0 {
+        return error_response(stream, "duration_seconds must be >= 0").await;
+    }
+
+    let columns = vec![
+        ("db", DataType::Text),
+        ("user", DataType::Text),
+        ("role", DataType::Text),
+        ("host", DataType::Text),
+    ];
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    for (id, pool) in get_all_pools().iter() {
+        for address in pool.get_addresses_from_host(host) {
+            if !pool.is_banned(&address) {
+                pool.ban(&address, BanReason::AdminBan(duration_seconds), -1);
+                res.put(data_row(&vec![
+                    id.db.clone(),
+                    id.user.clone(),
+                    address.role.to_string(),
+                    address.host,
+                ]));
+            }
+        }
+    }
+
+    res.put(command_complete("BAN"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
+/// Clear a host for use
+async fn unban<T>(stream: &mut T, tokens: Vec<&str>) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let host = match tokens.get(1) {
+        Some(host) => host,
+        None => return error_response(stream, "UNBAN command requires a hostname to unban").await,
+    };
+
+    let columns = vec![
+        ("db", DataType::Text),
+        ("user", DataType::Text),
+        ("role", DataType::Text),
+        ("host", DataType::Text),
+    ];
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    for (id, pool) in get_all_pools().iter() {
+        for address in pool.get_addresses_from_host(host) {
+            if pool.is_banned(&address) {
+                pool.unban(&address);
+                res.put(data_row(&vec![
+                    id.db.clone(),
+                    id.user.clone(),
+                    address.role.to_string(),
+                    address.host,
+                ]));
+            }
+        }
+    }
+
+    res.put(command_complete("UNBAN"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
+/// Shows all the bans
+async fn show_bans<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let columns = vec![
+        ("db", DataType::Text),
+        ("user", DataType::Text),
+        ("role", DataType::Text),
+        ("host", DataType::Text),
+        ("reason", DataType::Text),
+        ("ban_time", DataType::Text),
+        ("ban_duration_seconds", DataType::Text),
+        ("ban_remaining_seconds", DataType::Text),
+    ];
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    // The block should be pretty quick so we cache the time outside
+    let now = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .expect("Time went backwards")
+        .as_secs() as i64;
+
+    for (id, pool) in get_all_pools().iter() {
+        for (address, (ban_reason, ban_time)) in pool.get_bans().iter() {
+            let ban_duration = match ban_reason {
+                BanReason::AdminBan(duration) => *duration,
+                _ => pool.settings.ban_time,
+            };
+            let remaining = ban_duration - (now - ban_time.timestamp());
+            if remaining <= 0 {
+                continue;
+            }
+            res.put(data_row(&vec![
+                id.db.clone(),
+                id.user.clone(),
+                address.role.to_string(),
+                address.host.clone(),
+                format!("{:?}", ban_reason),
+                ban_time.to_string(),
+                ban_duration.to_string(),
+                remaining.to_string(),
+            ]));
+        }
+    }
+
+    res.put(command_complete("SHOW BANS"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
 /// Reload the configuration file without restarting the process.
 async fn reload<T>(stream: &mut T, client_server_map: ClientServerMap) -> Result<(), Error>
 where
@@ -346,7 +541,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Shows current configuration.
@@ -392,7 +587,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Show shard and replicas statistics.
@@ -452,7 +647,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Show currently connected clients
@@ -502,7 +697,7 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
 }
 
 /// Show currently connected servers
@@ -556,5 +751,156 @@ where
     res.put_i32(5);
     res.put_u8(b'I');
 
-    write_all_half(stream, res).await
+    write_all_half(stream, &res).await
+}
+
+/// Pause a pool. It won't pass any more queries to the backends.
+async fn pause<T>(stream: &mut T, query: &str) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let parts: Vec<&str> = query.split(",").map(|part| part.trim()).collect();
+
+    if parts.len() != 2 {
+        error_response(
+            stream,
+            "PAUSE requires a database and a user, e.g. PAUSE my_db,my_user",
+        )
+        .await
+    } else {
+        let database = parts[0];
+        let user = parts[1];
+
+        match get_pool(database, user) {
+            Some(pool) => {
+                pool.pause();
+
+                let mut res = BytesMut::new();
+
+                res.put(command_complete(&format!("PAUSE {},{}", database, user)));
+
+                // ReadyForQuery
+                res.put_u8(b'Z');
+                res.put_i32(5);
+                res.put_u8(b'I');
+
+                write_all_half(stream, &res).await
+            }
+
+            None => {
+                error_response(
+                    stream,
+                    &format!(
+                        "No pool configured for database: {}, user: {}",
+                        database, user
+                    ),
+                )
+                .await
+            }
+        }
+    }
+}
+
+/// Resume a pool. Queries are allowed again.
+async fn resume<T>(stream: &mut T, query: &str) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let parts: Vec<&str> = query.split(",").map(|part| part.trim()).collect();
+
+    if parts.len() != 2 {
+        error_response(
+            stream,
+            "RESUME requires a database and a user, e.g. RESUME my_db,my_user",
+        )
+        .await
+    } else {
+        let database = parts[0];
+        let user = parts[1];
+
+        match get_pool(database, user) {
+            Some(pool) => {
+                pool.resume();
+
+                let mut res = BytesMut::new();
+
+                res.put(command_complete(&format!("RESUME {},{}", database, user)));
+
+                // ReadyForQuery
+                res.put_u8(b'Z');
+                res.put_i32(5);
+                res.put_u8(b'I');
+
+                write_all_half(stream, &res).await
+            }
+
+            None => {
+                error_response(
+                    stream,
+                    &format!(
+                        "No pool configured for database: {}, user: {}",
+                        database, user
+                    ),
+                )
+                .await
+            }
+        }
+    }
+}
+
+/// Send response packets for shutdown.
+async fn shutdown<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let mut res = BytesMut::new();
+
+    res.put(row_description(&vec![("success", DataType::Text)]));
+
+    let mut shutdown_success = "t";
+
+    let pid = std::process::id();
+    if signal::kill(Pid::from_raw(pid.try_into().unwrap()), Signal::SIGINT).is_err() {
+        error!("Unable to send SIGINT to PID: {}", pid);
+        shutdown_success = "f";
+    }
+
+    res.put(data_row(&vec![shutdown_success.to_string()]));
+
+    res.put(command_complete("SHUTDOWN"));
+
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
+/// Show Users.
+async fn show_users<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let mut res = BytesMut::new();
+
+    res.put(row_description(&vec![
+        ("name", DataType::Text),
+        ("pool_mode", DataType::Text),
+    ]));
+
+    for (user_pool, pool) in get_all_pools() {
+        let pool_config = &pool.settings;
+        res.put(data_row(&vec![
+            user_pool.user.clone(),
+            pool_config.pool_mode.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
 }

src/config.rs

@@ -2,14 +2,15 @@
 use arc_swap::ArcSwap;
 use log::{error, info};
 use once_cell::sync::Lazy;
+use regex::Regex;
 use serde_derive::{Deserialize, Serialize};
+use std::collections::hash_map::DefaultHasher;
 use std::collections::{BTreeMap, HashMap, HashSet};
-use std::hash::Hash;
+use std::hash::{Hash, Hasher};
 use std::path::Path;
 use std::sync::Arc;
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
-use toml;
 
 use crate::errors::Error;
 use crate::pool::{ClientServerMap, ConnectionPool};
@@ -28,6 +29,8 @@ pub enum Role {
     Primary,
     #[serde(alias = "replica", alias = "Replica")]
     Replica,
+    #[serde(alias = "mirror", alias = "Mirror")]
+    Mirror,
 }
 
 impl ToString for Role {
@@ -35,6 +38,7 @@ impl ToString for Role {
         match *self {
             Role::Primary => "primary".to_string(),
             Role::Replica => "replica".to_string(),
+            Role::Mirror => "mirror".to_string(),
         }
     }
 }
@@ -89,6 +93,9 @@ pub struct Address {
 
     /// The name of this pool (i.e. database name visible to the client).
     pub pool_name: String,
+
+    /// List of addresses to receive mirrored traffic.
+    pub mirrors: Vec<Address>,
 }
 
 impl Default for Address {
@@ -104,6 +111,7 @@ impl Default for Address {
             role: Role::Replica,
             username: String::from("username"),
             pool_name: String::from("pool_name"),
+            mirrors: Vec::new(),
         }
     }
 }
@@ -113,11 +121,14 @@ impl Address {
     pub fn name(&self) -> String {
         match self.role {
             Role::Primary => format!("{}_shard_{}_primary", self.pool_name, self.shard),
-
             Role::Replica => format!(
                 "{}_shard_{}_replica_{}",
                 self.pool_name, self.shard, self.replica_number
             ),
+            Role::Mirror => format!(
+                "{}_shard_{}_mirror_{}",
+                self.pool_name, self.shard, self.replica_number
+            ),
         }
     }
 }
@@ -158,6 +169,22 @@ pub struct General {
     #[serde(default = "General::default_connect_timeout")]
     pub connect_timeout: u64,
 
+    #[serde(default = "General::default_idle_timeout")]
+    pub idle_timeout: u64,
+
+    #[serde(default = "General::default_tcp_keepalives_idle")]
+    pub tcp_keepalives_idle: u64,
+    #[serde(default = "General::default_tcp_keepalives_count")]
+    pub tcp_keepalives_count: u32,
+    #[serde(default = "General::default_tcp_keepalives_interval")]
+    pub tcp_keepalives_interval: u64,
+
+    #[serde(default)] // False
+    pub log_client_connections: bool,
+
+    #[serde(default)] // False
+    pub log_client_disconnections: bool,
+
     #[serde(default = "General::default_shutdown_timeout")]
     pub shutdown_timeout: u64,
 
@@ -170,6 +197,9 @@ pub struct General {
     #[serde(default = "General::default_ban_time")]
     pub ban_time: i64,
 
+    #[serde(default = "General::default_worker_threads")]
+    pub worker_threads: usize,
+
     #[serde(default)] // False
     pub autoreload: bool,
 
@@ -192,6 +222,25 @@ impl General {
         1000
     }
 
+    // These keepalive defaults should detect a dead connection within 30 seconds.
+    // Tokio defaults to disabling keepalives which keeps dead connections around indefinitely.
+    // This can lead to permenant server pool exhaustion
+    pub fn default_tcp_keepalives_idle() -> u64 {
+        5 // 5 seconds
+    }
+
+    pub fn default_tcp_keepalives_count() -> u32 {
+        5 // 5 time
+    }
+
+    pub fn default_tcp_keepalives_interval() -> u64 {
+        5 // 5 seconds
+    }
+
+    pub fn default_idle_timeout() -> u64 {
+        60000 // 10 minutes
+    }
+
     pub fn default_shutdown_timeout() -> u64 {
         60000
     }
@@ -207,6 +256,10 @@ impl General {
     pub fn default_ban_time() -> i64 {
         60
     }
+
+    pub fn default_worker_threads() -> usize {
+        4
+    }
 }
 
 impl Default for General {
@@ -217,10 +270,17 @@ impl Default for General {
             enable_prometheus_exporter: Some(false),
             prometheus_exporter_port: 9930,
             connect_timeout: General::default_connect_timeout(),
+            idle_timeout: General::default_idle_timeout(),
             shutdown_timeout: Self::default_shutdown_timeout(),
             healthcheck_timeout: Self::default_healthcheck_timeout(),
             healthcheck_delay: Self::default_healthcheck_delay(),
             ban_time: Self::default_ban_time(),
+            worker_threads: Self::default_worker_threads(),
+            tcp_keepalives_idle: Self::default_tcp_keepalives_idle(),
+            tcp_keepalives_count: Self::default_tcp_keepalives_count(),
+            tcp_keepalives_interval: Self::default_tcp_keepalives_interval(),
+            log_client_connections: false,
+            log_client_disconnections: false,
             autoreload: false,
             tls_certificate: None,
             tls_private_key: None,
@@ -241,7 +301,6 @@ pub enum PoolMode {
     #[serde(alias = "session", alias = "Session")]
     Session,
 }
-
 impl ToString for PoolMode {
     fn to_string(&self) -> String {
         match *self {
@@ -251,11 +310,33 @@ impl ToString for PoolMode {
     }
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Copy, Hash)]
+pub enum LoadBalancingMode {
+    #[serde(alias = "random", alias = "Random")]
+    Random,
+
+    #[serde(alias = "loc", alias = "LOC", alias = "least_outstanding_connections")]
+    LeastOutstandingConnections,
+}
+impl ToString for LoadBalancingMode {
+    fn to_string(&self) -> String {
+        match *self {
+            LoadBalancingMode::Random => "random".to_string(),
+            LoadBalancingMode::LeastOutstandingConnections => {
+                "least_outstanding_connections".to_string()
+            }
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
 pub struct Pool {
     #[serde(default = "Pool::default_pool_mode")]
     pub pool_mode: PoolMode,
 
+    #[serde(default = "Pool::default_load_balancing_mode")]
+    pub load_balancing_mode: LoadBalancingMode,
+
     pub default_role: String,
 
     #[serde(default)] // False
@@ -266,17 +347,44 @@ pub struct Pool {
 
     pub connect_timeout: Option<u64>,
 
+    pub idle_timeout: Option<u64>,
+
     pub sharding_function: ShardingFunction,
+
+    #[serde(default = "Pool::default_automatic_sharding_key")]
+    pub automatic_sharding_key: Option<String>,
+
+    pub sharding_key_regex: Option<String>,
+    pub shard_id_regex: Option<String>,
+    pub regex_search_limit: Option<usize>,
+
     pub shards: BTreeMap<String, Shard>,
     pub users: BTreeMap<String, User>,
+    // Note, don't put simple fields below these configs. There's a compatability issue with TOML that makes it
+    // incompatible to have simple fields in TOML after complex objects. See
+    // https://users.rust-lang.org/t/why-toml-to-string-get-error-valueaftertable/85903
 }
 
 impl Pool {
+    pub fn hash_value(&self) -> u64 {
+        let mut s = DefaultHasher::new();
+        self.hash(&mut s);
+        s.finish()
+    }
+
     pub fn default_pool_mode() -> PoolMode {
         PoolMode::Transaction
     }
 
-    pub fn validate(&self) -> Result<(), Error> {
+    pub fn default_load_balancing_mode() -> LoadBalancingMode {
+        LoadBalancingMode::Random
+    }
+
+    pub fn default_automatic_sharding_key() -> Option<String> {
+        None
+    }
+
+    pub fn validate(&mut self) -> Result<(), Error> {
         match self.default_role.as_ref() {
             "any" => (),
             "primary" => (),
@@ -304,6 +412,37 @@ impl Pool {
             shard.validate()?;
         }
 
+        for (option, name) in [
+            (&self.shard_id_regex, "shard_id_regex"),
+            (&self.sharding_key_regex, "sharding_key_regex"),
+        ] {
+            if let Some(regex) = option {
+                if let Err(parse_err) = Regex::new(regex.as_str()) {
+                    error!("{} is not a valid Regex: {}", name, parse_err);
+                    return Err(Error::BadConfig);
+                }
+            }
+        }
+
+        self.automatic_sharding_key = match &self.automatic_sharding_key {
+            Some(key) => {
+                // No quotes in the key so we don't have to compare quoted
+                // to unquoted idents.
+                let key = key.replace("\"", "");
+
+                if key.split(".").count() != 2 {
+                    error!(
+                        "automatic_sharding_key '{}' must be fully qualified, e.g. t.{}`",
+                        key, key
+                    );
+                    return Err(Error::BadConfig);
+                }
+
+                Some(key)
+            }
+            None => None,
+        };
+
         Ok(())
     }
 }
@@ -312,13 +451,19 @@ impl Default for Pool {
     fn default() -> Pool {
         Pool {
             pool_mode: Self::default_pool_mode(),
+            load_balancing_mode: Self::default_load_balancing_mode(),
             shards: BTreeMap::from([(String::from("1"), Shard::default())]),
             users: BTreeMap::default(),
             default_role: String::from("any"),
             query_parser_enabled: false,
             primary_reads_enabled: false,
             sharding_function: ShardingFunction::PgBigintHash,
+            automatic_sharding_key: None,
             connect_timeout: None,
+            idle_timeout: None,
+            sharding_key_regex: None,
+            shard_id_regex: None,
+            regex_search_limit: Some(1000),
         }
     }
 }
@@ -330,10 +475,18 @@ pub struct ServerConfig {
     pub role: Role,
 }
 
+#[derive(Clone, PartialEq, Serialize, Deserialize, Debug, Hash, Eq)]
+pub struct MirrorServerConfig {
+    pub host: String,
+    pub port: u16,
+    pub mirroring_target_index: usize,
+}
+
 /// Shard configuration.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Hash, Eq)]
 pub struct Shard {
     pub database: String,
+    pub mirrors: Option<Vec<MirrorServerConfig>>,
     pub servers: Vec<ServerConfig>,
 }
 
@@ -344,7 +497,7 @@ impl Shard {
         let mut dup_check = HashSet::new();
         let mut primary_count = 0;
 
-        if self.servers.len() == 0 {
+        if self.servers.is_empty() {
             error!("Shard {} has no servers configured", self.database);
             return Err(Error::BadConfig);
         }
@@ -353,10 +506,9 @@ impl Shard {
             dup_check.insert(server);
 
             // Check that we define only zero or one primary.
-            match server.role {
-                Role::Primary => primary_count += 1,
-                _ => (),
-            };
+            if server.role == Role::Primary {
+                primary_count += 1
+            }
         }
 
         if primary_count > 1 {
@@ -384,6 +536,7 @@ impl Default for Shard {
                 port: 5432,
                 role: Role::Primary,
             }],
+            mirrors: None,
             database: String::from("postgres"),
         }
     }
@@ -437,6 +590,10 @@ impl From<&Config> for std::collections::HashMap<String, String> {
                         format!("pools.{}.pool_mode", pool_name),
                         pool.pool_mode.to_string(),
                     ),
+                    (
+                        format!("pools.{}.load_balancing_mode", pool_name),
+                        pool.load_balancing_mode.to_string(),
+                    ),
                     (
                         format!("pools.{}.primary_reads_enabled", pool_name),
                         pool.primary_reads_enabled.to_string(),
@@ -481,6 +638,10 @@ impl From<&Config> for std::collections::HashMap<String, String> {
                 "connect_timeout".to_string(),
                 config.general.connect_timeout.to_string(),
             ),
+            (
+                "idle_timeout".to_string(),
+                config.general.idle_timeout.to_string(),
+            ),
             (
                 "healthcheck_timeout".to_string(),
                 config.general.healthcheck_timeout.to_string(),
@@ -505,11 +666,21 @@ impl Config {
     /// Print current configuration.
     pub fn show(&self) {
         info!("Ban time: {}s", self.general.ban_time);
+        info!("Worker threads: {}", self.general.worker_threads);
         info!(
             "Healthcheck timeout: {}ms",
             self.general.healthcheck_timeout
         );
         info!("Connection timeout: {}ms", self.general.connect_timeout);
+        info!("Idle timeout: {}ms", self.general.idle_timeout);
+        info!(
+            "Log client connections: {}",
+            self.general.log_client_connections
+        );
+        info!(
+            "Log client disconnections: {}",
+            self.general.log_client_disconnections
+        );
         info!("Shutdown timeout: {}ms", self.general.shutdown_timeout);
         info!("Healthcheck delay: {}ms", self.general.healthcheck_delay);
         match self.general.tls_certificate.clone() {
@@ -547,6 +718,10 @@ impl Config {
                 "[pool: {}] Pool mode: {:?}",
                 pool_name, pool_config.pool_mode
             );
+            info!(
+                "[pool: {}] Load Balancing mode: {:?}",
+                pool_name, pool_config.load_balancing_mode
+            );
             let connect_timeout = match pool_config.connect_timeout {
                 Some(connect_timeout) => connect_timeout,
                 None => self.general.connect_timeout,
@@ -555,6 +730,11 @@ impl Config {
                 "[pool: {}] Connection timeout: {}ms",
                 pool_name, connect_timeout
             );
+            let idle_timeout = match pool_config.idle_timeout {
+                Some(idle_timeout) => idle_timeout,
+                None => self.general.idle_timeout,
+            };
+            info!("[pool: {}] Idle timeout: {}ms", pool_name, idle_timeout);
             info!(
                 "[pool: {}] Sharding function: {}",
                 pool_name,
@@ -596,22 +776,17 @@ impl Config {
         // Validate TLS!
         match self.general.tls_certificate.clone() {
             Some(tls_certificate) => {
-                match load_certs(&Path::new(&tls_certificate)) {
+                match load_certs(Path::new(&tls_certificate)) {
                     Ok(_) => {
                         // Cert is okay, but what about the private key?
                         match self.general.tls_private_key.clone() {
-                            Some(tls_private_key) => {
-                                match load_keys(&Path::new(&tls_private_key)) {
-                                    Ok(_) => (),
-                                    Err(err) => {
-                                        error!(
-                                            "tls_private_key is incorrectly configured: {:?}",
-                                            err
-                                        );
-                                        return Err(Error::BadConfig);
-                                    }
+                            Some(tls_private_key) => match load_keys(Path::new(&tls_private_key)) {
+                                Ok(_) => (),
+                                Err(err) => {
+                                    error!("tls_private_key is incorrectly configured: {:?}", err);
+                                    return Err(Error::BadConfig);
                                 }
-                            }
+                            },
 
                             None => {
                                 error!("tls_certificate is set, but the tls_private_key is not");
@@ -629,7 +804,7 @@ impl Config {
             None => (),
         };
 
-        for (_, pool) in &mut self.pools {
+        for pool in self.pools.values_mut() {
             pool.validate()?;
         }
 
@@ -714,8 +889,10 @@ mod test {
         assert_eq!(get_config().path, "pgcat.toml".to_string());
 
         assert_eq!(get_config().general.ban_time, 60);
+        assert_eq!(get_config().general.idle_timeout, 30000);
         assert_eq!(get_config().pools.len(), 2);
         assert_eq!(get_config().pools["sharded_db"].shards.len(), 3);
+        assert_eq!(get_config().pools["sharded_db"].idle_timeout, Some(40000));
         assert_eq!(get_config().pools["simple_db"].shards.len(), 1);
         assert_eq!(get_config().pools["sharded_db"].users.len(), 2);
         assert_eq!(get_config().pools["simple_db"].users.len(), 1);

src/errors.rs

@@ -3,14 +3,16 @@
 /// Various errors.
 #[derive(Debug, PartialEq)]
 pub enum Error {
-    SocketError,
+    SocketError(String),
     ClientBadStartup,
-    ProtocolSyncError,
+    ProtocolSyncError(String),
+    BadQuery(String),
     ServerError,
     BadConfig,
     AllServersDown,
-    ClientError,
+    ClientError(String),
     TlsError,
     StatementTimeout,
     ShuttingDown,
+    ParseBytesError(String),
 }

src/lib.rs

@@ -2,6 +2,8 @@ pub mod config;
 pub mod constants;
 pub mod errors;
 pub mod messages;
+pub mod mirrors;
+pub mod multi_logger;
 pub mod pool;
 pub mod scram;
 pub mod server;

src/main.rs

@@ -37,14 +37,22 @@ extern crate tokio;
 extern crate tokio_rustls;
 extern crate toml;
 
-use log::{error, info, warn};
+#[cfg(not(target_env = "msvc"))]
+use jemallocator::Jemalloc;
+
+#[cfg(not(target_env = "msvc"))]
+#[global_allocator]
+static GLOBAL: Jemalloc = Jemalloc;
+
+use log::{debug, error, info, warn};
 use parking_lot::Mutex;
 use pgcat::format_duration;
 use tokio::net::TcpListener;
-use tokio::{
-    signal::unix::{signal as unix_signal, SignalKind},
-    sync::mpsc,
-};
+#[cfg(not(windows))]
+use tokio::signal::unix::{signal as unix_signal, SignalKind};
+#[cfg(windows)]
+use tokio::signal::windows as win_signal;
+use tokio::{runtime::Builder, sync::mpsc};
 
 use std::collections::HashMap;
 use std::net::SocketAddr;
@@ -58,6 +66,8 @@ mod config;
 mod constants;
 mod errors;
 mod messages;
+mod mirrors;
+mod multi_logger;
 mod pool;
 mod prometheus;
 mod query_router;
@@ -68,14 +78,12 @@ mod stats;
 mod tls;
 
 use crate::config::{get_config, reload_config, VERSION};
-use crate::errors::Error;
 use crate::pool::{ClientServerMap, ConnectionPool};
 use crate::prometheus::start_metric_server;
 use crate::stats::{Collector, Reporter, REPORTER};
 
-#[tokio::main(worker_threads = 4)]
-async fn main() {
-    env_logger::builder().format_timestamp_micros().init();
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    multi_logger::MultiLogger::init().unwrap();
 
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
 
@@ -92,213 +100,249 @@ async fn main() {
         String::from("pgcat.toml")
     };
 
-    match config::parse(&config_file).await {
-        Ok(_) => (),
-        Err(err) => {
-            error!("Config parse error: {:?}", err);
-            std::process::exit(exitcode::CONFIG);
-        }
-    };
+    // Create a transient runtime for loading the config for the first time.
+    {
+        let runtime = Builder::new_multi_thread().worker_threads(1).build()?;
 
-    let config = get_config();
-
-    if let Some(true) = config.general.enable_prometheus_exporter {
-        let http_addr_str = format!(
-            "{}:{}",
-            config.general.host, config.general.prometheus_exporter_port
-        );
-        let http_addr = match SocketAddr::from_str(&http_addr_str) {
-            Ok(addr) => addr,
-            Err(err) => {
-                error!("Invalid http address: {}", err);
-                std::process::exit(exitcode::CONFIG);
-            }
-        };
-        tokio::task::spawn(async move {
-            start_metric_server(http_addr).await;
+        runtime.block_on(async {
+            match config::parse(&config_file).await {
+                Ok(_) => (),
+                Err(err) => {
+                    error!("Config parse error: {:?}", err);
+                    std::process::exit(exitcode::CONFIG);
+                }
+            };
         });
     }
 
-    let addr = format!("{}:{}", config.general.host, config.general.port);
+    let config = get_config();
 
-    let listener = match TcpListener::bind(&addr).await {
-        Ok(sock) => sock,
-        Err(err) => {
-            error!("Listener socket error: {:?}", err);
-            std::process::exit(exitcode::CONFIG);
+    // Create the runtime now we know required worker_threads.
+    let runtime = Builder::new_multi_thread()
+        .worker_threads(config.general.worker_threads)
+        .enable_all()
+        .build()?;
+
+    runtime.block_on(async move {
+
+        if let Some(true) = config.general.enable_prometheus_exporter {
+            let http_addr_str = format!(
+                "{}:{}",
+                config.general.host, config.general.prometheus_exporter_port
+            );
+
+            let http_addr = match SocketAddr::from_str(&http_addr_str) {
+                Ok(addr) => addr,
+                Err(err) => {
+                    error!("Invalid http address: {}", err);
+                    std::process::exit(exitcode::CONFIG);
+                }
+            };
+
+            tokio::task::spawn(async move {
+                start_metric_server(http_addr).await;
+            });
         }
-    };
 
-    info!("Running on {}", addr);
+        let addr = format!("{}:{}", config.general.host, config.general.port);
 
-    config.show();
+        let listener = match TcpListener::bind(&addr).await {
+            Ok(sock) => sock,
+            Err(err) => {
+                error!("Listener socket error: {:?}", err);
+                std::process::exit(exitcode::CONFIG);
+            }
+        };
 
-    // Tracks which client is connected to which server for query cancellation.
-    let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
+        info!("Running on {}", addr);
 
-    // Statistics reporting.
-    let (stats_tx, stats_rx) = mpsc::channel(100_000);
-    REPORTER.store(Arc::new(Reporter::new(stats_tx.clone())));
+        config.show();
 
-    // Connection pool that allows to query all shards and replicas.
-    match ConnectionPool::from_config(client_server_map.clone()).await {
-        Ok(_) => (),
-        Err(err) => {
-            error!("Pool error: {:?}", err);
-            std::process::exit(exitcode::CONFIG);
-        }
-    };
+        // Tracks which client is connected to which server for query cancellation.
+        let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
 
-    tokio::task::spawn(async move {
-        let mut stats_collector = Collector::new(stats_rx, stats_tx.clone());
-        stats_collector.collect().await;
-    });
+        // Statistics reporting.
+        let (stats_tx, stats_rx) = mpsc::channel(500_000);
+        REPORTER.store(Arc::new(Reporter::new(stats_tx.clone())));
 
-    info!("Config autoreloader: {}", config.general.autoreload);
+        // Connection pool that allows to query all shards and replicas.
+        match ConnectionPool::from_config(client_server_map.clone()).await {
+            Ok(_) => (),
+            Err(err) => {
+                error!("Pool error: {:?}", err);
+                std::process::exit(exitcode::CONFIG);
+            }
+        };
 
-    let mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000));
-    let autoreload_client_server_map = client_server_map.clone();
-    tokio::task::spawn(async move {
-        loop {
-            autoreload_interval.tick().await;
-            if config.general.autoreload {
-                info!("Automatically reloading config");
+        tokio::task::spawn(async move {
+            let mut stats_collector = Collector::new(stats_rx, stats_tx.clone());
+            stats_collector.collect().await;
+        });
 
-                match reload_config(autoreload_client_server_map.clone()).await {
-                    Ok(changed) => {
+        info!("Config autoreloader: {}", config.general.autoreload);
+
+        let mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000));
+        let autoreload_client_server_map = client_server_map.clone();
+
+        tokio::task::spawn(async move {
+            loop {
+                autoreload_interval.tick().await;
+                if config.general.autoreload {
+                    info!("Automatically reloading config");
+
+                    if let Ok(changed) = reload_config(autoreload_client_server_map.clone()).await {
                         if changed {
                             get_config().show()
                         }
-                    }
-                    Err(_) => (),
-                };
+                    };
+                }
             }
-        }
-    });
+        });
 
-    let mut term_signal = unix_signal(SignalKind::terminate()).unwrap();
-    let mut interrupt_signal = unix_signal(SignalKind::interrupt()).unwrap();
-    let mut sighup_signal = unix_signal(SignalKind::hangup()).unwrap();
-    let (shutdown_tx, _) = broadcast::channel::<()>(1);
-    let (drain_tx, mut drain_rx) = mpsc::channel::<i32>(2048);
-    let (exit_tx, mut exit_rx) = mpsc::channel::<()>(1);
+        #[cfg(windows)]
+        let mut term_signal = win_signal::ctrl_close().unwrap();
+        #[cfg(windows)]
+        let mut interrupt_signal = win_signal::ctrl_c().unwrap();
+        #[cfg(windows)]
+        let mut sighup_signal = win_signal::ctrl_shutdown().unwrap();
 
-    info!("Waiting for clients");
+        #[cfg(not(windows))]
+        let mut term_signal = unix_signal(SignalKind::terminate()).unwrap();
+        #[cfg(not(windows))]
+        let mut interrupt_signal = unix_signal(SignalKind::interrupt()).unwrap();
+        #[cfg(not(windows))]
+        let mut sighup_signal = unix_signal(SignalKind::hangup()).unwrap();
+        let (shutdown_tx, _) = broadcast::channel::<()>(1);
+        let (drain_tx, mut drain_rx) = mpsc::channel::<i32>(2048);
+        let (exit_tx, mut exit_rx) = mpsc::channel::<()>(1);
+        let mut admin_only = false;
+        let mut total_clients = 0;
 
-    let mut admin_only = false;
-    let mut total_clients = 0;
+        info!("Waiting for clients");
 
-    loop {
-        tokio::select! {
-            // Reload config:
-            // kill -SIGHUP $(pgrep pgcat)
-            _ = sighup_signal.recv() => {
-                info!("Reloading config");
+        loop {
+            tokio::select! {
+                // Reload config:
+                // kill -SIGHUP $(pgrep pgcat)
+                _ = sighup_signal.recv() => {
+                    info!("Reloading config");
 
-                match reload_config(client_server_map.clone()).await {
-                    Ok(_) => (),
-                    Err(_) => (),
-                };
+                    _ = reload_config(client_server_map.clone()).await;
 
-                get_config().show();
-            },
+                    get_config().show();
+                },
 
-            // Initiate graceful shutdown sequence on sig int
-            _ = interrupt_signal.recv() => {
-                info!("Got SIGINT, waiting for client connection drain now");
-                admin_only = true;
+                // Initiate graceful shutdown sequence on sig int
+                _ = interrupt_signal.recv() => {
+                    info!("Got SIGINT");
 
-                // Broadcast that client tasks need to finish
-                let _ = shutdown_tx.send(());
-                let exit_tx = exit_tx.clone();
-                let _ = drain_tx.send(0).await;
-
-                tokio::task::spawn(async move {
-                    let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
-
-                    // First tick fires immediately.
-                    interval.tick().await;
-
-                    // Second one in the interval time.
-                    interval.tick().await;
-
-                    // We're done waiting.
-                    error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
-
-                    let _ = exit_tx.send(()).await;
-                });
-            },
-
-            _ = term_signal.recv() => {
-                info!("Got SIGTERM, closing with {} clients active", total_clients);
-                break;
-            },
-
-            new_client = listener.accept() => {
-                let (socket, addr) = match new_client {
-                    Ok((socket, addr)) => (socket, addr),
-                    Err(err) => {
-                        error!("{:?}", err);
+                    // Don't want this to happen more than once
+                    if admin_only {
                         continue;
                     }
-                };
 
-                let shutdown_rx = shutdown_tx.subscribe();
-                let drain_tx = drain_tx.clone();
-                let client_server_map = client_server_map.clone();
+                    admin_only = true;
 
-                tokio::task::spawn(async move {
-                    let start = chrono::offset::Utc::now().naive_utc();
+                    // Broadcast that client tasks need to finish
+                    let _ = shutdown_tx.send(());
+                    let exit_tx = exit_tx.clone();
+                    let _ = drain_tx.send(0).await;
 
-                    match client::client_entrypoint(
-                        socket,
-                        client_server_map,
-                        shutdown_rx,
-                        drain_tx,
-                        admin_only,
-                    )
-                    .await
-                    {
-                        Ok(()) => {
+                    tokio::task::spawn(async move {
+                        let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
 
-                            let duration = chrono::offset::Utc::now().naive_utc() - start;
+                        // First tick fires immediately.
+                        interval.tick().await;
 
-                            info!(
-                                "Client {:?} disconnected, session duration: {}",
-                                addr,
-                                format_duration(&duration)
-                            );
-                        }
+                        // Second one in the interval time.
+                        interval.tick().await;
 
+                        // We're done waiting.
+                        error!("Graceful shutdown timed out. {} active clients being closed", total_clients);
+
+                        let _ = exit_tx.send(()).await;
+                    });
+                },
+
+                _ = term_signal.recv() => {
+                    info!("Got SIGTERM, closing with {} clients active", total_clients);
+                    break;
+                },
+
+                new_client = listener.accept() => {
+                    let (socket, addr) = match new_client {
+                        Ok((socket, addr)) => (socket, addr),
                         Err(err) => {
-                            match err {
-                                // Don't count the clients we rejected.
-                                Error::ShuttingDown => (),
-                                _ => {
-                                    // drain_tx.send(-1).await.unwrap();
+                            error!("{:?}", err);
+                            continue;
+                        }
+                    };
+
+                    let shutdown_rx = shutdown_tx.subscribe();
+                    let drain_tx = drain_tx.clone();
+                    let client_server_map = client_server_map.clone();
+
+                    let tls_certificate = config.general.tls_certificate.clone();
+
+                    tokio::task::spawn(async move {
+                        let start = chrono::offset::Utc::now().naive_utc();
+
+                        match client::client_entrypoint(
+                            socket,
+                            client_server_map,
+                            shutdown_rx,
+                            drain_tx,
+                            admin_only,
+                            tls_certificate.clone(),
+                            config.general.log_client_connections,
+                        )
+                        .await
+                        {
+                            Ok(()) => {
+                                let duration = chrono::offset::Utc::now().naive_utc() - start;
+
+                                if config.general.log_client_disconnections {
+                                    info!(
+                                        "Client {:?} disconnected, session duration: {}",
+                                        addr,
+                                        format_duration(&duration)
+                                    );
+                                } else {
+                                    debug!(
+                                        "Client {:?} disconnected, session duration: {}",
+                                        addr,
+                                        format_duration(&duration)
+                                    );
                                 }
                             }
 
-                            warn!("Client disconnected with error {:?}", err);
-                        }
-                    };
-                });
-            }
+                            Err(err) => {
+                                match err {
+                                    errors::Error::ClientBadStartup => debug!("Client disconnected with error {:?}", err),
+                                    _ => warn!("Client disconnected with error {:?}", err),
+                                }
 
-            _ = exit_rx.recv() => {
-                break;
-            }
+                            }
+                        };
+                    });
+                }
 
-            client_ping = drain_rx.recv() => {
-                let client_ping = client_ping.unwrap();
-                total_clients += client_ping;
+                _ = exit_rx.recv() => {
+                    break;
+                }
 
-                if total_clients == 0 && admin_only {
-                    let _ = exit_tx.send(()).await;
+                client_ping = drain_rx.recv() => {
+                    let client_ping = client_ping.unwrap();
+                    total_clients += client_ping;
+
+                    if total_clients == 0 && admin_only {
+                        let _ = exit_tx.send(()).await;
+                    }
                 }
             }
         }
-    }
 
     info!("Shutting down...");
+    });
+    Ok(())
 }

src/messages.rs

@@ -1,13 +1,18 @@
 /// Helper functions to send one-off protocol messages
 /// and handle TcpStream (TCP socket).
 use bytes::{Buf, BufMut, BytesMut};
+use log::error;
 use md5::{Digest, Md5};
+use socket2::{SockRef, TcpKeepalive};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
 use tokio::net::TcpStream;
 
+use crate::config::get_config;
 use crate::errors::Error;
 use std::collections::HashMap;
+use std::io::{BufRead, Cursor};
 use std::mem;
+use std::time::Duration;
 
 /// Postgres data type mappings
 /// used in RowDescription ('T') message.
@@ -38,7 +43,7 @@ where
     auth_ok.put_i32(8);
     auth_ok.put_i32(0);
 
-    Ok(write_all(stream, auth_ok).await?)
+    write_all(stream, auth_ok).await
 }
 
 /// Generate md5 password challenge.
@@ -79,7 +84,7 @@ where
     key_data.put_i32(backend_id);
     key_data.put_i32(secret_key);
 
-    Ok(write_all(stream, key_data).await?)
+    write_all(stream, key_data).await
 }
 
 /// Construct a `Q`: Query message.
@@ -88,7 +93,7 @@ pub fn simple_query(query: &str) -> BytesMut {
     let query = format!("{}\0", query);
 
     res.put_i32(query.len() as i32 + 4);
-    res.put_slice(&query.as_bytes());
+    res.put_slice(query.as_bytes());
 
     res
 }
@@ -106,7 +111,7 @@ where
     bytes.put_i32(5);
     bytes.put_u8(b'I'); // Idle
 
-    Ok(write_all(stream, bytes).await?)
+    write_all(stream, bytes).await
 }
 
 /// Send the startup packet the server. We're pretending we're a Pg client.
@@ -118,12 +123,12 @@ pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Resu
 
     // User
     bytes.put(&b"user\0"[..]);
-    bytes.put_slice(&user.as_bytes());
+    bytes.put_slice(user.as_bytes());
     bytes.put_u8(0);
 
     // Database
     bytes.put(&b"database\0"[..]);
-    bytes.put_slice(&database.as_bytes());
+    bytes.put_slice(database.as_bytes());
     bytes.put_u8(0);
     bytes.put_u8(0); // Null terminator
 
@@ -136,7 +141,12 @@ pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Resu
 
     match stream.write_all(&startup).await {
         Ok(_) => Ok(()),
-        Err(_) => return Err(Error::SocketError),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error writing startup to server socket - Error: {:?}",
+                err
+            )))
+        }
     }
 }
 
@@ -155,7 +165,7 @@ pub fn parse_params(mut bytes: BytesMut) -> Result<HashMap<String, String>, Erro
             c = bytes.get_u8();
         }
 
-        if tmp.len() > 0 {
+        if !tmp.is_empty() {
             buf.push(tmp.clone());
             tmp.clear();
         }
@@ -234,7 +244,7 @@ where
     message.put_i32(password.len() as i32 + 4);
     message.put_slice(&password[..]);
 
-    Ok(write_all(stream, message).await?)
+    write_all(stream, message).await
 }
 
 /// Implements a response to our custom `SET SHARDING KEY`
@@ -254,7 +264,7 @@ where
     res.put_i32(len);
     res.put_slice(&set_complete[..]);
 
-    write_all_half(stream, res).await?;
+    write_all_half(stream, &res).await?;
     ready_for_query(stream).await
 }
 
@@ -292,7 +302,7 @@ where
 
     // The short error message.
     error.put_u8(b'M');
-    error.put_slice(&format!("{}\0", message).as_bytes());
+    error.put_slice(format!("{}\0", message).as_bytes());
 
     // No more fields follow.
     error.put_u8(0);
@@ -304,7 +314,7 @@ where
     res.put_i32(error.len() as i32 + 4);
     res.put(error);
 
-    Ok(write_all_half(stream, res).await?)
+    write_all_half(stream, &res).await
 }
 
 pub async fn wrong_password<S>(stream: &mut S, user: &str) -> Result<(), Error>
@@ -327,7 +337,7 @@ where
 
     // The short error message.
     error.put_u8(b'M');
-    error.put_slice(&format!("password authentication failed for user \"{}\"\0", user).as_bytes());
+    error.put_slice(format!("password authentication failed for user \"{}\"\0", user).as_bytes());
 
     // No more fields follow.
     error.put_u8(0);
@@ -366,7 +376,7 @@ where
     // CommandComplete
     res.put(command_complete("SELECT 1"));
 
-    write_all_half(stream, res).await?;
+    write_all_half(stream, &res).await?;
     ready_for_query(stream).await
 }
 
@@ -379,7 +389,7 @@ pub fn row_description(columns: &Vec<(&str, DataType)>) -> BytesMut {
 
     for (name, data_type) in columns {
         // Column name
-        row_desc.put_slice(&format!("{}\0", name).as_bytes());
+        row_desc.put_slice(format!("{}\0", name).as_bytes());
 
         // Doesn't belong to any table
         row_desc.put_i32(0);
@@ -423,7 +433,7 @@ pub fn data_row(row: &Vec<String>) -> BytesMut {
     for column in row {
         let column = column.as_bytes();
         data_row.put_i32(column.len() as i32);
-        data_row.put_slice(&column);
+        data_row.put_slice(column);
     }
 
     res.put_u8(b'D');
@@ -450,18 +460,28 @@ where
 {
     match stream.write_all(&buf).await {
         Ok(_) => Ok(()),
-        Err(_) => return Err(Error::SocketError),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error writing to socket - Error: {:?}",
+                err
+            )))
+        }
     }
 }
 
 /// Write all the data in the buffer to the TcpStream, write owned half (see mpsc).
-pub async fn write_all_half<S>(stream: &mut S, buf: BytesMut) -> Result<(), Error>
+pub async fn write_all_half<S>(stream: &mut S, buf: &BytesMut) -> Result<(), Error>
 where
     S: tokio::io::AsyncWrite + std::marker::Unpin,
 {
-    match stream.write_all(&buf).await {
+    match stream.write_all(buf).await {
         Ok(_) => Ok(()),
-        Err(_) => return Err(Error::SocketError),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error writing to socket - Error: {:?}",
+                err
+            )))
+        }
     }
 }
 
@@ -472,26 +492,51 @@ where
 {
     let code = match stream.read_u8().await {
         Ok(code) => code,
-        Err(_) => return Err(Error::SocketError),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error reading message code from socket - Error {:?}",
+                err
+            )))
+        }
     };
 
     let len = match stream.read_i32().await {
         Ok(len) => len,
-        Err(_) => return Err(Error::SocketError),
-    };
-
-    let mut buf = vec![0u8; len as usize - 4];
-
-    match stream.read_exact(&mut buf).await {
-        Ok(_) => (),
-        Err(_) => return Err(Error::SocketError),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error reading message len from socket - Code: {:?}, Error: {:?}",
+                code, err
+            )))
+        }
     };
 
     let mut bytes = BytesMut::with_capacity(len as usize + 1);
 
     bytes.put_u8(code);
     bytes.put_i32(len);
-    bytes.put_slice(&buf);
+
+    bytes.resize(bytes.len() + len as usize - mem::size_of::<i32>(), b'0');
+
+    let slice_start = mem::size_of::<u8>() + mem::size_of::<i32>();
+    let slice_end = slice_start + len as usize - mem::size_of::<i32>();
+
+    // Avoids a panic
+    if slice_end < slice_start {
+        return Err(Error::SocketError(format!(
+            "Error reading message from socket - Code: {:?} - Length {:?}, Error: {:?}",
+            code, len, "Unexpected length value for message"
+        )));
+    }
+
+    match stream.read_exact(&mut bytes[slice_start..slice_end]).await {
+        Ok(_) => (),
+        Err(err) => {
+            return Err(Error::SocketError(format!(
+                "Error reading message from socket - Code: {:?}, Error: {:?}",
+                code, err
+            )))
+        }
+    };
 
     Ok(bytes)
 }
@@ -510,5 +555,41 @@ pub fn server_parameter_message(key: &str, value: &str) -> BytesMut {
     server_info.put_slice(value.as_bytes());
     server_info.put_bytes(0, 1);
 
-    return server_info;
+    server_info
+}
+
+pub fn configure_socket(stream: &TcpStream) {
+    let sock_ref = SockRef::from(stream);
+    let conf = get_config();
+
+    match sock_ref.set_keepalive(true) {
+        Ok(_) => {
+            match sock_ref.set_tcp_keepalive(
+                &TcpKeepalive::new()
+                    .with_interval(Duration::from_secs(conf.general.tcp_keepalives_interval))
+                    .with_retries(conf.general.tcp_keepalives_count)
+                    .with_time(Duration::from_secs(conf.general.tcp_keepalives_idle)),
+            ) {
+                Ok(_) => (),
+                Err(err) => error!("Could not configure socket: {}", err),
+            }
+        }
+        Err(err) => error!("Could not configure socket: {}", err),
+    }
+}
+
+pub trait BytesMutReader {
+    fn read_string(&mut self) -> Result<String, Error>;
+}
+
+impl BytesMutReader for Cursor<&BytesMut> {
+    /// Should only be used when reading strings from the message protocol.
+    /// Can be used to read multiple strings from the same message which are separated by the null byte
+    fn read_string(&mut self) -> Result<String, Error> {
+        let mut buf = vec![];
+        match self.read_until(b'\0', &mut buf) {
+            Ok(_) => Ok(String::from_utf8_lossy(&buf[..buf.len() - 1]).to_string()),
+            Err(err) => return Err(Error::ParseBytesError(err.to_string())),
+        }
+    }
 }

src/mirrors.rs

@@ -0,0 +1,179 @@
+/// A mirrored PostgreSQL client.
+/// Packets arrive to us through a channel from the main client and we send them to the server.
+use bb8::Pool;
+use bytes::{Bytes, BytesMut};
+
+use crate::config::{get_config, Address, Role, User};
+use crate::pool::{ClientServerMap, ServerPool};
+use crate::stats::get_reporter;
+use log::{error, info, trace, warn};
+use tokio::sync::mpsc::{channel, Receiver, Sender};
+
+pub struct MirroredClient {
+    address: Address,
+    user: User,
+    database: String,
+    bytes_rx: Receiver<Bytes>,
+    disconnect_rx: Receiver<()>,
+}
+
+impl MirroredClient {
+    async fn create_pool(&self) -> Pool<ServerPool> {
+        let config = get_config();
+        let default = std::time::Duration::from_millis(10_000).as_millis() as u64;
+        let (connection_timeout, idle_timeout) = match config.pools.get(&self.address.pool_name) {
+            Some(cfg) => (
+                cfg.connect_timeout.unwrap_or(default),
+                cfg.idle_timeout.unwrap_or(default),
+            ),
+            None => (default, default),
+        };
+
+        let manager = ServerPool::new(
+            self.address.clone(),
+            self.user.clone(),
+            self.database.as_str(),
+            ClientServerMap::default(),
+            get_reporter(),
+        );
+
+        Pool::builder()
+            .max_size(1)
+            .connection_timeout(std::time::Duration::from_millis(connection_timeout))
+            .idle_timeout(Some(std::time::Duration::from_millis(idle_timeout)))
+            .test_on_check_out(false)
+            .build(manager)
+            .await
+            .unwrap()
+    }
+
+    pub fn start(mut self) {
+        tokio::spawn(async move {
+            let pool = self.create_pool().await;
+            let address = self.address.clone();
+            loop {
+                let mut server = match pool.get().await {
+                    Ok(server) => server,
+                    Err(err) => {
+                        error!(
+                            "Failed to get connection from pool, Discarding message {:?}, {:?}",
+                            err,
+                            address.clone()
+                        );
+                        continue;
+                    }
+                };
+
+                tokio::select! {
+                    // Exit channel events
+                    _ = self.disconnect_rx.recv() => {
+                        info!("Got mirror exit signal, exiting {:?}", address.clone());
+                        break;
+                    }
+
+                    // Incoming data from server (we read to clear the socket buffer and discard the data)
+                    recv_result = server.recv() => {
+                        match recv_result {
+                            Ok(message) => trace!("Received from mirror: {} {:?}", String::from_utf8_lossy(&message[..]), address.clone()),
+                            Err(err) => {
+                                server.mark_bad();
+                                error!("Failed to receive from mirror {:?} {:?}", err, address.clone());
+                            }
+                        }
+                    }
+
+                    // Messages to send to the server
+                    message = self.bytes_rx.recv() => {
+                        match message {
+                            Some(bytes) => {
+                                match server.send(&BytesMut::from(&bytes[..])).await {
+                                    Ok(_) => trace!("Sent to mirror: {} {:?}", String::from_utf8_lossy(&bytes[..]), address.clone()),
+                                    Err(err) => {
+                                        server.mark_bad();
+                                        error!("Failed to send to mirror, Discarding message {:?}, {:?}", err, address.clone())
+                                    }
+                                }
+                            }
+                            None => {
+                                info!("Mirror channel closed, exiting {:?}", address.clone());
+                                break;
+                            },
+                        }
+                    }
+                }
+            }
+        });
+    }
+}
+pub struct MirroringManager {
+    pub byte_senders: Vec<Sender<Bytes>>,
+    pub disconnect_senders: Vec<Sender<()>>,
+}
+impl MirroringManager {
+    pub fn from_addresses(
+        user: User,
+        database: String,
+        addresses: Vec<Address>,
+    ) -> MirroringManager {
+        let mut byte_senders: Vec<Sender<Bytes>> = vec![];
+        let mut exit_senders: Vec<Sender<()>> = vec![];
+
+        addresses.iter().for_each(|mirror| {
+            let (bytes_tx, bytes_rx) = channel::<Bytes>(10);
+            let (exit_tx, exit_rx) = channel::<()>(1);
+            let mut addr = mirror.clone();
+            addr.role = Role::Mirror;
+            let client = MirroredClient {
+                user: user.clone(),
+                database: database.to_owned(),
+                address: addr,
+                bytes_rx,
+                disconnect_rx: exit_rx,
+            };
+            exit_senders.push(exit_tx.clone());
+            byte_senders.push(bytes_tx.clone());
+            client.start();
+        });
+
+        Self {
+            byte_senders: byte_senders,
+            disconnect_senders: exit_senders,
+        }
+    }
+
+    pub fn send(self: &mut Self, bytes: &BytesMut) {
+        // We want to avoid performing an allocation if we won't be able to send the message
+        // There is a possibility of a race here where we check the capacity and then the channel is
+        // closed or the capacity is reduced to 0, but mirroring is best effort anyway
+        if self
+            .byte_senders
+            .iter()
+            .all(|sender| sender.capacity() == 0 || sender.is_closed())
+        {
+            return;
+        }
+        let immutable_bytes = bytes.clone().freeze();
+        self.byte_senders.iter_mut().for_each(|sender| {
+            match sender.try_send(immutable_bytes.clone()) {
+                Ok(_) => {}
+                Err(err) => {
+                    warn!("Failed to send bytes to a mirror channel {}", err);
+                }
+            }
+        });
+    }
+
+    pub fn disconnect(self: &mut Self) {
+        self.disconnect_senders
+            .iter_mut()
+            .for_each(|sender| match sender.try_send(()) {
+                Ok(_) => {}
+                Err(err) => {
+                    warn!(
+                        "Failed to send disconnect signal to a mirror channel {}",
+                        err
+                    );
+                }
+            });
+    }
+}

src/multi_logger.rs

@@ -0,0 +1,80 @@
+use log::{Level, Log, Metadata, Record, SetLoggerError};
+
+// This is a special kind of logger that allows sending logs to different
+// targets depending on the log level.
+//
+// By default, if nothing is set, it acts as a regular env_log logger,
+// it sends everything to standard error.
+//
+// If the Env variable `STDOUT_LOG` is defined, it will be used for
+// configuring the standard out logger.
+//
+// The behavior is:
+//   - If it is an error, the message is written to standard error.
+//   - If it is not, and it matches the log level of the standard output logger (`STDOUT_LOG` env var), it will be send to standard output.
+//   - If the above is not true, it is sent to the stderr logger that will log it or not depending on the value
+//     of the RUST_LOG env var.
+//
+// So to summarize, if no `STDOUT_LOG` env var is present, the logger is the default logger. If `STDOUT_LOG` is set, everything
+// but errors, that matches the log level set in the `STDOUT_LOG` env var is sent to stdout. You can have also some esoteric configuration
+// where you set `RUST_LOG=debug` and `STDOUT_LOG=info`, in here, erros will go to stderr, warns and infos to stdout and debugs to stderr.
+//
+pub struct MultiLogger {
+    stderr_logger: env_logger::Logger,
+    stdout_logger: env_logger::Logger,
+}
+
+impl MultiLogger {
+    fn new() -> Self {
+        let stderr_logger = env_logger::builder().format_timestamp_micros().build();
+        let stdout_logger = env_logger::Builder::from_env("STDOUT_LOG")
+            .format_timestamp_micros()
+            .target(env_logger::Target::Stdout)
+            .build();
+
+        Self {
+            stderr_logger,
+            stdout_logger,
+        }
+    }
+
+    pub fn init() -> Result<(), SetLoggerError> {
+        let logger = Self::new();
+
+        log::set_max_level(logger.stderr_logger.filter());
+        log::set_boxed_logger(Box::new(logger))
+    }
+}
+
+impl Log for MultiLogger {
+    fn enabled(&self, metadata: &Metadata) -> bool {
+        self.stderr_logger.enabled(metadata) && self.stdout_logger.enabled(metadata)
+    }
+
+    fn log(&self, record: &Record) {
+        if record.level() == Level::Error {
+            self.stderr_logger.log(record);
+        } else {
+            if self.stdout_logger.matches(record) {
+                self.stdout_logger.log(record);
+            } else {
+                self.stderr_logger.log(record);
+            }
+        }
+    }
+
+    fn flush(&self) {
+        self.stderr_logger.flush();
+        self.stdout_logger.flush();
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_init() {
+        MultiLogger::init().unwrap();
+    }
+}

src/pool.rs

@@ -1,18 +1,23 @@
 use arc_swap::ArcSwap;
 use async_trait::async_trait;
 use bb8::{ManageConnection, Pool, PooledConnection};
-use bytes::BytesMut;
+use bytes::{BufMut, BytesMut};
 use chrono::naive::NaiveDateTime;
 use log::{debug, error, info, warn};
 use once_cell::sync::Lazy;
 use parking_lot::{Mutex, RwLock};
 use rand::seq::SliceRandom;
 use rand::thread_rng;
-use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
+use regex::Regex;
+use std::collections::HashMap;
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
 use std::time::Instant;
+use tokio::sync::Notify;
 
-use crate::config::{get_config, Address, PoolMode, Role, User};
+use crate::config::{get_config, Address, General, LoadBalancingMode, PoolMode, Role, User};
 use crate::errors::Error;
 
 use crate::server::Server;
@@ -24,7 +29,7 @@ pub type SecretKey = i32;
 pub type ServerHost = String;
 pub type ServerPort = u16;
 
-pub type BanList = Arc<RwLock<Vec<HashMap<Address, NaiveDateTime>>>>;
+pub type BanList = Arc<RwLock<Vec<HashMap<Address, (BanReason, NaiveDateTime)>>>>;
 pub type ClientServerMap =
     Arc<Mutex<HashMap<(ProcessId, SecretKey), (ProcessId, SecretKey, ServerHost, ServerPort)>>>;
 pub type PoolMap = HashMap<PoolIdentifier, ConnectionPool>;
@@ -32,8 +37,17 @@ pub type PoolMap = HashMap<PoolIdentifier, ConnectionPool>;
 /// This is atomic and safe and read-optimized.
 /// The pool is recreated dynamically when the config is reloaded.
 pub static POOLS: Lazy<ArcSwap<PoolMap>> = Lazy::new(|| ArcSwap::from_pointee(HashMap::default()));
-static POOLS_HASH: Lazy<ArcSwap<HashSet<crate::config::Pool>>> =
-    Lazy::new(|| ArcSwap::from_pointee(HashSet::default()));
+
+// Reasons for banning a server.
+#[derive(Debug, PartialEq, Clone)]
+pub enum BanReason {
+    FailedHealthCheck,
+    MessageSendFailed,
+    MessageReceiveFailed,
+    FailedCheckout,
+    StatementTimeout,
+    AdminBan(i64),
+}
 
 /// An identifier for a PgCat pool,
 /// a database visible to clients.
@@ -56,12 +70,21 @@ impl PoolIdentifier {
     }
 }
 
+impl From<&Address> for PoolIdentifier {
+    fn from(address: &Address) -> PoolIdentifier {
+        PoolIdentifier::new(&address.database, &address.username)
+    }
+}
+
 /// Pool settings.
 #[derive(Clone, Debug)]
 pub struct PoolSettings {
     /// Transaction or Session.
     pub pool_mode: PoolMode,
 
+    /// Random or LeastOutstandingConnections.
+    pub load_balancing_mode: LoadBalancingMode,
+
     // Number of shards.
     pub shards: usize,
 
@@ -79,18 +102,47 @@ pub struct PoolSettings {
 
     // Sharding function.
     pub sharding_function: ShardingFunction,
+
+    // Sharding key
+    pub automatic_sharding_key: Option<String>,
+
+    // Health check timeout
+    pub healthcheck_timeout: u64,
+
+    // Health check delay
+    pub healthcheck_delay: u64,
+
+    // Ban time
+    pub ban_time: i64,
+
+    // Regex for searching for the sharding key in SQL statements
+    pub sharding_key_regex: Option<Regex>,
+
+    // Regex for searching for the shard id in SQL statements
+    pub shard_id_regex: Option<Regex>,
+
+    // Limit how much of each query is searched for a potential shard regex match
+    pub regex_search_limit: usize,
 }
 
 impl Default for PoolSettings {
     fn default() -> PoolSettings {
         PoolSettings {
             pool_mode: PoolMode::Transaction,
+            load_balancing_mode: LoadBalancingMode::Random,
             shards: 1,
             user: User::default(),
             default_role: None,
             query_parser_enabled: false,
             primary_reads_enabled: true,
             sharding_function: ShardingFunction::PgBigintHash,
+            automatic_sharding_key: None,
+            healthcheck_delay: General::default_healthcheck_delay(),
+            healthcheck_timeout: General::default_healthcheck_timeout(),
+            ban_time: General::default_ban_time(),
+            sharding_key_regex: None,
+            shard_id_regex: None,
+            regex_search_limit: 1000,
         }
     }
 }
@@ -116,10 +168,23 @@ pub struct ConnectionPool {
     /// The server information (K messages) have to be passed to the
     /// clients on startup. We pre-connect to all shards and replicas
     /// on pool creation and save the K messages here.
-    server_info: BytesMut,
+    server_info: Arc<RwLock<BytesMut>>,
 
     /// Pool configuration.
     pub settings: PoolSettings,
+
+    /// If not validated, we need to double check the pool is available before allowing a client
+    /// to use it.
+    validated: Arc<AtomicBool>,
+
+    /// Hash value for the pool configs. It is used to compare new configs
+    /// against current config to decide whether or not we need to recreate
+    /// the pool after a RELOAD command
+    pub config_hash: u64,
+
+    /// If the pool has been paused or not.
+    paused: Arc<AtomicBool>,
+    paused_waiter: Arc<Notify>,
 }
 
 impl ConnectionPool {
@@ -128,32 +193,32 @@ impl ConnectionPool {
         let config = get_config();
 
         let mut new_pools = HashMap::new();
-        let mut address_id = 0;
-
-        let mut pools_hash = (*(*POOLS_HASH.load())).clone();
+        let mut address_id: usize = 0;
 
         for (pool_name, pool_config) in &config.pools {
-            let changed = pools_hash.insert(pool_config.clone());
+            let new_pool_hash_value = pool_config.hash_value();
 
             // There is one pool per database/user pair.
-            for (_, user) in &pool_config.users {
-                // If the pool hasn't changed, get existing reference and insert it into the new_pools.
-                // We replace all pools at the end, but if the reference is kept, the pool won't get re-created (bb8).
-                if !changed {
-                    match get_pool(&pool_name, &user.username) {
-                        Some(pool) => {
+            for user in pool_config.users.values() {
+                let old_pool_ref = get_pool(pool_name, &user.username);
+
+                match old_pool_ref {
+                    Some(pool) => {
+                        // If the pool hasn't changed, get existing reference and insert it into the new_pools.
+                        // We replace all pools at the end, but if the reference is kept, the pool won't get re-created (bb8).
+                        if pool.config_hash == new_pool_hash_value {
                             info!(
                                 "[pool: {}][user: {}] has not changed",
                                 pool_name, user.username
                             );
                             new_pools.insert(
-                                PoolIdentifier::new(&pool_name, &user.username),
+                                PoolIdentifier::new(pool_name, &user.username),
                                 pool.clone(),
                             );
                             continue;
                         }
-                        None => (),
                     }
+                    None => (),
                 }
 
                 info!(
@@ -168,7 +233,6 @@ impl ConnectionPool {
                     .shards
                     .clone()
                     .into_keys()
-                    .map(|x| x.to_string())
                     .collect::<Vec<String>>();
 
                 // Sort by shard number to ensure consistency.
@@ -178,10 +242,35 @@ impl ConnectionPool {
                     let shard = &pool_config.shards[shard_idx];
                     let mut pools = Vec::new();
                     let mut servers = Vec::new();
-                    let mut address_index = 0;
                     let mut replica_number = 0;
 
-                    for server in shard.servers.iter() {
+                    // Load Mirror settings
+                    for (address_index, server) in shard.servers.iter().enumerate() {
+                        let mut mirror_addresses = vec![];
+                        if let Some(mirror_settings_vec) = &shard.mirrors {
+                            for (mirror_idx, mirror_settings) in
+                                mirror_settings_vec.iter().enumerate()
+                            {
+                                if mirror_settings.mirroring_target_index != address_index {
+                                    continue;
+                                }
+                                mirror_addresses.push(Address {
+                                    id: address_id,
+                                    database: shard.database.clone(),
+                                    host: mirror_settings.host.clone(),
+                                    port: mirror_settings.port,
+                                    role: server.role,
+                                    address_index: mirror_idx,
+                                    replica_number,
+                                    shard: shard_idx.parse::<usize>().unwrap(),
+                                    username: user.username.clone(),
+                                    pool_name: pool_name.clone(),
+                                    mirrors: vec![],
+                                });
+                                address_id += 1;
+                            }
+                        }
+
                         let address = Address {
                             id: address_id,
                             database: shard.database.clone(),
@@ -193,10 +282,10 @@ impl ConnectionPool {
                             shard: shard_idx.parse::<usize>().unwrap(),
                             username: user.username.clone(),
                             pool_name: pool_name.clone(),
+                            mirrors: mirror_addresses,
                         };
 
                         address_id += 1;
-                        address_index += 1;
 
                         if server.role == Role::Replica {
                             replica_number += 1;
@@ -215,9 +304,15 @@ impl ConnectionPool {
                             None => config.general.connect_timeout,
                         };
 
+                        let idle_timeout = match pool_config.idle_timeout {
+                            Some(idle_timeout) => idle_timeout,
+                            None => config.general.idle_timeout,
+                        };
+
                         let pool = Pool::builder()
                             .max_size(user.pool_size)
                             .connection_timeout(std::time::Duration::from_millis(connect_timeout))
+                            .idle_timeout(Some(std::time::Duration::from_millis(idle_timeout)))
                             .test_on_check_out(false)
                             .build(manager)
                             .await
@@ -234,14 +329,16 @@ impl ConnectionPool {
 
                 assert_eq!(shards.len(), addresses.len());
 
-                let mut pool = ConnectionPool {
+                let pool = ConnectionPool {
                     databases: shards,
-                    addresses: addresses,
+                    addresses,
                     banlist: Arc::new(RwLock::new(banlist)),
                     stats: get_reporter(),
-                    server_info: BytesMut::new(),
+                    config_hash: new_pool_hash_value,
+                    server_info: Arc::new(RwLock::new(BytesMut::new())),
                     settings: PoolSettings {
                         pool_mode: pool_config.pool_mode,
+                        load_balancing_mode: pool_config.load_balancing_mode,
                         // shards: pool_config.shards.clone(),
                         shards: shard_ids.len(),
                         user: user.clone(),
@@ -251,30 +348,42 @@ impl ConnectionPool {
                             "primary" => Some(Role::Primary),
                             _ => unreachable!(),
                         },
-                        query_parser_enabled: pool_config.query_parser_enabled.clone(),
+                        query_parser_enabled: pool_config.query_parser_enabled,
                         primary_reads_enabled: pool_config.primary_reads_enabled,
                         sharding_function: pool_config.sharding_function,
+                        automatic_sharding_key: pool_config.automatic_sharding_key.clone(),
+                        healthcheck_delay: config.general.healthcheck_delay,
+                        healthcheck_timeout: config.general.healthcheck_timeout,
+                        ban_time: config.general.ban_time,
+                        sharding_key_regex: pool_config
+                            .sharding_key_regex
+                            .clone()
+                            .map(|regex| Regex::new(regex.as_str()).unwrap()),
+                        shard_id_regex: pool_config
+                            .shard_id_regex
+                            .clone()
+                            .map(|regex| Regex::new(regex.as_str()).unwrap()),
+                        regex_search_limit: pool_config.regex_search_limit.unwrap_or(1000),
                     },
+                    validated: Arc::new(AtomicBool::new(false)),
+                    paused: Arc::new(AtomicBool::new(false)),
+                    paused_waiter: Arc::new(Notify::new()),
                 };
 
                 // Connect to the servers to make sure pool configuration is valid
                 // before setting it globally.
-                match pool.validate().await {
-                    Ok(_) => (),
-                    Err(err) => {
-                        error!("Could not validate connection pool: {:?}", err);
-                        return Err(err);
-                    }
-                };
+                // Do this async and somewhere else, we don't have to wait here.
+                let mut validate_pool = pool.clone();
+                tokio::task::spawn(async move {
+                    let _ = validate_pool.validate().await;
+                });
 
                 // There is one pool per database/user pair.
-                new_pools.insert(PoolIdentifier::new(&pool_name, &user.username), pool);
+                new_pools.insert(PoolIdentifier::new(pool_name, &user.username), pool);
             }
         }
 
         POOLS.store(Arc::new(new_pools.clone()));
-        POOLS_HASH.store(Arc::new(pools_hash.clone()));
-
         Ok(())
     }
 
@@ -283,66 +392,110 @@ impl ConnectionPool {
     /// when they connect.
     /// This also warms up the pool for clients that connect when
     /// the pooler starts up.
-    async fn validate(&mut self) -> Result<(), Error> {
-        let mut server_infos = Vec::new();
+    pub async fn validate(&mut self) -> Result<(), Error> {
+        let mut futures = Vec::new();
+        let validated = Arc::clone(&self.validated);
+
         for shard in 0..self.shards() {
             for server in 0..self.servers(shard) {
-                let connection = match self.databases[shard][server].get().await {
-                    Ok(conn) => conn,
-                    Err(err) => {
-                        error!("Shard {} down or misconfigured: {:?}", shard, err);
-                        continue;
-                    }
-                };
+                let databases = self.databases.clone();
+                let validated = Arc::clone(&validated);
+                let pool_server_info = Arc::clone(&self.server_info);
 
-                let proxy = connection;
-                let server = &*proxy;
-                let server_info = server.server_info();
+                let task = tokio::task::spawn(async move {
+                    let connection = match databases[shard][server].get().await {
+                        Ok(conn) => conn,
+                        Err(err) => {
+                            error!("Shard {} down or misconfigured: {:?}", shard, err);
+                            return;
+                        }
+                    };
 
-                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",
-                            proxy.address()
-                        );
-                    }
-                }
+                    let proxy = connection;
+                    let server = &*proxy;
+                    let server_info = server.server_info();
 
-                server_infos.push(server_info);
+                    let mut guard = pool_server_info.write();
+                    guard.clear();
+                    guard.put(server_info.clone());
+                    validated.store(true, Ordering::Relaxed);
+                });
+
+                futures.push(task);
             }
         }
 
+        futures::future::join_all(futures).await;
+
         // TODO: compare server information to make sure
         // all shards are running identical configurations.
-        if server_infos.len() == 0 {
+        if self.server_info.read().is_empty() {
+            error!("Could not validate connection pool");
             return Err(Error::AllServersDown);
         }
 
-        // We're assuming all servers are identical.
-        // TODO: not true.
-        self.server_info = server_infos[0].clone();
-
         Ok(())
     }
 
+    /// The pool can be used by clients.
+    ///
+    /// If not, we need to validate it first by connecting to servers.
+    /// Call `validate()` to do so.
+    pub fn validated(&self) -> bool {
+        self.validated.load(Ordering::Relaxed)
+    }
+
+    /// Pause the pool, allowing no more queries and make clients wait.
+    pub fn pause(&self) {
+        self.paused.store(true, Ordering::Relaxed);
+    }
+
+    /// Resume the pool, allowing queries and resuming any pending queries.
+    pub fn resume(&self) {
+        self.paused.store(false, Ordering::Relaxed);
+        self.paused_waiter.notify_waiters();
+    }
+
+    /// Check if the pool is paused.
+    pub fn paused(&self) -> bool {
+        self.paused.load(Ordering::Relaxed)
+    }
+
+    /// Check if the pool is paused and wait until it's resumed.
+    pub async fn wait_paused(&self) -> bool {
+        let waiter = self.paused_waiter.notified();
+        let paused = self.paused.load(Ordering::Relaxed);
+
+        if paused {
+            waiter.await;
+        }
+
+        paused
+    }
+
     /// Get a connection from the pool.
     pub async fn get(
         &self,
-        shard: usize,       // shard number
-        role: Option<Role>, // primary or replica
-        process_id: i32,    // client id
+        shard: usize,           // shard number
+        role: Option<Role>,     // primary or replica
+        client_process_id: i32, // client id
     ) -> Result<(PooledConnection<'_, ServerPool>, Address), Error> {
         let mut candidates: Vec<&Address> = self.addresses[shard]
             .iter()
             .filter(|address| address.role == role)
             .collect();
 
-        // Random load balancing
+        // We shuffle even if least_outstanding_queries is used to avoid imbalance
+        // in cases where all candidates have more or less the same number of outstanding
+        // queries
         candidates.shuffle(&mut thread_rng());
-
-        let healthcheck_timeout = get_config().general.healthcheck_timeout;
-        let healthcheck_delay = get_config().general.healthcheck_delay as u128;
+        if self.settings.load_balancing_mode == LoadBalancingMode::LeastOutstandingConnections {
+            candidates.sort_by(|a, b| {
+                self.busy_connection_count(b)
+                    .partial_cmp(&self.busy_connection_count(a))
+                    .unwrap()
+            });
+        }
 
         while !candidates.is_empty() {
             // Get the next candidate
@@ -351,14 +504,20 @@ impl ConnectionPool {
                 None => break,
             };
 
-            if self.is_banned(&address, role) {
-                debug!("Address {:?} is banned", address);
-                continue;
+            let mut force_healthcheck = false;
+
+            if self.is_banned(address) {
+                if self.try_unban(&address).await {
+                    force_healthcheck = true;
+                } else {
+                    debug!("Address {:?} is banned", address);
+                    continue;
+                }
             }
 
             // Indicate we're waiting on a server connection from a pool.
             let now = Instant::now();
-            self.stats.client_waiting(process_id);
+            self.stats.client_waiting(client_process_id);
 
             // Check if we can connect
             let mut conn = match self.databases[address.shard][address.address_index]
@@ -368,8 +527,9 @@ impl ConnectionPool {
                 Ok(conn) => conn,
                 Err(err) => {
                     error!("Banning instance {:?}, error: {:?}", address, err);
-                    self.ban(&address, process_id);
-                    self.stats.client_checkout_error(process_id, address.id);
+                    self.ban(address, BanReason::FailedCheckout, client_process_id);
+                    self.stats
+                        .client_checkout_error(client_process_id, address.id);
                     continue;
                 }
             };
@@ -378,138 +538,122 @@ impl ConnectionPool {
             let server = &mut *conn;
 
             // Will return error if timestamp is greater than current system time, which it should never be set to
-            let require_healthcheck =
-                server.last_activity().elapsed().unwrap().as_millis() > healthcheck_delay;
+            let require_healthcheck = force_healthcheck
+                || server.last_activity().elapsed().unwrap().as_millis()
+                    > self.settings.healthcheck_delay as u128;
 
             // Do not issue a health check unless it's been a little while
             // since we last checked the server is ok.
             // Health checks are pretty expensive.
             if !require_healthcheck {
+                self.stats.checkout_time(
+                    now.elapsed().as_micros(),
+                    client_process_id,
+                    server.server_id(),
+                );
                 self.stats
-                    .checkout_time(now.elapsed().as_micros(), process_id, server.server_id());
-                self.stats.server_active(process_id, server.server_id());
+                    .server_active(client_process_id, server.server_id());
                 return Ok((conn, address.clone()));
             }
 
-            debug!("Running health check on server {:?}", address);
-
-            self.stats.server_tested(server.server_id());
-
-            match tokio::time::timeout(
-                tokio::time::Duration::from_millis(healthcheck_timeout),
-                server.query(";"), // Cheap query as it skips the query planner
-            )
-            .await
+            if self
+                .run_health_check(address, server, now, client_process_id)
+                .await
             {
-                // Check if health check succeeded.
-                Ok(res) => match res {
-                    Ok(_) => {
-                        self.stats.checkout_time(
-                            now.elapsed().as_micros(),
-                            process_id,
-                            conn.server_id(),
-                        );
-                        self.stats.server_active(process_id, conn.server_id());
-                        return Ok((conn, address.clone()));
-                    }
-
-                    // Health check failed.
-                    Err(err) => {
-                        error!(
-                            "Banning instance {:?} because of failed health check, {:?}",
-                            address, err
-                        );
-
-                        // Don't leave a bad connection in the pool.
-                        server.mark_bad();
-
-                        self.ban(&address, process_id);
-                        continue;
-                    }
-                },
-
-                // Health check timed out.
-                Err(err) => {
-                    error!(
-                        "Banning instance {:?} because of health check timeout, {:?}",
-                        address, err
-                    );
-                    // Don't leave a bad connection in the pool.
-                    server.mark_bad();
-
-                    self.ban(&address, process_id);
-                    continue;
-                }
+                return Ok((conn, address.clone()));
+            } else {
+                continue;
             }
         }
 
         Err(Error::AllServersDown)
     }
 
+    async fn run_health_check(
+        &self,
+        address: &Address,
+        server: &mut Server,
+        start: Instant,
+        client_process_id: i32,
+    ) -> bool {
+        debug!("Running health check on server {:?}", address);
+
+        self.stats.server_tested(server.server_id());
+
+        match tokio::time::timeout(
+            tokio::time::Duration::from_millis(self.settings.healthcheck_timeout),
+            server.query(";"), // Cheap query as it skips the query planner
+        )
+        .await
+        {
+            // Check if health check succeeded.
+            Ok(res) => match res {
+                Ok(_) => {
+                    self.stats.checkout_time(
+                        start.elapsed().as_micros(),
+                        client_process_id,
+                        server.server_id(),
+                    );
+                    self.stats
+                        .server_active(client_process_id, server.server_id());
+                    return true;
+                }
+
+                // Health check failed.
+                Err(err) => {
+                    error!(
+                        "Banning instance {:?} because of failed health check, {:?}",
+                        address, err
+                    );
+                }
+            },
+
+            // Health check timed out.
+            Err(err) => {
+                error!(
+                    "Banning instance {:?} because of health check timeout, {:?}",
+                    address, err
+                );
+            }
+        }
+
+        // Don't leave a bad connection in the pool.
+        server.mark_bad();
+
+        self.ban(&address, BanReason::FailedHealthCheck, client_process_id);
+        return false;
+    }
+
     /// Ban an address (i.e. replica). It no longer will serve
     /// traffic for any new transactions. Existing transactions on that replica
     /// will finish successfully or error out to the clients.
-    pub fn ban(&self, address: &Address, client_id: i32) {
-        error!("Banning {:?}", address);
-        self.stats.client_ban_error(client_id, address.id);
+    pub fn ban(&self, address: &Address, reason: BanReason, client_id: i32) {
+        // Primary can never be banned
+        if address.role == Role::Primary {
+            return;
+        }
 
         let now = chrono::offset::Utc::now().naive_utc();
         let mut guard = self.banlist.write();
-        guard[address.shard].insert(address.clone(), now);
+        error!("Banning {:?}", address);
+        self.stats.client_ban_error(client_id, address.id);
+        guard[address.shard].insert(address.clone(), (reason, now));
     }
 
     /// Clear the replica to receive traffic again. Takes effect immediately
     /// for all new transactions.
-    pub fn _unban(&self, address: &Address) {
+    pub fn unban(&self, address: &Address) {
         let mut guard = self.banlist.write();
         guard[address.shard].remove(address);
     }
 
-    /// Check if a replica can serve traffic. If all replicas are banned,
-    /// we unban all of them. Better to try then not to.
-    pub fn is_banned(&self, address: &Address, role: Option<Role>) -> bool {
-        let replicas_available = match role {
-            Some(Role::Replica) => self.addresses[address.shard]
-                .iter()
-                .filter(|addr| addr.role == Role::Replica)
-                .count(),
-            None => self.addresses[address.shard].len(),
-            Some(Role::Primary) => return false, // Primary cannot be banned.
-        };
-
-        debug!("Available targets for {:?}: {}", role, replicas_available);
-
+    /// Check if address is banned
+    /// true if banned, false otherwise
+    pub fn is_banned(&self, address: &Address) -> bool {
         let guard = self.banlist.read();
 
-        // Everything is banned = nothing is banned.
-        if guard[address.shard].len() == replicas_available {
-            drop(guard);
-            let mut guard = self.banlist.write();
-            guard[address.shard].clear();
-            drop(guard);
-            warn!("Unbanning all replicas.");
-            return false;
-        }
-
-        // I expect this to miss 99.9999% of the time.
         match guard[address.shard].get(address) {
-            Some(timestamp) => {
-                let now = chrono::offset::Utc::now().naive_utc();
-                let config = get_config();
-
-                // Ban expired.
-                if now.timestamp() - timestamp.timestamp() > config.general.ban_time {
-                    drop(guard);
-                    warn!("Unbanning {:?}", address);
-                    let mut guard = self.banlist.write();
-                    guard[address.shard].remove(address);
-                    false
-                } else {
-                    debug!("{:?} is banned", address);
-                    true
-                }
-            }
-
+            Some(_) => true,
             None => {
                 debug!("{:?} is ok", address);
                 false
@@ -517,11 +661,92 @@ impl ConnectionPool {
         }
     }
 
+    /// Determines trying to unban this server was successful
+    pub async fn try_unban(&self, address: &Address) -> bool {
+        // If somehow primary ends up being banned we should return true here
+        if address.role == Role::Primary {
+            return true;
+        }
+
+        // Check if all replicas are banned, in that case unban all of them
+        let replicas_available = self.addresses[address.shard]
+            .iter()
+            .filter(|addr| addr.role == Role::Replica)
+            .count();
+
+        debug!("Available targets: {}", replicas_available);
+
+        let read_guard = self.banlist.read();
+        let all_replicas_banned = read_guard[address.shard].len() == replicas_available;
+        drop(read_guard);
+
+        if all_replicas_banned {
+            let mut write_guard = self.banlist.write();
+            warn!("Unbanning all replicas.");
+            write_guard[address.shard].clear();
+
+            return true;
+        }
+
+        // Check if ban time is expired
+        let read_guard = self.banlist.read();
+        let exceeded_ban_time = match read_guard[address.shard].get(address) {
+            Some((ban_reason, timestamp)) => {
+                let now = chrono::offset::Utc::now().naive_utc();
+                match ban_reason {
+                    BanReason::AdminBan(duration) => {
+                        now.timestamp() - timestamp.timestamp() > *duration
+                    }
+                    _ => now.timestamp() - timestamp.timestamp() > self.settings.ban_time,
+                }
+            }
+            None => return true,
+        };
+        drop(read_guard);
+
+        if exceeded_ban_time {
+            warn!("Unbanning {:?}", address);
+            let mut write_guard = self.banlist.write();
+            write_guard[address.shard].remove(address);
+            drop(write_guard);
+
+            true
+        } else {
+            debug!("{:?} is banned", address);
+            false
+        }
+    }
+
     /// Get the number of configured shards.
     pub fn shards(&self) -> usize {
         self.databases.len()
     }
 
+    pub fn get_bans(&self) -> Vec<(Address, (BanReason, NaiveDateTime))> {
+        let mut bans: Vec<(Address, (BanReason, NaiveDateTime))> = Vec::new();
+        let guard = self.banlist.read();
+        for banlist in guard.iter() {
+            for (address, (reason, timestamp)) in banlist.iter() {
+                bans.push((address.clone(), (reason.clone(), timestamp.clone())));
+            }
+        }
+        return bans;
+    }
+
+    /// Get the address from the host url
+    pub fn get_addresses_from_host(&self, host: &str) -> Vec<Address> {
+        let mut addresses = Vec::new();
+        for shard in 0..self.shards() {
+            for server in 0..self.servers(shard) {
+                let address = self.address(shard, server);
+                if address.host == host {
+                    addresses.push(address.clone());
+                }
+            }
+        }
+        addresses
+    }
+
     /// Get the number of servers (primary and replicas)
     /// configured for a shard.
     pub fn servers(&self, shard: usize) -> usize {
@@ -548,7 +773,21 @@ impl ConnectionPool {
     }
 
     pub fn server_info(&self) -> BytesMut {
-        self.server_info.clone()
+        self.server_info.read().clone()
+    }
+
+    fn busy_connection_count(&self, address: &Address) -> u32 {
+        let state = self.pool_state(address.shard, address.address_index);
+        let idle = state.idle_connections;
+        let provisioned = state.connections;
+
+        if idle > provisioned {
+            // Unlikely but avoids an overflow panic if this ever happens
+            return 0;
+        }
+        let busy = provisioned - idle;
+        debug!("{:?} has {:?} busy connections", address, busy);
+        return busy;
     }
 }
 
@@ -570,11 +809,11 @@ impl ServerPool {
         stats: Reporter,
     ) -> ServerPool {
         ServerPool {
-            address: address,
-            user: user,
+            address,
+            user,
             database: database.to_string(),
-            client_server_map: client_server_map,
-            stats: stats,
+            client_server_map,
+            stats,
         }
     }
 }
@@ -633,15 +872,14 @@ impl ManageConnection for ServerPool {
 
 /// Get the connection pool
 pub fn get_pool(db: &str, user: &str) -> Option<ConnectionPool> {
-    match get_all_pools().get(&PoolIdentifier::new(&db, &user)) {
-        Some(pool) => Some(pool.clone()),
-        None => None,
-    }
+    (*(*POOLS.load()))
+        .get(&PoolIdentifier::new(db, user))
+        .cloned()
 }
 
 /// Get a pointer to all configured pools.
 pub fn get_all_pools() -> HashMap<PoolIdentifier, ConnectionPool> {
-    return (*(*POOLS.load())).clone();
+    (*(*POOLS.load())).clone()
 }
 
 /// How many total servers we have in the config.

src/prometheus.rs

@@ -8,7 +8,7 @@ use std::net::SocketAddr;
 
 use crate::config::Address;
 use crate::pool::get_all_pools;
-use crate::stats::get_address_stats;
+use crate::stats::{get_address_stats, get_pool_stats, get_server_stats, ServerInformation};
 
 struct MetricHelpType {
     help: &'static str,
@@ -19,109 +19,141 @@ struct MetricHelpType {
 // counters only increase
 // gauges can arbitrarily increase or decrease
 static METRIC_HELP_AND_TYPES_LOOKUP: phf::Map<&'static str, MetricHelpType> = phf_map! {
-    "total_query_count" => MetricHelpType {
+    "stats_total_query_count" => MetricHelpType {
         help: "Number of queries sent by all clients",
         ty: "counter",
     },
-    "total_query_time" => MetricHelpType {
+    "stats_total_query_time" => MetricHelpType {
         help: "Total amount of time for queries to execute",
         ty: "counter",
     },
-    "total_received" => MetricHelpType {
+    "stats_total_received" => MetricHelpType {
         help: "Number of bytes received from the server",
         ty: "counter",
     },
-    "total_sent" => MetricHelpType {
+    "stats_total_sent" => MetricHelpType {
         help: "Number of bytes sent to the server",
         ty: "counter",
     },
-    "total_xact_count" => MetricHelpType {
+    "stats_total_xact_count" => MetricHelpType {
         help: "Total number of transactions started by the client",
         ty: "counter",
     },
-    "total_xact_time" => MetricHelpType {
+    "stats_total_xact_time" => MetricHelpType {
         help: "Total amount of time for all transactions to execute",
         ty: "counter",
     },
-    "total_wait_time" => MetricHelpType {
+    "stats_total_wait_time" => MetricHelpType {
         help: "Total time client waited for a server connection",
         ty: "counter",
     },
-    "avg_query_count" => MetricHelpType {
+    "stats_avg_query_count" => MetricHelpType {
         help: "Average of total_query_count every 15 seconds",
         ty: "gauge",
     },
-    "avg_query_time" => MetricHelpType {
+    "stats_avg_query_time" => MetricHelpType {
         help: "Average time taken for queries to execute every 15 seconds",
         ty: "gauge",
     },
-    "avg_recv" => MetricHelpType {
+    "stats_avg_recv" => MetricHelpType {
         help: "Average of total_received bytes every 15 seconds",
         ty: "gauge",
     },
-    "avg_sent" => MetricHelpType {
+    "stats_avg_sent" => MetricHelpType {
         help: "Average of total_sent bytes every 15 seconds",
         ty: "gauge",
     },
-    "avg_xact_count" => MetricHelpType {
+    "stats_avg_errors" => MetricHelpType {
+        help: "Average number of errors every 15 seconds",
+        ty: "gauge",
+    },
+    "stats_avg_xact_count" => MetricHelpType {
         help: "Average of total_xact_count every 15 seconds",
         ty: "gauge",
     },
-    "avg_xact_time" => MetricHelpType {
+    "stats_avg_xact_time" => MetricHelpType {
         help: "Average of total_xact_time every 15 seconds",
         ty: "gauge",
     },
-    "avg_wait_time" => MetricHelpType {
+    "stats_avg_wait_time" => MetricHelpType {
         help: "Average of total_wait_time every 15 seconds",
         ty: "gauge",
     },
-    "maxwait_us" => MetricHelpType {
+    "pools_maxwait_us" => MetricHelpType {
         help: "The time a client waited for a server connection in microseconds",
         ty: "gauge",
     },
-    "maxwait" => MetricHelpType {
+    "pools_maxwait" => MetricHelpType {
         help: "The time a client waited for a server connection in seconds",
         ty: "gauge",
     },
-    "cl_waiting" => MetricHelpType {
+    "pools_cl_waiting" => MetricHelpType {
         help: "How many clients are waiting for a connection from the pool",
         ty: "gauge",
     },
-    "cl_active" => MetricHelpType {
+    "pools_cl_active" => MetricHelpType {
         help: "How many clients are actively communicating with a server",
         ty: "gauge",
     },
-    "cl_idle" => MetricHelpType {
+    "pools_cl_idle" => MetricHelpType {
         help: "How many clients are idle",
         ty: "gauge",
     },
-    "sv_idle" => MetricHelpType {
+    "pools_sv_idle" => MetricHelpType {
         help: "How many server connections are idle",
         ty: "gauge",
     },
-    "sv_active" => MetricHelpType {
+    "pools_sv_active" => MetricHelpType {
         help: "How many server connections are actively communicating with a client",
         ty: "gauge",
     },
-    "sv_login" => MetricHelpType {
+    "pools_sv_login" => MetricHelpType {
         help: "How many server connections are currently being created",
         ty: "gauge",
     },
-    "sv_tested" => MetricHelpType {
+    "pools_sv_tested" => MetricHelpType {
         help: "How many server connections are currently waiting on a health check to succeed",
         ty: "gauge",
     },
+    "servers_bytes_received" => MetricHelpType {
+        help: "Volume in bytes of network traffic received by server",
+        ty: "gauge",
+    },
+    "servers_bytes_sent" => MetricHelpType {
+        help: "Volume in bytes of network traffic sent by server",
+        ty: "gauge",
+    },
+    "servers_transaction_count" => MetricHelpType {
+        help: "Number of transactions executed by server",
+        ty: "gauge",
+    },
+    "servers_query_count" => MetricHelpType {
+        help: "Number of queries executed by server",
+        ty: "gauge",
+    },
+    "servers_error_count" => MetricHelpType {
+        help: "Number of errors",
+        ty: "gauge",
+    },
+    "databases_pool_size" => MetricHelpType {
+        help: "Maximum number of server connections",
+        ty: "gauge",
+    },
+    "databases_current_connections" => MetricHelpType {
+        help: "Current number of connections for this database",
+        ty: "gauge",
+    },
 };
 
-struct PrometheusMetric {
+struct PrometheusMetric<Value: fmt::Display> {
     name: String,
     help: String,
     ty: String,
     labels: HashMap<&'static str, String>,
-    value: i64,
+    value: Value,
 }
 
-impl fmt::Display for PrometheusMetric {
+impl<Value: fmt::Display> fmt::Display for PrometheusMetric<Value> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let formatted_labels = self
             .labels
@@ -141,50 +173,81 @@ impl fmt::Display for PrometheusMetric {
     }
 }
 
-impl PrometheusMetric {
-    fn new(address: &Address, name: &str, value: i64) -> Option<PrometheusMetric> {
+impl<Value: fmt::Display> PrometheusMetric<Value> {
+    fn from_name<V: fmt::Display>(
+        name: &str,
+        value: V,
+        labels: HashMap<&'static str, String>,
+    ) -> Option<PrometheusMetric<V>> {
+        METRIC_HELP_AND_TYPES_LOOKUP
+            .get(name)
+            .map(|metric| PrometheusMetric::<V> {
+                name: name.to_owned(),
+                help: metric.help.to_owned(),
+                ty: metric.ty.to_owned(),
+                value,
+                labels,
+            })
+    }
+
+    fn from_database_info(
+        address: &Address,
+        name: &str,
+        value: u32,
+    ) -> Option<PrometheusMetric<u32>> {
         let mut labels = HashMap::new();
         labels.insert("host", address.host.clone());
         labels.insert("shard", address.shard.to_string());
         labels.insert("role", address.role.to_string());
+        labels.insert("pool", address.pool_name.clone());
         labels.insert("database", address.database.to_string());
 
-        METRIC_HELP_AND_TYPES_LOOKUP
-            .get(name)
-            .map(|metric| PrometheusMetric {
-                name: name.to_owned(),
-                help: metric.help.to_owned(),
-                ty: metric.ty.to_owned(),
-                labels,
-                value,
-            })
+        Self::from_name(&format!("databases_{}", name), value, labels)
+    }
+
+    fn from_server_info(
+        address: &Address,
+        name: &str,
+        value: u64,
+    ) -> Option<PrometheusMetric<u64>> {
+        let mut labels = HashMap::new();
+        labels.insert("host", address.host.clone());
+        labels.insert("shard", address.shard.to_string());
+        labels.insert("role", address.role.to_string());
+        labels.insert("pool", address.pool_name.clone());
+        labels.insert("database", address.database.to_string());
+
+        Self::from_name(&format!("servers_{}", name), value, labels)
+    }
+
+    fn from_address(address: &Address, name: &str, value: i64) -> Option<PrometheusMetric<i64>> {
+        let mut labels = HashMap::new();
+        labels.insert("host", address.host.clone());
+        labels.insert("shard", address.shard.to_string());
+        labels.insert("pool", address.pool_name.clone());
+        labels.insert("role", address.role.to_string());
+        labels.insert("database", address.database.to_string());
+
+        Self::from_name(&format!("stats_{}", name), value, labels)
+    }
+
+    fn from_pool(pool: &(String, String), name: &str, value: i64) -> Option<PrometheusMetric<i64>> {
+        let mut labels = HashMap::new();
+        labels.insert("pool", pool.0.clone());
+        labels.insert("user", pool.1.clone());
+
+        Self::from_name(&format!("pools_{}", name), value, labels)
     }
 }
 
 async fn prometheus_stats(request: Request<Body>) -> Result<Response<Body>, hyper::http::Error> {
     match (request.method(), request.uri().path()) {
         (&Method::GET, "/metrics") => {
-            let stats: HashMap<usize, HashMap<String, i64>> = get_address_stats();
-
             let mut lines = Vec::new();
-            for (_, pool) in get_all_pools() {
-                for shard in 0..pool.shards() {
-                    for server in 0..pool.servers(shard) {
-                        let address = pool.address(shard, server);
-                        if let Some(address_stats) = stats.get(&address.id) {
-                            for (key, value) in address_stats.iter() {
-                                if let Some(prometheus_metric) =
-                                    PrometheusMetric::new(address, key, *value)
-                                {
-                                    lines.push(prometheus_metric.to_string());
-                                } else {
-                                    warn!("Metric {} not implemented for {}", key, address.name());
-                                }
-                            }
-                        }
-                    }
-                }
-            }
+            push_address_stats(&mut lines);
+            push_pool_stats(&mut lines);
+            push_server_stats(&mut lines);
+            push_database_stats(&mut lines);
 
             Response::builder()
                 .header("content-type", "text/plain; version=0.0.4")
@@ -196,10 +259,113 @@ async fn prometheus_stats(request: Request<Body>) -> Result<Response<Body>, hype
     }
 }
 
+// Adds metrics shown in a SHOW STATS admin command.
+fn push_address_stats(lines: &mut Vec<String>) {
+    let address_stats: HashMap<usize, HashMap<String, i64>> = get_address_stats();
+    for (_, pool) in get_all_pools() {
+        for shard in 0..pool.shards() {
+            for server in 0..pool.servers(shard) {
+                let address = pool.address(shard, server);
+                if let Some(address_stats) = address_stats.get(&address.id) {
+                    for (key, value) in address_stats.iter() {
+                        if let Some(prometheus_metric) =
+                            PrometheusMetric::<i64>::from_address(address, key, *value)
+                        {
+                            lines.push(prometheus_metric.to_string());
+                        } else {
+                            warn!("Metric {} not implemented for {}", key, address.name());
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+// Adds relevant metrics shown in a SHOW POOLS admin command.
+fn push_pool_stats(lines: &mut Vec<String>) {
+    let pool_stats = get_pool_stats();
+    for (pool, stats) in pool_stats.iter() {
+        for (name, value) in stats.iter() {
+            if let Some(prometheus_metric) = PrometheusMetric::<i64>::from_pool(pool, name, *value)
+            {
+                lines.push(prometheus_metric.to_string());
+            } else {
+                warn!(
+                    "Metric {} not implemented for ({},{})",
+                    name, pool.0, pool.1
+                );
+            }
+        }
+    }
+}
+
+// Adds relevant metrics shown in a SHOW DATABASES admin command.
+fn push_database_stats(lines: &mut Vec<String>) {
+    for (_, pool) in get_all_pools() {
+        let pool_config = pool.settings.clone();
+        for shard in 0..pool.shards() {
+            for server in 0..pool.servers(shard) {
+                let address = pool.address(shard, server);
+                let pool_state = pool.pool_state(shard, server);
+
+                let metrics = vec![
+                    ("pool_size", pool_config.user.pool_size),
+                    ("current_connections", pool_state.connections),
+                ];
+                for (key, value) in metrics {
+                    if let Some(prometheus_metric) =
+                        PrometheusMetric::<u32>::from_database_info(address, key, value)
+                    {
+                        lines.push(prometheus_metric.to_string());
+                    } else {
+                        warn!("Metric {} not implemented for {}", key, address.name());
+                    }
+                }
+            }
+        }
+    }
+}
+
+// Adds relevant metrics shown in a SHOW SERVERS admin command.
+fn push_server_stats(lines: &mut Vec<String>) {
+    let server_stats = get_server_stats();
+    let mut server_stats_by_addresses = HashMap::<String, ServerInformation>::new();
+    for (_, info) in server_stats {
+        server_stats_by_addresses.insert(info.address_name.clone(), info);
+    }
+
+    for (_, pool) in get_all_pools() {
+        for shard in 0..pool.shards() {
+            for server in 0..pool.servers(shard) {
+                let address = pool.address(shard, server);
+                if let Some(server_info) = server_stats_by_addresses.get(&address.name()) {
+                    let metrics = [
+                        ("bytes_received", server_info.bytes_received),
+                        ("bytes_sent", server_info.bytes_sent),
+                        ("transaction_count", server_info.transaction_count),
+                        ("query_count", server_info.query_count),
+                        ("error_count", server_info.error_count),
+                    ];
+                    for (key, value) in metrics {
+                        if let Some(prometheus_metric) =
+                            PrometheusMetric::<u64>::from_server_info(address, key, value)
+                        {
+                            lines.push(prometheus_metric.to_string());
+                        } else {
+                            warn!("Metric {} not implemented for {}", key, address.name());
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
 pub async fn start_metric_server(http_addr: SocketAddr) {
     let http_service_factory =
         make_service_fn(|_conn| async { Ok::<_, hyper::Error>(service_fn(prometheus_stats)) });
-    let server = Server::bind(&http_addr.into()).serve(http_service_factory);
+    let server = Server::bind(&http_addr).serve(http_service_factory);
     info!(
         "Exposing prometheus metrics on http://{}/metrics.",
         http_addr

src/scram.rs

@@ -2,6 +2,7 @@
 // https://github.com/sfackler/rust-postgres/
 // SASL implementation.
 
+use base64::{engine::general_purpose, Engine as _};
 use bytes::BytesMut;
 use hmac::{Hmac, Mac};
 use rand::{self, Rng};
@@ -57,7 +58,7 @@ impl ScramSha256 {
 
     /// Used for testing.
     pub fn from_nonce(password: &str, nonce: &str) -> ScramSha256 {
-        let message = BytesMut::from(&format!("{}n=,r={}", "n,,", nonce).as_bytes()[..]);
+        let message = BytesMut::from(format!("{}n=,r={}", "n,,", nonce).as_bytes());
 
         ScramSha256 {
             password: password.to_string(),
@@ -78,16 +79,16 @@ impl ScramSha256 {
         let server_message = Message::parse(message)?;
 
         if !server_message.nonce.starts_with(&self.nonce) {
-            return Err(Error::ProtocolSyncError);
+            return Err(Error::ProtocolSyncError(format!("SCRAM")));
         }
 
-        let salt = match base64::decode(&server_message.salt) {
+        let salt = match general_purpose::STANDARD.decode(&server_message.salt) {
             Ok(salt) => salt,
-            Err(_) => return Err(Error::ProtocolSyncError),
+            Err(_) => return Err(Error::ProtocolSyncError(format!("SCRAM"))),
         };
 
         let salted_password = Self::hi(
-            &normalize(&self.password.as_bytes()[..]),
+            &normalize(self.password.as_bytes()),
             &salt,
             server_message.iterations,
         );
@@ -111,7 +112,7 @@ impl ScramSha256 {
         let mut cbind_input = vec![];
         cbind_input.extend("n,,".as_bytes());
 
-        let cbind_input = base64::encode(&cbind_input);
+        let cbind_input = general_purpose::STANDARD.encode(&cbind_input);
 
         self.message.clear();
 
@@ -149,7 +150,11 @@ impl ScramSha256 {
             *proof ^= signature;
         }
 
-        match write!(&mut self.message, ",p={}", base64::encode(&*client_proof)) {
+        match write!(
+            &mut self.message,
+            ",p={}",
+            general_purpose::STANDARD.encode(&*client_proof)
+        ) {
             Ok(_) => (),
             Err(_) => return Err(Error::ServerError),
         };
@@ -161,9 +166,9 @@ impl ScramSha256 {
     pub fn finish(&mut self, message: &BytesMut) -> Result<(), Error> {
         let final_message = FinalMessage::parse(message)?;
 
-        let verifier = match base64::decode(&final_message.value) {
+        let verifier = match general_purpose::STANDARD.decode(&final_message.value) {
             Ok(verifier) => verifier,
-            Err(_) => return Err(Error::ProtocolSyncError),
+            Err(_) => return Err(Error::ProtocolSyncError(format!("SCRAM"))),
         };
 
         let mut hmac = match Hmac::<Sha256>::new_from_slice(&self.salted_password) {
@@ -181,7 +186,7 @@ impl ScramSha256 {
 
         match hmac.verify_slice(&verifier) {
             Ok(_) => Ok(()),
-            Err(_) => return Err(Error::ServerError),
+            Err(_) => Err(Error::ServerError),
         }
     }
 
@@ -220,19 +225,19 @@ impl Message {
     /// Parse the server SASL challenge.
     fn parse(message: &BytesMut) -> Result<Message, Error> {
         let parts = String::from_utf8_lossy(&message[..])
-            .split(",")
+            .split(',')
             .map(|s| s.to_string())
             .collect::<Vec<String>>();
 
         if parts.len() != 3 {
-            return Err(Error::ProtocolSyncError);
+            return Err(Error::ProtocolSyncError(format!("SCRAM")));
         }
 
         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),
+            Err(_) => return Err(Error::ProtocolSyncError(format!("SCRAM"))),
         };
 
         Ok(Message {
@@ -252,7 +257,7 @@ 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);
+            return Err(Error::ProtocolSyncError(format!("SCRAM")));
         }
 
         Ok(FinalMessage {
@@ -268,7 +273,7 @@ mod test {
     #[test]
     fn parse_server_first_message() {
         let message = BytesMut::from(
-            &"r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,s=QSXCR+Q6sek8bf92,i=4096".as_bytes()[..],
+            "r=fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j,s=QSXCR+Q6sek8bf92,i=4096".as_bytes(),
         );
         let message = Message::parse(&message).unwrap();
         assert_eq!(message.nonce, "fyko+d2lbbFgONRv9qkxdawL3rfcNHYJY1ZVvWVs7j");
@@ -279,7 +284,7 @@ mod test {
     #[test]
     fn parse_server_last_message() {
         let f = FinalMessage::parse(&BytesMut::from(
-            &"v=U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw".as_bytes()[..],
+            "v=U+ppxD5XUKtradnv8e2MkeupiA8FU87Sg8CXzXHDAzw".as_bytes(),
         ))
         .unwrap();
         assert_eq!(
@@ -309,12 +314,12 @@ mod test {
         assert_eq!(std::str::from_utf8(&message).unwrap(), client_first);
 
         let result = scram
-            .update(&BytesMut::from(&server_first.as_bytes()[..]))
+            .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()[..]))
+            .finish(&BytesMut::from(server_final.as_bytes()))
             .unwrap();
     }
 }

src/server.rs

@@ -14,6 +14,7 @@ use crate::config::{Address, User};
 use crate::constants::*;
 use crate::errors::Error;
 use crate::messages::*;
+use crate::mirrors::MirroringManager;
 use crate::pool::ClientServerMap;
 use crate::scram::ScramSha256;
 use crate::stats::Reporter;
@@ -68,6 +69,8 @@ pub struct Server {
 
     // Last time that a successful server send or response happened
     last_activity: SystemTime,
+
+    mirror_manager: Option<MirroringManager>,
 }
 
 impl Server {
@@ -86,9 +89,13 @@ impl Server {
                 Ok(stream) => stream,
                 Err(err) => {
                     error!("Could not connect to server: {}", err);
-                    return Err(Error::SocketError);
+                    return Err(Error::SocketError(format!(
+                        "Could not connect to server: {}",
+                        err
+                    )));
                 }
             };
+        configure_socket(&stream);
 
         trace!("Sending StartupMessage");
 
@@ -106,12 +113,12 @@ impl Server {
         loop {
             let code = match stream.read_u8().await {
                 Ok(code) => code as char,
-                Err(_) => return Err(Error::SocketError),
+                Err(_) => return Err(Error::SocketError(format!("Error reading message code on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
             };
 
             let len = match stream.read_i32().await {
                 Ok(len) => len,
-                Err(_) => return Err(Error::SocketError),
+                Err(_) => return Err(Error::SocketError(format!("Error reading message len on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
             };
 
             trace!("Message: {}", code);
@@ -122,7 +129,7 @@ impl Server {
                     // Determine which kind of authentication is required, if any.
                     let auth_code = match stream.read_i32().await {
                         Ok(auth_code) => auth_code,
-                        Err(_) => return Err(Error::SocketError),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading auth code on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
 
                     trace!("Auth: {}", auth_code);
@@ -135,7 +142,7 @@ impl Server {
 
                             match stream.read_exact(&mut salt).await {
                                 Ok(_) => (),
-                                Err(_) => return Err(Error::SocketError),
+                                Err(_) => return Err(Error::SocketError(format!("Error reading salt on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                             };
 
                             md5_password(&mut stream, &user.username, &user.password, &salt[..])
@@ -151,7 +158,7 @@ impl Server {
 
                             match stream.read_exact(&mut sasl_auth).await {
                                 Ok(_) => (),
-                                Err(_) => return Err(Error::SocketError),
+                                Err(_) => return Err(Error::SocketError(format!("Error reading sasl message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                             };
 
                             let sasl_type = String::from_utf8_lossy(&sasl_auth[..sasl_len - 2]);
@@ -175,7 +182,7 @@ impl Server {
                                         + sasl_response.len() as i32, // length of SASL response
                                 );
 
-                                res.put_slice(&format!("{}\0", SCRAM_SHA_256).as_bytes()[..]);
+                                res.put_slice(format!("{}\0", SCRAM_SHA_256).as_bytes());
                                 res.put_i32(sasl_response.len() as i32);
                                 res.put(sasl_response);
 
@@ -193,7 +200,7 @@ impl Server {
 
                             match stream.read_exact(&mut sasl_data).await {
                                 Ok(_) => (),
-                                Err(_) => return Err(Error::SocketError),
+                                Err(_) => return Err(Error::SocketError(format!("Error reading sasl cont message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                             };
 
                             let msg = BytesMut::from(&sasl_data[..]);
@@ -214,7 +221,7 @@ impl Server {
                             let mut sasl_final = vec![0u8; len as usize - 8];
                             match stream.read_exact(&mut sasl_final).await {
                                 Ok(_) => (),
-                                Err(_) => return Err(Error::SocketError),
+                                Err(_) => return Err(Error::SocketError(format!("Error reading sasl final message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                             };
 
                             match scram.finish(&BytesMut::from(&sasl_final[..])) {
@@ -240,7 +247,7 @@ impl Server {
                 'E' => {
                     let error_code = match stream.read_u8().await {
                         Ok(error_code) => error_code,
-                        Err(_) => return Err(Error::SocketError),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading error code message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
 
                     trace!("Error: {}", error_code);
@@ -256,7 +263,7 @@ impl Server {
 
                             match stream.read_exact(&mut error).await {
                                 Ok(_) => (),
-                                Err(_) => return Err(Error::SocketError),
+                                Err(_) => return Err(Error::SocketError(format!("Error reading error message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                             };
 
                             // TODO: the error message contains multiple fields; we can decode them and
@@ -275,7 +282,7 @@ impl Server {
 
                     match stream.read_exact(&mut param).await {
                         Ok(_) => (),
-                        Err(_) => return Err(Error::SocketError),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading parameter status message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
 
                     // Save the parameter so we can pass it to the client later.
@@ -292,12 +299,12 @@ impl Server {
                     // 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),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading process id message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
 
                     secret_key = match stream.read_i32().await {
                         Ok(id) => id,
-                        Err(_) => return Err(Error::SocketError),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading secret key message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
                 }
 
@@ -307,7 +314,7 @@ impl Server {
 
                     match stream.read_exact(&mut idle).await {
                         Ok(_) => (),
-                        Err(_) => return Err(Error::SocketError),
+                        Err(_) => return Err(Error::SocketError(format!("Error reading transaction status message on server startup {{ username: {:?}, database: {:?} }}", user.username, database))),
                     };
 
                     let (read, write) = stream.into_split();
@@ -315,21 +322,29 @@ impl Server {
                     let mut server = Server {
                         address: address.clone(),
                         read: BufReader::new(read),
-                        write: write,
+                        write,
                         buffer: BytesMut::with_capacity(8196),
-                        server_info: server_info,
-                        server_id: server_id,
-                        process_id: process_id,
-                        secret_key: secret_key,
+                        server_info,
+                        server_id,
+                        process_id,
+                        secret_key,
                         in_transaction: false,
                         data_available: false,
                         bad: false,
                         needs_cleanup: false,
-                        client_server_map: client_server_map,
+                        client_server_map,
                         connected_at: chrono::offset::Utc::now().naive_utc(),
-                        stats: stats,
+                        stats,
                         application_name: String::new(),
                         last_activity: SystemTime::now(),
+                        mirror_manager: match address.mirrors.len() {
+                            0 => None,
+                            _ => Some(MirroringManager::from_addresses(
+                                user.clone(),
+                                database.to_owned(),
+                                address.mirrors.clone(),
+                            )),
+                        },
                     };
 
                     server.set_name("pgcat").await?;
@@ -341,7 +356,10 @@ impl Server {
                 // Means we implemented the protocol wrong or we're not talking to a Postgres server.
                 _ => {
                     error!("Unknown code: {}", code);
-                    return Err(Error::ProtocolSyncError);
+                    return Err(Error::ProtocolSyncError(format!(
+                        "Unknown server code: {}",
+                        code
+                    )));
                 }
             };
         }
@@ -359,9 +377,10 @@ impl Server {
             Ok(stream) => stream,
             Err(err) => {
                 error!("Could not connect to server: {}", err);
-                return Err(Error::SocketError);
+                return Err(Error::SocketError(format!("Error reading cancel message")));
             }
         };
+        configure_socket(&stream);
 
         debug!("Sending CancelRequest");
 
@@ -371,11 +390,12 @@ impl Server {
         bytes.put_i32(process_id);
         bytes.put_i32(secret_key);
 
-        Ok(write_all(&mut stream, bytes).await?)
+        write_all(&mut stream, bytes).await
     }
 
     /// Send messages to the server from the client.
-    pub async fn send(&mut self, messages: BytesMut) -> Result<(), Error> {
+    pub async fn send(&mut self, messages: &BytesMut) -> Result<(), Error> {
+        self.mirror_send(messages);
         self.stats.data_sent(messages.len(), self.server_id);
 
         match write_all_half(&mut self.write, messages).await {
@@ -438,7 +458,10 @@ impl Server {
                         // Something totally unexpected, this is not a Postgres server we know.
                         _ => {
                             self.bad = true;
-                            return Err(Error::ProtocolSyncError);
+                            return Err(Error::ProtocolSyncError(format!(
+                                "Unknown transaction state: {}",
+                                transaction_state
+                            )));
                         }
                     };
 
@@ -457,7 +480,17 @@ impl Server {
                             // which can leak between clients. This is a best effort to block bad clients
                             // from poisoning a transaction-mode pool by setting inappropriate session variables
                             match command_tag.as_str() {
-                                "SET\0" | "PREPARE\0" => {
+                                "SET\0" => {
+                                    // We don't detect set statements in transactions
+                                    // No great way to differentiate between set and set local
+                                    // As a result, we will miss cases when set statements are used in transactions
+                                    // This will reduce amount of discard statements sent
+                                    if !self.in_transaction {
+                                        debug!("Server connection marked for clean up");
+                                        self.needs_cleanup = true;
+                                    }
+                                }
+                                "PREPARE\0" => {
                                     debug!("Server connection marked for clean up");
                                     self.needs_cleanup = true;
                                 }
@@ -491,9 +524,13 @@ impl Server {
                     break;
                 }
 
-                // CopyData: we are not buffering this one because there will be many more
-                // and we don't know how big this packet could be, best not to take a risk.
-                'd' => break,
+                // CopyData
+                'd' => {
+                    // Don't flush yet, buffer until we reach limit
+                    if self.buffer.len() >= 8196 {
+                        break;
+                    }
+                }
 
                 // CopyDone
                 // Buffer until ReadyForQuery shows up, so don't exit the loop yet.
@@ -523,6 +560,7 @@ impl Server {
     /// If the server is still inside a transaction.
     /// If the client disconnects while the server is in a transaction, we will clean it up.
     pub fn in_transaction(&self) -> bool {
+        debug!("Server in transaction: {}", self.in_transaction);
         self.in_transaction
     }
 
@@ -570,7 +608,7 @@ impl Server {
     pub async fn query(&mut self, query: &str) -> Result<(), Error> {
         let query = simple_query(query);
 
-        self.send(query).await?;
+        self.send(&query).await?;
 
         loop {
             let _ = self.recv().await?;
@@ -595,7 +633,7 @@ impl Server {
             self.query("ROLLBACK").await?;
         }
 
-        // Client disconnected but it perfromed session-altering operations such as
+        // Client disconnected but it performed session-altering operations such as
         // SET statement_timeout to 1 or create a prepared statement. We clear that
         // to avoid leaking state between clients. For performance reasons we only
         // send `DISCARD ALL` if we think the session is altered instead of just sending
@@ -606,7 +644,7 @@ impl Server {
             self.needs_cleanup = false;
         }
 
-        return Ok(());
+        Ok(())
     }
 
     /// A shorthand for `SET application_name = $1`.
@@ -621,7 +659,7 @@ impl Server {
                 .query(&format!("SET application_name = '{}'", name))
                 .await?);
             self.needs_cleanup = needs_cleanup_before;
-            return result;
+            result
         } else {
             Ok(())
         }
@@ -648,6 +686,20 @@ impl Server {
     pub fn mark_dirty(&mut self) {
         self.needs_cleanup = true;
     }
+
+    pub fn mirror_send(&mut self, bytes: &BytesMut) {
+        match self.mirror_manager.as_mut() {
+            Some(manager) => manager.send(bytes),
+            None => (),
+        }
+    }
+
+    pub fn mirror_disconnect(&mut self) {
+        match self.mirror_manager.as_mut() {
+            Some(manager) => manager.disconnect(),
+            None => (),
+        }
+    }
 }
 
 impl Drop for Server {
@@ -655,6 +707,7 @@ impl Drop for Server {
     /// the socket is in non-blocking mode, so it may not be ready
     /// for a write.
     fn drop(&mut self) {
+        self.mirror_disconnect();
         self.stats.server_disconnecting(self.server_id);
 
         let mut bytes = BytesMut::with_capacity(4);

src/sharding.rs

@@ -133,7 +133,7 @@ impl Sharder {
     #[inline]
     fn combine(mut a: u64, b: u64) -> u64 {
         a ^= b
-            .wrapping_add(0x49a0f4dd15e5a8e3 as u64)
+            .wrapping_add(0x49a0f4dd15e5a8e3_u64)
             .wrapping_add(a << 54)
             .wrapping_add(a >> 7);
         a
@@ -141,7 +141,7 @@ impl Sharder {
 
     #[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 a: u32 = 0x9e3779b9_u32 + std::mem::size_of::<u32>() as u32 + 3923095_u32;
         let mut b = a;
         let c = a;
 

src/stats.rs

@@ -245,7 +245,7 @@ impl Default for Reporter {
 impl Reporter {
     /// Create a new Reporter instance.
     pub fn new(tx: Sender<Event>) -> Reporter {
-        Reporter { tx: tx }
+        Reporter { tx }
     }
 
     /// Send statistics to the task keeping track of stats.
@@ -338,9 +338,9 @@ impl Reporter {
         let event = Event {
             name: EventName::ClientRegistered {
                 client_id,
-                pool_name: pool_name.clone(),
-                username: username.clone(),
-                application_name: app_name.clone(),
+                pool_name,
+                username,
+                application_name: app_name,
             },
             value: 1,
         };
@@ -582,7 +582,7 @@ impl Collector {
 
                             let address_stats = address_stat_lookup
                                 .entry(server_info.address_id)
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
                             let counter = address_stats
                                 .entry("total_query_count".to_string())
                                 .or_insert(0);
@@ -618,7 +618,7 @@ impl Collector {
 
                             let address_stats = address_stat_lookup
                                 .entry(server_info.address_id)
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
                             let counter = address_stats
                                 .entry("total_xact_count".to_string())
                                 .or_insert(0);
@@ -636,7 +636,7 @@ impl Collector {
 
                             let address_stats = address_stat_lookup
                                 .entry(server_info.address_id)
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
                             let counter =
                                 address_stats.entry("total_sent".to_string()).or_insert(0);
                             *counter += stat.value;
@@ -653,7 +653,7 @@ impl Collector {
 
                             let address_stats = address_stat_lookup
                                 .entry(server_info.address_id)
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
                             let counter = address_stats
                                 .entry("total_received".to_string())
                                 .or_insert(0);
@@ -683,7 +683,7 @@ impl Collector {
 
                             let address_stats = address_stat_lookup
                                 .entry(server_info.address_id)
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
                             let counter = address_stats
                                 .entry("total_wait_time".to_string())
                                 .or_insert(0);
@@ -694,7 +694,7 @@ impl Collector {
                                     server_info.pool_name.clone(),
                                     server_info.username.clone(),
                                 ))
-                                .or_insert(HashMap::default());
+                                .or_insert_with(HashMap::default);
 
                             // We record max wait in microseconds, we do the pgbouncer second/microsecond split on admin
                             let old_microseconds =
@@ -750,7 +750,7 @@ impl Collector {
                     // Update address aggregation stats
                     let address_stats = address_stat_lookup
                         .entry(address_id)
-                        .or_insert(HashMap::default());
+                        .or_insert_with(HashMap::default);
                     let counter = address_stats.entry("total_errors".to_string()).or_insert(0);
                     *counter += stat.value;
                 }
@@ -770,7 +770,7 @@ impl Collector {
                     // Update address aggregation stats
                     let address_stats = address_stat_lookup
                         .entry(address_id)
-                        .or_insert(HashMap::default());
+                        .or_insert_with(HashMap::default);
                     let counter = address_stats.entry("total_errors".to_string()).or_insert(0);
                     *counter += stat.value;
                 }
@@ -891,7 +891,7 @@ impl Collector {
                 } => {
                     let pool_stats = pool_stat_lookup
                         .entry((pool_name.clone(), username.clone()))
-                        .or_insert(HashMap::default());
+                        .or_insert_with(HashMap::default);
 
                     // These are re-calculated every iteration of the loop, so we don't want to add values
                     // from the last iteration.
@@ -964,17 +964,17 @@ impl Collector {
                     // Clear maxwait after reporting
                     pool_stat_lookup
                         .entry((pool_name.clone(), username.clone()))
-                        .or_insert(HashMap::default())
+                        .or_insert_with(HashMap::default)
                         .insert("maxwait_us".to_string(), 0);
                 }
 
                 EventName::UpdateAverages { address_id } => {
                     let stats = address_stat_lookup
                         .entry(address_id)
-                        .or_insert(HashMap::default());
+                        .or_insert_with(HashMap::default);
                     let old_stats = address_old_stat_lookup
                         .entry(address_id)
-                        .or_insert(HashMap::default());
+                        .or_insert_with(HashMap::default);
 
                     // Calculate averages
                     for stat in &[

src/tls.rs

@@ -1,6 +1,7 @@
 // Stream wrapper.
 
-use rustls_pemfile::{certs, rsa_private_keys};
+use rustls_pemfile::{certs, read_one, Item};
+use std::iter;
 use std::path::Path;
 use std::sync::Arc;
 use tokio_rustls::rustls::{self, Certificate, PrivateKey};
@@ -17,9 +18,17 @@ pub fn load_certs(path: &Path) -> std::io::Result<Vec<Certificate>> {
 }
 
 pub fn load_keys(path: &Path) -> std::io::Result<Vec<PrivateKey>> {
-    rsa_private_keys(&mut std::io::BufReader::new(std::fs::File::open(path)?))
-        .map_err(|_| std::io::Error::new(std::io::ErrorKind::InvalidInput, "invalid key"))
-        .map(|mut keys| keys.drain(..).map(PrivateKey).collect())
+    let mut rd = std::io::BufReader::new(std::fs::File::open(path)?);
+
+    iter::from_fn(|| read_one(&mut rd).transpose())
+        .filter_map(|item| match item {
+            Err(err) => Some(Err(err)),
+            Ok(Item::RSAKey(key)) => Some(Ok(PrivateKey(key))),
+            Ok(Item::ECKey(key)) => Some(Ok(PrivateKey(key))),
+            Ok(Item::PKCS8Key(key)) => Some(Ok(PrivateKey(key))),
+            _ => None,
+        })
+        .collect()
 }
 
 pub struct Tls {
@@ -30,12 +39,12 @@ impl Tls {
     pub fn new() -> Result<Self, Error> {
         let config = get_config();
 
-        let certs = match load_certs(&Path::new(&config.general.tls_certificate.unwrap())) {
+        let certs = match load_certs(Path::new(&config.general.tls_certificate.unwrap())) {
             Ok(certs) => certs,
             Err(_) => return Err(Error::TlsError),
         };
 
-        let mut keys = match load_keys(&Path::new(&config.general.tls_private_key.unwrap())) {
+        let mut keys = match load_keys(Path::new(&config.general.tls_private_key.unwrap())) {
             Ok(keys) => keys,
             Err(_) => return Err(Error::TlsError),
         };

tests/docker/Dockerfile

@@ -1,5 +1,8 @@
 FROM rust:bullseye
 
-RUN apt-get update && apt-get install llvm-11 psmisc postgresql-contrib postgresql-client ruby ruby-dev libpq-dev python3 python3-pip lcov sudo curl -y
+RUN apt-get update && apt-get install llvm-11 psmisc postgresql-contrib postgresql-client ruby ruby-dev libpq-dev python3 python3-pip lcov curl sudo iproute2 -y
 RUN cargo install cargo-binutils rustfilt
 RUN rustup component add llvm-tools-preview
+RUN sudo gem install bundler
+RUN wget -O toxiproxy-2.4.0.deb https://github.com/Shopify/toxiproxy/releases/download/v2.4.0/toxiproxy_2.4.0_linux_$(dpkg --print-architecture).deb && \
+	sudo dpkg -i toxiproxy-2.4.0.deb

tests/docker/docker-compose.yml

@@ -7,8 +7,8 @@ services:
       POSTGRES_USER: postgres
       POSTGRES_DB: postgres
       POSTGRES_PASSWORD: postgres
-      POSTGRES_HOST_AUTH_METHOD: scram-sha-256
-    command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "5432"]
+      POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
+    command: ["postgres", "-p", "5432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
   pg2:
     image: postgres:14
     network_mode: "service:main"
@@ -16,8 +16,8 @@ services:
       POSTGRES_USER: postgres
       POSTGRES_DB: postgres
       POSTGRES_PASSWORD: postgres
-      POSTGRES_HOST_AUTH_METHOD: scram-sha-256
-    command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "7432"]
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+    command: ["postgres", "-p", "7432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
   pg3:
     image: postgres:14
     network_mode: "service:main"
@@ -25,8 +25,8 @@ services:
       POSTGRES_USER: postgres
       POSTGRES_DB: postgres
       POSTGRES_PASSWORD: postgres
-      POSTGRES_HOST_AUTH_METHOD: scram-sha-256
-    command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "8432"]
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+    command: ["postgres", "-p", "8432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
   pg4:
     image: postgres:14
     network_mode: "service:main"
@@ -34,14 +34,11 @@ services:
       POSTGRES_USER: postgres
       POSTGRES_DB: postgres
       POSTGRES_PASSWORD: postgres
-      POSTGRES_HOST_AUTH_METHOD: scram-sha-256
-    command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "9432"]
+      POSTGRES_INITDB_ARGS: --auth-local=scram-sha-256 --auth-host=scram-sha-256 --auth=scram-sha-256
+    command: ["postgres", "-p", "9432", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
   main:
     build: .
     command: ["bash", "/app/tests/docker/run.sh"]
-    environment:
-      RUSTFLAGS: "-C instrument-coverage"
-      LLVM_PROFILE_FILE: "pgcat-%m.profraw"
     volumes:
       - ../../:/app/
       - /app/target/

tests/docker/run.sh

@@ -1,21 +1,37 @@
 #!/bin/bash
 
+rm -rf /app/target/ || true
 rm /app/*.profraw || true
 rm /app/pgcat.profdata || true
 rm -rf /app/cov || true
 
-cd /app/
+export LLVM_PROFILE_FILE="/app/pgcat-%m-%p.profraw"
+export RUSTC_BOOTSTRAP=1
+export CARGO_INCREMENTAL=0
+export RUSTFLAGS="-Zprofile -Ccodegen-units=1 -Copt-level=0 -Clink-dead-code -Coverflow-checks=off -Zpanic_abort_tests -Cpanic=abort -Cinstrument-coverage"
+export RUSTDOCFLAGS="-Cpanic=abort"
 
+cd /app/
+cargo clean
 cargo build
 cargo test --tests
 
 bash .circleci/run_tests.sh
 
-rust-profdata merge -sparse pgcat-*.profraw -o pgcat.profdata
+TEST_OBJECTS=$( \
+    for file in $(cargo test --no-run 2>&1 | grep "target/debug/deps/pgcat-[[:alnum:]]\+" -o); \
+    do \
+        printf "%s %s " --object $file; \
+    done \
+)
 
-rust-cov export -ignore-filename-regex="rustc|registry" -Xdemangler=rustfilt -instr-profile=pgcat.profdata --object ./target/debug/pgcat --format lcov > ./lcov.info
+echo "Generating coverage report"
 
-genhtml lcov.info --output-directory cov --prefix $(pwd)
+rust-profdata merge -sparse /app/pgcat-*.profraw -o /app/pgcat.profdata
+
+bash -c "rust-cov export -ignore-filename-regex='rustc|registry' -Xdemangler=rustfilt -instr-profile=/app/pgcat.profdata $TEST_OBJECTS --object ./target/debug/pgcat --format lcov > ./lcov.info"
+
+genhtml lcov.info --title "PgCat Code Coverage" --css-file ./cov-style.css --highlight --no-function-coverage --ignore-errors source --legend  --output-directory cov --prefix $(pwd)
 
 rm /app/*.profraw
 rm /app/pgcat.profdata

tests/python/tests.py

@@ -110,6 +110,37 @@ def test_shutdown_logic():
     cleanup_conn(conn, cur)
     pg_cat_send_signal(signal.SIGTERM)
 
+    # - - - - - - - - - - - - - - - - - -
+    # NO ACTIVE QUERIES ADMIN SHUTDOWN COMMAND
+
+    # Start pgcat
+    pgcat_start()
+
+    # Create client connection and begin transaction
+    conn, cur = connect_db()
+    admin_conn, admin_cur = connect_db(admin=True)
+
+    cur.execute("BEGIN;")
+    cur.execute("SELECT 1;")
+    cur.execute("COMMIT;")
+
+    # Send SHUTDOWN command pgcat while not in transaction
+    admin_cur.execute("SHUTDOWN;")
+    time.sleep(1)
+
+    # Check that any new queries fail after SHUTDOWN command since server should close with no active transactions
+    try:
+        cur.execute("SELECT 1;")
+    except psycopg2.OperationalError as e:
+        pass
+    else:
+        # Fail if query execution succeeded
+        raise Exception("Server not closed after sigint")
+
+    cleanup_conn(conn, cur)
+    cleanup_conn(admin_conn, admin_cur)
+    pg_cat_send_signal(signal.SIGTERM)
+
     # - - - - - - - - - - - - - - - - - -
     # HANDLE TRANSACTION WITH SIGINT
 
@@ -136,6 +167,36 @@ def test_shutdown_logic():
     cleanup_conn(conn, cur)
     pg_cat_send_signal(signal.SIGTERM)
 
+    # - - - - - - - - - - - - - - - - - -
+    # HANDLE TRANSACTION WITH ADMIN SHUTDOWN COMMAND
+
+    # Start pgcat
+    pgcat_start()
+
+    # Create client connection and begin transaction
+    conn, cur = connect_db()
+    admin_conn, admin_cur = connect_db(admin=True)
+
+    cur.execute("BEGIN;")
+    cur.execute("SELECT 1;")
+
+    # Send SHUTDOWN command pgcat while still in transaction
+    admin_cur.execute("SHUTDOWN;")
+    if admin_cur.fetchall()[0][0] != "t":
+        raise Exception("PgCat unable to send signal")
+    time.sleep(1)
+
+    # Check that any new queries succeed after SHUTDOWN command since server should still allow transaction to complete
+    try:
+        cur.execute("SELECT 1;")
+    except psycopg2.OperationalError as e:
+        # Fail if query fails since server closed
+        raise Exception("Server closed while in transaction", e.pgerror)
+
+    cleanup_conn(conn, cur)
+    cleanup_conn(admin_conn, admin_cur)
+    pg_cat_send_signal(signal.SIGTERM)
+
     # - - - - - - - - - - - - - - - - - -
     # NO NEW NON-ADMIN CONNECTIONS DURING SHUTDOWN
     # Start pgcat

tests/ruby/Gemfile.lock

@@ -1,12 +1,12 @@
 GEM
   remote: https://rubygems.org/
   specs:
-    activemodel (7.0.3.1)
-      activesupport (= 7.0.3.1)
-    activerecord (7.0.3.1)
-      activemodel (= 7.0.3.1)
-      activesupport (= 7.0.3.1)
-    activesupport (7.0.3.1)
+    activemodel (7.0.4.1)
+      activesupport (= 7.0.4.1)
+    activerecord (7.0.4.1)
+      activemodel (= 7.0.4.1)
+      activesupport (= 7.0.4.1)
+    activesupport (7.0.4.1)
       concurrent-ruby (~> 1.0, >= 1.0.2)
       i18n (>= 1.6, < 2)
       minitest (>= 5.1)
@@ -14,9 +14,9 @@ GEM
     ast (2.4.2)
     concurrent-ruby (1.1.10)
     diff-lcs (1.5.0)
-    i18n (1.11.0)
+    i18n (1.12.0)
       concurrent-ruby (~> 1.0)
-    minitest (5.16.2)
+    minitest (5.17.0)
     parallel (1.22.1)
     parser (3.1.2.0)
       ast (~> 2.4.1)
@@ -53,7 +53,7 @@ GEM
     toml (0.3.0)
       parslet (>= 1.8.0, < 3.0.0)
     toxiproxy (2.0.1)
-    tzinfo (2.0.4)
+    tzinfo (2.0.5)
       concurrent-ruby (~> 1.0)
     unicode-display_width (2.1.0)
 

tests/ruby/admin_spec.rb

@@ -221,7 +221,7 @@ describe "Admin" do
         results = admin_conn.async_exec("SHOW POOLS")[0]
 
         expect(results["maxwait"]).to eq("1")
-        expect(results["maxwait_us"].to_i).to be_within(100_000).of(500_000)
+        expect(results["maxwait_us"].to_i).to be_within(200_000).of(500_000)
 
         sleep(4.5) # Allow time for stats to update
         results = admin_conn.async_exec("SHOW POOLS")[0]
@@ -286,4 +286,84 @@ describe "Admin" do
       connections.map(&:close)
     end
   end
+
+  describe "Manual Banning" do
+    let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 10) }
+    before do
+      new_configs = processes.pgcat.current_config
+      # Prevent immediate unbanning when we ban localhost
+      new_configs["pools"]["sharded_db"]["shards"]["0"]["servers"][0][0] = "127.0.0.1"
+      new_configs["pools"]["sharded_db"]["shards"]["0"]["servers"][1][0] = "127.0.0.1"
+      processes.pgcat.update_config(new_configs)
+      processes.pgcat.reload_config
+    end
+
+    describe "BAN/UNBAN and SHOW BANS" do
+      it "bans/unbans hosts" do
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+
+        # Returns a list of the banned addresses
+        results = admin_conn.async_exec("BAN localhost 10").to_a
+        expect(results.count).to eq(2)
+        expect(results.map{ |r| r["host"] }.uniq).to eq(["localhost"])
+
+        # Subsequent calls should yield no results
+        results = admin_conn.async_exec("BAN localhost 10").to_a
+        expect(results.count).to eq(0)
+
+        results = admin_conn.async_exec("SHOW BANS").to_a
+        expect(results.count).to eq(2)
+        expect(results.map{ |r| r["host"] }.uniq).to eq(["localhost"])
+
+        # Returns a list of the unbanned addresses
+        results = admin_conn.async_exec("UNBAN localhost").to_a
+        expect(results.count).to eq(2)
+        expect(results.map{ |r| r["host"] }.uniq).to eq(["localhost"])
+
+        # Subsequent calls should yield no results
+        results = admin_conn.async_exec("UNBAN localhost").to_a
+        expect(results.count).to eq(0)
+
+        results = admin_conn.async_exec("SHOW BANS").to_a
+        expect(results.count).to eq(0)
+      end
+
+      it "honors ban duration" do
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+
+        # Returns a list of the banned addresses
+        results = admin_conn.async_exec("BAN localhost 1").to_a
+        expect(results.count).to eq(2)
+        expect(results.map{ |r| r["host"] }.uniq).to eq(["localhost"])
+
+        sleep(2)
+
+        # After 2 seconds the ban should be lifted
+        results = admin_conn.async_exec("SHOW BANS").to_a
+        expect(results.count).to eq(0)
+      end
+
+      it "can handle bad input" do
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+
+        expect { admin_conn.async_exec("BAN").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("BAN a").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("BAN a a").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("BAN a -5").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("BAN a 0").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("BAN a a a").to_a }.to raise_error(PG::SystemError)
+        expect { admin_conn.async_exec("UNBAN").to_a }.to raise_error(PG::SystemError)
+      end
+    end
+  end
+
+  describe "SHOW users" do
+    it "returns the right users" do
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      results = admin_conn.async_exec("SHOW USERS")[0]
+      admin_conn.close
+      expect(results["name"]).to eq("sharding_user")
+      expect(results["pool_mode"]).to eq("transaction")
+    end
+  end
 end

tests/ruby/helpers/pg_instance.rb

@@ -38,6 +38,8 @@ class PgInstance
   def reset
     reset_toxics
     reset_stats
+    drop_connections
+    sleep 0.1
   end
 
   def toxiproxy
@@ -66,12 +68,22 @@ class PgInstance
 
   def reset_toxics
     Toxiproxy[@toxiproxy_name].toxics.each(&:destroy)
+    sleep 0.1
   end
 
   def reset_stats
     with_connection { |c| c.async_exec("SELECT pg_stat_statements_reset()") }
   end
 
+  def drop_connections
+    username = with_connection { |c| c.async_exec("SELECT current_user")[0]["current_user"] }
+    with_connection { |c| c.async_exec("SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE pid <> pg_backend_pid() AND usename='#{username}'") }
+  end
+
+  def count_connections
+    with_connection { |c| c.async_exec("SELECT COUNT(*) as count FROM pg_stat_activity")[0]["count"].to_i }
+  end
+
   def count_query(query)
     with_connection { |c| c.async_exec("SELECT SUM(calls) FROM pg_stat_statements WHERE query = '#{query}'")[0]["sum"].to_i }
   end

tests/ruby/helpers/pgcat_helper.rb

@@ -5,7 +5,7 @@ require_relative 'pg_instance'
 
 module Helpers
   module Pgcat
-    def self.three_shard_setup(pool_name, pool_size, pool_mode="transaction")
+    def self.three_shard_setup(pool_name, pool_size, pool_mode="transaction", lb_mode="random", log_level="info")
       user = {
         "password" => "sharding_user",
         "pool_size" => pool_size,
@@ -13,7 +13,7 @@ module Helpers
         "username" => "sharding_user"
       }
 
-      pgcat    = PgcatProcess.new("info")
+      pgcat    = PgcatProcess.new(log_level)
       primary0 = PgInstance.new(5432, user["username"], user["password"], "shard0")
       primary1 = PgInstance.new(7432, user["username"], user["password"], "shard1")
       primary2 = PgInstance.new(8432, user["username"], user["password"], "shard2")
@@ -23,8 +23,10 @@ module Helpers
         "#{pool_name}" => {
           "default_role" => "any",
           "pool_mode" => pool_mode,
-          "primary_reads_enabled" => false,
-          "query_parser_enabled" => false,
+          "load_balancing_mode" => lb_mode,
+          "primary_reads_enabled" => true,
+          "query_parser_enabled" => true,
+          "automatic_sharding_key" => "data.id",
           "sharding_function" => "pg_bigint_hash",
           "shards" => {
             "0" => { "database" => "shard0", "servers" => [["localhost", primary0.port.to_s, "primary"]] },
@@ -46,7 +48,7 @@ module Helpers
       end
     end
 
-    def self.single_instance_setup(pool_name, pool_size, pool_mode="transaction")
+    def self.single_instance_setup(pool_name, pool_size, pool_mode="transaction", lb_mode="random", log_level="trace")
       user = {
         "password" => "sharding_user",
         "pool_size" => pool_size,
@@ -54,7 +56,7 @@ module Helpers
         "username" => "sharding_user"
       }
 
-      pgcat = PgcatProcess.new("trace")
+      pgcat = PgcatProcess.new(log_level)
       pgcat_cfg = pgcat.current_config
 
       primary  = PgInstance.new(5432, user["username"], user["password"], "shard0")
@@ -64,6 +66,7 @@ module Helpers
         "#{pool_name}" => {
           "default_role" => "primary",
           "pool_mode" => pool_mode,
+          "load_balancing_mode" => lb_mode,
           "primary_reads_enabled" => false,
           "query_parser_enabled" => false,
           "sharding_function" => "pg_bigint_hash",
@@ -90,7 +93,7 @@ module Helpers
       end
     end
 
-    def self.single_shard_setup(pool_name, pool_size, pool_mode="transaction")
+    def self.single_shard_setup(pool_name, pool_size, pool_mode="transaction", lb_mode="random", log_level="info")
       user = {
         "password" => "sharding_user",
         "pool_size" => pool_size,
@@ -98,7 +101,7 @@ module Helpers
         "username" => "sharding_user"
       }
 
-      pgcat = PgcatProcess.new("info")
+      pgcat = PgcatProcess.new(log_level)
       pgcat_cfg = pgcat.current_config
 
       primary  = PgInstance.new(5432, user["username"], user["password"], "shard0")
@@ -111,6 +114,7 @@ module Helpers
         "#{pool_name}" => {
           "default_role" => "any",
           "pool_mode" => pool_mode,
+          "load_balancing_mode" => lb_mode,
           "primary_reads_enabled" => false,
           "query_parser_enabled" => false,
           "sharding_function" => "pg_bigint_hash",

tests/ruby/helpers/pgcat_process.rb

@@ -8,7 +8,11 @@ class PgcatProcess
   attr_reader :pid
 
   def self.finalize(pid, log_filename, config_filename)
-    `kill #{pid}`
+    if pid
+      Process.kill("TERM", pid)
+      Process.wait(pid)
+    end
+
     File.delete(config_filename) if File.exist?(config_filename)
     File.delete(log_filename) if File.exist?(log_filename)
   end
@@ -20,7 +24,13 @@ class PgcatProcess
     @log_filename = "/tmp/pgcat_log_#{SecureRandom.urlsafe_base64}.log"
     @config_filename = "/tmp/pgcat_cfg_#{SecureRandom.urlsafe_base64}.toml"
 
-    @command = "../../target/debug/pgcat #{@config_filename}"
+    command_path = if ENV['CARGO_TARGET_DIR'] then
+                     "#{ENV['CARGO_TARGET_DIR']}/debug/pgcat"
+                   else
+                     '../../target/debug/pgcat'
+                   end
+
+    @command = "#{command_path} #{@config_filename}"
 
     FileUtils.cp("../../pgcat.toml", @config_filename)
     cfg = current_config
@@ -38,18 +48,20 @@ class PgcatProcess
     @original_config = current_config
     output_to_write = TOML::Generator.new(config_hash).body
     output_to_write = output_to_write.gsub(/,\s*["|'](\d+)["|']\s*,/, ',\1,')
+    output_to_write = output_to_write.gsub(/,\s*["|'](\d+)["|']\s*\]/, ',\1]')
     File.write(@config_filename, output_to_write)
   end
 
   def current_config
-    old_cfg = File.read(@config_filename)
-    loadable_string = old_cfg.gsub(/,\s*(\d+)\s*,/, ', "\1",')
+    loadable_string = File.read(@config_filename)
+    loadable_string = loadable_string.gsub(/,\s*(\d+)\s*,/,  ', "\1",')
+    loadable_string = loadable_string.gsub(/,\s*(\d+)\s*\]/, ', "\1"]')
     TOML.load(loadable_string)
   end
 
   def reload_config
     `kill -s HUP #{@pid}`
-    sleep 0.1
+    sleep 0.5
   end
 
   def start
@@ -75,8 +87,11 @@ class PgcatProcess
   end
 
   def stop
-    `kill #{@pid}`
-    sleep 0.1
+    return unless @pid
+
+    Process.kill("TERM", @pid)
+    Process.wait(@pid)
+    @pid = nil
   end
 
   def shutdown

tests/ruby/load_balancing_spec.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 require_relative 'spec_helper'
 
-describe "Load Balancing" do
+describe "Random Load Balancing" do
   let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5) }
   after do
     processes.all_databases.map(&:reset)
@@ -46,7 +46,110 @@ describe "Load Balancing" do
         end
       end
 
-      expect(failed_count).to eq(2)
+      processes.all_databases.each do |instance|
+        queries_routed = instance.count_select_1_plus_2
+        if processes.replicas[0..1].include?(instance)
+          expect(queries_routed).to eq(0)
+        else
+          expect(queries_routed).to be_within(expected_share * MARGIN_OF_ERROR).of(expected_share)
+        end
+      end
+    end
+  end
+end
+
+describe "Least Outstanding Queries Load Balancing" do
+  let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 1, "transaction", "loc") }
+  after do
+    processes.all_databases.map(&:reset)
+    processes.pgcat.shutdown
+  end
+
+  context "under homogenous load" do
+    it "balances query volume between all instances" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+
+      query_count = QUERY_COUNT
+      expected_share = query_count / processes.all_databases.count
+      failed_count = 0
+
+      query_count.times do
+        conn.async_exec("SELECT 1 + 2")
+      rescue
+        failed_count += 1
+      end
+
+      expect(failed_count).to eq(0)
+      processes.all_databases.map(&:count_select_1_plus_2).each do |instance_share|
+        expect(instance_share).to be_within(expected_share * MARGIN_OF_ERROR).of(expected_share)
+      end
+    end
+  end
+
+  context "under heterogeneous load" do
+    xit "balances query volume between all instances based on how busy they are" do
+      slow_query_count = 2
+      threads = Array.new(slow_query_count) do
+        Thread.new do
+          conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("BEGIN")
+        end
+      end
+
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+
+      query_count = QUERY_COUNT
+      expected_share = query_count / (processes.all_databases.count - slow_query_count)
+      failed_count = 0
+
+      query_count.times do
+        conn.async_exec("SELECT 1 + 2")
+      rescue
+        failed_count += 1
+      end
+
+      expect(failed_count).to eq(0)
+      # Under LOQ, we expect replicas running the slow pg_sleep
+      # to get no selects
+      expect(
+        processes.
+          all_databases.
+          map(&:count_select_1_plus_2).
+          count { |instance_share| instance_share == 0 }
+      ).to eq(slow_query_count)
+
+      # We also expect the quick queries to be spread across
+      # the idle servers only
+      processes.
+        all_databases.
+        map(&:count_select_1_plus_2).
+        reject { |instance_share| instance_share == 0 }.
+        each do |instance_share|
+          expect(instance_share).to be_within(expected_share * MARGIN_OF_ERROR).of(expected_share)
+      end
+
+      threads.map(&:join)
+    end
+  end
+
+  context "when some replicas are down" do
+    it "balances query volume between working instances" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      expected_share = QUERY_COUNT / (processes.all_databases.count - 2)
+      failed_count = 0
+
+      processes[:replicas][0].take_down do
+        processes[:replicas][1].take_down do
+          QUERY_COUNT.times do
+            conn.async_exec("SELECT 1 + 2")
+          rescue
+            conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+            failed_count += 1
+          end
+        end
+      end
+
+      expect(failed_count).to be <= 2
       processes.all_databases.each do |instance|
         queries_routed = instance.count_select_1_plus_2
         if processes.replicas[0..1].include?(instance)

tests/ruby/mirrors_spec.rb

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+require 'uri'
+require_relative 'spec_helper'
+
+describe "Query Mirroing" do
+  let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 10) }
+  let(:mirror_pg) { PgInstance.new(8432, "sharding_user", "sharding_user", "shard2")}
+  let(:pgcat_conn_str) { processes.pgcat.connection_string("sharded_db", "sharding_user") }
+  let(:mirror_host) { "localhost" }
+
+  before do
+    new_configs = processes.pgcat.current_config
+    new_configs["pools"]["sharded_db"]["shards"]["0"]["mirrors"] = [
+      [mirror_host, mirror_pg.port.to_s, "0"],
+      [mirror_host, mirror_pg.port.to_s, "0"],
+      [mirror_host, mirror_pg.port.to_s, "0"],
+    ]
+    processes.pgcat.update_config(new_configs)
+    processes.pgcat.reload_config
+  end
+
+  after do
+    processes.all_databases.map(&:reset)
+    mirror_pg.reset
+    processes.pgcat.shutdown
+  end
+
+  it "can mirror a query" do
+    conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+    runs = 15
+    runs.times { conn.async_exec("SELECT 1 + 2") }
+    sleep 0.5
+    expect(processes.all_databases.first.count_select_1_plus_2).to eq(runs)
+    expect(mirror_pg.count_select_1_plus_2).to eq(runs * 3)
+  end
+
+  context "when main server connection is closed" do
+    it "closes the mirror connection" do
+      baseline_count = processes.all_databases.first.count_connections
+      5.times do |i|
+        # Force pool cycling to detect zombie mirror connections
+        new_configs = processes.pgcat.current_config
+        new_configs["pools"]["sharded_db"]["idle_timeout"] = 5000 + i
+        new_configs["pools"]["sharded_db"]["shards"]["0"]["mirrors"] = [
+          [mirror_host, mirror_pg.port.to_s, "0"],
+          [mirror_host, mirror_pg.port.to_s, "0"],
+          [mirror_host, mirror_pg.port.to_s, "0"],
+        ]
+        processes.pgcat.update_config(new_configs)
+        processes.pgcat.reload_config
+      end
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      conn.async_exec("SELECT 1 + 2")
+      sleep 0.5
+      # Expect same number of connection even after pool cycling
+      expect(processes.all_databases.first.count_connections).to be < baseline_count + 2
+    end
+  end
+
+  xcontext "when mirror server goes down temporarily" do
+    it "continues to transmit queries after recovery" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      mirror_pg.take_down do
+        conn.async_exec("SELECT 1 + 2")
+        sleep 0.1
+      end
+      10.times { conn.async_exec("SELECT 1 + 2") }
+      sleep 1
+      expect(mirror_pg.count_select_1_plus_2).to be >= 2
+    end
+  end
+
+  context "when a mirror is down" do
+    let(:mirror_host) { "badhost" }
+
+    it "does not fail to send the main query" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      # No Errors here
+      conn.async_exec("SELECT 1 + 2")
+      expect(processes.all_databases.first.count_select_1_plus_2).to eq(1)
+    end
+
+    it "does not fail to send the main query (even after thousands of mirror attempts)" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      # No Errors here
+      1000.times { conn.async_exec("SELECT 1 + 2") }
+      expect(processes.all_databases.first.count_select_1_plus_2).to eq(1000)
+    end
+  end
+end

tests/ruby/misc_spec.rb

@@ -8,6 +8,100 @@ describe "Miscellaneous" do
     processes.pgcat.shutdown
   end
 
+  context "when adding then removing instance using RELOAD" do
+    it "works correctly" do
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+
+      current_configs = processes.pgcat.current_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+
+      extra_replica = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].last.clone
+      extra_replica[0] = "127.0.0.1"
+      current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"] << extra_replica
+
+      processes.pgcat.update_config(current_configs) # with replica added
+      processes.pgcat.reload_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+
+      current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].pop
+
+      processes.pgcat.update_config(current_configs) # with replica removed again
+      processes.pgcat.reload_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+    end
+  end
+
+  context "when removing then adding instance back using RELOAD" do
+    it "works correctly" do
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+
+      current_configs = processes.pgcat.current_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+
+      removed_replica = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].pop
+      processes.pgcat.update_config(current_configs) # with replica removed
+      processes.pgcat.reload_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+
+      current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"] << removed_replica
+
+      processes.pgcat.update_config(current_configs) # with replica added again
+      processes.pgcat.reload_config
+      correct_count = current_configs["pools"]["sharded_db"]["shards"]["0"]["servers"].count
+      expect(admin_conn.async_exec("SHOW DATABASES").count).to eq(correct_count)
+    end
+  end
+
+  describe "TCP Keepalives" do
+    # Ideally, we should block TCP traffic to the database using
+    # iptables to mimic passive (connection is dropped without a RST packet)
+    # but we cannot do this in CircleCI because iptables requires NET_ADMIN
+    # capability that we cannot enable in CircleCI
+    # Toxiproxy won't work either because it does not block keepalives
+    # so our best bet is to query the OS keepalive params set on the socket
+
+    context "default settings" do
+      it "applies default keepalive settings" do
+        # We query ss command to verify that we have correct keepalive values set
+        # we can only verify the keepalives_idle parameter but that's good enough
+        # example output
+        #Recv-Q Send-Q Local Address:Port  Peer Address:Port Process
+        #0      0          127.0.0.1:60526    127.0.0.1:18432 timer:(keepalive,1min59sec,0)
+        #0      0          127.0.0.1:60664    127.0.0.1:19432 timer:(keepalive,4.123ms,0)
+
+        port_search_criteria = processes.all_databases.map { |d| "dport = :#{d.port}"}.join(" or ")
+        results = `ss -t4 state established -o -at '( #{port_search_criteria}  )'`.lines
+        results.shift
+        results.each { |line| expect(line).to match(/timer:\(keepalive,.*ms,0\)/) }
+      end
+    end
+
+    context "changed settings" do
+      it "applies keepalive settings from config" do
+        new_configs = processes.pgcat.current_config
+
+        new_configs["general"]["tcp_keepalives_idle"] = 120
+        new_configs["general"]["tcp_keepalives_count"] = 1
+        new_configs["general"]["tcp_keepalives_interval"] = 1
+        processes.pgcat.update_config(new_configs)
+        # We need to kill the old process that was using the default configs
+        processes.pgcat.stop
+        processes.pgcat.start
+        processes.pgcat.wait_until_ready
+
+        port_search_criteria = processes.all_databases.map { |d| "dport = :#{d.port}"}.join(" or ")
+        results = `ss -t4 state established -o -at '( #{port_search_criteria}  )'`.lines
+        results.shift
+        results.each { |line| expect(line).to include("timer:(keepalive,1min") }
+      end
+    end
+  end
+
   describe "Extended Protocol handling" do
     it "does not send packets that client does not expect during extended protocol sequence" do
       new_configs = processes.pgcat.current_config
@@ -189,5 +283,30 @@ describe "Miscellaneous" do
         expect(processes.primary.count_query("DISCARD ALL")).to eq(10)
       end
     end
+
+    context "transaction mode with transactions" do
+      let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5, "transaction") }
+      it "Does not clear set statement state when declared in a transaction" do
+        10.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("SET SERVER ROLE to 'primary'")
+          conn.async_exec("BEGIN")
+          conn.async_exec("SET statement_timeout to 1000")
+          conn.async_exec("COMMIT")
+          conn.close
+        end
+        expect(processes.primary.count_query("DISCARD ALL")).to eq(0)
+
+        10.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("SET SERVER ROLE to 'primary'")
+          conn.async_exec("BEGIN")
+          conn.async_exec("SET LOCAL statement_timeout to 1000")
+          conn.async_exec("COMMIT")
+          conn.close
+        end
+        expect(processes.primary.count_query("DISCARD ALL")).to eq(0)
+      end
+    end
   end
 end

tests/ruby/sharding_spec.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+require_relative 'spec_helper'
+
+
+describe "Sharding" do
+  let(:processes) { Helpers::Pgcat.three_shard_setup("sharded_db", 5) }
+
+  before do
+    conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+
+    # Setup the sharding data
+    3.times do |i|
+      conn.exec("SET SHARD TO '#{i}'")
+      conn.exec("DELETE FROM data WHERE id > 0")
+    end
+
+    18.times do |i|
+      i = i + 1
+      conn.exec("SET SHARDING KEY TO '#{i}'")
+      conn.exec("INSERT INTO data (id, value) VALUES (#{i}, 'value_#{i}')")
+    end
+  end
+
+  after do
+
+    processes.all_databases.map(&:reset)
+    processes.pgcat.shutdown
+  end
+
+  describe "automatic routing of extended procotol" do
+    it "can do it" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      conn.exec("SET SERVER ROLE TO 'auto'")
+
+      18.times do |i|
+        result = conn.exec_params("SELECT * FROM data WHERE id = $1", [i + 1])
+        expect(result.ntuples).to eq(1)
+      end
+    end
+
+    it "can do it with multiple parameters" do
+      conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      conn.exec("SET SERVER ROLE TO 'auto'")
+
+      18.times do |i|
+        result = conn.exec_params("SELECT * FROM data WHERE id = $1 AND id = $2", [i + 1, i + 1])
+        expect(result.ntuples).to eq(1)
+      end
+    end
+  end
+end

tests/ruby/spec_helper.rb

@@ -4,7 +4,7 @@ require 'pg'
 require_relative 'helpers/pgcat_helper'
 
 QUERY_COUNT = 300
-MARGIN_OF_ERROR = 0.30
+MARGIN_OF_ERROR = 0.35
 
 def with_captured_stdout_stderr
   sout = STDOUT.clone

utilities/generate_config_docs.py

@@ -0,0 +1,92 @@
+import re
+import tomli
+
+class DocGenerator:
+    def __init__(self, filename):
+        self.doc = []
+        self.current_section = ""
+        self.current_comment = []
+        self.current_field_name = ""
+        self.current_field_value = []
+        self.current_field_unset = False
+        self.filename = filename
+
+    def write(self):
+        with open("../CONFIG.md", "w") as text_file:
+            text_file.write("# PgCat Configurations \n")
+            for entry in self.doc:
+                if entry["name"] == "__section__":
+                    text_file.write("## `" + entry["section"] + "` Section" + "\n")
+                    text_file.write("\n")
+                    continue
+                text_file.write("### " + entry["name"]+ "\n")
+                text_file.write("```"+ "\n")
+                text_file.write("path: " + entry["fqdn"]+ "\n")
+                text_file.write("default: " + entry["defaults"].strip()+ "\n")
+                if entry["example"] is not None:
+                    text_file.write("example: " + entry["example"].strip()+ "\n")
+                text_file.write("```"+ "\n")
+                text_file.write("\n")
+                text_file.write(entry["comment"]+ "\n")
+                text_file.write("\n")
+
+    def save_entry(self):
+        if len(self.current_field_name) == 0:
+            return
+        if len(self.current_comment) == 0:
+            return
+        self.current_section = self.current_section.replace("sharded_db", "<pool_name>")
+        self.current_section = self.current_section.replace("simple_db", "<pool_name>")
+        self.current_section = self.current_section.replace("users.0", "users.<user_index>")
+        self.current_section = self.current_section.replace("users.1", "users.<user_index>")
+        self.current_section = self.current_section.replace("shards.0", "shards.<shard_index>")
+        self.current_section = self.current_section.replace("shards.1", "shards.<shard_index>")
+        self.doc.append(
+            {
+                "name": self.current_field_name,
+                "fqdn": self.current_section + "." + self.current_field_name,
+                "section": self.current_section,
+                "comment": "\n".join(self.current_comment),
+                "defaults": self.current_field_value if not self.current_field_unset else "<UNSET>",
+                "example": self.current_field_value  if self.current_field_unset  else None
+            }
+        )
+        self.current_comment = []
+        self.current_field_name = ""
+        self.current_field_value = []
+    def parse(self):
+        with open("../pgcat.toml", "r") as f:
+            for line in f.readlines():
+                line = line.strip()
+                if len(line) == 0:
+                    self.save_entry()
+
+                if line.startswith("["):
+                    self.current_section = line[1:-1]
+                    self.current_field_name = "__section__"
+                    self.current_field_unset = False
+                    self.save_entry()
+
+                elif line.startswith("#"):
+                    results = re.search("^#\s*([A-Za-z0-9_]+)\s*=(.+)$", line)
+                    if results is not None:
+                        self.current_field_name = results.group(1)
+                        self.current_field_value = results.group(2)
+                        self.current_field_unset = True
+                        self.save_entry()
+                    else:
+                        self.current_comment.append(line[1:].strip())
+                else:
+                    results = re.search("^\s*([A-Za-z0-9_]+)\s*=(.+)$", line)
+                    if results is None:
+                        continue
+                    self.current_field_name = results.group(1)
+                    self.current_field_value = results.group(2)
+                    self.current_field_unset = False
+                    self.save_entry()
+        self.save_entry()
+        return self
+
+
+DocGenerator("../pgcat.toml").parse().write()
+