Exception handling for controllers using ActionController::Live

Any exceptions that occured at the view or controller level for a
controller using ActionController::Live would cause the server to either
hang with an open socket indefinitely, or immediately crash (depending
on whether the server was launched with rails s or directly). Changed
the behavior of exceptions to act the same as streaming templates for
html requests, and allow for an on_error callback if needed.

2 files changed, 103 insertions, 1 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index fb664a69dd..0a77352e57 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -56,6 +56,14 @@ module ActionController
         super
         @buf.push nil
       end
+
+      def on_error(&block)
+        @error_callback = block
+      end
+
+      def call_on_error
+        @error_callback.call
+      end
     end
 
     class Response < ActionDispatch::Response #:nodoc: all
@@ -121,6 +129,16 @@ module ActionController
 
         begin
           super(name)
+        rescue => e
+          begin
+            @_response.stream.write(ActionView::Base.streaming_completion_on_exception) if request.format == :html
+            @_response.stream.call_on_error
+          rescue => exceptionception
+            log_error(exceptionception)
+          ensure
+            log_error(e)
+            @_response.stream.close
+          end
         ensure
           @_response.commit!
         end
@@ -129,6 +147,16 @@ module ActionController
       @_response.await_commit
     end
 
+    def log_error(exception)
+      logger = ActionController::Base.logger
+      return unless logger
+
+      message = "\n#{exception.class} (#{exception.message}):\n"
+      message << exception.annoted_source_code.to_s if exception.respond_to?(:annoted_source_code)
+      message << "  " << exception.backtrace.join("\n  ")
+      logger.fatal("#{message}\n\n")
+    end
+
     def response_body=(body)
       super
       response.stream.close if response
diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb
index 5755444a65..34164a19f0 100644
--- a/actionpack/test/controller/live_stream_test.rb
+++ b/actionpack/test/controller/live_stream_test.rb
@@ -52,6 +52,29 @@ module ActionController
       def with_stale
         render :text => 'stale' if stale?(:etag => "123")
       end
+
+      def exception_in_view
+        render 'doesntexist'
+      end
+
+      def exception_with_callback
+        response.headers['Content-Type'] = 'text/event-stream'
+
+        response.stream.on_error do
+          response.stream.write %(data: "500 Internal Server Error"\n\n)
+          response.stream.close
+        end
+
+        raise 'An exception occurred...'
+      end
+
+      def exception_in_exception_callback
+        response.headers['Content-Type'] = 'text/event-stream'
+        response.stream.on_error do
+          raise 'We need to go deeper.'
+        end
+        response.stream.write params[:widget][:didnt_check_for_nil]
+      end
     end
 
     tests TestController
@@ -66,6 +89,21 @@ module ActionController
       TestResponse.new
     end
 
+    def assert_stream_closed
+      assert response.stream.closed?, 'stream should be closed'
+    end
+
+    def capture_log_output
+      output = StringIO.new
+      old_logger, ActionController::Base.logger = ActionController::Base.logger, ActiveSupport::Logger.new(output)
+
+      begin
+        yield output
+      ensure
+        ActionController::Base.logger = old_logger
+      end
+    end
+
     def test_set_response!
       @controller.set_response!(@request)
       assert_kind_of(Live::Response, @controller.response)
@@ -119,7 +157,43 @@ module ActionController
     def test_render_text
       get :render_text
       assert_equal 'zomg', response.body
-      assert response.stream.closed?, 'stream should be closed'
+      assert_stream_closed
+    end
+
+    def test_exception_handling_html
+      capture_log_output do |output|
+        get :exception_in_view
+        assert_match %r((window\.location = "/500\.html"</script></html>)$), response.body
+        assert_match 'Missing template test/doesntexist', output.rewind && output.read
+        assert_stream_closed
+      end
+    end
+
+    def test_exception_handling_plain_text
+      capture_log_output do |output|
+        get :exception_in_view, format: :json
+        assert_equal '', response.body
+        assert_match 'Missing template test/doesntexist', output.rewind && output.read
+        assert_stream_closed
+      end
+    end
+
+    def test_exception_callback
+      capture_log_output do |output|
+        get :exception_with_callback, format: 'text/event-stream'
+        assert_equal %(data: "500 Internal Server Error"\n\n), response.body
+        assert_match 'An exception occurred...', output.rewind && output.read
+        assert_stream_closed
+      end
+    end
+
+    def test_exceptions_raised_handling_exceptions
+      capture_log_output do |output|
+        get :exception_in_exception_callback, format: 'text/event-stream'
+        assert_equal '', response.body
+        assert_match 'We need to go deeper', output.rewind && output.read
+        assert_stream_closed
+      end
     end
 
     def test_stale_without_etag