query routing docs

* Toxiproxy

* up-to-date config

* debug

* hm

* more

* mroe

* more

* hmm

* aha

* less logs

* cleaner

* hmm

* we test these now

* update readme

.circleci/pgcat.toml

@@ -0,0 +1,108 @@
+#
+# PgCat config example.
+#
+
+#
+# General pooler settings
+[general]
+
+# What IP to run on, 0.0.0.0 means accessible from everywhere.
+host = "0.0.0.0"
+
+# Port to run on, same as PgBouncer used in this example.
+port = 6432
+
+# How many connections to allocate per server.
+pool_size = 15
+
+# Pool mode (see PgBouncer docs for more).
+# session: one server connection per connected client
+# transaction: one server connection per client transaction
+pool_mode = "transaction"
+
+# How long to wait before aborting a server connection (ms).
+connect_timeout = 100
+
+# How much time to give `SELECT 1` health check query to return with a result (ms).
+healthcheck_timeout = 100
+
+# For how long to ban a server if it fails a health check (seconds).
+ban_time = 60 # Seconds
+
+# Stats will be sent here
+statsd_address = "127.0.0.1:8125"
+
+#
+# User to use for authentication against the server.
+[user]
+name = "sharding_user"
+password = "sharding_user"
+
+
+#
+# Shards in the cluster
+[shards]
+
+# Shard 0
+[shards.0]
+
+# [ host, port, role ]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5433, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
+]
+# Database name (e.g. "postgres")
+database = "shard0"
+
+[shards.1]
+# [ host, port, role ]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5433, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
+]
+database = "shard1"
+
+[shards.2]
+# [ host, port, role ]
+servers = [
+    [ "127.0.0.1", 5432, "primary" ],
+    [ "localhost", 5433, "replica" ],
+    # [ "127.0.1.1", 5432, "replica" ],
+]
+database = "shard2"
+
+
+# Settings for our query routing layer.
+[query_router]
+
+# If the client doesn't specify, route traffic to
+# this role by default.
+#
+# any: round-robin between primary and replicas,
+# replica: round-robin between replicas only without touching the primary,
+# primary: all queries go to the primary unless otherwise specified.
+default_role = "any"
+
+
+# Query parser. If enabled, we'll attempt to parse
+# every incoming query to determine if it's a read or a write.
+# If it's a read query, we'll direct it to a replica. Otherwise, if it's a write,
+# we'll direct it to the primary.
+query_parser_enabled = false
+
+# If the query parser is enabled and this setting is enabled, the primary will be part of the pool of databases used for
+# load balancing of read queries. Otherwise, the primary will only be used for write
+# queries. The primary can always be explicitely selected with our custom protocol.
+primary_reads_enabled = true
+
+# So what if you wanted to implement a different hashing function,
+# or you've already built one and you want this pooler to use it?
+#
+# Current options:
+#
+# pg_bigint_hash: PARTITION BY HASH (Postgres hashing function)
+# sha1: A hashing function based on SHA1
+#
+sharding_function = "pg_bigint_hash"

.circleci/run_tests.sh

@@ -3,20 +3,34 @@
 set -e
 set -o xtrace
 
+# Start PgCat with a particular log level
+# for inspection.
+function start_pgcat() {
+    kill -s SIGINT $(pgrep pgcat) || true
+    RUST_LOG=${1} ./target/debug/pgcat .circleci/pgcat.toml &
+    sleep 1
+}
+
+# Setup the database with shards and user
 psql -e -h 127.0.0.1 -p 5432 -U postgres -f tests/sharding/query_routing_setup.sql
 
-./target/debug/pgcat &
+# Install Toxiproxy to simulate a downed/slow database
+wget -O toxiproxy-2.1.4.deb https://github.com/Shopify/toxiproxy/releases/download/v2.1.4/toxiproxy_2.1.4_amd64.deb
+sudo dpkg -i toxiproxy-2.1.4.deb
 
+# Start Toxiproxy
+toxiproxy-server &
 sleep 1
 
-# Setup PgBench
-pgbench -i -h 127.0.0.1 -p 6432
+# Create a database at port 5433, forward it to Postgres
+toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica
 
-# Run it
-pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple
+start_pgcat "info"
 
-# Extended protocol
-pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
+# pgbench test
+pgbench -i -h 127.0.0.1 -p 6432 && \
+    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol simple && \
+    pgbench -h 127.0.0.1 -p 6432 -t 500 -c 2 --protocol extended
 
 # COPY TO STDOUT test
 psql -h 127.0.0.1 -p 6432 -c 'COPY (SELECT * FROM pgbench_accounts LIMIT 15) TO STDOUT;' > /dev/null
@@ -35,18 +49,37 @@ psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_select.sql > /
 psql -e -h 127.0.0.1 -p 6432 -f tests/sharding/query_routing_test_primary_replica.sql > /dev/null
 
 #
-# ActiveRecord tests!
+# ActiveRecord tests
 #
-cd tests/ruby
-sudo gem install bundler
-bundle install
-ruby tests.rb
+cd tests/ruby && \
+    sudo gem install bundler && \
+    bundle install && \
+    ruby tests.rb && \
+cd ../..
+
+# Start PgCat in debug to demonstrate failover better
+start_pgcat "debug"
+
+# Add latency to the replica at port 5433 slightly above the healthcheck timeout
+toxiproxy-cli toxic add -t latency -a latency=300 postgres_replica
+sleep 1
+
+# Note the failover in the logs
+timeout 5 psql -e -h 127.0.0.1 -p 6432 <<-EOF
+SELECT 1;
+SELECT 1;
+SELECT 1;
+EOF
+
+# Remove latency
+toxiproxy-cli toxic remove --toxicName latency_downstream postgres_replica
+
+start_pgcat "info"
 
-cd ../../
 # Test session mode (and config reload)
 sed -i 's/pool_mode = "transaction"/pool_mode = "session"/' pgcat.toml
 
-# Reload config
+# Reload config test
 kill -SIGHUP $(pgrep pgcat)
 
 # Prepared statements that will only work in session mode

Cargo.lock

