From f3722a30b63f45add2e6e03c4e7498b46628b60b Mon Sep 17 00:00:00 2001
From: Emilio Tagua <miloops@gmail.com>
Date: Thu, 14 Oct 2010 15:39:57 -0300
Subject: [PATCH] Use IdentityMap middleware to flush map on each request.

---
 activerecord/lib/active_record/identity_map.rb       | 12 ++++++++++++
 activerecord/lib/active_record/railtie.rb            |  3 +++
 .../test/application/initializers/frameworks_test.rb |  2 +-
 railties/test/application/middleware_test.rb         |  2 ++
 4 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index 9f956858161e0..0062682555517 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -95,5 +95,17 @@ def identity_map
         ActiveRecord::IdentityMap
       end
     end
+
+    class Middleware
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        @app.call(env)
+      ensure
+        ActiveRecord::IdentityMap.clear
+      end
+    end
   end
 end
diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb
index dfe255ad7cc7f..5d0424a543874 100644
--- a/activerecord/lib/active_record/railtie.rb
+++ b/activerecord/lib/active_record/railtie.rb
@@ -22,6 +22,9 @@ class Railtie < Rails::Railtie
     config.app_middleware.insert_after "::ActionDispatch::Callbacks",
       "ActiveRecord::ConnectionAdapters::ConnectionManagement"
 
+    config.app_middleware.insert_after "::ActionDispatch::Callbacks",
+      "ActiveRecord::IdentityMap::Middleware"
+
     rake_tasks do
       load "active_record/railties/databases.rake"
     end
diff --git a/railties/test/application/initializers/frameworks_test.rb b/railties/test/application/initializers/frameworks_test.rb
index 475091f7894ea..9a340e6484145 100644
--- a/railties/test/application/initializers/frameworks_test.rb
+++ b/railties/test/application/initializers/frameworks_test.rb
@@ -166,7 +166,7 @@ def from_bar_helper
 
       require "#{app_path}/config/environment"
 
-      expects = [ActiveRecord::ConnectionAdapters::ConnectionManagement, ActiveRecord::QueryCache, ActiveRecord::SessionStore]
+      expects = [ActiveRecord::IdentityMap::Middleware, ActiveRecord::ConnectionAdapters::ConnectionManagement, ActiveRecord::QueryCache, ActiveRecord::SessionStore]
       middleware = Rails.application.config.middleware.map { |m| m.klass }
       assert_equal expects, middleware & expects
     end
diff --git a/railties/test/application/middleware_test.rb b/railties/test/application/middleware_test.rb
index 173ac40b12245..9820a92918c61 100644
--- a/railties/test/application/middleware_test.rb
+++ b/railties/test/application/middleware_test.rb
@@ -28,6 +28,7 @@ def app
         "ActionDispatch::RemoteIp",
         "Rack::Sendfile",
         "ActionDispatch::Callbacks",
+        "ActiveRecord::IdentityMap::Middleware",
         "ActiveRecord::ConnectionAdapters::ConnectionManagement",
         "ActiveRecord::QueryCache",
         "ActionDispatch::Cookies",
@@ -55,6 +56,7 @@ def app
       boot!
       assert !middleware.include?("ActiveRecord::ConnectionAdapters::ConnectionManagement")
       assert !middleware.include?("ActiveRecord::QueryCache")
+      assert !middleware.include?("ActiveRecord::IdentityMap::Middleware")
     end
 
     test "removes lock if allow concurrency is set" do