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,27 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
  self._stmt_cache = collections.OrderedDict()
  self._stmts_to_close = set()
 
+ if command_timeout is not None:
+ if isinstance(command_timeout, bool):
+ raise ValueError(
+ 'invalid command_timeout value: '
+ 'expected non-negative float (got {!r})'.format(
+ command_timeout))
+
+ try:
+ command_timeout = float(command_timeout)
+ except ValueError:
+ raise ValueError(
+ 'invalid command_timeout value: '
+ 'expected non-negative float (got {!r})'.format(
+ command_timeout)) from None
+
+ if command_timeout < 0:
+ raise ValueError(
+ 'invalid command_timeout value: '
+ 'expected non-negative float (got {!r})'.format(
+ command_timeout))
+
  self._command_timeout = command_timeout
 
  self._listeners = {}
@@ -187,7 +209,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 +230,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 +302,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 +318,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 +332,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 +451,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,14 +577,29 @@ 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)
+ if timeout is not None:
+ return self._do_execute_with_timeout(query, executor, timeout)
+ else:
+ return self._do_execute(query, executor)
+
+ def _executemany(self, query, args, timeout):
+ executor = lambda stmt, timeout: self._protocol.bind_execute_many(
+ stmt, args, '', timeout)
+ timeout = self._protocol._get_timeout(timeout)
+ if timeout is not None:
+ return self._do_execute_with_timeout(query, executor, timeout)
+ else:
+ return self._do_execute(query, executor)
 
- try:
- result = await self._protocol.bind_execute(
- stmt, args, '', limit, return_status, timeout)
+ async def _do_execute(self, query, executor, retry=True):
+ stmt = await self._get_statement(query, None)
 
+ try:
+ result = await executor(stmt, None)
  except exceptions.InvalidCachedStatementError as e:
  # PostgreSQL will raise an exception when it detects
  # that the result type of the query has changed from
@@ -586,13 +624,38 @@ async def _do_execute(self, query, args, limit, timeout,
  # for discussion.
  #
  self._drop_global_statement_cache()
+ if self._protocol.is_in_transaction() or not retry:
+ raise
+ else:
+ result = await self._do_execute(
+ query, executor, retry=False)
+
+ return result
+
+ async def _do_execute_with_timeout(self, query, executor, timeout,
+ retry=True):
+ before = time.monotonic()
+ stmt = await self._get_statement(query, timeout)
+ after = time.monotonic()
+ timeout -= after - before
+ before = after
+
+ try:
+ try:
+ result = await executor(stmt, timeout)
+ finally:
+ after = time.monotonic()
+ timeout -= after - before
+
+ except exceptions.InvalidCachedStatementError as e:
+ # See comment in _do_execute().
+ 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_with_timeout(
+ 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,34 @@ 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 type(timeout) is bool:
+ raise ValueError(
+ 'invalid timeout value: expected non-negative float '
+ '(got {!r})'.format(timeout))
+ elif timeout is not None:
+ try:
+ 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 +401,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,6 +109,22 @@ 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):
 
@@ -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)