Banning & connection pool fixes (#512)

Author: levkk

integration/failover/dev-server.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -e
+THIS_SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+source ${THIS_SCRIPT_DIR}/../toxi/setup.sh
+pushd ${THIS_SCRIPT_DIR}
+cargo run
+popd

integration/failover/pgdog.toml

@@ -0,0 +1,43 @@
+[general]
+checkout_timeout = 1_000
+connect_timeout = 1_000
+ban_timeout = 30_000
+query_timeout = 1_000
+idle_healthcheck_interval = 1_000
+client_login_timeout = 1_000
+load_balancing_algorithm = "round_robin"
+
+[[databases]]
+name = "failover"
+host = "127.0.0.1"
+port = 5435
+role = "primary"
+database_name = "pgdog"
+
+[[databases]]
+name = "failover"
+host = "127.0.0.1"
+port = 5436
+role = "replica"
+database_name = "pgdog"
+read_only = true
+
+[[databases]]
+name = "failover"
+host = "127.0.0.1"
+port = 5437
+role = "replica"
+database_name = "pgdog"
+read_only = true
+
+[[databases]]
+name = "failover"
+host = "127.0.0.1"
+port = 5438
+role = "replica"
+database_name = "pgdog"
+read_only = true
+
+[admin]
+password = "pgdog"
+user = "admin"

integration/failover/psql.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+export PGPASSWORD=pgdog
+psql -h 127.0.0.1 -p 6432 -U pgdog ${1}

integration/failover/users.toml

@@ -0,0 +1,4 @@
+[[users]]
+database = "failover"
+name = "pgdog"
+password = "pgdog"

integration/rust/tests/integration/ban.rs

@@ -70,30 +70,4 @@ async fn test_ban_unban() {
     }
 
     ensure_client_state("idle").await;
-
-    for (pool, database) in conns
-        .into_iter()
-        .zip(["pgdog", "pgdog_sharded"].into_iter())
-    {
-        for _ in 0..25 {
-            pool.execute("SELECT 1").await.unwrap();
-        }
-
-        ban_unban(database, true, false).await;
-
-        for _ in 0..25 {
-            let err = pool.execute("CREATE TABLE test (id BIGINT)").await;
-            assert!(err.err().unwrap().to_string().contains("pool is banned"));
-        }
-
-        ban_unban(database, false, false).await;
-
-        let mut t = pool.begin().await.unwrap();
-        t.execute("CREATE TABLE test_ban_unban (id BIGINT)")
-            .await
-            .unwrap();
-        t.rollback().await.unwrap();
-
-        pool.close().await;
-    }
 }

integration/schema_sync/pgdog.toml

@@ -0,0 +1,9 @@
+[[databases]]
+name = "source"
+host = "127.0.0.1"
+database_name = "pgdog"
+
+[[databases]]
+name = "destination"
+host = "127.0.0.1"
+database_name = "pgdog1"

integration/schema_sync/users.toml

@@ -0,0 +1,11 @@
+[[users]]
+name = "pgdog-4"
+database = "source"
+password = "pgdog"
+schema_admin = true
+
+[[users]]
+name = "pgdog-4"
+database = "destination"
+password = "pgdog"
+schema_admin = true

integration/setup.sh

@@ -24,7 +24,7 @@ fi
 export PGPASSWORD='pgdog'
 export PGHOST=127.0.0.1
 export PGPORT=5432
-#export PGUSER='pgdog'
+export PGUSER='pgdog'
 
 for db in pgdog shard_0 shard_1; do
     psql -c "DROP DATABASE $db" || true

integration/toxi/toxi_spec.rb

@@ -43,9 +43,8 @@ def warm_up
 
   it 'some connections survive' do
     threads = []
-    errors = 0
+    errors = Concurrent::AtomicFixnum.new(0)
     sem = Concurrent::Semaphore.new(0)
-    (5.0 / 25 * 25.0).ceil
     25.times do
       t = Thread.new do
         c = 1
@@ -54,13 +53,13 @@ def warm_up
           c = conn
           break
         rescue StandardError
