PgCat Query Mirroring (#341)

This is an implementation of Query mirroring in PgCat (outlined here #302)

In configs, we match mirror hosts with the servers handling the traffic. A mirror host will receive the same protocol messages as the main server it was matched with.

This is done by creating an async task for each mirror server, it communicates with the main server through two channels, one for the protocol messages and one for the exit signal. The mirror server sends the protocol packets to the underlying PostgreSQL server. We receive from the underlying PostgreSQL server as soon as the data is available and we immediately discard it. We use bb8 to manage the life cycle of the connection, not for pooling since each mirror server handler is more or less single-threaded.

We don't have any connection pooling in the mirrors. Matching each mirror connection to an actual server connection guarantees that we will not have more connections to any of the mirrors than the parent pool would allow.

tests/docker/docker-compose.yml

@@ -8,7 +8,7 @@ services:
       POSTGRES_DB: postgres
       POSTGRES_PASSWORD: postgres
       POSTGRES_INITDB_ARGS: --auth-local=md5 --auth-host=md5 --auth=md5
-    command: ["postgres", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "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:
     image: postgres:14
     network_mode: "service:main"
@@ -17,7 +17,7 @@ services:
       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", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "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:
     image: postgres:14
     network_mode: "service:main"
@@ -26,7 +26,7 @@ services:
       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", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "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:
     image: postgres:14
     network_mode: "service:main"
@@ -35,7 +35,7 @@ services:
       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", "-c", "shared_preload_libraries=pg_stat_statements", "-c", "pg_stat_statements.track=all", "-p", "9432"]
+    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"]

tests/ruby/helpers/pg_instance.rb

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

tests/ruby/helpers/pgcat_process.rb

@@ -29,7 +29,7 @@ class PgcatProcess
                    else
                      '../../target/debug/pgcat'
                    end
-                     
+
     @command = "#{command_path} #{@config_filename}"
 
     FileUtils.cp("../../pgcat.toml", @config_filename)
@@ -48,12 +48,14 @@ class PgcatProcess
     @original_config = current_config
     output_to_write = TOML::Generator.new(config_hash).body
     output_to_write = output_to_write.gsub(/,\s*["|'](\d+)["|']\s*,/, ',\1,')
+    output_to_write = output_to_write.gsub(/,\s*["|'](\d+)["|']\s*\]/, ',\1]')
     File.write(@config_filename, output_to_write)
   end
 
   def current_config
-    old_cfg = File.read(@config_filename)
-    loadable_string = old_cfg.gsub(/,\s*(\d+)\s*,/, ', "\1",')
+    loadable_string = File.read(@config_filename)
+    loadable_string = loadable_string.gsub(/,\s*(\d+)\s*,/,  ', "\1",')
+    loadable_string = loadable_string.gsub(/,\s*(\d+)\s*\]/, ', "\1"]')
     TOML.load(loadable_string)
   end
 

tests/ruby/load_balancing_spec.rb

@@ -46,7 +46,6 @@ describe "Random Load Balancing" do
         end
       end
 
-      expect(failed_count).to be <= 2
       processes.all_databases.each do |instance|
         queries_routed = instance.count_select_1_plus_2
         if processes.replicas[0..1].include?(instance)

tests/ruby/mirrors_spec.rb

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