From 998350091be30fd84e102755f009304cc45e26dc Mon Sep 17 00:00:00 2001
From: John Hawthorn <john.hawthorn@gmail.com>
Date: Mon, 6 Jul 2015 16:19:47 -0700
Subject: [PATCH] Wait for all server requests to finish on session reset

---
 lib/capybara/server.rb  | 31 +++++++++++++++++++++++++++++++
 lib/capybara/session.rb |  1 +
 spec/server_spec.rb     | 23 +++++++++++++++++++++++
 3 files changed, 55 insertions(+)

diff --git a/lib/capybara/server.rb b/lib/capybara/server.rb
index 28df7bc6..2999374a 100644
--- a/lib/capybara/server.rb
+++ b/lib/capybara/server.rb
@@ -4,22 +4,47 @@ require 'rack'
 
 module Capybara
   class Server
+    class Counter
+      attr_reader :value
+
+      def initialize
+        @value = 0
+        @mutex = Mutex.new
+      end
+
+      def increment
+        @mutex.synchronize { @value += 1 }
+      end
+
+      def decrement
+        @mutex.synchronize { @value -= 1 }
+      end
+    end
+
     class Middleware
       attr_accessor :error
 
       def initialize(app)
         @app = app
+        @counter = Counter.new
+      end
+
+      def pending_requests?
+        @counter.value > 0
       end
 
       def call(env)
         if env["PATH_INFO"] == "/__identify__"
           [200, {}, [@app.object_id.to_s]]
         else
+          @counter.increment
           begin
             @app.call(env)
           rescue *Capybara.server_errors => e
             @error = e unless @error
             raise e
+          ensure
+            @counter.decrement
           end
         end
       end
@@ -62,6 +87,12 @@ module Capybara
       return false
     end
 
+    def wait_for_pending_requests
+      Timeout.timeout(60) { sleep(0.01) while @middleware.pending_requests? }
+    rescue Timeout::Error
+      raise "Requests did not finish in 60 seconds"
+    end
+
     def boot
       unless responsive?
         Capybara::Server.ports[@app.object_id] = @port
diff --git a/lib/capybara/session.rb b/lib/capybara/session.rb
index 147fc81a..800f4ab6 100644
--- a/lib/capybara/session.rb
+++ b/lib/capybara/session.rb
@@ -109,6 +109,7 @@ module Capybara
         assert_no_selector :xpath, "/html/body/*" if driver.browser_initialized?
         @touched = false
       end
+      @server.wait_for_pending_requests if @server
       raise_server_error!
     end
     alias_method :cleanup!, :reset!
diff --git a/spec/server_spec.rb b/spec/server_spec.rb
index 711e09f6..e91ec07e 100644
--- a/spec/server_spec.rb
+++ b/spec/server_spec.rb
@@ -113,4 +113,27 @@ RSpec.describe Capybara::Server do
     expect(Net::HTTP).to receive(:start).and_raise(SystemCallError.allocate)
     expect(server.responsive?).to eq false
   end
+
+  it "can detect and wait for pending requests" do
+    done = false
+    app = proc do |env|
+      sleep 0.2
+      done = true
+      [200, {}, ["Hello Server!"]]
+    end
+    server = Capybara::Server.new(app).boot
+
+    # Start request, but don't wait for it to finish
+    socket = TCPSocket.new(server.host, server.port)
+    socket.write "GET / HTTP/1.0\r\n\r\n"
+    socket.close
+    sleep 0.1
+
+    expect(done).to be false
+
+    server.wait_for_pending_requests
+
+    # Ensure server was allowed to finish
+    expect(done).to be true
+  end
 end