2 files changed, 47 insertions, 3 deletions
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 4ad0b0d205..647f9e11d2 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -98,14 +98,32 @@ 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
+        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
diff --git a/activerecord/test/cases/identity_map/middleware_test.rb b/activerecord/test/cases/identity_map/middleware_test.rb
index 1004e958e5..130f370da6 100644
--- a/activerecord/test/cases/identity_map/middleware_test.rb
+++ b/activerecord/test/cases/identity_map/middleware_test.rb
@@ -29,6 +29,32 @@ module ActiveRecord
         }
         mw.call({})
       end
+
+      class Enum < Struct.new(:iter)
+        def each(&b)
+          iter.call(&b)
+        end
+      end
+
+      def test_im_enabled_during_body_each
+        mw = Middleware.new lambda { |env|
+          [200, {}, Enum.new(lambda { |&b|
+            assert IdentityMap.enabled?, 'identity map should be enabled'
+            b.call "hello"
+          })]
+        }
+        body = mw.call({}).last
+        body.each { |x| assert_equal 'hello', x }
+      end
+
+      def test_im_disabled_after_body_close
+        mw = Middleware.new lambda { |env| [200, {}, []] }
+        body = mw.call({}).last
+        assert IdentityMap.enabled?, 'identity map should be enabled'
+        body.close
+        assert !IdentityMap.enabled?, 'identity map should be disabled'
+      end
+
     end
   end
 end