use normal handle for timeouts

Author: fafhrd91

aiohttp/client.py

@@ -19,7 +19,7 @@
 from .client_ws import ClientWebSocketResponse
 from .cookiejar import CookieJar
 from .errors import WSServerHandshakeError
-from .helpers import TimeService
+from .helpers import TimeoutHandle, TimerContext, TimeService
 
 __all__ = ('ClientSession', 'request', 'get', 'options', 'head',
            'delete', 'post', 'put', 'patch', 'ws_connect')
@@ -195,8 +195,10 @@ def _request(self, method, url, *,
 
         # timeout is cumulative for all request operations
         # (request, redirects, responses, data consuming)
-        timer = self._time_service.timeout(timeout)
+        tm = TimeoutHandle(timeout)
+        handle = tm.handle(self._loop)
 
+        timer = TimerContext(self._loop, tm)
         with timer:
             while True:
                 url = URL(url).with_fragment(None)
@@ -273,6 +275,10 @@ def _request(self, method, url, *,
 
                 break
 
+        if resp.connection is not None:
+            resp.connection.add_callback(handle.cancel)
+        else:
+            handle.cancel()
         resp._history = tuple(history)
         return resp
 

aiohttp/client_reqrep.py

@@ -14,7 +14,7 @@
 import aiohttp
 
 from . import hdrs, helpers, streams
-from .helpers import HeadersMixin, SimpleCookie, _TimeServiceTimeoutNoop
+from .helpers import HeadersMixin, SimpleCookie, TimerNoop
 from .log import client_logger
 from .multipart import MultipartWriter
 from .protocol import HttpMessage
@@ -86,7 +86,7 @@ def __init__(self, method, url, *,
         self.compress = compress
         self.loop = loop
         self.response_class = response_class or ClientResponse
-        self._timer = timer if timer is not None else _TimeServiceTimeoutNoop()
+        self._timer = timer if timer is not None else TimerNoop()
 
         if loop.get_debug():
             self._source_traceback = traceback.extract_stack(sys._getframe(1))
@@ -526,7 +526,7 @@ def __init__(self, method, url, *,
         self._closed = False
         self._should_close = True  # override by message.should_close later
         self._history = ()
-        self._timer = timer if timer is not None else _TimeServiceTimeoutNoop()
+        self._timer = timer if timer is not None else TimerNoop()
         self.cookies = SimpleCookie()
 
     @property

aiohttp/connector.py

@@ -42,6 +42,7 @@ def __init__(self, connector, key, request, transport, protocol, loop):
         self._transport = transport
         self._protocol = protocol
         self._loop = loop
+        self._callbacks = []
         self.reader = protocol.reader
         self.writer = protocol.writer
 
@@ -68,25 +69,44 @@ def __del__(self, _warnings=warnings):
                 context['source_traceback'] = self._source_traceback
             self._loop.call_exception_handler(context)
 
+    def add_callback(self, callback):
+        if callback is not None:
+            self._callbacks.append(callback)
+
+    def release_callbacks(self):
+        callbacks, self._callbacks = self._callbacks[:], []
+
+        for cb in callbacks:
+            try:
+                cb()
+            except:
+                pass
+
     @property
     def loop(self):
         return self._loop
 
     def close(self):
+        self.release_callbacks()
+
         if self._transport is not None:
             self._connector._release(
                 self._key, self._request, self._transport, self._protocol,
                 should_close=True)
             self._transport = None
 
     def release(self):
+        self.release_callbacks()
+
         if self._transport is not None:
             self._connector._release(
                 self._key, self._request, self._transport, self._protocol,
                 should_close=False)
             self._transport = None
 
     def detach(self):
+        self.release_callbacks()
+
         if self._transport is not None:
             self._connector._release_acquired(self._key, self._transport)
         self._transport = None

aiohttp/helpers.py

@@ -15,6 +15,7 @@
 import warnings
 from collections import MutableSequence, namedtuple
 from functools import total_ordering
+from math import ceil
 from pathlib import Path
 from time import gmtime
 from urllib.parse import urlencode
@@ -716,18 +717,18 @@ def timeout(self, timeout):
 
         timeout - value in seconds or None to disable timeout logic
         """
-        if timeout:
-            ctx = _TimeServiceTimeoutContext(self._loop)
+        if timeout is not None and timeout > 0:
+            ctx = TimerContext(self._loop)
             when = self._loop_time + timeout
-            timer = TimerHandle(when, ctx.cancel, (), self._loop)
+            timer = TimerHandle(when, ctx.timeout, (), self._loop)
             heapq.heappush(self._scheduled, timer)
         else:
-            ctx = _TimeServiceTimeoutNoop()
+            ctx = TimerNoop()
 
         return ctx
 
 
-class _TimeServiceTimeoutNoop:
+class TimerNoop:
 
     def __enter__(self):
         return self
@@ -736,14 +737,45 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         return False
 
 
-class _TimeServiceTimeoutContext:
+class TimeoutHandle:
+    """ Timeout handle """
+
+    def __init__(self, timeout):
+        self._timeout = timeout
+        self._callbacks = []
+
+    def register(self, callback, *args, **kwargs):
+        self._callbacks.append((callback, args, kwargs))
+
+    def close(self):
+        self._callbacks.clear()
+
+    def handle(self, loop):
+        if self._timeout is not None and self._timeout > 0:
+            at = ceil(loop.time() + self._timeout)
+            return loop.call_at(at, self.__call__)
+
+    def __call__(self):
+        for cb, args, kwargs in self._callbacks:
+            try:
+                cb(*args, **kwargs)
+            except:
+                pass
+
+        self._callbacks.clear()
+
+
+class TimerContext:
     """ Low resolution timeout context manager """
 
-    def __init__(self, loop):
+    def __init__(self, loop, tm=None):
         self._loop = loop
         self._tasks = []
         self._cancelled = False
 
+        if tm is not None:
+            tm.register(self.timeout)
+
     def __enter__(self):
         task = asyncio.Task.current_task(loop=self._loop)
         if task is None:
@@ -759,17 +791,23 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         if self._tasks:
-            self._tasks.pop()
+            task = self._tasks.pop()
+        else:
+            task = None
 
         if exc_type is asyncio.CancelledError and self._cancelled:
+            for task in self._tasks:
+                task.cancel()
             raise asyncio.TimeoutError from None
 
-    def cancel(self):
+        if exc_type is None and self._cancelled and task is not None:
+            task.cancel()
+
+    def timeout(self):
         if not self._cancelled:
             for task in self._tasks:
                 task.cancel()
 
-            self._tasks = []
             self._cancelled = True
 
 

tests/test_client_functional.py

@@ -533,7 +533,11 @@ def handler(request):
 
 
 @asyncio.coroutine
-def test_timeout_on_reading_headers(loop, test_client):
+def test_timeout_on_reading_headers(loop, test_client, mocker):
+    def ceil(val):
+        return val
+
+    mocker.patch('aiohttp.helpers.ceil').side_effect = ceil
 
     @asyncio.coroutine
     def handler(request):
@@ -551,7 +555,7 @@ def handler(request):
 
 
 @asyncio.coroutine
-def test_timeout_on_conn_reading_headers(loop, test_client):
+def test_timeout_on_conn_reading_headers(loop, test_client, mocker):
     # tests case where user did not set a connection timeout
 
     @asyncio.coroutine
@@ -567,12 +571,17 @@ def handler(request):
     conn = aiohttp.TCPConnector(loop=loop)
     client = yield from test_client(app, connector=conn)
 
+    def ceil(val):
+        return val
+
+    mocker.patch('aiohttp.helpers.ceil').side_effect = ceil
+
     with pytest.raises(asyncio.TimeoutError):
         yield from client.get('/', timeout=0.01)
 
 
 @asyncio.coroutine
-def test_timeout_on_session_read_timeout(loop, test_client):
+def test_timeout_on_session_read_timeout(loop, test_client, mocker):
     @asyncio.coroutine
     def handler(request):
         resp = web.StreamResponse()
@@ -586,12 +595,22 @@ def handler(request):
     conn = aiohttp.TCPConnector(loop=loop)
     client = yield from test_client(app, connector=conn, read_timeout=0.01)
 
+    def ceil(val):
+        return val
+
+    mocker.patch('aiohttp.helpers.ceil').side_effect = ceil
+
     with pytest.raises(asyncio.TimeoutError):
         yield from client.get('/', timeout=None)
 
 
 @asyncio.coroutine
-def test_timeout_on_reading_data(loop, test_client):
+def test_timeout_on_reading_data(loop, test_client, mocker):
+
+    def ceil(val):
+        return val
+
+    mocker.patch('aiohttp.helpers.ceil').side_effect = ceil
 
     @asyncio.coroutine
     def handler(request):