-          errors += 1
+          errors.increment
         end
         25.times do
           c.exec 'SELECT 1'
         rescue PG::SystemError
           c = conn # reconnect
-          errors += 1
+          errors.increment
         end
       end
       threads << t
@@ -69,7 +68,7 @@ def warm_up
       sem.release(25)
       threads.each(&:join)
     end
-    expect(errors).to be < 25 # 5% error rate (instead of 100%)
+    expect(errors.value).to be < 25 # 5% error rate (instead of 100%)
   end
 
   it 'active record works' do
@@ -78,22 +77,93 @@ def warm_up
     # Connect (the pool is lazy)
     Sharded.where(id: 1).first
     errors = 0
+    ok = 0
     # Can't ban primary because it issues SET queries
     # that we currently route to primary.
     Toxiproxy[role].toxic(toxic).apply do
       25.times do
         Sharded.where(id: 1).first
+        ok += 1
       rescue StandardError
         errors += 1
       end
     end
-    expect(errors).to eq(1)
+    expect(errors).to be <= 1
+    expect(25 - ok).to eq(errors)
+  end
+end
+
+describe 'healthcheck' do
+  before :each do
+    admin_conn = admin
+    admin_conn.exec 'RECONNECT'
+    admin_conn.exec "SET read_write_split TO 'exclude_primary'"
+    admin_conn.exec 'SET ban_timeout TO 1'
+  end
+
+  describe 'will heal itself' do
+    def health(role, field = 'healthy')
+      admin.exec('SHOW POOLS').select do |pool|
+        pool['database'] == 'failover' && pool['role'] == role
+      end.map { |pool| pool[field] }
+    end
+
+    10.times do
+      it 'replica' do
+        # Cache connect params.
+        conn.exec 'SELECT 1'
+
+        Toxiproxy[:replica].toxic(:reset_peer).apply do
+          errors = 0
+          4.times do
+            conn.exec 'SELECT 1'
+          rescue PG::Error
+            errors += 1
+          end
+          expect(errors).to be >= 1
+          expect(health('replica')).to include('f')
+          sleep(0.4) # ban maintenance runs every 333ms
+          expect(health('replica', 'banned')).to include('t')
+        end
+
+        4.times do
+          conn.exec 'SELECT 1'
+        end
+
+        admin.exec 'HEALTHCHECK'
+        sleep(0.4)
+
+        expect(health('replica')).to eq(%w[t t t])
+        expect(health('replica', 'banned')).to eq(%w[f f f])
+      end
+    end
+
+    it 'primary' do
+      # Cache connect params.
+      conn.exec 'DELETE FROM sharded'
+
+      Toxiproxy[:primary].toxic(:reset_peer).apply do
+        begin
+          conn.exec 'DELETE FROM sharded'
+        rescue PG::Error
+        end
+        expect(health('primary')).to eq(['f'])
+      end
+
+      conn.exec 'DELETE FROM sharded'
+
+      expect(health('primary')).to eq(%w[t])
+    end
+  end
+
+  after do
+    admin.exec 'RELOAD'
   end
 end
 
 describe 'tcp' do
   around :each do |example|
-    Timeout.timeout(10) do
+    Timeout.timeout(30) do
       example.run
     end
   end
@@ -138,8 +208,23 @@ def warm_up
     end
 
     describe 'both down' do
-      it 'unbans all pools' do
-        25.times do
+      10.times do
+        it 'unbans all pools' do
+          rw_config = admin.exec('SHOW CONFIG').select do |config|
+            config['name'] == 'read_write_split'
+          end[0]['value']
+          expect(rw_config).to eq('include_primary')
+
+          def pool_stat(field, value)
+            failover = admin.exec('SHOW POOLS').select do |pool|
+              pool['database'] == 'failover'
+            end
+            entries = failover.select { |item| item[field] == value }
+            entries.size
+          end
+
+          admin.exec 'SET checkout_timeout TO 100'
+
           Toxiproxy[:primary].toxic(:reset_peer).apply do
             Toxiproxy[:replica].toxic(:reset_peer).apply do
               Toxiproxy[:replica2].toxic(:reset_peer).apply do
