From a2a525bbb622849fa73802cb489679d3d4cadc1d Mon Sep 17 00:00:00 2001 From: eileencodes Date: Fri, 1 Nov 2019 14:27:20 -0400 Subject: [PATCH] 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. :smile: **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 --- .../lib/active_record/connection_adapters.rb | 1 + .../abstract/connection_pool.rb | 39 ++++--- .../connection_adapters/role_manager.rb | 27 +++++ .../lib/active_record/test_fixtures.rb | 6 +- .../connection_handlers_multi_role_test.rb | 106 ++++++++++++++++++ activerecord/test/cases/fixtures_test.rb | 2 +- activerecord/test/cases/query_cache_test.rb | 2 +- 7 files changed, 164 insertions(+), 19 deletions(-) create mode 100644 activerecord/lib/active_record/connection_adapters/role_manager.rb create mode 100644 activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb diff --git a/activerecord/lib/active_record/connection_adapters.rb b/activerecord/lib/active_record/connection_adapters.rb index 0d4dc941fc..d7626150c9 100644 --- a/activerecord/lib/active_record/connection_adapters.rb +++ b/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 diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb index 41041e9548..3b5c23d2fb 100644 --- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb +++ b/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 diff --git a/activerecord/lib/active_record/connection_adapters/role_manager.rb b/activerecord/lib/active_record/connection_adapters/role_manager.rb new file mode 100644 index 0000000000..3837cdbb0e --- /dev/null +++ b/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 diff --git a/activerecord/lib/active_record/test_fixtures.rb b/activerecord/lib/active_record/test_fixtures.rb index 60c9a84121..30d3804ce3 100644 --- a/activerecord/lib/active_record/test_fixtures.rb +++ b/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 diff --git a/activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb b/activerecord/test/cases/connection_adapters/connection_handlers_multi_role_test.rb new file mode 100644 index 0000000000..292472fc36 --- /dev/null +++ b/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 diff --git a/activerecord/test/cases/fixtures_test.rb b/activerecord/test/cases/fixtures_test.rb index ce48560c0d..9c0e9e5ea6 100644 --- a/activerecord/test/cases/fixtures_test.rb +++ b/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 diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb index c418816877..e35bc7010b 100644 --- a/activerecord/test/cases/query_cache_test.rb +++ b/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