Defer body.close until final chunk is written. (#53)

Author: samuel-williams-shopify

lib/protocol/http1/connection.rb

@@ -614,32 +614,34 @@ def write_fixed_length_body(body, length, head)
 
 				if head
 					@stream.flush
+				else
+					@stream.flush unless body.ready?
 
-					body.close
+					chunk_length = 0
+					# Use a manual read loop (not body.each) so that body.close runs after the response is fully written and flushed. This ensures completion callbacks (e.g. rack.response_finished) don't delay the client.
+					while chunk = body.read
+						chunk_length += chunk.bytesize
+						
+						if chunk_length > length
+							raise ContentLengthError, "Trying to write #{chunk_length} bytes, but content length was #{length} bytes!"
+						end
+						
+						@stream.write(chunk)
+						@stream.flush unless body.ready?
+					end
 
-					return
-				end
-				
-				@stream.flush unless body.ready?
-				
-				chunk_length = 0
-				body.each do |chunk|
-					chunk_length += chunk.bytesize
+					@stream.flush
 
-					if chunk_length > length
-						raise ContentLengthError, "Trying to write #{chunk_length} bytes, but content length was #{length} bytes!"
+					if chunk_length != length
+						raise ContentLengthError, "Wrote #{chunk_length} bytes, but content length was #{length} bytes!"
 					end
-					
-					@stream.write(chunk)
-					@stream.flush unless body.ready?
-				end
-				
-				@stream.flush
-				
-				if chunk_length != length
-					raise ContentLengthError, "Wrote #{chunk_length} bytes, but content length was #{length} bytes!"
 				end
+			rescue => error
+				raise
 			ensure
+				# Close the body after the response is fully flushed, so that completion callbacks run after the client has received the response:
+				body.close(error)
+				
 				self.send_end_stream!
 			end
 
@@ -657,34 +659,36 @@ def write_chunked_body(body, head, trailer = nil)
 
 				if head
 					@stream.flush
+				else
+					@stream.flush unless body.ready?
 
-					body.close
-					
-					return
-				end
-				
-				@stream.flush unless body.ready?
-				
-				body.each do |chunk|
-					next if chunk.size == 0
+					# Use a manual read loop (not body.each) so that body.close runs after the terminal chunk is written. With body.each, the ensure { close } fires before the terminal "0\r\n\r\n" is sent, delaying the client.
+					while chunk = body.read
+						next if chunk.size == 0
+						
+						@stream.write("#{chunk.bytesize.to_s(16).upcase}\r\n")
+						@stream.write(chunk)
+						@stream.write(CRLF)
+						
+						@stream.flush unless body.ready?
+					end
 
-					@stream.write("#{chunk.bytesize.to_s(16).upcase}\r\n")
-					@stream.write(chunk)
-					@stream.write(CRLF)
+					if trailer&.any?
+						@stream.write("0\r\n")
+						write_headers(trailer)
+						@stream.write("\r\n")
+					else
+						@stream.write("0\r\n\r\n")
+					end
 
-					@stream.flush unless body.ready?
-				end
-				
-				if trailer&.any?
-					@stream.write("0\r\n")
-					write_headers(trailer)
-					@stream.write("\r\n")
-				else
-					@stream.write("0\r\n\r\n")
+					@stream.flush
 				end
-				
-				@stream.flush
+			rescue => error
+				raise
 			ensure
+				# Close the body after the complete chunked response (including terminal chunk) is flushed, so that completion callbacks don't block the client from seeing the response as complete:
+				body.close(error)
+				
 				self.send_end_stream!
 			end
 
@@ -697,12 +701,11 @@ def write_body_and_close(body, head)
 				@persistent = false
 
 				@stream.write("\r\n")
-				@stream.flush unless body.ready?
 
-				if head
-					body.close
-				else
-					body.each do |chunk|
+				unless head
+					@stream.flush unless body.ready?
+					
+					while chunk = body.read
 						@stream.write(chunk)
 
 						@stream.flush unless body.ready?
@@ -711,7 +714,12 @@ def write_body_and_close(body, head)
 
 				@stream.flush
 				@stream.close_write
+			rescue => error
+				raise
 			ensure
+				# Close the body after the stream is fully written and half-closed, so that completion callbacks run after the client has received the full response:
+				body.close(error)
+				
 				self.send_end_stream!
 			end
 

releases.md

@@ -1,5 +1,9 @@
 # Releases
 
+## Unreleased
+
+  - Defer `body.close` in `write_chunked_body`, `write_fixed_length_body`, and `write_body_and_close` until after the response is fully written and flushed. Previously, `body.each` called `close` in its `ensure` block before the terminal chunk (chunked encoding) or final flush was written, causing `rack.response_finished` callbacks to delay the client-visible response completion.
+
 ## v0.37.0
 
   - `Protocol::HTTP1::BadRequest` now includes `Protocol::HTTP::BadRequest` for better interoperability and handling of bad request errors across different HTTP protocol implementations.

test/protocol/http1/connection/body_close_ordering.rb

@@ -0,0 +1,89 @@
+# frozen_string_literal: true
+
+# Released under the MIT License.
+# Copyright, 2026, by Samuel Williams.
+
+require "protocol/http1/connection"
+require "protocol/http/body/buffered"
+require "protocol/http/body/wrapper"
+
+require "connection_context"
+
+# A body wrapper whose close callback reads the client side of the socket to verify
+# that the response framing is already complete at close time.
+class StreamCheckingBody < Protocol::HTTP::Body::Wrapper
+	attr_reader :client_data_at_close
+	
+	def initialize(body, client_socket)
+		super(body)
+		@client_socket = client_socket
+		@client_data_at_close = nil
+	end
+	
+	def close(error = nil)
+		# Non-blocking read of everything available on the client socket right now.
+		# If the response was fully flushed before close, all framing bytes are here.
+		@client_data_at_close = +""
+		
+		loop do
+			@client_data_at_close << @client_socket.read_nonblock(65536)
+		rescue IO::WaitReadable, EOFError
+			break
+		end
+		
+		super
+	end
+end
+
+describe Protocol::HTTP1::Connection do
+	include_context ConnectionContext
+	
+	before do
+		server.open!
+		client.open!
+	end
+	
+	with "#write_body_and_close" do
+		it "flushes all data to the stream before closing body" do
+			body = StreamCheckingBody.new(
+				Protocol::HTTP::Body::Buffered.new(["Hello", " ", "World"]),
+				sockets.first,
+			)
+			
+			server.write_body_and_close(body, false)
+			
+			expect(body.client_data_at_close).not.to be_nil
+			expect(body.client_data_at_close).to be(:include?, "Hello")
+			expect(body.client_data_at_close).to be(:include?, "World")
+		end
+	end
+	
+	with "#write_chunked_body" do
+		it "writes terminal chunk before closing body" do
+			body = StreamCheckingBody.new(
+				Protocol::HTTP::Body::Buffered.new(["Hello", " ", "World"]),
+				sockets.first,
+			)
+			
+			server.write_chunked_body(body, false)
+			
+			expect(body.client_data_at_close).not.to be_nil
+			# Terminal chunk should already be on the wire:
+			expect(body.client_data_at_close).to be(:include?, "\r\n0\r\n\r\n")
+		end
+	end
+	
+	with "#write_fixed_length_body" do
+		it "writes all data before closing body" do
+			body = StreamCheckingBody.new(
+				Protocol::HTTP::Body::Buffered.new(["Hello", " ", "World"]),
+				sockets.first,
+			)
+			
+			server.write_fixed_length_body(body, 11, false)
+			
+			expect(body.client_data_at_close).not.to be_nil
+			expect(body.client_data_at_close).to be(:include?, "Hello World")
+		end
+	end
+end

test/protocol/http1/hijack.rb

@@ -32,7 +32,7 @@
 			server.open!
 
 			expect(body).to receive(:ready?).and_return(false)
-			expect(body).to receive(:each).and_return(nil)
+			expect(body).to receive(:read).and_return(nil)
 			server.write_response(response_version, 101, {"upgrade" => "websocket"})
 			server.write_body(response_version, body)