From 0f7c970e4f1cf0f3bcc01c22a6a3038cb3e34668 Mon Sep 17 00:00:00 2001
From: John Firebaugh <john_firebaugh@us.ibm.com>
Date: Tue, 23 Nov 2010 18:04:05 -0800
Subject: [PATCH] Introduce ActionDispatch::Reloader
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Based on the implementation on the 2-3-stable branch, patches by Hongli
Lai <hongli@phusion.nl>, and helpful suggestions from José Valim.

Hongli Lai's patches included locking around the request cycle; this is
now handled by Rack::Lock (https://github.com/rack/rack/issues/issue/87/).

[#2873]

Signed-off-by: José Valim <jose.valim@gmail.com>
---
 actionpack/lib/action_dispatch.rb             |   1 +
 .../action_dispatch/middleware/reloader.rb    |  82 +++++++++++
 actionpack/test/dispatch/reloader_test.rb     | 139 ++++++++++++++++++
 railties/lib/rails/application.rb             |   1 +
 railties/test/application/middleware_test.rb  |   7 +
 5 files changed, 230 insertions(+)
 create mode 100644 actionpack/lib/action_dispatch/middleware/reloader.rb
 create mode 100644 actionpack/test/dispatch/reloader_test.rb

diff --git a/actionpack/lib/action_dispatch.rb b/actionpack/lib/action_dispatch.rb
index 50faf666e6..50d0d191c1 100644
--- a/actionpack/lib/action_dispatch.rb
+++ b/actionpack/lib/action_dispatch.rb
@@ -53,6 +53,7 @@ module ActionDispatch
     autoload :Flash
     autoload :Head
     autoload :ParamsParser
+    autoload :Reloader
     autoload :RemoteIp
     autoload :Rescue
     autoload :ShowExceptions
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
diff --git a/actionpack/test/dispatch/reloader_test.rb b/actionpack/test/dispatch/reloader_test.rb
new file mode 100644
index 0000000000..58ef346b5c
--- /dev/null
+++ b/actionpack/test/dispatch/reloader_test.rb
@@ -0,0 +1,139 @@
+require 'abstract_unit'
+
+class ReloaderTest < Test::Unit::TestCase
+  Reloader = ActionDispatch::Reloader
+
+  def test_prepare_callbacks
+    a = b = c = nil
+    Reloader.to_prepare { |*args| a = b = c = 1 }
+    Reloader.to_prepare { |*args| b = c = 2 }
+    Reloader.to_prepare { |*args| c = 3 }
+
+    # Ensure to_prepare callbacks are not run when defined
+    assert_nil a || b || c
+
+    # Run callbacks
+    call_and_return_body
+
+    assert_equal 1, a
+    assert_equal 2, b
+    assert_equal 3, c
+  end
+
+  def test_to_prepare_with_identifier_replaces
+    a = b = 0
+    Reloader.to_prepare(:unique_id) { |*args| a = b = 1 }
+    Reloader.to_prepare(:unique_id) { |*args| a = 2 }
+
+    call_and_return_body
+    assert_equal 2, a
+    assert_equal 0, b
+  end
+
+  class MyBody < Array
+    def initialize(&block)
+      @on_close = block
+    end
+
+    def foo
+      "foo"
+    end
+
+    def bar
+      "bar"
+    end
+
+    def close
+      @on_close.call if @on_close
+    end
+  end
+
+  def test_returned_body_object_always_responds_to_close
+    body = call_and_return_body
+    assert body.respond_to?(:close)
+  end
+
+  def test_returned_body_object_behaves_like_underlying_object
+    body = call_and_return_body do
+      b = MyBody.new
+      b << "hello"
+      b << "world"
+      [200, { "Content-Type" => "text/html" }, b]
+    end
+    assert_equal 2, body.size
+    assert_equal "hello", body[0]
+    assert_equal "world", body[1]
+    assert_equal "foo", body.foo
+    assert_equal "bar", body.bar
+  end
+
+  def test_it_calls_close_on_underlying_object_when_close_is_called_on_body
+    close_called = false
+    body = call_and_return_body do
+      b = MyBody.new do
+        close_called = true
+      end
+      [200, { "Content-Type" => "text/html" }, b]
+    end
+    body.close
+    assert close_called
+  end
+
+  def test_returned_body_object_responds_to_all_methods_supported_by_underlying_object
+    body = call_and_return_body do
+      [200, { "Content-Type" => "text/html" }, MyBody.new]
+    end
+    assert body.respond_to?(:size)
+    assert body.respond_to?(:each)
+    assert body.respond_to?(:foo)
+    assert body.respond_to?(:bar)
+  end
+
+  def test_cleanup_callbacks_are_called_when_body_is_closed
+    cleaned = false
+    Reloader.to_cleanup { cleaned = true }
+
+    body = call_and_return_body
+    assert !cleaned
+
+    body.close
+    assert cleaned
+  end
+
+  def test_prepare_callbacks_arent_called_when_body_is_closed
+    prepared = false
+    Reloader.to_prepare { prepared = true }
+
+    body = call_and_return_body
+    prepared = false
+
+    body.close
+    assert !prepared
+  end
+
+  def test_manual_reloading
+    prepared = cleaned = false
+    Reloader.to_prepare { prepared = true }
+    Reloader.to_cleanup { cleaned  = true }
+
+    Reloader.prepare!
+    assert prepared
+    assert !cleaned
+
+    prepared = cleaned = false
+    Reloader.cleanup!
+    assert !prepared
+    assert cleaned
+
+    prepared = cleaned = false
+    Reloader.reload!
+    assert prepared
+    assert cleaned
+  end
+
+  private
+    def call_and_return_body(&block)
+      @reloader ||= Reloader.new(block || proc {[200, {}, 'response']})
+      @reloader.call({'rack.input' => StringIO.new('')})[2]
+    end
+end
diff --git a/railties/lib/rails/application.rb b/railties/lib/rails/application.rb
index b248bc737c..3b7505ce37 100644
--- a/railties/lib/rails/application.rb
+++ b/railties/lib/rails/application.rb
@@ -156,6 +156,7 @@ module Rails
         middleware.use ::ActionDispatch::ShowExceptions, config.consider_all_requests_local if config.action_dispatch.show_exceptions
         middleware.use ::ActionDispatch::RemoteIp, config.action_dispatch.ip_spoofing_check, config.action_dispatch.trusted_proxies
         middleware.use ::Rack::Sendfile, config.action_dispatch.x_sendfile_header
+        middleware.use ::ActionDispatch::Reloader unless config.cache_classes
         middleware.use ::ActionDispatch::Callbacks, !config.cache_classes
         middleware.use ::ActionDispatch::Cookies
 
diff --git a/railties/test/application/middleware_test.rb b/railties/test/application/middleware_test.rb
index 173ac40b12..a2217888e4 100644
--- a/railties/test/application/middleware_test.rb
+++ b/railties/test/application/middleware_test.rb
@@ -27,6 +27,7 @@ module ApplicationTests
         "ActionDispatch::ShowExceptions",
         "ActionDispatch::RemoteIp",
         "Rack::Sendfile",
+        "ActionDispatch::Reloader",
         "ActionDispatch::Callbacks",
         "ActiveRecord::ConnectionAdapters::ConnectionManagement",
         "ActiveRecord::QueryCache",
@@ -81,6 +82,12 @@ module ApplicationTests
       assert !middleware.include?("ActionDispatch::ShowExceptions")
     end
 
+    test "removes ActionDispatch::Reloader if cache_classes is true" do
+      add_to_config "config.cache_classes = true"
+      boot!
+      assert !middleware.include?("ActionDispatch::Reloader")
+    end
+
     test "use middleware" do
       use_frameworks []
       add_to_config "config.middleware.use Rack::Config"