Merge pull request #173 from KeepSafe/client_connector_refactoring

Client connector refactoring

Author: asvetlov

aiohttp/connector.py

@@ -35,7 +35,9 @@ def loop(self):
 
     def close(self):
         if self._transport is not None:
-            self._transport.close()
+            self._connector._release(
+                self._key, self._request, self._transport, self._protocol,
+                should_close=True)
             self._transport = None
             self._wr = None
 
@@ -173,21 +175,22 @@ def _get(self, key):
 
         return None, None
 
-    def _release(self, key, req, transport, protocol):
+    def _release(self, key, req, transport, protocol, *, should_close=False):
         resp = req.response
-        should_close = False
 
-        if resp is not None:
-            if resp.message is None:
-                should_close = True
-            else:
-                should_close = resp.message.should_close
+        if not should_close:
+            if resp is not None:
+                if resp.message is None:
+                    should_close = True
+                else:
+                    should_close = resp.message.should_close
 
-        if self._force_close:
-            should_close = True
+            if self._force_close:
+                should_close = True
 
         reader = protocol.reader
         if should_close or (reader.output and not reader.output.at_eof()):
+            self._conns.pop(key, None)
             transport.close()
         else:
             conns = self._conns.get(key)

aiohttp/protocol.py

@@ -204,6 +204,11 @@ class HttpResponseParser(HttpParser):
     Returns RawResponseMessage"""
 
     def __call__(self, out, buf):
+        try:
+            yield from buf.wait(1)
+        except aiohttp.EofStream:
+            raise errors.ClientConnectionError(
+                'Connection closed by server') from None
         try:
             # read http message (response line + headers)
             try:

tests/test_client_functional.py

@@ -1020,6 +1020,102 @@ def go(url):
             url = httpd.url('keepalive')
             self.loop.run_until_complete(go(url))
 
+    def test_server_close_keepalive_connection(self):
+
+        class Proto(asyncio.Protocol):
+
+            def connection_made(self, transport):
+                self.transp = transport
+                self.data = b''
+
+            def data_received(self, data):
+                self.data += data
+                if data.endswith(b'\r\n\r\n'):
+                    self.transp.write(
+                        b'HTTP/1.1 200 OK\r\n'
+                        b'CONTENT-LENGTH: 2\r\n'
+                        b'CONNECTION: close\r\n'
+                        b'\r\n'
+                        b'ok')
+                    self.transp.close()
+
+            def connection_lost(self, exc):
+                self.transp = None
+
+        @asyncio.coroutine
+        def go():
+            server = yield from self.loop.create_server(
+                Proto, '127.0.0.1')
+
+            addr = server.sockets[0].getsockname()
+
+            connector = aiohttp.TCPConnector(loop=self.loop)
+
+            url = 'http://{}:{}/'.format(*addr)
+            for i in range(2):
+                r = yield from client.request('GET', url,
+                                              connector=connector,
+                                              loop=self.loop)
+                yield from r.read()
+                self.assertEqual(0, len(connector._conns))
+            connector.close()
+            server.close()
+            yield from server.wait_closed()
+
+        self.loop.run_until_complete(go())
+
+    def test_handle_keepalive_on_closed_connection(self):
+
+        class Proto(asyncio.Protocol):
+
+            def connection_made(self, transport):
+                self.transp = transport
+                self.data = b''
+
+            def data_received(self, data):
+                self.data += data
+                if data.endswith(b'\r\n\r\n'):
+                    self.transp.write(
+                        b'HTTP/1.1 200 OK\r\n'
+                        b'CONTENT-LENGTH: 2\r\n'
+                        b'\r\n'
+                        b'ok')
+                    self.transp.close()
+
+            def connection_lost(self, exc):
+                self.transp = None
+
+        @asyncio.coroutine
+        def go():
+            server = yield from self.loop.create_server(
+                Proto, '127.0.0.1')
+
+            addr = server.sockets[0].getsockname()
+
+            connector = aiohttp.TCPConnector(loop=self.loop)
+
+            url = 'http://{}:{}/'.format(*addr)
+
+            r = yield from client.request('GET', url,
+                                          connector=connector,
+                                          loop=self.loop)
+            yield from r.read()
+            self.assertEqual(1, len(connector._conns))
+
+            with self.assertRaisesRegex(
+                    aiohttp.ClientConnectionError,
+                    'Connection closed by server'):
+                yield from client.request('GET', url,
+                                          connector=connector,
+                                          loop=self.loop)
+            self.assertEqual(0, len(connector._conns))
+
+            connector.close()
+            server.close()
+            yield from server.wait_closed()
+
+        self.loop.run_until_complete(go())
+
 
 class Functional(test_utils.Router):
 

tests/test_connector.py

@@ -39,8 +39,10 @@ def test_close(self):
             self.connector, self.key, self.request,
             self.transport, self.protocol, self.loop)
         conn.close()
-        self.assertTrue(self.transport.close.called)
         self.assertIsNone(conn._transport)
+        self.connector._release.assert_called_with(
+            self.key, self.request, self.transport, self.protocol,
+            should_close=True)
 
     def test_release(self):
         conn = Connection(