368 lines
13 KiB
Python
368 lines
13 KiB
Python
# testing/cancellation.py
|
|
# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
|
|
# <see AUTHORS file>
|
|
#
|
|
# This module is part of SQLAlchemy and is released under
|
|
# the MIT License: https://www.opensource.org/licenses/mit-license.php
|
|
# mypy: ignore-errors
|
|
|
|
"""Deterministic cancellation testing for the asyncio dialects.
|
|
|
|
An asyncio task may be cancelled at any ``await``, and reproducing a
|
|
specific one of those points by racing ``task.cancel()`` against real
|
|
network traffic is unreliable. Here the choice is made exact.
|
|
:func:`.greenlet_spawn` drives every await performed by an asyncio-driven
|
|
SQLAlchemy operation, ``connect()`` included, so instrumenting the greenlet
|
|
it switches into gives a stable enumeration of a scenario's IO points, and
|
|
cancelling at point ``N`` is reproducible run to run.
|
|
|
|
The cancellation is delivered to the task the application is running,
|
|
which is not necessarily the task performing the await; SQLAlchemy runs
|
|
some operations, such as the ``close()`` performed by an async context
|
|
manager on exit, in a task of its own.
|
|
|
|
Two things are asserted afterwards:
|
|
|
|
* Every DBAPI connection the pool creates is eventually closed through
|
|
``Pool._close_connection()``. A connection that is never closed is
|
|
leaked; for an asyncio driver the garbage collector cannot close it
|
|
either, since closing requires the event loop.
|
|
|
|
* Once the application's task has finished and garbage has been
|
|
collected, no pool slot is still checked out. A slot still checked out
|
|
at that point is lost: either nothing remains that could return it, or
|
|
it is being returned by a task SQLAlchemy left running in the
|
|
background, which the application has no way to wait for. In the
|
|
latter case, an application that disposes of its engine at that point
|
|
has the connection returned into the disposed pool afterwards, where
|
|
nothing closes it.
|
|
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import gc
|
|
from unittest import mock
|
|
|
|
from .. import exc
|
|
from ..util import concurrency
|
|
|
|
|
|
class _InstrumentedAwaitable:
|
|
"""An awaitable that invokes ``hook`` when it is awaited."""
|
|
|
|
__slots__ = ("awaitable", "hook")
|
|
|
|
def __init__(self, awaitable, hook):
|
|
self.awaitable = awaitable
|
|
self.hook = hook
|
|
|
|
def __await__(self):
|
|
self.hook(self.awaitable)
|
|
return self.awaitable.__await__()
|
|
|
|
|
|
class AwaitPoints:
|
|
"""Counts await points, and cancels the running task at each of the
|
|
ones given.
|
|
|
|
``cancel_at`` is a tuple of await point numbers; each one is a separate
|
|
``task.cancel()`` call. asyncio delivers a given cancellation exactly
|
|
once, so a second number models a second, independent canceller
|
|
reaching a task that is still cleaning up after the first.
|
|
|
|
"""
|
|
|
|
def __init__(self, cancel_at=()):
|
|
self._count = 0
|
|
self._cancel_at = cancel_at
|
|
self._cancelled_at = []
|
|
self._task = None
|
|
|
|
@property
|
|
def count(self):
|
|
"""The number of await points counted so far."""
|
|
|
|
return self._count
|
|
|
|
@property
|
|
def cancelled_at(self):
|
|
"""Descriptions of the awaitables at which a cancellation was
|
|
delivered."""
|
|
|
|
return tuple(self._cancelled_at)
|
|
|
|
def create_task(self, coro):
|
|
"""Run ``coro`` in a new task, which is the task that cancellations
|
|
are delivered to."""
|
|
|
|
self._task = asyncio.create_task(coro)
|
|
return self._task
|
|
|
|
def __call__(self, awaitable):
|
|
self._count += 1
|
|
if self._count not in self._cancel_at:
|
|
return
|
|
|
|
self._cancelled_at.append(
|
|
getattr(awaitable, "__qualname__", repr(awaitable))
|
|
)
|
|
|
|
# cancel the task the application is running, which is not
|
|
# necessarily the one performing this await; SQLAlchemy runs some
|
|
# operations, such as the close() performed by an async context
|
|
# manager on exit, in a task of their own
|
|
task = self._task if self._task is not None else asyncio.current_task()
|
|
assert task is not None, (
|
|
"cancellation can only be delivered to an operation running "
|
|
"within an asyncio task"
|
|
)
|
|
|
|
# cancel the task rather than raise CancelledError from here: the
|
|
# awaitable then still runs, and asyncio interrupts it at its own
|
|
# first suspension point, as it does in production
|
|
task.cancel()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def await_points(cancel_at=()):
|
|
"""Count the await points performed within the block, cancelling the
|
|
running task at each of ``cancel_at``."""
|
|
|
|
points = AwaitPoints(cancel_at)
|
|
greenlet_cls = concurrency._concurrency_shim._AsyncIoGreenlet
|
|
real_switch = greenlet_cls.switch
|
|
real_throw = greenlet_cls.throw
|
|
|
|
def instrument(greenlet, result):
|
|
# a greenlet that has ended hands back the return value of the
|
|
# function it ran, which nothing awaits
|
|
if greenlet.dead:
|
|
return result
|
|
return _InstrumentedAwaitable(result, points)
|
|
|
|
def switch(self, *arg, **kw):
|
|
return instrument(self, real_switch(self, *arg, **kw))
|
|
|
|
def throw(self, *arg):
|
|
return instrument(self, real_throw(self, *arg))
|
|
|
|
with (
|
|
mock.patch.object(greenlet_cls, "switch", switch),
|
|
mock.patch.object(greenlet_cls, "throw", throw),
|
|
):
|
|
yield points
|
|
|
|
|
|
class ConnectionAccounting:
|
|
"""Records which DBAPI connections a pool created and which it managed
|
|
to close, and how many of the pool's slots were lost.
|
|
|
|
Creation is observed at ``Pool._invoke_creator`` rather than through the
|
|
``connect`` pool event, because the ``connect`` event is itself where
|
|
``dialect.initialize()`` runs; a listener added there does not fire for
|
|
precisely the connections that a failure in that event strands.
|
|
|
|
"""
|
|
|
|
def __init__(self, pool):
|
|
self._pool = pool
|
|
self._real_invoke_creator = pool._invoke_creator
|
|
self._real_close_connection = pool._close_connection
|
|
self._lost_slots = 0
|
|
self._created = set()
|
|
self._closed = set()
|
|
|
|
@contextlib.contextmanager
|
|
def start(self):
|
|
"""Account for the pool's connections within the block."""
|
|
|
|
with (
|
|
mock.patch.object(
|
|
self._pool, "_invoke_creator", self._invoke_creator
|
|
),
|
|
mock.patch.object(
|
|
self._pool, "_close_connection", self._close_connection
|
|
),
|
|
):
|
|
yield self
|
|
|
|
def _invoke_creator(self, rec):
|
|
dbapi_connection = self._real_invoke_creator(rec)
|
|
self._created.add(dbapi_connection)
|
|
return dbapi_connection
|
|
|
|
def _close_connection(self, dbapi_connection, *, terminate=False):
|
|
try:
|
|
self._real_close_connection(dbapi_connection, terminate=terminate)
|
|
except asyncio.CancelledError:
|
|
# AsyncAdapt_terminate force-closes the connection before
|
|
# letting a cancellation of its graceful close propagate
|
|
if terminate and self._pool._dialect.has_terminate:
|
|
self._closed.add(dbapi_connection)
|
|
raise
|
|
self._closed.add(dbapi_connection)
|
|
|
|
@property
|
|
def leaked(self):
|
|
"""Connections that were created but never successfully closed."""
|
|
|
|
return self._created - self._closed
|
|
|
|
def record_lost_slots(self):
|
|
"""Record the pool slots still checked out.
|
|
|
|
To be called once the application's task is done, and before the
|
|
engine is disposed. Every :class:`._ConnectionFairy` the task made
|
|
is gone once garbage is collected, and :func:`._finalize_fairy` has
|
|
returned its record; a record still checked out after that is lost.
|
|
|
|
"""
|
|
gc.collect()
|
|
self._lost_slots = self._pool.checkedout()
|
|
|
|
def get_failure_messages(self):
|
|
"""Return a list of strings describing each leaked connection and
|
|
lost pool slot, empty if there were none."""
|
|
|
|
failures = []
|
|
if self.leaked:
|
|
failures.append(
|
|
f"{len(self.leaked)} of {len(self._created)} connection(s) "
|
|
"were left open with nothing able to close them"
|
|
)
|
|
if self._lost_slots > 0:
|
|
failures.append(
|
|
f"{self._lost_slots} pool slot(s) were left checked out "
|
|
"with nothing able to return them"
|
|
)
|
|
return failures
|
|
|
|
|
|
def connection_accounting(engine):
|
|
"""Account for the DBAPI connections created by ``engine``'s pool.
|
|
|
|
Returns a context manager yielding a :class:`.ConnectionAccounting`.
|
|
|
|
"""
|
|
return ConnectionAccounting(engine.pool).start()
|
|
|
|
|
|
async def run_scenario(
|
|
engine_factory, scenario, cancel_at=(), warm=False, recycle=False
|
|
):
|
|
"""Run ``scenario`` against a fresh engine, optionally cancelling it.
|
|
|
|
``engine_factory`` is a zero-argument callable returning a new
|
|
:class:`_asyncio.AsyncEngine`; ``scenario`` is a coroutine function
|
|
taking that engine. If ``warm`` is True the scenario is run once
|
|
un-cancelled first, so that the await points counted are those of
|
|
checkout, execution and return-to-pool rather than those of
|
|
``connect()``. If ``recycle`` is also True, the warmed pool is then
|
|
invalidated, so that the checkout counted closes the old connection
|
|
before opening its replacement.
|
|
|
|
Returns the :class:`.AwaitPoints` for the run, and a description of
|
|
what the cancellation cost the pool, or ``None`` if it cost nothing.
|
|
|
|
"""
|
|
assert warm or not recycle, "recycle requires 'warm'"
|
|
|
|
engine = engine_factory()
|
|
with connection_accounting(engine) as accounting:
|
|
if warm:
|
|
await scenario(engine)
|
|
if recycle:
|
|
engine.pool._invalidate(None)
|
|
|
|
with await_points(cancel_at) as points:
|
|
try:
|
|
await points.create_task(scenario(engine))
|
|
except BaseException:
|
|
# the cancellation itself, or whatever the dialect raised
|
|
# in response to it. What the scenario raised is not the
|
|
# thing under test; what became of the connection is.
|
|
#
|
|
# if nothing here asked for a cancellation, this is not
|
|
# ours; let it out. asyncio.Runner delivers the first
|
|
# Ctrl-C as a cancellation of the main task rather than as
|
|
# KeyboardInterrupt, so swallowing everything here would
|
|
# otherwise absorb it and carry on sweeping
|
|
if not points.cancelled_at:
|
|
raise
|
|
|
|
# the scenario is over as far as the application can tell
|
|
accounting.record_lost_slots()
|
|
|
|
# the application may now dispose of the engine
|
|
try:
|
|
await engine.dispose()
|
|
except exc.SQLAlchemyError:
|
|
# a pool holding damaged connections may not be able to close
|
|
# them either; that is reported as a leak below
|
|
pass
|
|
|
|
# let anything that was left running in the background finish, so
|
|
# that a connection it returns to the disposed pool is counted as
|
|
# the leak that it is
|
|
others = asyncio.all_tasks() - {asyncio.current_task()}
|
|
if others:
|
|
await asyncio.wait(others, timeout=5)
|
|
gc.collect()
|
|
|
|
return points, "; ".join(accounting.get_failure_messages()) or None
|
|
|
|
|
|
async def sweep(
|
|
engine_factory, scenario, warm=False, recycle=False, cancels=1
|
|
):
|
|
"""Cancel ``scenario`` at each of its await points in turn.
|
|
|
|
With ``cancels=2``, each run that was cancelled at point ``N`` is then
|
|
repeated with a second cancellation at each await point after ``N``,
|
|
which is to say at each point of the cleanup the first cancellation
|
|
set off.
|
|
|
|
Returns a list of strings describing what went wrong, empty if every
|
|
cancellation point was handled cleanly. The strings are meant to be
|
|
read in an assertion failure::
|
|
|
|
eq_(await sweep(engine_factory, scenario), [])
|
|
|
|
"""
|
|
|
|
reported = []
|
|
|
|
async def _sweep(cancel_at):
|
|
points, _ = await run_scenario(
|
|
engine_factory,
|
|
scenario,
|
|
cancel_at=cancel_at,
|
|
warm=warm,
|
|
recycle=recycle,
|
|
)
|
|
total = points.count
|
|
start = cancel_at[-1] + 1 if cancel_at else 1
|
|
for index in range(start, total + 1):
|
|
points, failure = await run_scenario(
|
|
engine_factory,
|
|
scenario,
|
|
cancel_at=cancel_at + (index,),
|
|
warm=warm,
|
|
recycle=recycle,
|
|
)
|
|
if failure is not None:
|
|
indexes = ", ".join(str(i) for i in cancel_at + (index,))
|
|
awaitables = ", ".join(points.cancelled_at)
|
|
reported.append(
|
|
f"cancel at await point(s) {indexes} of {total} "
|
|
f"({awaitables}): {failure}"
|
|
)
|
|
if len(cancel_at) + 1 < cancels:
|
|
await _sweep(cancel_at + (index,))
|
|
|
|
await _sweep(())
|
|
return reported
|