Close servers that havent finished reading/writing (#566)

* Close servers that havent finished reading/writing

* fix read cancellation safety

* Test

* More tests

* fix test

* remove the bad test

* reset bans

* Refactor cleanup to be cleaner

* Fix sharded copy to stdout

* preserve health check failure on io error

* Remove

* remove unrecheable code

* more confusing code

* more dead code

Author: levkk

integration/python/test_asyncpg.py

@@ -5,6 +5,7 @@
 import random
 import string
 import pytest_asyncio
+from io import BytesIO
 
 
 @pytest_asyncio.fixture
@@ -231,13 +232,24 @@ async def test_copy(conns):
     records = 250
     for i in range(50):
         for conn in conns:
+            # Test COPY FROM (TO table)
             rows = [[x, f"value_{x}", datetime.now()] for x in range(records)]
             await conn.copy_records_to_table(
                 "sharded", records=rows, columns=["id", "value", "created_at"]
             )
             count = await conn.fetch("SELECT COUNT(*) FROM sharded")
             assert len(count) == 1
             assert count[0][0] == records
+
+            # Test COPY TO STDOUT
+            buffer = BytesIO()
+            copied_data = await conn.copy_from_table(
+                "sharded", columns=["id", "value", "created_at"], output=buffer
+            )
+            buffer.seek(0)
+            lines = buffer.read().decode('utf-8').strip().split('\n')
+            assert len(lines) == records, f"expected {records} lines in COPY output, got {len(lines)}"
+
             await conn.execute("DELETE FROM sharded")
 
 

integration/ruby/ar_spec.rb

@@ -2,6 +2,7 @@
 
 require_relative 'rspec_helper'
 require 'pp'
+require 'timeout'
 
 class Sharded < ActiveRecord::Base
   self.table_name = 'sharded'
@@ -187,4 +188,129 @@ def conn(db, prepared)
       end
     end
   end
+
+  describe 'chaos testing with interrupted queries' do
+    before do
+      conn('failover', false)
+      ActiveRecord::Base.connection.execute 'DROP TABLE IF EXISTS sharded'
+      ActiveRecord::Base.connection.execute 'CREATE TABLE sharded (id BIGSERIAL PRIMARY KEY, value TEXT)'
+    end
+
+    it 'handles interrupted queries and continues operating normally' do
+      interrupted_count = 0
+      successful_count = 0
+      mutex = Mutex.new
+
+      # Apply latency toxic to slow down query transmission,
+      # making it easier to interrupt queries mid-flight
+      Toxiproxy[:primary].toxic(:latency, latency: 100, jitter: 50).apply do
+        # Phase 1: Chaos - interrupt queries randomly with thread kills
+        chaos_threads = []
+        killer_threads = []
+
+        # Start 10 query threads
+        10.times do |thread_id|
+          t = Thread.new do
+            100.times do |i|
+              begin
+                case rand(3)
+                when 0
+                  # SELECT query
+                  Sharded.where('id > ?', 0).limit(10).to_a
+                when 1
+                  # INSERT query
+                  Sharded.create value: "thread_#{thread_id}_iter_#{i}"
+                when 2
+                  # Transaction with multiple operations
+                  Sharded.transaction do
+                    rec = Sharded.create value: "tx_#{thread_id}_#{i}"
+                    Sharded.where(id: rec.id).first if rec.id
+                  end
+                end
+                mutex.synchronize { successful_count += 1 }
+              rescue StandardError => e
+                # Killed mid-query or other error
+                mutex.synchronize { interrupted_count += 1 }
+              end
+            end
+          end
+          chaos_threads << t
+        end
+
+        # Start killer thread that randomly kills query threads
+        killer = Thread.new do
+          50.times do
+            sleep(rand(0.01..0.05))
+            alive_threads = chaos_threads.select(&:alive?)
+            if alive_threads.any?
+              victim = alive_threads.sample
+              victim.kill
+              mutex.synchronize { interrupted_count += 1 }
+            end
+          end
+        end
+        killer_threads << killer
+
+        # Wait for killer to finish
+        killer_threads.each(&:join)
+
+        # Wait for remaining threads (with timeout)
+        chaos_threads.each { |t| t.join(0.1) }
+
+        puts "Chaos phase complete: #{successful_count} successful, #{interrupted_count} interrupted"
+        expect(interrupted_count).to be > 0
+      end # End toxiproxy latency
+
+      # Give PgDog time to clean up broken connections
+      sleep(0.5)
+
+      # Disconnect all connections to clear bad state
+      ActiveRecord::Base.connection_pool.disconnect!
+
+      # Wait a bit more for cleanup
+      sleep(0.5)
+
+      # Phase 2: Verify database continues to operate normally
+      verification_errors = []
+      errors_mutex = Mutex.new
+
+      verification_threads = 10.times.map do |thread_id|
+        Thread.new do
+          20.times do |i|
+            begin
+              # Simple queries that don't depend on finding specific records
+              # INSERT
+              rec = Sharded.create value: "verify_#{thread_id}_#{i}"
+              expect(rec.id).to be > 0
+
+              # SELECT with basic query
+              results = Sharded.where('value LIKE ?', 'verify_%').limit(5).to_a
+              expect(results).to be_a(Array)
+
+              # COUNT query
+              count = Sharded.where('id > ?', 0).count
+              expect(count).to be >= 0
+            rescue PG::Error => e
+              # PG errors should fail the test
+              raise
+            rescue StandardError => e
+              errors_mutex.synchronize { verification_errors << e }
+            end
+          end
+        end
+      end
+
+      verification_threads.each(&:join)
+
+      # Verify no errors occurred during verification
+      expect(verification_errors).to be_empty, "Verification errors: #{verification_errors.map(&:message).join(', ')}"
+
+      # Verify we can still execute basic queries
+      ActiveRecord::Base.connection.execute('SELECT 1')
+
+      # Verify count works
+      count = Sharded.count
+      expect(count).to be >= 0
+    end
+  end
 end

integration/ruby/lb_spec.rb

@@ -12,6 +12,8 @@
   describe 'random' do
     it 'distributes traffic evenly' do
       conn = failover
+      # Reset stats and bans
+      admin.exec "RECONNECT"
 
       before = admin_stats('failover')
       250.times do

pgdog/src/admin/show_pools.rs

@@ -40,6 +40,7 @@ impl Command for ShowPools {
             Field::numeric("errors"),
             Field::numeric("re_synced"),
             Field::numeric("out_of_sync"),
+            Field::numeric("force_closed"),
             Field::bool("online"),
             Field::text("replica_lag"),
             Field::bool("schema_admin"),
@@ -73,6 +74,7 @@ impl Command for ShowPools {
                         .add(state.errors)
                         .add(state.re_synced)
                         .add(state.out_of_sync)
+                        .add(state.force_close)
                         .add(state.online)
                         .add(state.replica_lag.simple_display())
                         .add(cluster.schema_admin());

pgdog/src/admin/tests/mod.rs

@@ -114,6 +114,7 @@ async fn show_pools_reports_schema_admin_flag() {
         "errors",
         "re_synced",
         "out_of_sync",
+        "force_closed",
         "online",
         "replica_lag",
         "schema_admin",

pgdog/src/backend/error.rs

@@ -84,6 +84,9 @@ pub enum Error {
     #[error("protocol is out of sync")]
     ProtocolOutOfSync,
 
+    #[error("rollback left server in inconsistent state")]
+    RollbackFailed,
+
     #[error("decoder is missing required data to decode row")]
     DecoderRowError,
 

pgdog/src/backend/pool/connection/binding.rs

@@ -97,6 +97,7 @@ impl Binding {
                             }
 
                             let message = server.read().await?;
+
                             read = true;
                             if let Some(message) = state.forward(message)? {
                                 return Ok(message);

pgdog/src/backend/pool/connection/mirror/mod.rs

@@ -54,7 +54,7 @@ impl Mirror {
             prepared_statements: PreparedStatements::new(),
             params: params.clone(),
             timeouts: Timeouts::from_config(&config.config.general),
-            stream: Stream::DevNull,
+            stream: Stream::dev_null(),
             transaction: None,
             cross_shard_disabled: config.config.general.cross_shard_disabled,
         }

pgdog/src/backend/pool/connection/mod.rs

@@ -229,12 +229,14 @@ impl Connection {
     /// Only await this future inside a `select!`. One of the conditions
     /// suspends this loop indefinitely and expects another `select!` branch
     /// to cancel it.
+    ///
     pub(crate) async fn read(&mut self) -> Result<Message, Error> {
         select! {
             notification = self.pub_sub.recv() => {
                 Ok(notification.ok_or(Error::ProtocolOutOfSync)?.message()?)
             }
 
+            // BUG: This is not cancellation-safe.
             message = self.binding.read() => {
                 message
             }

pgdog/src/backend/pool/connection/multi_shard/mod.rs

@@ -39,6 +39,9 @@ struct Counters {
     bind_complete: usize,
     command_complete: Option<Message>,
     transaction_error: bool,
+    copy_done: usize,
+    copy_out: usize,
+    copy_data: usize,
 }
 
 /// Multi-shard state.
@@ -246,6 +249,25 @@ impl MultiShard {
                 }
             }
 
+            'c' => {
+                self.counters.copy_done += 1;
+                if self.counters.copy_done.is_multiple_of(self.shards) {
+                    forward = Some(message);
+                }
+            }
+
+            'd' => {
+                self.counters.copy_data += 1;
+                forward = Some(message);
+            }
+
+            'H' => {
+                self.counters.copy_out += 1;
+                if self.counters.copy_out.is_multiple_of(self.shards) {
+                    forward = Some(message);
+                }
+            }
+
             't' => {
                 self.counters.parameter_description += 1;
                 if self