Use inheritance chain instead of callbacks to increment counter caches after create

3 files changed, 31 insertions, 16 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..b0820f662a 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
+        with_cache_name { |name| decrement_counter name }
+      end
+
+      def increment_counters
+        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..bd2ee1a929 100644
--- a/activerecord/lib/active_record/associations/builder/belongs_to.rb
+++ b/activerecord/lib/active_record/associations/builder/belongs_to.rb
@@ -26,16 +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
@@ -74,10 +67,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)
       }
diff --git a/activerecord/lib/active_record/counter_cache.rb b/activerecord/lib/active_record/counter_cache.rb
index a5897edf03..163da8e870 100644
--- a/activerecord/lib/active_record/counter_cache.rb
+++ b/activerecord/lib/active_record/counter_cache.rb
@@ -131,6 +131,16 @@ module ActiveRecord
 
     private
 
+      def _create_record(*)
+        id = super
+
+        each_counter_cached_associations do |association|
+          association.increment_counters
+        end
+
+        id
+      end
+
       def destroy_row
         affected_rows = super
 
@@ -139,5 +149,11 @@ module ActiveRecord
         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