From 44c03cc49346b104bbed42d556029f2e78b3b20d Mon Sep 17 00:00:00 2001 From: Daniil Fajnberg Date: Wed, 16 Mar 2022 16:57:03 +0100 Subject: [PATCH] catching exceptions in `_queue_consumer` meta task --- setup.cfg | 4 ++-- src/asyncio_taskpool/pool.py | 11 +++++++++-- tests/test_pool.py | 22 ++++++++++++++-------- 3 files changed, 25 insertions(+), 12 deletions(-) diff --git a/setup.cfg b/setup.cfg index 5689541..34fc2f8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = asyncio-taskpool -version = 0.7.0 +version = 0.7.1 author = Daniil Fajnberg author_email = mail@daniil.fajnberg.de description = Dynamically manage pools of asyncio tasks @@ -9,7 +9,7 @@ long_description_content_type = text/markdown keywords = asyncio, concurrency, tasks, coroutines, asynchronous, server url = https://git.fajnberg.de/daniil/asyncio-taskpool project_urls = - Bug Tracker = https://git.fajnberg.de/daniil/asyncio-taskpool/issues + Bug Tracker = https://github.com/daniil-berg/asyncio-taskpool/issues classifiers = Development Status :: 3 - Alpha Programming Language :: Python :: 3 diff --git a/src/asyncio_taskpool/pool.py b/src/asyncio_taskpool/pool.py index 7abcdee..2198697 100644 --- a/src/asyncio_taskpool/pool.py +++ b/src/asyncio_taskpool/pool.py @@ -775,8 +775,15 @@ class TaskPool(BaseTaskPool): if next_arg is self._QUEUE_END_SENTINEL: # The `_queue_producer()` either reached the last argument or was cancelled. return - await self._start_task(star_function(func, next_arg, arg_stars=arg_stars), group_name=group_name, - ignore_lock=True, end_callback=release_cb, cancel_callback=cancel_callback) + try: + await self._start_task(star_function(func, next_arg, arg_stars=arg_stars), group_name=group_name, + ignore_lock=True, end_callback=release_cb, cancel_callback=cancel_callback) + except Exception as e: + # This means an exception occurred during task **creation**, meaning no task has been created. + # It does not imply an error within the task itself. + log.exception("%s occurred while trying to create task: %s(%s%s)", + str(e.__class__.__name__), func.__name__, '*' * arg_stars, str(next_arg)) + map_semaphore.release() async def _map(self, group_name: str, group_size: int, func: CoroutineFunc, arg_iter: ArgsT, arg_stars: int, end_callback: EndCB = None, cancel_callback: CancelCB = None) -> None: diff --git a/tests/test_pool.py b/tests/test_pool.py index e467d57..7b62cc1 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -603,28 +603,34 @@ class TaskPoolTestCase(CommonTestCase): @patch.object(pool, 'star_function') @patch.object(pool.TaskPool, '_start_task') - @patch.object(pool, 'Semaphore') @patch.object(pool.TaskPool, '_get_map_end_callback') - async def test__queue_consumer(self, mock__get_map_end_callback: MagicMock, mock_semaphore_cls: MagicMock, + @patch.object(pool, 'Semaphore') + async def test__queue_consumer(self, mock_semaphore_cls: MagicMock, mock__get_map_end_callback: MagicMock, mock__start_task: AsyncMock, mock_star_function: MagicMock): - mock__get_map_end_callback.return_value = map_cb = MagicMock() mock_semaphore_cls.return_value = semaphore = Semaphore(3) - mock_star_function.return_value = awaitable = 'totally an awaitable' - arg1, arg2 = 123456789, 'function argument' + mock__get_map_end_callback.return_value = map_cb = MagicMock() + awaitable = 'totally an awaitable' + mock_star_function.side_effect = [awaitable, awaitable, Exception()] + arg1, arg2, bad = 123456789, 'function argument', None mock_q_maxsize = 3 - mock_q = MagicMock(__aenter__=AsyncMock(side_effect=[arg1, arg2, pool.TaskPool._QUEUE_END_SENTINEL]), + mock_q = MagicMock(__aenter__=AsyncMock(side_effect=[arg1, arg2, bad, pool.TaskPool._QUEUE_END_SENTINEL]), __aexit__=AsyncMock(), maxsize=mock_q_maxsize) - group_name, mock_func, stars = 'whatever', MagicMock(), 3 + group_name, mock_func, stars = 'whatever', MagicMock(__name__="mock"), 3 end_cb, cancel_cb = MagicMock(), MagicMock() self.assertIsNone(await self.task_pool._queue_consumer(mock_q, group_name, mock_func, stars, end_cb, cancel_cb)) + # We expect the semaphore to be acquired 3 times, then be released once after the exception occurs, then + # acquired once more when the `_QUEUE_END_SENTINEL` is reached. Since we initialized it with a value of 3, + # at the end of the loop, we expect it be locked. self.assertTrue(semaphore.locked()) + mock_semaphore_cls.assert_called_once_with(mock_q_maxsize) mock__get_map_end_callback.assert_called_once_with(semaphore, actual_end_callback=end_cb) mock__start_task.assert_has_awaits(2 * [ call(awaitable, group_name=group_name, ignore_lock=True, end_callback=map_cb, cancel_callback=cancel_cb) ]) mock_star_function.assert_has_calls([ call(mock_func, arg1, arg_stars=stars), - call(mock_func, arg2, arg_stars=stars) + call(mock_func, arg2, arg_stars=stars), + call(mock_func, bad, arg_stars=stars) ]) @patch.object(pool, 'create_task')