tasks: refactor run_all_tasks to TaskPlan

The run_all_tasks() function was getting large and unwieldy. Upcoming
task features (support for including files) will require even more
complicated logic. To prepare, refactor this function into a class which
contains the shared state, and multiple methods that implement parts of
the process.

One nice benefit here is that we can now do the initial checking
(loading dependencies, checking for conflicts) prior to launching an
instance. This further helps avoid the case where bad task configuration
could cause the instance setup to fail after an otherwise successful
instance launch.

Another nice benefit is that we can print the plan for execution as part
of the --dry-run, which means we can also add a --dry-run option for "yo
task run".

Aside from these two benefits, no other functional change is intended.

Signed-off-by: Stephen Brennan <[email protected]>

Author: brenns10

yo/main.py

@@ -115,10 +115,10 @@
 from yo.ssh import SSH_OPTIONS
 from yo.ssh import wait_for_ssh_access
 from yo.tasks import list_tasks
-from yo.tasks import run_all_tasks
 from yo.tasks import task_get_status
 from yo.tasks import task_join
 from yo.tasks import task_status_to_table
+from yo.tasks import TaskPlan
 from yo.tasks import YoTask
 from yo.util import current_yo_version
 from yo.util import fmt_allow_deny
@@ -1232,11 +1232,21 @@ def add_args(self, parser: argparse.ArgumentParser) -> None:
             action="store_true",
             help="should we wait until the task is finished?",
         )
+        parser.add_argument(
+            "--dry-run",
+            action="store_true",
+            help="just print what would happen",
+        )
 
     def run_for_instance(self, inst: YoInstance) -> None:
-        run_all_tasks(self.c, inst, self.args.task)
+        plan = TaskPlan(self.args.task)
+        plan.prepare()
+        if self.args.dry_run:
+            plan.dry_run_print()
+            return
+        plan.run(self.c, inst)
         if self.args.wait:
-            task_join(self.c, inst, wait_tasks=self.args.task)
+            plan.join(self.c, inst)
             send_notification(
                 self.c,
                 f"Task {self.args.task} complete on instance {inst.name}",
@@ -1842,17 +1852,18 @@ def run(self) -> None:
         # Load tasks before we launch. That way an invalid configuration is
         # detected ASAP, and the user could correct the config before we've
         # actually launched.
-        tasks = set(profile.tasks + self.args.tasks)
+        task_plan = TaskPlan(profile.tasks + self.args.tasks)
+        task_plan.prepare()
 
         self.c.con.log(f"Launching instance [blue]{name}[/blue]")
         if self.args.dry_run:
             self.c.con.log("DRY RUN. Args below:")
             self.c.con.log(create_args)
-            self.c.con.log(f"Will launch tasks: {tasks}")
+            task_plan.dry_run_print()
             return
         inst = self.c.launch_instance(create_args)
 
-        self.standardize_wait(bool(tasks))
+        self.standardize_wait(task_plan.have_tasks())
 
         if not self.args.wait:
             return
@@ -1868,9 +1879,9 @@ def run(self) -> None:
                 self.c.con.log("Maybe you're not connected to VPN?")
                 return
 
-        run_all_tasks(self.c, inst, tasks)
-        if tasks:
-            task_join(self.c, inst)
+        task_plan.run(self.c, inst)
+        if task_plan.have_tasks():
+            task_plan.join(self.c, inst)
 
         send_notification(self.c, f"Instance {inst.name} is ready!")
 

yo/tasks.py

@@ -60,6 +60,8 @@
     os.path.expanduser("~/.oci/yo-tasks"),
     # This should be installed with the package
     os.path.join(os.path.abspath(os.path.dirname(__file__)), "data/yo-tasks"),
+    # Uncomment for testing, if running from a git checkout:
+    # os.path.join(os.path.abspath(os.path.dirname(__file__)), "../test-tasks"),
 ]
 
 
@@ -200,72 +202,108 @@ def _task_run(ctx: "YoCtx", inst: YoInstance, task: YoTask) -> None:
     ssh_into(ip, user, ctx, extra_args=["-q"], cmds=[commands], quiet=True)
 
 
