add
This commit is contained in:
macbook_max
@@ -0,0 +1,367 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user