@@ -148,19 +233,16 @@ def warm_up
                     conn.exec_params 'SELECT $1::bigint', [1]
                   rescue StandardError
                   end
-                  banned = admin.exec('SHOW POOLS').select do |pool|
-                    pool['database'] == 'failover'
-                  end.select { |item| item['banned'] == 't' }
-                  expect(banned.size).to eq(4)
+
+                  expect(pool_stat('healthy', 'f')).to eq(4)
                 end
               end
             end
           end
-          conn.exec 'SELECT $1::bigint', [25]
-          banned = admin.exec('SHOW POOLS').select do |pool|
-            pool['database'] == 'failover'
-          end.select { |item| item['banned'] == 't' }
-          expect(banned.size).to eq(0)
+
+          4.times do
+            conn.exec 'SELECT $1::bigint', [25]
+          end
         end
       end
     end
@@ -172,25 +254,21 @@ def warm_up
       Toxiproxy[:primary].toxic(:reset_peer).apply do
         c = conn
         c.exec 'BEGIN'
-        c.exec 'CREATE TABLE test(id BIGINT)'
+        c.exec 'CREATE TABLE IF NOT EXISTS test(id BIGINT)'
         c.exec 'ROLLBACK'
       rescue StandardError
       end
+
       banned = admin.exec('SHOW POOLS').select do |pool|
         pool['database'] == 'failover' && pool['role'] == 'primary'
       end
-      expect(banned[0]['banned']).to eq('t')
+      expect(banned[0]['healthy']).to eq('f')
 
       c = conn
       c.exec 'BEGIN'
-      c.exec 'CREATE TABLE test(id BIGINT)'
+      c.exec 'CREATE TABLE IF NOT EXISTS test(id BIGINT)'
       c.exec 'SELECT * FROM test'
       c.exec 'ROLLBACK'
-
-      banned = admin.exec('SHOW POOLS').select do |pool|
-        pool['database'] == 'failover' && pool['role'] == 'primary'
-      end
-      expect(banned[0]['banned']).to eq('f')
     end
 
     it 'active record works' do
@@ -199,16 +277,58 @@ def warm_up
       # Connect (the pool is lazy)
       Sharded.where(id: 1).first
       errors = 0
+      ok = 0
       # Can't ban primary because it issues SET queries
       # that we currently route to primary.
       Toxiproxy[:primary].toxic(:reset_peer).apply do
         25.times do
           Sharded.where(id: 1).first
+          ok += 1
         rescue StandardError
           errors += 1
         end
       end
-      expect(errors).to eq(1)
+      expect(errors).to be <= 1
+      expect(25 - ok).to eq(errors)
+    end
+
+    it 'clients can connect when all servers are down after caching connection params' do
+      # First, establish a connection to cache connection parameters
+      c = conn
+      c.exec 'SELECT 1'
+      c.close
+
+      # Verify initial state - all pools should be healthy before toxics
+      pools = admin.exec('SHOW POOLS').select do |pool|
+        pool['database'] == 'failover'
+      end
+      expect(pools.all? { |p| p['healthy'] == 't' }).to be true
+
+      # Now bring down all servers
+      Toxiproxy[:primary].toxic(:reset_peer).apply do
+        Toxiproxy[:replica].toxic(:reset_peer).apply do
+          Toxiproxy[:replica2].toxic(:reset_peer).apply do
+            Toxiproxy[:replica3].toxic(:reset_peer).apply do
+              # Try to establish many connections
+              connections = []
+              50.times do
+                c = conn
+                expect(c).not_to be_nil
+                connections << c
+              end
+
+              # Check internal state - verify we have active client connections
+              clients = admin.exec('SHOW CLIENTS').select do |client|
+                client['database'] == 'failover'
+              end
+              expect(clients.size).to be >= 50
+
+              # Clean up connections without executing queries to avoid timeouts
+              connections.each { |c| c.close rescue nil }
+            end
+          end
+        end
+      end
     end
   end
 end