@@ -364,6 +364,7 @@ dependencies = [
  "md-5",
  "num_cpus",
  "once_cell",
+ "parking_lot",
  "rand",
  "regex",
  "serde",

Cargo.toml

@@ -25,3 +25,4 @@ sqlparser = "0.14"
 log = "0.4"
 arc-swap = "1"
 env_logger = "0.9"
+parking_lot = "0.11"

README.md

@@ -11,14 +11,14 @@ Meow. PgBouncer rewritten in Rust, with sharding, load balancing and failover su
 ## Features
 | **Feature**                    | **Status**            | **Comments**                                                                                                                                          |
 |--------------------------------|-----------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------|
-| Transaction pooling            | :heavy_check_mark:    | Identical to PgBouncer.                                                                                                                               |
-| Session pooling                | :heavy_check_mark:    | Identical to PgBouncer.                                                                                                                               |
-| `COPY` support                 | :heavy_check_mark:    | Both `COPY TO` and `COPY FROM` are supported.                                                                                                         |
-| Query cancellation             | :heavy_check_mark:    | Supported both in transaction and session pooling modes.                                                                                              |
-| Load balancing of read queries | :heavy_check_mark:    | Using round-robin between replicas. Primary is included when `primary_reads_enabled` is enabled (default).                                            |
-| Sharding                       | :heavy_check_mark:    | Transactions are sharded using `SET SHARD TO` and `SET SHARDING KEY TO` syntax extensions; see examples below.                                        |
-| Failover                       | :heavy_check_mark:    | Replicas are tested with a health check. If a health check fails, remaining replicas are attempted; see below for algorithm description and examples. |
-| Statistics reporting           | :heavy_check_mark:    | Statistics similar to PgBouncers are reported via StatsD.                                                                                             |
+| Transaction pooling            | :white_check_mark:    | Identical to PgBouncer.                                                                                                                               |
+| Session pooling                | :white_check_mark:    | Identical to PgBouncer.                                                                                                                               |
+| `COPY` support                 | :white_check_mark:    | Both `COPY TO` and `COPY FROM` are supported.                                                                                                         |
+| Query cancellation             | :white_check_mark:    | Supported both in transaction and session pooling modes.                                                                                              |
+| Load balancing of read queries | :white_check_mark:    | Using round-robin between replicas. Primary is included when `primary_reads_enabled` is enabled (default).                                            |
+| Sharding                       | :white_check_mark:    | Transactions are sharded using `SET SHARD TO` and `SET SHARDING KEY TO` syntax extensions; see examples below.                                        |
+| Failover                       | :white_check_mark:    | Replicas are tested with a health check. If a health check fails, remaining replicas are attempted; see below for algorithm description and examples. |
+| Statistics reporting           | :white_check_mark:    | Statistics similar to PgBouncers are reported via StatsD.                                                                                             |
 | Live configuration reloading   | :construction_worker: | Reload config with a `SIGHUP` to the process, e.g. `kill -s SIGHUP $(pgrep pgcat)`. Not all settings can be reloaded without a restart.               |
 | Client authentication          | :x: :wrench:          | On the roadmap; currently all clients are allowed to connect and one user is used to connect to Postgres.                                             |
 
@@ -75,15 +75,15 @@ See [sharding README](./tests/sharding/README.md) for sharding logic testing.
 
 | **Feature**           | **Tested in CI**   | **Tested manually** | **Comments**                                                                                                             |
 |-----------------------|--------------------|---------------------|--------------------------------------------------------------------------------------------------------------------------|
-| Transaction pooling   | :heavy_check_mark: | :heavy_check_mark:  | Used by default for all tests.                                                                                           |
-| Session pooling       | :heavy_check_mark: | :heavy_check_mark:  | Tested by running pgbench with `--protocol prepared` which only works in session mode.                                   |
-| `COPY`                | :heavy_check_mark: | :heavy_check_mark:  | `pgbench -i` uses `COPY`. `COPY FROM` is tested as well.                                                                 |
-| Query cancellation    | :heavy_check_mark: | :heavy_check_mark:  | `psql -c 'SELECT pg_sleep(1000);'` and press `Ctrl-C`.                                                                   |
-| Load balancing        | :x:                | :heavy_check_mark:  | We could test this by emitting statistics for each replica and compare them.                                             |
-| Failover              | :x:                | :heavy_check_mark:  | Misconfigure a replica in `pgcat.toml` and watch it forward queries to spares. CI testing could include using Toxiproxy. |
-| Sharding              | :heavy_check_mark: | :heavy_check_mark:  | See `tests/sharding` and `tests/ruby` for an Rails/ActiveRecord example.                                                 |
-| Statistics reporting  | :x:                | :heavy_check_mark:  | Run `nc -l -u 8125` and watch the stats come in every 15 seconds.                                                        |
-| Live config reloading | :heavy_check_mark: | :heavy_check_mark:  | Run `kill -s SIGHUP $(pgrep pgcat)` and watch the config reload.                                                         |
+| Transaction pooling   | :white_check_mark: | :white_check_mark:  | Used by default for all tests.                                                                                           |
+| Session pooling       | :white_check_mark: | :white_check_mark:  | Tested by running pgbench with `--protocol prepared` which only works in session mode.                                   |
+| `COPY`                | :white_check_mark: | :white_check_mark:  | `pgbench -i` uses `COPY`. `COPY FROM` is tested as well.                                                                 |
+| Query cancellation    | :white_check_mark: | :white_check_mark:  | `psql -c 'SELECT pg_sleep(1000);'` and press `Ctrl-C`.                                                                   |
+| Load balancing        | :white_check_mark: | :white_check_mark:  | We could test this by emitting statistics for each replica and compare them.                                             |
+| Failover              | :white_check_mark: | :white_check_mark:  | Misconfigure a replica in `pgcat.toml` and watch it forward queries to spares. CI testing is using Toxiproxy.            |
+| Sharding              | :white_check_mark: | :white_check_mark:  | See `tests/sharding` and `tests/ruby` for an Rails/ActiveRecord example.                                                 |
+| Statistics reporting  | :x:                | :white_check_mark:  | Run `nc -l -u 8125` and watch the stats come in every 15 seconds.                                                        |
+| Live config reloading | :white_check_mark: | :white_check_mark:  | Run `kill -s SIGHUP $(pgrep pgcat)` and watch the config reload.                                                         |
 
 ## Usage
 
@@ -133,6 +133,31 @@ All servers are checked with a `SELECT 1` query before being given to a client.
 
 The ban time can be changed with `ban_time`. The default is 60 seconds.
 
+Failover behavior can get pretty interesting (read complex) when multiple configurations and factors are involved. The table below will try to explain what PgCat does in each scenario:
+
+| **Query**                 | **`SET SERVER ROLE TO`** | **`query_parser_enabled`** | **`primary_reads_enabled`** | **Target state** | **Outcome**                                                                                                                                                          |
+|---------------------------|--------------------------|----------------------------|-----------------------------|------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| Read query, i.e. `SELECT` | unset (any)              | false                      | false                       | up               | Query is routed to the first instance in the round-robin loop.                                                                                                       |
+| Read query                | unset (any)              | true                       | false                       | up               | Query is routed to the first replica instance in the round-robin loop.                                                                                               |
+| Read query                | unset (any)              | true                       | true                        | up               | Query is routed to the first instance in the round-robin loop.                                                                                                       |
+| Read query                | replica                  | false                      | false                       | up               | Query is routed to the first replica instance in the round-robin loop.                                                                                               |
+| Read query                | primary                  | false                      | false                       | up               | Query is routed to the primary.                                                                                                                                      |
+| Read query                | unset (any)              | false                      | false                       | down             | First instance is banned for reads. Next target in the round-robin loop is attempted.                                                                                |
+| Read query                | unset (any)              | true                       | false                       | down             | First replica instance is banned. Next replica instance is attempted in the round-robin loop.                                                                        |
+| Read query                | unset (any)              | true                       | true                        | down             | First instance (even if primary) is banned for reads. Next instance is attempted in the round-robin loop.                                                            |
+| Read query                | replica                  | false                      | false                       | down             | First replica instance is banned. Next replica instance is attempted in the round-robin loop.                                                                        |
+| Read query                | primary                  | false                      | false                       | down             | The query is attempted against the primary and fails. The client receives an error.                                                                                  |
+|                           |                          |                            |                             |                  |                                                                                                                                                                      |
+| Write query e.g. `INSERT` | unset (any)              | false                      | false                       | up               | The query is attempted against the first available instance in the round-robin loop. If the instance is a replica, the query fails and the client receives an error. |
+| Write query               | unset (any)              | true                       | false                       | up               | The query is routed to the primary.                                                                                                                                  |
+| Write query               | unset (any)              | true                       | true                        | up               | The query is routed to the primary.                                                                                                                                  |
+| Write query               | primary                  | false                      | false                       | up               | The query is routed to the primary.                                                                                                                                  |
+| Write query               | replica                  | false                      | false                       | up               | The query is routed to the replica and fails. The client receives an error.                                                                                          |
+| Write query               | unset (any)              | true                       | false                       | down             | The query is routed to the primary and fails. The client receives an error.                                                                                          |
+| Write query               | unset (any)              | true                       | true                        | down             | The query is routed to the primary and fails. The client receives an error.                                                                                          |
+| Write query               | primary                  | false                      | false                       | down             | The query is routed to the primary and fails. The client receives an error.                                                                                          |
+|                           |                          |                            |                             |                  |                                                                                                                                                                      |
+
 ### Sharding
 We use the `PARTITION BY HASH` hashing function, the same as used by Postgres for declarative partitioning. This allows to shard the database using Postgres partitions and place the partitions on different servers (shards). Both read and write queries can be routed to the shards using this pooler.
 

pgcat.toml

@@ -96,3 +96,13 @@ query_parser_enabled = false
 # load balancing of read queries. Otherwise, the primary will only be used for write
 # queries. The primary can always be explicitely selected with our custom protocol.
 primary_reads_enabled = true
+
+# So what if you wanted to implement a different hashing function,
+# or you've already built one and you want this pooler to use it?
+#
+# Current options:
+#
+# pg_bigint_hash: PARTITION BY HASH (Postgres hashing function)
+# sha1: A hashing function based on SHA1
+#
+sharding_function = "pg_bigint_hash"

src/client.rs

@@ -2,7 +2,7 @@
 /// We are pretending to the server in this scenario,
 /// and this module implements that.
 use bytes::{Buf, BufMut, BytesMut};
-use log::error;
+use log::{debug, error, trace};
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
     tcp::{OwnedReadHalf, OwnedWriteHalf},
@@ -70,6 +70,8 @@ impl Client {
         let transaction_mode = config.general.pool_mode.starts_with("t");
         drop(config);
         loop {
+            trace!("Waiting for StartupMessage");
+
             // Could be StartupMessage or SSLRequest
             // which makes this variable length.
             let len = match stream.read_i32().await {
@@ -91,6 +93,8 @@ impl Client {
             match code {
                 // Client wants SSL. We don't support it at the moment.
                 SSL_REQUEST_CODE => {
+                    trace!("Rejecting SSLRequest");
+
                     let mut no = BytesMut::with_capacity(1);
                     no.put_u8(b'N');
 
@@ -99,6 +103,8 @@ impl Client {
 
                 // Regular startup message.
                 PROTOCOL_VERSION_NUMBER => {
+                    trace!("Got StartupMessage");
+
                     // TODO: perform actual auth.
                     let parameters = parse_startup(bytes.clone())?;
 
@@ -110,6 +116,7 @@ impl Client {
                     write_all(&mut stream, server_info).await?;
                     backend_key_data(&mut stream, process_id, secret_key).await?;
                     ready_for_query(&mut stream).await?;
+                    trace!("Startup OK");
 
                     // Split the read and write streams
                     // so we can control buffering.
@@ -161,8 +168,10 @@ impl Client {
     pub async fn handle(&mut self, mut pool: ConnectionPool) -> Result<(), Error> {
         // The client wants to cancel a query it has issued previously.
         if self.cancel_mode {
+            trace!("Sending CancelRequest");
+
             let (process_id, secret_key, address, port) = {
-                let guard = self.client_server_map.lock().unwrap();
+                let guard = self.client_server_map.lock();
 
                 match guard.get(&(self.process_id, self.secret_key)) {
                     // Drop the mutex as soon as possible.
@@ -193,6 +202,8 @@ impl Client {
         // We expect the client to either start a transaction with regular queries
         // or issue commands for our sharding and server selection protocols.
         loop {
+            trace!("Client idle, waiting for message");
+
             // Client idle, waiting for messages.
             self.stats.client_idle(self.process_id);
 
@@ -203,6 +214,12 @@ impl Client {
             // SET SHARDING KEY TO 'bigint';
             let mut message = read_message(&mut self.read).await?;
 
+            // Avoid taking a server if the client just wants to disconnect.
+            if message[0] as char == 'X' {
+                trace!("Client disconnecting");
+                return Ok(());
+            }
+
             // Handle all custom protocol commands here.
             match query_router.try_execute_command(message.clone()) {
                 // Normal query
@@ -212,11 +229,17 @@ impl Client {
                     }
                 }
 
-                Some((Command::SetShard, _)) | Some((Command::SetShardingKey, _)) => {
+                Some((Command::SetShard, _)) => {
                     custom_protocol_response_ok(&mut self.write, &format!("SET SHARD")).await?;
                     continue;
                 }
 
+                Some((Command::SetShardingKey, _)) => {
+                    custom_protocol_response_ok(&mut self.write, &format!("SET SHARDING KEY"))
+                        .await?;
+                    continue;
+                }
+
                 Some((Command::SetServerRole, _)) => {
                     custom_protocol_response_ok(&mut self.write, "SET SERVER ROLE").await?;
                     continue;
@@ -250,9 +273,14 @@ impl Client {
             // Waiting for server connection.
             self.stats.client_waiting(self.process_id);
 
+            debug!("Waiting for connection from pool");
+
             // Grab a server from the pool: the client issued a regular query.
             let connection = match pool.get(query_router.shard(), query_router.role()).await {
-                Ok(conn) => conn,
+                Ok(conn) => {
+                    debug!("Got connection from pool");
+                    conn
+                }
                 Err(err) => {
                     error!("Could not get connection from pool: {:?}", err);
                     error_response(&mut self.write, "could not get connection from the pool")
@@ -272,11 +300,19 @@ impl Client {
             self.stats.client_active(self.process_id);
             self.stats.server_active(server.process_id());
 
+            debug!(
+                "Client {:?} talking to server {:?}",
+                self.write.peer_addr().unwrap(),
+                server.address()
+            );
+
             // Transaction loop. Multiple queries can be issued by the client here.
             // The connection belongs to the client until the transaction is over,
             // or until the client disconnects if we are in session mode.
             loop {
                 let mut message = if message.len() == 0 {
+                    trace!("Waiting for message inside transaction or in session mode");
+
                     match read_message(&mut self.read).await {
                         Ok(message) => message,
                         Err(err) => {
@@ -303,9 +339,13 @@ impl Client {
                 let code = message.get_u8() as char;
                 let _len = message.get_i32() as usize;
 
+                trace!("Message: {}", code);
+
                 match code {
                     // ReadyForQuery
                     'Q' => {
+                        debug!("Sending query to server");
+
                         // TODO: implement retries here for read-only transactions.
                         server.send(original).await?;
 
@@ -387,6 +427,8 @@ impl Client {
                     // Sync
                     // Frontend (client) is asking for the query result now.
                     'S' => {
+                        debug!("Sending query to server");
+
                         self.buffer.put(&original[..]);
 
                         // TODO: retries for read-only transactions.
@@ -471,13 +513,14 @@ impl Client {
             }
 
             // The server is no longer bound to us, we can't cancel it's queries anymore.
+            debug!("Releasing server back into the pool");
             self.release();
         }
     }
 
     /// Release the server from being mine. I can't cancel its queries anymore.
     pub fn release(&self) {
-        let mut guard = self.client_server_map.lock().unwrap();
+        let mut guard = self.client_server_map.lock();
         guard.remove(&(self.process_id, self.secret_key));
     }
 }

src/config.rs

@@ -118,6 +118,7 @@ pub struct QueryRouter {
     pub default_role: String,
     pub query_parser_enabled: bool,
     pub primary_reads_enabled: bool,
+    pub sharding_function: String,
 }
 
 impl Default for QueryRouter {
@@ -126,6 +127,7 @@ impl Default for QueryRouter {
             default_role: String::from("any"),
             query_parser_enabled: false,
             primary_reads_enabled: true,
+            sharding_function: "pg_bigint_hash".to_string(),
         }
     }
 }
@@ -159,6 +161,8 @@ impl Config {
             self.general.healthcheck_timeout
         );
         info!("Connection timeout: {}ms", self.general.connect_timeout);
+        info!("Sharding function: {}", self.query_router.sharding_function);
+        info!("Number of shards: {}", self.shards.len());
     }
 }
 
@@ -193,6 +197,18 @@ pub async fn parse(path: &str) -> Result<(), Error> {
         }
     };
 
+    match config.query_router.sharding_function.as_ref() {
+        "pg_bigint_hash" => (),
+        "sha1" => (),
+        _ => {
+            error!(
+                "Supported sharding functions are: 'pg_bigint_hash', 'sha1', got: '{}'",
+                config.query_router.sharding_function
+            );
+            return Err(Error::BadConfig);
+        }
+    };
+
     // Quick config sanity check.
     for shard in &config.shards {
         // We use addresses as unique identifiers,

src/main.rs

@@ -36,6 +36,7 @@ extern crate tokio;
 extern crate toml;
 
 use log::{error, info};
+use parking_lot::Mutex;
 use tokio::net::TcpListener;
 use tokio::{
     signal,
@@ -44,7 +45,7 @@ use tokio::{
 };
 
 use std::collections::HashMap;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 
 mod client;
 mod config;
@@ -111,9 +112,9 @@ async fn main() {
 
     // Collect statistics and send them to StatsD
     let (tx, rx) = mpsc::channel(100);
-
+    let collector_tx = tx.clone();
     tokio::task::spawn(async move {
-        let mut stats_collector = Collector::new(rx);
+        let mut stats_collector = Collector::new(rx, collector_tx);
         stats_collector.collect().await;
     });
 

src/messages.rs

@@ -38,7 +38,7 @@ pub async fn backend_key_data(
     Ok(write_all(stream, key_data).await?)
 }
 
-#[allow(dead_code)]
+/// Construct a `Q`: Query message.
 pub fn simple_query(query: &str) -> BytesMut {
     let mut res = BytesMut::from(&b"Q"[..]);
     let query = format!("{}\0", query);

src/pool.rs

@@ -3,7 +3,8 @@ use async_trait::async_trait;
 use bb8::{ManageConnection, Pool, PooledConnection};
 use bytes::BytesMut;
 use chrono::naive::NaiveDateTime;
-use log::{error, info, warn};
+use log::{debug, error, info, warn};
+use parking_lot::{Mutex, RwLock};
 
 use crate::config::{get_config, Address, Role, User};
 use crate::errors::Error;
@@ -11,11 +12,11 @@ use crate::server::Server;
 use crate::stats::Reporter;
 
 use std::collections::HashMap;
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 use std::time::Instant;
 
 // Banlist: bad servers go in here.
-pub type BanList = Arc<Mutex<Vec<HashMap<Address, NaiveDateTime>>>>;
+pub type BanList = Arc<RwLock<Vec<HashMap<Address, NaiveDateTime>>>>;
 pub type ClientServerMap = Arc<Mutex<HashMap<(i32, i32), (i32, i32, String, String)>>>;
 
 #[derive(Clone, Debug)]
@@ -101,7 +102,7 @@ impl ConnectionPool {
             databases: shards,
             addresses: addresses,
             round_robin: rand::random::<usize>() % address_len, // Start at a random replica
-            banlist: Arc::new(Mutex::new(banlist)),
+            banlist: Arc::new(RwLock::new(banlist)),
             stats: stats,
         }
     }
@@ -161,6 +162,8 @@ impl ConnectionPool {
             _ => addresses.len(),
         };
 
+        debug!("Allowed attempts for {:?}: {}", role, allowed_attempts);
+
         let exists = match role {
             Some(role) => addresses.iter().filter(|addr| addr.role == role).count() > 0,
             None => true,
@@ -251,14 +254,14 @@ impl ConnectionPool {
     pub fn ban(&self, address: &Address, shard: usize) {
         error!("Banning {:?}", address);
         let now = chrono::offset::Utc::now().naive_utc();
-        let mut guard = self.banlist.lock().unwrap();
+        let mut guard = self.banlist.write();
         guard[shard].insert(address.clone(), now);
     }
 
     /// Clear the replica to receive traffic again. Takes effect immediately
     /// for all new transactions.
     pub fn _unban(&self, address: &Address, shard: usize) {
-        let mut guard = self.banlist.lock().unwrap();
+        let mut guard = self.banlist.write();
         guard[shard].remove(address);
     }
 
@@ -274,12 +277,14 @@ impl ConnectionPool {
             Some(Role::Primary) => return false, // Primary cannot be banned.
         };
 
-        // If you're not asking for the primary,
-        // all databases are treated as replicas.
-        let mut guard = self.banlist.lock().unwrap();
+        debug!("Available targets for {:?}: {}", role, replicas_available);
+
+        let guard = self.banlist.read();
 
         // Everything is banned = nothing is banned.
         if guard[shard].len() == replicas_available {
+            drop(guard);
+            let mut guard = self.banlist.write();
             guard[shard].clear();
             drop(guard);
             warn!("Unbanning all replicas.");
@@ -291,16 +296,24 @@ impl ConnectionPool {
             Some(timestamp) => {
                 let now = chrono::offset::Utc::now().naive_utc();
                 let config = get_config();
+
                 // Ban expired.
                 if now.timestamp() - timestamp.timestamp() > config.general.ban_time {
+                    drop(guard);
+                    warn!("Unbanning {:?}", address);
+                    let mut guard = self.banlist.write();
                     guard[shard].remove(address);
                     false
                 } else {
+                    debug!("{:?} is banned", address);
                     true
                 }
             }
 
-            None => false,
+            None => {
+                debug!("{:?} is ok", address);
+                false
+            }
         }
     }
 

src/query_router.rs

@@ -1,21 +1,21 @@
 use crate::config::{get_config, Role};
-use crate::sharding::Sharder;
+use crate::sharding::{Sharder, ShardingFunction};
 /// Route queries automatically based on explicitely requested
 /// or implied query characteristics.
 use bytes::{Buf, BytesMut};
 use log::{debug, error};
 use once_cell::sync::OnceCell;
-use regex::RegexSet;
+use regex::{Regex, RegexSet};
 use sqlparser::ast::Statement::{Query, StartTransaction};
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
 
 const CUSTOM_SQL_REGEXES: [&str; 5] = [
-    r"(?i)SET SHARDING KEY TO '[0-9]+'",
-    r"(?i)SET SHARD TO '[0-9]+'",
-    r"(?i)SHOW SHARD",
-    r"(?i)SET SERVER ROLE TO '(PRIMARY|REPLICA|ANY|AUTO|DEFAULT)'",
-    r"(?i)SHOW SERVER ROLE",
+    r"(?i)^ *SET SHARDING KEY TO '?([0-9]+)'? *;? *$",
+    r"(?i)^ *SET SHARD TO '?([0-9]+)'? *;? *$",
+    r"(?i)^ *SHOW SHARD *;? *$",
+    r"(?i)^ *SET SERVER ROLE TO '(PRIMARY|REPLICA|ANY|AUTO|DEFAULT)' *;? *$",
+    r"(?i)^ *SHOW SERVER ROLE *;? *$",
 ];
 
 #[derive(PartialEq, Debug)]
@@ -27,8 +27,12 @@ pub enum Command {
     ShowServerRole,
 }
 
+// Quick test
 static CUSTOM_SQL_REGEX_SET: OnceCell<RegexSet> = OnceCell::new();
 
+// Capture value
+static CUSTOM_SQL_REGEX_LIST: OnceCell<Vec<Regex>> = OnceCell::new();
+
 pub struct QueryRouter {
     // By default, queries go here, unless we have better information
     // about what the client wants.
@@ -48,6 +52,9 @@ pub struct QueryRouter {
 
     // Should we try to parse queries?
     query_parser_enabled: bool,
+
+    // Which sharding function are we using?
+    sharding_function: ShardingFunction,
 }
 
 impl QueryRouter {
@@ -60,6 +67,21 @@ impl QueryRouter {
             }
         };
 
+        let list: Vec<_> = CUSTOM_SQL_REGEXES
+            .iter()
+            .map(|rgx| Regex::new(rgx).unwrap())
+            .collect();
+
+        // Impossible
+        if list.len() != set.len() {
+            return false;
+        }
+
+        match CUSTOM_SQL_REGEX_LIST.set(list) {
+            Ok(_) => true,
+            Err(_) => return false,
+        };
+
         match CUSTOM_SQL_REGEX_SET.set(set) {
             Ok(_) => true,
             Err(_) => false,
@@ -76,6 +98,12 @@ impl QueryRouter {
             _ => unreachable!(),
         };
 
+        let sharding_function = match config.query_router.sharding_function.as_ref() {
+            "pg_bigint_hash" => ShardingFunction::PgBigintHash,
+            "sha1" => ShardingFunction::Sha1,
+            _ => unreachable!(),
+        };
+
         QueryRouter {
             default_server_role: default_server_role,
             shards: config.shards.len(),
@@ -84,6 +112,7 @@ impl QueryRouter {
             active_shard: None,
             primary_reads_enabled: config.query_router.primary_reads_enabled,
             query_parser_enabled: config.query_router.query_parser_enabled,
+            sharding_function,
         }
     }
 
@@ -103,6 +132,11 @@ impl QueryRouter {
             None => return None,
         };
 
+        let regex_list = match CUSTOM_SQL_REGEX_LIST.get() {
+            Some(regex_list) => regex_list,
+            None => return None,
+        };
+
         let matches: Vec<_> = regex_set.matches(&query).into_iter().collect();
 
         if matches.len() != 1 {
@@ -120,7 +154,19 @@ impl QueryRouter {
 
         let mut value = match command {
             Command::SetShardingKey | Command::SetShard | Command::SetServerRole => {
-                query.split("'").collect::<Vec<&str>>()[1].to_string()
+                // Capture value. I know this re-runs the regex engine, but I haven't
+                // figured out a better way just yet. I think I can write a single Regex
+                // that matches all 5 custom SQL patterns, but maybe that's not very legible?
+                //
+                // I think this is faster than running the Regex engine 5 times, so
+                // this is a strong maybe for me so far.
+                match regex_list[matches[0]].captures(&query) {
+                    Some(captures) => match captures.get(1) {
+                        Some(value) => value.as_str().to_string(),
+                        None => return None,
+                    },
+                    None => return None,
+                }
             }
 
             Command::ShowShard => self.shard().to_string(),
@@ -139,8 +185,8 @@ impl QueryRouter {
 
         match command {
             Command::SetShardingKey => {
-                let sharder = Sharder::new(self.shards);
-                let shard = sharder.pg_bigint_hash(value.parse::<i64>().unwrap());
+                let sharder = Sharder::new(self.shards, self.sharding_function);
+                let shard = sharder.shard(value.parse::<i64>().unwrap());
                 self.active_shard = Some(shard);
                 value = shard.to_string();
             }
@@ -194,7 +240,12 @@ impl QueryRouter {
         let len = buf.get_i32() as usize;
 
         let query = match code {
-            'Q' => String::from_utf8_lossy(&buf[..len - 5]).to_string(),
+            'Q' => {
+                let query = String::from_utf8_lossy(&buf[..len - 5]).to_string();
+                debug!("Query: '{}'", query);
+                query
+            }
+
             'P' => {
                 let mut start = 0;
                 let mut end;
@@ -213,6 +264,8 @@ impl QueryRouter {
 
                 let query = String::from_utf8_lossy(&buf[start..end]).to_string();
 
+                debug!("Prepared statement: '{}'", query);
+
                 query.replace("$", "") // Remove placeholders turning them into "values"
             }
             _ => return false,
@@ -221,7 +274,7 @@ impl QueryRouter {
         let ast = match Parser::parse_sql(&PostgreSqlDialect {}, &query) {
             Ok(ast) => ast,
             Err(err) => {
-                debug!("{:?}, query: {}", err, query);
+                debug!("{}", err.to_string());
                 return false;
             }
         };
@@ -394,14 +447,38 @@ mod test {
             "set server role to 'any'",
             "set server role to 'auto'",
             "show server role",
+            // No quotes
+            "SET SHARDING KEY TO 11235",
+            "SET SHARD TO 15",
+            // Spaces and semicolon
+            "  SET SHARDING KEY TO 11235  ; ",
+            "  SET SHARD TO 15;   ",
+            "  SET SHARDING KEY TO 11235  ;",
+            " SET SERVER ROLE TO 'primary';   ",
+            "    SET SERVER ROLE TO 'primary'  ; ",
+            "  SET SERVER ROLE TO 'primary'  ;",
         ];
 
+        // Which regexes it'll match to in the list
+        let matches = [
+            0, 1, 2, 3, 3, 3, 3, 4, 0, 1, 2, 3, 3, 3, 3, 4, 0, 1, 0, 1, 0, 3, 3, 3,
+        ];
+
+        let list = CUSTOM_SQL_REGEX_LIST.get().unwrap();
         let set = CUSTOM_SQL_REGEX_SET.get().unwrap();
 
-        for test in &tests {
-            let matches: Vec<_> = set.matches(test).into_iter().collect();
+        for (i, test) in tests.iter().enumerate() {
+            assert!(list[matches[i]].is_match(test));
+            assert_eq!(set.matches(test).into_iter().collect::<Vec<_>>().len(), 1);
+        }
 
-            assert_eq!(matches.len(), 1);
+        let bad = [
+            "SELECT * FROM table",
+            "SELECT * FROM table WHERE value = 'set sharding key to 5'", // Don't capture things in the middle of the query
+        ];
+
+        for query in &bad {
+            assert_eq!(set.matches(query).into_iter().collect::<Vec<_>>().len(), 0);
         }
     }
 
@@ -411,7 +488,7 @@ mod test {
         let mut qr = QueryRouter::new();
 
         // SetShardingKey
-        let query = simple_query("SET SHARDING KEY TO '13'");
+        let query = simple_query("SET SHARDING KEY TO 13");
         assert_eq!(
             qr.try_execute_command(query),
             Some((Command::SetShardingKey, String::from("1")))

src/server.rs

@@ -1,7 +1,7 @@
 use bytes::{Buf, BufMut, BytesMut};
 ///! Implementation of the PostgreSQL server (database) protocol.
 ///! Here we are pretending to the a Postgres client.
-use log::{error, info};
+use log::{debug, error, info, trace};
 use tokio::io::{AsyncReadExt, BufReader};
 use tokio::net::{
     tcp::{OwnedReadHalf, OwnedWriteHalf},
@@ -75,6 +75,8 @@ impl Server {
                 }
             };
 
+        trace!("Sending StartupMessage");
+
         // Send the startup packet telling the server we're a normal Postgres client.
         startup(&mut stream, &user.name, database).await?;
 
@@ -95,6 +97,8 @@ impl Server {
                 Err(_) => return Err(Error::SocketError),
             };
 
+            trace!("Message: {}", code);
+
             match code {
                 // Authentication
                 'R' => {
@@ -104,6 +108,8 @@ impl Server {
                         Err(_) => return Err(Error::SocketError),
                     };
 
+                    trace!("Auth: {}", auth_code);
+
                     match auth_code {
                         MD5_ENCRYPTED_PASSWORD => {
                             // The salt is 4 bytes.
@@ -135,6 +141,8 @@ impl Server {
                         Err(_) => return Err(Error::SocketError),
                     };
 
+                    trace!("Error: {}", error_code);
+
                     match error_code {
                         // No error message is present in the message.
                         MESSAGE_TERMINATOR => (),
@@ -247,6 +255,8 @@ impl Server {
             }
         };
 
+        debug!("Sending CancelRequest");
+
         let mut bytes = BytesMut::with_capacity(16);
         bytes.put_i32(16);
         bytes.put_i32(CANCEL_REQUEST_CODE);
@@ -290,6 +300,8 @@ impl Server {
             let code = message.get_u8() as char;
             let _len = message.get_i32();
 
+            trace!("Message: {}", code);
+
             match code {
                 // ReadyForQuery
                 'Z' => {
@@ -403,7 +415,7 @@ impl Server {
 
     /// Claim this server as mine for the purposes of query cancellation.
     pub fn claim(&mut self, process_id: i32, secret_key: i32) {
-        let mut guard = self.client_server_map.lock().unwrap();
+        let mut guard = self.client_server_map.lock();
         guard.insert(
             (process_id, secret_key),
             (
@@ -419,18 +431,9 @@ impl Server {
     /// It will use the simple query protocol.
     /// Result will not be returned, so this is useful for things like `SET` or `ROLLBACK`.
     pub async fn query(&mut self, query: &str) -> Result<(), Error> {
-        let mut query = BytesMut::from(&query.as_bytes()[..]);
-        query.put_u8(0); // C-string terminator (NULL character).
+        let query = simple_query(query);
 
-        let len = query.len() as i32 + 4;
-
-        let mut msg = BytesMut::with_capacity(len as usize + 1);
-
-        msg.put_u8(b'Q');
-        msg.put_i32(len);
-        msg.put_slice(&query[..]);
-
-        self.send(msg).await?;
+        self.send(query).await?;
 
         loop {
             let _ = self.recv().await?;

src/sharding.rs

@@ -1,26 +1,62 @@
+use sha1::{Digest, Sha1};
+
 // https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/include/catalog/partition.h#L20
 const PARTITION_HASH_SEED: u64 = 0x7A5B22367996DCFD;
 
+#[derive(Debug, PartialEq, Copy, Clone)]
+pub enum ShardingFunction {
+    PgBigintHash,
+    Sha1,
+}
+
 pub struct Sharder {
     shards: usize,
+    sharding_function: ShardingFunction,
 }
 
 impl Sharder {
-    pub fn new(shards: usize) -> Sharder {
-        Sharder { shards: shards }
+    pub fn new(shards: usize, sharding_function: ShardingFunction) -> Sharder {
+        Sharder {
+            shards,
+            sharding_function,
+        }
+    }
+
+    pub fn shard(&self, key: i64) -> usize {
+        match self.sharding_function {
+            ShardingFunction::PgBigintHash => self.pg_bigint_hash(key),
+            ShardingFunction::Sha1 => self.sha1(key),
+        }
     }
 
     /// Hash function used by Postgres to determine which partition
     /// to put the row in when using HASH(column) partitioning.
     /// Source: https://github.com/postgres/postgres/blob/27b77ecf9f4d5be211900eda54d8155ada50d696/src/common/hashfn.c#L631
     /// Supports only 1 bigint at the moment, but we can add more later.
-    pub fn pg_bigint_hash(&self, key: i64) -> usize {
+    fn pg_bigint_hash(&self, key: i64) -> usize {
         let mut lohalf = key as u32;
         let hihalf = (key >> 32) as u32;
         lohalf ^= if key >= 0 { hihalf } else { !hihalf };
         Self::combine(0, Self::pg_u32_hash(lohalf)) as usize % self.shards
     }
 
+    /// Example of a hashing function based on SHA1.
+    fn sha1(&self, key: i64) -> usize {
+        let mut hasher = Sha1::new();
+
+        hasher.update(&key.to_string().as_bytes());
+
+        let result = hasher.finalize();
+
+        // Convert the SHA1 hash into hex so we can parse it as a large integer.
+        let hex = format!("{:x}", result);
+
+        // Parse the last 8 bytes as an integer (8 bytes = bigint).
+        let key = i64::from_str_radix(&hex[hex.len() - 8..], 16).unwrap() as usize;
+
+        key % self.shards
+    }
+
     #[inline]
     fn rot(x: u32, k: u32) -> u32 {
         (x << k) | (x >> (32 - k))
@@ -109,36 +145,51 @@ mod test {
     // confirming that we implemented Postgres BIGINT hashing correctly.
     #[test]
     fn test_pg_bigint_hash() {
-        let sharder = Sharder::new(5);
+        let sharder = Sharder::new(5, ShardingFunction::PgBigintHash);
 
         let shard_0 = vec![1, 4, 5, 14, 19, 39, 40, 46, 47, 53];
 
         for v in shard_0 {
-            assert_eq!(sharder.pg_bigint_hash(v), 0);
+            assert_eq!(sharder.shard(v), 0);
         }
 
         let shard_1 = vec![2, 3, 11, 17, 21, 23, 30, 49, 51, 54];
 
         for v in shard_1 {
-            assert_eq!(sharder.pg_bigint_hash(v), 1);
+            assert_eq!(sharder.shard(v), 1);
         }
 
         let shard_2 = vec![6, 7, 15, 16, 18, 20, 25, 28, 34, 35];
 
         for v in shard_2 {
-            assert_eq!(sharder.pg_bigint_hash(v), 2);
+            assert_eq!(sharder.shard(v), 2);
         }
 
         let shard_3 = vec![8, 12, 13, 22, 29, 31, 33, 36, 41, 43];
 
         for v in shard_3 {
-            assert_eq!(sharder.pg_bigint_hash(v), 3);
+            assert_eq!(sharder.shard(v), 3);
         }
 
         let shard_4 = vec![9, 10, 24, 26, 27, 32, 37, 38, 42, 45];
 
         for v in shard_4 {
-            assert_eq!(sharder.pg_bigint_hash(v), 4);
+            assert_eq!(sharder.shard(v), 4);
+        }
+    }
+
+    #[test]
+    fn test_sha1_hash() {
+        let sharder = Sharder::new(12, ShardingFunction::Sha1);
+        let ids = vec![
+            0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+        ];
+        let shards = vec![
+            4, 7, 8, 3, 6, 0, 0, 10, 3, 11, 1, 7, 4, 4, 11, 2, 5, 0, 8, 3,
+        ];
+
+        for (i, id) in ids.iter().enumerate() {
+            assert_eq!(sharder.shard(*id), shards[i]);
         }
     }
 }

src/stats.rs

@@ -4,7 +4,6 @@ use statsd::Client;
 use tokio::sync::mpsc::{Receiver, Sender};
 
 use std::collections::HashMap;
-use std::time::Instant;
 
 use crate::config::get_config;
 
@@ -24,6 +23,7 @@ enum EventName {
     ServerTested,
     ServerLogin,
     ServerDisconnecting,
+    FlushStatsToStatsD,
 }
 
 #[derive(Debug)]
@@ -44,155 +44,167 @@ impl Reporter {
     }
 
     pub fn query(&self) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::Query,
             value: 1,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn transaction(&self) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::Transaction,
             value: 1,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn data_sent(&self, amount: usize) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::DataSent,
             value: amount as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn data_received(&self, amount: usize) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::DataReceived,
             value: amount as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn checkout_time(&self, ms: u128) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::CheckoutTime,
             value: ms as i64,
             process_id: None,
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_waiting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientWaiting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_active(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientActive,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_idle(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientIdle,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn client_disconnecting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ClientDisconnecting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_active(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerActive,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_idle(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerIdle,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_login(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerLogin,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_tested(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerTested,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
 
     pub fn server_disconnecting(&self, process_id: i32) {
-        let statistic = Event {
+        let event = Event {
             name: EventName::ServerDisconnecting,
             value: 1,
             process_id: Some(process_id),
         };
 
-        let _ = self.tx.try_send(statistic);
+        let _ = self.tx.try_send(event);
     }
+
+    // pub fn flush_to_statsd(&self) {
+    //     let event = Event {
+    //         name: EventName::FlushStatsToStatsD,
+    //         value: 0,
+    //         process_id: None,
+    //     };
+
+    //     let _ = self.tx.try_send(event);
+    // }
 }
 
 pub struct Collector {
     rx: Receiver<Event>,
+    tx: Sender<Event>,
     client: Client,
 }
 
 impl Collector {
-    pub fn new(rx: Receiver<Event>) -> Collector {
+    pub fn new(rx: Receiver<Event>, tx: Sender<Event>) -> Collector {
         Collector {
-            rx: rx,
+            rx,
+            tx,
             client: Client::new(&get_config().general.statsd_address, "pgcat").unwrap(),
         }
     }
@@ -218,8 +230,19 @@ impl Collector {
         ]);
 
         let mut client_server_states: HashMap<i32, EventName> = HashMap::new();
+        let tx = self.tx.clone();
 
-        let mut now = Instant::now();
+        tokio::task::spawn(async move {
+            let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(15000));
+            loop {
+                interval.tick().await;
+                let _ = tx.try_send(Event {
+                    name: EventName::FlushStatsToStatsD,
+                    value: 0,
+                    process_id: None,
+                });
+            }
+        });
 
         loop {
             let stat = match self.rx.recv().await {
@@ -284,65 +307,61 @@ impl Collector {
                 EventName::ClientDisconnecting | EventName::ServerDisconnecting => {
                     client_server_states.remove(&stat.process_id.unwrap());
                 }
+
+                EventName::FlushStatsToStatsD => {
+                    for (_, state) in &client_server_states {
+                        match state {
+                            EventName::ClientActive => {
+                                let counter = stats.entry("cl_active").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ClientWaiting => {
+                                let counter = stats.entry("cl_waiting").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ClientIdle => {
+                                let counter = stats.entry("cl_idle").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerIdle => {
+                                let counter = stats.entry("sv_idle").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerActive => {
+                                let counter = stats.entry("sv_active").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerTested => {
+                                let counter = stats.entry("sv_tested").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            EventName::ServerLogin => {
+                                let counter = stats.entry("sv_login").or_insert(0);
+                                *counter += 1;
+                            }
+
+                            _ => unreachable!(),
+                        };
+                    }
+
+                    info!("{:?}", stats);
+
+                    let mut pipeline = self.client.pipeline();
+
+                    for (key, value) in stats.iter_mut() {
+                        pipeline.gauge(key, *value as f64);
+                        *value = 0;
+                    }
+
+                    pipeline.send(&self.client);
+                }
             };
-
-            // It's been 15 seconds. If there is no traffic, it won't publish anything,
-            // but it also doesn't matter then.
-            if now.elapsed().as_secs() > 15 {
-                for (_, state) in &client_server_states {
-                    match state {
-                        EventName::ClientActive => {
-                            let counter = stats.entry("cl_active").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ClientWaiting => {
-                            let counter = stats.entry("cl_waiting").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ClientIdle => {
-                            let counter = stats.entry("cl_idle").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerIdle => {
-                            let counter = stats.entry("sv_idle").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerActive => {
-                            let counter = stats.entry("sv_active").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerTested => {
-                            let counter = stats.entry("sv_tested").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        EventName::ServerLogin => {
-                            let counter = stats.entry("sv_login").or_insert(0);
-                            *counter += 1;
-                        }
-
-                        _ => unreachable!(),
-                    };
-                }
-
-                info!("{:?}", stats);
-
-                let mut pipeline = self.client.pipeline();
-
-                for (key, value) in stats.iter_mut() {
-                    pipeline.gauge(key, *value as f64);
-                    *value = 0;
-                }
-
-                pipeline.send(&self.client);
-
-                now = Instant::now();
-            }
         }
     }
 }