MONGOID-5531 Add transaction callbacks (#5519)

Author: comandeo-mongo

lib/mongoid/clients/sessions.rb

@@ -5,7 +5,6 @@ module Clients
 
     # Encapsulates behavior for using sessions and transactions.
     module Sessions
-
       def self.included(base)
         base.include(ClassMethods)
       end
@@ -46,8 +45,7 @@ def with_session(options = {})
           end
         rescue Mongo::Error::OperationFailure => ex
           if (ex.code == 40415 && ex.server_message =~ /startTransaction/) ||
-            (ex.code == 20 && ex.server_message =~ /Transaction/)
-          then
+             (ex.code == 20 && ex.server_message =~ /Transaction/)
             raise Mongoid::Errors::TransactionsNotSupported.new
           else
             raise ex
@@ -82,27 +80,66 @@ def transaction(options = {}, session_options: {})
             begin
               session.start_transaction(options)
               yield
-              session.commit_transaction
+              commit_transaction(session)
             rescue Mongoid::Errors::Rollback
-              session.abort_transaction
+              abort_transaction(session)
             rescue Mongoid::Errors::InvalidSessionNesting
               # Session should be ended here.
               raise Mongoid::Errors::InvalidTransactionNesting.new
             rescue Mongo::Error::InvalidSession, Mongo::Error::InvalidTransactionOperation => e
-              session.abort_transaction
+              abort_transaction(session)
               raise Mongoid::Errors::TransactionError(e)
             rescue StandardError => e
-              session.abort_transaction
+              abort_transaction(session)
               raise e
             end
           end
         end
 
         private
 
+        # @return [ Mongo::Session ] Session for the current client.
         def _session
           Threaded.get_session(client: persistence_context.client)
         end
+
+        # This method should be used to detect whether a persistence operation
+        # is executed inside transaction or not.
+        #
+        # Currently this method is used to detect when +after_commit+ callbacks
+        # should be triggered. If we introduce implicit transactions and
+        # therefore do not need to handle two different ways of triggering callbacks,
+        # we may want to remove this method.
+        #
+        # @return [ true | false ] Whether there is a session for the current
+        #   client, and there is a transaction in progress for this session.
+        def in_transaction?
+          _session&.in_transaction? || false
+        end
+
+        # Commits the active transaction on the session, and calls
+        # after_commit callbacks on modified documents.
+        #
+        # @param [ Mongo::Session ] session Session on which
+        #   a transaction is started.
+        def commit_transaction(session)
+          session.commit_transaction
+          Threaded.clear_modified_documents(session).each do |doc|
+            doc.run_after_callbacks(:commit)
+          end
+        end
+
+        # Aborts the active transaction on the session, and calls
+        # after_rollback callbacks on modified documents.
+        #
+        # @param [ Mongo::Session ] session Session on which
+        #   a transaction is started.
+        def abort_transaction(session)
+          session.abort_transaction
+          Threaded.clear_modified_documents(session).each do |doc|
+            doc.run_after_callbacks(:rollback)
+          end
+        end
       end
     end
   end

lib/mongoid/interceptable.rb

@@ -43,6 +43,8 @@ module Interceptable
       # @api private
       define_model_callbacks :persist_parent
 
+      define_model_callbacks :commit, :rollback, only: :after
+
       attr_accessor :before_callback_halted
     end
 
@@ -115,8 +117,14 @@ def run_before_callbacks(*kinds)
     #   end
     #
     # @param [ Symbol ] kind The type of callback to execute.
-    # @param [ true | false ] with_children Flag specifies whether callbacks of embedded document should be run.
-    def run_callbacks(kind, with_children: true, &block)
+    # @param [ true | false ] with_children Flag specifies whether callbacks
+    #   of embedded document should be run.
+    # @param [ Proc | nil ] skip_if If this proc returns true, the callbacks
+    #   will not be triggered, while the given block will be still called.
+    def run_callbacks(kind, with_children: true, skip_if: nil, &block)
+      if skip_if&.call
+        return block&.call
+      end
       if with_children
         cascadable_children(kind).each do |child|
           if child.run_callbacks(child_callback_type(kind, child), with_children: with_children) == false

lib/mongoid/persistable.rb

@@ -167,6 +167,9 @@ def executing_atomically?
     def post_process_persist(result, options = {})
       post_persist unless result == false
       errors.clear unless performing_validations?(options)
+      if in_transaction?
+        Threaded.add_modified_document(_session, self)
+      end
       true
     end
 

lib/mongoid/persistable/creatable.rb

@@ -103,15 +103,17 @@ def prepare_insert(options = {})
         raise Errors::ReadonlyDocument.new(self.class) if readonly? && !Mongoid.legacy_readonly
         return self if performing_validations?(options) &&
           invalid?(options[:context] || :create)
-        run_callbacks(:save, with_children: false) do
-          run_callbacks(:create, with_children: false) do
-            run_callbacks(:persist_parent, with_children: false) do
-              _mongoid_run_child_callbacks(:save) do
-                _mongoid_run_child_callbacks(:create) do
-                  result = yield(self)
-                  if !result.is_a?(Document) || result.errors.empty?
-                    post_process_insert
-                    post_process_persist(result, options)
+        run_callbacks(:commit, with_children: true, skip_if: -> { in_transaction? }) do
+          run_callbacks(:save, with_children: false) do
+            run_callbacks(:create, with_children: false) do
+              run_callbacks(:persist_parent, with_children: false) do
+                _mongoid_run_child_callbacks(:save) do
+                  _mongoid_run_child_callbacks(:create) do
+                    result = yield(self)
+                    if !result.is_a?(Document) || result.errors.empty?
+                      post_process_insert
+                      post_process_persist(result, options)
+                    end
                   end
                 end
               end

lib/mongoid/persistable/deletable.rb

@@ -99,7 +99,7 @@ def notifying_parent?(options = {})
       #     collection.find(atomic_selector).remove
       #   end
       #
-      # @return [ Object ] The result of the block.
+      # @return [ true ] If the object was deleted successfully.
       def prepare_delete
         raise Errors::ReadonlyDocument.new(self.class) if readonly?
         yield(self)

lib/mongoid/persistable/destroyable.rb

@@ -18,11 +18,17 @@ module Destroyable
       def destroy(options = nil)
         raise Errors::ReadonlyDocument.new(self.class) if readonly?
         self.flagged_for_destroy = true
-        result = run_callbacks(:destroy) do
-          if catch(:abort) { apply_destroy_dependencies! }
-            delete(options || {})
-          else
-            false
+        result = run_callbacks(:commit, skip_if: -> { in_transaction? }) do
+          run_callbacks(:destroy) do
+            if catch(:abort) { apply_destroy_dependencies! }
+              delete(options || {}).tap do |res|
+                if res && in_transaction?
+                  Threaded.add_modified_document(_session, self)
+                end
+              end
+            else
+              false
+            end
           end
         end
         self.flagged_for_destroy = false

lib/mongoid/persistable/updatable.rb

@@ -102,15 +102,17 @@ def prepare_update(options = {})
         process_flagged_destroys
         update_children = cascadable_children(:update)
         process_touch_option(options, update_children)
-        run_callbacks(:save, with_children: false) do
-          run_callbacks(:update, with_children: false) do
-            run_callbacks(:persist_parent, with_children: false) do
-              _mongoid_run_child_callbacks(:save) do
-                _mongoid_run_child_callbacks(:update, children: update_children) do
-                  result = yield(self)
-                  self.previously_new_record = false
-                  post_process_persist(result, options)
-                  true
+        run_callbacks(:commit, with_children: true, skip_if: -> { in_transaction? }) do
+          run_callbacks(:save, with_children: false) do
+            run_callbacks(:update, with_children: false) do
+              run_callbacks(:persist_parent, with_children: false) do
+                _mongoid_run_child_callbacks(:save) do
+                  _mongoid_run_child_callbacks(:update, children: update_children) do
+                    result = yield(self)
+                    self.previously_new_record = false
+                    post_process_persist(result, options)
+                    true
+                  end
                 end
               end
             end

lib/mongoid/threaded.rb

@@ -29,6 +29,9 @@ module Threaded
     # The key for the current thread's sessions.
     SESSIONS_KEY="[mongoid]:sessions"
 
+    # The key for storing documents modified inside transactions.
+    MODIFIED_DOCUMENTS_KEY="[mongoid]:modified-documents"
+
     extend self
 
     # Begin entry into a named thread local stack.
@@ -353,9 +356,41 @@ def clear_session(client: nil)
       sessions.delete(client.object_id)&.end_session
     end
 
+    # Store a reference to the document that was modified inside a transaction
+    # associated with the session.
+    #
+    # @param [ Mongo::Session ] session Session in scope of which the document
+    #   was modified.
+    # @param [ Mongoid::Document ] document Mongoid document that was modified.
+    def add_modified_document(session, document)
+      if session&.in_transaction?
+        modified_documents[session] << document
+      end
+    end
+
+    # Clears the set of modified documents for the given session, and return the
+    # content of the set before the clearance.
+    # @param [ Mongo::Session ] session Session for which the modified documents
+    #   set should be cleared.
+    #
+    # @return [ Set<Mongoid::Document> ] Collection of modified documents before
+    #   it was cleared.
+    def clear_modified_documents(session)
+      modified_documents[session].dup
+    ensure
+      modified_documents[session].clear
+    end
+
     # @api private
     def sessions
       Thread.current[SESSIONS_KEY] ||= {}
     end
+
+    # @api private
+    def modified_documents
+      Thread.current[MODIFIED_DOCUMENTS_KEY] ||= Hash.new do |h, k|
+        h[k] = Set.new
+      end
+    end
   end
 end