can ignore duplicates in OrderedTaskPreparation

Author: carver

tests/trinity/utils/test_ordered_task_preparation.py

@@ -202,6 +202,21 @@ async def test_no_prereq_tasks():
     assert finished == (2, 3)
 
 
+@pytest.mark.asyncio
+async def test_ignore_duplicates():
+    ti = OrderedTaskPreparation(NoPrerequisites, identity, lambda x: x - 1)
+    ti.set_finished_dependency(1)
+    ti.register_tasks((2, ))
+    # this will ignore the 2 task:
+    ti.register_tasks((2, 3), ignore_duplicates=True)
+    # this will be completely ignored:
+    ti.register_tasks((2, 3), ignore_duplicates=True)
+
+    # with no prerequisites, tasks are *immediately* finished, as long as they are in order
+    finished = await wait(ti.ready_tasks())
+    assert finished == (2, 3)
+
+
 @pytest.mark.asyncio
 async def test_register_out_of_order():
     ti = OrderedTaskPreparation(OnePrereq, identity, lambda x: x - 1, accept_dangling_tasks=True)

trinity/utils/datastructures.py

@@ -503,7 +503,7 @@ def set_finished_dependency(self, finished_task: TTask) -> None:
         self._declared_finished.add(task_id)
         # note that this task is intentionally *not* added to self._unready
 
-    def register_tasks(self, tasks: Tuple[TTask, ...]) -> None:
+    def register_tasks(self, tasks: Tuple[TTask, ...], ignore_duplicates: bool = False) -> None:
         """
         Initiate a task into tracking. By default, each task must be registered
         *after* its dependency has been registered.
@@ -512,23 +512,25 @@ def register_tasks(self, tasks: Tuple[TTask, ...]) -> None:
         initialize this intance with: ``accept_dangling_tasks=True``.
 
         :param tasks: the tasks to register, in iteration order
+        :param ignore_duplicates: any tasks that have already been registered will be ignored,
+            whether ready or not
         """
-        task_meta_info = tuple(
-            (self._prereq_tracker(task), self._id_of(task), self._dependency_of(task))
-            for task in tasks
-        )
-
-        duplicates = tuple(
-            tracker.task for tracker, task_id, _ in task_meta_info
-            if task_id in self._tasks
-        )
+        identified_tasks = tuple((self._id_of(task), task) for task in tasks)
+        duplicates = tuple(task for task_id, task in identified_tasks if task_id in self._tasks)
 
-        if duplicates:
+        if duplicates and not ignore_duplicates:
             raise DuplicateTasks(
                 f"Cannot re-register tasks: {duplicates!r} for completion",
                 duplicates,
             )
 
+        task_meta_info = tuple(
+            (self._prereq_tracker(task), task_id, self._dependency_of(task))
+            for task_id, task in identified_tasks
+            # when ignoring duplicates, must not try to re-add them
+            if task_id not in self._tasks
+        )
+
         for prereq_tracker, task_id, dependency_id in task_meta_info:
             if not self._accept_dangling_tasks and dependency_id not in self._tasks:
                 raise MissingDependency(