Repository: MatterMiners/tardis

Issue:
Missing task done call in AsyncBulkCall
It seems that `AsyncBulkCall` is missing one `task_done` call. The simplify the logic, we can assume that `size` is set to 1. 

In that case the first task is fetched in 

https://github.com/MatterMiners/tardis/blob/75b60e3a6c91caf4774d6638b2faf26cc8b6b3a7/tardis/utilities/asyncbulkcall.py#L134

the `while` loop in https://github.com/MatterMiners/tardis/blob/75b60e3a6c91caf4774d6638b2faf26cc8b6b3a7/tardis/utilities/asyncbulkcall.py#L136 is skipped. 

and therefore `task_done` in https://github.com/MatterMiners/tardis/blob/75b60e3a6c91caf4774d6638b2faf26cc8b6b3a7/tardis/utilities/asyncbulkcall.py#L148 in never called.

More generally speaking there is always one `task_done` call missing independent of `size`. 

Discussion:

Patch:
diff --git a/.flake8 b/.flake8
index d461ede..d098624 100644
--- a/.flake8
+++ b/.flake8
@@ -1,6 +1,6 @@
 [flake8]
 statistics = True
 max-line-length = 80
-ignore = E501, B008, B011, W503
+ignore = E501, B008, B011, B019, W503
 select = C,E,F,W,B,B9
 exclude = docs,.svn,CVS,.bzr,.hg,.git,__pycache__,.tox,.eggs,*.egg
diff --git a/docs/source/changelog.rst b/docs/source/changelog.rst
index a09a9a3..28cb254 100644
--- a/docs/source/changelog.rst
+++ b/docs/source/changelog.rst
@@ -1,4 +1,4 @@
-.. Created by changelog.py at 2022-03-09, command
+.. Created by changelog.py at 2022-03-07, command
    '/Users/giffler/.cache/pre-commit/repor6pnmwlm/py_env-default/bin/changelog docs/source/changes compile --output=docs/source/changelog.rst'
    based on the format of 'https://keepachangelog.com/'
 
@@ -6,7 +6,7 @@
 CHANGELOG
 #########
 
-[Unreleased] - 2022-03-09
+[Unreleased] - 2022-03-07
 =========================
 
 Added
