Implement conistent handling of timeouts.

The timeout logic is currently a bit of a mess.  This commit attempts to
tidy things up.

Most importantly, the timeout budget is now applied consistently to the
_entire_ call, whereas previously multiple consecutive operations used
the same timeout value, making it possible for the overall run time to
exceed the timeout.

Secondly, tighten the validation for timeouts: booleans are not
accepted, and neither any value that cannot be converted to float.

Author: elprans

asyncpg/connection.py

@@ -11,6 +11,7 @@
 import os
 import socket
 import struct
+import time
 import urllib.parse
 
 from . import cursor
@@ -60,6 +61,22 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
         self._stmt_cache = collections.OrderedDict()
         self._stmts_to_close = set()
 
+        if command_timeout is not None:
+            try:
+                if isinstance(command_timeout, bool):
+                    raise ValueError
+
+                command_timeout = float(command_timeout)
+
+                if command_timeout < 0:
+                    raise ValueError
+
+            except ValueError:
+                raise ValueError(
+                    'invalid command_timeout value: '
+                    'expected non-negative float (got {!r})'.format(
+                        command_timeout)) from None
+
         self._command_timeout = command_timeout
 
         self._listeners = {}
@@ -187,7 +204,7 @@ async def execute(self, query: str, *args, timeout: float=None) -> str:
         if not args:
             return await self._protocol.query(query, timeout)
 
-        _, status, _ = await self._do_execute(query, args, 0, timeout, True)
+        _, status, _ = await self._execute(query, args, 0, timeout, True)
         return status.decode()
 
     async def executemany(self, command: str, args, timeout: float=None):
@@ -208,8 +225,7 @@ async def executemany(self, command: str, args, timeout: float=None):
 
         .. versionadded:: 0.7.0
         """
-        stmt = await self._get_statement(command, timeout)
-        return await self._protocol.bind_execute_many(stmt, args, '', timeout)
+        return await self._executemany(command, args, timeout)
 
     async def _get_statement(self, query, timeout):
         cache = self._stmt_cache_max_size > 0
@@ -281,7 +297,7 @@ async def fetch(self, query, *args, timeout=None) -> list:
 
         :return list: A list of :class:`Record` instances.
         """
-        return await self._do_execute(query, args, 0, timeout)
+        return await self._execute(query, args, 0, timeout)
 
     async def fetchval(self, query, *args, column=0, timeout=None):
         """Run a query and return a value in the first row.
@@ -297,7 +313,7 @@ async def fetchval(self, query, *args, column=0, timeout=None):
 
         :return: The value of the specified column of the first record.
         """
-        data = await self._do_execute(query, args, 1, timeout)
+        data = await self._execute(query, args, 1, timeout)
         if not data:
             return None
         return data[0][column]
