post_configure: ensure asyncio tasks created by plugins are awaited

* Await any background tasks started by a plugin before continuing.
* Addresses cylc/cylc-rose#274
* Centralise the plugin loading/running/reporting logic.
* Fix the installation test for back-compat-mode.

cylc/flow/async_util.py

@@ -16,6 +16,7 @@
 """Utilities for use with asynchronous code."""
 
 import asyncio
+from contextlib import asynccontextmanager
 from functools import partial, wraps
 from inspect import signature
 import os
@@ -478,3 +479,48 @@ async def _fcn(*args, executor=None, **kwargs):
 
 
 async_listdir = make_async(os.listdir)
+
+
+@asynccontextmanager
+async def async_block():
+    """Ensure all tasks started within the context are awaited when it closes.
+
+    Normally, you would await a task e.g:
+
+    await three()
+
+    If it's possible to await the task, do that, however, this isn't always an
+    option. This interface exists is to help patch over issues where async code
+    (one) calls sync code (two) which calls async code (three) e.g:
+
+    async def one():
+        two()
+
+    def two():
+        # this breaks - event loop is already running
+        asyncio.get_event_loop().run_until_complete(three())
+
+    async def three():
+        await asyncio.sleep(1)
+
+    This code will error because you can't nest asyncio (without nest-asyncio)
+    which means you can schedule tasks the tasks in "two", but you can't await
+    them.
+
+    def two():
+        # this works, but it doesn't wait for three() to complete
+        asyncio.create_task(three())
+
+    This interface allows you to await the tasks
+
+    async def one()
+        async with async_block():
+            two()
+        # any tasks two() started will have been awaited by now
+    """
+    # make a list of all tasks running before we enter the context manager
+    tasks_before = asyncio.all_tasks()
+    # run the user code
+    yield
+    # await any new tasks
+    await asyncio.gather(*(asyncio.all_tasks() - tasks_before))

cylc/flow/parsec/fileparse.py

@@ -36,14 +36,14 @@
 import sys
 import typing as t
 
-from cylc.flow import __version__, iter_entry_points
+from cylc.flow import __version__
 from cylc.flow import LOG
-from cylc.flow.exceptions import PluginError
 from cylc.flow.parsec.exceptions import (
     FileParseError, ParsecError, TemplateVarLanguageClash
 )
-from cylc.flow.parsec.OrderedDict import OrderedDictWithDefaults
 from cylc.flow.parsec.include import inline
+from cylc.flow.parsec.OrderedDict import OrderedDictWithDefaults
+from cylc.flow.plugins import run_plugins
 from cylc.flow.parsec.util import itemstr
 from cylc.flow.templatevars import get_template_vars_from_db
 from cylc.flow.workflow_files import (
@@ -271,23 +271,11 @@ def process_plugins(fpath, opts):
     }
 
     # Run entry point pre_configure items, trying to merge values with each.:
