App cleanup context

CHANGES/2747.feature

@@ -0,0 +1 @@
+Implement application cleanup context (``app.cleanup_ctx`` property).

aiohttp/web.py

@@ -9,7 +9,7 @@
                web_middlewares, web_protocol, web_request, web_response,
                web_runner, web_server, web_urldispatcher, web_ws)
 from .log import access_logger
-from .web_app import Application  # noqa
+from .web_app import *  # noqa
 from .web_exceptions import *  # noqa
 from .web_fileresponse import *  # noqa
 from .web_middlewares import *  # noqa

aiohttp/web_app.py

@@ -16,7 +16,7 @@
 from .web_urldispatcher import PrefixedSubAppResource, UrlDispatcher
 
 
-__all__ = ('Application',)
+__all__ = ('Application', 'CleanupError')
 
 
 class Application(MutableMapping):
@@ -25,7 +25,7 @@ class Application(MutableMapping):
         '_middlewares', '_middlewares_handlers', '_run_middlewares',
         '_state', '_frozen', '_subapps',
         '_on_response_prepare', '_on_startup', '_on_shutdown',
-        '_on_cleanup', '_client_max_size'])
+        '_on_cleanup', '_client_max_size', '_cleanup_ctx'])
 
     def __init__(self, *,
                  logger=web_logger,
@@ -60,6 +60,9 @@ def __init__(self, *,
         self._on_startup = Signal(self)
         self._on_shutdown = Signal(self)
         self._on_cleanup = Signal(self)
+        self._cleanup_ctx = CleanupContext()
+        self._on_startup.append(self._cleanup_ctx._on_startup)
+        self._on_cleanup.append(self._cleanup_ctx._on_cleanup)
         self._client_max_size = client_max_size
 
     def __init_subclass__(cls):
@@ -139,6 +142,7 @@ def freeze(self):
         self._middlewares.freeze()
         self._router.freeze()
         self._on_response_prepare.freeze()
+        self._cleanup_ctx.freeze()
         self._on_startup.freeze()
         self._on_shutdown.freeze()
         self._on_cleanup.freeze()
@@ -213,6 +217,10 @@ def on_shutdown(self):
     def on_cleanup(self):
         return self._on_cleanup
 
+    @property
+    def cleanup_ctx(self):
+        return self._cleanup_ctx
+
     @property
     def router(self):
         return self._router
@@ -327,3 +335,40 @@ def __call__(self):
 
     def __repr__(self):
         return "<Application 0x{:x}>".format(id(self))
+
+
+class CleanupError(RuntimeError):
+    @property
+    def exceptions(self):
+        return self.args[1]
+
+
+class CleanupContext(FrozenList):
+
+    def __init__(self):
+        super().__init__()
+        self._exits = []
+
+    async def _on_startup(self, app):
+        for cb in self:
+            it = cb(app).__aiter__()
+            await it.__anext__()
+            self._exits.append(it)
+
+    async def _on_cleanup(self, app):
+        errors = []
+        for it in reversed(self._exits):
+            try:
+                await it.__anext__()
+            except StopAsyncIteration:
+                pass
+            except Exception as exc:
+                errors.append(exc)
+            else:
+                errors.append(RuntimeError("{!r} has more than one 'yield'"
+                                           .format(it)))
+        if errors:
+            if len(errors) == 1:
+                raise errors[0]
+            else:
+                raise CleanupError("Multiple errors on cleanup stage", errors)

docs/web_advanced.rst

@@ -549,6 +549,53 @@ is asynchronous, it will be awaited before calling next one.
    object creation is subject to change. As long as you are not creating new
    signals, but simply reusing existing ones, you will not be affected.
 
+.. _aiohttp-web-cleanup-ctx:
+
+Cleanup Context
+---------------
+
+Bare :attr:`Application.on_startup` / :attr:`Application.on_cleanup`
+pair still has a pitfall: signals handlers are independent on each other.
+
+E.g. we have ``[create_pg, create_redis]`` in *startup* signal and
+``[dispose_pg, dispose_redis]`` in *cleanup*.
+
+If, for example, ``create_pg(app)`` call fails ``create_redis(app)``
+is not called. But on application cleanup both ``dispose_pg(app)`` and
+``dispose_redis(app)`` are still called: *cleanup signal* has no
+knowledge about startup/cleanup pairs and their execution state.
+
+
+The solution is :attr:`Application.cleanup_ctx` usage::
+
+    async def pg_engine(app):
+        app['pg_engine'] = await create_engine(
+            user='postgre',
+            database='postgre',
+            host='localhost',
+            port=5432,
+            password=''
+        )
+        yield
+        app['pg_engine'].close()
+        await app['pg_engine'].wait_closed()
+
+    app.cleanup_ctx.append(pg_engine)
+
+The attribute is a list of *asynchronous generators*, a code *before*
+``yield`` is an initialization stage (called on *startup*), a code
+*after* ``yield`` is executed on *cleanup*. The generator must have only
+one ``yield``.
+
+*aiohttp* guarantees that *cleanup code* is called if and only if
+*startup code* was successfully finished.
+
+Asynchronous generators are supported by Python 3.6+, on Python 3.5
+please use `async_generator <https://pypi.org/project/async_generator/>`_
+library.
+
+.. versionadded:: 3.1
+
 .. _aiohttp-web-nested-applications:
 
 Nested applications

docs/web_reference.rst

@@ -1274,7 +1274,7 @@ duplicated like one using :meth:`Application.copy`.
           async def on_startup(app):
               pass
 
-      .. seealso:: :ref:`aiohttp-web-background-tasks`.
+      .. seealso:: :ref:`aiohttp-web-signals`.
 
    .. attribute:: on_shutdown
 
@@ -1308,7 +1308,22 @@ duplicated like one using :meth:`Application.copy`.
           async def on_cleanup(app):
               pass
 
-      .. seealso:: :ref:`aiohttp-web-graceful-shutdown` and :attr:`on_shutdown`.
+      .. seealso:: :ref:`aiohttp-web-signals` and :attr:`on_shutdown`.
+
+   .. attribute:: cleanup_ctx
+
+      A list of *context generators* for *startup*/*cleanup* handling.
+
+      Signal handlers should have the following signature::
+
+          async def context(app):
+              # do startup stuff
+              yield
+              # do cleanup
+
+      .. versionadded:: 3.1
+
+      .. seealso:: :ref:`aiohttp-web-cleanup-ctx`.
 
    .. method:: add_subapp(prefix, subapp)
 

requirements/ci-wheel.txt

@@ -1,4 +1,5 @@
 attrs==17.4.0
+async-generator==1.5
 async-timeout==2.0.0
 brotlipy==0.7.0
 cchardet==2.1.1

tests/test_web_app.py

@@ -2,6 +2,7 @@
 from unittest import mock
 
 import pytest
+from async_generator import async_generator, yield_
 
 from aiohttp import log, web
 from aiohttp.abc import AbstractAccessLogger, AbstractRouter
@@ -259,3 +260,132 @@ def test_app_custom_attr():
     app = web.Application()
     with pytest.warns(DeprecationWarning):
         app.custom = None
+
+
+async def test_cleanup_ctx():
+    app = web.Application()
+    out = []
+
+    def f(num):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            await yield_(None)
+            out.append('post_' + str(num))
+        return inner
+
+    app.cleanup_ctx.append(f(1))
+    app.cleanup_ctx.append(f(2))
+    app.freeze()
+    await app.startup()
+    assert out == ['pre_1', 'pre_2']
+    await app.cleanup()
+    assert out == ['pre_1', 'pre_2', 'post_2', 'post_1']
+
+
+async def test_cleanup_ctx_exception_on_startup():
+    app = web.Application()
+    out = []
+
+    exc = Exception('fail')
+
+    def f(num, fail=False):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            if fail:
+                raise exc
+            await yield_(None)
+            out.append('post_' + str(num))
+        return inner
+
+    app.cleanup_ctx.append(f(1))
+    app.cleanup_ctx.append(f(2, True))
+    app.cleanup_ctx.append(f(3))
+    app.freeze()
+    with pytest.raises(Exception) as ctx:
+        await app.startup()
+    assert ctx.value is exc
+    assert out == ['pre_1', 'pre_2']
+    await app.cleanup()
+    assert out == ['pre_1', 'pre_2', 'post_1']
+
+
+async def test_cleanup_ctx_exception_on_cleanup():
+    app = web.Application()
+    out = []
+
+    exc = Exception('fail')
+
+    def f(num, fail=False):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            await yield_(None)
+            out.append('post_' + str(num))
+            if fail:
+                raise exc
+        return inner
+
+    app.cleanup_ctx.append(f(1))
+    app.cleanup_ctx.append(f(2, True))
+    app.cleanup_ctx.append(f(3))
+    app.freeze()
+    await app.startup()
+    assert out == ['pre_1', 'pre_2', 'pre_3']
+    with pytest.raises(Exception) as ctx:
+        await app.cleanup()
+    assert ctx.value is exc
+    assert out == ['pre_1', 'pre_2', 'pre_3', 'post_3', 'post_2', 'post_1']
+
+
+async def test_cleanup_ctx_exception_on_cleanup_multiple():
+    app = web.Application()
+    out = []
+
+    def f(num, fail=False):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            await yield_(None)
+            out.append('post_' + str(num))
+            if fail:
+                raise Exception('fail_' + str(num))
+        return inner
+
+    app.cleanup_ctx.append(f(1))
+    app.cleanup_ctx.append(f(2, True))
+    app.cleanup_ctx.append(f(3, True))
+    app.freeze()
+    await app.startup()
+    assert out == ['pre_1', 'pre_2', 'pre_3']
+    with pytest.raises(web.CleanupError) as ctx:
+        await app.cleanup()
+    exc = ctx.value
+    assert len(exc.exceptions) == 2
+    assert str(exc.exceptions[0]) == 'fail_3'
+    assert str(exc.exceptions[1]) == 'fail_2'
+    assert out == ['pre_1', 'pre_2', 'pre_3', 'post_3', 'post_2', 'post_1']
+
+
+async def test_cleanup_ctx_multiple_yields():
+    app = web.Application()
+    out = []
+
+    def f(num):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            await yield_(None)
+            out.append('post_' + str(num))
+            await yield_(None)
+        return inner
+
+    app.cleanup_ctx.append(f(1))
+    app.freeze()
+    await app.startup()
+    assert out == ['pre_1']
+    with pytest.raises(RuntimeError) as ctx:
+        await app.cleanup()
+    assert "has more than one 'yield'" in str(ctx.value)
+    assert out == ['pre_1', 'post_1']