2 files changed, 47 insertions, 30 deletions

diff --git a/activesupport/lib/active_support/cache/strategy/local_cache.rb b/activesupport/lib/active_support/cache/strategy/local_cache.rb
index cea7eee924..73c6b3cb88 100644
--- a/activesupport/lib/active_support/cache/strategy/local_cache.rb
+++ b/activesupport/lib/active_support/cache/strategy/local_cache.rb
@@ -1,5 +1,6 @@
 require 'active_support/core_ext/object/duplicable'
 require 'active_support/core_ext/string/inflections'
+require 'active_support/per_thread_registry'
 
 module ActiveSupport
   module Cache
@@ -8,6 +9,8 @@ module ActiveSupport
       # duration of a block. Repeated calls to the cache for the same key will hit the
       # in-memory cache for faster access.
       module LocalCache
+        autoload :Middleware, 'active_support/cache/strategy/local_cache_middleware'
+
         # Class for storing and registering the local caches.
         class LocalCacheRegistry # :nodoc:
           extend ActiveSupport::PerThreadRegistry
@@ -63,32 +66,6 @@ module ActiveSupport
         def with_local_cache
           use_temporary_local_cache(LocalStore.new) { yield }
         end
-
-        #--
-        # This class wraps up local storage for middlewares. Only the middleware method should
-        # construct them.
-        class Middleware # :nodoc:
-          attr_reader :name, :local_cache_key
-
-          def initialize(name, local_cache_key)
-            @name             = name
-            @local_cache_key = local_cache_key
-            @app              = nil
-          end
-
-          def new(app)
-            @app = app
-            self
-          end
-
-          def call(env)
-            LocalCacheRegistry.set_cache_for(local_cache_key, LocalStore.new)
-            @app.call(env)
-          ensure
-            LocalCacheRegistry.set_cache_for(local_cache_key, nil)
-          end
-        end
-
         # Middleware class can be inserted as a Rack handler to be local cache for the
         # duration of request.
         def middleware
@@ -109,13 +86,13 @@ module ActiveSupport
 
         def increment(name, amount = 1, options = nil) # :nodoc:
           value = bypass_local_cache{super}
-          increment_or_decrement(value, name, amount, options)
+          set_cache_value(value, name, amount, options)
           value
         end
 
         def decrement(name, amount = 1, options = nil) # :nodoc:
           value = bypass_local_cache{super}
-          increment_or_decrement(value, name, amount, options)
+          set_cache_value(value, name, amount, options)
           value
         end
 
@@ -143,8 +120,7 @@ module ActiveSupport
             super
           end
 
-        private
-          def increment_or_decrement(value, name, amount, options)
+          def set_cache_value(value, name, amount, options)
             if local_cache
               local_cache.mute do
                 if value
@@ -156,6 +132,8 @@ module ActiveSupport
             end
           end
 
+        private
+
           def local_cache_key
             @local_cache_key ||= "#{self.class.name.underscore}_local_cache_#{object_id}".gsub(/[\/-]/, '_').to_sym
           end
diff --git a/activesupport/lib/active_support/cache/strategy/local_cache_middleware.rb b/activesupport/lib/active_support/cache/strategy/local_cache_middleware.rb
new file mode 100644
index 0000000000..901c2e05a8
--- /dev/null
+++ b/activesupport/lib/active_support/cache/strategy/local_cache_middleware.rb
@@ -0,0 +1,39 @@
+require 'rack/body_proxy'
+module ActiveSupport
+  module Cache
+    module Strategy
+      module LocalCache
+
+        #--
+        # This class wraps up local storage for middlewares. Only the middleware method should
+        # construct them.
+        class Middleware # :nodoc:
+          attr_reader :name, :local_cache_key
+
+          def initialize(name, local_cache_key)
+            @name             = name
+            @local_cache_key = local_cache_key
+            @app              = nil
+          end
+
+          def new(app)
+            @app = app
+            self
+          end
+
+          def call(env)
+            LocalCacheRegistry.set_cache_for(local_cache_key, LocalStore.new)
+            response = @app.call(env)
+            response[2] = ::Rack::BodyProxy.new(response[2]) do
+              LocalCacheRegistry.set_cache_for(local_cache_key, nil)
+            end
+            response
+          rescue Exception
+            LocalCacheRegistry.set_cache_for(local_cache_key, nil)
+            raise
+          end
+        end
+      end
+    end
+  end
+end