-    for entry_point in iter_entry_points(
-        'cylc.pre_configure'
+    for entry_point, plugin_result in run_plugins(
+        'cylc.pre_configure',
+        srcdir=fpath,
+        opts=opts,
     ):
-        try:
-            # If you want it to work on sourcedirs you need to get the options
-            # to here.
-            plugin_result = entry_point.load()(
-                srcdir=fpath, opts=opts
-            )
-        except Exception as exc:
-            # NOTE: except Exception (purposefully vague)
-            # this is to separate plugin from core Cylc errors
-            raise PluginError(
-                'cylc.pre_configure',
-                entry_point.name,
-                exc
-            ) from None
         for section in ['env', 'template_variables']:
             if section in plugin_result and plugin_result[section] is not None:
                 # Raise error if multiple plugins try to update the same keys.

cylc/flow/plugins.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+# THIS FILE IS PART OF THE CYLC WORKFLOW ENGINE.
+# Copyright (C) NIWA & British Crown (Met Office) & Contributors.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Common functionality related to the loading and calling of plugins."""
+
+from time import time
+
+from cylc.flow import LOG, iter_entry_points
+from cylc.flow.exceptions import PluginError
+import cylc.flow.flags
+
+
+def get_entry_points(plugin_namespace):
+    """Yield all installed plugin entry points which match the given name."""
+    yield from iter_entry_points(plugin_namespace)
+
+
+def run_plugins(plugin_namespace, *args, **kwargs):
+    """Run all installed plugins for the given namespace.
+
+    This runs plugins in series, yielding the results one by one.
+
+    Args:
+        plugin_namespace:
+            The entry point namespace for the plugins to run,
+            e.g. "cylc.post_install".
+        args:
+            Any arguments to call plugins with.
+        kwargs:
+            Any kwargs to call plugins with.
+
+    Yields:
+        (entry_point, plugin_result)
+
+    Warning:
+        Remember to wrap "cylc.post_install" plugins with
+        "cylc.flow.async_util.async_block".
+
+        See https://github.com/cylc/cylc-rose/issues/274
+
+    """
+    for entry_point in get_entry_points(plugin_namespace):
+        try:
+            meth = entry_point.load()
+            start_time = time()
+            plugin_result = meth(*args, **kwargs)
+            LOG.debug(
+                f'ran {entry_point.name} in {time() - start_time:0.05f}s'
+            )
+            yield entry_point, plugin_result
+        except Exception as exc:
+            # NOTE: except Exception (purposefully vague)
+            # this is to separate plugin from core Cylc errors
+            if cylc.flow.flags.verbosity > 1:
+                # raise the full exception in debug mode
+                raise
+            # raise a user-friendly exception
+            raise PluginError(
+                plugin_namespace,
+                entry_point.name,
+                exc
+            ) from None

cylc/flow/scheduler_cli.py

@@ -369,6 +369,12 @@ async def scheduler_cli(
     functionality.
 
     """
+    if options.starttask:
+        options.starttask = upgrade_legacy_ids(
+            *options.starttask,
+            relative=True,
+        )
+
     # Parse workflow name but delay Cylc 7 suite.rc deprecation warning
     # until after the start-up splash is printed.
     # TODO: singleton
@@ -651,14 +657,4 @@ async def _run(scheduler: Scheduler) -> int:
 @cli_function(get_option_parser)
 def play(parser: COP, options: 'Values', id_: str):
     """Implement cylc play."""
-    return _play(parser, options, id_)
-
-
-def _play(parser: COP, options: 'Values', id_: str):
-    """Allows compound scripts to import play, but supply their own COP."""
-    if options.starttask:
-        options.starttask = upgrade_legacy_ids(
-            *options.starttask,
-            relative=True,
-        )
     return asyncio.run(scheduler_cli(options, id_))

cylc/flow/scripts/install.py

@@ -86,14 +86,9 @@
 from pathlib import Path
 from typing import Any, Dict, Optional, Tuple
 
-from cylc.flow.scripts.scan import (
-    get_pipe,
-    _format_plain,
-    FLOW_STATE_SYMBOLS,
-    FLOW_STATE_CMAP
-)
-from cylc.flow import LOG, iter_entry_points
-from cylc.flow.exceptions import PluginError, InputError
+from cylc.flow import LOG
+from cylc.flow.async_util import async_block
+from cylc.flow.exceptions import InputError
 from cylc.flow.loggingutil import CylcLogFormatter, set_timestamps
 from cylc.flow.option_parsers import (
     CylcOptionParser as COP,
@@ -105,12 +100,19 @@
     expand_path,
     get_workflow_run_dir
 )
+from cylc.flow.plugins import run_plugins
 from cylc.flow.install import (
     install_workflow,
     parse_cli_sym_dirs,
     search_install_source_dirs,
     check_deprecation,
 )
+from cylc.flow.scripts.scan import (
+    get_pipe,
+    _format_plain,
+    FLOW_STATE_SYMBOLS,
+    FLOW_STATE_CMAP
+)
 from cylc.flow.terminal import cli_function
 
 
@@ -268,22 +270,20 @@ def main(
     id_: Optional[str] = None
 ) -> None:
     """CLI wrapper."""
-    install_cli(opts, id_)
+    asyncio.run(install_cli(opts, id_))
 
 
-def install_cli(
+async def install_cli(
     opts: 'Values',
     id_: Optional[str] = None
 ) -> Tuple[str, str]:
     """Install workflow and scan for already-running instances."""
-    wf_name, wf_id = install(opts, id_)
-    asyncio.run(
-        scan(wf_name, not opts.no_ping)
-    )
+    wf_name, wf_id = await install(opts, id_)
+    await scan(wf_name, not opts.no_ping)
     return wf_name, wf_id
 
 
-def install(
+async def install(
     opts: 'Values', id_: Optional[str] = None
 ) -> Tuple[str, str]:
     set_timestamps(LOG, opts.log_timestamp and opts.verbosity > 1)
@@ -297,19 +297,12 @@ def install(
     # for compatibility mode:
     check_deprecation(source)
 
-    for entry_point in iter_entry_points(
-        'cylc.pre_configure'
+    for _entry_point, _plugin_result in run_plugins(
+        'cylc.pre_configure',
+        srcdir=source,
+        opts=opts,
     ):
-        try:
-            entry_point.load()(srcdir=source, opts=opts)
-        except Exception as exc:
-            # NOTE: except Exception (purposefully vague)
-            # this is to separate plugin from core Cylc errors
-            raise PluginError(
-                'cylc.pre_configure',
-                entry_point.name,
-                exc
-            ) from None
+        pass
 
     cli_symdirs: Optional[Dict[str, Dict[str, Any]]] = None
     if opts.symlink_dirs == '':
@@ -325,23 +318,14 @@ def install(
         cli_symlink_dirs=cli_symdirs
     )
 
-    for entry_point in iter_entry_points(
-        'cylc.post_install'
-    ):
-        try:
-            entry_point.load()(
-                srcdir=source_dir,
-                opts=opts,
-                rundir=str(rundir)
-            )
-        except Exception as exc:
-            # NOTE: except Exception (purposefully vague)
-            # this is to separate plugin from core Cylc errors
-            raise PluginError(
-                'cylc.post_install',
-                entry_point.name,
-                exc
-            ) from None
+    async with async_block():
+        for _entry_point, _plugin_result in run_plugins(
+            'cylc.post_install',
+            srcdir=source_dir,
+            opts=opts,
+            rundir=str(rundir),
+        ):
+            pass
 
     print(f'INSTALLED {workflow_id} from {source_dir}')