Merge pull request #14765 from byroot/refactor-counter-cache-create-and-destroy

Refactor counter cache create and destroy

3 files changed, 44 insertions, 34 deletions

diff --git a/activerecord/lib/active_record/associations/belongs_to_association.rb b/activerecord/lib/active_record/associations/belongs_to_association.rb
index 8272a5584c..1edd4fa3aa 100644
--- a/activerecord/lib/active_record/associations/belongs_to_association.rb
+++ b/activerecord/lib/active_record/associations/belongs_to_association.rb
@@ -31,6 +31,14 @@ module ActiveRecord
         @updated
       end
 
+      def decrement_counters # :nodoc:
+        with_cache_name { |name| decrement_counter name }
+      end
+
+      def increment_counters # :nodoc:
+        with_cache_name { |name| increment_counter name }
+      end
+
       private
 
         def find_target?
@@ -51,13 +59,15 @@ module ActiveRecord
           end
         end
 
-        def decrement_counters
-          with_cache_name { |name| decrement_counter name }
+        def decrement_counter(counter_cache_name)
+          if foreign_key_present?
+            klass.decrement_counter(counter_cache_name, target_id)
+          end
         end
 
-        def decrement_counter counter_cache_name
+        def increment_counter(counter_cache_name)
           if foreign_key_present?
-            klass.decrement_counter(counter_cache_name, target_id)
+            klass.increment_counter(counter_cache_name, target_id)
           end
         end
 
diff --git a/activerecord/lib/active_record/associations/builder/belongs_to.rb b/activerecord/lib/active_record/associations/builder/belongs_to.rb
index 11be92ae01..47cc1f4b34 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -26,29 +26,9 @@ module ActiveRecord::Associations::Builder
     private
 
     def self.add_counter_cache_methods(mixin)
-      return if mixin.method_defined? :belongs_to_counter_cache_after_create
+      return if mixin.method_defined? :belongs_to_counter_cache_after_update
 
       mixin.class_eval do
-        def belongs_to_counter_cache_after_create(reflection)
-          if record = send(reflection.name)
-            cache_column = reflection.counter_cache_column
-            record.class.increment_counter(cache_column, record.id)
-            @_after_create_counter_called = true
-          end
-        end
-
-        def belongs_to_counter_cache_after_destroy(reflection)
-          foreign_key = reflection.foreign_key.to_sym
-          unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == foreign_key
-            record = send reflection.name
-            if record && self.actually_destroyed?
-              cache_column = reflection.counter_cache_column
-              record.class.decrement_counter(cache_column, record.id)
-              self.clear_destroy_state
-            end
-          end
-        end
-
         def belongs_to_counter_cache_after_update(reflection)
           foreign_key  = reflection.foreign_key
           cache_column = reflection.counter_cache_column
@@ -74,14 +54,6 @@ module ActiveRecord::Associations::Builder
     def self.add_counter_cache_callbacks(model, reflection)
       cache_column = reflection.counter_cache_column
 
-      model.after_create lambda { |record|
-        record.belongs_to_counter_cache_after_create(reflection)
-      }
-
-      model.after_destroy lambda { |record|
-        record.belongs_to_counter_cache_after_destroy(reflection)
-      }
-
       model.after_update lambda { |record|
         record.belongs_to_counter_cache_after_update(reflection)
       }
diff --git a/activerecord/lib/active_record/counter_cache.rb b/activerecord/lib/active_record/counter_cache.rb
index a5897edf03..b7b790322a 100644
--- a/activerecord/lib/active_record/counter_cache.rb
+++ b/activerecord/lib/active_record/counter_cache.rb
@@ -131,13 +131,41 @@ module ActiveRecord
 
     private
 
+      def _create_record(*)
+        id = super
+
+        each_counter_cached_associations do |association|
+          if send(association.reflection.name)
+            association.increment_counters
+            @_after_create_counter_called = true
+          end
+        end
+
+        id
+      end
+
       def destroy_row
         affected_rows = super
 
-        @_actually_destroyed = affected_rows > 0
+        if affected_rows > 0
+          each_counter_cached_associations do |association|
+            foreign_key = association.reflection.foreign_key.to_sym
+            unless destroyed_by_association && destroyed_by_association.foreign_key.to_sym == foreign_key
+              if send(association.reflection.name)
+                association.decrement_counters
+              end
+            end
+          end
+        end
 
         affected_rows
       end
 
+      def each_counter_cached_associations
+        reflections.each do |name, reflection|
+          yield association(name) if reflection.belongs_to? && reflection.counter_cache_column
+        end
+      end
+
   end
 end