CockroachDB TransactionManagerMonkeyPatch monkey patch :D

crdb_transaction_retry.rb

# this lives in config/initializers 

require 'active_record/connection_adapters/cockroachdb/transaction_manager'
require 'pg'

module ActiveRecord
  module ConnectionAdapters
    module CockroachDB
      module TransactionManagerMonkeyPatch
        alias_method :old_retryable?, :retryable?
        alias_method :old_within_new_transaction, :within_new_transaction

        def within_new_transaction(isolation: nil, joinable: true, attempts: 0)
          old_within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts) { yield }
        rescue ActiveRecord::ConnectionNotEstablished => e
          unless retryable? e
            ec = e.class
            cc = e.cause ? e.cause.class : 'no cause'
            warn "raising due to not retryable. Error=(#{ec}), cause=(#{cc})"
            raise
          end

          if attempts >= @connection.max_transaction_retries
            warn "max retry attempts reached (#{attempts}/#{@connection.max_transaction_retries}), raising"
            raise
          end

          attempts += 1
          sleep_seconds = (2 ** attempts + rand) / 10
          sleep(sleep_seconds)

          unless @connection.active?
            warn "connection isn't active, reconnecting"
            @connection.reconnect!
          end

          within_new_transaction(isolation: isolation, joinable: joinable, attempts: attempts) { yield }
        end

        def retryable?(error)
          if serialization_error?(error)
            warn "retryable serialization error detected, retrying transaction"
            return true
          end

          old_retryable?(error)
        end

        private

        def serialization_error?(error)
          errors = [error]
          errors << error.cause if error.cause
          errors.any? {|e| e.is_a? PG::TRSerializationFailure }
        end

        def warn(message)
          Rails.logger.warn "TransactionManagerMonkeyPatch: #{message}"
        end
      end
    end
  end
end