Add retry handling to connection establishment too

Author: matthewd

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -557,21 +557,42 @@ def all_foreign_keys_valid?
       def active?
       end
 
-      # Disconnects from the database if already connected, and establishes a
-      # new connection with the database. Implementors should call super
-      # immediately after establishing the new connection (and while still
-      # holding @lock).
+      # Disconnects from the database if already connected, and establishes a new
+      # connection with the database. Implementors should define private #reconnect
+      # instead.
       def reconnect!(restore_transactions: false)
-        enable_lazy_transactions!
-        @raw_connection_dirty = false
-        @verified = true
+        retries_available = connection_retries
+
+        @lock.synchronize do
+          reconnect
+
+          enable_lazy_transactions!
+          @raw_connection_dirty = false
+          @verified = true
+
+          reset_transaction(restore: restore_transactions) do
+            clear_cache!(new_connection: true)
+            configure_connection
+          end
+        rescue => original_exception
+          translated_exception = translate_exception_class(original_exception, nil, nil)
+
+          if retries_available > 0
+            retries_available -= 1
+
+            if retryable_connection_error?(translated_exception)
+              backoff(connection_retries - retries_available)
+              retry
+            end
+          end
 
-        reset_transaction(restore: restore_transactions) do
-          clear_cache!(new_connection: true)
-          configure_connection
+          @verified = false
+
+          raise translated_exception
         end
       end
 
+
       # Disconnects from the database if already connected. Otherwise, this
       # method does nothing.
       def disconnect!
@@ -845,14 +866,24 @@ def with_raw_connection(allow_retry: false, uses_transaction: true)
           @lock.synchronize do
             materialize_transactions if uses_transaction
 
-            retries_available = 0
+            retries_available = allow_retry ? connection_retries : 0
 
-            if reconnect_can_restore_state?
+            if @verified
+              # Cool, we're confident the connection's ready to use. (Note this might have
+              # become true during the above #materialize_transactions.)
+            elsif reconnect_can_restore_state?
               if allow_retry
-                retries_available = connection_retries
-              elsif !@verified
+                # Not sure about the connection yet, but if anything goes wrong we can
+                # just reconnect and re-run our query
+              else
+                # We can reconnect if needed, but we don't trust the upcoming query to be
+                # safely re-runnable: let's verify the connection to be sure
                 verify!
               end
+            else
+              # We don't know whether the connection is okay, but it also doesn't matter:
+              # we wouldn't be able to reconnect anyway. We're just going to run our query
+              # and hope for the best.
             end
 
             begin
@@ -868,7 +899,8 @@ def with_raw_connection(allow_retry: false, uses_transaction: true)
                 if retryable_query_error?(translated_exception)
                   backoff(connection_retries - retries_available)
                   retry
-                elsif retryable_connection_error?(translated_exception)
+                elsif retryable_connection_error?(translated_exception) &&
+                    reconnect_can_restore_state?
                   reconnect!(restore_transactions: true)
                   retry
                 end
@@ -894,6 +926,10 @@ def backoff(counter)
           sleep 0.1 * counter
         end
 
+        def reconnect
+          raise NotImplementedError
+        end
+
         # Returns a raw connection for internal use with methods that are known
         # to both be thread-safe and not rely upon actual server communication.
         # This is useful for e.g. string escaping methods.

activerecord/lib/active_record/connection_adapters/mysql2_adapter.rb

@@ -123,13 +123,6 @@ def active?
         @raw_connection.ping
       end
 
-      def reconnect!(restore_transactions: false)
-        @lock.synchronize do
-          @raw_connection.close
-          connect
-          super
-        end
-      end
       alias :reset! :reconnect!
 
       # Disconnects from the database if already connected.
@@ -158,6 +151,11 @@ def connect
           @raw_connection = self.class.new_client(@config)
         end
 
+        def reconnect
+          @raw_connection.close
+          connect
+        end
+
         def configure_connection
           @raw_connection.query_options[:as] = :array
           @raw_connection.query_options[:database_timezone] = default_timezone

activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb

@@ -325,19 +325,6 @@ def reload_type_map # :nodoc:
         end
       end
 
-      # Close then reopen the connection.
-      def reconnect!(restore_transactions: false)
-        @lock.synchronize do
-          begin
-            @raw_connection.reset
-          rescue PG::ConnectionBad
-            connect
-          end
-
-          super
-        end
-      end
-
       def reset!
         @lock.synchronize do
           unless @raw_connection.transaction_status == ::PG::PQTRANS_IDLE
@@ -885,6 +872,12 @@ def connect
           @raw_connection = self.class.new_client(@connection_parameters)
         end
 
+        def reconnect
+          @raw_connection.reset
+        rescue PG::ConnectionBad
+          connect
+        end
+
         # Configures the encoding, verbosity, schema search path, and time zone of the connection.
         # This is called by #connect and should not be called manually.
         def configure_connection

activerecord/lib/active_record/connection_adapters/sqlite3_adapter.rb

@@ -174,17 +174,6 @@ def active?
         !@raw_connection.closed?
       end
 
-      def reconnect!(restore_transactions: false)
-        @lock.synchronize do
-          if active?
-            @raw_connection.rollback rescue nil
-          else
-            connect
-          end
-
-          super
-        end
-      end
       alias :reset! :reconnect!
 
       # Disconnects from the database if already connected. Otherwise, this
@@ -633,6 +622,14 @@ def connect
           )
         end
 
+        def reconnect
+          if active?
+            @raw_connection.rollback rescue nil
+          else
+            connect
+          end
+        end
+
         def configure_connection
           @raw_connection.busy_timeout(self.class.type_cast_config_to_integer(@config[:timeout])) if @config[:timeout]