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,20 +9,43 @@ 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
-        # auth:
-        #   username: mydockerhub-user
-        #   password: $DOCKERHUB_PASSWORD
+        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", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+        environment:
+          POSTGRES_USER: postgres
+          POSTGRES_DB: postgres
+          POSTGRES_PASSWORD: postgres
+          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", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+        environment:
+          POSTGRES_USER: postgres
+          POSTGRES_DB: postgres
+          POSTGRES_PASSWORD: postgres
+          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", "-c", "pg_stat_statements.track=all", "-c", "pg_stat_statements.max=100000"]
+        environment:
+          POSTGRES_USER: postgres
+          POSTGRES_DB: postgres
+          POSTGRES_PASSWORD: postgres
+          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
     steps:
@@ -32,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
 
@@ -108,7 +114,6 @@ servers = [
 ]
 # Database name (e.g. "postgres")
 database = "shard0"
-search_path = "\"$user\",public"
 
 [pools.sharded_db.shards.1]
 servers = [

.circleci/run_tests.sh

@@ -16,17 +16,16 @@ function start_pgcat() {
 
 # Setup the database with shards and user
 PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
+PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 7432 -U postgres -f tests/sharding/query_routing_setup.sql
+PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 8432 -U postgres -f tests/sharding/query_routing_setup.sql
+PGPASSWORD=postgres psql -e -h 127.0.0.1 -p 9432 -U postgres -f tests/sharding/query_routing_setup.sql
 
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard0 -i
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard1 -i
 PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i
 
-# Install Toxiproxy to simulate a downed/slow database
-wget -O toxiproxy-2.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
-toxiproxy-server &
+LOG_LEVEL=error toxiproxy-server &
 sleep 1
 
 # Create a database at port 5433, forward it to Postgres
@@ -53,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'
 
@@ -82,12 +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
-ruby tests.rb
+sudo bundle install
+bundle exec ruby tests.rb --format documentation || exit 1
+bundle exec rspec *_spec.rb --format documentation || exit 1
 cd ../..
 
 #
@@ -95,7 +105,7 @@ cd ../..
 # These tests will start and stop the pgcat server so it will need to be restarted after the tests
 #
 pip3 install -r tests/python/requirements.txt
-python3 tests/python/tests.py
+python3 tests/python/tests.py || exit 1
 
 start_pgcat "info"
 
@@ -105,9 +115,9 @@ psql -U admin_user -e -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW STATS' > /dev/n
 psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'RELOAD' > /dev/null
 psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW CONFIG' > /dev/null
 psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW DATABASES' > /dev/null
-psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW LISTS' > /dev/null
-psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW POOLS' > /dev/null
-psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c 'SHOW VERSION' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgcat -c 'SHOW LISTS' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgcat -c 'SHOW POOLS' > /dev/null
+psql -U admin_user -h 127.0.0.1 -p 6432 -d pgcat -c 'SHOW VERSION' > /dev/null
 psql -U admin_user -h 127.0.0.1 -p 6432 -d pgbouncer -c "SET client_encoding TO 'utf8'" > /dev/null # will ignore
 (! psql -U admin_user -e -h 127.0.0.1 -p 6432 -d random_db -c 'SHOW STATS' > /dev/null)
 export PGPASSWORD=sharding_user

.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

@@ -9,28 +9,34 @@ edition = "2021"
 tokio = { version = "1", features = ["full"] }
 bytes = "1"
 md-5 = "0.10"
-bb8 = "0.7"
+bb8 = "0.8.0"
 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.14"
+sqlparser = "0.32.0"
 log = "0.4"
 arc-swap = "1"
-env_logger = "0.9"
-parking_lot = "0.11"
+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.10", features = ["macros"] }
+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,8 +1,9 @@
-# PgCat
+##### 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)
-
-![PgCat](./pgcat3.png)
+<a href="https://discord.gg/DmyJP3qJ7U" target="_blank">
+    <img src="https://img.shields.io/discord/1013868243036930099" alt="Join our Discord!" />
+</a>
 
 PostgreSQL pooler (like PgBouncer) with sharding, load balancing and failover support.
 
@@ -33,39 +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`                        |
-|                              |                                                                                                                                            |                                  |
-| **`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
 
@@ -87,6 +61,14 @@ pgbench -t 1000 -p 6432 -h 127.0.0.1 --protocol extended
 
 See [sharding README](./tests/sharding/README.md) for sharding logic testing.
 
+Run `cargo test` to run Rust tests.
+
+Run the following commands to run Integration tests locally.
+```
+cd tests/docker/
+docker compose up --exit-code-from main # This will also produce coverage report under ./cov/
+```
+
 | **Feature**           | **Tested in CI**   | **Tested manually** | **Comments**                                                                                                             |
 |-----------------------|--------------------|---------------------|--------------------------------------------------------------------------------------------------------------------------|
 | Transaction pooling   | :white_check_mark: | :white_check_mark:  | Used by default for all tests.                                                                                           |
@@ -99,6 +81,17 @@ See [sharding README](./tests/sharding/README.md) for sharding logic testing.
 | 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
@@ -178,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
@@ -447,7 +440,7 @@ Always good to have a base line.
 
 ```
 $ pgbench -t 1000 -c 16 -j 2 -p 5432 -h 127.0.0.1 -S --protocol extended shard0
-Password: 
+Password:
 starting vacuum...end.
 transaction type: <builtin: select only>
 scaling factor: 1
@@ -461,7 +454,7 @@ tps = 139443.955722 (including connections establishing)
 tps = 142314.859075 (excluding connections establishing)
 
 $ pgbench -t 1000 -c 32 -j 2 -p 5432 -h 127.0.0.1 -S --protocol extended shard0
-Password: 
+Password:
 starting vacuum...end.
 transaction type: <builtin: select only>
 scaling factor: 1
@@ -475,7 +468,7 @@ tps = 150644.840891 (including connections establishing)
 tps = 152218.499430 (excluding connections establishing)
 
 $ pgbench -t 1000 -c 64 -j 2 -p 5432 -h 127.0.0.1 -S --protocol extended shard0
-Password: 
+Password:
 starting vacuum...end.
 transaction type: <builtin: select only>
 scaling factor: 1
@@ -489,7 +482,7 @@ tps = 152517.663404 (including connections establishing)
 tps = 153319.188482 (excluding connections establishing)
 
 $ pgbench -t 1000 -c 128 -j 2 -p 5432 -h 127.0.0.1 -S --protocol extended shard0
-Password: 
+Password:
 starting vacuum...end.
 transaction type: <builtin: select only>
 scaling factor: 1

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,25 +1,32 @@
+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::stats::get_stats;
+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,
+};
 use crate::ClientServerMap;
 
 pub fn generate_server_info_for_admin() -> BytesMut {
     let mut server_info = BytesMut::new();
 
-    server_info.put(server_paramater_message("application_name", ""));
-    server_info.put(server_paramater_message("client_encoding", "UTF8"));
-    server_info.put(server_paramater_message("server_encoding", "UTF8"));
-    server_info.put(server_paramater_message("server_version", VERSION));
-    server_info.put(server_paramater_message("DateStyle", "ISO, MDY"));
+    server_info.put(server_parameter_message("application_name", ""));
+    server_info.put(server_parameter_message("client_encoding", "UTF8"));
+    server_info.put(server_parameter_message("server_encoding", "UTF8"));
+    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.
@@ -34,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
@@ -55,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
@@ -72,6 +104,14 @@ where
                 trace!("SHOW POOLS");
                 show_pools(stream).await
             }
+            "CLIENTS" => {
+                trace!("SHOW CLIENTS");
+                show_clients(stream).await
+            }
+            "SERVERS" => {
+                trace!("SHOW SERVERS");
+                show_servers(stream).await
+            }
             "STATS" => {
                 trace!("SHOW STATS");
                 show_stats(stream).await
@@ -80,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,
@@ -91,7 +135,8 @@ async fn show_lists<T>(stream: &mut T) -> Result<(), Error>
 where
     T: tokio::io::AsyncWrite + std::marker::Unpin,
 {
-    let stats = get_stats();
+    let client_stats = get_client_stats();
+    let server_stats = get_server_stats();
 
     let columns = vec![("list", DataType::Text), ("items", DataType::Int4)];
 
@@ -111,18 +156,18 @@ where
     res.put(data_row(&vec!["pools".to_string(), databases.to_string()]));
     res.put(data_row(&vec![
         "free_clients".to_string(),
-        stats
+        client_stats
             .keys()
-            .map(|address_id| stats[&address_id]["cl_idle"])
-            .sum::<i64>()
+            .filter(|client_id| client_stats.get(client_id).unwrap().state == ClientState::Idle)
+            .count()
             .to_string(),
     ]));
     res.put(data_row(&vec![
         "used_clients".to_string(),
-        stats
+        client_stats
             .keys()
-            .map(|address_id| stats[&address_id]["cl_active"])
-            .sum::<i64>()
+            .filter(|client_id| client_stats.get(client_id).unwrap().state == ClientState::Active)
+            .count()
             .to_string(),
     ]));
     res.put(data_row(&vec![
@@ -131,18 +176,18 @@ where
     ]));
     res.put(data_row(&vec![
         "free_servers".to_string(),
-        stats
+        server_stats
             .keys()
-            .map(|address_id| stats[&address_id]["sv_idle"])
-            .sum::<i64>()
+            .filter(|server_id| server_stats.get(server_id).unwrap().state == ServerState::Idle)
+            .count()
             .to_string(),
     ]));
     res.put(data_row(&vec![
         "used_servers".to_string(),
-        stats
+        server_stats
             .keys()
-            .map(|address_id| stats[&address_id]["sv_active"])
-            .sum::<i64>()
+            .filter(|server_id| server_stats.get(server_id).unwrap().state == ServerState::Active)
+            .count()
             .to_string(),
     ]));
     res.put(data_row(&vec!["dns_names".to_string(), "0".to_string()]));
@@ -156,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.
@@ -167,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.
@@ -182,11 +227,12 @@ async fn show_pools<T>(stream: &mut T) -> Result<(), Error>
 where
     T: tokio::io::AsyncWrite + std::marker::Unpin,
 {
-    let stats = get_stats();
+    let all_pool_stats = get_pool_stats();
 
     let columns = vec![
         ("database", DataType::Text),
         ("user", DataType::Text),
+        ("pool_mode", DataType::Text),
         ("cl_idle", DataType::Numeric),
         ("cl_active", DataType::Numeric),
         ("cl_waiting", DataType::Numeric),
@@ -198,32 +244,33 @@ where
         ("sv_login", DataType::Numeric),
         ("maxwait", DataType::Numeric),
         ("maxwait_us", DataType::Numeric),
-        ("pool_mode", DataType::Text),
     ];
 
     let mut res = BytesMut::new();
     res.put(row_description(&columns));
-    for (_, pool) in get_all_pools() {
+    for (user_pool, pool) in get_all_pools() {
+        let def = HashMap::default();
+        let pool_stats = all_pool_stats
+            .get(&(user_pool.db.clone(), user_pool.user.clone()))
+            .unwrap_or(&def);
+
         let pool_config = &pool.settings;
-        for shard in 0..pool.shards() {
-            for server in 0..pool.servers(shard) {
-                let address = pool.address(shard, server);
-                let stats = match stats.get(&address.id) {
-                    Some(stats) => stats.clone(),
-                    None => HashMap::new(),
-                };
-
-                let mut row = vec![address.name(), pool_config.user.username.clone()];
-
-                for column in &columns[2..columns.len() - 1] {
-                    let value = stats.get(column.0).unwrap_or(&0).to_string();
-                    row.push(value);
+        let mut row = vec![
+            user_pool.db.clone(),
+            user_pool.user.clone(),
+            pool_config.pool_mode.to_string(),
+        ];
+        for column in &columns[3..columns.len()] {
+            let value = match column.0 {
+                "maxwait" => (pool_stats.get("maxwait_us").unwrap_or(&0) / 1_000_000).to_string(),
+                "maxwait_us" => {
+                    (pool_stats.get("maxwait_us").unwrap_or(&0) % 1_000_000).to_string()
                 }
-
-                row.push(pool_config.pool_mode.to_string());
-                res.put(data_row(&row));
-            }
+                _other_values => pool_stats.get(column.0).unwrap_or(&0).to_string(),
+            };
+            row.push(value);
         }
+        res.put(data_row(&row));
     }
 
     res.put(command_complete("SHOW"));
@@ -233,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.
@@ -269,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
@@ -283,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(),
@@ -300,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.
@@ -312,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
@@ -332,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.
@@ -378,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.
@@ -387,6 +596,7 @@ where
     T: tokio::io::AsyncWrite + std::marker::Unpin,
 {
     let columns = vec![
+        ("instance", DataType::Text),
         ("database", DataType::Text),
         ("user", DataType::Text),
         ("total_xact_count", DataType::Numeric),
@@ -396,32 +606,32 @@ where
         ("total_xact_time", DataType::Numeric),
         ("total_query_time", DataType::Numeric),
         ("total_wait_time", DataType::Numeric),
+        ("total_errors", DataType::Numeric),
         ("avg_xact_count", DataType::Numeric),
         ("avg_query_count", DataType::Numeric),
         ("avg_recv", DataType::Numeric),
         ("avg_sent", DataType::Numeric),
+        ("avg_errors", DataType::Numeric),
         ("avg_xact_time", DataType::Numeric),
         ("avg_query_time", DataType::Numeric),
         ("avg_wait_time", DataType::Numeric),
     ];
 
-    let stats = get_stats();
+    let all_stats = get_address_stats();
     let mut res = BytesMut::new();
     res.put(row_description(&columns));
 
-    for ((_db_name, username), pool) in get_all_pools() {
+    for (user_pool, pool) in get_all_pools() {
         for shard in 0..pool.shards() {
             for server in 0..pool.servers(shard) {
                 let address = pool.address(shard, server);
-                let stats = match stats.get(&address.id) {
+                let stats = match all_stats.get(&address.id) {
                     Some(stats) => stats.clone(),
                     None => HashMap::new(),
                 };
 
-                let mut row = vec![address.name()];
-                row.push(username.clone());
-
-                for column in &columns[2..] {
+                let mut row = vec![address.name(), user_pool.db.clone(), user_pool.user.clone()];
+                for column in &columns[3..] {
                     row.push(stats.get(column.0).unwrap_or(&0).to_string());
                 }
 
@@ -437,5 +647,260 @@ 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
+async fn show_clients<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let columns = vec![
+        ("client_id", DataType::Text),
+        ("database", DataType::Text),
+        ("user", DataType::Text),
+        ("application_name", DataType::Text),
+        ("state", DataType::Text),
+        ("transaction_count", DataType::Numeric),
+        ("query_count", DataType::Numeric),
+        ("error_count", DataType::Numeric),
+        ("age_seconds", DataType::Numeric),
+    ];
+
+    let new_map = get_client_stats();
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    for (_, client) in new_map {
+        let row = vec![
+            format!("{:#010X}", client.client_id),
+            client.pool_name,
+            client.username,
+            client.application_name.clone(),
+            client.state.to_string(),
+            client.transaction_count.to_string(),
+            client.query_count.to_string(),
+            client.error_count.to_string(),
+            Instant::now()
+                .duration_since(client.connect_time)
+                .as_secs()
+                .to_string(),
+        ];
+
+        res.put(data_row(&row));
+    }
+
+    res.put(command_complete("SHOW"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
+/// Show currently connected servers
+async fn show_servers<T>(stream: &mut T) -> Result<(), Error>
+where
+    T: tokio::io::AsyncWrite + std::marker::Unpin,
+{
+    let columns = vec![
+        ("server_id", DataType::Text),
+        ("database_name", DataType::Text),
+        ("user", DataType::Text),
+        ("address_id", DataType::Text),
+        ("application_name", DataType::Text),
+        ("state", DataType::Text),
+        ("transaction_count", DataType::Numeric),
+        ("query_count", DataType::Numeric),
+        ("bytes_sent", DataType::Numeric),
+        ("bytes_received", DataType::Numeric),
+        ("age_seconds", DataType::Numeric),
+    ];
+
+    let new_map = get_server_stats();
+    let mut res = BytesMut::new();
+    res.put(row_description(&columns));
+
+    for (_, server) in new_map {
+        let row = vec![
+            format!("{:#010X}", server.server_id),
+            server.pool_name,
+            server.username,
+            server.address_name,
+            server.application_name,
+            server.state.to_string(),
+            server.transaction_count.to_string(),
+            server.query_count.to_string(),
+            server.bytes_sent.to_string(),
+            server.bytes_received.to_string(),
+            Instant::now()
+                .duration_since(server.connect_time)
+                .as_secs()
+                .to_string(),
+        ];
+
+        res.put(data_row(&row));
+    }
+
+    res.put(command_complete("SHOW"));
+
+    // ReadyForQuery
+    res.put_u8(b'Z');
+    res.put_i32(5);
+    res.put_u8(b'I');
+
+    write_all_half(stream, &res).await
+}
+
+/// 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,18 +2,20 @@
 use arc_swap::ArcSwap;
 use log::{error, info};
 use once_cell::sync::Lazy;
+use regex::Regex;
 use serde_derive::{Deserialize, Serialize};
-use std::collections::{HashMap, HashSet};
-use std::hash::Hash;
+use std::collections::hash_map::DefaultHasher;
+use std::collections::{BTreeMap, HashMap, HashSet};
+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};
+use crate::sharding::ShardingFunction;
 use crate::tls::{load_certs, load_keys};
-use crate::{ClientServerMap, ConnectionPool};
 
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");
 
@@ -23,8 +25,12 @@ static CONFIG: Lazy<ArcSwap<Config>> = Lazy::new(|| ArcSwap::from_pointee(Config
 /// Server role: primary or replica.
 #[derive(Clone, PartialEq, Serialize, Deserialize, Hash, std::cmp::Eq, Debug, Copy)]
 pub enum Role {
+    #[serde(alias = "primary", alias = "Primary")]
     Primary,
+    #[serde(alias = "replica", alias = "Replica")]
     Replica,
+    #[serde(alias = "mirror", alias = "Mirror")]
+    Mirror,
 }
 
 impl ToString for Role {
@@ -32,6 +38,7 @@ impl ToString for Role {
         match *self {
             Role::Primary => "primary".to_string(),
             Role::Replica => "replica".to_string(),
+            Role::Mirror => "mirror".to_string(),
         }
     }
 }
@@ -72,9 +79,6 @@ pub struct Address {
     /// The name of the Postgres database.
     pub database: String,
 
-    /// Default search_path.
-    pub search_path: Option<String>,
-
     /// Server role: replica, primary.
     pub role: Role,
 
@@ -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 {
@@ -101,10 +108,10 @@ impl Default for Address {
             address_index: 0,
             replica_number: 0,
             database: String::from("database"),
-            search_path: None,
             role: Role::Replica,
             username: String::from("username"),
             pool_name: String::from("pool_name"),
+            mirrors: Vec::new(),
         }
     }
 }
@@ -114,21 +121,25 @@ 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
+            ),
         }
     }
 }
 
 /// PostgreSQL user.
-#[derive(Clone, PartialEq, Hash, std::cmp::Eq, Serialize, Deserialize, Debug)]
+#[derive(Clone, PartialEq, Hash, Eq, Serialize, Deserialize, Debug)]
 pub struct User {
     pub username: String,
     pub password: String,
     pub pool_size: u32,
+    #[serde(default)] // 0
     pub statement_timeout: u64,
 }
 
@@ -146,34 +157,130 @@ impl Default for User {
 /// General configuration.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct General {
+    #[serde(default = "General::default_host")]
     pub host: String,
+
+    #[serde(default = "General::default_port")]
     pub port: i16,
+
     pub enable_prometheus_exporter: Option<bool>,
     pub prometheus_exporter_port: i16,
+
+    #[serde(default = "General::default_connect_timeout")]
     pub connect_timeout: u64,
-    pub healthcheck_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,
+
+    #[serde(default = "General::default_healthcheck_timeout")]
+    pub healthcheck_timeout: u64,
+
+    #[serde(default = "General::default_healthcheck_delay")]
     pub healthcheck_delay: u64,
+
+    #[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,
+
     pub tls_certificate: Option<String>,
     pub tls_private_key: Option<String>,
     pub admin_username: String,
     pub admin_password: String,
 }
 
+impl General {
+    pub fn default_host() -> String {
+        "0.0.0.0".into()
+    }
+
+    pub fn default_port() -> i16 {
+        5432
+    }
+
+    pub fn default_connect_timeout() -> u64 {
+        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
+    }
+
+    pub fn default_healthcheck_timeout() -> u64 {
+        1000
+    }
+
+    pub fn default_healthcheck_delay() -> u64 {
+        30000
+    }
+
+    pub fn default_ban_time() -> i64 {
+        60
+    }
+
+    pub fn default_worker_threads() -> usize {
+        4
+    }
+}
+
 impl Default for General {
     fn default() -> General {
         General {
-            host: String::from("localhost"),
-            port: 5432,
+            host: Self::default_host(),
+            port: Self::default_port(),
             enable_prometheus_exporter: Some(false),
             prometheus_exporter_port: 9930,
-            connect_timeout: 5000,
-            healthcheck_timeout: 1000,
-            shutdown_timeout: 60000,
-            healthcheck_delay: 30000,
-            ban_time: 60,
+            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,
@@ -182,52 +289,259 @@ impl Default for General {
         }
     }
 }
-#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
-pub struct Pool {
-    pub pool_mode: String,
-    pub default_role: String,
-    pub query_parser_enabled: bool,
-    pub primary_reads_enabled: bool,
-    pub sharding_function: String,
-    pub shards: HashMap<String, Shard>,
-    pub users: HashMap<String, User>,
+
+/// Pool mode:
+/// - transaction: server serves one transaction,
+/// - session: server is attached to the client.
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Copy, Hash)]
+pub enum PoolMode {
+    #[serde(alias = "transaction", alias = "Transaction")]
+    Transaction,
+
+    #[serde(alias = "session", alias = "Session")]
+    Session,
 }
-impl Default for Pool {
-    fn default() -> Pool {
-        Pool {
-            pool_mode: String::from("transaction"),
-            shards: HashMap::from([(String::from("1"), Shard::default())]),
-            users: HashMap::default(),
-            default_role: String::from("any"),
-            query_parser_enabled: false,
-            primary_reads_enabled: true,
-            sharding_function: "pg_bigint_hash".to_string(),
+impl ToString for PoolMode {
+    fn to_string(&self) -> String {
+        match *self {
+            PoolMode::Transaction => "transaction".to_string(),
+            PoolMode::Session => "session".to_string(),
         }
     }
 }
 
+#[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
+    pub query_parser_enabled: bool,
+
+    #[serde(default)] // False
+    pub primary_reads_enabled: bool,
+
+    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 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" => (),
+            "replica" => (),
+            other => {
+                error!(
+                    "Query router default_role must be 'primary', 'replica', or 'any', got: '{}'",
+                    other
+                );
+                return Err(Error::BadConfig);
+            }
+        };
+
+        for (shard_idx, shard) in &self.shards {
+            match shard_idx.parse::<usize>() {
+                Ok(_) => (),
+                Err(_) => {
+                    error!(
+                        "Shard '{}' is not a valid number, shards must be numbered starting at 0",
+                        shard_idx
+                    );
+                    return Err(Error::BadConfig);
+                }
+            };
+            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(())
+    }
+}
+
+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),
+        }
+    }
+}
+
+#[derive(Clone, PartialEq, Serialize, Deserialize, Debug, Hash, Eq)]
+pub struct ServerConfig {
+    pub host: String,
+    pub port: u16,
+    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)]
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Hash, Eq)]
 pub struct Shard {
     pub database: String,
-    pub search_path: Option<String>,
-    pub servers: Vec<(String, u16, String)>,
+    pub mirrors: Option<Vec<MirrorServerConfig>>,
+    pub servers: Vec<ServerConfig>,
+}
+
+impl Shard {
+    pub fn validate(&self) -> Result<(), Error> {
+        // We use addresses as unique identifiers,
+        // let's make sure they are unique in the config as well.
+        let mut dup_check = HashSet::new();
+        let mut primary_count = 0;
+
+        if self.servers.is_empty() {
+            error!("Shard {} has no servers configured", self.database);
+            return Err(Error::BadConfig);
+        }
+
+        for server in &self.servers {
+            dup_check.insert(server);
+
+            // Check that we define only zero or one primary.
+            if server.role == Role::Primary {
+                primary_count += 1
+            }
+        }
+
+        if primary_count > 1 {
+            error!(
+                "Shard {} has more than on primary configured",
+                self.database
+            );
+            return Err(Error::BadConfig);
+        }
+
+        if dup_check.len() != self.servers.len() {
+            error!("Shard {} contains duplicate server configs", self.database);
+            return Err(Error::BadConfig);
+        }
+
+        Ok(())
+    }
 }
 
 impl Default for Shard {
     fn default() -> Shard {
         Shard {
-            servers: vec![(String::from("localhost"), 5432, String::from("primary"))],
-            search_path: None,
+            servers: vec![ServerConfig {
+                host: String::from("localhost"),
+                port: 5432,
+                role: Role::Primary,
+            }],
+            mirrors: None,
             database: String::from("postgres"),
         }
     }
 }
 
-fn default_path() -> String {
-    String::from("pgcat.toml")
-}
-
 /// Configuration wrapper.
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
 pub struct Config {
@@ -235,24 +549,30 @@ pub struct Config {
     // so we should always put simple fields before nested fields
     // in all serializable structs to avoid ValueAfterTable errors
     // These errors occur when the toml serializer is about to produce
-    // ambigous toml structure like the one below
+    // ambiguous toml structure like the one below
     // [main]
     // field1_under_main = 1
     // field2_under_main = 2
     // [main.subconf]
     // field1_under_subconf = 1
     // field3_under_main = 3 # This field will be interpreted as being under subconf and not under main
-    #[serde(default = "default_path")]
+    #[serde(default = "Config::default_path")]
     pub path: String,
 
     pub general: General,
     pub pools: HashMap<String, Pool>,
 }
 
+impl Config {
+    pub fn default_path() -> String {
+        String::from("pgcat.toml")
+    }
+}
+
 impl Default for Config {
     fn default() -> Config {
         Config {
-            path: String::from("pgcat.toml"),
+            path: Self::default_path(),
             general: General::default(),
             pools: HashMap::default(),
         }
@@ -268,7 +588,11 @@ impl From<&Config> for std::collections::HashMap<String, String> {
                 [
                     (
                         format!("pools.{}.pool_mode", pool_name),
-                        pool.pool_mode.clone(),
+                        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),
@@ -284,7 +608,7 @@ impl From<&Config> for std::collections::HashMap<String, String> {
                     ),
                     (
                         format!("pools.{}.sharding_function", pool_name),
-                        pool.sharding_function.clone(),
+                        pool.sharding_function.to_string(),
                     ),
                     (
                         format!("pools.{:?}.shard_count", pool_name),
@@ -314,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(),
@@ -338,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() {
@@ -376,10 +714,31 @@ impl Config {
                     .sum::<u32>()
                     .to_string()
             );
-            info!("[pool: {}] Pool mode: {}", pool_name, pool_config.pool_mode);
+            info!(
+                "[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,
+            };
+            info!(
+                "[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, pool_config.sharding_function
+                pool_name,
+                pool_config.sharding_function.to_string()
             );
             info!(
                 "[pool: {}] Primary reads: {}",
@@ -412,6 +771,45 @@ impl Config {
             }
         }
     }
+
+    pub fn validate(&mut self) -> Result<(), Error> {
+        // Validate TLS!
+        match self.general.tls_certificate.clone() {
+            Some(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);
+                                }
+                            },
+
+                            None => {
+                                error!("tls_certificate is set, but the tls_private_key is not");
+                                return Err(Error::BadConfig);
+                            }
+                        };
+                    }
+
+                    Err(err) => {
+                        error!("tls_certificate is incorrectly configured: {:?}", err);
+                        return Err(Error::BadConfig);
+                    }
+                }
+            }
+            None => (),
+        };
+
+        for pool in self.pools.values_mut() {
+            pool.validate()?;
+        }
+
+        Ok(())
+    }
 }
 
 /// Get a read-only instance of the configuration
@@ -448,132 +846,7 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
-    // Validate TLS!
-    match config.general.tls_certificate.clone() {
-        Some(tls_certificate) => {
-            match load_certs(&Path::new(&tls_certificate)) {
-                Ok(_) => {
-                    // Cert is okay, but what about the private key?
-                    match config.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);
-                            }
-                        },
-
-                        None => {
-                            error!("tls_certificate is set, but the tls_private_key is not");
-                            return Err(Error::BadConfig);
-                        }
-                    };
-                }
-
-                Err(err) => {
-                    error!("tls_certificate is incorrectly configured: {:?}", err);
-                    return Err(Error::BadConfig);
-                }
-            }
-        }
-        None => (),
-    };
-
-    for (pool_name, pool) in &config.pools {
-        match pool.sharding_function.as_ref() {
-            "pg_bigint_hash" => (),
-            "sha1" => (),
-            _ => {
-                error!(
-                    "Supported sharding functions are: 'pg_bigint_hash', 'sha1', got: '{}' in pool {} settings",
-                    pool.sharding_function,
-                    pool_name
-                );
-                return Err(Error::BadConfig);
-            }
-        };
-
-        match pool.default_role.as_ref() {
-            "any" => (),
-            "primary" => (),
-            "replica" => (),
-            other => {
-                error!(
-                    "Query router default_role must be 'primary', 'replica', or 'any', got: '{}'",
-                    other
-                );
-                return Err(Error::BadConfig);
-            }
-        };
-
-        match pool.pool_mode.as_ref() {
-            "transaction" => (),
-            "session" => (),
-            other => {
-                error!(
-                    "pool_mode can be 'session' or 'transaction', got: '{}'",
-                    other
-                );
-                return Err(Error::BadConfig);
-            }
-        };
-
-        for shard in &pool.shards {
-            // We use addresses as unique identifiers,
-            // let's make sure they are unique in the config as well.
-            let mut dup_check = HashSet::new();
-            let mut primary_count = 0;
-
-            match shard.0.parse::<usize>() {
-                Ok(_) => (),
-                Err(_) => {
-                    error!(
-                        "Shard '{}' is not a valid number, shards must be numbered starting at 0",
-                        shard.0
-                    );
-                    return Err(Error::BadConfig);
-                }
-            };
-
-            if shard.1.servers.len() == 0 {
-                error!("Shard {} has no servers configured", shard.0);
-                return Err(Error::BadConfig);
-            }
-
-            for server in &shard.1.servers {
-                dup_check.insert(server);
-
-                // Check that we define only zero or one primary.
-                match server.2.as_ref() {
-                    "primary" => primary_count += 1,
-                    _ => (),
-                };
-
-                // Check role spelling.
-                match server.2.as_ref() {
-                    "primary" => (),
-                    "replica" => (),
-                    _ => {
-                        error!(
-                            "Shard {} server role must be either 'primary' or 'replica', got: '{}'",
-                            shard.0, server.2
-                        );
-                        return Err(Error::BadConfig);
-                    }
-                };
-            }
-
-            if primary_count > 1 {
-                error!("Shard {} has more than on primary configured", &shard.0);
-                return Err(Error::BadConfig);
-            }
-
-            if dup_check.len() != shard.1.servers.len() {
-                error!("Shard {} contains duplicate server configs", &shard.0);
-                return Err(Error::BadConfig);
-            }
-        }
-    }
+    config.validate()?;
 
     config.path = path.to_string();
 
@@ -595,7 +868,7 @@ pub async fn reload_config(client_server_map: ClientServerMap) -> Result<bool, E
     let new_config = get_config();
 
     if old_config.pools != new_config.pools {
-        info!("Pool configuration changed, re-creating server pools");
+        info!("Pool configuration changed");
         ConnectionPool::from_config(client_server_map).await?;
         Ok(true)
     } else if old_config != new_config {
@@ -616,19 +889,21 @@ 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);
 
         assert_eq!(
-            get_config().pools["sharded_db"].shards["0"].servers[0].0,
+            get_config().pools["sharded_db"].shards["0"].servers[0].host,
             "127.0.0.1"
         );
         assert_eq!(
-            get_config().pools["sharded_db"].shards["1"].servers[0].2,
-            "primary"
+            get_config().pools["sharded_db"].shards["1"].servers[0].role,
+            Role::Primary
         );
         assert_eq!(
             get_config().pools["sharded_db"].shards["1"].database,
@@ -646,11 +921,11 @@ mod test {
         assert_eq!(get_config().pools["sharded_db"].default_role, "any");
 
         assert_eq!(
-            get_config().pools["simple_db"].shards["0"].servers[0].0,
+            get_config().pools["simple_db"].shards["0"].servers[0].host,
             "127.0.0.1"
         );
         assert_eq!(
-            get_config().pools["simple_db"].shards["0"].servers[0].1,
+            get_config().pools["simple_db"].shards["0"].servers[0].port,
             5432
         );
         assert_eq!(

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

@@ -0,0 +1,34 @@
+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;
+pub mod sharding;
+pub mod stats;
+pub mod tls;
+
+/// Format chrono::Duration to be more human-friendly.
+///
+/// # Arguments
+///
+/// * `duration` - A duration of time
+pub fn format_duration(duration: &chrono::Duration) -> String {
+    let milliseconds = format!("{:0>3}", duration.num_milliseconds() % 1000);
+
+    let seconds = format!("{:0>2}", duration.num_seconds() % 60);
+
+    let minutes = format!("{:0>2}", duration.num_minutes() % 60);
+
+    let hours = format!("{:0>2}", duration.num_hours() % 24);
+
+    let days = duration.num_days().to_string();
+
+    format!(
+        "{}d {}:{}:{}.{}",
+        days, hours, minutes, seconds, milliseconds
+    )
+}

src/main.rs

@@ -37,13 +37,22 @@ extern crate tokio;
 extern crate tokio_rustls;
 extern crate toml;
 
-use log::{debug, error, info};
+#[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;
@@ -57,6 +66,8 @@ mod config;
 mod constants;
 mod errors;
 mod messages;
+mod mirrors;
+mod multi_logger;
 mod pool;
 mod prometheus;
 mod query_router;
@@ -67,14 +78,13 @@ 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::init();
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    multi_logger::MultiLogger::init().unwrap();
+
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
 
     if !query_router::QueryRouter::setup() {
@@ -90,244 +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 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 mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000));
-    let (shutdown_tx, _) = broadcast::channel::<()>(1);
-    let (drain_tx, mut drain_rx) = mpsc::channel::<i8>(2048);
-    let (exit_tx, mut exit_rx) = mpsc::channel::<()>(1);
+        tokio::task::spawn(async move {
+            let mut stats_collector = Collector::new(stats_rx, stats_tx.clone());
+            stats_collector.collect().await;
+        });
 
-    info!("Waiting for clients");
+        info!("Config autoreloader: {}", config.general.autoreload);
 
-    let mut admin_only = false;
-    let mut total_clients = 0;
+        let mut autoreload_interval = tokio::time::interval(tokio::time::Duration::from_millis(15_000));
+        let autoreload_client_server_map = client_server_map.clone();
 
-    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(_) => (),
-                };
-
-                get_config().show();
-            },
-
-            _ = autoreload_interval.tick() => {
+        tokio::task::spawn(async move {
+            loop {
+                autoreload_interval.tick().await;
                 if config.general.autoreload {
                     info!("Automatically reloading config");
 
-                    match reload_config(client_server_map.clone()).await {
-                        Ok(changed) => {
-                            if changed {
-                                get_config().show()
-                            }
+                    if let Ok(changed) = reload_config(autoreload_client_server_map.clone()).await {
+                        if changed {
+                            get_config().show()
                         }
-                        Err(_) => (),
                     };
                 }
-            },
+            }
+        });
 
-            // Initiate graceful shutdown sequence on sig int
-            _ = interrupt_signal.recv() => {
-                info!("Got SIGINT, waiting for client connection drain now");
-                admin_only = true;
+        #[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();
 
-                // Broadcast that client tasks need to finish
-                let _ = shutdown_tx.send(());
-                let exit_tx = exit_tx.clone();
-                let _ = drain_tx.send(0).await;
+        #[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;
 
-                tokio::task::spawn(async move {
-                    let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(config.general.shutdown_timeout));
+        info!("Waiting for clients");
 
-                    // First tick fires immediately.
-                    interval.tick().await;
+        loop {
+            tokio::select! {
+                // Reload config:
+                // kill -SIGHUP $(pgrep pgcat)
+                _ = sighup_signal.recv() => {
+                    info!("Reloading config");
 
-                    // Second one in the interval time.
-                    interval.tick().await;
+                    _ = reload_config(client_server_map.clone()).await;
 
-                    // We're done waiting.
-                    error!("Timed out waiting for clients");
+                    get_config().show();
+                },
 
-                    let _ = exit_tx.send(()).await;
-                });
-            },
+                // Initiate graceful shutdown sequence on sig int
+                _ = interrupt_signal.recv() => {
+                    info!("Got SIGINT");
 
-            _ = term_signal.recv() => 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)
+                                    );
                                 }
                             }
 
-                            debug!("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...");
-}
-
-/// Format chrono::Duration to be more human-friendly.
-///
-/// # Arguments
-///
-/// * `duration` - A duration of time
-fn format_duration(duration: &chrono::Duration) -> String {
-    let seconds = {
-        let seconds = duration.num_seconds() % 60;
-        if seconds < 10 {
-            format!("0{}", seconds)
-        } else {
-            format!("{}", seconds)
-        }
-    };
-
-    let minutes = {
-        let minutes = duration.num_minutes() % 60;
-        if minutes < 10 {
-            format!("0{}", minutes)
-        } else {
-            format!("{}", minutes)
-        }
-    };
-
-    let hours = {
-        let hours = duration.num_hours() % 24;
-        if hours < 10 {
-            format!("0{}", hours)
-        } else {
-            format!("{}", hours)
-        }
-    };
-
-    let days = duration.num_days().to_string();
-
-    format!("{}d {}:{}:{}", days, hours, minutes, seconds)
+    });
+    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,41 +111,25 @@ 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.
 /// This tells the server which user we are and what database we want.
-pub async fn startup(
-    stream: &mut TcpStream,
-    user: &str,
-    database: &str,
-    search_path: Option<&String>,
-) -> Result<(), Error> {
+pub async fn startup(stream: &mut TcpStream, user: &str, database: &str) -> Result<(), Error> {
     let mut bytes = BytesMut::with_capacity(25);
 
     bytes.put_i32(196608); // Protocol number
 
     // 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);
-
-    // search_path
-    match search_path {
-        Some(search_path) => {
-            bytes.put(&b"options\0"[..]);
-            bytes.put_slice(&format!("-c search_path={}", search_path).as_bytes());
-            bytes.put_u8(0);
-        }
-        None => (),
-    };
-
     bytes.put_u8(0); // Null terminator
 
     let len = bytes.len() as i32 + 4i32;
@@ -152,7 +141,12 @@ pub async fn startup(
 
     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
+            )))
+        }
     }
 }
 
@@ -171,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();
         }
@@ -250,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`
@@ -270,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
 }
 
@@ -308,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);
@@ -320,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>
@@ -343,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);
@@ -382,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
 }
 