@@ -311,7 +327,7 @@ async def fetchrow(self, query, *args, timeout=None):
 
         :return: The first row as a :class:`Record` instance.
         """
-        data = await self._do_execute(query, args, 1, timeout)
+        data = await self._execute(query, args, 1, timeout)
         if not data:
             return None
         return data[0]
@@ -430,7 +446,9 @@ async def _cleanup_stmts(self):
         to_close = self._stmts_to_close
         self._stmts_to_close = set()
         for stmt in to_close:
-            await self._protocol.close_statement(stmt, False)
+            # It is imperative that statements are cleaned properly,
+            # so we ignore the timeout.
+            await self._protocol.close_statement(stmt, protocol.NO_TIMEOUT)
 
     def _request_portal_name(self):
         return self._get_unique_id()
@@ -554,13 +572,37 @@ def _drop_global_statement_cache(self):
         else:
             self._drop_local_statement_cache()
 
-    async def _do_execute(self, query, args, limit, timeout,
-                          return_status=False):
-        stmt = await self._get_statement(query, timeout)
+    def _execute(self, query, args, limit, timeout, return_status=False):
+        executor = lambda stmt, timeout: self._protocol.bind_execute(
+            stmt, args, '', limit, return_status, timeout)
+        timeout = self._protocol._get_timeout(timeout)
+        return self._do_execute(query, executor, timeout)
+
+    def _executemany(self, query, args, timeout):
+        executor = lambda stmt, timeout: self._protocol.bind_execute_many(
+            stmt, args, '', timeout)
+        timeout = self._protocol._get_timeout(timeout)
+        return self._do_execute(query, executor, timeout)
+
+    async def _do_execute(self, query, executor, timeout, retry=True):
+        if timeout is None:
+            stmt = await self._get_statement(query, None)
+        else:
+            before = time.monotonic()
+            stmt = await self._get_statement(query, timeout)
+            after = time.monotonic()
+            timeout -= after - before
+            before = after
 
         try:
-            result = await self._protocol.bind_execute(
-                stmt, args, '', limit, return_status, timeout)
+            if timeout is None:
+                result = await executor(stmt, None)
+            else:
+                try:
+                    result = await executor(stmt, timeout)
+                finally:
+                    after = time.monotonic()
+                    timeout -= after - before
 
         except exceptions.InvalidCachedStatementError as e:
             # PostgreSQL will raise an exception when it detects
@@ -586,13 +628,11 @@ async def _do_execute(self, query, args, limit, timeout,
             # for discussion.
             #
             self._drop_global_statement_cache()
-
-            if self._protocol.is_in_transaction():
+            if self._protocol.is_in_transaction() or not retry:
                 raise
             else:
-                stmt = await self._get_statement(query, timeout)
-                result = await self._protocol.bind_execute(
-                    stmt, args, '', limit, return_status, timeout)
+                result = await self._do_execute(
+                    query, executor, timeout, retry=False)
 
         return result
 

asyncpg/protocol/__init__.py

@@ -5,4 +5,4 @@
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
 
-from .protocol import Protocol, Record  # NOQA
+from .protocol import Protocol, Record, NO_TIMEOUT  # NOQA

asyncpg/protocol/protocol.pxd

@@ -47,7 +47,8 @@ cdef class BaseProtocol(CoreProtocol):
 
         PreparedStatementState statement
 
-    cdef _ensure_clear_state(self)
+    cdef _get_timeout_impl(self, timeout)
+    cdef _check_state(self)
     cdef _new_waiter(self, timeout)
 
     cdef _on_result__connect(self, object waiter)

asyncpg/protocol/protocol.pyx

@@ -79,6 +79,9 @@ include "coreproto.pyx"
 include "prepared_stmt.pyx"
 
 
+NO_TIMEOUT = object()
+
+
 cdef class BaseProtocol(CoreProtocol):
     def __init__(self, addr, connected_fut, con_args, loop):
         CoreProtocol.__init__(self, con_args)
@@ -132,7 +135,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         if stmt_name is None:
             self.uid_counter += 1
@@ -154,7 +158,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         self._bind_execute(
             portal_name,
@@ -178,7 +183,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         # Make sure the argument sequence is encoded lazily with
         # this generator expression to keep the memory pressure under
@@ -209,7 +215,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         self._bind(
             portal_name,
@@ -231,7 +238,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         self._execute(
             portal_name,
@@ -251,7 +259,11 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        # query() needs to call _get_timeout instead of _get_timeout_impl
+        # for consistent validation, as it is called differently from
+        # prepare/bind/execute methods.
+        timeout = self._get_timeout(timeout)
 
         self._simple_query(query)
         self.last_query = query
@@ -266,7 +278,8 @@ cdef class BaseProtocol(CoreProtocol):
             await self.cancel_sent_waiter
             self.cancel_sent_waiter = None
 
-        self._ensure_clear_state()
+        self._check_state()
+        timeout = self._get_timeout_impl(timeout)
 
         if state.refs != 0:
             raise RuntimeError(
@@ -348,7 +361,32 @@ cdef class BaseProtocol(CoreProtocol):
     cdef _set_server_parameter(self, name, val):
         self.settings.add_setting(name, val)
 
-    cdef _ensure_clear_state(self):
+    def _get_timeout(self, timeout):
+        if timeout is not None:
+            try:
+                if type(timeout) is bool:
+                    raise ValueError
+                timeout = float(timeout)
+            except ValueError:
+                raise ValueError(
+                    'invalid timeout value: expected non-negative float '
+                    '(got {!r})'.format(timeout)) from None
+
+        return self._get_timeout_impl(timeout)
+
+    cdef inline _get_timeout_impl(self, timeout):
+        if timeout is None:
+            timeout = self.connection._command_timeout
+        elif timeout is NO_TIMEOUT:
+            timeout = None
+        else:
+            timeout = float(timeout)
+
+        if timeout is not None and timeout <= 0:
+            raise asyncio.TimeoutError()
+        return timeout
+
+    cdef _check_state(self):
         if self.cancel_waiter is not None:
             raise apg_exc.InterfaceError(
                 'cannot perform operation: another operation is cancelling')
@@ -361,11 +399,9 @@ cdef class BaseProtocol(CoreProtocol):
 
     cdef _new_waiter(self, timeout):
         self.waiter = self.create_future()
-        if timeout is not False:
-            timeout = timeout or self.connection._command_timeout
-            if timeout is not None and timeout > 0:
-                self.timeout_handle = self.connection._loop.call_later(
-                    timeout, self.timeout_callback, self.waiter)
+        if timeout is not None:
+            self.timeout_handle = self.connection._loop.call_later(
+                timeout, self.timeout_callback, self.waiter)
         self.waiter.add_done_callback(self.completed_callback)
         return self.waiter
 

tests/test_timeout.py

@@ -6,8 +6,9 @@
 
 
 import asyncio
-import asyncpg
 
+import asyncpg
+from asyncpg import connection as pg_connection
 from asyncpg import _testbase as tb
 
 
@@ -108,12 +109,28 @@ async def test_timeout_06(self):
 
         self.assertEqual(await self.con.fetch('select 1'), [(1,)])
 
+    async def test_invalid_timeout(self):
+        for command_timeout in ('a', False, -1):
+            with self.subTest(command_timeout=command_timeout):
+                with self.assertRaisesRegex(ValueError,
+                                            'invalid command_timeout'):
+                    await self.cluster.connect(
+                        database='postgres', loop=self.loop,
+                        command_timeout=command_timeout)
+
+        # Note: negative timeouts are OK for method calls.
+        for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
+            for timeout in ('a', False):
+                with self.subTest(timeout=timeout):
+                    with self.assertRaisesRegex(ValueError, 'invalid timeout'):
+                        await self.con.execute('SELECT 1', timeout=timeout)
+
 
 class TestConnectionCommandTimeout(tb.ConnectedTestCase):
 
     def getExtraConnectOptions(self):
         return {
-            'command_timeout': 0.02
+            'command_timeout': 0.2
         }
 
     async def test_command_timeout_01(self):
@@ -123,3 +140,25 @@ async def test_command_timeout_01(self):
                 meth = getattr(self.con, methname)
                 await meth('select pg_sleep(10)')
             self.assertEqual(await self.con.fetch('select 1'), [(1,)])
+
+
+class SlowPrepareConnection(pg_connection.Connection):
+    """Connection class to test timeouts."""
+    async def _get_statement(self, query, timeout):
+        await asyncio.sleep(0.15, loop=self._loop)
+        return await super()._get_statement(query, timeout)
+
+
+class TestTimeoutCoversPrepare(tb.ConnectedTestCase):
+
+    def getExtraConnectOptions(self):
+        return {
+            '__connection_class__': SlowPrepareConnection,
+            'command_timeout': 0.3
+        }
+
+    async def test_timeout_covers_prepare_01(self):
+        for methname in {'fetch', 'fetchrow', 'fetchval', 'execute'}:
+            with self.assertRaises(asyncio.TimeoutError):
+                meth = getattr(self.con, methname)
+                await meth('select pg_sleep($1)', 0.2)