Merge lp:~tcole/wsgi-oops/result-iteration-fixes into lp:wsgi-oops

[Tim Cole]

Properly handle closing of the result generator

This gracefully handles the WSGI container closing the result generator. Ordinarily this should only happen when the client drops the connection before we've finished writing all the data to it. This fix prevents oopses in such cases.

[lp:~tcole/wsgi-oops/result-iteration-fixes]

56. By Tim Cole on 2010-07-06

drop unused helper function

[Philip Fibiger (pfibiger)]

looks good.

review: Approve

[John O'Brien (jdobrien)]

Looks good. Great testing.

review: Approve

[Roman Yepishev (rye)]

Ran tests, tested that web interface continues working so approving.

review: Approve

=== modified file 'canonical/oops/serializer.py'
--- canonical/oops/serializer.py	2010-05-06 19:07:16 +0000
+++ canonical/oops/serializer.py	2010-07-06 15:47:25 +0000
@@ -56,7 +56,14 @@
 import zipfile
 
 from pprint import pformat
-from canonical.versioninfo import version_info
+try:
+    from canonical.versioninfo import version_info
+except ImportError:
+    version_info = {
+        'branch_nick': 'unknown',
+        'revno': '0',
+        'date': '1970-01-01'
+    }
 
 UTC = pytz.utc
 EPOCH = datetime(2006, 01, 01)

=== modified file 'canonical/oops/tests/test_wsgi.py'
--- canonical/oops/tests/test_wsgi.py	2010-01-05 16:23:19 +0000
+++ canonical/oops/tests/test_wsgi.py	2010-07-06 15:47:25 +0000
@@ -41,6 +41,68 @@
         self.assertFalse(self.oops.wanted)
 
 
+class TestResultIteration(testcase.WSGIOopsTestCase):
+    """Test iteration on results."""
+
+    def test_immediate_start(self):
+        """Test that we start without waiting for body consumption."""
+
+        def dummy_app(environ, start_response):
+            """Produce a response with infinite garbage."""
+            start_response("200 OK", {}, exc_info=None)
+            return ["foobar"]
+
+        started = [False]
+
+        def dummy_start_response(status, headers, exc_info=None):
+            """Sets flag when response is started"""
+            started[0] = True
+
+        app = wsgi.OopsWare(dummy_app, self.tmp_dir)
+        app({}, dummy_start_response)
+        self.assertTrue(started[0])
+
+    def _generate_garbage(self):
+        """Generate an infinite sequence of garbage until closed."""
+        while True:
+            yield "xyzzy"
+
+    def test_close_no_oops(self):
+        """Test that closing the generator doesn't oops"""
+
+        def dummy_app(environ, start_response):
+            """Produce a response with infinite garbage."""
+            start_response("200 OK", {}, exc_info=None)
+            return self._generate_garbage()
+
+        def dummy_start_response(status, headers, exc_info=None):
+            """Do nothing."""
+
+        interceptor = testcase.OopsWareInterceptor(dummy_app)
+        app = wsgi.OopsWare(interceptor, self.tmp_dir)
+        gen = app({}, dummy_start_response)
+        gen.close()
+        self.assertFalse(interceptor.oops.wanted)
+
+    def test_close_closes_inner_generator(self):
+        """Test that closing the outer generator closes the inner one."""
+        garbage = self._generate_garbage()
+
+        def dummy_app(environ, start_response):
+            """Produce a response with the specific garbage generator."""
+            start_response("200 OK", {}, exc_info=None)
+            return garbage
+
+        def dummy_start_response(status, headers, exc_info=None):
+            """Do nothing."""
+
+        interceptor = testcase.OopsWareInterceptor(dummy_app)
+        app = wsgi.OopsWare(interceptor, self.tmp_dir)
+        gen = app({}, dummy_start_response)
+        gen.close()
+        self.assertRaises(StopIteration, garbage.next)
+
+
 class TestEnvironWSGIOops(testcase.WSGIOopsTestCase):
     """Create the oops enviroment for testing"""
     def create_test_app(self):
@@ -61,6 +123,7 @@
         """ test that OOPSID is in the environ """
         res = self.testapp.get('?ok=all_ok', status=200, expect_errors=False)
 
+
 class TestWSGIOopsKey(testcase.WSGIOopsTestCase):
     """ Check that the OOPSSerializer key can be configured when instantiating
         an OopsWare.

=== modified file 'canonical/oops/tests/testcase.py'
--- canonical/oops/tests/testcase.py	2010-03-23 15:49:56 +0000
+++ canonical/oops/tests/testcase.py	2010-07-06 15:47:25 +0000
@@ -102,6 +102,6 @@
                 return ["An error was requested"]
         else:
             start_response("200 OK", [('Content-Type', 'text/html'), ])
-            return ["Hooary, all OK!"]
+            return ["Hooray, all OK!"]
 
 

=== modified file 'canonical/oops/wsgi.py'
--- canonical/oops/wsgi.py	2010-03-23 15:49:56 +0000
+++ canonical/oops/wsgi.py	2010-07-06 15:47:25 +0000
@@ -68,6 +68,12 @@
                                         key, oops_dir, fh=fh)
 
     def __call__(self, environ, start_response):
+        """Call do_request and pump once to start request processing."""
+        gen = self.do_request(environ, start_response)
+        gen.next() # first value yielded is meant to be discarded
+        return gen
+
+    def do_request(self, environ, start_response):
         """
         Capture exceptions during WSGI application execution and dump
         oopses as needed.
@@ -119,9 +125,13 @@
         metadata["request-start"] = time.time()
         try:
             body = self.app(environ, trap_response)
-            # trap exceptions during lazy content generation also
+            yield "" # thrown away by __call__
             for chunk in body:
                 yield chunk
+        except GeneratorExit:
+            close = getattr(body, 'close', None)
+            if close:
+                close()
         except:
             oops.keep()
             exc_info = sys.exc_info()