1 files changed, 95 insertions, 18 deletions

diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb
index 0dd788645b..706ce04062 100644
--- a/actionpack/lib/action_controller/metal/live.rb
+++ b/actionpack/lib/action_controller/metal/live.rb
@@ -48,7 +48,7 @@ module ActionController
     #   the server will receive a +Last-Event-ID+ header with value equal to +id+.
     #
     # After setting an option in the constructor of the SSE object, all future
-    # SSEs sent accross the stream will use those options unless overridden.
+    # SSEs sent across the stream will use those options unless overridden.
     #
     # Example Usage:
     #
@@ -102,13 +102,30 @@ module ActionController
             end
           end
 
-          @stream.write "data: #{json}\n\n"
+          message = json.gsub("\n", "\ndata: ")
+          @stream.write "data: #{message}\n\n"
         end
     end
 
+    class ClientDisconnected < RuntimeError
+    end
+
     class Buffer < ActionDispatch::Response::Buffer #:nodoc:
+      include MonitorMixin
+
+      # Ignore that the client has disconnected.
+      #
+      # If this value is `true`, calling `write` after the client
+      # disconnects will result in the written content being silently
+      # discarded. If this value is `false` (the default), a
+      # ClientDisconnected exception will be raised.
+      attr_accessor :ignore_disconnect
+
       def initialize(response)
-        @error_callback = nil
+        @error_callback = lambda { true }
+        @cv = new_cond
+        @aborted = false
+        @ignore_disconnect = false
         super(response, SizedQueue.new(10))
       end
 
@@ -119,17 +136,63 @@ module ActionController
         end
 
         super
+
+        unless connected?
+          @buf.clear
+
+          unless @ignore_disconnect
+            # Raise ClientDisconnected, which is a RuntimeError (not an
+            # IOError), because that's more appropriate for something beyond
+            # the developer's control.
+            raise ClientDisconnected, "client disconnected"
+          end
+        end
       end
 
       def each
+        @response.sending!
         while str = @buf.pop
           yield str
         end
+        @response.sent!
       end
 
+      # Write a 'close' event to the buffer; the producer/writing thread
+      # uses this to notify us that it's finished supplying content.
+      #
+      # See also #abort.
       def close
-        super
-        @buf.push nil
+        synchronize do
+          super
+          @buf.push nil
+          @cv.broadcast
+        end
+      end
+
+      # Inform the producer/writing thread that the client has
+      # disconnected; the reading thread is no longer interested in
+      # anything that's being written.
+      #
+      # See also #close.
+      def abort
+        synchronize do
+          @aborted = true
+          @buf.clear
+        end
+      end
+
+      # Is the client still connected and waiting for content?
+      #
+      # The result of calling `write` when this is `false` is determined
+      # by `ignore_disconnect`.
+      def connected?
+        !@aborted
+      end
+
+      def await_close
+        synchronize do
+          @cv.wait_until { @closed }
+        end
       end
 
       def on_error(&block)
@@ -142,7 +205,7 @@ module ActionController
     end
 
     class Response < ActionDispatch::Response #:nodoc: all
-      class Header < DelegateClass(Hash)
+      class Header < DelegateClass(Hash) # :nodoc:
         def initialize(response, header)
           @response = response
           super(header)
@@ -165,12 +228,20 @@ module ActionController
         end
       end
 
-      def commit!
-        headers.freeze
+      private
+
+      def before_committed
         super
+        jar = request.cookie_jar
+        # The response can be committed multiple times
+        jar.write self unless committed?
       end
 
-      private
+      def before_sending
+        super
+        request.cookie_jar.commit!
+        headers.freeze
+      end
 
       def build_buffer(response, body)
         buf = Live::Buffer.new response
@@ -191,6 +262,7 @@ module ActionController
       t1 = Thread.current
       locals = t1.keys.map { |key| [key, t1[key]] }
 
+      error = nil
       # This processes the action in a child thread. It lets us return the
       # response code and headers back up the rack stack, and still process
       # the body in parallel with sending data to the client
@@ -205,14 +277,18 @@ 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 => exception
-            log_error(exception)
-          ensure
-            log_error(e)
-            @_response.stream.close
+          if @_response.committed?
+            begin
+              @_response.stream.write(ActionView::Base.streaming_completion_on_exception) if request.format == :html
+              @_response.stream.call_on_error
+            rescue => exception
+              log_error(exception)
+            ensure
+              log_error(e)
+              @_response.stream.close
+            end
+          else
+            error = e
           end
         ensure
           @_response.commit!
@@ -220,6 +296,7 @@ module ActionController
       }
 
       @_response.await_commit
+      raise error if error
     end
 
     def log_error(exception)
@@ -234,7 +311,7 @@ module ActionController
 
     def response_body=(body)
       super
-      response.stream.close if response
+      response.close if response
     end
 
     def set_response!(request)