New TaskQueue tests: bug when adding existing task

Author: carver

setup.py

@@ -44,11 +44,11 @@
         "web3==4.4.1",
     ],
     'test': [
-        "hypothesis==3.44.26",
+        "hypothesis==3.69.5",
         # pinned to <3.7 until async fixtures work again
         # https://github.com/pytest-dev/pytest-asyncio/issues/89
         "pytest>=3.6,<3.7",
-        "pytest-asyncio==0.8.0",
+        "pytest-asyncio==0.9.0",
         "pytest-cov==2.5.1",
         "pytest-watch>=4.1.0,<5",
         "pytest-xdist==1.18.1",

tests/trinity/utils/test_task_queue.py

@@ -1,15 +1,132 @@
 import asyncio
+from asyncio import (
+    Event,
+)
+from contextlib import contextmanager
+import functools
 import pytest
+import random
 
+from cancel_token import CancelToken, OperationCancelled
 from eth_utils import ValidationError
+from hypothesis import (
+    example,
+    given,
+    strategies as st,
+)
 
 from trinity.utils.datastructures import TaskQueue
 
+DEFAULT_TIMEOUT = 0.05
+
 
-async def wait(coro, timeout=0.05):
+async def wait(coro, timeout=DEFAULT_TIMEOUT):
     return await asyncio.wait_for(coro, timeout=timeout)
 
 
+@contextmanager
+def trap_operation_cancelled():
+    try:
+        yield
+    except OperationCancelled:
+        pass
+
+
+def run_in_event_loop(async_func):
+    @functools.wraps(async_func)
+    def wrapped(operations, queue_size, add_size, get_size, event_loop):
+        event_loop.run_until_complete(asyncio.ensure_future(
+            async_func(operations, queue_size, add_size, get_size, event_loop),
+            loop=event_loop,
+        ))
+    return wrapped
+
+
+@given(
+    operations=st.lists(
+        elements=st.tuples(st.integers(min_value=0, max_value=5), st.booleans()),
+        min_size=10,
+        max_size=30,
+    ),
+    queue_size=st.integers(min_value=1, max_value=20),
+    add_size=st.integers(min_value=1, max_value=20),
+    get_size=st.integers(min_value=1, max_value=20),
+)
+@example(
+    # try having two adders alternate a couple times quickly
+    operations=[(0, False), (1, False), (0, False), (1, True), (2, False), (2, False), (2, False)],
+    queue_size=5,
+    add_size=2,
+    get_size=5,
+)
+@run_in_event_loop
+async def test_no_asyncio_exception_leaks(operations, queue_size, add_size, get_size, event_loop):
+    """
+    This could be made much more general, at the cost of simplicity.
+    For now, this mimics real usage enough to hopefully catch the big issues.
+
+    Some examples for more generality:
+
+    - different get sizes on each call
+    - complete varying amounts of tasks at each call
+    """
+
+    async def getter(queue, num_tasks, get_event, complete_event, cancel_token):
+        with trap_operation_cancelled():
+            # wait to run the get
+            await cancel_token.cancellable_wait(get_event.wait())
+
+            batch, tasks = await cancel_token.cancellable_wait(
+                queue.get(num_tasks)
+            )
+            get_event.clear()
+
+            # wait to run the completion
+            await cancel_token.cancellable_wait(complete_event.wait())
+
+            queue.complete(batch, tasks)
+            complete_event.clear()
+
+    async def adder(queue, add_size, add_event, cancel_token):
+        with trap_operation_cancelled():
+            # wait to run the add
+            await cancel_token.cancellable_wait(add_event.wait())
+
+            await cancel_token.cancellable_wait(
+                queue.add(tuple(random.randint(0, 2 ** 32) for _ in range(add_size)))
+            )
+            add_event.clear()
+
+    async def operation_order(operations, events, cancel_token):
+        for operation_id, pause in operations:
+            events[operation_id].set()
+            if pause:
+                await asyncio.sleep(0)
+
+        await asyncio.sleep(0)
+        cancel_token.trigger()
+
+    q = TaskQueue(queue_size)
+    events = tuple(Event() for _ in range(6))
+    add_event, add2_event, get_event, get2_event, complete_event, complete2_event = events
+    cancel_token = CancelToken('end test')
+
+    done, pending = await asyncio.wait([
+        getter(q, get_size, get_event, complete_event, cancel_token),
+        getter(q, get_size, get2_event, complete2_event, cancel_token),
+        adder(q, add_size, add_event, cancel_token),
+        adder(q, add_size, add2_event, cancel_token),
+        operation_order(operations, events, cancel_token),
+    ], return_when=asyncio.FIRST_EXCEPTION)
+
+    for task in done:
+        exc = task.exception()
+        if exc:
+            raise exc
+
+    assert not pending
+
+
 @pytest.mark.asyncio
 async def test_queue_size_reset_after_complete():
     q = TaskQueue(maxsize=2)
@@ -63,7 +180,7 @@ async def test_default_priority_order():
 
 @pytest.mark.asyncio
 async def test_custom_priority_order():
-    q = TaskQueue(maxsize=4, order_fn=lambda x: 0-x)
+    q = TaskQueue(maxsize=4, order_fn=lambda x: 0 - x)
 
     await wait(q.add((2, 1, 3)))
     (batch, tasks) = await wait(q.get())
@@ -108,6 +225,25 @@ async def test_wait_empty_queue():
         assert False, "should not return from get() when nothing is available on queue"
 
 
+@pytest.mark.asyncio
+async def test_cannot_complete_batch_with_wrong_task():
+    q = TaskQueue()
+
+    await wait(q.add((1, 2)))
+
+    batch, tasks = await wait(q.get())
+
+    # cannot complete a valid task with a task it wasn't given
+    with pytest.raises(ValidationError):
+        q.complete(batch, (3, 4))
+
+    # partially invalid completion calls leave the valid task in an incomplete state
+    with pytest.raises(ValidationError):
+        q.complete(batch, (1, 3))
+
+    assert 1 in q
+
+
 @pytest.mark.asyncio
 async def test_cannot_complete_batch_unless_pending():
     q = TaskQueue()
@@ -156,10 +292,9 @@ async def test_two_pending_adds_one_release():
     assert len(tasks) in {0, 1}
 
     if len(tasks) == 1:
-        batch2, tasks2 = await wait(q.get())
+        _, tasks2 = await wait(q.get())
         all_tasks = tuple(sorted(tasks + tasks2))
     elif len(tasks) == 2:
-        batch2 = None
         all_tasks = tasks
 
     assert all_tasks == (0, 3)
@@ -186,12 +321,20 @@ async def test_queue_get_cap(start_tasks, get_max, expected, remainder):
     assert tasks == expected
 
     if remainder:
-        batch2, tasks2 = await wait(q.get())
+        _, tasks2 = await wait(q.get())
         assert tasks2 == remainder
     else:
         try:
-            batch2, tasks2 = await wait(q.get())
+            _, tasks2 = await wait(q.get())
         except asyncio.TimeoutError:
             pass
         else:
             assert False, f"No more tasks to get, but got {tasks2!r}"
+
+
+@pytest.mark.asyncio
+async def test_cannot_readd_same_task():
+    q = TaskQueue()
+    await q.add((1, 2))
+    with pytest.raises(ValidationError):
+        await q.add((2,))

trinity/utils/datastructures.py

@@ -76,6 +76,12 @@ async def add(self, tasks: Tuple[TTask, ...]) -> None:
         if not isinstance(tasks, tuple):
             raise ValidationError(f"must pass a tuple of tasks to add(), but got {tasks!r}")
 
+        already_pending = self._tasks.intersection(tasks)
+        if already_pending:
+            raise ValidationError(
+                f"Can't readd a task to queue. {already_pending!r} are already present"
+            )
+
         # make sure to insert the highest-priority items first, in case queue fills up
         remaining = tuple(sorted((self._order_fn(task), task) for task in tasks))
 
@@ -102,24 +108,15 @@ async def add(self, tasks: Tuple[TTask, ...]) -> None:
                 # There will always be room in _open_queue until _maxsize is reached
                 try:
                     self._open_queue.put_nowait(task)
-                except QueueFull:
+                except QueueFull as exc:
                     task_idx = queueing.index(task)
-                    # TODO remove once this bug is tracked down
-                    import logging; logging.error(
-                        'TaskQueue unsuccessful in adding task %r because qsize=%d, '
-                        'num_tasks=%d, _maxsize=%d, open_slots=%d, num queueing=%d, '
-                        'len(_tasks)=%d, task_idx=%d, queuing=%r',
-                        task,
-                        self._open_queue.qsize(),
-                        num_tasks,
-                        self._maxsize,
-                        open_slots,
-                        len(queueing),
-                        len(self._tasks),
-                        task_idx,
-                        queueing,
+                    qsize = self._open_queue.qsize()
+                    raise QueueFull(
+                        f'TaskQueue unsuccessful in adding task {task[1]!r} because qsize={qsize}, '
+                        f'num_tasks={num_tasks}, maxsize={self._maxsize}, open_slots={open_slots}, '
+                        f'num queueing={len(queueing)}, len(_tasks)={len(self._tasks)}, task_idx='
+                        f'{task_idx}, queuing={queueing}, original msg: {exc}',
                     )
-                    raise
 
             unranked_queued = tuple(task for _rank, task in queueing)
             self._tasks.update(unranked_queued)
@@ -171,9 +168,16 @@ def complete(self, batch_id: int, completed: Tuple[TTask, ...]) -> None:
 
         attempted = self._in_progress.pop(batch_id)
 
-        remaining = set(attempted).difference(completed)
+        unrecognized_tasks = set(completed).difference(attempted)
+        if unrecognized_tasks:
+            self._in_progress[batch_id] = attempted
+            raise ValidationError(
+                f"cannot complete tasks {unrecognized_tasks!r} in this batch, only {attempted!r}"
+            )
+
+        incomplete = set(attempted).difference(completed)
 
-        for task in remaining:
+        for task in incomplete:
             # These tasks are already counted in the total task count, so there will be room
             self._open_queue.put_nowait((self._order_fn(task), task))