retry ConnectionNotEstablished errors

Author: ecwall

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 7.0.1 - 2023-03-24
+
+- Reconnect on retryable connection errors.
+
 ## 7.0.0 - 2022-06-02
 
 - Add support for Active Record 7.0.3

lib/active_record/connection_adapters/cockroachdb/transaction_manager.rb

@@ -10,21 +10,44 @@ module TransactionManagerMonkeyPatch
         # exceeded.
         def within_new_transaction(isolation: nil, joinable: true, attempts: 0)
           super(isolation: isolation, joinable: joinable)
+        rescue ActiveRecord::ConnectionNotEstablished => error
+          raise unless retryable? error
+          raise if attempts >= @connection.max_transaction_retries
+
+          sleep_rand_seconds(attempts)
+
+          unless @connection.active?
+            warn "connection isn't active, reconnecting"
+            @connection.reconnect!
+          end
+
+          within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts + 1) { yield }
         rescue ActiveRecord::StatementInvalid => error
           raise unless retryable? error
           raise if attempts >= @connection.max_transaction_retries
 
-          attempts += 1
-          sleep_seconds = (2 ** attempts + rand) / 10
-          sleep(sleep_seconds)
-          within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts) { yield }
+          sleep_rand_seconds(attempts)
+
+          within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts + 1) { yield }
         end
 
         def retryable?(error)
+          return true if serialization_error?(error)
           return true if error.is_a? ActiveRecord::SerializationFailure
           return retryable? error.cause if error.cause
           false
         end
+
+        def serialization_error?(error)
+          errors = [error]
+          errors << error.cause if error.cause
+          errors.any? {|e| e.is_a? PG::TRSerializationFailure }
+        end
+
+        def sleep_rand_seconds(attempts)
+          sleep_seconds = (2 ** attempts + rand) / 10
+          sleep(sleep_seconds)
+        end
       end
     end
 

lib/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ActiveRecord
-  COCKROACH_DB_ADAPTER_VERSION = "7.0.0"
+  COCKROACH_DB_ADAPTER_VERSION = "7.0.1"
 end