Add support for on_close and testing premature closes, sugars up the …

…case when a user closes the connection prior to a response.

lib/sinatra/async.rb

@@ -136,6 +136,14 @@ def ahalt(*args)
         invoke { error_block! response.status }
         body response.body
       end
+
+      # The given block will be executed if the user closes the connection
+      # prematurely (before we've sent a response). This is good for
+      # deregistering callbacks that would otherwise send the body (for
+      # example channel subscriptions).
+      def on_close(&blk)
+        env['async.close'].callback(&blk)
+      end
     end
 
     def self.registered(app) #:nodoc:

lib/sinatra/async/test.rb

@@ -8,10 +8,31 @@ def async?
 end
 
 class Sinatra::Async::Test
+
   class AsyncSession < Rack::MockSession
+    class AsyncCloser
+      def initialize
+        @callbacks, @errbacks = [], []
+      end
+      def callback(&b)
+        @callbacks << b
+      end
+      def errback(&b)
+        @errbacks << b
+      end
+      def fail
+        @errbacks.each { |cb| cb.call }
+        @errbacks.clear
+      end
+      def succeed
+        @callbacks.each { |cb| cb.call }
+        @callbacks.clear
+      end
+    end
+
     def request(uri, env)
       env['async.callback'] = lambda { |r| s,h,b = *r; handle_last_response(uri, env, s,h,b) }
-      env['async.close'] = lambda { raise 'close connection' } # XXX deal with this
+      env['async.close'] = AsyncCloser.new
       catch(:async) { super }
       @last_response ||= Rack::MockResponse.new(-1, {}, [], env["rack.errors"].flush)
     end
@@ -48,6 +69,12 @@ def assert_async
       assert last_response.async?
     end
 
+    # Simulate a user closing the connection before a response is sent.
+    def async_close
+      raise ArgumentError, 'please make a request first' unless last_request
+      current_session.last_request.env['async.close'].succeed
+    end
+
     def async_continue
       while b = app.options.async_schedules.shift
         b.call

test/test_async.rb

@@ -12,6 +12,12 @@ class TestApp < Sinatra::Base
     set :environment, :test
     register Sinatra::Async
 
+    # Hack for storing some global data accessible in tests (normally you
+    # shouldn't need to do this!)
+    def self.singletons
+      @singletons ||= []
+    end
+
     error 401 do
       '401'
     end
@@ -47,6 +53,18 @@ class TestApp < Sinatra::Base
     aget '/a401' do
       ahalt 401
     end
+
+    aget '/async_close' do
+      # don't call body here, the 'user' is going to 'disconnect' before we do
+      env['async.close'].callback { self.class.singletons << 'async_closed' }
+    end
+
+    aget '/on_close' do
+      # sugared version of the above
+      on_close do
+        self.class.singletons << 'async_close_cleaned_up'
+      end
+    end
   end
 
   def app
@@ -113,4 +131,20 @@ def test_error_blocks_async
     assert_equal 401, last_response.status
     assert_equal '401', last_response.body
   end
+
+  def test_async_close
+    get '/async_close'
+    assert_async
+    async_continue
+    async_close
+    assert_equal 'async_closed', TestApp.singletons.shift
+  end
+
+  def test_on_close
+    get '/on_close'
+    assert_async
+    async_continue
+    async_close
+    assert_equal 'async_close_cleaned_up', TestApp.singletons.shift
+  end
 end