@@ -395,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);
@@ -439,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');
@@ -466,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
+            )))
+        }
     }
 }
 
@@ -488,31 +492,56 @@ 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)
 }
 
-pub fn server_paramater_message(key: &str, value: &str) -> BytesMut {
+pub fn server_parameter_message(key: &str, value: &str) -> BytesMut {
     let mut server_info = BytesMut::new();
 
     let null_byte_size = 1;
@@ -526,5 +555,41 @@ pub fn server_paramater_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/prometheus.rs

@@ -8,7 +8,7 @@ use std::net::SocketAddr;
 
 use crate::config::Address;
 use crate::pool::get_all_pools;
-use crate::stats::get_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 = get_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

@@ -1,7 +1,8 @@
 /// Implementation of the PostgreSQL server (database) protocol.
 /// Here we are pretending to the a Postgres client.
 use bytes::{Buf, BufMut, BytesMut};
-use log::{debug, error, info, trace};
+use log::{debug, error, info, trace, warn};
+use std::io::Read;
 use std::time::SystemTime;
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
@@ -13,12 +14,15 @@ 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;
-use crate::ClientServerMap;
 
 /// Server state.
 pub struct Server {
+    server_id: i32,
+
     /// Server host, e.g. localhost,
     /// port, e.g. 5432, and role, e.g. primary or replica.
     address: Address,
@@ -48,6 +52,9 @@ pub struct Server {
     /// Is the server broken? We'll remote it from the pool if so.
     bad: bool,
 
+    /// If server connection requires a DISCARD ALL before checkin
+    needs_cleanup: bool,
+
     /// Mapping of clients and servers used for query cancellation.
     client_server_map: ClientServerMap,
 
@@ -62,12 +69,15 @@ pub struct Server {
 
     // Last time that a successful server send or response happened
     last_activity: SystemTime,
+
+    mirror_manager: Option<MirroringManager>,
 }
 
 impl Server {
     /// Pretend to be the Postgres client and connect to the server given host, port and credentials.
     /// Perform the authentication and return the server in a ready for query state.
     pub async fn startup(
+        server_id: i32,
         address: &Address,
         user: &User,
         database: &str,
@@ -79,20 +89,18 @@ 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");
 
         // StartupMessage
-        startup(
-            &mut stream,
-            &user.username,
-            database,
-            address.search_path.as_ref(),
-        )
-        .await?;
+        startup(&mut stream, &user.username, database).await?;
 
         let mut server_info = BytesMut::new();
         let mut process_id: i32 = 0;
@@ -105,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);
@@ -121,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);
@@ -134,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[..])
@@ -150,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]);
@@ -174,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);
 
