Patch ActiveRecord::Base.clear_query_caches_for_current_thread to restore underlying database

Otherwise it's possible for a Mysql2Adapter instance's database to be changed out from under it during an operation (e.g. insert, update, delete, truncate, etc).

CONTEXT OF THE ISSUE

This issue presented itself with a combination of using active_record_shards and active_record_host_pool. It was representing the sharded and unsharded database connections within the same connection pool. This problem did not manifest itself until Rails 6 when cache clearing for cache dirtying operations—e.g. insert, update, truncate, delete, etc—was updated to loop over the app's connection handlers, connection pools to search for currently active connections.

This happens mid-operation (inside of ActiveRecord::ConnectionAdapters::Mysql2Adapter instance) just before the operation is carried out. The result was that referencing the connection was changing the currently selected database and it was not being restored. This led to issues where the operation would try to be carried out on a different database that did not have the same tables/columns.

It is important to note that this issue existed in active_record_host_pool for years, but no code paths exercised it. The assumption that connections would not be re-referenced from within an operation on a connection was safe for years until ActiveRecord 6.

WHY THIS FIX

The active_record_host_pool library is optimized to reduce the number of database connections by pooling them by host/port/socket/user/slave. This means that different databases on the same host will use the same connection pool. This requires that this library provide a mechanism for selecting the correct database when executing queries.

This fix maintains that optimization and fixes the currently known problematic code path introduced in ActiveRecord 6.

BENEFITS OF THIS CHANGE

Preserves host pool optimization.

CONSEQUENCES OF THIS CHANGE

This does not fix the issue, it just fixes the one spot where we know ActiveRecord 6.0 can expose it.

REFERENCES

The problematic line that causes the database to switch is: https://github.com/zendesk/active_record_host_pool/blob/981fdeb8dfa52f390bd91471adc19049283c6fb8/lib/active_record_host_pool/connection_proxy.rb#L16

The change in Rails 6 that exercised a code path to expose this surface the issue: https://github.com/rails/rails/pull/35089/files#diff-193e6e4a86bf5d360a5001bce59334cbR180-R186

For an alternate solution see: #60

Author: zdennis

lib/active_record_host_pool.rb

@@ -4,6 +4,7 @@
 require 'active_record/base'
 require 'active_record/connection_adapters/abstract_adapter'
 
+require 'active_record_host_pool/connection_handling'
 require 'active_record_host_pool/connection_proxy'
 require 'active_record_host_pool/pool_proxy'
 require 'active_record_host_pool/connection_adapter_mixin'

lib/active_record_host_pool/connection_handling.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module ActiveRecordHostPool
+  # ActiveRecord 6.0 introduces multiple database support. This introduces a
+  # change where dirtying operations on a Mysql2Adapter instance calls out to
+  # ActiveRecord::Base.clear_query_caches_for_current_thread. This will
+  # for any active connections whose query caches need to be cleared. This
+  # involves iterating over every connection handler's connection pool
+  # looking for active active connections.
+  #
+  # This exposes an issue with active active_record_host_pool since by
+  # referencing a connection it causes the underlying Mysql2Adapter instance
+  # to potentially change its _host_pool_current_database. This means
+  # mid-way into carrying out an operation like inserting a record the
+  # act of clearing the query cache may select a new database and leave it
+  # changed. This can result in the insert SQL then being sent to a new
+  # database.
+  #
+  # This module wraps ActiveRecord::Base.clear_query_caches_for_current_thread
+  # in order to restore the current connection's database to what it was
+  # before the caches were attempted to be cleared.
+  #
+  # Before ActiveRecord 6 this potential issue existed in the
+  # active_record_host_pool library but no code paths exercised it. This
+  # module does not resolve the general issue and the underlying issue still
+  # exists, but currently clearing the cache is the only known code path that
+  # exercises this problem so that is all that is patched.
+  module ResetActiveDatabaseAfterClearingCache
+    def clear_query_caches_for_current_thread
+      host_pool_current_database_was = connection.unproxied._host_pool_current_database
+      super
+    ensure
+      # restore in case clearing the cache changed the database
+      connection.unproxied._host_pool_current_database = host_pool_current_database_was
+    end
+  end
+end
+
+ActiveRecord::Base.singleton_class.prepend ActiveRecordHostPool::ResetActiveDatabaseAfterClearingCache

