Merge pull request #238 from Pylons/close-app-iter-on-disconnect

interrupt the app_iter if it tries to write to a closed socket

4 files changed, 72 insertions, 19 deletions

diff --git a/CHANGES.txt b/CHANGES.txt
index 4815499..3521864 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -1,3 +1,14 @@
+unreleased
+----------
+
+Bugfixes
+~~~~~~~~
+
+- Stop early and close the ``app_iter`` when attempting to write to a closed
+  socket due to a client disconnect. This should notify a long-lived streaming
+  response when a client hangs up.
+  See https://github.com/Pylons/waitress/pull/238
+
 1.2.1 (2019-01-25)
 ------------------
 
diff --git a/waitress/channel.py b/waitress/channel.py
index 7ed3461..6416725 100644
--- a/waitress/channel.py
+++ b/waitress/channel.py
@@ -32,6 +32,9 @@ from waitress.utilities import InternalServerError
 
 from . import wasyncore
 
+class ClientDisconnected(Exception):
+    """ Raised when attempting to write to a closed socket."""
+
 class HTTPChannel(wasyncore.dispatcher, object):
     """
     Setting self.requests = [somerequest] prevents more requests from being
@@ -305,6 +308,10 @@ class HTTPChannel(wasyncore.dispatcher, object):
     #
 
     def write_soon(self, data):
+        if not self.connected:
+            # if the socket is closed then interrupt the task so that it
+            # can cleanup possibly before the app_iter is exhausted
+            raise ClientDisconnected
         if data:
             # the async mainloop might be popping data off outbuf; we can
             # block here waiting for it because we're in a task thread
@@ -334,6 +341,10 @@ class HTTPChannel(wasyncore.dispatcher, object):
                     task = self.task_class(self, request)
                 try:
                     task.service()
+                except ClientDisconnected:
+                    self.logger.warn('Client disconnected when serving %s' %
+                                     task.request.path)
+                    task.close_on_finish = True
                 except:
                     self.logger.exception('Exception when serving %s' %
                                           task.request.path)
diff --git a/waitress/task.py b/waitress/task.py
index 8e14b4f..dc283ee 100644
--- a/waitress/task.py
+++ b/waitress/task.py
@@ -451,25 +451,25 @@ class WSGITask(Task):
         # Call the application to handle the request and write a response
         app_iter = self.channel.server.application(env, start_response)
 
-        if app_iter.__class__ is ReadOnlyFileBasedBuffer:
-            # NB: do not put this inside the below try: finally: which closes
-            # the app_iter; we need to defer closing the underlying file.  It's
-            # intention that we don't want to call ``close`` here if the
-            # app_iter is a ROFBB; the buffer (and therefore the file) will
-            # eventually be closed within channel.py's _flush_some or
-            # handle_close instead.
-            cl = self.content_length
-            size = app_iter.prepare(cl)
-            if size:
-                if cl != size:
-                    if cl is not None:
-                        self.remove_content_length_header()
-                    self.content_length = size
-                self.write(b'') # generate headers
-                self.channel.write_soon(app_iter)
-                return
-
+        can_close_app_iter = True
         try:
+            if app_iter.__class__ is ReadOnlyFileBasedBuffer:
+                cl = self.content_length
+                size = app_iter.prepare(cl)
+                if size:
+                    if cl != size:
+                        if cl is not None:
+                            self.remove_content_length_header()
+                        self.content_length = size
+                    self.write(b'') # generate headers
+                    # if the write_soon below succeeds then the channel will
+                    # take over closing the underlying file via the channel's
+                    # _flush_some or handle_close so we intentionally avoid
+                    # calling close in the finally block
+                    self.channel.write_soon(app_iter)
+                    can_close_app_iter = False
+                    return
+
             first_chunk_len = None
             for chunk in app_iter:
                 if first_chunk_len is None:
@@ -503,7 +503,7 @@ class WSGITask(Task):
                                 self.content_bytes_written, cl),
                         )
         finally:
-            if hasattr(app_iter, 'close'):
+            if can_close_app_iter and hasattr(app_iter, 'close'):
                 app_iter.close()
 
     def parse_proxy_headers(
diff --git a/waitress/tests/test_channel.py b/waitress/tests/test_channel.py
index afe6e51..d550f87 100644
--- a/waitress/tests/test_channel.py
+++ b/waitress/tests/test_channel.py
@@ -225,6 +225,12 @@ class TestHTTPChannel(unittest.TestCase):
         self.assertEqual(outbufs[1], wrapper)
         self.assertEqual(outbufs[2].__class__.__name__, 'OverflowableBuffer')
 
+    def test_write_soon_disconnected(self):
+        from waitress.channel import ClientDisconnected
+        inst, sock, map = self._makeOneWithMap()
+        inst.connected = False
+        self.assertRaises(ClientDisconnected, lambda: inst.write_soon(b'stuff'))
+
     def test__flush_some_empty_outbuf(self):
         inst, sock, map = self._makeOneWithMap()
         result = inst._flush_some()
@@ -558,6 +564,27 @@ class TestHTTPChannel(unittest.TestCase):
         self.assertTrue(inst.close_when_flushed)
         self.assertTrue(request.closed)
 
+    def test_service_with_request_raises_disconnect(self):
+        from waitress.channel import ClientDisconnected
+
+        inst, sock, map = self._makeOneWithMap()
+        inst.adj.expose_tracebacks = False
+        inst.server = DummyServer()
+        request = DummyRequest()
+        inst.requests = [request]
+        inst.task_class = DummyTaskClass(ClientDisconnected)
+        inst.error_task_class = DummyTaskClass()
+        inst.logger = DummyLogger()
+        inst.service()
+        self.assertTrue(request.serviced)
+        self.assertEqual(inst.requests, [])
+        self.assertEqual(len(inst.logger.warnings), 1)
+        self.assertTrue(inst.force_flush)
+        self.assertTrue(inst.last_activity)
+        self.assertFalse(inst.will_close)
+        self.assertEqual(inst.error_task_class.serviced, False)
+        self.assertTrue(request.closed)
+
     def test_cancel_no_requests(self):
         inst, sock, map = self._makeOneWithMap()
         inst.requests = ()
@@ -699,6 +726,10 @@ class DummyLogger(object):
 
     def __init__(self):
         self.exceptions = []
+        self.warnings = []
+
+    def warn(self, msg):
+        self.warnings.append(msg)
 
     def exception(self, msg):
         self.exceptions.append(msg)