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 rails#45450

Author: eileencodes

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

@@ -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)
-
-        unless pool && pool.db_config == db_config
-          remove_connection_pool(pool_config.connection_specification_name, role: role, shard: shard)
+        existing_pool_config = pool_manager.get_pool_config(role, 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.
         #