From 5359428a142239578b4d1dfb43dd8c417ab57b5c Mon Sep 17 00:00:00 2001
From: Eugene Kenny <elkenny@gmail.com>
Date: Wed, 16 May 2018 22:40:20 +0100
Subject: Finalize transaction record state after real transaction

After a real (non-savepoint) transaction has committed or rolled back,
the original persistence-related state for all records modified in that
transaction is discarded or restored, respectively.

When the model has transactional callbacks, this happens synchronously
in the `committed!` or `rolled_back!` methods; otherwise, it happens
lazily the next time the record's persistence-related state is accessed.

The synchronous code path always finalizes the state of the record, but
the lazy code path only pops one "level" from the transaction counter,
assuming it will always reach zero immediately after a real transaction.
As the test cases included here demonstrate, that isn't always the case.

By using the same logic as the synchronous code path, we ensure that the
record's state is always updated after a real transaction has finished.
---
 .../connection_adapters/abstract/transaction.rb    | 37 ++++++++++++++--------
 activerecord/lib/active_record/transactions.rb     |  3 +-
 activerecord/test/cases/transactions_test.rb       | 29 +++++++++++++++++
 3 files changed, 54 insertions(+), 15 deletions(-)

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb b/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
index 0ce3796829..b59df2fff7 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/transaction.rb
@@ -17,11 +17,19 @@ module ActiveRecord
       end
 
       def committed?
-        @state == :committed
+        @state == :committed || @state == :fully_committed
+      end
+
+      def fully_committed?
+        @state == :fully_committed
       end
 
       def rolledback?
-        @state == :rolledback
+        @state == :rolledback || @state == :fully_rolledback
+      end
+
+      def fully_rolledback?
+        @state == :fully_rolledback
       end
 
       def fully_completed?
@@ -55,10 +63,19 @@ module ActiveRecord
         @state = :rolledback
       end
 
+      def full_rollback!
+        @children.each { |c| c.rollback! }
+        @state = :fully_rolledback
+      end
+
       def commit!
         @state = :committed
       end
 
+      def full_commit!
+        @state = :fully_committed
+      end
+
       def nullify!
         @state = nil
       end
@@ -88,10 +105,6 @@ module ActiveRecord
         records << record
       end
 
-      def rollback
-        @state.rollback!
-      end
-
       def rollback_records
         ite = records.uniq
         while record = ite.shift
@@ -103,10 +116,6 @@ module ActiveRecord
         end
       end
 
-      def commit
-        @state.commit!
-      end
-
       def before_commit_records
         records.uniq.each(&:before_committed!) if @run_commit_callbacks
       end
@@ -145,12 +154,12 @@ module ActiveRecord
 
       def rollback
         connection.rollback_to_savepoint(savepoint_name)
-        super
+        @state.rollback!
       end
 
       def commit
         connection.release_savepoint(savepoint_name)
-        super
+        @state.commit!
       end
 
       def full_rollback?; false; end
@@ -168,12 +177,12 @@ module ActiveRecord
 
       def rollback
         connection.rollback_db_transaction
-        super
+        @state.full_rollback!
       end
 
       def commit
         connection.commit_db_transaction
-        super
+        @state.full_commit!
       end
     end
 
diff --git a/activerecord/lib/active_record/transactions.rb b/activerecord/lib/active_record/transactions.rb
index 20603aaf2d..ccb2d57b03 100644
--- a/activerecord/lib/active_record/transactions.rb
+++ b/activerecord/lib/active_record/transactions.rb
@@ -472,7 +472,8 @@ module ActiveRecord
 
       def update_attributes_from_transaction_state(transaction_state)
         if transaction_state && transaction_state.finalized?
-          restore_transaction_record_state if transaction_state.rolledback?
+          restore_transaction_record_state(transaction_state.fully_rolledback?) if transaction_state.rolledback?
+          force_clear_transaction_record_state if transaction_state.fully_committed?
           clear_transaction_record_state if transaction_state.fully_completed?
         end
       end
diff --git a/activerecord/test/cases/transactions_test.rb b/activerecord/test/cases/transactions_test.rb
index 17aeabe34e..0d7857c0cf 100644
--- a/activerecord/test/cases/transactions_test.rb
+++ b/activerecord/test/cases/transactions_test.rb
@@ -679,6 +679,35 @@ class TransactionTest < ActiveRecord::TestCase
     assert_not_predicate topic, :frozen?
   end
 
+  def test_restore_new_record_after_double_save
+    topic = Topic.new
+
+    Topic.transaction do
+      topic.save!
+      topic.save!
+      raise ActiveRecord::Rollback
+    end
+
+    assert_predicate topic, :new_record?
+  end
+
+  def test_dont_restore_new_record_in_subsequent_transaction
+    topic = Topic.new
+
+    Topic.transaction do
+      topic.save!
+      topic.save!
+    end
+
+    Topic.transaction do
+      topic.save!
+      raise ActiveRecord::Rollback
+    end
+
+    assert_predicate topic, :persisted?
+    assert_not_predicate topic, :new_record?
+  end
+
   def test_restore_id_after_rollback
     topic = Topic.new
 
-- 
cgit v1.2.3