Wait for HTTP/2 requests to finish

lib/async/http/protocol/http2/server.rb

@@ -6,6 +6,7 @@
 require_relative "connection"
 require_relative "request"
 
+require "async/barrier"
 require "protocol/http2/server"
 
 module Async
@@ -51,9 +52,10 @@ def close(error = nil)
 					# 	@parameter request [Request] The incoming HTTP/2 request.
 					def each(task: Task.current)
 						task.annotate("Reading #{version} requests for #{self.class}.")
+						barrier = Async::Barrier.new(parent: task)
 
 						# It's possible the connection has died before we get here...
-						@requests.async do |task, request|
+						@requests.async(parent: barrier) do |task, request|
 							task.annotate("Incoming request: #{request.method} #{request.path.inspect}.")
 
 							response = nil
@@ -75,7 +77,9 @@ def each(task: Task.current)
 							end
 						end
 
-						# Maybe we should add some synchronisation here - i.e. only exit once all requests are finished.
+						barrier.wait
+					ensure
+						barrier&.stop
 					end
 				end
 			end

test/async/http/protocol/http2/server.rb

@@ -13,6 +13,13 @@
 	let(:sockets) {Socket.pair(Socket::PF_UNIX, Socket::SOCK_STREAM)}
 	let(:stream) {IO::Stream(sockets.first)}
 	let(:server) {subject.new(stream)}
+	let(:request) do
+		Object.new.tap do |request|
+			def request.method = "GET"
+			def request.path = "/"
+			def request.send_response(response) = nil
+		end
+	end
 
 	it "closes the request queue" do
 		request = Object.new
@@ -28,4 +35,33 @@
 			server.requests.enqueue(Object.new)
 		end.to raise_exception(Async::Queue::ClosedError)
 	end
+
+	it "waits for active requests to finish" do
+		handler_started = Async::Promise.new
+		handler_release = Async::Promise.new
+		each_finished = Async::Promise.new
+
+		server.requests.enqueue(request)
+
+		each_task = Async do
+			server.each do |incoming_request|
+				expect(incoming_request).to be == request
+				handler_started.resolve(nil)
+				handler_release.wait
+			end
+
+			each_finished.resolve(nil)
+		end
+
+		handler_started.wait
+		server.requests.close
+		Fiber.scheduler.yield
+
+		expect(each_finished).not.to be(:resolved?)
+
+		handler_release.resolve(nil)
+		each_task.wait
+
+		expect(each_finished).to be(:resolved?)
+	end
 end