Add an intermediary called RoleManager to manage connections

This PR is an alternate solution to #37388. While there are benefits
to merging #37388 it changes the public API and swaps around existing
concepts for how connection management works. The changes are
backwards-incompatible and pretty major. This will have a negative impact
on gems and applications relying on how conn management currently works.

**Background:**

Shopify and other applications need sharding but Rails has
made it impossible to do this because a handler can only hold one
connection pool per class. Sharded apps need to hold multiple
connections per handler per class.

This PR aims to solve only that problem.

**What this PR does:**

In this PR we've added a `RoleManager` class that can hold multiple
`Roles`. Each `Role` holds the `db_config`,
`connection_specification_name`, `schema_cache` and `pool`. By default
the `RoleManager` holds a single reference from a `default` key to the
`Role` instance. A sharded/multi-tenant app can pass an optional second
argument to `remove_connection`, `retrieve_connection_pool`,
`establish_connection` and `connected?` on the handler, thus allowing
for multiple connections belonging to the same class/handler without
breaking backwards compatibility.

By using the `RoleManager` we can avoid altering the public API, moving
around handler/role concepts, and achieve the internal needs for
establishing multiple connections per handler per class.

**A note about why we opened this PR:**

We very much appreciate the work that went into #37388 and in no way mean
to diminish that work. However, it breaks the following public APIs:

* `#retrieve_connection`, `#connected?`, and `#remove_connection` are
public methods on handler and can't be changed from taking a spec to a
role.
* The knowledge that the handler keys are symbols relating to a role
(`:writing`/`:reading`) is public - changing how handlers are accessed
will break apps/libraries.

In addition it doesn't solve the problem of mapping a single connection
to a single class since it has a 1:1 mapping of `class (handler) -> role
(writing) -> db_config`. Multiple pools in a writing role can't exist
in that implementation.

The new PR solves this by using the `RoleManager` to hold multiple connection
objects for the same class. This lets a handler hold a role manager
which can hold as many roles for that writer as the app needs.

**Regarding the `Role` name:**