@@ -192,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[..]);
@@ -213,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[..])) {
@@ -239,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);
@@ -255,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
@@ -274,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.
@@ -291,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))),
                     };
                 }
 
@@ -306,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();
@@ -314,19 +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,
-                        process_id: process_id,
-                        secret_key: secret_key,
+                        server_info,
+                        server_id,
+                        process_id,
+                        secret_key,
                         in_transaction: false,
                         data_available: false,
                         bad: false,
-                        client_server_map: client_server_map,
+                        needs_cleanup: false,
+                        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?;
@@ -338,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
+                    )));
                 }
             };
         }
@@ -356,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");
 
@@ -368,13 +390,13 @@ 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> {
-        self.stats
-            .data_sent(messages.len(), self.process_id, self.address.id);
+    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 {
             Ok(_) => {
@@ -436,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
+                            )));
                         }
                     };
 
@@ -446,6 +471,39 @@ impl Server {
                     break;
                 }
 
+                // CommandComplete
+                'C' => {
+                    let mut command_tag = String::new();
+                    match message.reader().read_to_string(&mut command_tag) {
+                        Ok(_) => {
+                            // Non-exhaustive list of commands that are likely to change session variables/resources
+                            // 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" => {
+                                    // 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;
+                                }
+                                _ => (),
+                            }
+                        }
+
+                        Err(err) => {
+                            warn!("Encountered an error while parsing CommandTag {}", err);
+                        }
+                    }
+                }
+
                 // DataRow
                 'D' => {
                     // More data is available after this message, this is not the end of the reply.
@@ -466,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.
@@ -483,8 +545,7 @@ impl Server {
         let bytes = self.buffer.clone();
 
         // Keep track of how much data we got from the server for stats.
-        self.stats
-            .data_received(bytes.len(), self.process_id, self.address.id);
+        self.stats.data_received(bytes.len(), self.server_id);
 
         // Clear the buffer for next query.
         self.buffer.clear();
@@ -499,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
     }
 
@@ -546,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?;
@@ -559,14 +621,45 @@ impl Server {
         Ok(())
     }
 
+    /// Perform any necessary cleanup before putting the server
+    /// connection back in the pool
+    pub async fn checkin_cleanup(&mut self) -> Result<(), Error> {
+        // Client disconnected with an open transaction on the server connection.
+        // Pgbouncer behavior is to close the server connection but that can cause
+        // server connection thrashing if clients repeatedly do this.
+        // Instead, we ROLLBACK that transaction before putting the connection back in the pool
+        if self.in_transaction() {
+            warn!("Server returned while still in transaction, rolling back transaction");
+            self.query("ROLLBACK").await?;
+        }
+
+        // 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
+        // it before each checkin.
+        if self.needs_cleanup {
+            warn!("Server returned with session state altered, discarding state");
+            self.query("DISCARD ALL").await?;
+            self.needs_cleanup = false;
+        }
+
+        Ok(())
+    }
+
     /// A shorthand for `SET application_name = $1`.
-    #[allow(dead_code)]
     pub async fn set_name(&mut self, name: &str) -> Result<(), Error> {
         if self.application_name != name {
             self.application_name = name.to_string();
-            Ok(self
+            // We don't want `SET application_name` to mark the server connection
+            // as needing cleanup
+            let needs_cleanup_before = self.needs_cleanup;
+
+            let result = Ok(self
                 .query(&format!("SET application_name = '{}'", name))
-                .await?)
+                .await?);
+            self.needs_cleanup = needs_cleanup_before;
+            result
         } else {
             Ok(())
         }
@@ -578,15 +671,35 @@ impl Server {
         self.address.clone()
     }
 
-    /// Get the server's unique identifier.
-    pub fn process_id(&self) -> i32 {
-        self.process_id
+    /// Get the server connection identifier
+    /// Used to uniquely identify connection in statistics
+    pub fn server_id(&self) -> i32 {
+        self.server_id
     }
 
     // Get server's latest response timestamp
     pub fn last_activity(&self) -> SystemTime {
         self.last_activity
     }
+
+    // Marks a connection as needing DISCARD ALL at checkin
+    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 {
@@ -594,8 +707,8 @@ impl Drop for Server {
     /// the socket is in non-blocking mode, so it may not be ready
     /// for a write.
     fn drop(&mut self) {
-        self.stats
-            .server_disconnecting(self.process_id(), self.address.id);
+        self.mirror_disconnect();
+        self.stats.server_disconnecting(self.server_id);
 
         let mut bytes = BytesMut::with_capacity(4);
         bytes.put_u8(b'X');
@@ -613,7 +726,8 @@ impl Drop for Server {
         let duration = now - self.connected_at;
 
         info!(
-            "Server connection closed, session duration: {}",
+            "Server connection closed {:?}, session duration: {}",
+            self.address,
             crate::format_duration(&duration)
         );
     }

src/sharding.rs

@@ -1,3 +1,4 @@
+use serde_derive::{Deserialize, Serialize};
 /// Implements various sharding functions.
 use sha1::{Digest, Sha1};
 
@@ -5,12 +6,23 @@ use sha1::{Digest, Sha1};
 const PARTITION_HASH_SEED: u64 = 0x7A5B22367996DCFD;
 
 /// The sharding functions we support.
-#[derive(Debug, PartialEq, Copy, Clone)]
+#[derive(Debug, PartialEq, Copy, Clone, Serialize, Deserialize, Hash, std::cmp::Eq)]
 pub enum ShardingFunction {
+    #[serde(alias = "pg_bigint_hash", alias = "PgBigintHash")]
     PgBigintHash,
+    #[serde(alias = "sha1", alias = "Sha1")]
     Sha1,
 }
 
+impl ToString for ShardingFunction {
+    fn to_string(&self) -> String {
+        match *self {
+            ShardingFunction::PgBigintHash => "pg_bigint_hash".to_string(),
+            ShardingFunction::Sha1 => "sha1".to_string(),
+        }
+    }
+}
+
 /// The sharder.
 pub struct Sharder {
     /// Number of shards in the cluster.
@@ -121,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
@@ -129,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/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

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

@@ -0,0 +1,44 @@
+version: "3"
+services:
+  pg1:
+    image: postgres:14
+    network_mode: "service:main"
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_DB: postgres
+      POSTGRES_PASSWORD: postgres
+      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"
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_DB: postgres
+      POSTGRES_PASSWORD: postgres
+      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"
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_DB: postgres
+      POSTGRES_PASSWORD: postgres
+      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"
+    environment:
+      POSTGRES_USER: postgres
+      POSTGRES_DB: postgres
+      POSTGRES_PASSWORD: postgres
+      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"]
+    volumes:
+      - ../../:/app/
+      - /app/target/

tests/docker/run.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+rm -rf /app/target/ || true
+rm /app/*.profraw || true
+rm /app/pgcat.profdata || true
+rm -rf /app/cov || true
+
+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
+
+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 \
+)
+
+echo "Generating coverage report"
+
+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

@@ -18,9 +18,14 @@ def pgcat_start():
 
 
 def pg_cat_send_signal(signal: signal.Signals):
-    for proc in psutil.process_iter(["pid", "name"]):
-        if "pgcat" == proc.name():
-            os.kill(proc.pid, signal)
+    try:
+        for proc in psutil.process_iter(["pid", "name"]):
+            if "pgcat" == proc.name():
+                os.kill(proc.pid, signal)
+    except Exception as e:
+        # The process can be gone when we send this signal
+        print(e)
+
     if signal == signal.SIGTERM:
         # Returns 0 if pgcat process exists
         time.sleep(2)
@@ -105,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
 
@@ -131,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

@@ -1,6 +1,8 @@
 source "https://rubygems.org"
 
 gem "pg"
-gem "activerecord"
+gem "toml"
+gem "rspec"
 gem "rubocop"
-gem "toml", "~> 0.3.0"
+gem "toxiproxy"
+gem "activerecord"

tests/ruby/Gemfile.lock

@@ -1,21 +1,22 @@
 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)
       tzinfo (~> 2.0)
     ast (2.4.2)
     concurrent-ruby (1.1.10)
-    i18n (1.11.0)
+    diff-lcs (1.5.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)
@@ -24,6 +25,19 @@ GEM
     rainbow (3.1.1)
     regexp_parser (2.3.1)
     rexml (3.2.5)
+    rspec (3.11.0)
+      rspec-core (~> 3.11.0)
+      rspec-expectations (~> 3.11.0)
+      rspec-mocks (~> 3.11.0)
+    rspec-core (3.11.0)
+      rspec-support (~> 3.11.0)
+    rspec-expectations (3.11.0)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.11.0)
+    rspec-mocks (3.11.1)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.11.0)
+    rspec-support (3.11.0)
     rubocop (1.29.0)
       parallel (~> 1.10)
       parser (>= 3.1.0.0)
@@ -38,19 +52,23 @@ GEM
     ruby-progressbar (1.11.0)
     toml (0.3.0)
       parslet (>= 1.8.0, < 3.0.0)
-    tzinfo (2.0.4)
+    toxiproxy (2.0.1)
+    tzinfo (2.0.5)
       concurrent-ruby (~> 1.0)
     unicode-display_width (2.1.0)
 
 PLATFORMS
+  aarch64-linux
   arm64-darwin-21
   x86_64-linux
 
 DEPENDENCIES
   activerecord
   pg
+  rspec
   rubocop
-  toml (~> 0.3.0)
+  toml
+  toxiproxy
 
 BUNDLED WITH
-   2.3.7
+   2.3.21

tests/ruby/admin_spec.rb

@@ -0,0 +1,369 @@
+# frozen_string_literal: true
+require 'uri'
+require_relative 'spec_helper'
+
+describe "Admin" do
+  let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 10) }
+  let(:pgcat_conn_str) { processes.pgcat.connection_string("sharded_db", "sharding_user") }
+
+  after do
+    processes.all_databases.map(&:reset)
+    processes.pgcat.shutdown
+  end
+
+  describe "SHOW STATS" do
+    context "clients connect and make one query" do
+      it "updates *_query_time and *_wait_time" do
+        connection = PG::connect("#{pgcat_conn_str}?application_name=one_query")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.async_exec("SELECT pg_sleep(0.25)")
+        connection.close
+
+        # wait for averages to be calculated, we shouldn't do this too often
+        sleep(15.5)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW STATS")[0]
+        admin_conn.close
+        expect(results["total_query_time"].to_i).to be_within(200).of(750)
+        expect(results["avg_query_time"].to_i).to_not eq(0)
+
+        expect(results["total_wait_time"].to_i).to_not eq(0)
+        expect(results["avg_wait_time"].to_i).to_not eq(0)
+      end
+    end
+  end
+
+  describe "SHOW POOLS" do
+    context "bad credentials" do
+      it "does not change any stats" do
+        bad_passsword_url = URI(pgcat_conn_str)
+        bad_passsword_url.password = "wrong"
+        expect { PG::connect("#{bad_passsword_url.to_s}?application_name=bad_password") }.to raise_error(PG::ConnectionBad)
+
+        sleep(1)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+
+        expect(results["sv_idle"]).to eq("1")
+      end
+    end
+
+    context "bad database name" do
+      it "does not change any stats" do
+        bad_db_url = URI(pgcat_conn_str)
+        bad_db_url.path = "/wrong_db"
+        expect { PG::connect("#{bad_db_url.to_s}?application_name=bad_db") }.to raise_error(PG::ConnectionBad)
+
+        sleep(1)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+
+        expect(results["sv_idle"]).to eq("1")
+      end
+    end
+
+    context "client connects but issues no queries" do
+      it "only affects cl_idle stats" do
+        connections = Array.new(20) { PG::connect(pgcat_conn_str) }
+        sleep(1)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["cl_idle"]).to eq("20")
+        expect(results["sv_idle"]).to eq("1")
+
+        connections.map(&:close)
+        sleep(1.1)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_active cl_idle cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["sv_idle"]).to eq("1")
+      end
+    end
+
+    context "clients connect and make one query" do
+      it "only affects cl_idle, sv_idle stats" do
+        connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
+        connections.each do |c|
+          Thread.new { c.async_exec("SELECT pg_sleep(2.5)") }
+        end
+
+        sleep(1.1)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_waiting cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["cl_active"]).to eq("5")
+        expect(results["sv_active"]).to eq("5")
+
+        sleep(3)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["cl_idle"]).to eq("5")
+        expect(results["sv_idle"]).to eq("5")
+
+        connections.map(&:close)
+        sleep(1)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["sv_idle"]).to eq("5")
+      end
+    end
+
+    context "client connects and opens a transaction and closes connection uncleanly" do
+      it "produces correct statistics" do
+        connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
+        connections.each do |c|
+          Thread.new do
+            c.async_exec("BEGIN")
+            c.async_exec("SELECT pg_sleep(0.01)")
+            c.close
+          end
+        end
+
+        sleep(1.1)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["sv_idle"]).to eq("5")
+      end
+    end
+
+    context "client fail to checkout connection from the pool" do
+      it "counts clients as idle" do
+        new_configs = processes.pgcat.current_config
+        new_configs["general"]["connect_timeout"] = 500
+        new_configs["general"]["ban_time"] = 1
+        new_configs["general"]["shutdown_timeout"] = 1
+        new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
+        processes.pgcat.update_config(new_configs)
+        processes.pgcat.reload_config
+
+        threads = []
+        connections = Array.new(5) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
+        connections.each do |c|
+          threads << Thread.new { c.async_exec("SELECT pg_sleep(1)") rescue PG::SystemError }
+        end
+
+        sleep(2)
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["cl_idle"]).to eq("5")
+        expect(results["sv_idle"]).to eq("1")
+
+        threads.map(&:join)
+        connections.map(&:close)
+      end
+    end
+
+    context "clients overwhelm server pools" do
+      let(:processes) { Helpers::Pgcat.single_instance_setup("sharded_db", 2) }
+
+      it "cl_waiting is updated to show it" do
+        threads = []
+        connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
+        connections.each do |c|
+          threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") }
+        end
+
+        sleep(1.1) # Allow time for stats to update
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_idle cl_cancel_req sv_idle sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+
+        expect(results["cl_waiting"]).to eq("2")
+        expect(results["cl_active"]).to eq("2")
+        expect(results["sv_active"]).to eq("2")
+
+        sleep(2.5) # Allow time for stats to update
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+        %w[cl_active cl_waiting cl_cancel_req sv_active sv_used sv_tested sv_login maxwait].each do |s|
+          raise StandardError, "Field #{s} was expected to be 0 but found to be #{results[s]}" if results[s] != "0"
+        end
+        expect(results["cl_idle"]).to eq("4")
+        expect(results["sv_idle"]).to eq("2")
+
+        threads.map(&:join)
+        connections.map(&:close)
+      end
+
+      it "show correct max_wait" do
+        threads = []
+        connections = Array.new(4) { PG::connect("#{pgcat_conn_str}?application_name=one_query") }
+        connections.each do |c|
+          threads << Thread.new { c.async_exec("SELECT pg_sleep(1.5)") }
+        end
+
+        sleep(2.5) # Allow time for stats to update
+        admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+        results = admin_conn.async_exec("SHOW POOLS")[0]
+
+        expect(results["maxwait"]).to eq("1")
+        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]
+        expect(results["maxwait"]).to eq("0")
+
+        threads.map(&:join)
+        connections.map(&:close)
+      end
+    end
+  end
+
+  describe "SHOW CLIENTS" do
+    it "reports correct number and application names" do
+      conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
+      connections = Array.new(20) { |i| PG::connect("#{conn_str}?application_name=app#{i % 5}") }
+
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      sleep(1) # Wait for stats to be updated
+
+      results = admin_conn.async_exec("SHOW CLIENTS")
+      expect(results.count).to eq(21) # count admin clients
+      expect(results.select { |c| c["application_name"] == "app3" ||  c["application_name"] == "app4" }.count).to eq(8)
+      expect(results.select { |c| c["database"] == "pgcat" }.count).to eq(1)
+
+      connections[0..5].map(&:close)
+      sleep(1) # Wait for stats to be updated
+      results = admin_conn.async_exec("SHOW CLIENTS")
+      expect(results.count).to eq(15)
+
+      connections[6..].map(&:close)
+      sleep(1) # Wait for stats to be updated
+      expect(admin_conn.async_exec("SHOW CLIENTS").count).to eq(1)
+      admin_conn.close
+    end
+
+    it "reports correct number of queries and transactions" do
+      conn_str = processes.pgcat.connection_string("sharded_db", "sharding_user")
+
+      connections = Array.new(2) { |i| PG::connect("#{conn_str}?application_name=app#{i}") }
+      connections.each do |c|
+        c.async_exec("SELECT 1")
+        c.async_exec("SELECT 2")
+        c.async_exec("SELECT 3")
+        c.async_exec("BEGIN")
+        c.async_exec("SELECT 4")
+        c.async_exec("SELECT 5")
+        c.async_exec("COMMIT")
+      end
+
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      sleep(1) # Wait for stats to be updated
+
+      results = admin_conn.async_exec("SHOW CLIENTS")
+      expect(results.count).to eq(3)
+      normal_client_results = results.reject { |r| r["database"] == "pgcat" }
+      expect(normal_client_results[0]["transaction_count"]).to eq("4")
+      expect(normal_client_results[1]["transaction_count"]).to eq("4")
+      expect(normal_client_results[0]["query_count"]).to eq("7")
+      expect(normal_client_results[1]["query_count"]).to eq("7")
+
+      admin_conn.close
+      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

@@ -0,0 +1,94 @@
+require 'pg'
+require 'toxiproxy'
+
+class PgInstance
+  attr_reader :port
+  attr_reader :username
+  attr_reader :password
+  attr_reader :database_name
+
+  def initialize(port, username, password, database_name)
+    @original_port = port
+    @toxiproxy_port = 10000 + port.to_i
+    @port = @toxiproxy_port
+
+    @username = username
+    @password = password
+    @database_name = database_name
+    @toxiproxy_name = "database_#{@original_port}"
+    Toxiproxy.populate([{
+      name: @toxiproxy_name,
+      listen: "0.0.0.0:#{@toxiproxy_port}",
+      upstream: "localhost:#{@original_port}",
+    }])
+
+    # Toxiproxy server will outlive our PgInstance objects
+    # so we want to destroy our proxies before exiting
+    # Ruby finalizer is ideal for doing this
+    ObjectSpace.define_finalizer(@toxiproxy_name, proc { Toxiproxy[@toxiproxy_name].destroy })
+  end
+
+  def with_connection
+    conn = PG.connect("postgres://#{@username}:#{@password}@localhost:#{port}/#{database_name}")
+    yield conn
+  ensure
+    conn&.close
+  end
+
+  def reset
+    reset_toxics
+    reset_stats
+    drop_connections
+    sleep 0.1
+  end
+
+  def toxiproxy
+    Toxiproxy[@toxiproxy_name]
+  end
+
+  def take_down
+    if block_given?
+      Toxiproxy[@toxiproxy_name].toxic(:limit_data, bytes: 5).apply { yield }
+    else
+      Toxiproxy[@toxiproxy_name].toxic(:limit_data, bytes: 5).toxics.each(&:save)
+    end
+  end
+
+  def add_latency(latency)
+    if block_given?
+      Toxiproxy[@toxiproxy_name].toxic(:latency, latency: latency).apply { yield }
+    else
+      Toxiproxy[@toxiproxy_name].toxic(:latency, latency: latency).toxics.each(&:save)
+    end
+  end
+
+  def delete_proxy
+    Toxiproxy[@toxiproxy_name].delete
+  end
+
+  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
+
+  def count_select_1_plus_2
+    with_connection { |c| c.async_exec("SELECT SUM(calls) FROM pg_stat_statements WHERE query = 'SELECT $1 + $2'")[0]["sum"].to_i }
+  end
+end

tests/ruby/helpers/pgcat_helper.rb

@@ -0,0 +1,148 @@
+require 'json'
+require 'ostruct'
+require_relative 'pgcat_process'
+require_relative 'pg_instance'
+
+module Helpers
+  module Pgcat
+    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,
+        "statement_timeout" => 0,
+        "username" => "sharding_user"
+      }
+
+      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")
+
+      pgcat_cfg = pgcat.current_config
+      pgcat_cfg["pools"] = {
+        "#{pool_name}" => {
+          "default_role" => "any",
+          "pool_mode" => pool_mode,
+          "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"]] },
+            "1" => { "database" => "shard1", "servers" => [["localhost", primary1.port.to_s, "primary"]] },
+            "2" => { "database" => "shard2", "servers" => [["localhost", primary2.port.to_s, "primary"]] },
+          },
+          "users" => { "0" => user }
+        }
+      }
+      pgcat.update_config(pgcat_cfg)
+
+      pgcat.start
+      pgcat.wait_until_ready
+
+      OpenStruct.new.tap do |struct|
+        struct.pgcat = pgcat
+        struct.shards = [primary0, primary1, primary2]
+        struct.all_databases = [primary0, primary1, primary2]
+      end
+    end
+
+    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,
+        "statement_timeout" => 0,
+        "username" => "sharding_user"
+      }
+
+      pgcat = PgcatProcess.new(log_level)
+      pgcat_cfg = pgcat.current_config
+
+      primary  = PgInstance.new(5432, user["username"], user["password"], "shard0")
+
+      # Main proxy configs
+      pgcat_cfg["pools"] = {
+        "#{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",
+          "shards" => {
+            "0" => {
+              "database" => "shard0",
+              "servers" => [
+                ["localhost", primary.port.to_s, "primary"]
+              ]
+            },
+          },
+          "users" => { "0" => user }
+        }
+      }
+      pgcat_cfg["general"]["port"] = pgcat.port
+      pgcat.update_config(pgcat_cfg)
+      pgcat.start
+      pgcat.wait_until_ready
+
+      OpenStruct.new.tap do |struct|
+        struct.pgcat = pgcat
+        struct.primary = primary
+        struct.all_databases = [primary]
+      end
+    end
+
+    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,
+        "statement_timeout" => 0,
+        "username" => "sharding_user"
+      }
+
+      pgcat = PgcatProcess.new(log_level)
+      pgcat_cfg = pgcat.current_config
+
+      primary  = PgInstance.new(5432, user["username"], user["password"], "shard0")
+      replica0 = PgInstance.new(7432, user["username"], user["password"], "shard0")
+      replica1 = PgInstance.new(8432, user["username"], user["password"], "shard0")
+      replica2 = PgInstance.new(9432, user["username"], user["password"], "shard0")
+
+      # Main proxy configs
+      pgcat_cfg["pools"] = {
+        "#{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",
+          "shards" => {
+            "0" => {
+              "database" => "shard0",
+              "servers" => [
+                ["localhost", primary.port.to_s, "primary"],
+                ["localhost", replica0.port.to_s, "replica"],
+                ["localhost", replica1.port.to_s, "replica"],
+                ["localhost", replica2.port.to_s, "replica"]
+              ]
+            },
+          },
+          "users" => { "0" => user }
+        }
+      }
+      pgcat_cfg["general"]["port"] = pgcat.port
+      pgcat.update_config(pgcat_cfg)
+      pgcat.start
+      pgcat.wait_until_ready
+
+      OpenStruct.new.tap do |struct|
+        struct.pgcat = pgcat
+        struct.primary = primary
+        struct.replicas = [replica0, replica1, replica2]
+        struct.all_databases = [primary, replica0, replica1, replica2]
+      end
+    end
+  end
+end

tests/ruby/helpers/pgcat_process.rb

@@ -0,0 +1,131 @@
+require 'pg'
+require 'toml'
+require 'fileutils'
+require 'securerandom'
+
+class PgcatProcess
+  attr_reader :port
+  attr_reader :pid
+
+  def self.finalize(pid, log_filename, config_filename)
+    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
+
+  def initialize(log_level)
+    @env = {"RUST_LOG" => log_level}
+    @port = rand(20000..32760)
+    @log_level = log_level
+    @log_filename = "/tmp/pgcat_log_#{SecureRandom.urlsafe_base64}.log"
+    @config_filename = "/tmp/pgcat_cfg_#{SecureRandom.urlsafe_base64}.toml"
+
+    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
+    cfg["general"]["port"] = @port.to_i
+    cfg["general"]["enable_prometheus_exporter"] = false
+
+    update_config(cfg)
+  end
+
+  def logs
+    File.read(@log_filename)
+  end
+
+  def update_config(config_hash)
+    @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
+    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.5
+  end
+
+  def start
+    raise StandardError, "Process is already started" unless @pid.nil?
+    @pid = Process.spawn(@env, @command, err: @log_filename, out: @log_filename)
+    ObjectSpace.define_finalizer(@log_filename, proc { PgcatProcess.finalize(@pid, @log_filename, @config_filename) })
+
+    return self
+  end
+
+  def wait_until_ready
+    exc = nil
+    10.times do
+      PG::connect(example_connection_string).close
+
+      return self
+    rescue => e
+      exc = e
+      sleep(0.5)
+    end
+    puts exc
+    raise StandardError, "Process #{@pid} never became ready. Logs #{logs}"
+  end
+
+  def stop
+    return unless @pid
+
+    Process.kill("TERM", @pid)
+    Process.wait(@pid)
+    @pid = nil
+  end
+
+  def shutdown
+    stop
+    File.delete(@config_filename) if File.exist?(@config_filename)
+    File.delete(@log_filename) if File.exist?(@log_filename)
+  end
+
+  def admin_connection_string
+    cfg = current_config
+    username = cfg["general"]["admin_username"]
+    password = cfg["general"]["admin_password"]
+
+    "postgresql://#{username}:#{password}@0.0.0.0:#{@port}/pgcat"
+  end
+
+  def connection_string(pool_name, username)
+    cfg = current_config
+
+    user_idx, user_obj = cfg["pools"][pool_name]["users"].detect { |k, user| user["username"] == username }
+    password = user_obj["password"]
+
+    "postgresql://#{username}:#{password}@0.0.0.0:#{@port}/#{pool_name}"
+  end
+
+  def example_connection_string
+    cfg = current_config
+    first_pool_name = cfg["pools"].keys[0]
+
+    db_name = first_pool_name
+
+    username = cfg["pools"][first_pool_name]["users"]["0"]["username"]
+    password = cfg["pools"][first_pool_name]["users"]["0"]["password"]
+
+    "postgresql://#{username}:#{password}@0.0.0.0:#{@port}/#{db_name}?application_name=example_app"
+  end
+end

tests/ruby/load_balancing_spec.rb

@@ -0,0 +1,164 @@
+# frozen_string_literal: true
+require_relative 'spec_helper'
+
+describe "Random Load Balancing" do
+  let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5) }
+  after do
+    processes.all_databases.map(&:reset)
+    processes.pgcat.shutdown
+  end
+
+  context "under regular circumstances" 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 "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
+
+      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)
+          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
+

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

@@ -0,0 +1,312 @@
+# frozen_string_literal: true
+require_relative 'spec_helper'
+
+describe "Miscellaneous" do
+  let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5) }
+  after do
+    processes.all_databases.map(&:reset)
+    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
+
+      new_configs["general"]["connect_timeout"] = 500
+      new_configs["general"]["ban_time"] = 1
+      new_configs["general"]["shutdown_timeout"] = 1
+      new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
+
+      processes.pgcat.update_config(new_configs)
+      processes.pgcat.reload_config
+
+      25.times do
+        Thread.new do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("SELECT pg_sleep(5)") rescue PG::SystemError
+        ensure
+          conn&.close
+        end
+      end
+
+      sleep(0.5)
+      conn_under_test = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      stdout, stderr = with_captured_stdout_stderr do
+        15.times do |i|
+          conn_under_test.async_exec("SELECT 1") rescue PG::SystemError
+          conn_under_test.exec_params("SELECT #{i} + $1", [i]) rescue PG::SystemError
+          sleep 1
+        end
+      end
+
+      raise StandardError, "Libpq got unexpected messages while idle" if stderr.include?("arrived from server while idle")
+    end
+  end
+
+  describe "Pool recycling after config reload" do
+    let(:processes) { Helpers::Pgcat.three_shard_setup("sharded_db", 5) }
+
+    it "should update pools for new clients and clients that are no longer in transaction" do
+      server_conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      server_conn.async_exec("BEGIN")
+
+      # No config change yet, client should set old configs
+      current_datebase_from_pg = server_conn.async_exec("SELECT current_database();")[0]["current_database"]
+      expect(current_datebase_from_pg).to eq('shard0')
+
+      # Swap shards
+      new_config = processes.pgcat.current_config
+      shard0 = new_config["pools"]["sharded_db"]["shards"]["0"]
+      shard1 = new_config["pools"]["sharded_db"]["shards"]["1"]
+      new_config["pools"]["sharded_db"]["shards"]["0"] = shard1
+      new_config["pools"]["sharded_db"]["shards"]["1"] = shard0
+
+      # Reload config
+      processes.pgcat.update_config(new_config)
+      processes.pgcat.reload_config
+      sleep 0.5
+
+      # Config changed but transaction is in progress, client should set old configs
+      current_datebase_from_pg = server_conn.async_exec("SELECT current_database();")[0]["current_database"]
+      expect(current_datebase_from_pg).to eq('shard0')
+      server_conn.async_exec("COMMIT")
+
+      # Transaction finished, client should get new configs
+      current_datebase_from_pg = server_conn.async_exec("SELECT current_database();")[0]["current_database"]
+      expect(current_datebase_from_pg).to eq('shard1')
+
+      # New connection should get new configs
+      server_conn.close()
+      server_conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      current_datebase_from_pg = server_conn.async_exec("SELECT current_database();")[0]["current_database"]
+      expect(current_datebase_from_pg).to eq('shard1')
+    end
+  end
+
+  describe "Clients closing connection in the middle of transaction" do
+    it "sends a rollback to the server" 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.close
+
+      expect(processes.primary.count_query("ROLLBACK")).to eq(1)
+    end
+  end
+
+  describe "Server version reporting" do
+    it "reports correct version for normal and admin databases" do
+      server_conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+      expect(server_conn.server_version).not_to eq(0)
+      server_conn.close
+
+      admin_conn = PG::connect(processes.pgcat.admin_connection_string)
+      expect(admin_conn.server_version).not_to eq(0)
+      admin_conn.close
+    end
+  end
+
+  describe "State clearance" do
+    context "session mode" do
+      let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5, "session") }
+
+      it "Clears state before connection checkin" do
+        # Both modes of operation should not raise
+        # ERROR:  prepared statement "prepared_q" already exists
+        15.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("PREPARE prepared_q (int) AS SELECT $1")
+          conn.close
+        end
+
+        conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+        initial_value = conn.async_exec("SHOW statement_timeout")[0]["statement_timeout"]
+        conn.async_exec("SET statement_timeout to 1000")
+        current_value = conn.async_exec("SHOW statement_timeout")[0]["statement_timeout"]
+        expect(conn.async_exec("SHOW statement_timeout")[0]["statement_timeout"]).to eq("1s")
+        conn.close
+      end
+
+      it "Does not send DISCARD ALL unless necessary" 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("SELECT 1")
+          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("SELECT 1")
+          conn.async_exec("SET statement_timeout to 5000")
+          conn.close
+        end
+
+        expect(processes.primary.count_query("DISCARD ALL")).to eq(10)
+      end
+    end
+
+    context "transaction mode" do
+      let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5, "transaction") }
+      it "Clears state before connection checkin" do
+        # Both modes of operation should not raise
+        # ERROR:  prepared statement "prepared_q" already exists
+        15.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.async_exec("PREPARE prepared_q (int) AS SELECT $1")
+          conn.close
+        end
+
+        15.times do
+          conn = PG::connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+          conn.prepare("prepared_q", "SELECT $1")
+          conn.close
+        end
+      end
+
+      it "Does not send DISCARD ALL unless necessary" 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("SELECT 1")
+          conn.exec_params("SELECT $1", [1])
+          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("SELECT 1")
+          conn.async_exec("SET statement_timeout to 5000")
+          conn.close
+        end
+
+        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/routing_spec.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+require_relative 'spec_helper'
+
+
+describe "Routing" do
+  let(:processes) { Helpers::Pgcat.single_shard_setup("sharded_db", 5) }
+  after do
+    processes.all_databases.map(&:reset)
+    processes.pgcat.shutdown
+  end
+
+  describe "SET ROLE" do
+    context "primary" do
+      it "routes queries only to primary" do
+        conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+        conn.async_exec("SET SERVER ROLE to 'primary'")
+
+        query_count = 30
+        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.replicas.map(&:count_select_1_plus_2).each do |instance_share|
+          expect(instance_share).to eq(0)
+        end
+
+        expect(processes.primary.count_select_1_plus_2).to eq(query_count)
+      end
+    end
+    context "replica" do
+      it "routes queries only to replicas" do
+        conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+        conn.async_exec("SET SERVER ROLE to 'replica'")
+
+        expected_share = QUERY_COUNT / processes.replicas.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.replicas.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
+
+        expect(processes.primary.count_select_1_plus_2).to eq(0)
+      end
+    end
+
+    context "any" do
+      it "routes queries to all instances" do
+        conn = PG.connect(processes.pgcat.connection_string("sharded_db", "sharding_user"))
+        conn.async_exec("SET SERVER ROLE to 'any'")
+
+        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
+  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

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require 'pg'
+require_relative 'helpers/pgcat_helper'
+
+QUERY_COUNT = 300
+MARGIN_OF_ERROR = 0.35
+
+def with_captured_stdout_stderr
+  sout = STDOUT.clone
+  serr = STDERR.clone
+  STDOUT.reopen("/tmp/out.txt", "w+")
+  STDERR.reopen("/tmp/err.txt", "w+")
+  STDOUT.sync = true
+  STDERR.sync = true
+  yield
+  return File.read('/tmp/out.txt'), File.read('/tmp/err.txt')
+ensure
+  STDOUT.reopen(sout)
+  STDERR.reopen(serr)
+end

tests/ruby/tests.rb

@@ -1,93 +1,6 @@
 # frozen_string_literal: true
-
-require 'active_record'
 require 'pg'
-require 'toml'
-
-$stdout.sync = true
-$stderr.sync = true
-
-class ConfigEditor
-  def initialize
-    @original_config_text = File.read('../../.circleci/pgcat.toml')
-    text_to_load = @original_config_text.gsub("5432", "\"5432\"")
-
-    @original_configs = TOML.load(text_to_load)
-  end
-
-  def original_configs
-    TOML.load(TOML::Generator.new(@original_configs).body)
-  end
-
-  def with_modified_configs(new_configs)
-    text_to_write = TOML::Generator.new(new_configs).body
-    text_to_write = text_to_write.gsub("\"5432\"", "5432")
-    File.write('../../.circleci/pgcat.toml', text_to_write)
-    yield
-  ensure
-    File.write('../../.circleci/pgcat.toml', @original_config_text)
-  end
-end
-
-def with_captured_stdout_stderr
-  sout = STDOUT.clone
-  serr = STDERR.clone
-  STDOUT.reopen("/tmp/out.txt", "w+")
-  STDERR.reopen("/tmp/err.txt", "w+")
-  STDOUT.sync = true
-  STDERR.sync = true
-  yield
-  return File.read('/tmp/out.txt'), File.read('/tmp/err.txt')
-ensure
-  STDOUT.reopen(sout)
-  STDERR.reopen(serr)
-end
-
-
-def test_extended_protocol_pooler_errors
-  admin_conn = PG::connect("postgres://admin_user:admin_pass@127.0.0.1:6432/pgcat")
-
-  conf_editor = ConfigEditor.new
-  new_configs = conf_editor.original_configs
-
-  # shorter timeouts
-  new_configs["general"]["connect_timeout"] = 500
-  new_configs["general"]["ban_time"] = 1
-  new_configs["general"]["shutdown_timeout"] = 1
-  new_configs["pools"]["sharded_db"]["users"]["0"]["pool_size"] = 1
-  new_configs["pools"]["sharded_db"]["users"]["1"]["pool_size"] = 1
-
-  conf_editor.with_modified_configs(new_configs) { admin_conn.async_exec("RELOAD") }
-
-  conn_str = "postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db"
-  10.times do
-    Thread.new do
-      conn = PG::connect(conn_str)
-      conn.async_exec("SELECT pg_sleep(5)") rescue PG::SystemError
-    ensure
-      conn&.close
-    end
-  end
-
-  sleep(0.5)
-  conn_under_test = PG::connect(conn_str)
-  stdout, stderr = with_captured_stdout_stderr do
-    5.times do |i|
-      conn_under_test.async_exec("SELECT 1") rescue PG::SystemError
-      conn_under_test.exec_params("SELECT #{i} + $1", [i]) rescue PG::SystemError
-      sleep 1
-    end
-  end
-
-  raise StandardError, "Libpq got unexpected messages while idle" if stderr.include?("arrived from server while idle")
-  puts "Pool checkout errors not breaking clients passed"
-ensure
-  sleep 1
-  admin_conn.async_exec("RELOAD") # Reset state
-  conn_under_test&.close
-end
-
-test_extended_protocol_pooler_errors
+require 'active_record'
 
 # Uncomment these two to see all queries.
 # ActiveRecord.verbose_query_logs = true
@@ -198,68 +111,3 @@ begin
 rescue ActiveRecord::StatementInvalid
   puts 'OK'
 end
-
-# Test evil clients
-def poorly_behaved_client
-  conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db?application_name=testing_pgcat")
-  conn.async_exec 'BEGIN'
-  conn.async_exec 'SELECT 1'
-
-  conn.close
-  puts 'Bad client ok'
-end
-
-25.times do
-  poorly_behaved_client
-end
-
-
-def test_server_parameters
-  server_conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db?application_name=testing_pgcat")
-  raise StandardError, "Bad server version" if server_conn.server_version == 0
-  server_conn.close
-
-  admin_conn = PG::connect("postgres://admin_user:admin_pass@127.0.0.1:6432/pgcat")
-  raise StandardError, "Bad server version" if admin_conn.server_version == 0
-  admin_conn.close
-
-  puts 'Server parameters ok'
-end
-
-
-def test_reload_pool_recycling
-  admin_conn = PG::connect("postgres://admin_user:admin_pass@127.0.0.1:6432/pgcat")
-  server_conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db?application_name=testing_pgcat")
-
-  server_conn.async_exec("BEGIN")
-  conf_editor = ConfigEditor.new
-  new_configs = conf_editor.original_configs
-
-  # swap shards
-  new_configs["pools"]["sharded_db"]["shards"]["0"]["database"] = "shard1"
-  new_configs["pools"]["sharded_db"]["shards"]["1"]["database"] = "shard0"
-
-  raise StandardError if server_conn.async_exec("SELECT current_database();")[0]["current_database"] != 'shard0'
-  conf_editor.with_modified_configs(new_configs) { admin_conn.async_exec("RELOAD") }
-  raise StandardError if server_conn.async_exec("SELECT current_database();")[0]["current_database"] != 'shard0'
-  server_conn.async_exec("COMMIT;")
-
-  # Transaction finished, client should get new configs
-  raise StandardError if server_conn.async_exec("SELECT current_database();")[0]["current_database"] != 'shard1'
-  server_conn.close()
-
-  # New connection should get new configs
-  server_conn = PG::connect("postgres://sharding_user:sharding_user@127.0.0.1:6432/sharded_db?application_name=testing_pgcat")
-  raise StandardError if server_conn.async_exec("SELECT current_database();")[0]["current_database"] != 'shard1'
-
-ensure
-  admin_conn.async_exec("RELOAD") # Go back to old state
-  admin_conn.close
-  server_conn.close
-  puts "Pool Recycling okay!"
-end
-
-test_reload_pool_recycling
-
-
-

tests/sharding/query_routing_setup.sql

@@ -70,23 +70,35 @@ GRANT CONNECT ON DATABASE shard2  TO other_user;
 GRANT CONNECT ON DATABASE some_db TO simple_user;
 
 \c shard0
+CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO sharding_user;
 GRANT ALL ON SCHEMA public TO sharding_user;
 GRANT ALL ON TABLE data TO sharding_user;
 GRANT ALL ON SCHEMA public TO other_user;
 GRANT ALL ON TABLE data TO other_user;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO other_user;
 
 \c shard1
+CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO sharding_user;
 GRANT ALL ON SCHEMA public TO sharding_user;
 GRANT ALL ON TABLE data TO sharding_user;
 GRANT ALL ON SCHEMA public TO other_user;
 GRANT ALL ON TABLE data TO other_user;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO other_user;
+
 
 \c shard2
+CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO sharding_user;
 GRANT ALL ON SCHEMA public TO sharding_user;
 GRANT ALL ON TABLE data TO sharding_user;
 GRANT ALL ON SCHEMA public TO other_user;
 GRANT ALL ON TABLE data TO other_user;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO other_user;
 
 \c some_db
+CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
+GRANT EXECUTE ON FUNCTION pg_stat_statements_reset TO simple_user;
 GRANT ALL ON SCHEMA public TO simple_user;
 GRANT ALL ON TABLE data TO simple_user;

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()
+