Added method to IM to remove objects by class and id. Then used it to remove objects when updating counters.

2 files changed, 6 insertions, 0 deletions

diff --git a/activerecord/lib/active_record/counter_cache.rb b/activerecord/lib/active_record/counter_cache.rb
index ed0d4aef7f..d48ff92b6c 100644
--- a/activerecord/lib/active_record/counter_cache.rb
+++ b/activerecord/lib/active_record/counter_cache.rb
@@ -73,6 +73,8 @@ module ActiveRecord
         "#{quoted_column} = COALESCE(#{quoted_column}, 0) #{operator} #{value.abs}"
       end
 
+      IdentityMap.remove_by_id(symbolized_base_class, id) if IdentityMap.enabled?
+
       update_all(updates.join(', '), primary_key => id )
     end
 
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 30a2a1eb0f..7eb792fec4 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -59,6 +59,10 @@ module ActiveRecord
         repository[record.class.symbolized_base_class].delete(record.id)
       end
 
+      def remove_by_id(symbolized_base_class, id)
+        repository[symbolized_base_class].delete(id)
+      end
+
       def clear
         repository.clear
       end