diff --git a/lib/active_record_proxy_adapters/primary_replica_proxy.rb b/lib/active_record_proxy_adapters/primary_replica_proxy.rb index 086376fc..381167c9 100644 --- a/lib/active_record_proxy_adapters/primary_replica_proxy.rb +++ b/lib/active_record_proxy_adapters/primary_replica_proxy.rb @@ -173,15 +173,6 @@ def connection_for(role, sql_string) update_primary_latest_write_timestamp if !replica_connection?(connection) && write_statement?(sql_string) result - ensure - # Check the connection back into its own pool, not whatever pool currently - # answers to the :reading role. In Rails system tests, - # ActiveRecord::TestFixtures#setup_shared_connection_pool reassigns the - # :reading role's pool_config on every transactional-fixture test's - # before_setup, so re-resolving `replica_pool` here can return a pool - # different from the one the connection was checked out of — which would - # check a replica adapter into the primary pool and poison it. - connection.pool.checkin(connection) if replica_connection?(connection) end def connected_to(role:, &block) @@ -195,7 +186,7 @@ def replica_connection?(connection) end def checkout_replica_connection - replica_pool.checkout(proxy_checkout_timeout) + replica_pool.lease_connection # rescue NoDatabaseError to avoid crashing when running db:create rake task # rescue ConnectionNotEstablished to handle connectivity issues in the replica # (for example, replication delay) diff --git a/spec/active_record_proxy_adapters/primary_replica_proxy_spec.rb b/spec/active_record_proxy_adapters/primary_replica_proxy_spec.rb deleted file mode 100644 index 12a0abac..00000000 --- a/spec/active_record_proxy_adapters/primary_replica_proxy_spec.rb +++ /dev/null @@ -1,71 +0,0 @@ -# frozen_string_literal: true - -RSpec.describe ActiveRecordProxyAdapters::PrimaryReplicaProxy do - describe "#connection_for" do - # In Rails' system tests, ActiveRecord::TestFixtures#setup_shared_connection_pool - # reassigns the :reading role's pool_config to the :writing pool_config on every - # transactional-fixture test's before_setup. Because PrimaryReplicaProxy resolves - # `replica_pool` on every call via retrieve_connection_pool(..., role: :reading), - # a swap that occurs between the initial checkout (in #checkout_replica_connection) - # and the ensure-block checkin causes the replica connection to be checked into - # the wrong (primary) pool, where a later ConnectionPool#pin_connection! picks it - # up as the fixture-transaction connection and every subsequent write on the - # primary class raises ActiveRecord::ReadOnlyError. - # - # This spec exercises that race deterministically: we stub `replica_pool` on the - # proxy to return the real replica pool on the checkout call and the primary pool - # on the ensure-block call, mirroring what the handler does when - # setup_shared_connection_pool runs mid-SELECT. After the query completes, the - # connection must still land in its original (replica) pool, not the primary pool. - context "when the :reading pool_config is reassigned between checkout and the ensure-block checkin" do - subject(:run_query) { proxy.execute(sql) } - - let(:primary_pool) { TestHelper.sqlite3_primary_pool } - let(:replica_pool) { TestHelper.sqlite3_replica_pool } - let(:proxy) { ActiveRecordProxyAdapters::SQLite3Proxy.new(primary_adapter) } - let(:sql) { "SELECT 1" } - - around do |example| - primary_pool.with_connection do |connection| - @primary_adapter = connection - example.run - @primary_adapter = nil - end - end - - attr_reader :primary_adapter - - def primary_available_connections - primary_pool.instance_variable_get(:@available).instance_variable_get(:@queue) - end - - before do - # A single `connection_for` resolves `replica_pool` twice on the - # checkout side: - # 1. `replica_pool_unavailable?` — to decide whether to use a replica - # 2. `checkout_replica_connection` — for `replica_pool.checkout` - # Any subsequent resolution would be the ensure-block re-resolving the - # pool to check the connection back in — which is exactly the call that - # `setup_shared_connection_pool` corrupts when it swaps the :reading - # pool_config between checkout and checkin. We reproduce that swap by - # returning the real replica pool for the two checkout-side calls and - # the primary pool for anything after. - checkout_side_calls = 2 - call_count = 0 - allow(proxy).to receive(:replica_pool).and_wrap_original do |original, *args, **kwargs| - call_count += 1 - after_checkout = call_count > checkout_side_calls - after_checkout ? primary_pool : original.call(*args, **kwargs) - end - end - - it "does not leave the replica connection in the primary pool" do - run_query - - leaked = primary_available_connections.find { |c| c.respond_to?(:replica?) && c.replica? } - - expect(leaked).to be_nil, "replica adapter leaked into the primary pool (#{leaked.inspect})" - end - end - end -end diff --git a/spec/shared_examples/a_sql_statement.rb b/spec/shared_examples/a_sql_statement.rb index 775ea846..f8fff470 100644 --- a/spec/shared_examples/a_sql_statement.rb +++ b/spec/shared_examples/a_sql_statement.rb @@ -3,22 +3,21 @@ require "shared_contexts/a_proxied_method_setup" RSpec.shared_examples_for "a SQL read statement" do - it "checks out a connection from the replica pool" do - allow(replica_pool).to receive(:checkout).and_call_original + it "leases a connection from the replica pool" do + allow(replica_pool).to receive(:lease_connection).and_call_original run_test - expect(replica_pool).to have_received(:checkout).once + expect(replica_pool).to have_received(:lease_connection).once end - it "checks replica connection back in to the pool" do + it "uses the leased replica connection" do conn = instance_double(adapter_class, method_name => nil, pool: replica_pool) - allow(replica_pool).to receive(:checkout).and_return(conn) - allow(replica_pool).to receive(:checkin) + allow(replica_pool).to receive(:lease_connection).and_return(conn) run_test - expect(replica_pool).to have_received(:checkin).with(conn).once + expect(conn).to have_received(method_name).once end context "when a transaction is open" do @@ -31,11 +30,11 @@ end it "does not checkout a connection from the replica pool" do - allow(replica_pool).to receive(:checkout).and_call_original + allow(replica_pool).to receive(:lease_connection).and_call_original primary_adapter.transaction { run_test } - expect(replica_pool).not_to have_received(:checkout) + expect(replica_pool).not_to have_received(:lease_connection) end end @@ -49,22 +48,22 @@ end it "does not checkout a connection from the replica pool" do - allow(replica_pool).to receive(:checkout).and_call_original + allow(replica_pool).to receive(:lease_connection).and_call_original model_class.connected_to(role: TestHelper.writing_role) { run_test } - expect(replica_pool).not_to have_received(:checkout) + expect(replica_pool).not_to have_received(:lease_connection) end end end RSpec.shared_examples_for "a SQL write statement" do it "does not checkout a connection from replica pool" do - allow(replica_pool).to receive(:checkout).and_call_original + allow(replica_pool).to receive(:lease_connection).and_call_original run_test - expect(replica_pool).not_to have_received(:checkout) + expect(replica_pool).not_to have_received(:lease_connection) end it "sends query to primary connection" do