Maintain cleanup context and startup insertion order (#3191) (#3194)

CHANGES/3191.feature

@@ -0,0 +1 @@
+You can mix `app.on_startup` and `app.cleanup_ctx` usage, the insertion order is preserved.

CONTRIBUTORS.txt

@@ -211,6 +211,7 @@ Vasyl Baran
 Victor Kovtun
 Vikas Kawadia
 Viktor Danyliuk
+Vincent Maillol
 Vitalik Verhovodov
 Vitaly Haritonsky
 Vitaly Magerya

aiohttp/web_app.py

@@ -96,8 +96,7 @@ def __init__(self, *,
         self._on_startup = Signal(self)  # type: _AppSignal
         self._on_shutdown = Signal(self)  # type: _AppSignal
         self._on_cleanup = Signal(self)  # type: _AppSignal
-        self._cleanup_ctx = CleanupContext()
-        self._on_startup.append(self._cleanup_ctx._on_startup)
+        self._cleanup_ctx = CleanupContext(self._on_startup)
         self._on_cleanup.append(self._cleanup_ctx._on_cleanup)
         self._client_max_size = client_max_size
 
@@ -413,28 +412,49 @@ def exceptions(self):
 
 class CleanupContext(FrozenList):
 
-    def __init__(self):
-        super().__init__()
-        self._exits = []
+    class CleanupContextItem:
+        """
+        CleanupContext uses this class to wrap an asynchronous generator
+        before adding to itself.
+        """
+
+        def __init__(self, cb):
+            self._cb = cb
+            self._iterator_at_exit = None
 
-    async def _on_startup(self, app):
-        for cb in self:
-            it = cb(app).__aiter__()
+        async def _enter(self, app):
+            it = self._cb(app).__aiter__()
             await it.__anext__()
-            self._exits.append(it)
+            self._iterator_at_exit = it
 
-    async def _on_cleanup(self, app):
-        errors = []
-        for it in reversed(self._exits):
+        async def _exit(self, app):
+            if self._iterator_at_exit is None:
+                return
             try:
-                await it.__anext__()
+                await self._iterator_at_exit.__anext__()
             except StopAsyncIteration:
                 pass
+            else:
+                raise RuntimeError("{!r} has more than one 'yield'"
+                                   .format(self._iterator_at_exit))
+
+    def __init__(self, on_startup):
+        super().__init__()
+        self._on_startup = on_startup
+
+    def append(self, item):
+        cleanup_ctx_item = self.CleanupContextItem(item)
+        super().append(cleanup_ctx_item._exit)
+        self._on_startup.append(cleanup_ctx_item._enter)
+
+    async def _on_cleanup(self, app):
+        errors = []
+        for cb in reversed(self):
+            try:
+                await cb(app)
             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]

docs/web_advanced.rst

@@ -619,6 +619,13 @@ one ``yield``.
 *aiohttp* guarantees that *cleanup code* is called if and only if
 *startup code* was successfully finished.
 
+*statup codes* are called in order of insertion. The order is preserved between
+callbacks in *cleanup* signal and *statup codes*.
+
+*cleanup codes* are called in reverse order of insertion before any
+callback in *cleanup* signal.
+
+
 Asynchronous generators are supported by Python 3.6+, on Python 3.5
 please use `async_generator <https://pypi.org/project/async_generator/>`_
 library.

tests/test_web_app.py

@@ -415,6 +415,43 @@ async def inner(app):
     assert out == ['pre_1', 'post_1']
 
 
+async def test_mixe_cleanup_ctx_on_startup_and_on_cleanup() -> None:
+    app = web.Application()
+    out = []
+
+    def startup(num):
+        async def inner(app):
+            out.append('pre_' + str(num))
+        return inner
+
+    def cleanup(num):
+        async def inner(app):
+            out.append('post_' + str(num))
+        return inner
+
+    def cleanup_ctx(num):
+        @async_generator
+        async def inner(app):
+            out.append('pre_' + str(num))
+            await yield_(None)
+            out.append('post_' + str(num))
+        return inner
+
+    app.on_startup.append(startup(1))
+    app.cleanup_ctx.append(cleanup_ctx(2))
+    app.on_startup.append(startup(3))
+    app.cleanup_ctx.append(cleanup_ctx(4))
+    app.on_startup.append(startup(5))
+
+    app.freeze()
+    await app.startup()
+    assert out == ['pre_1', 'pre_2', 'pre_3', 'pre_4', 'pre_5']
+
+    del out[:]
+    await app.cleanup()
+    assert out == ['post_4', 'post_2']
+
+
 async def test_subapp_chained_config_dict_visibility(aiohttp_client) -> None:
 
     async def main_handler(request):