Merge pull request rails#49391 from ipc103/transaction-tracking-on-reconnect

Ensure transaction tracking finishes when reconnecting

Author: rafaelfranca

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

@@ -82,8 +82,11 @@ def initialize(payload = {})
         @base_payload = payload
       end
 
+      class InstrumentationNotStartedError < ActiveRecordError; end
+      class InstrumentationAlreadyStartedError < ActiveRecordError; end
+
       def start
-        return if @started
+        raise InstrumentationAlreadyStartedError.new("Called start on an already started transaction") if @started
         @started = true
 
         @payload = @base_payload.dup
@@ -92,7 +95,7 @@ def start
       end
 
       def finish(outcome)
-        return unless @started
+        raise InstrumentationNotStartedError.new("Called finish on a transaction that hasn't started") unless @started
         @started = false
 
         @payload[:outcome] = outcome
@@ -166,7 +169,7 @@ def restartable?
       end
 
       def incomplete!
-        @instrumenter.finish(:incomplete)
+        @instrumenter.finish(:incomplete) if materialized?
       end
 
       def materialize!
@@ -180,6 +183,7 @@ def materialized?
 
       def restore!
         if materialized?
+          incomplete!
           @materialized = false
           materialize!
         end
@@ -348,13 +352,13 @@ def rollback
           connection.rollback_to_savepoint(savepoint_name) if materialized?
         end
         @state.rollback!
-        @instrumenter.finish(:rollback)
+        @instrumenter.finish(:rollback) if materialized?
       end
 
       def commit
         connection.release_savepoint(savepoint_name) if materialized?
         @state.commit!
-        @instrumenter.finish(:commit)
+        @instrumenter.finish(:commit) if materialized?
       end
 
       def full_rollback?; false; end
@@ -389,13 +393,13 @@ def restart
       def rollback
         connection.rollback_db_transaction if materialized?
         @state.full_rollback!
-        @instrumenter.finish(:rollback)
+        @instrumenter.finish(:rollback) if materialized?
       end
 
       def commit
         connection.commit_db_transaction if materialized?
         @state.full_commit!
-        @instrumenter.finish(:commit)
+        @instrumenter.finish(:commit) if materialized?
       end
     end
 

activerecord/test/cases/transaction_instrumentation_test.rb

@@ -129,6 +129,27 @@ def test_transaction_instrumentation_with_unmaterialized_restart_parent_transact
     ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
   end
 
+  def test_transaction_instrumentation_with_materialized_restart_parent_transactions
+    topic = topics(:fifth)
+    events = []
+    subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
+      events << event
+    end
+
+    ActiveRecord::Base.transaction do
+      topic.update(title: "Sinatra")
+      ActiveRecord::Base.transaction(requires_new: true) do
+        raise ActiveRecord::Rollback
+      end
+    end
+
+    assert_equal 1, events.count
+    event = events.first
+    assert_equal :commit, event.payload[:outcome]
+  ensure
+    ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+  end
+
   def test_transaction_instrumentation_with_restart_savepoint_parent_transactions
     topic = topics(:fifth)
 
@@ -156,6 +177,27 @@ def test_transaction_instrumentation_with_restart_savepoint_parent_transactions
     ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
   end
 
+  def test_transaction_instrumentation_with_restart_savepoint_parent_transactions_on_commit
+    topic = topics(:fifth)
+
+    events = []
+    subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
+      events << event
+    end
+
+    ActiveRecord::Base.transaction do
+      topic.update(title: "Sinatra")
+      ActiveRecord::Base.transaction(requires_new: true) do
+      end
+    end
+
+    assert_equal 1, events.count
+    event = events.first
+    assert_equal :commit, event.payload[:outcome]
+  ensure
+    ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+  end
+
   def test_transaction_instrumentation_only_fires_if_materialized
     notified = false
     subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
@@ -170,6 +212,36 @@ def test_transaction_instrumentation_only_fires_if_materialized
     ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
   end
 
+  def test_transaction_instrumentation_only_fires_on_rollback_if_materialized
+    notified = false
+    subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
+      notified = true
+    end
+
+    ActiveRecord::Base.transaction do
+      raise ActiveRecord::Rollback
+    end
+
+    assert_not notified
+  ensure
+    ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+  end
+
+  def test_reconnecting_after_materialized_transaction_starts_new_event
+    events = []
+    subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
+      events << event
+    end
+    Topic.transaction do
+      Topic.connection.materialize_transactions
+      Topic.connection.reconnect!(restore_transactions: true)
+    end
+
+    assert_equal 2, events.count
+  ensure
+    ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+  end
+
   def test_transaction_instrumentation_fires_before_after_commit_callbacks
     notified = false
     after_commit_triggered = false
@@ -269,6 +341,26 @@ def test_transaction_instrumentation_on_failed_rollback
     ensure
       ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
     end
+
+    def test_transaction_instrumentation_on_failed_rollback_when_unmaterialized
+      notified = false
+      subscriber = ActiveSupport::Notifications.subscribe("transaction.active_record") do |event|
+        notified = true
+      end
+
+      error = Class.new(StandardError)
+      assert_raises error do
+        # Stubbing this method simulates an error that occurs when the transaction is still unmaterilized.
+        Topic.connection.transaction_manager.stub(:rollback_transaction, -> (*) { raise error }) do
+          Topic.transaction do
+            raise ActiveRecord::Rollback
+          end
+        end
+      end
+      assert_not notified
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+    end
   end
 
   def test_transaction_instrumentation_on_broken_subscription