1 files changed, 8 insertions, 20 deletions
diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 680d9ffea0..d9777bb2f6 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -123,24 +123,6 @@ 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
@@ -148,8 +130,14 @@ module ActiveRecord
       def call(env)
         enabled = IdentityMap.enabled
         IdentityMap.enabled = true
-        status, headers, body = @app.call(env)
-        [status, headers, Body.new(body, enabled)]
+
+        response = @app.call(env)
+        response[2] = Rack::BodyProxy.new(response[2]) do
+          IdentityMap.enabled = enabled
+          IdentityMap.clear
+        end
+
+        response
       end
     end
   end