1 files changed, 30 insertions, 17 deletions
diff --git a/actionpack/lib/action_dispatch/middleware/reloader.rb b/actionpack/lib/action_dispatch/middleware/reloader.rb
index 29289a76b4..2f6968eb2e 100644
--- a/actionpack/lib/action_dispatch/middleware/reloader.rb
+++ b/actionpack/lib/action_dispatch/middleware/reloader.rb
@@ -18,10 +18,10 @@ module ActionDispatch
   # classes before they are unloaded.
   #
   # By default, ActionDispatch::Reloader is included in the middleware stack
-  # only in the development environment; specifically, when config.cache_classes
+  # only in the development environment; specifically, when +config.cache_classes+
   # is false. Callbacks may be registered even when it is not included in the
-  # middleware stack, but are executed only when +ActionDispatch::Reloader.prepare!+
-  # or +ActionDispatch::Reloader.cleanup!+ are called manually.
+  # middleware stack, but are executed only when <tt>ActionDispatch::Reloader.prepare!</tt>
+  # or <tt>ActionDispatch::Reloader.cleanup!</tt> are called manually.
   #
   class Reloader
     include ActiveSupport::Callbacks
@@ -43,34 +43,47 @@ module ActionDispatch
 
     # Execute all prepare callbacks.
     def self.prepare!
-      new(nil).run_callbacks :prepare
+      new(nil).prepare!
     end
 
     # Execute all cleanup callbacks.
     def self.cleanup!
-      new(nil).run_callbacks :cleanup
+      new(nil).cleanup!
     end
 
-    def initialize(app)
+    def initialize(app, condition=nil)
       @app = app
-    end
-
-    module CleanupOnClose
-      def close
-        super if defined?(super)
-      ensure
-        ActionDispatch::Reloader.cleanup!
-      end
+      @condition = condition || lambda { true }
+      @validated = true
     end
 
     def call(env)
-      run_callbacks :prepare
+      @validated = @condition.call
+      prepare!
+
       response = @app.call(env)
-      response[2].extend(CleanupOnClose)
+      response[2] = ::Rack::BodyProxy.new(response[2]) { cleanup! }
+
       response
     rescue Exception
-      run_callbacks :cleanup
+      cleanup!
       raise
     end
+
+    def prepare! #:nodoc:
+      run_callbacks :prepare if validated?
+    end
+
+    def cleanup! #:nodoc:
+      run_callbacks :cleanup if validated?
+    ensure
+      @validated = true
+    end
+
+    private
+
+    def validated? #:nodoc:
+      @validated
+    end
   end
 end