1 files changed, 31 insertions, 9 deletions
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 95a8e5cff7..9eb47ad99f 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -49,12 +49,15 @@ module ActiveRecord
       end
 
       def get(klass, primary_key)
-        obj = repository[klass.symbolized_base_class][primary_key]
-        if obj.is_a?(klass)
-          if ActiveRecord::Base.logger
-            ActiveRecord::Base.logger.debug "#{klass} with ID = #{primary_key} loaded from Identity Map"
-          end
-          obj
+        record = repository[klass.symbolized_base_class][primary_key]
+
+        if record.is_a?(klass)
+          ActiveSupport::Notifications.instrument("identity.active_record",
+            :line => "From Identity Map (id: #{primary_key})", 
+            :name => "#{klass} Loaded",
+            :connection_id => object_id)
+
+          record
         else
           nil
         end
@@ -95,14 +98,33 @@ module ActiveRecord
     end
 
     class Middleware
+      class Body #:nodoc:
+        def initialize(target, original)
+          @target   = target
+          @original = original
+        end
+
+        def each(&block)
+          @target.each(&block)
+        end
+
+        def close
+          @target.close if @target.respond_to?(:close)
+        ensure
+          IdentityMap.enabled = @original
+          IdentityMap.clear
+        end
+      end
+
       def initialize(app)
         @app = app
       end
 
       def call(env)
-        ActiveRecord::IdentityMap.use do
-          @app.call(env)
-        end
+        enabled = IdentityMap.enabled
+        IdentityMap.enabled = true
+        status, headers, body = @app.call(env)
+        [status, headers, Body.new(body, enabled)]
       end
     end
   end