Use block syntax in IdentityMap middleware.

2 files changed, 16 insertions, 3 deletions

diff --git a/activerecord/lib/active_record/identity_map.rb b/activerecord/lib/active_record/identity_map.rb
index cca4198c69..30a2a1eb0f 100644
--- a/activerecord/lib/active_record/identity_map.rb
+++ b/activerecord/lib/active_record/identity_map.rb
@@ -27,7 +27,12 @@ module ActiveRecord
       end
 
       def use
+        old, self.enabled = self.enabled, true
+
         yield if block_given?
+      ensure
+        self.enabled = old
+        ActiveRecord::IdentityMap.clear
       end
 
       def without
@@ -93,9 +98,9 @@ module ActiveRecord
       end
 
       def call(env)
-        @app.call(env)
-      ensure
-        ActiveRecord::IdentityMap.clear
+        ActiveRecord::IdentityMap.use do
+          @app.call(env)
+        end
       end
     end
   end
diff --git a/activerecord/test/cases/identity_map_test.rb b/activerecord/test/cases/identity_map_test.rb
index 1588633e1d..df77bc5af2 100644
--- a/activerecord/test/cases/identity_map_test.rb
+++ b/activerecord/test/cases/identity_map_test.rb
@@ -39,6 +39,14 @@ class IdentityMapTest < ActiveRecord::TestCase
     end
   end
 
+  def test_find_id_use_identity_map
+    ActiveRecord::IdentityMap.enabled = false
+    ActiveRecord::IdentityMap.use do
+      assert_same(Client.find(3), Client.find(3))
+    end
+    ActiveRecord::IdentityMap.enabled = true
+  end
+
   def test_find_pkey
     assert_same(
       Subscriber.find('swistak'),