Ensure transaction tracking finishes when reconnecting

We noticed a couple of spots where transaction tracking events
were potentially incorrect.

When the connection reconnects with `restore_transactions: true`,
we were keeping the original start time for the transaction. In this case
it seems more accurate to treat these as separte transactions where the first
one finishes as incomplete.

Instead of forcing the adapter to tell the transaction manager to
suspend (or `detach` or `lost`) the transactions, we can keep that
logic encapsulated inside of the Transaction class for when to broadcast
a finish event. This means that we won't capture a finish event if you
manually call `reconnect!` without `restore_transactions: true`, but
that might be worthy of the tradeoff since this is a rare use-case anyway.

We also start raising here if the TransactionInstrumenter receives `#start`
when already started, or `#finish` when not started. This ensures that we don't
quietly end up in any weird states.

When marking that a transaction is incomplete, we also want to check
that the transaction is materialized to avoid finishing our instrumentation
if it hasn't already started. Also added a test to simulate losing a connection
without ever materializing a transaction.

Co-authored-by: Daniel Colson <[email protected]>

Author: ipc103

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