diff --git a/tardis/resources/dronestates.py b/tardis/resources/dronestates.py
index a4a45ab..5c7f88b 100644
--- a/tardis/resources/dronestates.py
+++ b/tardis/resources/dronestates.py
@@ -143,7 +143,7 @@ class AvailableState(State):
     transition = {
         ResourceStatus.Running: lambda: {
             MachineStatus.Available: lambda: AvailableState(),
-            MachineStatus.NotAvailable: lambda: ShutDownState(),
+            MachineStatus.NotAvailable: lambda: IntegratingState(),
             MachineStatus.Draining: lambda: DrainingState(),
             MachineStatus.Drained: lambda: DisintegrateState(),
         },
diff --git a/tardis/utilities/asyncbulkcall.py b/tardis/utilities/asyncbulkcall.py
index 1e0f993..4fb5192 100644
--- a/tardis/utilities/asyncbulkcall.py
+++ b/tardis/utilities/asyncbulkcall.py
@@ -1,4 +1,4 @@
-from typing import TypeVar, Generic, Iterable, List, Tuple, Optional
+from typing import TypeVar, Generic, Iterable, List, Tuple, Optional, Set
 from typing_extensions import Protocol
 import asyncio
 import time
@@ -73,7 +73,9 @@ class AsyncBulkCall(Generic[T, R]):
         self._delay = delay
         self._concurrency = sys.maxsize if concurrent is None else concurrent
         # task handling dispatch from queue to command execution
-        self._master_worker: Optional[asyncio.Task] = None
+        self._dispatch_task: Optional[asyncio.Task] = None
+        # tasks handling individual command executions
+        self._bulk_tasks: Set[asyncio.Task] = set()
         self._verify_settings()
 
     @cached_property
@@ -100,18 +102,16 @@ class AsyncBulkCall(Generic[T, R]):
     async def __call__(self, __task: T) -> R:
         """Queue a ``task`` for bulk execution and return the result when available"""
         result: "asyncio.Future[R]" = asyncio.get_event_loop().create_future()
-        await self._queue.put((__task, result))
-        self._ensure_worker()
+        # queue item first so that the dispatch task does not finish before
+        self._queue.put_nowait((__task, result))
+        # ensure there is a worker to dispatch items for command execution
+        if self._dispatch_task is None:
+            self._dispatch_task = asyncio.ensure_future(self._bulk_dispatch())
         return await result
 
-    def _ensure_worker(self):
-        """Ensure there is a worker to dispatch tasks for command execution"""
-        if self._master_worker is None:
-            self._master_worker = asyncio.ensure_future(self._bulk_dispatch())
-
     async def _bulk_dispatch(self):
         """Collect tasks into bulks and dispatch them for command execution"""
-        while True:
+        while not self._queue.empty():
             bulk = list(zip(*(await self._get_bulk())))
             if not bulk:
                 continue
@@ -122,9 +122,14 @@ class AsyncBulkCall(Generic[T, R]):
             await self._concurrent.acquire()
             task = asyncio.ensure_future(self._bulk_execute(tuple(tasks), futures))
             task.add_done_callback(lambda _: self._concurrent.release)
+            # track tasks via strong references to avoid them being garbage collected.
+            # see bpo#44665
+            self._bulk_tasks.add(task)
+            task.add_done_callback(lambda _: self._bulk_tasks.discard(task))
             # yield to the event loop so that the `while True` loop does not arbitrarily
             # delay other tasks on the fast paths for `_get_bulk` and `acquire`.
             await asyncio.sleep(0)
+        self._dispatch_task = None
 
     async def _get_bulk(self) -> "List[Tuple[T, asyncio.Future[R]]]":
         """Fetch the next bulk from the internal queue"""
@@ -132,6 +137,7 @@ class AsyncBulkCall(Generic[T, R]):
         # always pull in at least one item asynchronously
         # this avoids stalling for very low delays and efficiently waits for items
         results = [await queue.get()]
+        queue.task_done()
         deadline = time.monotonic() + self._delay
         while len(results) < max_items and time.monotonic() < deadline:
             try:


Test patch:
diff --git a/tests/resources_t/test_dronestates.py b/tests/resources_t/test_dronestates.py
index 755069c..ac62cab 100644
--- a/tests/resources_t/test_dronestates.py
+++ b/tests/resources_t/test_dronestates.py
@@ -197,7 +197,7 @@ class TestDroneStates(TestCase):
     def test_available_state(self):
         matrix = [
             (ResourceStatus.Running, MachineStatus.Available, AvailableState),
-            (ResourceStatus.Running, MachineStatus.NotAvailable, ShutDownState),
+            (ResourceStatus.Running, MachineStatus.NotAvailable, IntegratingState),
             (ResourceStatus.Running, MachineStatus.Draining, DrainingState),
             (ResourceStatus.Running, MachineStatus.Drained, DisintegrateState),
             (ResourceStatus.Booting, MachineStatus.NotAvailable, BootingState),
diff --git a/tests/utilities_t/test_asyncbulkcall.py b/tests/utilities_t/test_asyncbulkcall.py
index 7fbcf4f..6ffd6a0 100644
--- a/tests/utilities_t/test_asyncbulkcall.py
+++ b/tests/utilities_t/test_asyncbulkcall.py
@@ -61,6 +61,22 @@ class TestAsyncBulkCall(TestCase):
         result = run_async(self.execute, execution, count=2048)
         self.assertEqual(result, [(i, i) for i in range(2048)])
 
+    def test_restart(self):
+        """Test that calls work after pausing"""
+        run_async(self.check_restart)
+
+    async def check_restart(self):
+        bunch_size = 4
+        # use large delay to only trigger on size
+        execution = AsyncBulkCall(CallCounter(), size=bunch_size // 2, delay=256)
+        for repeat in range(6):
+            result = await self.execute(execution, bunch_size)
+            self.assertEqual(
+                result, [(i, i // 2 + repeat * 2) for i in range(bunch_size)]
+            )
+            await asyncio.sleep(0.01)  # pause to allow for cleanup
+            assert execution._dispatch_task is None
+
     def test_sanity_checks(self):
         """Test against illegal settings"""
         for wrong_size in (0, -1, 0.5, 2j, "15"):