Ensure the request store is active during a streamed response.

Adds Rack::BodyProxy in order to hook into the `close`
callback that is fired when the request terminates after
generating all chunks of the streamed response.

Without this patch, the middleware cleans up the request
prematurely, after the controller returns, but before
all view rendering has completed.

lib/request_store/middleware.rb

@@ -1,3 +1,12 @@
+require 'rack/body_proxy'
+
+# A middleware that ensures the RequestStore stays around until
+# the last part of the body is rendered. This is useful when
+# using streaming.
+#
+# Uses Rack::BodyProxy, adapted from Rack::Lock's usage of the
+# same pattern.
+
 module RequestStore
   class Middleware
     def initialize(app)
@@ -6,10 +15,18 @@ def initialize(app)
 
     def call(env)
       RequestStore.begin!
-      @app.call(env)
+
+      response = @app.call(env)
+
+      returned = response << Rack::BodyProxy.new(response.pop) do
+        RequestStore.end!
+        RequestStore.clear!
+      end
     ensure
-      RequestStore.end!
-      RequestStore.clear!
+      unless returned
+        RequestStore.end!
+        RequestStore.clear!
+      end
     end
   end
 end

request_store.gemspec

@@ -18,6 +18,8 @@ Gem::Specification.new do |gem|
   gem.test_files    = gem.files.grep(%r{^(test|spec|features)/})
   gem.require_paths = ["lib"]
 
+  gem.add_dependency "rack", ">= 1.4"
+
   gem.add_development_dependency "rake", "~> 10.5"
   gem.add_development_dependency "minitest", "~> 5.0"
 end

test/middleware_test.rb

@@ -9,37 +9,57 @@ def setup
     @middleware = RequestStore::Middleware.new(@app)
   end
 
+  def call_middleware(opts = {})
+    _, _, proxy = @middleware.call(opts)
+    proxy.close
+  end
+
   def test_middleware_resets_store
-    2.times { @middleware.call({}) }
+    2.times do
+      call_middleware
+    end
 
     assert_equal 1, @app.last_value
     assert_equal({}, RequestStore.store)
   end
 
   def test_middleware_resets_store_on_error
     e = assert_raises RuntimeError do
-      @middleware.call({:error => true})
+      call_middleware error: true
     end
 
     assert_equal 'FAIL', e.message
     assert_equal({}, RequestStore.store)
   end
 
   def test_middleware_begins_store
-    @middleware.call({})
+    call_middleware
     assert_equal true, @app.store_active
   end
 
   def test_middleware_ends_store
-    @middleware.call({})
+    call_middleware
+
     assert_equal false, RequestStore.active?
   end
 
   def test_middleware_ends_store_on_error
     assert_raises RuntimeError do
-      @middleware.call({:error => true})
+      call_middleware error: true
     end
 
     assert_equal false, RequestStore.active?
   end
+
+  def test_middleware_stores_until_proxy_closes
+    _, _, proxy = @middleware.call({})
+
+    assert_equal 1, @app.last_value
+    assert RequestStore.active?
+
+    proxy.close
+
+    refute RequestStore.active?
+    refute RequestStore.store[:foo]
+  end
 end

test/test_helper.rb

@@ -7,5 +7,7 @@ def call(env)
     @last_value = RequestStore.store[:foo]
     @store_active = RequestStore.active?
     raise 'FAIL' if env[:error]
+
+    [200, {}, ["response"]]
   end
 end