post_configure: ensure asyncio tasks created by plugins are awaited

* Await any background tasks started by a plugin before continuing.
  See 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/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}')
 

cylc/flow/scripts/reinstall.py

@@ -75,9 +75,8 @@
 
 from ansimarkup import parse as cparse
 
-from cylc.flow import iter_entry_points
+from cylc.flow.async_util import async_block
 from cylc.flow.exceptions import (
-    PluginError,
     ServiceFileError,
     WorkflowFilesError,
 )
@@ -90,8 +89,8 @@
     OptionSettings,
     ID_MULTI_ARG_DOC
 )
-
 from cylc.flow.pathutil import get_workflow_run_dir
+from cylc.flow.plugins import run_plugins
 from cylc.flow.workflow_files import (
     get_workflow_source_dir,
     load_contact_file,
@@ -198,7 +197,7 @@ async def reinstall_cli(
         if is_terminal() and not opts.skip_interactive:
             # interactive mode - perform dry-run and prompt
             # dry-mode reinstall
-            if not reinstall(
+            if not await reinstall(
                 opts,
                 workflow_id,
                 source,
@@ -231,7 +230,7 @@ async def reinstall_cli(
 
     if usr == 'y':
         # reinstall for real
-        reinstall(opts, workflow_id, source, run_dir, dry_run=False)
+        await reinstall(opts, workflow_id, source, run_dir, dry_run=False)
         print(cparse('<green>Successfully reinstalled.</green>'))
         if print_reload_tip:
             display_cylc_reload_tip(workflow_id)
@@ -245,7 +244,7 @@ async def reinstall_cli(
         return False
 
 
-def reinstall(
+async def reinstall(
     opts: 'Values',
     workflow_id: str,
     src_dir: Path,
@@ -273,7 +272,12 @@ def reinstall(
     # run pre_configure plugins
     if not dry_run:
         # don't run plugins in dry-mode
-        pre_configure(opts, src_dir)
+        for _entry_point, _plugin_result in run_plugins(
+            'cylc.pre_configure',
+            srcdir=src_dir,
+            opts=opts,
+        ):
+            pass
 
     # reinstall from src_dir (will raise WorkflowFilesError on error)
     stdout: str = reinstall_workflow(
@@ -295,7 +299,14 @@ def reinstall(
     # run post_install plugins
     if not dry_run:
         # don't run plugins in dry-mode
-        post_install(opts, src_dir, run_dir)
+        async with async_block():
+            for _entry_point, _plugin_result in run_plugins(
+                'cylc.post_install',
+                srcdir=src_dir,
+                opts=opts,
+                rundir=str(run_dir),
+            ):
+                pass
 
     return True
 
@@ -338,45 +349,6 @@ def format_rsync_out(out: str) -> List[str]:
     return lines
 
 
-def pre_configure(opts: 'Values', src_dir: Path) -> None:
-    """Run pre_configure plugins."""
-    # don't run plugins in dry-mode
-    for entry_point in iter_entry_points(
-        'cylc.pre_configure'
-    ):
-        try:
-            entry_point.load()(srcdir=src_dir, 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
-
-
-def post_install(opts: 'Values', src_dir: Path, run_dir: Path) -> None:
-    """Run post_install plugins."""
-    for entry_point in iter_entry_points(
-        'cylc.post_install'
-    ):
-        try:
-            entry_point.load()(
-                srcdir=src_dir,
-                opts=opts,
-                rundir=str(run_dir)
-            )
-        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
-
-
 def display_rose_warning(src_dir: Path) -> None:
     """Explain why rose installed files are marked as deleted."""
     if (src_dir / 'rose-suite.conf').is_file():

cylc/flow/scripts/validate.py

@@ -133,14 +133,10 @@ def get_option_parser():
 
 @cli_function(get_option_parser)
 def main(parser: COP, options: 'Values', workflow_id: str) -> None:
-    _main(parser, options, workflow_id)
+    asyncio.run(run(parser, options, workflow_id))
 
 
-def _main(parser: COP, options: 'Values', workflow_id: str) -> None:
-    asyncio.run(wrapped_main(parser, options, workflow_id))
-
-
-async def wrapped_main(
+async def run(
     parser: COP, options: 'Values', workflow_id: str
 ) -> None:
     """cylc validate CLI."""