When I originally designed the API for multiple databases, it wasn't
accidental that handler and role are the same concept. Handler is the
internal concept (since that's what was there already) and Role was the
public external concept. Meaning, role and handler were meant to
be the same thing. The concept here means that when you switch a
handler/role, Rails automatically can pick up the connection on the
other role by knowing the specification name. Changing this would mean not
just that we need to rework how GitHub and many many gems work, but also
means retraining users of Rails 6.0 that all these concepts changed.

Since this PR doesn't move around the concepts in connection
management and instead creates an intermediary between `handler` and
`role` to manage the connection data (`db_config`, `schema_cache`,
`pool`, and `connection_specification`) we think that `Role` and
`RoleManager` are the wrong name.

We didn't change it yet in this PR because we wanted to keep change
churn low for initial review. We also haven't come up with a better
name yet. 😄

**What this PR does not solve:**

Our PR here solves a small portion of the problem - it allows models to
have multiple connections on a class. It doesn't aim to solve any other
problems than that. Going forward we'll need to still solve the
following problems:

* `DatabaseConfig` doesn't support a sharding configuration
* `connects_to`/`connected_to` still needs a way to switch connections for shards
* Automatic switching of shards
* `connection_specification_name` still exists

**The End**

Thanks for reading this far. These problems aren't easy to solve. John
and I spent a lot of time trying different things and so I hope that
this doesn't come across as if we think we know better. I would have
commented on the other PR what changes to make but we needed to try out
different solutions in order to get here.

Ultimately we're aiming to change as little as the API as possible. Even
if the handler/role -> manager -> db_config/pool/etc isn't how we'd
design connection management if we could start over, we also don't want
to break public APIs. It's important that we make things better while
maintaining compatibility.

The `RoleManager` class makes it possible for us to fix the underlying
problem while maintaining all the backwards compatibility in the public
API.

We all have the same goal; to add sharding support to Rails. Let me know
your thoughts on this change in lieu of #37388 and if you have questions.

Co-authored-by: John Crepezzi <seejohnrun@github.com>

activerecord/lib/active_record/connection_adapters.rb

@@ -10,6 +10,7 @@ module ActiveRecord
 
     autoload :Column
     autoload :Role
+    autoload :RoleManager
     autoload :Resolver
 
     autoload_at "active_record/connection_adapters/abstract/schema_definitions" do

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -1004,7 +1004,7 @@ module ActiveRecord
 
       def initialize
         # These caches are keyed by role.connection_specification_name (Role#connection_specification_name).
-        @owner_to_role = Concurrent::Map.new(initial_capacity: 2)
+        @owner_to_role_manager = Concurrent::Map.new(initial_capacity: 2)
 
         # Backup finalizer: if the forked child skipped Kernel#fork the early discard has not occurred
         ObjectSpace.define_finalizer self, FINALIZER
@@ -1031,20 +1031,20 @@ module ActiveRecord
       end
 
       def connection_pool_names # :nodoc:
-        owner_to_role.keys
+        owner_to_role_manager.keys
       end
 
       def connection_pool_list
-        owner_to_role.values.compact.map(&:pool)
+        owner_to_role_manager.values.compact.flat_map { |rc| rc.roles.map(&:pool) }
       end
       alias :connection_pools :connection_pool_list
 
-      def establish_connection(config)
+      def establish_connection(config, role_name = :default)
         resolver = Resolver.new(Base.configurations)
         role = resolver.resolve_role(config)
         db_config = role.db_config
 
-        remove_connection(role.connection_specification_name)
+        remove_connection(role.connection_specification_name, role_name)
 
         message_bus = ActiveSupport::Notifications.instrumenter
         payload = {
@@ -1055,7 +1055,9 @@ module ActiveRecord
           payload[:config] = db_config.configuration_hash
         end
 
-        owner_to_role[role.connection_specification_name] = role
+        owner_to_role_manager[role.connection_specification_name] ||= RoleManager.new
+        role_manager = owner_to_role_manager[role.connection_specification_name]
+        role_manager.set_role(role_name, role)
 
         message_bus.instrument("!connection.active_record", payload) do
           role.pool
@@ -1114,8 +1116,8 @@ module ActiveRecord
 
       # Returns true if a connection that's accessible to this class has
       # already been opened.
-      def connected?(spec_name)
-        pool = retrieve_connection_pool(spec_name)
+      def connected?(spec_name, role_name = :default)
+        pool = retrieve_connection_pool(spec_name, role_name)
         pool && pool.connected?
       end
 
@@ -1123,22 +1125,27 @@ module ActiveRecord
       # connection and the defined connection (if they exist). The result
       # can be used as an argument for #establish_connection, for easily
       # re-establishing the connection.
-      def remove_connection(spec_name)
-        if role = owner_to_role.delete(spec_name)
-          role.disconnect!
-          role.db_config.configuration_hash
+      def remove_connection(spec_name, role_name = :default)
+        if role_manager = owner_to_role_manager[spec_name]
+          role = role_manager.remove_role(role_name)
+
+          if role
+            role.disconnect!
+            role.db_config.configuration_hash
+          end
         end
       end
 
-      # Retrieving the connection pool happens a lot, so we cache it in @owner_to_role.
+      # Retrieving the connection pool happens a lot, so we cache it in @owner_to_role_manager.
       # This makes retrieving the connection pool O(1) once the process is warm.
       # When a connection is established or removed, we invalidate the cache.
-      def retrieve_connection_pool(spec_name)
-        owner_to_role[spec_name]&.pool
+      def retrieve_connection_pool(spec_name, role_name = :default)
+        role = owner_to_role_manager[spec_name]&.get_role(role_name)
+        role&.pool
       end
 
       private
-        attr_reader :owner_to_role
+        attr_reader :owner_to_role_manager
     end
   end
 end

activerecord/lib/active_record/connection_adapters/role_manager.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    class RoleManager # :nodoc:
+      def initialize
+        @name_to_role = {}
+      end
+
+      def roles
+        @name_to_role.values
+      end
+
+      def remove_role(name)
+        @name_to_role.delete(name)
+      end
+
+      def get_role(name)
+        @name_to_role[name]
+      end
+
+      def set_role(name, role)
+        @name_to_role[name] = role
+      end
+    end
+  end
+end

activerecord/lib/active_record/test_fixtures.rb

@@ -192,7 +192,11 @@ module ActiveRecord
         ActiveRecord::Base.connection_handlers.values.each do |handler|
           if handler != writing_handler
             handler.connection_pool_names.each do |name|
-              handler.send(:owner_to_role)[name] = writing_handler.send(:owner_to_role)[name]
+              writing_role_manager = writing_handler.send(:owner_to_role_manager)[name]
+              writing_role = writing_role_manager.get_role(:default)
+
+              role_manager = handler.send(:owner_to_role_manager)[name]
+              role_manager.set_role(:default, writing_role)
             end
           end
         end

activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/person"
+
+module ActiveRecord
+  module ConnectionAdapters
+    class ConnectionHandlersMultiRoleTest < ActiveRecord::TestCase
+      self.use_transactional_tests = false
+
+      fixtures :people
+
+      def setup
+        @handlers = { writing: ConnectionHandler.new }
+        @rw_handler = @handlers[:writing]
+        @spec_name = "primary"
+        @writing_handler = ActiveRecord::Base.connection_handlers[:writing]
+      end
+
+      def teardown
+        ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler }
+      end
+
+      unless in_memory_db?
+        def test_establish_connection_with_roles
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          @writing_handler.establish_connection(:primary)
+          @writing_handler.establish_connection(:primary, :role_two)
+
+          default_pool = @writing_handler.retrieve_connection_pool("primary", :default)
+          other_pool = @writing_handler.retrieve_connection_pool("primary", :role_two)
+
+          assert_not_nil default_pool
+          assert_not_equal default_pool, other_pool
+
+          # :default if passed with no key
+          assert_equal default_pool, @writing_handler.retrieve_connection_pool("primary")
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
+
+        def test_remove_connection
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          @writing_handler.establish_connection(:primary)
+          @writing_handler.establish_connection(:primary, :role_two)
+
+          # remove default
+          @writing_handler.remove_connection("primary")
+
+          assert_nil @writing_handler.retrieve_connection_pool("primary")
+          assert_not_nil @writing_handler.retrieve_connection_pool("primary", :role_two)
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
+
+        def test_connected?
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "db/primary.sqlite3" }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          @writing_handler.establish_connection(:primary)
+          @writing_handler.establish_connection(:primary, :role_two)
+
+          # connect to default
+          @writing_handler.connection_pool_list.first.checkout
+
+          assert @writing_handler.connected?("primary")
+          assert @writing_handler.connected?("primary", :default)
+          assert_not @writing_handler.connected?("primary", :role_two)
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
+      end
+    end
+  end
+end

activerecord/test/cases/fixtures_test.rb

@@ -1414,7 +1414,7 @@ class MultipleDatabaseFixturesTest < ActiveRecord::TestCase
 
   private
     def with_temporary_connection_pool
-      role = ActiveRecord::Base.connection_handler.send(:owner_to_role).fetch("primary")
+      role = ActiveRecord::Base.connection_handler.send(:owner_to_role_manager).fetch("primary").get_role(:default)
       new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(role)
 
       role.stub(:pool, new_pool) do

activerecord/test/cases/query_cache_test.rb

@@ -555,7 +555,7 @@ class QueryCacheTest < ActiveRecord::TestCase
 
   private
     def with_temporary_connection_pool
-      role = ActiveRecord::Base.connection_handler.send(:owner_to_role).fetch("primary")
+      role = ActiveRecord::Base.connection_handler.send(:owner_to_role_manager).fetch("primary").get_role(:default)
       new_pool = ActiveRecord::ConnectionAdapters::ConnectionPool.new(role)
 
       role.stub(:pool, new_pool) do