1 files changed, 82 insertions, 0 deletions
diff --git a/actionpack/lib/action_dispatch/middleware/reloader.rb b/actionpack/lib/action_dispatch/middleware/reloader.rb
new file mode 100644
index 0000000000..b84410fc62
--- /dev/null
+++ b/actionpack/lib/action_dispatch/middleware/reloader.rb
@@ -0,0 +1,82 @@
+module ActionDispatch
+  # ActionDispatch::Reloader provides to_prepare and to_cleanup callbacks.
+  # These are analogs of ActionDispatch::Callback's before and after
+  # callbacks, with the difference that to_cleanup is not called until the
+  # request is fully complete -- that is, after #close has been called on
+  # the request body. This is important for streaming responses such as the
+  # following:
+  #
+  #     self.response_body = lambda { |response, output|
+  #       # code here which refers to application models
+  #     }
+  #
+  # Cleanup callbacks will not be called until after the response_body lambda
+  # is evaluated, ensuring that it can refer to application models and other
+  # classes before they are unloaded.
+  #
+  # By default, ActionDispatch::Reloader is included in the middleware stack
+  # only in the development environment.
+  #
+  class Reloader
+    include ActiveSupport::Callbacks
+
+    define_callbacks :prepare, :scope => :name
+    define_callbacks :cleanup, :scope => :name
+
+    # Add a preparation callback. Preparation callbacks are run before each
+    # request.
+    #
+    # If a symbol with a block is given, the symbol is used as an identifier.
+    # That allows to_prepare to be called again with the same identifier to
+    # replace the existing callback. Passing an identifier is a suggested
+    # practice if the code adding a preparation block may be reloaded.
+    def self.to_prepare(*args, &block)
+      first_arg = args.first
+      if first_arg.is_a?(Symbol) && block_given?
+        remove_method :"__#{first_arg}" if method_defined?(:"__#{first_arg}")
+        define_method :"__#{first_arg}", &block
+        set_callback(:prepare, :"__#{first_arg}")
+      else
+        set_callback(:prepare, *args, &block)
+      end
+    end
+
+    # Add a cleanup callback. Cleanup callbacks are run after each request is
+    # complete (after #close is called on the response body).
+    def self.to_cleanup(&block)
+      set_callback(:cleanup, &block)
+    end
+
+    def self.prepare!
+      new(nil).send(:_run_prepare_callbacks)
+    end
+
+    def self.cleanup!
+      new(nil).send(:_run_cleanup_callbacks)
+    end
+
+    def self.reload!
+      prepare!
+      cleanup!
+    end
+
+    def initialize(app)
+      @app = app
+    end
+
+    module CleanupOnClose
+      def close
+        super if defined?(super)
+      ensure
+        ActionDispatch::Reloader.cleanup!
+      end
+    end
+
+    def call(env)
+      _run_prepare_callbacks
+      response = @app.call(env)
+      response[2].extend(CleanupOnClose)
+      response
+    end
+  end
+end