test/database.yml

@@ -66,3 +66,12 @@ test_host_1_db_not_there:
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   reconnect: true
+
+test_host_1_db_shard:
+  adapter: mysql2
+  encoding: utf8
+  database: arhp_test_1_shard
+  username: <%= mysql.user %>
+  password: <%= mysql.password %>
+  host: <%= mysql.host %>
+  reconnect: true

test/helper.rb

@@ -25,6 +25,13 @@ def arhp_create_models
     return if ARHPTestSetup.const_defined?('Test1')
 
     eval <<-RUBY
+      # The placement of the Test1Shard class is important so that its
+      # connection will not be the most recent connection established
+      # for test_host_1.
+      class Test1Shard < ::ActiveRecord::Base
+        establish_connection(:test_host_1_db_shard)
+      end
+
       class Test1 < ActiveRecord::Base
         self.table_name = "tests"
         establish_connection(:test_host_1_db_1)

test/schema.rb

@@ -2,7 +2,15 @@
 
 require_relative 'helper'
 ActiveRecord::Schema.define(version: 1) do
-  create_table 'tests' do |t|
+  create_table 'tests', force: true do |t|
     t.string   'val'
   end
+
+  # Add a table only the shard database will have. Conditional
+  # exists since Phenix loads the schema file for every database.
+  if ActiveRecord::Base.connection.current_database == 'arhp_test_1_shard'
+    create_table 'test1_shards' do |t|
+      t.string 'name'
+    end
+  end
 end

test/test_arhp.rb

@@ -19,7 +19,7 @@ def test_process_forking_with_connections
 
     # Verify that when we fork, the process doesn't crash
     pid = Process.fork do
-      if ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 2
+      if ActiveRecord.version >= Gem::Version.new('5.2')
         assert_equal(false, ActiveRecord::Base.connected?) # New to Rails 5.2
       else
         assert_equal(true, ActiveRecord::Base.connected?)
@@ -55,6 +55,22 @@ def test_should_insert_on_correct_database
     assert_action_uses_correct_database(:insert, "insert into tests values(NULL, 'foo')")
   end
 
+  def test_models_with_matching_hosts_and_non_matching_databases_should_share_a_connection
+    simulate_rails_app_active_record_railties
+    assert_equal(Test1.connection.raw_connection, Test1Shard.connection.raw_connection)
+  end
+
+  def test_models_with_matching_hosts_and_non_matching_databases_do_not_mix_up_underlying_database
+    simulate_rails_app_active_record_railties
+
+    # ActiveRecord 6.0 introduced a change that surfaced a problematic code
+    # path in active_record_host_pool: Clearing the cache across connection
+    # handlers can cause the shared host pool to change the underlying
+    # database out from under an in-progress operation.
+    # See lib/active_record_host_pool/connection_handling.rb for more
+    ActiveRecord::Base.cache { Test1Shard.create! }
+  end
+
   def test_connection_returns_a_proxy
     assert_kind_of ActiveRecordHostPool::ConnectionProxy, Test1.connection
   end
@@ -180,4 +196,16 @@ def assert_action_uses_correct_database(action, sql)
       assert_equal desired_db, current_database(klass)
     end
   end
+
+  def simulate_rails_app_active_record_railties
+    if ActiveRecord.version >= Gem::Version.new('6.0')
+      # Necessary for testing ActiveRecord 6.0 which uses the connection
+      # handlers when clearing query caches across all handlers when
+      # an operation that dirties the cache isinvolved (e.g. create/insert,
+      # update, delete/destroy, truncate, etc.)
+      ActiveRecord::Base.connection_handlers = {
+        ActiveRecord::Base.writing_role => ActiveRecord::Base.default_connection_handler
+      }
+    end
+  end
 end