Refactor connection handler's establish_connection

In the prior code we were getting the pool manager a bunch of times -
once in the retrieve_connection_pool, once in remove_connection_pool,
and once in the else conditional after setting it.

This change ensures we're only getting the pool manager once. If there
is no pool manager for the given key then we create a new one in the new
`set_pool_manager` method and use that.

The refactor also has a change in that we no longer instrument if a new
connection was not established. This instrumentation is private though
(denoted by the !) so it's safe to make this change.

Followup to #45450
This commit is contained in:
eileencodes 2022-06-30 08:58:52 -04:00
parent 462545c0a5
commit 0396fe918d
No known key found for this signature in database
GPG Key ID: BA5C575120BBE8DF

@ -106,29 +106,29 @@ def establish_connection(config, owner_name: Base, role: ActiveRecord::Base.curr
pool_config = resolve_pool_config(config, owner_name, role, shard)
db_config = pool_config.db_config
pool_manager = set_pool_manager(pool_config.connection_specification_name)
# If there is an existing pool with the same values as the pool_config
# don't remove the connection. Connections should only be removed if we are
# establishing a connection on a class that is already connected to a different
# configuration.
pool = retrieve_connection_pool(pool_config.connection_specification_name, role: role, shard: shard)
existing_pool_config = pool_manager.get_pool_config(role, shard)
unless pool && pool.db_config == db_config
remove_connection_pool(pool_config.connection_specification_name, role: role, shard: shard)
owner_to_pool_manager[pool_config.connection_specification_name] ||= PoolManager.new
pool_manager = get_pool_manager(pool_config.connection_specification_name)
if existing_pool_config && existing_pool_config.db_config == db_config
existing_pool_config.pool
else
disconnect_pool_from_pool_manager(pool_manager, role, shard)
pool_manager.set_pool_config(role, shard, pool_config)
pool = pool_config.pool
end
payload = {
spec_name: pool_config.connection_specification_name,
shard: shard,
config: db_config.configuration_hash
}
payload = {
spec_name: pool_config.connection_specification_name,
shard: shard,
config: db_config.configuration_hash
}
ActiveSupport::Notifications.instrumenter.instrument("!connection.active_record", payload) do
pool
ActiveSupport::Notifications.instrumenter.instrument("!connection.active_record", payload) do
pool_config.pool
end
end
end
@ -196,12 +196,7 @@ def connected?(spec_name, role: ActiveRecord::Base.current_role, shard: ActiveRe
def remove_connection_pool(owner, role: ActiveRecord::Base.current_role, shard: ActiveRecord::Base.current_shard)
if pool_manager = get_pool_manager(owner)
pool_config = pool_manager.remove_pool_config(role, shard)
if pool_config
pool_config.disconnect!
pool_config.db_config
end
disconnect_pool_from_pool_manager(pool_manager, role, shard)
end
end
@ -221,6 +216,20 @@ def get_pool_manager(owner)
owner_to_pool_manager[owner]
end
# Get the existing pool manager or initialize and assign a new one.
def set_pool_manager(owner)
owner_to_pool_manager[owner] ||= PoolManager.new
end
def disconnect_pool_from_pool_manager(pool_manager, role, shard)
pool_config = pool_manager.remove_pool_config(role, shard)
if pool_config
pool_config.disconnect!
pool_config.db_config
end
end
# Returns an instance of PoolConfig for a given adapter.
# Accepts a hash one layer deep that contains all connection information.
#