-def run_all_tasks(
-    ctx: "YoCtx", inst: YoInstance, tasks: t.Iterable[t.Union[YoTask, str]]
-) -> None:
-    # The caller may specify tasks as either strings or YoTask instances, for
-    # convenience. Let's get everything into a "name_to_task" dict.
-    name_to_task: t.Dict[str, YoTask] = {}
-    for task in tasks:
-        if isinstance(task, str):
-            name_to_task[task] = YoTask.load(task)
-        else:
-            name_to_task[task.name] = task
-
-    # Tasks may have dependencies. Let's go through every task, and their
-    # dependencies, and load them all. At this point, we're not yet checking
-    # whether there are any circular dependencies: just loading them.
-    tasks_to_load = list(name_to_task.values())
-    for task in tasks_to_load:
-        for name in task.dependencies:
-            if name not in name_to_task:
-                name_to_task[name] = YoTask.load(name)
-                tasks_to_load.append(name_to_task[name])
-
-    # Now we have loaded the complete set of tasks that should run. Some tasks
-    # may appoint themselves as "prerequisites" for another. We need to insert
-    # this dependency relationship so that the script is updated, and so that
-    # the circular dependency detection knows about it. We can also use this
-    # opportunity to detect conflicts.
-    for task in name_to_task.values():
-        for name in task.prereq_for:
-            if name in name_to_task:
-                name_to_task[name].insert_prereq(task.name)
-        for name in task.conflicts:
-            if name in name_to_task:
-                raise YoExc(f"Task {task.name} conflicts with {name}")
-
-    # Now all tasks are loaded, and prerequisites have been marked. Use a
-    # topological sort to verify that no circular dependencies are present. Here
-    # we use a recursive traversal because honestly, if you specify enough tasks
-    # to trigger a recursion error, then I would like to receive that bug
-    # report!
-    name_to_visit: t.Dict[str, int] = collections.defaultdict(int)
-    ordered_tasks: t.List[YoTask] = []
-
-    def visit(task: YoTask) -> None:
-        if name_to_visit[task.name] == 2:
-            # already completed, skip
+class TaskPlan:
+    """
+    A group of tasks to be run altogether on an instance
+
+    This plan encompasses the necessary files to copy over, as well as the order
+    of the tasks to run.
+    """
+
+    name_to_task: t.Dict[str, YoTask]
+    ordered_tasks: t.List[YoTask]
+
+    def _prepare_prereqs_check_conflicts(self) -> None:
+        # Now we have loaded the complete set of tasks that should run. Some tasks
+        # may appoint themselves as "prerequisites" for another. We need to insert
+        # this dependency relationship so that the script is updated, and so that
+        # the circular dependency detection knows about it. We can also use this
+        # opportunity to detect conflicts.
+        for task in self.name_to_task.values():
+            for name in task.prereq_for:
+                if name in self.name_to_task:
+                    self.name_to_task[name].insert_prereq(task.name)
+            for name in task.conflicts:
+                if name in self.name_to_task:
+                    raise YoExc(f"Task {task.name} conflicts with {name}")
+
+    def _create_execution_order(self) -> None:
+        # Now all tasks are loaded, and prerequisites have been marked. Use a
+        # topological sort to verify that no circular dependencies are present. Here
+        # we use a recursive traversal because honestly, if you specify enough tasks
+        # to trigger a recursion error, then I would like to receive that bug
+        # report!
+        name_to_visit: t.Dict[str, int] = collections.defaultdict(int)
+        self.ordered_tasks.clear()  # guard against running twice
+
+        def visit(task: YoTask) -> None:
+            if name_to_visit[task.name] == 2:
+                # already completed, skip
+                return
+            if name_to_visit[task.name] == 1:
+                # currently visiting, not a DAG
+                raise YoExc("Tasks express a circular dependency")
+
+            name_to_visit[task.name] = 1
+            for dep_name in task.dependencies:
+                visit(self.name_to_task[dep_name])
+            name_to_visit[task.name] = 2
+            self.ordered_tasks.append(task)
+
+        for task in self.name_to_task.values():
+            visit(task)
+
+    def __init__(self, tasks: t.Iterable[t.Union[YoTask, str]]) -> None:
+        # The caller may specify tasks as either strings or YoTask instances, for
+        # convenience. Let's get everything into a "name_to_task" dict.
+        self.name_to_task = {}
+        for task in tasks:
+            if isinstance(task, str):
+                self.name_to_task[task] = YoTask.load(task)
+            else:
+                self.name_to_task[task.name] = task
+
+        # Tasks may have dependencies. Let's go through every task, and their
+        # dependencies, and load them all. At this point, we're not yet checking
+        # whether there are any circular dependencies: just loading them.
+        tasks_to_load = list(self.name_to_task.values())
+        for task in tasks_to_load:
+            for name in task.dependencies:
+                if name not in self.name_to_task:
+                    self.name_to_task[name] = YoTask.load(name)
+                    tasks_to_load.append(self.name_to_task[name])
+
+        self.ordered_tasks = []
+
+    def prepare(self) -> None:
+        self._prepare_prereqs_check_conflicts()
+        self._create_execution_order()
+
+    def have_tasks(self) -> bool:
+        return bool(self.ordered_tasks)
+
+    def dry_run_print(self) -> None:
+        if not self.have_tasks():
+            print("No tasks to run.")
             return
-        if name_to_visit[task.name] == 1:
-            # currently visiting, not a DAG
-            raise YoExc("Tasks express a circular dependency")
-
-        name_to_visit[task.name] = 1
-        for dep_name in task.dependencies:
-            visit(name_to_task[dep_name])
-        name_to_visit[task.name] = 2
-        ordered_tasks.append(task)
-
-    for task in name_to_task.values():
-        visit(task)
-
-    # Now ordered_tasks contains the order in which we should launch them.  This
-    # is just a nice-to-have: even if we launched them out of order, the
-    # DEPENDS_ON function would enforce the order of execution. Regardless,
-    # let's start the tasks.
-    for task in ordered_tasks:
-        _task_run(ctx, inst, task)
+        print("Would start tasks in the following order:")
+        for task in self.ordered_tasks:
+            deps = ""
+            if task.dependencies:
+                deps = f" (depends on: {', '.join(task.dependencies)})"
+            print(f" 1. {task.name}{deps}")
+
+    def run(self, ctx: "YoCtx", inst: YoInstance) -> None:
+        # Now ordered_tasks contains the order in which we should launch them.  This
+        # is just a nice-to-have: even if we launched them out of order, the
+        # DEPENDS_ON function would enforce the order of execution. Regardless,
+        # let's start the tasks.
+        for task in self.ordered_tasks:
+            _task_run(ctx, inst, task)
+
+    def join(self, ctx: "YoCtx", inst: YoInstance) -> None:
+        wait_tasks = [t.name for t in self.ordered_tasks]
+        task_join(ctx, inst, wait_tasks)
 
 
 def task_get_status(