diff --git a/actionpack/lib/action_controller/metal/live.rb b/actionpack/lib/action_controller/metal/live.rb index 6055fde4f7..5d395cd8bd 100644 --- a/actionpack/lib/action_controller/metal/live.rb +++ b/actionpack/lib/action_controller/metal/live.rb @@ -163,14 +163,6 @@ module ActionController 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. # @@ -210,6 +202,14 @@ module ActionController def call_on_error @error_callback.call end + + private + + def each_chunk(&block) + while str = @buf.pop + yield str + end + end end class Response < ActionDispatch::Response #:nodoc: all diff --git a/actionpack/lib/action_controller/test_case.rb b/actionpack/lib/action_controller/test_case.rb index ecd21f29ce..ed2edcbe06 100644 --- a/actionpack/lib/action_controller/test_case.rb +++ b/actionpack/lib/action_controller/test_case.rb @@ -554,6 +554,8 @@ module ActionController end @request.query_string = '' + @response.sent! + @response end diff --git a/actionpack/lib/action_dispatch/http/response.rb b/actionpack/lib/action_dispatch/http/response.rb index fa4c54701a..1515d59df3 100644 --- a/actionpack/lib/action_dispatch/http/response.rb +++ b/actionpack/lib/action_dispatch/http/response.rb @@ -68,7 +68,13 @@ module ActionDispatch # :nodoc: alias_method :headers, :header delegate :[], :[]=, :to => :@header - delegate :each, :to => :@stream + + def each(&block) + sending! + x = @stream.each(&block) + sent! + x + end CONTENT_TYPE = "Content-Type".freeze SET_COOKIE = "Set-Cookie".freeze @@ -97,10 +103,10 @@ module ActionDispatch # :nodoc: def body @str_body ||= begin - buf = '' - each { |chunk| buf << chunk } - buf - end + buf = '' + each { |chunk| buf << chunk } + buf + end end def write(string) @@ -112,10 +118,13 @@ module ActionDispatch # :nodoc: end def each(&block) - @response.sending! - x = @buf.each(&block) - @response.sent! - x + if @str_body + return enum_for(:each) unless block_given? + + yield @str_body + else + each_chunk(&block) + end end def abort @@ -129,6 +138,12 @@ module ActionDispatch # :nodoc: def closed? @closed end + + private + + def each_chunk(&block) + @buf.each(&block) # extract into own method + end end def self.create(status = 200, header = {}, body = [], default_headers: self.default_headers) diff --git a/actionpack/test/controller/live_stream_test.rb b/actionpack/test/controller/live_stream_test.rb index a7759c080b..5977124594 100644 --- a/actionpack/test/controller/live_stream_test.rb +++ b/actionpack/test/controller/live_stream_test.rb @@ -246,7 +246,8 @@ module ActionController def assert_stream_closed assert response.stream.closed?, 'stream should be closed' - assert response.sent?, 'stream should be sent' + assert response.committed?, 'response should be committed' + assert response.sent?, 'response should be sent' end def capture_log_output diff --git a/actionpack/test/dispatch/live_response_test.rb b/actionpack/test/dispatch/live_response_test.rb index e4475f4233..de57c4be1d 100644 --- a/actionpack/test/dispatch/live_response_test.rb +++ b/actionpack/test/dispatch/live_response_test.rb @@ -65,7 +65,7 @@ module ActionController latch = Concurrent::CountDownLatch.new t = Thread.new { - @response.stream.each do + @response.each do latch.count_down end } diff --git a/actionpack/test/dispatch/response_test.rb b/actionpack/test/dispatch/response_test.rb index 658e0d004b..aa90433505 100644 --- a/actionpack/test/dispatch/response_test.rb +++ b/actionpack/test/dispatch/response_test.rb @@ -37,6 +37,39 @@ class ResponseTest < ActiveSupport::TestCase assert_equal "closed stream", e.message end + def test_each_isnt_called_if_str_body_is_written + # Controller writes and reads response body + each_counter = 0 + @response.body = Object.new.tap {|o| o.singleton_class.send(:define_method, :each) { |&block| each_counter += 1; block.call 'foo' } } + @response['X-Foo'] = @response.body + + assert_equal 1, each_counter, "#each was not called once" + + # Build response + status, headers, body = @response.to_a + + assert_equal 200, status + assert_equal "foo", headers['X-Foo'] + assert_equal "foo", body.each.to_a.join + + # Show that #each was not called twice + assert_equal 1, each_counter, "#each was not called once" + end + + def test_set_header_after_read_body_during_action + @response.body + + # set header after the action reads back @response.body + @response['x-header'] = "Best of all possible worlds." + + # the response can be built. + status, headers, body = @response.to_a + assert_equal 200, status + assert_equal "", body.body + + assert_equal "Best of all possible worlds.", headers['x-header'] + end + def test_read_body_during_action @response.body = "Hello, World!"